From efb2bf0928597078190bfea67a54a217ccbe52d8 Mon Sep 17 00:00:00 2001
From: Itay Alroy <ialroy@nvidia.com>
Date: Tue, 4 Nov 2025 17:10:37 +0200
Subject: [PATCH] EXAMPLES: Introduce NIXL-EP example

Add an example implementation of expert-parallel
dispatch and combine operations using the NIXL
device API.

Co-authored-by: Roey Azran <roeya@nvidia.com>
Co-authored-by: Micha Dery <mdery@nvidia.com>
Co-authored-by: Michal Shalev <mshalev@nvidia.com>

Signed-off-by: Itay Alroy <ialroy@nvidia.com>
---
 .github/workflows/clang-format.yml            |    2 +-
 .github/workflows/copyright-check.sh          |    6 +-
 .pre-commit-config.yaml                       |    1 +
 CONTRIBUTING.md                               |    4 +
 LICENSE                                       |    9 +
 examples/device/ep/LICENSE-DeepEP             |   21 +
 examples/device/ep/README.md                  |  102 ++
 examples/device/ep/csrc/config.hpp            |  130 ++
 examples/device/ep/csrc/event.hpp             |   65 +
 examples/device/ep/csrc/kernels/api.cuh       |  140 ++
 examples/device/ep/csrc/kernels/configs.cuh   |   74 ++
 examples/device/ep/csrc/kernels/exception.cuh |   73 ++
 examples/device/ep/csrc/kernels/launch.cuh    |   86 ++
 examples/device/ep/csrc/kernels/nixl_ep.cu    | 1154 +++++++++++++++++
 examples/device/ep/csrc/kernels/utils.cuh     |  467 +++++++
 examples/device/ep/csrc/nixl_ep.cpp           | 1055 +++++++++++++++
 examples/device/ep/csrc/nixl_ep.hpp           |  221 ++++
 examples/device/ep/meson.build                |  176 +++
 examples/device/ep/nixl_ep/__init__.py        |   24 +
 examples/device/ep/nixl_ep/buffer.py          |  378 ++++++
 examples/device/ep/nixl_ep/utils.py           |   81 ++
 examples/device/ep/scripts/reset_etcd.sh      |   23 +
 examples/device/ep/tests/elastic/README.md    |   34 +
 .../ep/tests/elastic/double_expansion.json    |    5 +
 examples/device/ep/tests/elastic/elastic.py   |  399 ++++++
 .../tests/elastic/expansion_contraction.json  |    6 +
 .../device/ep/tests/elastic/no_expansion.json |    3 +
 examples/device/ep/tests/elastic/plan.py      |   95 ++
 .../device/ep/tests/elastic/rank_server.py    |  141 ++
 .../ep/tests/elastic/single_expansion.json    |    4 +
 .../elastic/single_expansion_16_ranks.json    |    4 +
 .../device/ep/tests/elastic/single_rank.json  |    3 +
 examples/device/ep/tests/utils.py             |  247 ++++
 examples/device/meson.build                   |   18 +
 examples/meson.build                          |    1 +
 meson.build                                   |   10 +-
 meson_options.txt                             |    1 +
 src/core/meson.build                          |    3 +-
 src/plugins/ucx/meson.build                   |    3 +-
 src/utils/ucx/meson.build                     |    4 +-
 40 files changed, 5264 insertions(+), 9 deletions(-)
 create mode 100644 examples/device/ep/LICENSE-DeepEP
 create mode 100644 examples/device/ep/README.md
 create mode 100644 examples/device/ep/csrc/config.hpp
 create mode 100644 examples/device/ep/csrc/event.hpp
 create mode 100644 examples/device/ep/csrc/kernels/api.cuh
 create mode 100644 examples/device/ep/csrc/kernels/configs.cuh
 create mode 100644 examples/device/ep/csrc/kernels/exception.cuh
 create mode 100644 examples/device/ep/csrc/kernels/launch.cuh
 create mode 100644 examples/device/ep/csrc/kernels/nixl_ep.cu
 create mode 100644 examples/device/ep/csrc/kernels/utils.cuh
 create mode 100644 examples/device/ep/csrc/nixl_ep.cpp
 create mode 100644 examples/device/ep/csrc/nixl_ep.hpp
 create mode 100644 examples/device/ep/meson.build
 create mode 100644 examples/device/ep/nixl_ep/__init__.py
 create mode 100644 examples/device/ep/nixl_ep/buffer.py
 create mode 100644 examples/device/ep/nixl_ep/utils.py
 create mode 100755 examples/device/ep/scripts/reset_etcd.sh
 create mode 100644 examples/device/ep/tests/elastic/README.md
 create mode 100644 examples/device/ep/tests/elastic/double_expansion.json
 create mode 100644 examples/device/ep/tests/elastic/elastic.py
 create mode 100644 examples/device/ep/tests/elastic/expansion_contraction.json
 create mode 100644 examples/device/ep/tests/elastic/no_expansion.json
 create mode 100644 examples/device/ep/tests/elastic/plan.py
 create mode 100644 examples/device/ep/tests/elastic/rank_server.py
 create mode 100644 examples/device/ep/tests/elastic/single_expansion.json
 create mode 100644 examples/device/ep/tests/elastic/single_expansion_16_ranks.json
 create mode 100644 examples/device/ep/tests/elastic/single_rank.json
 create mode 100644 examples/device/ep/tests/utils.py
 create mode 100644 examples/device/meson.build

diff --git a/.github/workflows/clang-format.yml b/.github/workflows/clang-format.yml
index e1a0da0688..30c98a1ab5 100644
--- a/.github/workflows/clang-format.yml
+++ b/.github/workflows/clang-format.yml
@@ -18,7 +18,7 @@ jobs:
         FILE_PATTERN='\.(cpp|h|cc|c|cxx|hpp|cu|cuh)$'
 
         echo "### Modified C/C++ files:"
-        FILES=$(git diff --name-only HEAD^1 HEAD | grep -E "$FILE_PATTERN") || true
+        FILES=$(git diff --name-only HEAD^1 HEAD -- . ':(exclude)examples/device/ep' | grep -E "$FILE_PATTERN") || true
         [ -z "$FILES" ] && echo "(none)" || echo "$FILES"
 
         echo "### clang format errors:"
diff --git a/.github/workflows/copyright-check.sh b/.github/workflows/copyright-check.sh
index 1fe931983d..9669c8e819 100755
--- a/.github/workflows/copyright-check.sh
+++ b/.github/workflows/copyright-check.sh
@@ -21,7 +21,7 @@ for f in $(git ls-files); do
     *.png|*.jpg|*.jpeg|*.gif|*.ico|*.zip|*.rst|*.pyc|*.lock|*.md|*.svg|*.wrap|*.in|*.json|*.template|*.gitignore|*.python-version|*py.typed)
       continue
       ;;
-    CODEOWNERS|LICENSE|Doxyfile|.clang-format|.clang-tidy|.codespellrc)
+    CODEOWNERS|*LICENSE*|Doxyfile|.clang-format|.clang-tidy|.codespellrc)
       continue
       ;;
   esac
@@ -39,7 +39,7 @@ for f in $(git ls-files); do
 
   # Extract copyright years (handles YYYY or YYYY-YYYY)
   copyright_years=$(echo "$header" | \
-    grep -Eo 'Copyright \(c\) [0-9]{4}(-[0-9]{4})?' | \
+    grep NVIDIA | grep -Eo 'Copyright \(c\) [0-9]{4}(-[0-9]{4})?' | \
     sed -E 's/.* ([0-9]{4})(-[0-9]{4})?/\1\2/' || true)
 
   if [[ -z "$copyright_years" ]]; then
@@ -57,7 +57,7 @@ for f in $(git ls-files); do
   fi
 
   # License line must exist
-  if ! echo "$header" | grep -Eq '^[[:space:]]*(#|//|\*|/\*|<!--)[[:space:]]*SPDX-License-Identifier:[[:space:]]*Apache-2\.0'; then
+  if ! echo "$header" | grep -Eq '^[[:space:]]*(#|//|\*|/\*|<!--)[[:space:]]*SPDX-License-Identifier:.*Apache-2\.0'; then
     failures+=("$f (missing license)")
     continue
   fi
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 02aa92f993..7965f302ea 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -32,6 +32,7 @@ repos:
     hooks:
       - id: black
         types_or: [python, cython]
+        exclude: ^examples/device/ep/
 
   - repo: https://github.com/PyCQA/flake8
     rev: 7.1.2
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 997c3d8c67..bbafe4ea80 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -167,6 +167,10 @@ All code must adhere to the style guide in `docs/CodeStyle.md` and be formatted
 
 ## Contributing Process
 
+Contributions to the code under `./examples/device/ep` (which is derived
+from DeepEP and licensed under the MIT License) must be licensed under Apache
+2.0.
+
 Contributions that fix documentation errors or that make small changes
 to existing code can be contributed directly by following the rules
 below and submitting an appropriate PR.
diff --git a/LICENSE b/LICENSE
index 261eeb9e9f..0bd68627fc 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,3 +1,12 @@
+   NOTICE: The code under ./examples/device/ep is derived from DeepEP
+   (originally developed by DeepSeek) and has been modified for use in this
+   project. The original code was obtained from:
+   https://github.com/deepseek-ai/DeepEP (commit f0d34aabcb7bdcb3a05d022e7d11b3bf4ccf8ee8)
+   This code is licensed under the MIT License. The full text of the MIT
+   License can be found in ./examples/device/ep/LICENSE-DeepEP.
+   The rest of this codebase is licensed under the Apache License 2.0 as
+   described below.
+
                                  Apache License
                            Version 2.0, January 2004
                         http://www.apache.org/licenses/
diff --git a/examples/device/ep/LICENSE-DeepEP b/examples/device/ep/LICENSE-DeepEP
new file mode 100644
index 0000000000..5c48bdc9f6
--- /dev/null
+++ b/examples/device/ep/LICENSE-DeepEP
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 DeepSeek
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/examples/device/ep/README.md b/examples/device/ep/README.md
new file mode 100644
index 0000000000..21e0c8c83b
--- /dev/null
+++ b/examples/device/ep/README.md
@@ -0,0 +1,102 @@
+# NIXL EP: Expert-Parallel Communication Example
+
+## Overview
+
+NIXL EP is a complete example implementation of expert-parallel communication for Mixture of Experts (MoE) models built on top of [NIXL](https://github.com/ai-dynamo/nixl)'s device API. It provides elastic scaling capabilities, enabling dynamic addition and removal of processes (ranks) during runtime without disrupting existing connections, and leverages NIXL's RDMA and NVLink support for optimal performance.
+
+## Features
+- **Dispatch and Combine support**: Supports dispatch and combine operations for MoE inference
+- **RDMA and NVLink support**: Utilizes NIXL's abstractions to support both RDMA and NVLink transports for optimal performance
+- **Elastic Scaling**: Dynamically add or remove ranks during runtime
+
+## Buffer Initialization
+
+NIXL EP provides a flexible buffer initialization pattern that supports dynamic rank management:
+
+```python
+import nixl_ep
+
+# Initialize buffer with dynamic rank support
+buffer = nixl_ep.Buffer(rank, explicitly_destroy=True)
+buffer.update_memory_buffers(num_ranks, num_experts_per_rank, rdma_bytes)
+buffer.connect_ranks(initial_ranks)
+
+# Dispatch & Combine calls
+buffer.dispatch(...)
+buffer.combine(...)
+
+# Later: Connect new ranks dynamically
+buffer.connect_ranks(ranks)
+
+# Dispatch & Combine calls
+buffer.dispatch(...)
+buffer.combine(...)
+
+# Disconnect ranks when scaling down
+buffer.disconnect_ranks(ranks)
+```
+
+## Key APIs
+
+- `Buffer(rank_id, nvlink_backend, explicitly_destroy)`: Initialize the NIXL communication buffer
+- `update_memory_buffers(num_ranks, num_experts_per_rank, num_rdma_bytes)`: Prepare buffers for up to `num_ranks` ranks and `num_experts_per_rank` experts
+- `connect_ranks(remote_ranks)`: Establish NIXL connections to new peers (can be called multiple times)
+- `disconnect_ranks(remote_ranks)`: Clean up connections to departing peers
+
+## Testing
+
+The elastic test suite in `tests/elastic/` validates dynamic scaling capabilities:
+- Plan files define scaling phases (representing an orchestrator)
+- Tests validate correctness and measure bandwidth between scaling phases
+
+**Example Plan** (`expansion_contraction.json`):
+```json
+[
+  [0, 1, 2, 3],
+  [0, 1, 2, 3, 4, 5, 6, 7],
+  [0, 1, 2, 3, 4, 5]
+]
+```
+This plan defines three phases:
+- **Phase 0**: Initial state with ranks 0-3
+- **Phase 1**: Ranks 4-7 are added dynamically (launched independently from initial ranks)
+- **Phase 2**: Ranks 6-7 are removed dynamically
+
+## Getting Started
+
+#### Build NIXL with NIXL EP:
+
+First, configure the pkg-config paths (only needed when dependencies are installed to non-default paths)
+
+```bash
+export PKG_CONFIG_PATH=<path to rdma-core install>/lib/pkgconfig:$PKG_CONFIG_PATH
+export PKG_CONFIG_PATH=<path to UCX install>/lib/pkgconfig:$PKG_CONFIG_PATH
+export PKG_CONFIG_PATH=<path to DOCA install>/lib/x86_64-linux-gnu/pkgconfig:$PKG_CONFIG_PATH
+```
+
+Then, configure the NIXL plugin directory so it can find UCX plugin, and set the LD_LIBRARY_PATH so UCX can find rdma-core:
+```bash
+export NIXL_PLUGIN_DIR=<path to NIXL install directory>/lib/x86_64-linux-gnu/plugins
+export LD_LIBRARY_PATH=<path to rdma-core install>/lib:$LD_LIBRARY_PATH
+```
+
+Build and install:
+
+```bash
+meson setup build \
+    -Ducx_path=<path to UCX install> \
+    -Dprefix=<path to NIXL install directory> \
+    -Dbuildtype=release \
+    -Dbuild_nixl_ep=true
+
+cd build
+ninja install
+```
+
+
+Finally, configure PYTHONPATH to use NIXL EP:
+```bash
+export PYTHONPATH=<path to NIXL build directory>/examples/device/ep
+```
+
+Refer to [tests/elastic/README.md](tests/elastic/README.md) for detailed instructions on how to run the elastic test suite.
diff --git a/examples/device/ep/csrc/config.hpp b/examples/device/ep/csrc/config.hpp
new file mode 100644
index 0000000000..412e6e2f66
--- /dev/null
+++ b/examples/device/ep/csrc/config.hpp
@@ -0,0 +1,130 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * This file incorporates material from the DeepSeek project, licensed under the MIT License.
+ * The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+ *
+ * SPDX-License-Identifier: MIT AND Apache-2.0
+ *
+ * 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 "kernels/api.cuh"
+#include "kernels/exception.cuh"
+
+namespace nixl_ep {
+
+template <typename dtype_t>
+dtype_t ceil_div(dtype_t a, dtype_t b) {
+    return (a + b - 1) / b;
+}
+
+template <typename dtype_t>
+dtype_t align(dtype_t a, dtype_t b) {
+    return ceil_div<dtype_t>(a, b) * b;
+}
+
+struct EPBuffer {
+    int num_clean_int = 0;
+
+    void* dispatch_rdma_send_buffer = nullptr;
+    void* dispatch_rdma_recv_data_buffer = nullptr;
+    int* dispatch_rdma_recv_count_buffer = nullptr;
+
+    void* combine_rdma_send_buffer = nullptr;
+    void* combine_rdma_recv_data_buffer = nullptr;
+    int* combine_rdma_recv_flag_buffer = nullptr;
+
+    void* combine_rdma_send_buffer_data_start = nullptr;
+    size_t num_bytes_per_combine_msg = 0;
+
+    std::pair<int*, int> clean_meta() {
+        EP_HOST_ASSERT(dispatch_rdma_recv_count_buffer == combine_rdma_recv_flag_buffer);
+        return {dispatch_rdma_recv_count_buffer, num_clean_int};
+    }
+};
+
+struct EPLayout {
+    size_t total_bytes = 0;
+    EPBuffer buffers[2];
+
+    template <typename out_ptr_t = void*, typename count_ptr_t = uint8_t*, typename in_ptr_t = void*>
+    out_ptr_t advance(const in_ptr_t& ptr, size_t count) {
+        return reinterpret_cast<out_ptr_t>(reinterpret_cast<count_ptr_t>(ptr) + count);
+    }
+
+    EPLayout(void* rdma_buffer, int num_max_dispatch_tokens_per_rank, int hidden, int num_ranks, int num_experts) {
+        const int num_scales = hidden / 128;
+
+        // Dispatch and combine layout:
+        //  - 2 symmetric odd/even send buffer
+        //  - 2 symmetric odd/even receive buffers
+        //  - 2 symmetric odd/even signaling buffers
+
+        // Message sizes
+        // NOTES: you should add a control `int4` for combine messages if you want to do data transformation
+        // NOTES: `num_scales * sizeof(nv_bfloat162)` means the per-128-channel min/max
+        EP_HOST_ASSERT(num_scales * sizeof(float) <= hidden);
+        size_t num_bytes_per_dispatch_msg = sizeof(int4) + std::max(hidden * sizeof(nv_bfloat16), hidden + num_scales * sizeof(float));
+        size_t num_bytes_per_combine_msg = num_scales * sizeof(nv_bfloat162) + hidden * sizeof(nv_bfloat16);
+
+        // Send buffer
+        size_t dispatch_send_buffer_bytes = num_max_dispatch_tokens_per_rank * num_bytes_per_dispatch_msg;
+        size_t combine_send_buffer_bytes = num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_combine_msg;
+        size_t send_buffer_bytes = std::max(dispatch_send_buffer_bytes, combine_send_buffer_bytes);
+        EP_HOST_ASSERT(send_buffer_bytes % sizeof(int4) == 0);
+        total_bytes += send_buffer_bytes * 2;
+
+        // Symmetric receive buffers
+        // TODO: optimize memory usages
+        size_t dispatch_recv_data_buffer_bytes = num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_dispatch_msg;
+        size_t combine_recv_buffer_bytes = num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_combine_msg;
+        size_t recv_buffer_bytes = std::max(dispatch_recv_data_buffer_bytes, combine_recv_buffer_bytes);
+        EP_HOST_ASSERT(recv_buffer_bytes % sizeof(int4) == 0);
+        total_bytes += recv_buffer_bytes * 2;
+
+        // Symmetric signaling buffers
+        size_t dispatch_recv_count_buffer_bytes = num_experts * sizeof(int);
+        size_t combine_recv_flag_buffer_bytes = dispatch_recv_count_buffer_bytes;
+        size_t signaling_buffer_bytes = std::max(dispatch_recv_count_buffer_bytes, combine_recv_flag_buffer_bytes);
+        size_t signaling_buffer_bytes_aligned = align<size_t>(signaling_buffer_bytes, 128);
+        total_bytes += signaling_buffer_bytes_aligned * 2;
+
+        // Assign pointers
+        // NOTES: we still leave some space for distinguishing dispatch/combine buffer,
+        // so you may see some parameters are duplicated
+        for (int i = 0; i < 2; ++ i) {
+            buffers[i] = {
+                static_cast<int>(signaling_buffer_bytes / sizeof(int)),
+                advance(rdma_buffer, signaling_buffer_bytes_aligned * 2 + send_buffer_bytes * i),
+                advance(rdma_buffer, signaling_buffer_bytes_aligned * 2 + send_buffer_bytes * 2 + recv_buffer_bytes * i),
+                advance<int*>(rdma_buffer, signaling_buffer_bytes_aligned * i),
+                advance(rdma_buffer, signaling_buffer_bytes_aligned * 2 + send_buffer_bytes * i),
+                advance(rdma_buffer, signaling_buffer_bytes_aligned * 2 + send_buffer_bytes * 2 + recv_buffer_bytes * i),
+                advance<int*>(rdma_buffer, signaling_buffer_bytes_aligned * i),
+                advance(rdma_buffer, signaling_buffer_bytes_aligned * 2 + send_buffer_bytes * i),
+                num_bytes_per_combine_msg
+            };
+        }
+    }
+};
+
+size_t get_rdma_size_hint(int num_max_dispatch_tokens_per_rank, int hidden, int num_ranks, int num_experts) {
+    auto num_bytes = EPLayout(nullptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts).total_bytes;
+    return ((num_bytes + NUM_BUFFER_ALIGNMENT_BYTES) / NUM_BUFFER_ALIGNMENT_BYTES) * NUM_BUFFER_ALIGNMENT_BYTES;
+}
+
+} // namespace nixl_ep
diff --git a/examples/device/ep/csrc/event.hpp b/examples/device/ep/csrc/event.hpp
new file mode 100644
index 0000000000..1bf4762450
--- /dev/null
+++ b/examples/device/ep/csrc/event.hpp
@@ -0,0 +1,65 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * This file incorporates material from the DeepSeek project, licensed under the MIT License.
+ * The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+ *
+ * SPDX-License-Identifier: MIT AND Apache-2.0
+ *
+ * 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 <ATen/cuda/CUDAContext.h>
+#include <memory>
+
+#include "kernels/exception.cuh"
+
+namespace nixl_ep {
+
+struct EventHandle {
+    std::shared_ptr<torch::Event> event;
+
+    EventHandle() {
+        event = std::make_shared<torch::Event>(torch::kCUDA);
+        event->record(at::cuda::getCurrentCUDAStream());
+    }
+
+    explicit EventHandle(const at::cuda::CUDAStream& stream) {
+        event = std::make_shared<torch::Event>(torch::kCUDA);
+        event->record(stream);
+    }
+
+    EventHandle(const EventHandle& other) = default;
+
+    void current_stream_wait() const {
+        at::cuda::getCurrentCUDAStream().unwrap().wait(*event);
+    }
+};
+
+torch::Event create_event(const at::cuda::CUDAStream &s) {
+    auto event = torch::Event(torch::kCUDA);
+    event.record(s);
+    return event;
+}
+
+void stream_wait(const at::cuda::CUDAStream& s_0, const at::cuda::CUDAStream& s_1) {
+    EP_HOST_ASSERT(s_0.id() != s_1.id());
+    s_0.unwrap().wait(create_event(s_1));
+}
+
+void stream_wait(const at::cuda::CUDAStream& s, const EventHandle& event) {
+    s.unwrap().wait(*event.event);
+}
+
+} // namespace nixl_ep
diff --git a/examples/device/ep/csrc/kernels/api.cuh b/examples/device/ep/csrc/kernels/api.cuh
new file mode 100644
index 0000000000..b3ed93795c
--- /dev/null
+++ b/examples/device/ep/csrc/kernels/api.cuh
@@ -0,0 +1,140 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * This file incorporates material from the DeepSeek project, licensed under the MIT License.
+ * The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+ *
+ * SPDX-License-Identifier: MIT AND Apache-2.0
+ *
+ * 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 <vector>
+#include <iostream>
+#include "nixl_types.h"
+#include "exception.cuh"
+
+namespace nixl_ep {
+
+// EP kernels
+namespace ep_kernels {
+struct gpu_nixl_ctx {
+    uint64_t *local_counters; // [local_expert_id][src_rank]
+    uint64_t *clean_counters; // Counters to be cleaned for the next iteration
+    nixlGpuXferReqH *remote_counter_reqs; // [local_expert_id,dest_rank]
+    nixlGpuXferReqH *batch_reqs; // [local_expert_id,dest_rank]
+    int *local_barrier_buffer; // [src_rank]
+    nixlGpuXferReqH *remote_barrier_reqs; // [local_expert_idx,dest_rank]
+    void **rdma_p2p_ptrs; // [num_ranks]
+    uint64_t **counters_p2p_ptrs; // [num_ranks]
+    void *rdma_buffer_ptr;
+    int num_local_experts;
+    int num_ranks;
+    int rank;
+
+    /* Double buffering considerations are handled by the caller */
+    __device__ inline void *rdma_p2p_ptr_get(uint64_t ptr, int dst_rank) {
+        if (rdma_p2p_ptrs[dst_rank] == nullptr)
+            return nullptr;
+
+        return (void *)(reinterpret_cast<uint64_t>(rdma_p2p_ptrs[dst_rank]) + batch_offset_get(ptr));
+    }
+
+    /* Double buffering considerations are handled by nixl_ctx */
+    __device__ inline uint64_t *counter_p2p_ptr_get(int local_expert_idx, int dst_rank) {
+        if (counters_p2p_ptrs[dst_rank] == nullptr)
+            return nullptr;
+
+        return counters_p2p_ptrs[dst_rank] + (local_expert_idx * num_ranks + rank);
+    }
+
+    __device__ inline uint64_t *local_counter_get(int local_expert_idx, int src_rank) {
+        return &local_counters[local_expert_idx * num_ranks + src_rank];
+    }
+
+    __device__ inline nixlGpuXferReqH remote_counter_get(int local_expert_idx, int dest_rank) {
+        return remote_counter_reqs[local_expert_idx * num_ranks + dest_rank];
+    }
+
+    __device__ inline nixlGpuXferReqH remote_barrier_get(int local_expert_idx, int dest_rank) {
+        return remote_barrier_reqs[local_expert_idx * num_ranks + dest_rank];
+    }
+
+    __device__ inline int* local_barrier_buffer_get(int src_rank) {
+        return &local_barrier_buffer[src_rank];
+    }
+
+    __device__ inline nixlGpuXferReqH batch_get(int local_expert_idx, int dest_rank) {
+        return batch_reqs[local_expert_idx * num_ranks + dest_rank];
+    }
+
+    __device__ inline size_t batch_offset_get(uint64_t ptr) {
+        return ptr - reinterpret_cast<uint64_t>(rdma_buffer_ptr);
+    }
+
+    __device__ inline void clean_counters_warp(int lane_id) {
+#ifdef __CUDACC__
+        #pragma unroll
+#endif
+        for (int i = lane_id; i < num_ranks * num_local_experts; i += 32)
+            clean_counters[i] = 0;
+    }
+};
+
+void clean_buffer(int* clean_0, int num_clean_int_0,
+                              int* clean_1, int num_clean_int_1,
+                              int rank, int num_ranks, int* mask_buffer, int* sync_buffer,
+                              cudaStream_t stream);
+
+void dispatch(void* packed_recv_x, void* packed_recv_x_scales,
+              int* packed_recv_src_info, int64_t* packed_recv_layout_range,
+              int* packed_recv_count,
+              int* mask_buffer,
+              int* cumulative_local_expert_recv_stats,
+              int64_t* dispatch_wait_recv_cost_stats,
+              void* rdma_recv_x, int* rdma_recv_count, void* rdma_x,
+              const void* x, const int64_t* topk_idx,
+              int* next_clean, int num_next_clean_int,
+              int num_tokens, int hidden, int num_max_dispatch_tokens_per_rank,
+              int num_topk, int num_experts, int rank, int num_ranks,
+              bool use_fp8, bool round_scale, bool use_ue8m0,
+              void* workspace, int num_device_sms,
+              cudaStream_t stream, int phases, ep_kernels::gpu_nixl_ctx nixl_ctx);
+
+void combine(void* combined_x,
+             void* rdma_recv_x, int* rdma_recv_flag, void* rdma_send_x,
+             const void* x, const int64_t* topk_idx, const float* topk_weights,
+             const int* src_info, const int64_t* layout_range,
+             int* mask_buffer,
+             int64_t* combine_wait_recv_cost_stats,
+             int* next_clean, int num_next_clean_int,
+             int num_combined_tokens, int hidden, int num_max_dispatch_tokens_per_rank,
+             int num_topk, int num_experts, int rank, int num_ranks,
+             bool use_logfmt,
+             void* workspace, int num_device_sms,
+             cudaStream_t stream, int phases, bool zero_copy, ep_kernels::gpu_nixl_ctx nixl_ctx);
+
+void barrier(ep_kernels::gpu_nixl_ctx nixl_ctx, int* mask_buffer_ptr, int* sync_buffer_ptr, cudaStream_t stream);
+
+void query_mask_buffer(int* mask_buffer_ptr, int num_ranks, int* output_mask_tensor, cudaStream_t stream);
+
+void update_mask_buffer(int* mask_buffer_ptr, int rank_to_mask, bool mask, cudaStream_t stream);
+
+void clean_mask_buffer(int* mask_buffer_ptr, int num_ranks, cudaStream_t stream);
+
+} // namespace ep_kernels
+
+} // namespace nixl_ep
diff --git a/examples/device/ep/csrc/kernels/configs.cuh b/examples/device/ep/csrc/kernels/configs.cuh
new file mode 100644
index 0000000000..a3c781255f
--- /dev/null
+++ b/examples/device/ep/csrc/kernels/configs.cuh
@@ -0,0 +1,74 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * This file incorporates material from the DeepSeek project, licensed under the MIT License.
+ * The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+ *
+ * SPDX-License-Identifier: MIT AND Apache-2.0
+ *
+ * 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
+
+#define NUM_WORKSPACE_BYTES (32 * 1024 * 1024)
+#define NUM_MAX_LOCAL_EXPERTS 1024
+#define NUM_BUFFER_ALIGNMENT_BYTES 128
+
+#define FINISHED_SUM_TAG 1024
+
+#ifndef ENABLE_FAST_DEBUG
+#define NUM_CPU_TIMEOUT_SECS 100
+#define NUM_TIMEOUT_CYCLES 200000000000ull // 200G cycles ~= 100s
+#else
+#define NUM_CPU_TIMEOUT_SECS 10
+#define NUM_TIMEOUT_CYCLES 20000000000ull // 20G cycles ~= 10s
+#endif
+
+#define EP_SEND_PHASE 1
+#define EP_RECV_PHASE 2
+
+// Remove Torch restrictions
+#ifdef __CUDA_NO_HALF_CONVERSIONS__
+#undef __CUDA_NO_HALF_CONVERSIONS__
+#endif
+#ifdef __CUDA_NO_HALF_OPERATORS__
+#undef __CUDA_NO_HALF_OPERATORS__
+#endif
+#ifdef __CUDA_NO_HALF2_OPERATORS__
+#undef __CUDA_NO_HALF2_OPERATORS__
+#endif
+#ifdef __CUDA_NO_BFLOAT16_CONVERSIONS__
+#undef __CUDA_NO_BFLOAT16_CONVERSIONS__
+#endif
+#ifdef __CUDA_NO_BFLOAT162_OPERATORS__
+#undef __CUDA_NO_BFLOAT162_OPERATORS__
+#endif
+
+#include <cstdint>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#ifndef DISABLE_SM90_FEATURES
+#include <cuda_fp8.h>
+#else
+// Ampere does not support FP8 features
+#define __NV_E4M3 0
+#define __NV_E5M2 1
+typedef int __nv_fp8_interpretation_t;
+typedef int __nv_fp8x4_e4m3;
+typedef uint8_t __nv_fp8_storage_t;
+#endif
+
+#include <infiniband/mlx5dv.h>
diff --git a/examples/device/ep/csrc/kernels/exception.cuh b/examples/device/ep/csrc/kernels/exception.cuh
new file mode 100644
index 0000000000..93bdc68401
--- /dev/null
+++ b/examples/device/ep/csrc/kernels/exception.cuh
@@ -0,0 +1,73 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * This file incorporates material from the DeepSeek project, licensed under the MIT License.
+ * The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+ *
+ * SPDX-License-Identifier: MIT AND Apache-2.0
+ *
+ * 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 <string>
+#include <exception>
+
+#include "configs.cuh"
+
+#ifndef EP_STATIC_ASSERT
+#define EP_STATIC_ASSERT(cond, reason) static_assert(cond, reason)
+#endif
+
+class EPException: public std::exception {
+private:
+    std::string message = {};
+
+public:
+    explicit EPException(const char *name, const char* file, const int line, const std::string& error) {
+        message = std::string("Failed: ") + name + " error " + file + ":" + std::to_string(line) + " '" + error + "'";
+    }
+
+    const char *what() const noexcept override { return message.c_str(); }
+};
+
+#ifndef CUDA_CHECK
+#define CUDA_CHECK(cmd) \
+do { \
+    cudaError_t e = (cmd); \
+    if (e != cudaSuccess) { \
+        throw EPException("CUDA", __FILE__, __LINE__, cudaGetErrorString(e)); \
+    } \
+} while (0)
+#endif
+
+#ifndef EP_HOST_ASSERT
+#define EP_HOST_ASSERT(cond) \
+do { \
+    if (not (cond)) { \
+        throw EPException("Assertion", __FILE__, __LINE__, #cond); \
+    } \
+} while (0)
+#endif
+
+#ifndef EP_DEVICE_ASSERT
+#define EP_DEVICE_ASSERT(cond) \
+do { \
+    if (not (cond)) { \
+        printf("Assertion failed: %s:%d, condition: %s\n", __FILE__, __LINE__, #cond); \
+        asm("trap;"); \
+    } \
+} while (0)
+#endif
diff --git a/examples/device/ep/csrc/kernels/launch.cuh b/examples/device/ep/csrc/kernels/launch.cuh
new file mode 100644
index 0000000000..eb1f9230a4
--- /dev/null
+++ b/examples/device/ep/csrc/kernels/launch.cuh
@@ -0,0 +1,86 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * This file incorporates material from the DeepSeek project, licensed under the MIT License.
+ * The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+ *
+ * SPDX-License-Identifier: MIT AND Apache-2.0
+ *
+ * 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 "configs.cuh"
+#include "exception.cuh"
+
+#ifndef SETUP_LAUNCH_CONFIG
+#ifndef DISABLE_SM90_FEATURES
+#define SETUP_LAUNCH_CONFIG(num_sms, num_threads, stream) \
+    cudaLaunchConfig_t cfg = {(num_sms), (num_threads), 0, stream, nullptr, 0}; \
+    cudaLaunchAttribute attr[2]; \
+    attr[0].id = cudaLaunchAttributeCooperative; \
+    attr[0].val.cooperative = 1; \
+    attr[1].id = cudaLaunchAttributeClusterDimension; \
+    attr[1].val.clusterDim.x = (num_sms % 2 == 0 ? 2 : 1); \
+    attr[1].val.clusterDim.y = 1; \
+    attr[1].val.clusterDim.z = 1; \
+    cfg.attrs = attr; \
+    cfg.numAttrs = 2
+#else
+#define SETUP_LAUNCH_CONFIG(sms, threads, stream) \
+    int __num_sms = (sms); \
+    int __num_threads = (threads); \
+    auto __stream = (stream)
+#endif
+#endif
+
+#ifndef LAUNCH_KERNEL
+#ifndef DISABLE_SM90_FEATURES
+#define LAUNCH_KERNEL(config, kernel, ...) CUDA_CHECK(cudaLaunchKernelEx(config, kernel, ##__VA_ARGS__))
+#else
+#define LAUNCH_KERNEL(config, kernel, ...) \
+do { \
+    kernel<<<__num_sms, __num_threads, 0, __stream>>>(__VA_ARGS__); \
+    cudaError_t e = cudaGetLastError(); \
+    if (e != cudaSuccess) { \
+        EPException cuda_exception("CUDA", __FILE__, __LINE__, cudaGetErrorString(e)); \
+        fprintf(stderr, "%s\n", cuda_exception.what()); \
+        throw cuda_exception; \
+    } \
+} while (0)
+#endif
+#endif
+
+#ifndef SET_SHARED_MEMORY_FOR_TMA
+#ifndef DISABLE_SM90_FEATURES
+#define SET_SHARED_MEMORY_FOR_TMA(kernel) \
+EP_HOST_ASSERT(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size) == cudaSuccess); \
+cfg.dynamicSmemBytes = smem_size;
+#else
+#define SET_SHARED_MEMORY_FOR_TMA(kernel) void()
+#endif
+#endif
+
+#define SWITCH_HIDDEN(case_macro) \
+    switch (hidden) { \
+        case 2048: case_macro(2048); \
+        case 2560: case_macro(2560); \
+        case 4096: case_macro(4096); \
+        case 5120: case_macro(5120); \
+        case 6144: case_macro(6144); /* For qwen3 coder */ \
+        case 7168: case_macro(7168); \
+        case 8192: case_macro(8192); \
+        default: EP_HOST_ASSERT(false and "Unsupported hidden"); \
+    } while (false)
diff --git a/examples/device/ep/csrc/kernels/nixl_ep.cu b/examples/device/ep/csrc/kernels/nixl_ep.cu
new file mode 100644
index 0000000000..3005e89a03
--- /dev/null
+++ b/examples/device/ep/csrc/kernels/nixl_ep.cu
@@ -0,0 +1,1154 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * This file incorporates material from the DeepSeek project, licensed under the MIT License.
+ * The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+ *
+ * SPDX-License-Identifier: MIT AND Apache-2.0
+ *
+ * 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 "configs.cuh"
+#include "exception.cuh"
+#include "launch.cuh"
+
+#include "api.cuh"
+#include "nixl_device.cuh"
+#include "utils.cuh"
+#include <cooperative_groups.h>
+namespace cg = cooperative_groups;
+
+namespace nixl_ep {
+
+namespace ep_kernels {
+
+template<bool use_warp_sync = false>
+__forceinline__ __device__ bool is_rank_masked(int* mask_buffer_ptr, int rank) {
+    if (mask_buffer_ptr == nullptr) {
+        return false;
+    }
+    if constexpr (use_warp_sync) {
+        return __shfl_sync(0xffffffff, ld_acquire_global(mask_buffer_ptr + rank), 0) != 0;
+    } else {
+        return ld_acquire_global(mask_buffer_ptr + rank) != 0;
+    }
+}
+
+template <bool kUseFP8, bool kUseUE8M0, int kHidden>
+__global__ __launch_bounds__(1024, 1) void
+dispatch(void* packed_recv_x, void* packed_recv_x_scales,
+         int* packed_recv_src_info, int64_t* packed_recv_layout_range,
+         int* packed_recv_count,
+         int* mask_buffer_ptr,
+         int* cumulative_local_expert_recv_stats,
+         int64_t* dispatch_wait_recv_cost_stats,
+         void* rdma_recv_x, int* rdma_recv_count, void* rdma_x,
+         const void* x, const int64_t* topk_idx,
+         int* atomic_counter_per_expert, int* atomic_finish_counter_per_expert,
+         int* next_clean, int num_next_clean_int,
+         int num_tokens, int num_max_dispatch_tokens_per_rank,
+         int num_topk, int num_experts, int rank, int num_ranks,
+         int num_warp_groups, int num_warps_per_group,
+         bool round_scale, int phases, ep_kernels::gpu_nixl_ctx nixl_ctx) {
+    const auto sm_id = static_cast<int>(blockIdx.x);
+    const auto thread_id = static_cast<int>(threadIdx.x);
+    const auto warp_id = thread_id / 32, lane_id = get_lane_id();
+    const auto num_sms = static_cast<int>(gridDim.x);
+    const auto num_warps = num_warp_groups * num_warps_per_group;
+    const auto num_local_experts = num_experts / num_ranks;
+    const auto warp_group_id = warp_id / num_warps_per_group;
+    const auto sub_warp_id = warp_id % num_warps_per_group;
+    const auto responsible_expert_idx = sm_id * num_warp_groups + warp_group_id;
+
+    // May extract UE8M0 from the scales
+    using scale_t = std::conditional_t<kUseUE8M0, uint8_t, float>;
+    using packed_t = std::conditional_t<kUseUE8M0, uint32_t, float>;
+    EP_STATIC_ASSERT(sizeof(packed_t) % sizeof(scale_t) == 0, "Invalid vector length");
+
+    // FP8 staffs
+    constexpr int kNumPerChannels = 128;
+    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_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 = std::conditional_t<kUseFP8, int2, int4>;
+    const size_t num_bytes_per_msg = sizeof(int4) + (kUseFP8 ? (kHidden + num_scales * sizeof(float)) : (kHidden * sizeof(nv_bfloat16)));
+    const size_t num_int4_per_msg = num_bytes_per_msg / sizeof(int4);
+    EP_DEVICE_ASSERT(num_bytes_per_msg % sizeof(int4) == 0);
+
+    // Expert counts
+    constexpr int kNumMaxWarpGroups = 32;
+    __shared__ int shared_num_tokens_sent_per_expert[kNumMaxWarpGroups];
+
+    // Sending phase
+    if ((phases & EP_SEND_PHASE) == 0)
+        goto dispatch_RECV;
+
+    // There are 2 kinds of warps in this part:
+    // 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);
+        EP_STATIC_ASSERT(kHidden % (32 * kNumElemsPerRead) == 0, "Invalid hidden");
+        EP_STATIC_ASSERT(kNumElemsPerRead * 32 % kNumPerChannels == 0, "Invalid vectorization");
+        const auto num_threads = (num_warps - 1) * 32;
+        const size_t hidden_bf16_int4 = kHidden / kNumElemsPerRead;
+
+        for (int token_idx = sm_id; token_idx < num_tokens; token_idx += num_sms) {
+            const auto x_int4 = static_cast<const int4*>(x) + token_idx * hidden_bf16_int4;
+            const auto rdma_x_src_idx = reinterpret_cast<int*>(static_cast<uint8_t*>(rdma_x) + token_idx * num_bytes_per_msg);
+            const auto rdma_x_vec = reinterpret_cast<vec_t*>(reinterpret_cast<uint8_t*>(rdma_x_src_idx) + sizeof(int4));
+            const auto rdma_x_scales = reinterpret_cast<float*>(reinterpret_cast<uint8_t*>(rdma_x_vec) + hidden_bytes);
+
+            // Overlap top-k index read and source token index writes
+            auto dst_expert_idx = warp_id < num_topk ? static_cast<int>(__ldg(topk_idx + token_idx * num_topk + warp_id)) : -1;
+            thread_id == 0 ? (*rdma_x_src_idx = token_idx) : 0;
+
+            // FP8 cast
+            EP_STATIC_ASSERT(hidden_bf16_int4 % 32 == 0, "Must use the full warp to reduce");
+            #pragma unroll
+            for (int i = thread_id; i < hidden_bf16_int4; i += num_threads) {
+                // Read
+                auto int4_value = __ldg(x_int4 + i);
+
+                if constexpr (kUseFP8) {
+                    // Calculate local amax
+                    auto bf16_values = reinterpret_cast<nv_bfloat16*>(&int4_value);
+                    float fp32_values[kNumElemsPerRead];
+                    float amax = kFP8Margin, scale, scale_inv;
+                    #pragma unroll
+                    for (int j = 0; j < kNumElemsPerRead; ++ j) {
+                        fp32_values[j] = static_cast<float>(bf16_values[j]);
+                        amax = fmaxf(amax, fabsf(fp32_values[j]));
+                    }
+
+                    // Reduce amax and scale
+                    EP_STATIC_ASSERT(kNumElemsPerRead * 32 / kNumPerChannels == 2, "Invalid vectorization");
+                    amax = warp_reduce_max<16>(amax);
+                    calculate_fp8_scales(amax, scale, scale_inv, round_scale);
+                    if (lane_id == 0 or lane_id == 16)
+                        rdma_x_scales[i * kNumElemsPerRead / 128] = scale_inv;
+
+                    // Cast into send buffer
+                    vec_t int2_value;
+                    auto fp8x2_values = reinterpret_cast<__nv_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);
+                    }
+                    rdma_x_vec[i] = int2_value;
+                } else {
+                    // Reinterpret-cast is for C++14 compatibility
+                    rdma_x_vec[i] = *reinterpret_cast<vec_t*>(&int4_value);
+                }
+            }
+            asm volatile("bar.sync 1, %0;" :: "r"(num_threads));
+
+            // Issue NIXL 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_x) +
+                                     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;
+                const auto dst_p2p_ptr = nixl_ctx.rdma_p2p_ptr_get(dst_ptr, dst_rank);
+                if (not is_rank_masked<true>(mask_buffer_ptr, dst_rank)) {
+                    if (dst_p2p_ptr == 0) {
+                        nixlGpuXferReqH batch_req = nixl_ctx.batch_get(dst_expert_local_idx, dst_rank);
+                        size_t src_offset = nixl_ctx.batch_offset_get(src_ptr);
+                        size_t dst_offset = nixl_ctx.batch_offset_get(dst_ptr);
+                        EP_DEVICE_ASSERT(nixlGpuPostSingleWriteXferReq<nixl_gpu_level_t::WARP>(batch_req, 0, src_offset, dst_offset, num_bytes_per_msg, 0, (slot_idx + 1) % 4 == 0) == NIXL_IN_PROG);
+                    } else {
+                        // 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*>(dst_p2p_ptr);
+                        UNROLLED_WARP_COPY(8, lane_id, num_int4_per_msg, dst_int4_ptr, src_int4_ptr, ld_nc_global, st_na_global);
+                    }
+                }
+                // Increase counter after finishing
+                __syncwarp();
+                lane_id == 0 ? atomic_add_release_global(atomic_finish_counter_per_expert + dst_expert_idx, 1) : 0;
+            }
+        }
+    } else if (warp_id == num_warps - 1) {
+        EP_DEVICE_ASSERT(num_sms > 1);
+        if (sm_id == 0) {
+            // The first SM is responsible for cleaning the next buffer
+            nixl_ctx.clean_counters_warp(lane_id);
+
+            // Notify before executing `int_p`
+            __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);
+        }
+
+        // This SM should be responsible for some destination experts, read `topk_idx` for them
+        int expert_count[kNumMaxWarpGroups] = {0};
+        const auto expert_begin_idx = sm_id * num_warp_groups;
+        const auto expert_end_idx = min(expert_begin_idx + num_warp_groups, num_experts);
+
+        // Per lane count
+        #pragma unroll 8
+        for (int i = lane_id; i < num_tokens * num_topk; i += 32) {
+            auto idx = static_cast<int>(__ldg(topk_idx + i));
+            if (idx >= expert_begin_idx and idx < expert_end_idx)
+                expert_count[idx - expert_begin_idx] ++;
+        }
+
+        // Warp reduce
+        #pragma unroll
+        for (int i = expert_begin_idx; i < expert_end_idx; ++ i) {
+            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);
+            }
+        }
+    }
+    __syncthreads();
+
+    // Issue count sends
+    if (responsible_expert_idx < num_experts and sub_warp_id == 0 and lane_id == 0) {
+        const auto dst_rank = responsible_expert_idx / num_local_experts;
+        const auto dst_expert_local_idx = responsible_expert_idx % num_local_experts;
+        const auto num_tokens_sent = shared_num_tokens_sent_per_expert[responsible_expert_idx - sm_id * num_warp_groups];
+
+        // Wait local sends issued and send expert counts
+        while (ld_acquire_global(atomic_finish_counter_per_expert + responsible_expert_idx) != FINISHED_SUM_TAG * 2);
+        if (not is_rank_masked(mask_buffer_ptr, dst_rank)) {
+            uint64_t *p2p_counter = nixl_ctx.counter_p2p_ptr_get(dst_expert_local_idx, dst_rank);
+            if (p2p_counter == nullptr) {
+                nixlGpuXferReqH remote_counter = nixl_ctx.remote_counter_get(dst_expert_local_idx, dst_rank);
+                /* WARNING: original counter value is negative, I chose different encoding since our counters are uint64_ts */
+                EP_DEVICE_ASSERT(nixlGpuPostSignalXferReq<nixl_gpu_level_t::THREAD>(remote_counter, 0, num_tokens_sent + 1, 0) == NIXL_IN_PROG);
+            } else {
+                /* WARNING: original counter value is negative, I chose different encoding since our counters are uint64_ts */
+                st_release_sys_global(p2p_counter, static_cast<uint64_t>(num_tokens_sent + 1));
+            }
+        }
+
+        // Clean workspace for next use
+        atomic_counter_per_expert[responsible_expert_idx] = 0;
+        atomic_finish_counter_per_expert[responsible_expert_idx] = 0;
+
+        // Clean `packed_recv_count`
+        if (dst_rank == 0)
+            packed_recv_count[dst_expert_local_idx] = 0;
+    }
+    __syncwarp();
+
+    // Receiving phase
+    dispatch_RECV:
+    if ((phases & EP_RECV_PHASE) == 0)
+        return;
+
+    // For send-and-recv kernels, we need a grid sync for making `packed_recv_count` visible
+    if (phases & EP_SEND_PHASE)
+        cg::this_grid().sync();
+
+    // Receiving and packing
+    if (responsible_expert_idx < num_experts) {
+        const auto src_rank = responsible_expert_idx / num_local_experts;
+        const auto local_expert_idx = responsible_expert_idx % num_local_experts;
+        const auto rdma_recv_x_uint8 = static_cast<uint8_t*>(rdma_recv_x) +
+                local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_msg +
+                src_rank * num_max_dispatch_tokens_per_rank * num_bytes_per_msg;
+        const auto recv_x_int4 = static_cast<int4*>(packed_recv_x) +
+                local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * hidden_int4;
+        const auto recv_src_info = packed_recv_src_info + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank;
+        const auto recv_range = packed_recv_layout_range + local_expert_idx * num_ranks;
+        const auto num_aligned_scales = align<int>(num_scales, sizeof(float) / sizeof(scale_t));
+        const auto recv_x_scales = static_cast<scale_t*>(packed_recv_x_scales) + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_aligned_scales;
+
+        // Shared between sub-warps in warp groups
+        __shared__ int shared_num_recv_tokens[kNumMaxWarpGroups], shared_recv_token_begin_idx[kNumMaxWarpGroups];
+
+        // Wait tokens to arrive
+        // NOTES: using sub-warp 1 to overlap with sub-warp 0
+        int num_recv_tokens = 0, recv_token_begin_idx;
+        EP_DEVICE_ASSERT(num_warps_per_group > 1 and num_warp_groups < 15);
+        if (sub_warp_id == 1 and lane_id == 0) {
+            auto start_time = clock64();
+            uint64_t wait_recv_cost = 0;
+            if (not is_rank_masked(mask_buffer_ptr, src_rank)) {
+                uint64_t *local_counter = nixl_ctx.local_counter_get(local_expert_idx, src_rank);
+                int poll_count = 0;
+                uint64_t current_value = 0;
+                while ((current_value = ld_acquire_sys_global(local_counter)) == 0 && (wait_recv_cost = clock64() - start_time) <= NUM_TIMEOUT_CYCLES) {
+                    poll_count++;
+                    if (poll_count % 10000000 == 0) {
+                        printf("[NIXL_EP-LL-DISPATCH] RANK %2d | POLL_HANG: expert[%d] from rank[%d] still waiting after %d polls, counter_val=%lu\n",
+                               rank, local_expert_idx, src_rank, poll_count, current_value);
+                    }
+                }
+                num_recv_tokens = current_value - 1;
+            }
+
+            // Mask rank if timeout
+            if (wait_recv_cost > NUM_TIMEOUT_CYCLES) {
+                printf("Warning: NIXL_EP timeout for dispatch receive, rank %d, local_expert_idx %d, src_rank %d\n", rank, local_expert_idx, src_rank);
+                if (mask_buffer_ptr == nullptr)
+                    trap();
+                atomicExch(mask_buffer_ptr + src_rank, 1);
+            }
+
+            recv_token_begin_idx = atomicAdd(packed_recv_count + local_expert_idx, num_recv_tokens);
+            shared_num_recv_tokens[warp_group_id] = num_recv_tokens;
+            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);
+
+            // Add stats for diagnosis
+            if (cumulative_local_expert_recv_stats != nullptr)
+                atomicAdd(cumulative_local_expert_recv_stats + local_expert_idx, num_recv_tokens);
+            if (dispatch_wait_recv_cost_stats != nullptr)
+                atomicAdd(reinterpret_cast<unsigned long long*>(dispatch_wait_recv_cost_stats + src_rank), wait_recv_cost);
+        }
+        asm volatile("bar.sync %0, %1;" :: "r"(warp_group_id + 2), "r"(num_warps_per_group * 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);
+        for (int i = sub_warp_id; i < num_recv_tokens; i += num_warps_per_group) {
+            // 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);
+            __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);
+
+            // Copy scales
+            if constexpr (kUseFP8) {
+                // Equivalent CuTe layout:
+                //   (num_tokens, (num_packed, num_elems_per_pack)):(num_elems_per_pack, (num_tokens * num_elems_per_pack, 1))
+                const auto src_scales = reinterpret_cast<float*>(reinterpret_cast<uint8_t*>(src_data) + hidden_bytes);
+                const auto num_elems_per_pack = static_cast<int>(sizeof(packed_t) / sizeof(scale_t));
+                const auto token_idx = recv_token_begin_idx + i;
+                const auto token_stride = num_elems_per_pack;
+                const auto pack_stride = num_ranks * num_max_dispatch_tokens_per_rank * num_elems_per_pack;
+                if (lane_id < num_scales) {
+                    const auto pack_idx = lane_id / num_elems_per_pack;
+                    const auto elem_idx = lane_id % num_elems_per_pack;
+                    auto scale = extract_required_scale_format<kUseUE8M0>(ld_nc_global(src_scales + lane_id));
+                    recv_x_scales[token_idx * token_stride + pack_idx * pack_stride + elem_idx] = scale;
+                }
+                if (lane_id + 32 < num_scales) {
+                    const auto pack_idx = (lane_id + 32) / num_elems_per_pack;
+                    const auto elem_idx = (lane_id + 32) % num_elems_per_pack;
+                    auto scale = extract_required_scale_format<kUseUE8M0>(ld_nc_global(src_scales + lane_id + 32));
+                    recv_x_scales[token_idx * token_stride + pack_idx * pack_stride + elem_idx] = scale;
+                }
+            }
+        }
+    }
+}
+
+void dispatch(void* packed_recv_x, void* packed_recv_x_scales,
+              int* packed_recv_src_info, int64_t* packed_recv_layout_range,
+              int* packed_recv_count,
+              int* mask_buffer_ptr,
+              int* cumulative_local_expert_recv_stats,
+              int64_t* dispatch_wait_recv_cost_stats,
+              void* rdma_recv_x, int* rdma_recv_count, void* rdma_x,
+              const void* x, const int64_t* topk_idx,
+              int* next_clean, int num_next_clean_int,
+              int num_tokens, int hidden, int num_max_dispatch_tokens_per_rank,
+              int num_topk, int num_experts, int rank, int num_ranks,
+              bool use_fp8, bool round_scale, bool use_ue8m0,
+              void* workspace, int num_device_sms,
+              cudaStream_t stream, int phases, ep_kernels::gpu_nixl_ctx nixl_ctx) {
+    constexpr int kNumMaxTopK = 9;
+    const int num_warp_groups = ceil_div(num_experts, num_device_sms);
+    const int num_warps_per_group = 32 / num_warp_groups;
+    EP_HOST_ASSERT(num_warp_groups > 0 and num_warps_per_group > 0);
+    EP_HOST_ASSERT(kNumMaxTopK + 1 <= num_warp_groups * num_warps_per_group);
+
+    const auto num_warps = num_warp_groups * num_warps_per_group;
+    const auto num_sms = ceil_div(num_experts, num_warp_groups);
+    EP_HOST_ASSERT(num_topk <= kNumMaxTopK);
+
+    // Workspace checks
+    auto atomic_counter_per_expert = static_cast<int*>(workspace);
+    auto atomic_finish_counter_per_expert = atomic_counter_per_expert + num_experts;
+    EP_HOST_ASSERT(num_experts * sizeof(int) * 2 <= NUM_WORKSPACE_BYTES);
+
+    // FP8 checks
+    if (use_ue8m0)
+        EP_HOST_ASSERT(round_scale and "UE8M0 SF requires `round_scale=True`");
+
+#define DISPATCH_LAUNCH_CASE(hidden) { \
+auto dispatch_func = dispatch<false, false, hidden>; \
+if (use_fp8 and not use_ue8m0) \
+    dispatch_func = dispatch<true, false, hidden>; \
+if (use_fp8 and use_ue8m0) \
+    dispatch_func = dispatch<true, true, hidden>; \
+LAUNCH_KERNEL(&cfg, dispatch_func, \
+              packed_recv_x, packed_recv_x_scales, \
+              packed_recv_src_info, packed_recv_layout_range, \
+              packed_recv_count, \
+              mask_buffer_ptr, \
+              cumulative_local_expert_recv_stats, \
+              dispatch_wait_recv_cost_stats, \
+              rdma_recv_x, rdma_recv_count, rdma_x, \
+              x, topk_idx, \
+              atomic_counter_per_expert, atomic_finish_counter_per_expert, \
+              next_clean, num_next_clean_int, \
+              num_tokens, num_max_dispatch_tokens_per_rank, \
+              num_topk, num_experts, rank, num_ranks, \
+              num_warp_groups, num_warps_per_group, \
+              round_scale, phases, nixl_ctx); } break
+
+    SETUP_LAUNCH_CONFIG(num_sms, num_warps * 32, stream);
+    SWITCH_HIDDEN(DISPATCH_LAUNCH_CASE);
+#undef DISPATCH_LAUNCH_CASE
+}
+
+template <int kNumSendUnrolls>
+__forceinline__ __device__ int logfmt_encode(void* buffer, nv_bfloat162 *shared_amaxmin, const int& lane_id) {
+    constexpr int kNumElemsPerInt4 = sizeof(int4) / sizeof(nv_bfloat16);
+    constexpr float kLogThreshold = 0;
+    constexpr float kMinClip = 32; // `== log_2(2 ^ (2 ^ 5))`
+    constexpr int kNumBits = 10;
+    constexpr int kNumValues = 1 << (kNumBits - 1);
+
+    int4 int4_values[kNumSendUnrolls];
+    const auto& uint32_values = reinterpret_cast<uint32_t*>(int4_values);
+    const auto& bf162_values = reinterpret_cast<nv_bfloat162*>(int4_values);
+
+    // Calculate lane offset
+    const auto& ld_buffer = reinterpret_cast<uint32_t*>(static_cast<uint8_t*>(buffer) + lane_id * (kNumSendUnrolls * sizeof(int4)));
+    const auto& st_buffer = reinterpret_cast<uint32_t*>(static_cast<uint8_t*>(buffer) + lane_id * (kNumSendUnrolls * sizeof(int4) * 10 / 16));
+
+    // Local log amax
+    auto bf162_amax = __nv_bfloat162(CUDART_ZERO_BF16, CUDART_ZERO_BF16);
+    auto bf162_amin = __nv_bfloat162(CUDART_INF_BF16, CUDART_INF_BF16);
+    uint32_t local_signs = 0;
+    #pragma unroll
+    for (int k = 0; k < kNumSendUnrolls * kNumElemsPerInt4 / 2; ++ k) {
+        // TODO: eliminate bank conflicts
+        uint32_values[k] = ld_buffer[k];
+        local_signs |= ((uint32_values[k] >> 15) & 1) << (k * 2);
+        local_signs |= ((uint32_values[k] >> 31) & 1) << (k * 2 + 1);
+        uint32_values[k] &= 0x7fff7fff;
+
+        bf162_amax = __hmax2(bf162_amax, bf162_values[k]);
+        bf162_amin = __hmin2(bf162_amin, bf162_values[k]);
+    }
+
+    // Reduce per 128 channels
+    // TODO: figure out how hardware do 2-byte min/max
+    auto amax = std::max(static_cast<float>(bf162_amax.x), static_cast<float>(bf162_amax.y));
+    auto amin = std::min(static_cast<float>(bf162_amin.x), static_cast<float>(bf162_amin.y));
+    constexpr static int kNumLanesToReduce = 128 * sizeof(nv_bfloat16) / (kNumSendUnrolls * sizeof(int4));
+    amax = warp_reduce_max<kNumLanesToReduce>(amax);
+    amin = warp_reduce_min<kNumLanesToReduce>(amin);
+
+    // Write min/max into the shared memory
+    if (shared_amaxmin != nullptr)
+        *shared_amaxmin = __nv_bfloat162(amax, amin);
+    __syncwarp();
+
+    // Calculate log amin/amax float
+    const auto& log_amax = log2f_approx(amax);
+    const auto& log_amin = fmaxf(log2f_approx(amin), log_amax - kMinClip);
+    const bool& enable_cast = warp_reduce_and<kNumLanesToReduce, true>(log_amax < kLogThreshold and log_amin < log_amax);
+
+    // Case into LogFMT-10 if satisfied
+    if (enable_cast) {
+        const auto step = (log_amax - log_amin) / static_cast<float>(kNumValues - 2);
+        const auto step_inv = 1.0f / step;
+        const auto rounding = 2.0f - log2f_approx((1.0f + exp2f_approx(step)) * 0.5f) * step_inv;
+        const auto fused_rounding = rounding - log_amin * step_inv;
+
+        // Pack every 256 bits into 160 bits
+        EP_STATIC_ASSERT(kNumSendUnrolls == 2 or kNumSendUnrolls == 4, "kNumSendUnrolls == 2 or 4 only");
+        uint32_t encoded[kNumElemsPerInt4 * 2];
+        #pragma unroll 1
+        for (int i = 0; i < kNumSendUnrolls / 2; ++ i) {
+            #pragma unroll
+            for (int k = 0; k < kNumElemsPerInt4; ++ k) {
+                const auto& [x, y] = __bfloat1622float2(bf162_values[i * kNumElemsPerInt4 + k]);
+                encoded[k * 2 + 0] = __float2uint_rd(fmaxf(log2f_approx(x) * step_inv + fused_rounding, 0));
+                encoded[k * 2 + 1] = __float2uint_rd(fmaxf(log2f_approx(y) * step_inv + fused_rounding, 0));
+            }
+            st_buffer[i * 5 + 0] = (encoded[ 0] >> 0) | (encoded[ 1] << 9) | (encoded[ 2] << 18) | (encoded[ 3] << 27);
+            st_buffer[i * 5 + 1] = (encoded[ 3] >> 5) | (encoded[ 4] << 4) | (encoded[ 5] << 13) | (encoded[ 6] << 22) | (encoded[7]  << 31);
+            st_buffer[i * 5 + 2] = (encoded[ 7] >> 1) | (encoded[ 8] << 8) | (encoded[ 9] << 17) | (encoded[10] << 26);
+            st_buffer[i * 5 + 3] = (encoded[10] >> 6) | (encoded[11] << 3) | (encoded[12] << 12) | (encoded[13] << 21) | (encoded[14] << 30);
+            st_buffer[i * 5 + 4] = (encoded[14] >> 2) | (encoded[15] << 7) | ((i == 0) ? (local_signs << 16) : (local_signs & 0xffff0000u));
+        }
+        tma_store_fence();
+        __syncwarp();
+    }
+
+    // Return TMA copy bytes
+    return enable_cast ? (32 * (kNumSendUnrolls * sizeof(int4) * 8 * 10 / 16 / 8)):
+                         (32 * (kNumSendUnrolls * sizeof(int4)));
+}
+
+template <int kNumLanes, int kNumSendUnrolls, int kNumRecvUnrolls>
+__forceinline__ __device__ void logfmt_check_amaxmin(uint8_t* meta_buffer, float2* shared_log_amax,
+                                                     float2* shared_log_amin, int* shared_cast_info,
+                                                     const int lane_id) {
+    constexpr float kLogThreshold = 0;
+    constexpr float kMinClip = 32; // `== log_2(2 ^ (2 ^ 5))`
+
+    bool enable_cast = true;
+    if (lane_id < kNumLanes) {
+        // Calculate log amin/amax float
+        auto amaxmin2 = reinterpret_cast<uint64_t*>(meta_buffer)[lane_id];
+        const auto& bf162_amaxmin = reinterpret_cast<__nv_bfloat162*>(&amaxmin2);
+        float log_amax[2], log_amin[2];
+        #pragma unroll
+        for (int i = 0; i < 2; ++ i) {
+            auto amax = static_cast<float>(bf162_amaxmin[i].x);
+            auto amin = static_cast<float>(bf162_amaxmin[i].y);
+            log_amax[i] = log2f_approx(amax);
+            log_amin[i] = amin == 0 ? log_amax[i] - kMinClip : fmaxf(log2f_approx(amin), log_amax[i] - kMinClip);
+            enable_cast = enable_cast and log_amax[i] < kLogThreshold and log_amin[i] < log_amax[i];
+        }
+        shared_log_amax[lane_id] = make_float2(log_amax[0], log_amax[1]);
+        shared_log_amin[lane_id] = make_float2(log_amin[0], log_amin[1]);
+    }
+
+    const auto& casted = warp_reduce_and<kNumSendUnrolls>(enable_cast) ? 1u << (lane_id / kNumRecvUnrolls): 0u;
+    const auto& num_casted_prefix = __popc(warp_reduce_or<kNumRecvUnrolls, true>(casted) & ((1u << (lane_id / kNumRecvUnrolls)) - 1));
+
+    if (lane_id < kNumLanes and lane_id % kNumRecvUnrolls == 0)
+        shared_cast_info[lane_id / kNumRecvUnrolls] = (num_casted_prefix << 1) | (casted ? 1u : 0u);
+    __syncwarp();
+}
+
+template <int kNumRecvUnrolls>
+__forceinline__ __device__ void decode_and_accumulate(uint32_t* ld_buffer, float* accum,
+                                                      const float& log_amax, const float& log_amin,
+                                                      const bool& enable_cast, const float& weight) {
+    if (enable_cast) {
+        constexpr int kNumBits = 10;
+        constexpr int kNumValues = 1 << (kNumBits - 1);
+
+        const auto& step = (log_amax - log_amin) / static_cast<float>(kNumValues - 2);
+        auto decode = [=](const uint32_t &encoded, const uint32_t &sign) {
+            const auto decoded = encoded == 0 ? .0f : exp2f_approx((encoded - 1) * step + log_amin);
+            return sign ? -decoded : decoded;
+        };
+
+        EP_STATIC_ASSERT(kNumRecvUnrolls == 2 or kNumRecvUnrolls == 4, "kNumRecvUnrolls == 2 or 4 only");
+        #pragma unroll
+        for (int i = 0; i < kNumRecvUnrolls / 2; ++ i) {
+            uint32_t concat[6];
+            concat[0] = ld_buffer[i * 5];
+            #pragma unroll
+            for (int k = 1; k < 5; ++ k)
+                concat[k] = (ld_buffer[i * 5 + k - 1] >> (32 - k * 5)) | (ld_buffer[i * 5 + k] << (k * 5));
+            concat[5] = ld_buffer[i * 5 + 4] >> 7;
+
+            const uint32_t& local_signs = ld_buffer[i * 5 + 4] >> 16;
+            #pragma unroll
+            for (int k = 0; k < 5; ++ k) {
+                accum[i * 16 + k * 3 + 0] += decode((concat[k] >>  0) & 0x1ff, (local_signs >> (k * 3 + 0)) & 1) * weight;
+                accum[i * 16 + k * 3 + 1] += decode((concat[k] >>  9) & 0x1ff, (local_signs >> (k * 3 + 1)) & 1) * weight;
+                accum[i * 16 + k * 3 + 2] += decode((concat[k] >> 18) & 0x1ff, (local_signs >> (k * 3 + 2)) & 1) * weight;
+            }
+            accum[i * 16 + 15] += decode(concat[5] & 0x1ff, (local_signs >> 15) & 1) * weight;
+        }
+    } else {
+        #pragma unroll
+        for (int k = 0; k < kNumRecvUnrolls * 4; ++ k) {
+            auto bf16_pack = *reinterpret_cast<__nv_bfloat162*>(ld_buffer + k);
+            accum[k * 2 + 0] += static_cast<float>(bf16_pack.x) * weight;
+            accum[k * 2 + 1] += static_cast<float>(bf16_pack.y) * weight;
+        }
+    }
+}
+
+template <bool kUseLogFMT, int kHidden, int kNumMaxTopk, int kNumMaxUnrolls>
+__global__ __launch_bounds__(1024, 1) void
+combine(void* combined_x,
+        void* rdma_recv_x, int* rdma_recv_flag, void* rdma_send_x,
+        const void* x, const int64_t* topk_idx, const float* topk_weights,
+        const int* src_info, const int64_t* layout_range,
+        int* mask_buffer_ptr,
+        int64_t* combine_wait_recv_cost_stats,
+        int* next_clean, int num_next_clean_int,
+        int* atomic_clean_flag,
+        int num_combined_tokens, int hidden, int num_topk,
+        int num_max_dispatch_tokens_per_rank,
+        int num_experts, int rank, int num_ranks,
+        int num_warp_groups, int num_warps_per_group,
+        int phases, bool zero_copy, ep_kernels::gpu_nixl_ctx nixl_ctx) {
+    const auto sm_id = __shfl_sync(0xffffffff, static_cast<int>(blockIdx.x), 0);
+    const auto num_sms = __shfl_sync(0xffffffff, static_cast<int>(gridDim.x), 0);
+    const auto thread_id = static_cast<int>(threadIdx.x);
+    const auto num_threads = __shfl_sync(0xffffffff, static_cast<int>(blockDim.x), 0);
+    const auto warp_id = __shfl_sync(0xffffffff, thread_id / 32, 0), lane_id = get_lane_id();
+    const auto num_local_experts = num_experts / num_ranks;
+    const auto warp_group_id = warp_id / num_warps_per_group;
+    const auto sub_warp_id = warp_id % num_warps_per_group;
+    const auto responsible_expert_idx = sm_id * num_warp_groups + warp_group_id;
+
+    extern __shared__ __align__(1024) uint8_t smem_buffer[];
+
+    // Data type staffs
+    constexpr int kNumElemsPerInt4 = sizeof(int4) / sizeof(nv_bfloat16);
+    constexpr int64_t hidden_bf16_int4 = kHidden / kNumElemsPerInt4;
+
+    // Use different unroll factors for send and recv phases
+    constexpr int kNumSendUnrolls = kHidden % (32 * 4 * sizeof(int4) / sizeof(nv_bfloat16)) == 0 ? 4 : 2;
+    constexpr int kNumRecvUnrolls = 2;
+    constexpr int hidden_bf16_int4_pad = align(static_cast<int>(hidden_bf16_int4), 32 * kNumSendUnrolls);
+    EP_STATIC_ASSERT(kHidden % (32 * 2 * sizeof(int4) / sizeof(nv_bfloat16)) == 0, "Invalid hidden");
+    EP_STATIC_ASSERT(kNumSendUnrolls <= kNumMaxUnrolls and kNumRecvUnrolls <= kNumMaxUnrolls, "Invalid unrolls");
+    EP_STATIC_ASSERT(hidden_bf16_int4 % kNumSendUnrolls == 0, "Invalid hidden");
+    EP_STATIC_ASSERT(kNumSendUnrolls >= kNumRecvUnrolls, "Invalid unroll factors");
+
+    // Message package
+    EP_STATIC_ASSERT(kHidden % 128 == 0, "Invalid hidden");
+    constexpr int kNumDivisions = kHidden / 128;
+    constexpr int kNumMetaBytes = kNumDivisions * sizeof(nv_bfloat162);
+    constexpr size_t num_bytes_per_slot = kHidden * sizeof(nv_bfloat16) + kNumMetaBytes;
+    EP_STATIC_ASSERT(num_bytes_per_slot % sizeof(int4) == 0, "Invalid vectorization");
+
+    // Sending phase
+    if ((phases & EP_SEND_PHASE) == 0)
+        goto combine_RECV;
+
+    // Clean up next buffer
+    if (sm_id == 0 and warp_group_id == 0 and sub_warp_id == 0) {
+        nixl_ctx.clean_counters_warp(lane_id);
+
+        // Notify before executing `int_p`
+        __syncwarp();
+        if (lane_id == 0)
+            atomic_add_release_global(atomic_clean_flag, num_experts);
+    }
+
+    // Issue NIXL sends
+    if (responsible_expert_idx < num_experts) {
+        const auto dst_rank = responsible_expert_idx / num_local_experts;
+        const auto local_expert_idx = responsible_expert_idx % num_local_experts;
+        const auto global_expert_idx = rank * num_local_experts + local_expert_idx;
+        const auto layout = __ldg(layout_range + local_expert_idx * num_ranks + dst_rank);
+        const auto local_x = static_cast<const int4*>(x) +
+                local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * hidden_bf16_int4;
+        const auto local_src_info = src_info + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank;
+        const auto rdma_send_x_vec = static_cast<uint8_t*>(rdma_send_x) +
+                local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_slot;
+
+        // Unpack layout
+        int offset, num_tokens_to_send;
+        unpack2(layout, num_tokens_to_send, offset);
+
+        // TMA stuffs
+        constexpr int kNumTMABufferBytes = sizeof(int4) * 32 * kNumSendUnrolls;
+        constexpr int kNumStages = 3;
+        constexpr int kNumPrefetch = 1;
+        EP_STATIC_ASSERT(kNumStages == 3 and kNumPrefetch == 1, "Invalid stages");
+
+        auto smem_ptr = smem_buffer + warp_id * (kNumStages * (kNumTMABufferBytes + 16) + kNumMetaBytes);
+        uint32_t tma_phase = 0;
+        auto tma_buffers   = PatternVisitor([=](const int& i) { return reinterpret_cast<int4*>(smem_ptr + i * (kNumTMABufferBytes + 16)); });
+        auto full_barriers = PatternVisitor([=](const int& i) { return reinterpret_cast<uint64_t*>(smem_ptr + i * (kNumTMABufferBytes + 16) + kNumTMABufferBytes); });
+        auto meta_buffers  = kUseLogFMT ? reinterpret_cast<nv_bfloat162*>(smem_ptr + kNumStages * (kNumTMABufferBytes + 16)) : nullptr;
+        EP_STATIC_ASSERT(kNumSendUnrolls * kNumStages <= 12, "TMA buffer size exceed limit");
+
+        // Initialize m-barriers
+        if (lane_id < kNumStages) {
+            mbarrier_init(full_barriers[lane_id], 1);
+            fence_view_async_shared();
+            fence_barrier_init();
+        }
+        __syncwarp();
+
+        constexpr int kNumIters = hidden_bf16_int4_pad / (32 * kNumSendUnrolls);
+        auto tma_load_and_arrive = [&](const int& stage_idx, const int4* gmem_ptr, const int& num_bytes) {
+            tma_load_1d(tma_buffers[stage_idx], gmem_ptr, full_barriers[stage_idx], num_bytes);
+            mbarrier_arrive_and_expect_tx(full_barriers[stage_idx], num_bytes);
+        };
+        auto get_num_tma_bytes = [&](const int& offset_int4) {
+            return min(kNumTMABufferBytes, static_cast<int>((hidden_bf16_int4 - offset_int4) * sizeof(int4)));
+        };
+
+        // Issue NIXL send
+        if (not is_rank_masked<true>(mask_buffer_ptr, dst_rank)) {
+            for (int token_idx = offset + sub_warp_id; token_idx < offset + num_tokens_to_send; token_idx += num_warps_per_group) {
+                const auto x_int4 = local_x + token_idx * hidden_bf16_int4;
+                const auto rdma_send_type_row = reinterpret_cast<int*>(rdma_send_x_vec + token_idx * num_bytes_per_slot);
+                const auto rdma_send_x_vec_row = reinterpret_cast<uint8_t*>(rdma_send_type_row);
+
+                // Copy directly to local rank, or copy to buffer and issue RDMA
+                const auto src_idx = __shfl_sync(0xffffffff, __ldg(local_src_info + token_idx), 0);
+                const auto buf_ptr = reinterpret_cast<int64_t>(rdma_send_x_vec_row);
+                const auto dst_ptr = reinterpret_cast<uint64_t>(rdma_recv_x) + (global_expert_idx * num_max_dispatch_tokens_per_rank + src_idx) * num_bytes_per_slot;
+                const auto dst_p2p_ptr = nixl_ctx.rdma_p2p_ptr_get(dst_ptr, dst_rank);
+                int num_send_bytes = hidden * sizeof(nv_bfloat16);
+
+                if (not zero_copy or dst_p2p_ptr != 0) {
+                    // Read from `cpy_src_int4_ptr` and copy into `cpy_dst_int4_ptr`
+                    const auto cpy_src_int4_ptr = zero_copy ? reinterpret_cast<int4*>(buf_ptr) : x_int4;
+                    const auto cpy_dst_int4_ptr = dst_p2p_ptr == 0 ? reinterpret_cast<int4*>(buf_ptr) : reinterpret_cast<int4*>(dst_p2p_ptr);
+
+                    // Prefetch
+                    if (elect_one_sync(lane_id))
+                        tma_load_and_arrive(0, cpy_src_int4_ptr, get_num_tma_bytes(0));
+                    __syncwarp();
+
+                    int tma_offset_bytes = kNumMetaBytes;
+                    #pragma unroll
+                    for (int i = lane_id * kNumSendUnrolls, iter_idx = 0; i < hidden_bf16_int4_pad; i += 32 * kNumSendUnrolls, ++ iter_idx) {
+                        // Load the next iteration
+                        const int& stage_idx = iter_idx % kNumStages;
+                        const int& next_stage_idx = (iter_idx + 1) % kNumStages;
+                        if (iter_idx + 1 < kNumIters and elect_one_sync(lane_id)) {
+                            tma_store_wait<kNumStages - kNumPrefetch - 1>();
+                            const auto& offset_int4 = i + 32 * kNumSendUnrolls;
+                            tma_load_and_arrive(next_stage_idx, cpy_src_int4_ptr + offset_int4, get_num_tma_bytes(offset_int4));
+                        }
+                        __syncwarp();
+
+                        // Wait the current TMA arrival
+                        EP_STATIC_ASSERT(kNumStages < 32, "Too many stages");
+                        mbarrier_wait<true>(full_barriers[stage_idx], tma_phase, stage_idx);
+                        if constexpr (kUseLogFMT) {
+                            // Cast if possible
+                            constexpr int kNumInt4PerDivision = 128 / kNumElemsPerInt4;
+                            int num_tma_bytes = logfmt_encode<kNumSendUnrolls>(
+                                tma_buffers[stage_idx],
+                                // NOTES: only the leader lane will write the result
+                                (i % kNumInt4PerDivision == 0) ? meta_buffers + i / kNumInt4PerDivision : nullptr,
+                                lane_id);
+                            if (elect_one_sync(lane_id))
+                                tma_store_1d(tma_buffers[stage_idx], reinterpret_cast<uint8_t*>(cpy_dst_int4_ptr) + tma_offset_bytes, num_tma_bytes);
+                            tma_offset_bytes += num_tma_bytes;
+                        } else {
+                            // BF16 original values
+                            if (elect_one_sync(lane_id))
+                                tma_store_1d(tma_buffers[stage_idx], cpy_dst_int4_ptr + i, get_num_tma_bytes(i));
+                        }
+                        __syncwarp();
+                    }
+
+                    // Store metadata (min/max values) for LogFMT
+                    if constexpr (kUseLogFMT) {
+                        num_send_bytes = tma_offset_bytes;
+                        if (elect_one_sync(lane_id))
+                            tma_store_1d(meta_buffers, cpy_dst_int4_ptr, kNumMetaBytes);
+                    }
+
+                    // Flush all stores
+                    tma_store_wait();
+                    __syncwarp();
+                }
+
+                // Issue RDMA
+                // NOTES: for zero-copy mode, we assume the data is already in the send buffer
+                if (dst_p2p_ptr == 0) {
+                    nixlGpuXferReqH batch_req = nixl_ctx.batch_get(local_expert_idx, dst_rank);
+                    size_t src_offset = nixl_ctx.batch_offset_get(buf_ptr);
+                    size_t dst_offset = nixl_ctx.batch_offset_get(dst_ptr);
+                    EP_DEVICE_ASSERT(nixlGpuPostSingleWriteXferReq<nixl_gpu_level_t::WARP>(batch_req, 0, src_offset, dst_offset, num_send_bytes, 0, (token_idx - offset + 1) % 4 == 0) == NIXL_IN_PROG);
+                }
+            }
+        }
+
+        // Put the finishing flag
+        EP_DEVICE_ASSERT(num_warps_per_group > 1 and num_warp_groups < 16);
+        asm volatile("bar.sync %0, %1;" :: "r"(warp_group_id + 1), "r"(num_warps_per_group * 32));
+        if (sub_warp_id == 1 and lane_id == 0) {
+            while (ld_acquire_global(atomic_clean_flag) == 0);
+            if (not is_rank_masked<false>(mask_buffer_ptr, dst_rank)) {
+                uint64_t *p2p_counter = nixl_ctx.counter_p2p_ptr_get(local_expert_idx, dst_rank);
+                if (p2p_counter == nullptr) {
+                    nixlGpuXferReqH remote_counter = nixl_ctx.remote_counter_get(local_expert_idx, dst_rank);
+                    EP_DEVICE_ASSERT(nixlGpuPostSignalXferReq<nixl_gpu_level_t::THREAD>(remote_counter, 0, 1, 0) == NIXL_IN_PROG);
+                } else {
+                    st_release_sys_global(p2p_counter, 1);
+                }
+            }
+            atomic_add_release_global(atomic_clean_flag, -1);
+        }
+        __syncwarp();
+
+        // Destroy m-barriers
+        if (lane_id < kNumStages) {
+            mbarrier_inval(full_barriers[lane_id]);
+            fence_view_async_shared();
+            fence_barrier_init();
+        }
+        __syncwarp();
+    }
+
+    cg::this_grid().sync(); // TODO: revisit this grid sync
+
+    // Receiving phase
+    combine_RECV:
+    if ((phases & EP_RECV_PHASE) == 0)
+        return;
+
+    // Wait all ranks to arrive
+    if (responsible_expert_idx < num_experts) {
+        EP_DEVICE_ASSERT(num_warps_per_group > 1);
+        if (sub_warp_id == 0 and lane_id == 0) {
+            const auto src_rank = responsible_expert_idx / num_local_experts;
+            auto start_time = clock64();
+            uint64_t wait_recv_cost = 0;
+            if (not is_rank_masked(mask_buffer_ptr, src_rank)) {
+                auto local_expert_idx = responsible_expert_idx % num_local_experts;
+                uint64_t *local_counter = nixl_ctx.local_counter_get(local_expert_idx, src_rank);
+                int poll_count = 0;
+                uint64_t current_value = 0;
+                while ((current_value = ld_acquire_sys_global(local_counter)) == 0 && (wait_recv_cost = clock64() - start_time) <= NUM_TIMEOUT_CYCLES) {
+                    poll_count++;
+                    if (poll_count % 10000000 == 0) {
+                        printf("[NIXL_EP-LL-COMBINE] RANK %2d | POLL_HANG: expert[%d] from rank[%d] still waiting after %d polls, counter_val=%lu\n",
+                                rank, local_expert_idx, src_rank, poll_count, current_value);
+                    }
+                }
+            }
+            // Mask rank if timeout
+            if (wait_recv_cost > NUM_TIMEOUT_CYCLES) {
+                printf("Warning: NIXL_EP timeout for combine receive, rank %d, local_expert_idx %d, src_rank %d\n", rank, responsible_expert_idx % num_local_experts, src_rank);
+                if (mask_buffer_ptr == nullptr)
+                    trap();
+                atomicExch(mask_buffer_ptr + src_rank, 1);
+            }
+
+            if (combine_wait_recv_cost_stats != nullptr) {
+                atomicAdd(reinterpret_cast<unsigned long long*>(combine_wait_recv_cost_stats + src_rank), wait_recv_cost);
+            }
+        }
+    }
+    cg::this_grid().sync();
+
+    // Reassign warp groups
+    constexpr int kMaxNumGroups = 2;
+    const int num_decode_warps = hidden_bf16_int4_pad / (kNumRecvUnrolls * 32);
+    const int num_groups = min(kMaxNumGroups, (num_threads / 32) / (num_decode_warps + 1));
+    const int decode_warp_idx = __shfl_sync(0xffffffff, warp_id % (num_decode_warps + 1), 0);
+    const int group_idx = __shfl_sync(0xffffffff, warp_id / (num_decode_warps + 1), 0);
+    EP_STATIC_ASSERT(kHidden % (32 * kNumElemsPerInt4) == 0, "Invalid vectorization");
+    EP_DEVICE_ASSERT(num_topk <= 32);
+    EP_DEVICE_ASSERT(num_groups > 0);
+
+    if (group_idx < num_groups) {
+        constexpr int kNumStages = 3;
+        constexpr int kNumTMABufferBytes = 16 * 2 + kHidden * 2;
+        constexpr int kNumBF16PerWarpBytes = 32 * kNumRecvUnrolls * kNumElemsPerInt4 * 2;
+        constexpr int kNumLogFMTPerWarpBytes = kNumBF16PerWarpBytes / 16 * 10;
+        constexpr int kNumDivisionBytes = kNumDivisions * sizeof(uint32_t);
+        constexpr int kNumBytesPerGroup = kNumStages * kNumTMABufferBytes + kHidden * 2 + kNumStages * kNumDivisionBytes * 3;
+
+        // Reallocate shared memory
+        const auto smem_group_buffer = smem_buffer + kNumBytesPerGroup * group_idx;
+        auto full_barriers  = PatternVisitor([=](const int& i) { return reinterpret_cast<uint64_t*>(smem_group_buffer + i * kNumTMABufferBytes); });
+        auto empty_barriers = PatternVisitor([=](const int& i) { return reinterpret_cast<uint64_t*>(smem_group_buffer + i * kNumTMABufferBytes + 8); });
+        auto tma_ld_buffers = PatternVisitor([=](const int& i) { return reinterpret_cast<uint8_t* >(smem_group_buffer + i * kNumTMABufferBytes + 16); });
+        auto tma_st_buffers = PatternVisitor([=](const int& i) { return reinterpret_cast<uint32_t*>(smem_group_buffer + kNumStages * kNumTMABufferBytes + i * kNumBF16PerWarpBytes); });
+
+        // Redundant when logfmt is disabled
+        const auto smem_group_ptr = smem_group_buffer + kNumStages * kNumTMABufferBytes + kHidden * 2;
+        auto log_amax_buffers  = PatternVisitor([=](const int& i) { return reinterpret_cast<float*>(smem_group_ptr + i * kNumDivisionBytes); });
+        auto log_amin_buffers  = PatternVisitor([=](const int& i) { return reinterpret_cast<float*>(smem_group_ptr + kNumStages * kNumDivisionBytes + i * kNumDivisionBytes); });
+        auto cast_info_buffers = PatternVisitor([=](const int& i) { return reinterpret_cast<int*>  (smem_group_ptr + kNumStages * kNumDivisionBytes * 2 + i * kNumDivisionBytes); });
+
+        uint32_t tma_phase = 0;
+        EP_STATIC_ASSERT(kNumStages < 32, "Too many stages");
+        if (decode_warp_idx == num_decode_warps)
+            tma_phase = (1 << kNumStages) - 1;
+
+        // Initialize m-barriers
+        if (decode_warp_idx == num_decode_warps and lane_id < kNumStages) {
+            mbarrier_init(full_barriers[lane_id], 1);
+            mbarrier_init(empty_barriers[lane_id], num_decode_warps);
+        }
+        asm volatile("bar.sync %0, %1;" :: "r"(group_idx + 1), "r"((num_decode_warps + 1) * 32));
+
+        int stage_idx = 0, topk_idx_by_lane = 0;
+        EP_STATIC_ASSERT(kNumMaxTopk <= 32, "Invalid number of topks");
+        if (decode_warp_idx == num_decode_warps) {
+            // TMA load warp
+            for (int token_idx = sm_id + num_sms * group_idx; token_idx < num_combined_tokens; token_idx += num_sms * num_groups) {
+                if (lane_id < num_topk)
+                    topk_idx_by_lane = static_cast<int>(__ldg(topk_idx + token_idx * num_topk + lane_id));
+                for (int i = 0; i < num_topk; ++ i) {
+                    int topk_idx_reg = __shfl_sync(0xffffffff, topk_idx_by_lane, i);
+                    if (topk_idx_reg < 0)
+                        continue;
+                    if (is_rank_masked(mask_buffer_ptr, topk_idx_reg / num_local_experts))
+                        continue;
+
+                    mbarrier_wait<true>(empty_barriers[stage_idx], tma_phase, stage_idx);
+                    auto buffer = static_cast<uint8_t*>(rdma_recv_x) + (topk_idx_reg * num_max_dispatch_tokens_per_rank + token_idx) * num_bytes_per_slot;
+                    if constexpr (kUseLogFMT) {
+                        logfmt_check_amaxmin<kNumDivisions / 2, kNumSendUnrolls, kNumRecvUnrolls>(
+                            buffer, reinterpret_cast<float2*>(log_amax_buffers[stage_idx]),
+                            reinterpret_cast<float2*>(log_amin_buffers[stage_idx]), cast_info_buffers[stage_idx], lane_id);
+                    }
+                    if (elect_one_sync(lane_id)) {
+                        int num_casted = 0;
+                        if constexpr (kUseLogFMT) {
+                            const auto& info = cast_info_buffers[stage_idx][num_decode_warps - 1];
+                            num_casted = (info >> 1) + (info & 1);
+                        }
+                        int num_tma_bytes = num_casted * kNumLogFMTPerWarpBytes + (num_decode_warps - num_casted) * kNumBF16PerWarpBytes;
+                        tma_load_1d(tma_ld_buffers[stage_idx], buffer + (kUseLogFMT ? kNumMetaBytes : 0), full_barriers[stage_idx], num_tma_bytes);
+                        mbarrier_arrive_and_expect_tx(full_barriers[stage_idx], num_tma_bytes);
+                    }
+                    __syncwarp();
+                    stage_idx = (stage_idx + 1) % kNumStages;
+                }
+            }
+        } else {
+            // Reduction warps
+            float topk_weights_by_lane;
+            for (int token_idx = sm_id + num_sms * group_idx; token_idx < num_combined_tokens; token_idx += num_sms * num_groups) {
+                if (lane_id < num_topk) {
+                    topk_idx_by_lane = static_cast<int>(__ldg(topk_idx + token_idx * num_topk + lane_id));
+                    topk_weights_by_lane = __ldg(topk_weights + token_idx * num_topk + lane_id);
+                }
+                __syncwarp();
+
+                float combined_values[kNumElemsPerInt4 * kNumRecvUnrolls] = {0.0f};
+                for (int i = 0; i < num_topk; ++ i) {
+                    int topk_idx_reg = __shfl_sync(0xffffffff, topk_idx_by_lane, i);
+                    if (topk_idx_reg < 0)
+                        continue;
+                    if (is_rank_masked(mask_buffer_ptr, topk_idx_reg / num_local_experts))
+                        continue;
+                    const auto& topk_weight = __shfl_sync(0xffffffff, topk_weights_by_lane, i);
+
+                    mbarrier_wait<true>(full_barriers[stage_idx], tma_phase, stage_idx);
+                    if constexpr (kUseLogFMT) {
+                        const auto& info = cast_info_buffers[stage_idx][decode_warp_idx];
+                        bool enable_cast = info & 1;
+                        int num_casted_prefix = info >> 1;
+                        int tma_offset = kNumLogFMTPerWarpBytes * num_casted_prefix + kNumBF16PerWarpBytes * (decode_warp_idx - num_casted_prefix);
+                        int division_idx = decode_warp_idx * (kNumRecvUnrolls * 2) + lane_id * kNumRecvUnrolls / 16;
+                        decode_and_accumulate<kNumRecvUnrolls>(
+                            reinterpret_cast<uint32_t*>(tma_ld_buffers[stage_idx] + tma_offset + (enable_cast ? kNumLogFMTPerWarpBytes : kNumBF16PerWarpBytes) / 32 * lane_id),
+                            combined_values, log_amax_buffers[stage_idx][division_idx], log_amin_buffers[stage_idx][division_idx], enable_cast, topk_weight
+                        );
+                    } else {
+                        int tma_offset = kNumBF16PerWarpBytes * decode_warp_idx;
+                        decode_and_accumulate<kNumRecvUnrolls>(
+                            reinterpret_cast<uint32_t*>(tma_ld_buffers[stage_idx] + tma_offset + kNumBF16PerWarpBytes / 32 * lane_id),
+                            combined_values, 0, 0, false, topk_weight
+                        );
+                    }
+
+                    if (elect_one_sync(lane_id))
+                        mbarrier_arrive(empty_barriers[stage_idx]);
+                    stage_idx = (stage_idx + 1) % kNumStages;
+                }
+                tma_store_wait<0>();
+
+                #pragma unroll
+                for (int k = 0; k < kNumRecvUnrolls * 4; ++ k) {
+                    auto combined_pack = __nv_bfloat162(combined_values[k * 2], combined_values[k * 2 + 1]);
+                    tma_st_buffers[decode_warp_idx][kNumRecvUnrolls * 4 * lane_id + k] = *reinterpret_cast<uint32_t*>(&combined_pack);
+                }
+                tma_store_fence();
+                if (elect_one_sync(lane_id)) {
+                    tma_store_1d(tma_st_buffers[decode_warp_idx],
+                                 static_cast<int4*>(combined_x) + token_idx * hidden_bf16_int4 + decode_warp_idx * kNumRecvUnrolls * 32,
+                                 kNumBF16PerWarpBytes);
+                }
+                __syncwarp();
+            }
+        }
+
+        // Flush all stores
+        tma_store_wait<0>();
+    }
+}
+
+void combine(void* combined_x,
+             void* rdma_recv_x, int* rdma_recv_flag, void* rdma_send_x,
+             const void* x, const int64_t* topk_idx, const float* topk_weights,
+             const int* src_info, const int64_t* layout_range,
+             int* mask_buffer_ptr,
+             int64_t* combine_wait_recv_cost_stats,
+             int* next_clean, int num_next_clean_int,
+             int num_combined_tokens, int hidden, int num_max_dispatch_tokens_per_rank,
+             int num_topk, int num_experts, int rank, int num_ranks,
+             bool use_logfmt,
+             void* workspace, int num_device_sms,
+             cudaStream_t stream, int phases, bool zero_copy, ep_kernels::gpu_nixl_ctx nixl_ctx) {
+    constexpr int kNumMaxTopk = 9;
+    const int num_warp_groups = ceil_div(num_experts, num_device_sms);
+    const int num_warps_per_group = 32 / num_warp_groups;
+    const int num_recv_per_sm = ceil_div(num_combined_tokens, num_device_sms);
+    EP_HOST_ASSERT(num_warp_groups > 0 and num_warps_per_group > 0 and num_recv_per_sm >= 0);
+
+    const auto num_warps = num_warp_groups * num_warps_per_group;
+    const auto num_sms = max(ceil_div(num_experts, num_warp_groups),
+                             num_recv_per_sm == 0 ? 1 : ceil_div(num_combined_tokens, num_recv_per_sm));
+
+    // Check workspace
+    auto atomic_clean_flag = static_cast<int*>(workspace);
+    EP_HOST_ASSERT(sizeof(int) <= NUM_WORKSPACE_BYTES);
+    EP_HOST_ASSERT(num_topk <= kNumMaxTopk);
+
+    // Online cast cannot use zero-copy
+    EP_HOST_ASSERT(not (zero_copy and use_logfmt));
+
+    constexpr int kNumStages = 3;
+    constexpr int kNumMaxUnrolls = 4;
+    constexpr int kMaxNumGroups = 2;
+
+    // Send buffer size
+    const int num_meta_bytes = hidden / 128 * 4;
+    const int num_send_tma_bytes = 32 * sizeof(int4) * kNumMaxUnrolls + 16;
+    const int smem_send_size = num_warps * (kNumStages * num_send_tma_bytes + num_meta_bytes);
+
+    // Receive buffer size
+    const int num_recv_tma_bytes = 16 + hidden * 2;
+    const int smem_recv_size = kMaxNumGroups * (kNumStages * num_recv_tma_bytes + hidden * 2 + kNumStages * num_meta_bytes * 3);
+
+    // Total requirement
+    const int smem_size = max(smem_send_size, smem_recv_size);
+
+#define COMBINE_LAUNCH_CASE(hidden) { \
+auto combine_func = use_logfmt ? \
+    combine<true, hidden, kNumMaxTopk, kNumMaxUnrolls> : \
+    combine<false, hidden, kNumMaxTopk, kNumMaxUnrolls>; \
+SET_SHARED_MEMORY_FOR_TMA(combine_func); \
+LAUNCH_KERNEL(&cfg, combine_func, \
+              combined_x, \
+              rdma_recv_x, rdma_recv_flag, rdma_send_x, \
+              x, topk_idx, topk_weights, src_info, layout_range, \
+              mask_buffer_ptr, \
+              combine_wait_recv_cost_stats, \
+              next_clean, num_next_clean_int, \
+              atomic_clean_flag, \
+              num_combined_tokens, hidden, num_topk, \
+              num_max_dispatch_tokens_per_rank, \
+              num_experts, rank, num_ranks, \
+              num_warp_groups, num_warps_per_group, \
+              phases, zero_copy, nixl_ctx); } break
+
+    SETUP_LAUNCH_CONFIG(num_sms, num_warps * 32, stream);
+    SWITCH_HIDDEN(COMBINE_LAUNCH_CASE);
+#undef COMBINE_LAUNCH_CASE
+}
+
+template <int kNumThreads> __launch_bounds__(kNumThreads, 1)
+__global__ void query_mask_buffer(int* mask_buffer_ptr, int num_ranks, int* mask_tensor) {
+    const auto num_sms = static_cast<int>(gridDim.x);
+    const auto sm_id = static_cast<int>(blockIdx.x);
+    const auto num_threads = num_sms * kNumThreads;
+    const auto thread_id = sm_id * kNumThreads + static_cast<int>(threadIdx.x);
+    for (int rank_id = thread_id; rank_id < num_ranks; rank_id += num_threads) {
+        mask_tensor[rank_id] = mask_buffer_ptr[rank_id];
+    }
+}
+
+void query_mask_buffer(int* mask_buffer_ptr, int num_ranks, int* mask_tensor, cudaStream_t stream) {
+    constexpr int num_sms = 1;
+    constexpr int kNumThreads = 1024;
+    SETUP_LAUNCH_CONFIG(num_sms, kNumThreads, stream);
+    LAUNCH_KERNEL(&cfg, query_mask_buffer<kNumThreads>, mask_buffer_ptr, num_ranks, mask_tensor);
+}
+
+
+template <int kNumThreads> __launch_bounds__(kNumThreads, 1)
+__global__ void update_mask_buffer(int* mask_buffer_ptr, int rank_to_mask, bool mask) {
+    const auto sm_id = static_cast<int>(blockIdx.x);
+    const auto thread_id = static_cast<int>(threadIdx.x);
+    if (sm_id == 0 && thread_id == 0) {
+        atomicExch(mask_buffer_ptr + rank_to_mask, mask ? 1 : 0);
+    }
+}
+
+void update_mask_buffer(int* mask_buffer_ptr, int rank, bool mask, cudaStream_t stream) {
+    constexpr int num_sms = 1;
+    constexpr int kNumThreads = 32;
+    SETUP_LAUNCH_CONFIG(num_sms, kNumThreads, stream);
+    LAUNCH_KERNEL(&cfg, update_mask_buffer<kNumThreads>, mask_buffer_ptr, rank, mask);
+}
+
+
+template <int kNumThreads> __launch_bounds__(kNumThreads, 1)
+__global__ void clean_mask_buffer(int* mask_buffer_ptr, int num_ranks) {
+    auto thread_id = static_cast<int>(threadIdx.x);
+    #pragma unroll
+    for (int i = thread_id; i < num_ranks; i += kNumThreads)
+    mask_buffer_ptr[i] = 0;
+}
+
+void clean_mask_buffer(int* mask_buffer_ptr, int num_ranks, cudaStream_t stream) {
+    constexpr int num_sms = 1;
+    constexpr int kNumThreads = 32;
+    SETUP_LAUNCH_CONFIG(num_sms, kNumThreads, stream);
+    LAUNCH_KERNEL(&cfg, clean_mask_buffer<kNumThreads>, mask_buffer_ptr, num_ranks);
+}
+
+template <int kNumThreads>
+__global__ void barrier(int thread_id, int rank, int num_ranks,
+                                        int* mask_buffer_ptr, int* sync_buffer_ptr, ep_kernels::gpu_nixl_ctx nixl_ctx) {
+    EP_DEVICE_ASSERT(kNumThreads >= num_ranks);
+    if (thread_id == 0) atomicAdd(sync_buffer_ptr + rank, -1);
+    __syncthreads();
+
+    int cnt = sync_buffer_ptr[rank];
+    if (thread_id < num_ranks && thread_id != rank) {
+        const auto dst_rank = thread_id;
+        const auto dst_ptr = reinterpret_cast<uint64_t>(sync_buffer_ptr + rank);
+        if (not is_rank_masked(mask_buffer_ptr, dst_rank)) {
+            if (rank != dst_rank) {
+                nixlGpuXferReqH barrier_req = nixl_ctx.remote_barrier_get(0, dst_rank);
+                nixlGpuPostSingleWriteXferReq<nixl_gpu_level_t::THREAD>(barrier_req, 0, rank*sizeof(int), dst_rank*sizeof(int), sizeof(int), 0);
+            } else {
+                st_release_sys_global(reinterpret_cast<int*>(dst_ptr), cnt);
+            }
+
+            auto start_time = clock64();
+            uint64_t wait_recv_cost = 0;
+            while (ld_acquire_sys_global(reinterpret_cast<int*>(dst_ptr)) != cnt   // remote is not ready
+                   && (wait_recv_cost = clock64() - start_time) <= NUM_TIMEOUT_CYCLES               // not timeout
+            );
+            // Mask rank if timeout
+            if (wait_recv_cost > NUM_TIMEOUT_CYCLES) {
+                printf("Warning: NixlEP timeout for barrier, rank %d, dst_rank %d, expected value %d, actual value %d\n", rank, dst_rank, cnt, ld_acquire_global(sync_buffer_ptr + dst_rank));
+                if (mask_buffer_ptr == nullptr)
+                    trap();
+                atomicExch(mask_buffer_ptr + dst_rank, 1);
+            }
+        }
+    }
+    __syncthreads();
+}
+
+void barrier(ep_kernels::gpu_nixl_ctx nixl_ctx, int* mask_buffer_ptr, int* sync_buffer_ptr, cudaStream_t stream) {
+    constexpr int kNumThreads = 32;
+    SETUP_LAUNCH_CONFIG(1, kNumThreads, stream);
+    LAUNCH_KERNEL(&cfg, barrier<kNumThreads>, 0, nixl_ctx.rank, nixl_ctx.num_ranks, mask_buffer_ptr, sync_buffer_ptr, nixl_ctx);
+}
+} // namespace ep_kernels
+
+} // namespace nixl_ep
diff --git a/examples/device/ep/csrc/kernels/utils.cuh b/examples/device/ep/csrc/kernels/utils.cuh
new file mode 100644
index 0000000000..69a3c95a04
--- /dev/null
+++ b/examples/device/ep/csrc/kernels/utils.cuh
@@ -0,0 +1,467 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * This file incorporates material from the DeepSeek project, licensed under the MIT License.
+ * The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+ *
+ * SPDX-License-Identifier: MIT AND Apache-2.0
+ *
+ * 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 "exception.cuh"
+
+#define UNROLLED_WARP_COPY(UNROLL_FACTOR, LANE_ID, N, DST, SRC, LD_FUNC, ST_FUNC) \
+{ \
+    constexpr int kLoopStride = 32 * (UNROLL_FACTOR); \
+    typename std::remove_reference<decltype(LD_FUNC((SRC) + 0))>::type unrolled_values[(UNROLL_FACTOR)]; \
+    auto __src = (SRC); \
+    auto __dst = (DST); \
+    for (int __i = (LANE_ID); __i < ((N) / kLoopStride) * kLoopStride; __i += kLoopStride) { \
+        _Pragma("unroll") \
+        for (int __j = 0; __j < (UNROLL_FACTOR); ++ __j) \
+            unrolled_values[__j] = LD_FUNC(__src + __i + __j * 32); \
+        _Pragma("unroll") \
+        for (int __j = 0; __j < (UNROLL_FACTOR); ++ __j) \
+            ST_FUNC(__dst + __i + __j * 32, unrolled_values[__j]); \
+    } \
+    for (int __i = ((N) / kLoopStride) * kLoopStride + (LANE_ID); __i < (N); __i += 32) \
+        ST_FUNC(__dst + __i, LD_FUNC(__src + __i)); \
+}
+
+namespace nixl_ep {
+
+template <int kBytes>
+struct VecInt {};
+template<> struct VecInt<1> { using vec_t = int8_t; };
+template<> struct VecInt<2> { using vec_t = int16_t; };
+template<> struct VecInt<4> { using vec_t = int; };
+template<> struct VecInt<8> { using vec_t = int64_t; };
+template<> struct VecInt<16> { using vec_t = int4; };
+
+template <typename FuncT>
+struct PatternVisitor {
+    FuncT func;
+
+    __device__ __host__
+    explicit PatternVisitor(FuncT&& func): func(std::forward<FuncT>(func)) {}
+
+    __device__ __host__
+    auto operator [](const uint32_t& i) {
+        return func(i);
+    }
+};
+
+__device__ __forceinline__ void trap() {
+    asm("trap;");
+}
+
+__device__ __forceinline__ void memory_fence() {
+    asm volatile("fence.acq_rel.sys;":: : "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_sys_global(const uint64_t *ptr, uint64_t val) {
+    asm volatile("st.release.sys.global.u64 [%0], %1;"::"l"(ptr), "l"(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__ void st_release_cta(const volatile 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__ uint64_t ld_acquire_global(const uint64_t *ptr) {
+    uint64_t ret;
+    asm volatile("ld.acquire.gpu.global.u64 %0, [%1];" : "=l"(ret) : "l"(ptr));
+    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 atomic_add_release_global(const uint64_t* ptr, uint64_t value) {
+    uint64_t ret;
+    asm volatile("atom.add.release.gpu.global.u64 %0, [%1], %2;" : "=l"(ret) : "l"(ptr), "l"(value));
+    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;
+}
+
+// `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));
+}
+
+__device__  __forceinline__ float log2f_approx(const float &x) {
+    float ret;
+    asm volatile("lg2.approx.f32 %0, %1;" : "=f"(ret) : "f"(x));
+    return ret;
+}
+
+__device__  __forceinline__ float exp2f_approx(const float &x) {
+    float ret;
+    asm volatile("ex2.approx.f32 %0, %1;" : "=f"(ret) : "f"(x));
+    return ret;
+}
+
+// TMA PTX instructions
+#ifndef DISABLE_SM90_FEATURES
+
+__device__ __forceinline__ uint32_t elect_one_sync(int lane_id) {
+    uint32_t pred = 0;
+    asm volatile(
+      "{\n"
+      ".reg .b32 %%rx;\n"
+      ".reg .pred %%px;\n"
+      "      elect.sync %%rx|%%px, %2;\n"
+      "@%%px mov.s32 %1, 1;\n"
+      "      mov.s32 %0, %%rx;\n"
+      "}\n"
+      : "+r"(lane_id), "+r"(pred)
+      : "r"(0xffffffff));
+    return pred;
+}
+
+__device__ __forceinline__ void fence_view_async_shared() {
+    asm volatile("fence.proxy.async.shared::cta; \n" :: );
+}
+
+__device__ __forceinline__ void fence_barrier_init() {
+    asm volatile("fence.mbarrier_init.release.cluster; \n" :: );
+}
+
+__device__ __forceinline__ void mbarrier_init(uint64_t* mbar_ptr, uint32_t arrive_count) {
+    auto mbar_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_ptr));
+    asm volatile("mbarrier.init.shared::cta.b64 [%1], %0;" :: "r"(arrive_count), "r"(mbar_int_ptr));
+}
+
+__device__ __forceinline__ void mbarrier_inval(uint64_t* mbar_ptr) {
+    auto mbar_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_ptr));
+    asm volatile("mbarrier.inval.shared::cta.b64 [%0];" :: "r"(mbar_int_ptr));
+}
+
+template <bool kWithMultiStages = false>
+__device__ __forceinline__ void mbarrier_wait(uint64_t* mbar_ptr, uint32_t& phase, int stage_idx = 0) {
+    auto mbar_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_ptr));
+    const auto& wait = kWithMultiStages ? (phase >> stage_idx) & 1 : phase;
+    asm volatile("{\n\t"
+                 ".reg .pred       P1; \n\t"
+                 "LAB_WAIT: \n\t"
+                 "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2; \n\t"
+                 "@P1 bra DONE; \n\t"
+                 "bra     LAB_WAIT; \n\t"
+                 "DONE: \n\t"
+                 "}" :: "r"(mbar_int_ptr), "r"(wait), "r"(0x989680));
+    phase ^= kWithMultiStages ? (1 << stage_idx) : 1;
+}
+
+__device__ __forceinline__ void mbarrier_arrive_and_expect_tx(uint64_t* mbar_ptr, int num_bytes) {
+    auto mbar_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_ptr));
+    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%1], %0; \n\t" :: "r"(num_bytes), "r"(mbar_int_ptr));
+}
+
+__device__ __forceinline__ void mbarrier_arrive(uint64_t* mbar_ptr) {
+    auto mbar_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_ptr));
+    asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0]; \n\t" :: "r"(mbar_int_ptr));
+}
+
+__device__ __forceinline__ void tma_store_fence() {
+    asm volatile ("fence.proxy.async.shared::cta;");
+}
+
+constexpr uint64_t kEvictFirst = 0x12f0000000000000;
+constexpr uint64_t kEvictNormal = 0x1000000000000000;
+
+__device__ __forceinline__ void tma_load_1d(const void* smem_ptr, const void* gmem_ptr, uint64_t* mbar_ptr, int num_bytes,
+                                            bool evict_first = true) {
+    auto mbar_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_ptr));
+    auto smem_int_ptr  = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+    const auto cache_hint = evict_first ? kEvictFirst : kEvictNormal;
+    asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;\n"
+                 :: "r"(smem_int_ptr), "l"(gmem_ptr), "r"(num_bytes), "r"(mbar_int_ptr), "l"(cache_hint) : "memory");
+}
+
+__device__ __forceinline__ void tma_store_1d(const void* smem_ptr, const void* gmem_ptr, int num_bytes,
+                                             bool evict_first = true) {
+    auto smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+    const auto cache_hint = evict_first ? kEvictFirst : kEvictNormal;
+    asm volatile("cp.async.bulk.global.shared::cta.bulk_group.L2::cache_hint [%0], [%1], %2, %3;\n"
+                 :: "l"(gmem_ptr), "r"(smem_int_ptr), "r"(num_bytes), "l"(cache_hint) : "memory");
+    asm volatile("cp.async.bulk.commit_group;");
+}
+
+template <int N = 0>
+__device__ __forceinline__ void tma_store_wait() {
+    asm volatile("cp.async.bulk.wait_group.read %0;" :: "n"(N) : "memory");
+}
+
+#endif
+
+template <typename dtype_t>
+__host__ __device__ constexpr dtype_t ceil_div(dtype_t a, dtype_t b) {
+    return (a + b - 1) / b;
+}
+
+template <typename dtype_t>
+__host__ __device__ constexpr dtype_t align(dtype_t a, dtype_t b) {
+    return ceil_div<dtype_t>(a, b) * b;
+}
+
+__forceinline__ __device__ void get_channel_task_range(int num_tokens, int num_sms, int sm_id,
+                                                       int& token_start_idx, int& token_end_idx) {
+    int num_tokens_per_sm = ceil_div(num_tokens, num_sms);
+    token_start_idx = min(num_tokens_per_sm * sm_id, num_tokens);
+    token_end_idx = min(token_start_idx + num_tokens_per_sm, num_tokens);
+}
+
+template <typename dtype_a_t, typename dtype_b_t>
+__device__ __forceinline__ dtype_b_t pack2(const dtype_a_t& x, const dtype_a_t& y) {
+    EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), "Invalid dtypes");
+    dtype_b_t packed;
+    auto unpacked_ptr = reinterpret_cast<dtype_a_t*>(&packed);
+    unpacked_ptr[0] = x, unpacked_ptr[1] = y;
+    return packed;
+}
+
+template <typename dtype_a_t, typename dtype_b_t>
+__device__ __forceinline__ void unpack2(const dtype_b_t& packed, dtype_a_t& x, dtype_a_t& y) {
+    EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), "Invalid dtypes");
+    auto unpacked_ptr = reinterpret_cast<const dtype_a_t*>(&packed);
+    x = unpacked_ptr[0], y = unpacked_ptr[1];
+}
+
+template <typename dtype_t>
+__device__ __forceinline__ dtype_t broadcast(dtype_t& ptr, int src_lane_idx) {
+    EP_STATIC_ASSERT(sizeof(dtype_t) % sizeof(int) == 0, "");
+    auto send_int_values = reinterpret_cast<int*>(&ptr);
+    int recv_int_values[sizeof(dtype_t) / sizeof(int)];
+    #pragma unroll
+    for (int i = 0; i < sizeof(dtype_t) / sizeof(int); ++ i)
+        recv_int_values[i] = __shfl_sync(0xffffffff, send_int_values[i], src_lane_idx);
+    return *reinterpret_cast<dtype_t*>(recv_int_values);
+}
+
+__forceinline__ __device__ int get_lane_id() {
+    int lane_id;
+    asm("mov.s32 %0, %laneid;" : "=r"(lane_id));
+    return lane_id;
+}
+
+constexpr float kFP8Margin = 1e-4;
+constexpr float kFinfoAmaxE4M3 = 448.0f;
+constexpr float kFinfoAmaxInvE4M3 = 1 / 448.0f;
+
+__forceinline__ __device__ float fast_pow2(int x) {
+    // We can ensure `-126 <= x and x <= 127`
+    uint32_t bits_x = (x + 127) << 23;
+    return *reinterpret_cast<float*>(&bits_x);
+}
+
+__forceinline__ __device__ int fast_log2_ceil(float x) {
+    auto bits_x = *reinterpret_cast<uint32_t*>(&x);
+    auto exp_x = (bits_x >> 23) & 0xff;
+    auto man_bits = bits_x & ((1 << 23) - 1);
+    return exp_x - 127 + (man_bits != 0);
+}
+
+__forceinline__ __device__ void calculate_fp8_scales(float amax, float& scale, float& scale_inv, bool round_scale) {
+    if (round_scale) {
+        auto exp_scale_inv = fast_log2_ceil(amax * kFinfoAmaxInvE4M3);
+        scale = fast_pow2(-exp_scale_inv);
+        scale_inv = fast_pow2(exp_scale_inv);
+    } else {
+        scale_inv = amax * kFinfoAmaxInvE4M3;
+        scale = kFinfoAmaxE4M3 / amax;
+    }
+}
+
+template <bool kIsUE8M0, typename out_dtype_t = std::conditional_t<kIsUE8M0, uint8_t, float>>
+__forceinline__ __device__ out_dtype_t extract_required_scale_format(float value) {
+    if constexpr (kIsUE8M0) {
+        return static_cast<uint8_t>((*reinterpret_cast<uint32_t*>(&value)) >> 23);
+    } else {
+        return value;
+    }
+}
+
+// Operation functors
+template <typename T> struct ReduceSum { __device__ T operator()(T a, T b) const { return a + b; } };
+template <typename T> struct ReduceMax { __device__ T operator()(T a, T b) const { return a > b ? a : b; } };
+template <typename T> struct ReduceMin { __device__ T operator()(T a, T b) const { return a < b ? a : b; } };
+template <typename T> struct ReduceAnd { __device__ T operator()(T a, T b) const { return a & b; } };
+template <typename T> struct ReduceOr  { __device__ T operator()(T a, T b) const { return a | b; } };
+
+// Unified reduction function
+template <int kNumLanesPerGroup, bool kIntergroupReduce, typename T, typename Op>
+__forceinline__ __device__ T warp_reduce(T value, Op op) {
+    EP_STATIC_ASSERT(kNumLanesPerGroup == 32 or kNumLanesPerGroup == 16 or kNumLanesPerGroup == 8 or
+                     kNumLanesPerGroup ==  4 or kNumLanesPerGroup == 2  or kNumLanesPerGroup == 1,
+                     "Invalid number of lanes");
+    constexpr uint32_t mask = 0xffffffff;
+    if constexpr (kIntergroupReduce) {
+        if constexpr (kNumLanesPerGroup <=  1) value = op(value, __shfl_xor_sync(mask, value,  1));
+        if constexpr (kNumLanesPerGroup <=  2) value = op(value, __shfl_xor_sync(mask, value,  2));
+        if constexpr (kNumLanesPerGroup <=  4) value = op(value, __shfl_xor_sync(mask, value,  4));
+        if constexpr (kNumLanesPerGroup <=  8) value = op(value, __shfl_xor_sync(mask, value,  8));
+        if constexpr (kNumLanesPerGroup <= 16) value = op(value, __shfl_xor_sync(mask, value, 16));
+    } else {
+        if constexpr (kNumLanesPerGroup >= 32) value = op(value, __shfl_xor_sync(mask, value, 16));
+        if constexpr (kNumLanesPerGroup >= 16) value = op(value, __shfl_xor_sync(mask, value,  8));
+        if constexpr (kNumLanesPerGroup >=  8) value = op(value, __shfl_xor_sync(mask, value,  4));
+        if constexpr (kNumLanesPerGroup >=  4) value = op(value, __shfl_xor_sync(mask, value,  2));
+        if constexpr (kNumLanesPerGroup >=  2) value = op(value, __shfl_xor_sync(mask, value,  1));
+    }
+    return value;
+}
+
+// Convenience aliases
+template <int kNumLanesPerGroup = 32, bool kIntergroupReduce = false, typename T>
+__forceinline__ __device__ T warp_reduce_sum(T value) {
+    return warp_reduce<kNumLanesPerGroup, kIntergroupReduce, T>(value, ReduceSum<T>{});
+}
+
+template <int kNumLanesPerGroup = 32, bool kIntergroupReduce = false, typename T>
+__forceinline__ __device__ T warp_reduce_max(T value) {
+    return warp_reduce<kNumLanesPerGroup, kIntergroupReduce, T>(value, ReduceMax<T>{});
+}
+
+template <int kNumLanesPerGroup = 32, bool kIntergroupReduce = false, typename T>
+__forceinline__ __device__ T warp_reduce_min(T value) {
+    return warp_reduce<kNumLanesPerGroup, kIntergroupReduce, T>(value, ReduceMin<T>{});
+}
+
+template <int kNumLanesPerGroup = 32, bool kIntergroupReduce = false, typename T>
+__forceinline__ __device__ T warp_reduce_and(T value) {
+    return warp_reduce<kNumLanesPerGroup, kIntergroupReduce, T>(value, ReduceAnd<T>{});
+}
+
+template <int kNumLanesPerGroup = 32, bool kIntergroupReduce = false, typename T>
+__forceinline__ __device__ T warp_reduce_or(T value) {
+    return warp_reduce<kNumLanesPerGroup, kIntergroupReduce, T>(value, ReduceOr<T>{});
+}
+
+} // namespace nixl_ep
diff --git a/examples/device/ep/csrc/nixl_ep.cpp b/examples/device/ep/csrc/nixl_ep.cpp
new file mode 100644
index 0000000000..8b50313640
--- /dev/null
+++ b/examples/device/ep/csrc/nixl_ep.cpp
@@ -0,0 +1,1055 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * This file incorporates material from the DeepSeek project, licensed under the MIT License.
+ * The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+ *
+ * SPDX-License-Identifier: MIT AND Apache-2.0
+ *
+ * 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 <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/CUDADataType.h>
+#include <algorithm>
+#include <chrono>
+#include <cstddef>
+#include <cstdint>
+#include <cuda_runtime.h>
+#include <memory>
+#include <pybind11/functional.h>
+#include <torch/python.h>
+
+#include "nixl_ep.hpp"
+#include "kernels/api.cuh"
+#include "kernels/configs.cuh"
+#include <cstdio>
+#include <fstream>
+#include <unistd.h>
+#include <stdio.h>
+#include "nixl.h"
+#include <ifaddrs.h>
+#include <sys/socket.h>
+#include <netdb.h>
+#include <arpa/inet.h>
+#include <net/if.h>
+#include <sstream>
+
+#define NIXL_ETCD_WATCH_TIMEOUT std::chrono::microseconds(1000000000) // 1000 seconds
+
+#ifdef ENABLE_DEBUG_LOGS
+#define HOST_LOG_DEBUG(fmt, ...) printf("[DEBUG] " fmt "\n", ##__VA_ARGS__)
+#else
+#define HOST_LOG_DEBUG(...)
+#endif
+
+namespace nixl_ep {
+
+static void sleep_ms(int milliseconds) {
+    std::this_thread::sleep_for(std::chrono::milliseconds(milliseconds));
+}
+
+static std::string _get_local_ip() {
+    struct ifaddrs *ifaddr, *ifa;
+    char host[NI_MAXHOST];
+
+    if (getifaddrs(&ifaddr) == -1) {
+        perror("getifaddrs");
+        return "127.0.0.1";
+    }
+
+    for (ifa = ifaddr; ifa != NULL; ifa = ifa->ifa_next) {
+        if (ifa->ifa_addr == NULL || ifa->ifa_addr->sa_family != AF_INET)
+            continue;
+
+        if ((ifa->ifa_flags & IFF_UP) && !(ifa->ifa_flags & IFF_LOOPBACK)) {
+            if (getnameinfo(ifa->ifa_addr, sizeof(struct sockaddr_in), host, NI_MAXHOST, NULL, 0, NI_NUMERICHOST) == 0) {
+                freeifaddrs(ifaddr);
+                return std::string(host);
+            }
+        }
+    }
+
+    freeifaddrs(ifaddr);
+    return "127.0.0.1";
+}
+
+static std::string boot_id_get() {
+    std::ifstream boot_id_file("/proc/sys/kernel/random/boot_id");
+    if (!boot_id_file.is_open()) {
+        return "";
+    }
+
+    std::string boot_id;
+    std::getline(boot_id_file, boot_id);
+
+    if (!boot_id.empty() && boot_id.back() == '\n') {
+        boot_id.pop_back();
+    }
+
+    return boot_id;
+}
+
+static ino_t ipc_namespace_inode_get() {
+    struct stat st;
+    if (stat("/proc/self/ns/ipc", &st) != 0) {
+        return 0;
+    }
+    return st.st_ino;
+}
+
+void Buffer::update_memory_buffers(int num_ranks, int64_t num_rdma_bytes)
+{
+    if (!available) {
+        init(num_ranks, num_rdma_bytes);
+        available = true;
+    } else {
+        throw std::runtime_error("Multiple calls to update_memory_buffers are not supported");
+    }
+}
+
+Buffer::Buffer(int rank, bool explicitly_destroy, bool enable_shrink):
+        rank(rank), num_ranks(1),
+        explicitly_destroy(explicitly_destroy),
+        comm_stream(at::cuda::getStreamFromPool(true)),
+        dummy_src_dlist(VRAM_SEG),
+        enable_shrink(enable_shrink) {}
+
+void Buffer::init(int num_ranks, int64_t num_rdma_bytes)
+{
+    // Update buffer attributes
+    this->max_num_ranks = num_ranks;
+    this->num_rdma_bytes = num_rdma_bytes;
+
+    // Common checks
+    EP_HOST_ASSERT(num_rdma_bytes % NUM_BUFFER_ALIGNMENT_BYTES == 0);
+    EP_HOST_ASSERT(0 <= rank and rank < num_ranks);
+
+    // Get ranks
+    CUDA_CHECK(cudaGetDevice(&device_id));
+    // Get device info
+    cudaDeviceProp device_prop = {};
+    CUDA_CHECK(cudaGetDeviceProperties(&device_prop, device_id));
+    num_device_sms = device_prop.multiProcessorCount;
+
+    // Create 32 MiB workspace
+    CUDA_CHECK(cudaMalloc(&workspace, NUM_WORKSPACE_BYTES));
+    CUDA_CHECK(cudaMemsetAsync(workspace, 0, NUM_WORKSPACE_BYTES, comm_stream));
+
+    env_num_channels = std::getenv("NIXL_EP_NUM_CHANNELS") ? std::stoi(std::getenv("NIXL_EP_NUM_CHANNELS")) : 1;
+    EP_HOST_ASSERT(env_num_channels > 0);
+    num_counters = env_num_channels * max_num_ranks * 2 + max_num_ranks;
+    CUDA_CHECK(cudaMalloc(&rdma_buffer_ptr, num_rdma_bytes));
+    CUDA_CHECK(cudaMemset(rdma_buffer_ptr, 0, num_rdma_bytes));
+    CUDA_CHECK(cudaMalloc(&counters_buffer_ptr, num_counters * sizeof(uint64_t)));
+    CUDA_CHECK(cudaMemset(counters_buffer_ptr, 0, num_counters * sizeof(uint64_t)));
+    CUDA_CHECK(cudaMalloc(&wireup_buffer_ptr, sizeof(uint64_t)));
+    CUDA_CHECK(cudaMemset(wireup_buffer_ptr, 0, sizeof(uint64_t)));
+
+    /* Initialize dummy src dlist with a dummy device address */
+    dummy_src_dlist.addDesc(nixlBlobDesc((uintptr_t)counters_buffer_ptr, sizeof(uint64_t), device_id, ""));
+    // Allocate and clean shrink buffer
+    mask_buffer_ptr = nullptr;
+    sync_buffer_ptr = nullptr;
+    if (enable_shrink) {
+        int num_mask_buffer_bytes = max_num_ranks * sizeof(int);
+        CUDA_CHECK(cudaMalloc(&mask_buffer_ptr, num_mask_buffer_bytes));
+        CUDA_CHECK(cudaMemset(mask_buffer_ptr, 0xff, num_mask_buffer_bytes));
+        CUDA_CHECK(cudaMemset(mask_buffer_ptr + rank, 0, sizeof(int)));
+    }
+    int num_sync_buffer_bytes = max_num_ranks * sizeof(int);
+    CUDA_CHECK(cudaMalloc(&sync_buffer_ptr, num_sync_buffer_bytes));
+    CUDA_CHECK(cudaMemset(sync_buffer_ptr, 0, num_sync_buffer_bytes));
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    strncpy(my_peer_info.ip, _get_local_ip().c_str(), MAX_IP_LENGTH - 1);
+    my_peer_info.ip[MAX_IP_LENGTH - 1] = '\0';
+    my_peer_info.rdma_buffer_ptr = rdma_buffer_ptr;
+    my_peer_info.counters_buffer_ptr = counters_buffer_ptr;
+    my_peer_info.wireup_ptr = wireup_buffer_ptr;
+    my_peer_info.device_id = get_local_device_id();
+    my_peer_info.sync_buffer_ptr = sync_buffer_ptr;
+    my_peer_info.rank = rank;
+
+    // Create IPC handles for rdma buffer and counters
+    CUDA_CHECK(cudaIpcGetMemHandle(&my_peer_info.rdma_ipc_handle, rdma_buffer_ptr));
+    CUDA_CHECK(cudaIpcGetMemHandle(&my_peer_info.counters_ipc_handle, counters_buffer_ptr));
+
+    strncpy(my_peer_info.boot_id, boot_id_get().c_str(), MAX_BOOT_ID_LENGTH - 1);
+    my_peer_info.boot_id[MAX_BOOT_ID_LENGTH - 1] = '\0';
+    my_peer_info.ipc_namespace_inode = ipc_namespace_inode_get();
+
+    nixl_peer_info.resize(max_num_ranks);
+    nixl_peer_info[rank] = my_peer_info;
+
+    _nixl_agent_init();
+
+    _nixl_ep_init(std::vector<int>{rank});
+}
+
+Buffer::~Buffer() noexcept(false) {
+    if (not explicitly_destroy) {
+        destroy();
+    } else if (not destroyed) {
+        printf("WARNING: destroy() was not called before NIXL_EP buffer destruction, which can leak resources.\n");
+        fflush(stdout);
+    }
+}
+
+bool Buffer::is_available() const {
+    return available;
+}
+
+int Buffer::get_local_device_id() const {
+    return device_id;
+}
+
+torch::Tensor Buffer::get_local_buffer_tensor(const pybind11::object& dtype, int64_t offset) const {
+    torch::ScalarType casted_dtype = torch::python::detail::py_object_to_dtype(dtype);
+    auto element_bytes = static_cast<int64_t>(elementSize(casted_dtype));
+    auto base_ptr = static_cast<uint8_t*>(rdma_buffer_ptr) + offset;
+    return torch::from_blob(base_ptr, num_rdma_bytes / element_bytes, torch::TensorOptions().dtype(casted_dtype).device(at::kCUDA));
+}
+
+torch::Stream Buffer::get_comm_stream() const {
+    return comm_stream;
+}
+
+void Buffer::destroy() {
+    EP_HOST_ASSERT(not destroyed);
+
+    // Synchronize
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    cudaFree(counters_buffer_ptr);
+    cudaFree(wireup_buffer_ptr);
+    cudaFree(rdma_buffer_ptr);
+
+    if (nixl_agent_info and nixl_agent_info->agent != nullptr) {
+        nixl_agent_info->agent->invalidateLocalMD();
+    }
+
+    rdma_buffer_ptr = nullptr;
+    counters_buffer_ptr = nullptr;
+
+    if (enable_shrink) {
+        cudaFree(mask_buffer_ptr);
+        cudaFree(sync_buffer_ptr);
+    }
+
+    // Free workspace
+    CUDA_CHECK(cudaFree(workspace));
+
+    destroyed = true;
+    available = false;
+}
+
+void Buffer::barrier() {
+    auto compute_stream = at::cuda::getCurrentCUDAStream();
+    ep_kernels::barrier(nixl_ctx->gpu[0],mask_buffer_ptr, sync_buffer_ptr, compute_stream);
+}
+
+void Buffer::_nixl_agents_connect(const std::vector<int>& ranks) {
+    EP_HOST_ASSERT(!ranks.empty());
+
+    // Assuming ranks vector does not include current rank and has only new ranks
+    remote_ranks.insert(remote_ranks.end(), ranks.begin(), ranks.end());
+    for (int remote_rank : ranks) {
+        nixl_agent_info->remote_agent_names[remote_rank] = std::to_string(remote_rank);
+    }
+
+    for (int remote_rank : ranks) {
+        nixl_status_t fetch_status = nixl_agent_info->agent->fetchRemoteMD(nixl_agent_info->remote_agent_names[remote_rank]);
+        if (fetch_status != NIXL_SUCCESS) {
+            throw std::runtime_error("Failed to fetch metadata for remote agent " + std::to_string(remote_rank) +
+                                    ", status: " + std::to_string(fetch_status));
+        }
+
+        // Wait for remote metadata to be available
+        nixl_xfer_dlist_t empty_descs(VRAM_SEG);
+        while (nixl_agent_info->agent->checkRemoteMD(std::to_string(remote_rank), empty_descs) != NIXL_SUCCESS) {
+            sleep_ms(10);
+        }
+    }
+}
+
+void Buffer::_nixl_agents_peer_info_gather(std::vector<int>& ranks) {
+    for (int remote_rank : ranks) {
+        std::string my_peer_info_str(reinterpret_cast<const char*>(&my_peer_info), sizeof(NixlPeerInfo));
+        nixl_agent_info->agent->genNotif(std::to_string(remote_rank), my_peer_info_str);
+    }
+
+    for (int remote_rank : ranks) {
+        do {
+            nixl_notifs_t notif_map;
+            nixl_agent_info->agent->getNotifs(notif_map);
+            for (auto &notif : notif_map) {
+                std::string my_peer_info_str = notif.second[0];
+                NixlPeerInfo remote_peer_info;
+                memcpy(&remote_peer_info, my_peer_info_str.c_str(), sizeof(NixlPeerInfo));
+                nixl_peer_info[remote_peer_info.rank] = remote_peer_info;
+                nixl_agent_info->wire_up_done[remote_peer_info.rank] = true;
+            }
+        } while (!nixl_agent_info->wire_up_done[remote_rank]);
+    }
+}
+
+// This is a workaround to NIXL/UCX wireup issue and should be removed once it is fixed
+void Buffer::_nixl_agents_wireup(std::vector<int>& ranks) {
+    for (int worker_id = 0; worker_id < env_num_channels; worker_id++) {
+        for (int remote_rank : ranks) {
+            nixl_opt_args_t wireup_params = {};
+            wireup_params.backends.push_back(nixl_agent_info->backend);
+            wireup_params.customParam = "worker_id=" + std::to_string(worker_id);
+
+            nixlXferReqH *wireup_req = nullptr;
+            nixl_xfer_dlist_t dummy_dst_dlist(VRAM_SEG);
+            dummy_dst_dlist.addDesc(nixlBlobDesc((uintptr_t)nixl_peer_info[remote_rank].wireup_ptr, sizeof(uint64_t), nixl_peer_info[remote_rank].device_id, ""));
+            EP_HOST_ASSERT(nixl_agent_info->agent->createXferReq(
+                NIXL_WRITE, dummy_src_dlist, dummy_dst_dlist,
+                nixl_agent_info->remote_agent_names[remote_rank], wireup_req, &wireup_params) == NIXL_SUCCESS);
+
+            nixl_status_t status = nixl_agent_info->agent->postXferReq(wireup_req);
+            EP_HOST_ASSERT(status == NIXL_SUCCESS || status == NIXL_IN_PROG);
+
+            while ((status = nixl_agent_info->agent->getXferStatus(wireup_req)) == NIXL_IN_PROG) {
+                sleep_ms(1);
+            }
+
+            EP_HOST_ASSERT(status == NIXL_SUCCESS);
+            EP_HOST_ASSERT(nixl_agent_info->agent->releaseXferReq(wireup_req) == NIXL_SUCCESS);
+        }
+    }
+}
+
+void Buffer::_nixl_ep_barrier_buffer_clear() {
+    CUDA_CHECK(cudaMemset(sync_buffer_ptr, 0, max_num_ranks * sizeof(int)));
+}
+
+void Buffer::connect_ranks(const std::vector<int>& remote_ranks_list) {
+    EP_HOST_ASSERT(!remote_ranks_list.empty());
+    std::vector<int> new_ranks;
+    int max_added_rank = std::max(rank, *std::max_element(remote_ranks_list.begin(), remote_ranks_list.end()));
+    num_ranks = std::max(num_ranks, max_added_rank + 1);
+
+    for (int remote_rank : remote_ranks_list) {
+        // Skip self and ranks we are already connected to
+        if (remote_rank == rank or std::find(remote_ranks.begin(), remote_ranks.end(), remote_rank) != remote_ranks.end())
+            continue;
+
+        new_ranks.push_back(remote_rank);
+        CUDA_CHECK(cudaMemset(mask_buffer_ptr + remote_rank, 0, sizeof(int))); // Reset mask buffer for new ranks
+    }
+
+    if (new_ranks.empty())
+        return;
+
+    _nixl_ep_barrier_buffer_clear();
+
+    _nixl_agents_connect(new_ranks);
+
+    _nixl_agents_peer_info_gather(new_ranks);
+
+    _nixl_agents_wireup(new_ranks);
+
+    _nixl_ep_init(new_ranks);
+
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    // Ready to use
+    available = true;
+}
+
+void Buffer::disconnect_ranks(const std::vector<int>& remote_ranks_list) {
+    EP_HOST_ASSERT(!remote_ranks_list.empty());
+    EP_HOST_ASSERT(remote_ranks_list.size() <= remote_ranks.size());
+
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    // Update mask buffer to mark ranks as inactive
+    for (int removed_rank : remote_ranks_list) {
+        update_mask_buffer(removed_rank, true);  // mask=true
+    }
+
+    _nixl_ep_cleanup(remote_ranks_list);
+
+    _nixl_agents_peer_info_cleanup(remote_ranks_list);
+
+    _nixl_agents_disconnect(remote_ranks_list);
+
+    // Remove ranks from remote_ranks vector (arbitrary order)
+    for (int removed_rank : remote_ranks_list) {
+        remote_ranks.erase(
+            std::remove(remote_ranks.begin(), remote_ranks.end(), removed_rank),
+            remote_ranks.end()
+        );
+    }
+
+    int max_rank = rank;  // Include self
+    if (!remote_ranks.empty()) {
+        max_rank = std::max(max_rank,
+                           *std::max_element(remote_ranks.begin(), remote_ranks.end()));
+    }
+    num_ranks = max_rank + 1;  // Sparse indexing maintained
+}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>, torch::Tensor, torch::Tensor, torch::Tensor, std::optional<EventHandle>, std::optional<std::function<void()>>>
+Buffer::dispatch(const torch::Tensor& x, const torch::Tensor& topk_idx,
+                             const std::optional<torch::Tensor>& cumulative_local_expert_recv_stats,
+                             const std::optional<torch::Tensor>& dispatch_wait_recv_cost_stats,
+                             int num_max_dispatch_tokens_per_rank, int num_experts,
+                             bool use_fp8, bool round_scale, bool use_ue8m0,
+                             bool async, bool return_recv_hook) {
+    // Tensor checks
+    // By default using `ptp128c` FP8 cast
+    EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous() and x.scalar_type() == torch::kBFloat16);
+    EP_HOST_ASSERT(x.size(1) % sizeof(int4) == 0 and x.size(1) % 128 == 0);
+    EP_HOST_ASSERT(topk_idx.dim() == 2 and topk_idx.is_contiguous());
+    EP_HOST_ASSERT(x.size(0) == topk_idx.size(0) and x.size(0) <= num_max_dispatch_tokens_per_rank);
+    EP_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt64);
+    EP_HOST_ASSERT(num_experts % num_ranks == 0);
+
+    // Diagnosis tensors
+    if (cumulative_local_expert_recv_stats.has_value()) {
+        EP_HOST_ASSERT(cumulative_local_expert_recv_stats->scalar_type() == torch::kInt);
+        EP_HOST_ASSERT(cumulative_local_expert_recv_stats->dim() == 1 and cumulative_local_expert_recv_stats->is_contiguous());
+        EP_HOST_ASSERT(cumulative_local_expert_recv_stats->size(0) == num_experts / num_ranks);
+    }
+    if (dispatch_wait_recv_cost_stats.has_value()) {
+        EP_HOST_ASSERT(dispatch_wait_recv_cost_stats->scalar_type() == torch::kInt64);
+        EP_HOST_ASSERT(dispatch_wait_recv_cost_stats->dim() == 1 and dispatch_wait_recv_cost_stats->is_contiguous());
+        EP_HOST_ASSERT(dispatch_wait_recv_cost_stats->size(0) == num_ranks);
+    }
+
+    auto num_tokens = static_cast<int>(x.size(0)), hidden = static_cast<int>(x.size(1));
+    auto num_topk = static_cast<int>(topk_idx.size(1));
+    int num_local_experts = num_experts / num_ranks;
+
+    // Buffer control
+    EPLayout layout(rdma_buffer_ptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts);
+    EP_HOST_ASSERT(layout.total_bytes <= num_rdma_bytes);
+    ep_kernels::gpu_nixl_ctx gpu_nixl_ctx = nixl_ctx->gpu[buffer_idx];
+    auto buffer = layout.buffers[buffer_idx];
+    auto next_buffer = layout.buffers[buffer_idx ^= 1];
+
+    // Wait previous tasks to be finished
+    // NOTES: the hook mode will always use the default stream
+    auto compute_stream = at::cuda::getCurrentCUDAStream();
+    auto launch_stream = return_recv_hook ? compute_stream : comm_stream;
+    EP_HOST_ASSERT(not (async and return_recv_hook));
+    if (not return_recv_hook)
+        stream_wait(launch_stream, compute_stream);
+
+    // Allocate packed tensors
+    auto packed_recv_x = torch::empty({num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank, hidden},
+                                      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));
+    auto packed_recv_layout_range = torch::empty({num_local_experts, num_ranks}, torch::dtype(torch::kInt64).device(torch::kCUDA));
+    auto packed_recv_count = torch::empty({num_local_experts}, torch::dtype(torch::kInt32).device(torch::kCUDA));
+
+    // Allocate column-majored scales
+    auto packed_recv_x_scales = std::optional<torch::Tensor>();
+    void* packed_recv_x_scales_ptr = nullptr;
+    EP_HOST_ASSERT((num_ranks * num_max_dispatch_tokens_per_rank) % 4 == 0 and "TMA requires the number of tokens to be multiple of 4");
+
+    if (use_fp8) {
+        // TODO: support unaligned cases
+        EP_HOST_ASSERT(hidden % 512 == 0);
+        if (not use_ue8m0) {
+            packed_recv_x_scales = torch::empty({num_local_experts, hidden / 128, num_ranks * num_max_dispatch_tokens_per_rank},
+                                                torch::dtype(torch::kFloat32).device(torch::kCUDA));
+        } else {
+            EP_HOST_ASSERT(round_scale);
+            packed_recv_x_scales = torch::empty({num_local_experts, hidden / 512, num_ranks * num_max_dispatch_tokens_per_rank},
+                                                torch::dtype(torch::kInt).device(torch::kCUDA));
+        }
+        packed_recv_x_scales = torch::transpose(packed_recv_x_scales.value(), 1, 2);
+        packed_recv_x_scales_ptr = packed_recv_x_scales->data_ptr();
+    }
+
+    // Kernel launch
+    auto next_clean_meta = next_buffer.clean_meta();
+    auto launcher = [=](int phases) {
+        ep_kernels::dispatch(packed_recv_x.data_ptr(), packed_recv_x_scales_ptr,
+                               packed_recv_src_info.data_ptr<int>(), packed_recv_layout_range.data_ptr<int64_t>(),
+                               packed_recv_count.data_ptr<int>(),
+                               mask_buffer_ptr,
+                               cumulative_local_expert_recv_stats.has_value() ? cumulative_local_expert_recv_stats->data_ptr<int>() : nullptr,
+                               dispatch_wait_recv_cost_stats.has_value() ? dispatch_wait_recv_cost_stats->data_ptr<int64_t>() : nullptr,
+                               buffer.dispatch_rdma_recv_data_buffer, buffer.dispatch_rdma_recv_count_buffer,
+                               buffer.dispatch_rdma_send_buffer,
+                               x.data_ptr(), topk_idx.data_ptr<int64_t>(),
+                               next_clean_meta.first, next_clean_meta.second,
+                               num_tokens, hidden, num_max_dispatch_tokens_per_rank,
+                               num_topk, num_experts, rank, num_ranks,
+                               use_fp8, round_scale, use_ue8m0,
+                               workspace, num_device_sms,
+                               launch_stream, phases, gpu_nixl_ctx);
+    };
+    launcher(return_recv_hook ? EP_SEND_PHASE : (EP_SEND_PHASE | EP_RECV_PHASE));
+
+    // Wait streams
+    std::optional<EventHandle> event;
+    if (async) {
+        // NOTES: we must ensure the all tensors will not be deallocated before the stream-wait happens,
+        // so in Python API, we must wrap all tensors into the event handle.
+        event = EventHandle(launch_stream);
+    } else if (not return_recv_hook) {
+        stream_wait(compute_stream, launch_stream);
+    }
+
+    // Receiver callback
+    std::optional<std::function<void()>> recv_hook = std::nullopt;
+    if (return_recv_hook)
+        recv_hook = [=]() { launcher(EP_RECV_PHASE); };
+
+    // Return values
+    return {packed_recv_x, packed_recv_x_scales, packed_recv_count, packed_recv_src_info, packed_recv_layout_range, event, recv_hook};
+}
+
+std::tuple<torch::Tensor, std::optional<EventHandle>, std::optional<std::function<void()>>>
+Buffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx, const torch::Tensor& topk_weights,
+                            const torch::Tensor& src_info, const torch::Tensor& layout_range,
+                            const std::optional<torch::Tensor>& combine_wait_recv_cost_stats,
+                            int num_max_dispatch_tokens_per_rank, int num_experts,
+                            bool use_logfmt, bool zero_copy, bool async, bool return_recv_hook,
+                            const std::optional<torch::Tensor>& out) {
+    // Tensor checks
+    EP_HOST_ASSERT(x.dim() == 3 and x.is_contiguous() and x.scalar_type() == torch::kBFloat16);
+    EP_HOST_ASSERT(x.size(0) == num_experts / num_ranks);
+    EP_HOST_ASSERT(x.size(1) == num_ranks * num_max_dispatch_tokens_per_rank);
+    EP_HOST_ASSERT(x.size(2) % sizeof(int4) == 0 and x.size(2) % 128 == 0);
+    EP_HOST_ASSERT(topk_idx.dim() == 2 and topk_idx.is_contiguous());
+    EP_HOST_ASSERT(topk_idx.size(0) == topk_weights.size(0) and topk_idx.size(1) == topk_weights.size(1));
+    EP_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt64);
+    EP_HOST_ASSERT(topk_weights.dim() == 2 and topk_weights.is_contiguous());
+    EP_HOST_ASSERT(topk_weights.size(0) <= num_max_dispatch_tokens_per_rank);
+    EP_HOST_ASSERT(topk_weights.scalar_type() == torch::kFloat32);
+    EP_HOST_ASSERT(src_info.dim() == 2 and src_info.is_contiguous());
+    EP_HOST_ASSERT(src_info.scalar_type() == torch::kInt32 and x.size(0) == src_info.size(0));
+    EP_HOST_ASSERT(layout_range.dim() == 2 and layout_range.is_contiguous());
+    EP_HOST_ASSERT(layout_range.scalar_type() == torch::kInt64);
+    EP_HOST_ASSERT(layout_range.size(0) == num_experts / num_ranks and layout_range.size(1) == num_ranks);
+
+    if (combine_wait_recv_cost_stats.has_value()) {
+        EP_HOST_ASSERT(combine_wait_recv_cost_stats->scalar_type() == torch::kInt64);
+        EP_HOST_ASSERT(combine_wait_recv_cost_stats->dim() == 1 and combine_wait_recv_cost_stats->is_contiguous());
+        EP_HOST_ASSERT(combine_wait_recv_cost_stats->size(0) == num_ranks);
+    }
+
+    auto hidden = static_cast<int>(x.size(2));
+    auto num_topk = static_cast<int>(topk_weights.size(1));
+    auto num_combined_tokens = static_cast<int>(topk_weights.size(0));
+
+    // Buffer control
+    EPLayout layout(rdma_buffer_ptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts);
+    EP_HOST_ASSERT(layout.total_bytes <= num_rdma_bytes);
+    ep_kernels::gpu_nixl_ctx gpu_nixl_ctx = nixl_ctx->gpu[buffer_idx];
+    auto buffer = layout.buffers[buffer_idx];
+    auto next_buffer = layout.buffers[buffer_idx ^= 1];
+
+    // Wait previous tasks to be finished
+    // NOTES: the hook mode will always use the default stream
+    auto compute_stream = at::cuda::getCurrentCUDAStream();
+    auto launch_stream = return_recv_hook ? compute_stream : comm_stream;
+    EP_HOST_ASSERT(not (async and return_recv_hook));
+    if (not return_recv_hook)
+        stream_wait(launch_stream, compute_stream);
+
+    // Allocate output tensor
+    torch::Tensor combined_x;
+    if (out.has_value()) {
+        EP_HOST_ASSERT(out->dim() == 2 and out->is_contiguous());
+        EP_HOST_ASSERT(out->size(0) == num_combined_tokens and out->size(1) == hidden);
+        EP_HOST_ASSERT(out->scalar_type() == x.scalar_type());
+        combined_x = out.value();
+    } else {
+        combined_x = torch::empty({num_combined_tokens, hidden}, x.options());
+    }
+
+    // Kernel launch
+    auto next_clean_meta = next_buffer.clean_meta();
+    auto launcher = [=](int phases) {
+        ep_kernels::combine(combined_x.data_ptr(),
+                              buffer.combine_rdma_recv_data_buffer, buffer.combine_rdma_recv_flag_buffer,
+                              buffer.combine_rdma_send_buffer,
+                              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>(),
+                              mask_buffer_ptr,
+                              combine_wait_recv_cost_stats.has_value() ? combine_wait_recv_cost_stats->data_ptr<int64_t>() : nullptr,
+                              next_clean_meta.first, next_clean_meta.second,
+                              num_combined_tokens, hidden, num_max_dispatch_tokens_per_rank,
+                              num_topk, num_experts, rank, num_ranks,
+                              use_logfmt,
+                              workspace, num_device_sms,
+                              launch_stream, phases, zero_copy, gpu_nixl_ctx);
+    };
+    launcher(return_recv_hook ? EP_SEND_PHASE : (EP_SEND_PHASE | EP_RECV_PHASE));
+
+    // Wait streams
+    std::optional<EventHandle> event;
+    if (async) {
+        // NOTES: we must ensure the all tensors will not be deallocated before the stream-wait happens,
+        // so in Python API, we must wrap all tensors into the event handle.
+        event = EventHandle(launch_stream);
+    } else if (not return_recv_hook) {
+        stream_wait(compute_stream, launch_stream);
+    }
+
+    // Receiver callback
+    std::optional<std::function<void()>> recv_hook = std::nullopt;
+    if (return_recv_hook)
+        recv_hook = [=]() { launcher(EP_RECV_PHASE); };
+
+    // Return values
+    return {combined_x, event, recv_hook};
+}
+
+torch::Tensor
+Buffer::get_next_combine_buffer(int num_max_dispatch_tokens_per_rank, int hidden, int num_experts) const {
+    EPLayout layout(rdma_buffer_ptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts);
+
+    auto buffer = layout.buffers[buffer_idx];
+    auto dtype = torch::kBFloat16;
+    auto num_msg_elems = static_cast<int>(buffer.num_bytes_per_combine_msg / elementSize(torch::kBFloat16));
+
+    EP_HOST_ASSERT(buffer.num_bytes_per_combine_msg % elementSize(torch::kBFloat16) == 0);
+    return torch::from_blob(buffer.combine_rdma_send_buffer_data_start,
+                            {num_experts / num_ranks, num_ranks * num_max_dispatch_tokens_per_rank, hidden},
+                            {num_ranks * num_max_dispatch_tokens_per_rank * num_msg_elems, num_msg_elems, 1},
+                            torch::TensorOptions().dtype(dtype).device(torch::kCUDA));
+}
+
+bool is_sm90_compiled() {
+#ifndef DISABLE_SM90_FEATURES
+    return true;
+#else
+    return false;
+#endif
+}
+
+void Buffer::update_mask_buffer(int rank_to_mask, bool mask) {
+    EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled");
+    EP_HOST_ASSERT(rank_to_mask >= 0 and rank_to_mask < max_num_ranks);
+    ep_kernels::update_mask_buffer(mask_buffer_ptr, rank_to_mask, mask, at::cuda::getCurrentCUDAStream());
+}
+
+void Buffer::query_mask_buffer(const torch::Tensor& mask_status) {
+    EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled");
+    EP_HOST_ASSERT(mask_status.numel() == max_num_ranks && mask_status.scalar_type() == torch::kInt32);
+
+    ep_kernels::query_mask_buffer(mask_buffer_ptr, max_num_ranks,
+                                    reinterpret_cast<int*>(mask_status.data_ptr()),
+                                    at::cuda::getCurrentCUDAStream());
+}
+
+void Buffer::clean_mask_buffer() {
+    EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled");
+    ep_kernels::clean_mask_buffer(mask_buffer_ptr, max_num_ranks, at::cuda::getCurrentCUDAStream());
+}
+
+void Buffer::_nixl_ep_gpu_ctx_update() {
+    int num_local_experts = env_num_channels;
+
+    /* Initialize local counter arrays */
+    nixl_ctx->gpu[0].local_counters = counters_buffer_ptr;
+    nixl_ctx->gpu[1].local_counters = counters_buffer_ptr + max_num_ranks * num_local_experts;
+
+    /* Each context cleans the counters of the other context */
+    nixl_ctx->gpu[0].clean_counters = nixl_ctx->gpu[1].local_counters;
+    nixl_ctx->gpu[1].clean_counters = nixl_ctx->gpu[0].local_counters;
+
+    /* Copy remote counter reqs to device */
+    if (nixl_ctx->gpu[0].remote_counter_reqs == nullptr) {
+        CUDA_CHECK(cudaMalloc(&nixl_ctx->gpu[0].remote_counter_reqs, num_local_experts * max_num_ranks * sizeof(nixlGpuXferReqH)));
+        CUDA_CHECK(cudaMalloc(&nixl_ctx->gpu[1].remote_counter_reqs, num_local_experts * max_num_ranks * sizeof(nixlGpuXferReqH)));
+    }
+
+    // Always copy the updated arrays to GPU (since new ranks may have been added)
+    CUDA_CHECK(cudaMemcpy(nixl_ctx->gpu[0].remote_counter_reqs, nixl_ctx->gpu_remote_counter_reqs_0.data(), num_local_experts * max_num_ranks * sizeof(nixlGpuXferReqH), cudaMemcpyHostToDevice));
+    CUDA_CHECK(cudaMemcpy(nixl_ctx->gpu[1].remote_counter_reqs, nixl_ctx->gpu_remote_counter_reqs_1.data(), num_local_experts * max_num_ranks * sizeof(nixlGpuXferReqH), cudaMemcpyHostToDevice));
+
+    /* Copy batch reqs to device */
+    if (nixl_ctx->gpu[0].batch_reqs == nullptr)
+        CUDA_CHECK(cudaMalloc(&nixl_ctx->gpu[0].batch_reqs, num_local_experts * max_num_ranks * sizeof(nixlGpuXferReqH)));
+
+    // Always copy the updated batch arrays to GPU (since new ranks may have been added)
+    for (int dest_expert_idx = 0; dest_expert_idx < num_local_experts; dest_expert_idx++)
+        CUDA_CHECK(cudaMemcpy(nixl_ctx->gpu[0].batch_reqs + dest_expert_idx * max_num_ranks, nixl_ctx->gpu_batch_reqs[dest_expert_idx].data(), max_num_ranks * sizeof(nixlGpuXferReqH), cudaMemcpyHostToDevice));
+    nixl_ctx->gpu[1].batch_reqs = nixl_ctx->gpu[0].batch_reqs; // Both contexts share the same batch handles, no need to duplicate them
+
+    if (nixl_ctx->gpu[0].remote_barrier_reqs == nullptr)
+        CUDA_CHECK(cudaMalloc(&nixl_ctx->gpu[0].remote_barrier_reqs, num_local_experts * max_num_ranks * sizeof(nixlGpuXferReqH*)));
+
+    for (int dest_expert_idx = 0; dest_expert_idx < num_local_experts; dest_expert_idx++) {
+        CUDA_CHECK(cudaMemcpy(nixl_ctx->gpu[0].remote_barrier_reqs + dest_expert_idx * max_num_ranks, nixl_ctx->gpu_barrier_reqs[dest_expert_idx].data(), max_num_ranks * sizeof(nixlGpuXferReqH), cudaMemcpyHostToDevice));
+    }
+    nixl_ctx->gpu[1].remote_barrier_reqs = nixl_ctx->gpu[0].remote_barrier_reqs; // Both contexts share the same batch handles, no need to duplicate them
+
+    /* Initialize counters P2P pointers */
+    if (nixl_ctx->gpu[0].counters_p2p_ptrs == nullptr)
+        CUDA_CHECK(cudaMalloc(&nixl_ctx->gpu[0].counters_p2p_ptrs, max_num_ranks * sizeof(uint64_t *)));
+
+    if (nixl_ctx->gpu[1].counters_p2p_ptrs == nullptr)
+        CUDA_CHECK(cudaMalloc(&nixl_ctx->gpu[1].counters_p2p_ptrs, max_num_ranks * sizeof(uint64_t *)));
+
+    CUDA_CHECK(cudaMemcpy(nixl_ctx->gpu[0].counters_p2p_ptrs, nixl_ctx->counters_p2p_ptrs.data(), num_ranks * sizeof(uint64_t *), cudaMemcpyHostToDevice));
+    for (int i = 0; i < num_ranks; i++) if (nixl_ctx->counters_p2p_ptrs[i] != 0) nixl_ctx->counters_p2p_ptrs[i] += max_num_ranks * num_local_experts;
+    CUDA_CHECK(cudaMemcpy(nixl_ctx->gpu[1].counters_p2p_ptrs, nixl_ctx->counters_p2p_ptrs.data(), num_ranks * sizeof(uint64_t *), cudaMemcpyHostToDevice));
+    for (int i = 0; i < num_ranks; i++) if (nixl_ctx->counters_p2p_ptrs[i] != 0) nixl_ctx->counters_p2p_ptrs[i] -= max_num_ranks * num_local_experts;
+
+    /* Initialize RDMA P2P pointers */
+    if (nixl_ctx->gpu[0].rdma_p2p_ptrs == nullptr)
+        CUDA_CHECK(cudaMalloc(&nixl_ctx->gpu[0].rdma_p2p_ptrs, max_num_ranks * sizeof(void *)));
+
+    if (nixl_ctx->gpu[1].rdma_p2p_ptrs == nullptr)
+        CUDA_CHECK(cudaMalloc(&nixl_ctx->gpu[1].rdma_p2p_ptrs, max_num_ranks * sizeof(void *)));
+
+    CUDA_CHECK(cudaMemcpy(nixl_ctx->gpu[0].rdma_p2p_ptrs, nixl_ctx->rdma_p2p_ptrs.data(), num_ranks * sizeof(void *), cudaMemcpyHostToDevice));
+    CUDA_CHECK(cudaMemcpy(nixl_ctx->gpu[1].rdma_p2p_ptrs, nixl_ctx->rdma_p2p_ptrs.data(), num_ranks * sizeof(void *), cudaMemcpyHostToDevice));
+
+    /* Initialize info fields */
+    nixl_ctx->gpu[0].rdma_buffer_ptr = rdma_buffer_ptr;
+    nixl_ctx->gpu[1].rdma_buffer_ptr = rdma_buffer_ptr;
+    nixl_ctx->gpu[0].num_local_experts = num_local_experts;
+    nixl_ctx->gpu[1].num_local_experts = num_local_experts;
+    nixl_ctx->gpu[0].local_barrier_buffer = sync_buffer_ptr;
+    nixl_ctx->gpu[1].local_barrier_buffer = sync_buffer_ptr;
+    nixl_ctx->gpu[0].num_ranks = max_num_ranks;
+    nixl_ctx->gpu[1].num_ranks = max_num_ranks;
+    nixl_ctx->gpu[0].rank = rank;
+    nixl_ctx->gpu[1].rank = rank;
+}
+
+void Buffer::_nixl_ep_context_init() {
+    int num_local_experts = env_num_channels;
+
+    nixl_ctx = std::make_unique<nixl_ep_ctx>();
+    nixl_ctx->cpu_remote_counter_reqs_0.resize(num_local_experts * max_num_ranks);
+    nixl_ctx->cpu_remote_counter_reqs_1.resize(num_local_experts * max_num_ranks);
+    nixl_ctx->gpu_remote_counter_reqs_0.resize(num_local_experts * max_num_ranks);
+    nixl_ctx->gpu_remote_counter_reqs_1.resize(num_local_experts * max_num_ranks);
+    nixl_ctx->gpu_batch_reqs.resize(num_local_experts, std::vector<nixlGpuXferReqH>(max_num_ranks));
+    nixl_ctx->cpu_batch_reqs.resize(num_local_experts, std::vector<nixlXferReqH*>(max_num_ranks));
+    nixl_ctx->cpu_barrier_reqs.resize(num_local_experts, std::vector<nixlXferReqH*>(max_num_ranks));
+    nixl_ctx->gpu_barrier_reqs.resize(num_local_experts, std::vector<nixlGpuXferReqH>(max_num_ranks));
+    nixl_ctx->rdma_p2p_ptrs.resize(max_num_ranks);
+    nixl_ctx->counters_p2p_ptrs.resize(max_num_ranks);
+}
+
+void Buffer::_nixl_ep_init(const std::vector<int>& ranks) {
+    EP_EXECUTE_ONCE(_nixl_ep_context_init());
+    _nixl_ep_counters_prepare(ranks);
+    _nixl_ep_batches_prepare(ranks);
+    _nixl_ep_p2p_ptrs_prepare(ranks);
+    _nixl_ep_gpu_ctx_update();
+}
+
+void Buffer::_nixl_agent_init() {
+    std::string agent_name = std::to_string(rank);
+    nixlAgentConfig cfg(true, false, 0,
+                        nixl_thread_sync_t::NIXL_THREAD_SYNC_RW, 1, 0, 100000, false, NIXL_ETCD_WATCH_TIMEOUT);
+    auto agent = std::make_shared<nixlAgent>(agent_name, cfg);
+
+    // Create UCX backend
+    nixl_mem_list_t mems;
+    nixl_b_params_t init_params;
+
+    nixl_status_t status = agent->getPluginParams("UCX", mems, init_params);
+    if (status != NIXL_SUCCESS) {
+        throw std::runtime_error("Failed to get UCX plugin parameters for agent " + agent_name +
+                                ", status: " + std::to_string(status));
+    }
+
+    // Set UCX-specific parameters
+    init_params["ucx_error_handling_mode"] = "none";
+    init_params["num_workers"] = std::to_string(env_num_channels);
+
+    nixlBackendH* ucx_backend = nullptr;
+    status = agent->createBackend("UCX", init_params, ucx_backend);
+    if (status != NIXL_SUCCESS || !ucx_backend) {
+        throw std::runtime_error("Failed to create UCX backend for agent " + agent_name +
+                                ", status: " + std::to_string(status));
+    }
+
+    nixl_agent_info = std::make_unique<NixlAgentInfo>(agent, ucx_backend, max_num_ranks);
+    nixl_agent_info->extra_params.backends.push_back(ucx_backend);
+    nixl_agent_info->agent_name = agent_name;
+
+    /* Register RDMA buffer */
+    nixl_reg_dlist_t rdma_ptr_dlist(VRAM_SEG);
+    rdma_ptr_dlist.addDesc(nixlBlobDesc((uintptr_t)(rdma_buffer_ptr), num_rdma_bytes, get_local_device_id(), ""));
+    EP_HOST_ASSERT(agent->registerMem(rdma_ptr_dlist) == NIXL_SUCCESS);
+
+    /* Register counters buffer */
+    nixl_reg_dlist_t counters_dlist(VRAM_SEG);
+    counters_dlist.addDesc(nixlBlobDesc((uintptr_t)(counters_buffer_ptr), num_counters * sizeof(uint64_t), get_local_device_id(), ""));
+    EP_HOST_ASSERT(agent->registerMem(counters_dlist) == NIXL_SUCCESS);
+
+    /* Register sync buffer */
+    if (sync_buffer_ptr) {
+        nixl_reg_dlist_t sync_dlist(VRAM_SEG);
+        sync_dlist.addDesc(nixlBlobDesc((uintptr_t)(sync_buffer_ptr), max_num_ranks * sizeof(int), get_local_device_id(), ""));
+        EP_HOST_ASSERT(agent->registerMem(sync_dlist) == NIXL_SUCCESS);
+    }
+
+    size_t signal_size = 0;
+    EP_HOST_ASSERT(nixl_agent_info->agent->getGpuSignalSize(signal_size, &nixl_agent_info->extra_params) == NIXL_SUCCESS);
+    EP_HOST_ASSERT(signal_size == sizeof(uint64_t));
+    EP_HOST_ASSERT(agent->prepGpuSignal(counters_dlist, &nixl_agent_info->extra_params) == NIXL_SUCCESS);
+
+    /* Register wireup buffer */
+    nixl_reg_dlist_t wireup_dlist(VRAM_SEG);
+    wireup_dlist.addDesc(nixlBlobDesc((uintptr_t)(wireup_buffer_ptr), sizeof(uint64_t), get_local_device_id(), ""));
+    EP_HOST_ASSERT(agent->registerMem(wireup_dlist) == NIXL_SUCCESS);
+
+    // Send local metadata
+    status = nixl_agent_info->agent->sendLocalMD();
+    if (status != NIXL_SUCCESS) {
+        throw std::runtime_error("Failed to send local metadata for agent " +
+                                nixl_agent_info->agent_name + ", status: " + std::to_string(status));
+    }
+}
+
+void Buffer::_nixl_ep_batches_prepare(const std::vector<int>& ranks) {
+    nixl_status_t status;
+
+    for (int i = 0; i < env_num_channels; ++i) {
+        for (int j : ranks) {
+            if (j == rank) continue; // Skip self
+            if (nixl_ctx->gpu_batch_reqs[i][j]) continue; // Skip if already exported
+            nixl_xfer_dlist_t src_vram(VRAM_SEG);
+            src_vram.addDesc(nixlBlobDesc((uintptr_t)(rdma_buffer_ptr), num_rdma_bytes, get_local_device_id(), ""));
+            nixl_xfer_dlist_t dst_vram(VRAM_SEG);
+            dst_vram.addDesc(nixlBlobDesc((uintptr_t)(nixl_peer_info[j].rdma_buffer_ptr), num_rdma_bytes, nixl_peer_info[j].device_id, ""));
+            nixl_opt_args_t extra_params = {};
+            extra_params.backends.push_back(nixl_agent_info->backend);
+            extra_params.customParam = "worker_id=" + std::to_string(i);
+            status = nixl_agent_info->agent->createXferReq(NIXL_WRITE, src_vram, dst_vram, nixl_agent_info->remote_agent_names[j], nixl_ctx->cpu_batch_reqs[i][j], &extra_params);
+            EP_HOST_ASSERT(status == NIXL_SUCCESS);
+            EP_HOST_ASSERT(nixl_agent_info->agent->createGpuXferReq(*nixl_ctx->cpu_batch_reqs[i][j], nixl_ctx->gpu_batch_reqs[i][j]) == NIXL_SUCCESS);
+
+            nixl_xfer_dlist_t src_vram_ll(VRAM_SEG);
+            src_vram_ll.addDesc(nixlBlobDesc((uintptr_t)(sync_buffer_ptr), max_num_ranks * sizeof(int), get_local_device_id(), ""));
+            nixl_xfer_dlist_t dst_vram_ll(VRAM_SEG);
+            dst_vram_ll.addDesc(nixlBlobDesc((uintptr_t)(nixl_peer_info[j].sync_buffer_ptr), max_num_ranks * sizeof(int), nixl_peer_info[j].device_id, ""));
+            EP_HOST_ASSERT(nixl_agent_info->agent->createXferReq(NIXL_WRITE, src_vram_ll, dst_vram_ll, nixl_agent_info->remote_agent_names[j], nixl_ctx->cpu_barrier_reqs[i][j], &extra_params) == NIXL_SUCCESS);
+            EP_HOST_ASSERT(nixl_agent_info->agent->createGpuXferReq(*nixl_ctx->cpu_barrier_reqs[i][j], nixl_ctx->gpu_barrier_reqs[i][j]) == NIXL_SUCCESS);
+        }
+    }
+}
+
+void Buffer::_nixl_ep_p2p_ptrs_prepare(const std::vector<int>& ranks) {
+    for (int i : ranks) {
+        if (i == rank) {
+            nixl_ctx->rdma_p2p_ptrs[i] = rdma_buffer_ptr;
+            nixl_ctx->counters_p2p_ptrs[i] = counters_buffer_ptr;
+        } else if (std::string(nixl_peer_info[i].boot_id) == std::string(my_peer_info.boot_id) &&
+                   nixl_peer_info[i].ipc_namespace_inode == my_peer_info.ipc_namespace_inode &&
+                   std::string(std::getenv("NIXL_EP_NVLINK_BACKEND_IPC")) == "1") {
+            CUDA_CHECK(cudaIpcOpenMemHandle((void **)&nixl_ctx->rdma_p2p_ptrs[i], nixl_peer_info[i].rdma_ipc_handle, cudaIpcMemLazyEnablePeerAccess));
+            CUDA_CHECK(cudaIpcOpenMemHandle((void **)&nixl_ctx->counters_p2p_ptrs[i], nixl_peer_info[i].counters_ipc_handle, cudaIpcMemLazyEnablePeerAccess));
+        } else {
+            nixl_ctx->rdma_p2p_ptrs[i] = nullptr;
+            nixl_ctx->counters_p2p_ptrs[i] = nullptr;
+        }
+    }
+}
+
+void Buffer::_nixl_ep_counters_prepare(const std::vector<int>& ranks) {
+    int num_local_experts = env_num_channels;
+
+    for (int expert_idx = 0; expert_idx < num_local_experts; expert_idx++) {
+        for (int remote_rank : ranks) {
+            if (remote_rank == rank)
+                continue;
+
+            int remote_counter_idx = expert_idx * max_num_ranks + rank; // remote rank's counters array is indexed by [local_expert_idx, src_rank]
+            int local_counter_idx = expert_idx * max_num_ranks + remote_rank; // remote_counter_reqs is indexed by [local_expert_idx, dst_rank]
+
+            if (nixl_peer_info[remote_rank].counters_buffer_ptr == nullptr) {
+                printf("[ERROR] _nixl_ep_counters_prepare: nixl_peer_info[%d].counters_buffer_ptr is NULL!\n", remote_rank);
+                exit(1);
+            }
+
+            // Fetch the first counter (double buffering)
+            uint64_t *remote_counter_addr = nixl_peer_info[remote_rank].counters_buffer_ptr + remote_counter_idx;
+            nixl_opt_args_t eparams = {};
+            eparams.backends.push_back(nixl_agent_info->backend);
+            eparams.customParam = "worker_id=" + std::to_string(expert_idx);
+            nixl_xfer_dlist_t dst_dlist(VRAM_SEG);
+            dst_dlist.addDesc(nixlBlobDesc((uintptr_t)remote_counter_addr, sizeof(uint64_t), nixl_peer_info[remote_rank].device_id, ""));
+            EP_HOST_ASSERT(nixl_agent_info->agent->createXferReq(NIXL_WRITE, dummy_src_dlist, dst_dlist, nixl_agent_info->remote_agent_names[remote_rank], nixl_ctx->cpu_remote_counter_reqs_0[local_counter_idx], &eparams) == NIXL_SUCCESS);
+            EP_HOST_ASSERT(nixl_agent_info->agent->createGpuXferReq(*nixl_ctx->cpu_remote_counter_reqs_0[local_counter_idx], nixl_ctx->gpu_remote_counter_reqs_0[local_counter_idx]) == NIXL_SUCCESS);
+
+            // Fetch the second counter (double buffering)
+            remote_counter_addr += max_num_ranks * num_local_experts;
+            nixl_xfer_dlist_t dst_dlist_2(VRAM_SEG);
+            dst_dlist_2.addDesc(nixlBlobDesc((uintptr_t)remote_counter_addr, sizeof(uint64_t), nixl_peer_info[remote_rank].device_id, ""));
+            EP_HOST_ASSERT(nixl_agent_info->agent->createXferReq(NIXL_WRITE, dummy_src_dlist, dst_dlist_2, nixl_agent_info->remote_agent_names[remote_rank], nixl_ctx->cpu_remote_counter_reqs_1[local_counter_idx], &eparams) == NIXL_SUCCESS);
+            EP_HOST_ASSERT(nixl_agent_info->agent->createGpuXferReq(*nixl_ctx->cpu_remote_counter_reqs_1[local_counter_idx], nixl_ctx->gpu_remote_counter_reqs_1[local_counter_idx]) == NIXL_SUCCESS);
+        }
+    }
+}
+
+void Buffer::_nixl_agents_disconnect(const std::vector<int>& ranks) {
+    for (int remote_rank : ranks) {
+        EP_HOST_ASSERT(remote_rank != rank);
+        EP_HOST_ASSERT(remote_rank < num_ranks);
+        nixl_xfer_dlist_t empty_descs(VRAM_SEG);
+        if(nixl_agent_info->agent->checkRemoteMD(nixl_agent_info->remote_agent_names[remote_rank], empty_descs) == NIXL_SUCCESS) {
+            nixl_status_t status = nixl_agent_info->agent->invalidateRemoteMD(nixl_agent_info->remote_agent_names[remote_rank]);
+            // NIXL watchers might invalidate peer metadata, so we ignore NIXL_ERR_NOT_FOUND errors
+            if (status != NIXL_SUCCESS && status != NIXL_ERR_NOT_FOUND) {
+                printf("WARNING: rank %d Failed to invalidate remote rank %d metadata for agent %s, status: %d\n",
+                    rank, remote_rank, std::to_string(remote_rank).c_str(), status); fflush(stdout);
+            }
+        }
+    }
+}
+
+void Buffer::_nixl_agents_peer_info_cleanup(const std::vector<int>& ranks) {
+    for (int remote_rank : ranks) {
+        nixl_agent_info->wire_up_done[remote_rank] = false;
+        // Clear nixl_peer_info for removed ranks (only do this once per rank, on first channel)
+        nixl_peer_info[remote_rank] = NixlPeerInfo{};
+    }
+}
+
+void Buffer::_nixl_ep_cleanup(const std::vector<int>& ranks) {
+    _nixl_ep_p2p_ptrs_cleanup(ranks);
+    _nixl_ep_batches_cleanup(ranks);
+    _nixl_ep_counters_cleanup(ranks);
+    _nixl_ep_gpu_ctx_update();
+}
+
+void Buffer::_nixl_ep_counters_cleanup(const std::vector<int>& ranks) {
+    int num_local_experts = env_num_channels;
+
+    for (int expert_idx = 0; expert_idx < num_local_experts; expert_idx++) {
+        for (int remote_rank : ranks) {
+            EP_HOST_ASSERT(remote_rank != rank);
+
+            int local_counter_idx = expert_idx * max_num_ranks + remote_rank;
+
+            // Clean up remote counter requests (double buffering)
+            if (nixl_ctx->cpu_remote_counter_reqs_0[local_counter_idx] != nullptr) {
+
+#ifndef EP_REMOVE_ONCE
+                nixl_agent_info->agent->releaseGpuXferReq(nixl_ctx->gpu_remote_counter_reqs_0[local_counter_idx]);
+                nixl_agent_info->agent->releaseXferReq(nixl_ctx->cpu_remote_counter_reqs_0[local_counter_idx]);
+#endif
+                nixl_ctx->cpu_remote_counter_reqs_0[local_counter_idx] = nullptr;
+                nixl_ctx->gpu_remote_counter_reqs_0[local_counter_idx] = nullptr;
+            }
+
+            if (nixl_ctx->cpu_remote_counter_reqs_1[local_counter_idx] != nullptr) {
+#ifndef EP_REMOVE_ONCE
+                nixl_agent_info->agent->releaseGpuXferReq(nixl_ctx->gpu_remote_counter_reqs_1[local_counter_idx]);
+                nixl_agent_info->agent->releaseXferReq(nixl_ctx->cpu_remote_counter_reqs_1[local_counter_idx]);
+#endif
+                nixl_ctx->cpu_remote_counter_reqs_1[local_counter_idx] = nullptr;
+                nixl_ctx->gpu_remote_counter_reqs_1[local_counter_idx] = nullptr;
+            }
+
+            if (nixl_ctx->cpu_barrier_reqs[expert_idx][remote_rank] != nullptr) {
+#ifndef EP_REMOVE_ONCE
+                nixl_agent_info->agent->releaseGpuXferReq(nixl_ctx->gpu_barrier_reqs[expert_idx][remote_rank]);
+                nixl_agent_info->agent->releaseXferReq(nixl_ctx->cpu_barrier_reqs[expert_idx][remote_rank]);
+#endif
+                nixl_ctx->cpu_barrier_reqs[expert_idx][remote_rank] = nullptr;
+                nixl_ctx->gpu_barrier_reqs[expert_idx][remote_rank] = nullptr;
+            }
+        }
+    }
+}
+
+void Buffer::_nixl_ep_batches_cleanup(const std::vector<int>& ranks) {
+    for (int channel = 0; channel < env_num_channels; ++channel) {
+        for (int remote_rank : ranks) {
+            if (remote_rank == rank) continue;
+
+            // Clean up cpu_batch_reqs and gpu_batch_reqs
+            if (remote_rank < nixl_ctx->cpu_batch_reqs[channel].size() &&
+                nixl_ctx->cpu_batch_reqs[channel][remote_rank] != nullptr) {
+
+                // Release GPU transfer request first
+                if (nixl_ctx->gpu_batch_reqs[channel][remote_rank] != nullptr) {
+#ifndef EP_REMOVE_ONCE
+                  nixl_agent_info->agent->releaseGpuXferReq(nixl_ctx->gpu_batch_reqs[channel][remote_rank]);
+#endif
+                    nixl_ctx->gpu_batch_reqs[channel][remote_rank] = nullptr;
+                }
+
+                // Release CPU transfer request
+#ifndef EP_REMOVE_ONCE
+                nixl_status_t status = nixl_agent_info->agent->releaseXferReq(nixl_ctx->cpu_batch_reqs[channel][remote_rank]);
+                if (status != NIXL_SUCCESS) {
+                    printf("[WARNING] %s: Failed to release CPU batch xfer req for rank %d on channel %d, status: %d\n",
+                           __func__, remote_rank, channel, status);
+                }
+#endif
+                nixl_ctx->cpu_batch_reqs[channel][remote_rank] = nullptr;
+            }
+        }
+    }
+}
+
+void Buffer::_nixl_ep_p2p_ptrs_cleanup(const std::vector<int>& ranks) {
+    for (int remote_rank : ranks) {
+        EP_HOST_ASSERT(remote_rank < num_ranks);
+        // Close P2P memory mappings if they exist
+        if (nixl_ctx->rdma_p2p_ptrs[remote_rank] != nullptr &&
+            nixl_ctx->rdma_p2p_ptrs[remote_rank] != rdma_buffer_ptr) {
+            CUDA_CHECK(cudaIpcCloseMemHandle(nixl_ctx->rdma_p2p_ptrs[remote_rank]));
+            nixl_ctx->rdma_p2p_ptrs[remote_rank] = nullptr;
+        }
+
+        if (nixl_ctx->counters_p2p_ptrs[remote_rank] != nullptr &&
+            nixl_ctx->counters_p2p_ptrs[remote_rank] != counters_buffer_ptr) {
+            CUDA_CHECK(cudaIpcCloseMemHandle(nixl_ctx->counters_p2p_ptrs[remote_rank]));
+            nixl_ctx->counters_p2p_ptrs[remote_rank] = nullptr;
+        }
+    }
+}
+
+} // namespace nixl_ep
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.doc() = "NIXL_EP: an efficient expert-parallel communication library";
+    m.def("get_rdma_size_hint", &nixl_ep::get_rdma_size_hint);
+
+    pybind11::class_<nixl_ep::EventHandle>(m, "EventHandle")
+        .def(pybind11::init<>())
+        .def("current_stream_wait", &nixl_ep::EventHandle::current_stream_wait);
+
+    pybind11::class_<nixl_ep::Buffer>(m, "Buffer")
+        .def(pybind11::init<int, bool, bool>())
+        .def("update_memory_buffers", &nixl_ep::Buffer::update_memory_buffers)
+        .def("barrier", &nixl_ep::Buffer::barrier)
+        .def("connect_ranks", &nixl_ep::Buffer::connect_ranks, py::arg("remote_ranks"))
+        .def("disconnect_ranks", &nixl_ep::Buffer::disconnect_ranks)
+        .def("is_available", &nixl_ep::Buffer::is_available)
+        .def("get_local_device_id", &nixl_ep::Buffer::get_local_device_id)
+        .def("get_local_buffer_tensor", &nixl_ep::Buffer::get_local_buffer_tensor)
+        .def("get_comm_stream", &nixl_ep::Buffer::get_comm_stream)
+        .def("destroy", &nixl_ep::Buffer::destroy)
+        .def("dispatch", &nixl_ep::Buffer::dispatch)
+        .def("combine", &nixl_ep::Buffer::combine)
+        .def("update_mask_buffer", &nixl_ep::Buffer::update_mask_buffer)
+        .def("query_mask_buffer", &nixl_ep::Buffer::query_mask_buffer)
+        .def("clean_mask_buffer", &nixl_ep::Buffer::clean_mask_buffer)
+        .def("get_next_combine_buffer", &nixl_ep::Buffer::get_next_combine_buffer);
+    m.def("is_sm90_compiled", nixl_ep::is_sm90_compiled);
+}
diff --git a/examples/device/ep/csrc/nixl_ep.hpp b/examples/device/ep/csrc/nixl_ep.hpp
new file mode 100644
index 0000000000..6e8ea223be
--- /dev/null
+++ b/examples/device/ep/csrc/nixl_ep.hpp
@@ -0,0 +1,221 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * This file incorporates material from the DeepSeek project, licensed under the MIT License.
+ * The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+ *
+ * SPDX-License-Identifier: MIT AND Apache-2.0
+ *
+ * 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
+
+// Forcibly disable NDEBUG
+#ifdef NDEBUG
+#undef NDEBUG
+#endif
+
+#include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
+#include <torch/types.h>
+#include <tuple>
+#include <vector>
+#include <string>
+
+#include <memory>
+#include "config.hpp"
+#include "event.hpp"
+#include "kernels/configs.cuh"
+#include "kernels/exception.cuh"
+
+#include "nixl.h"
+
+#define EP_EXECUTE_ONCE(func) do { static bool _ = ((func), true); } while(0)
+
+#ifndef TORCH_EXTENSION_NAME
+#define TORCH_EXTENSION_NAME nixl_ep_cpp
+#endif
+
+namespace nixl_ep {
+
+#define MAX_IP_LENGTH 16
+#define MAX_BOOT_ID_LENGTH 37
+
+struct NixlPeerInfo {
+    char ip[MAX_IP_LENGTH];
+    char boot_id[MAX_BOOT_ID_LENGTH];
+    ino_t ipc_namespace_inode;
+    void *rdma_buffer_ptr;
+    uint64_t *counters_buffer_ptr;
+    uint64_t *wireup_ptr;
+    cudaIpcMemHandle_t rdma_ipc_handle;
+    cudaIpcMemHandle_t counters_ipc_handle;
+    int* sync_buffer_ptr;
+    int device_id;
+    int rank;
+};
+
+struct NixlAgentInfo
+{
+    NixlAgentInfo(std::shared_ptr<nixlAgent> agent, nixlBackendH* backend, int max_num_ranks): agent(agent), backend(backend) {
+        wire_up_done.resize(max_num_ranks, false);
+        remote_agent_names.resize(max_num_ranks);
+    }
+
+    std::shared_ptr<nixlAgent> agent;
+    std::string agent_name;
+    std::vector<std::string> remote_agent_names;
+    nixl_opt_args_t extra_params;
+    nixlBackendH* backend;
+    std::vector<bool> wire_up_done; // [num_peers]
+};
+
+struct nixl_ep_ctx {
+    std::vector<nixlXferReqH *> cpu_remote_counter_reqs_0; // [dest_expert_id,remote_rank], cpu ptrs to nixlXferReqH
+    std::vector<nixlXferReqH *> cpu_remote_counter_reqs_1; // [dest_expert_id,remote_rank], cpu ptrs to nixlXferReqH
+    std::vector<nixlGpuXferReqH> gpu_remote_counter_reqs_0; // [dest_expert_id,remote_rank], gpu ptrs to nixlGpuXferReqH
+    std::vector<nixlGpuXferReqH> gpu_remote_counter_reqs_1; // [dest_expert_id,remote_rank], gpu ptrs to nixlGpuXferReqH
+    std::vector<std::vector<nixlXferReqH*>> cpu_batch_reqs; // [num_local_experts][num_peers]
+    std::vector<std::vector<nixlGpuXferReqH>> gpu_batch_reqs; // [num_local_experts][num_peers]
+    std::vector<std::vector<nixlXferReqH*>> cpu_barrier_reqs;
+    std::vector<std::vector<nixlGpuXferReqH>> gpu_barrier_reqs;
+
+    std::vector<void *> rdma_p2p_ptrs; // [num_ranks]
+    std::vector<uint64_t *> counters_p2p_ptrs; // [num_ranks]
+    ep_kernels::gpu_nixl_ctx gpu[2]; // Double buffering
+};
+
+struct Buffer {
+private:
+    int buffer_idx = 0; // Double buffering index
+
+    // RDMA Buffer
+    int64_t num_rdma_bytes;
+    void* rdma_buffer_ptr = nullptr;
+
+    // Shrink mode buffer
+    bool enable_shrink = false;
+    int *mask_buffer_ptr = nullptr;
+    int *sync_buffer_ptr = nullptr;
+
+    // Device info and communication
+    int device_id;
+    int num_device_sms;
+    int rank;
+    int num_ranks;
+    std::vector<int> remote_ranks; /* global ranks */
+
+    // Stream for communication
+    at::cuda::CUDAStream comm_stream;
+
+    // After synchronization, this flag will be true
+    bool available = false;
+
+    // Whether explicit `destroy()` is required.
+    bool explicitly_destroy;
+    // After `destroy()` be called, this flag will be true
+    bool destroyed = false;
+
+    // Workspace
+    void* workspace = nullptr;
+
+    std::unique_ptr<NixlAgentInfo> nixl_agent_info;
+    std::vector<NixlPeerInfo> nixl_peer_info;
+    uint64_t *counters_buffer_ptr = nullptr;
+    uint64_t *wireup_buffer_ptr = nullptr;
+    NixlPeerInfo my_peer_info;
+    uint64_t num_counters;
+    uint64_t max_num_ranks;
+    int env_num_channels;
+    nixl_xfer_dlist_t dummy_src_dlist; // TODO: Remove once NIXL supports null src dlist for signals
+    std::unique_ptr<nixl_ep_ctx> nixl_ctx = nullptr;
+
+    /* Common private funcs */
+    void _nixl_agent_init();
+    void _nixl_agents_connect(const std::vector<int>& ranks);
+    void _nixl_agents_disconnect(const std::vector<int>& ranks);
+    void _nixl_agents_peer_info_gather(std::vector<int>& ranks);
+    void _nixl_agents_peer_info_cleanup(const std::vector<int>& ranks);
+    void _nixl_agents_wireup(std::vector<int>& ranks);
+
+    /* NIXL EP private funcs */
+    void _nixl_ep_init(const std::vector<int>& ranks);
+    void _nixl_ep_context_init();
+    void _nixl_ep_counters_prepare(const std::vector<int>& ranks);
+    void _nixl_ep_batches_prepare(const std::vector<int>& ranks);
+    void _nixl_ep_p2p_ptrs_prepare(const std::vector<int>& ranks);
+    void _nixl_ep_gpu_ctx_update();
+
+    /* NIXL EP cleanup funcs */
+    void _nixl_ep_cleanup(const std::vector<int>& ranks_to_remove);
+    void _nixl_ep_counters_cleanup(const std::vector<int>& ranks_to_remove);
+    void _nixl_ep_batches_cleanup(const std::vector<int>& ranks_to_remove);
+    void _nixl_ep_p2p_ptrs_cleanup(const std::vector<int>& ranks_to_remove);
+    void _nixl_ep_barrier_buffer_clear();
+
+public:
+    Buffer(int rank, bool explicitly_destroy, bool enable_shrink);
+
+    void update_memory_buffers(int num_ranks, int64_t num_rdma_bytes);
+
+    void connect_ranks(const std::vector<int>& remote_ranks_list);
+
+    void disconnect_ranks(const std::vector<int>& remote_ranks_list);
+
+    void init(int num_ranks, int64_t num_rdma_bytes);
+
+    ~Buffer() noexcept(false);
+
+    bool is_available() const;
+
+    int get_local_device_id() const;
+
+    torch::Tensor get_local_buffer_tensor(const pybind11::object& dtype, int64_t offset) const;
+
+    torch::Stream get_comm_stream() const;
+
+    void destroy();
+
+    void clean_buffer(int num_max_dispatch_tokens_per_rank, int hidden, int num_experts);
+
+    std::tuple<torch::Tensor, std::optional<torch::Tensor>, torch::Tensor, torch::Tensor, torch::Tensor, std::optional<EventHandle>, std::optional<std::function<void()>>>
+    dispatch(const torch::Tensor& x, const torch::Tensor& topk_idx,
+                         const std::optional<torch::Tensor>& cumulative_local_expert_recv_stats,
+                         const std::optional<torch::Tensor>& dispatch_wait_recv_cost_stats,
+                         int num_max_dispatch_tokens_per_rank, int num_experts,
+                         bool use_fp8, bool round_scale, bool use_ue8m0,
+                         bool async, bool return_recv_hook);
+
+    std::tuple<torch::Tensor, std::optional<EventHandle>, std::optional<std::function<void()>>>
+    combine(const torch::Tensor& x, const torch::Tensor& topk_idx, const torch::Tensor& topk_weights,
+                        const torch::Tensor& src_info, const torch::Tensor& layout_range,
+                        const std::optional<torch::Tensor>& combine_wait_recv_cost_stats,
+                        int num_max_dispatch_tokens_per_rank, int num_experts,
+                        bool use_logfmt, bool zero_copy, bool async, bool return_recv_hook,
+                        const std::optional<torch::Tensor>& out = std::nullopt);
+
+    void barrier();
+
+    torch::Tensor
+    get_next_combine_buffer(int num_max_dispatch_tokens_per_rank, int hidden, int num_experts) const;
+
+    void update_mask_buffer(int rank_to_mask, bool mask);
+
+    void query_mask_buffer(const torch::Tensor& mask_status);
+
+    void clean_mask_buffer();
+};
+
+} // namespace nixl_ep
diff --git a/examples/device/ep/meson.build b/examples/device/ep/meson.build
new file mode 100644
index 0000000000..221bda61c7
--- /dev/null
+++ b/examples/device/ep/meson.build
@@ -0,0 +1,176 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 not cuda_dep.found()
+    warning('CUDA not found, skipping nixl_ep build')
+    subdir_done()
+endif
+
+py = import('python').find_installation('python3', pure: false)
+
+# Check if PyTorch is available
+torch_check = run_command(py, '-c', 'import torch', check: false)
+if torch_check.returncode() != 0
+    warning('PyTorch not found, skipping nixl_ep build')
+    subdir_done()
+endif
+
+# Get PyTorch base directory from Python
+torch_dir = run_command(py, '-c', 'import torch, os; print(os.path.dirname(torch.__file__))', check: true).stdout().strip()
+torch_include_base = join_paths(torch_dir, 'include')
+torch_lib_dir = join_paths(torch_dir, 'lib')
+
+# All Torch include directories
+torch_inc_dirs = [
+    include_directories(torch_include_base),
+    include_directories(join_paths(torch_include_base, 'torch')),
+    include_directories(join_paths(torch_include_base, 'torch', 'csrc', 'api', 'include')),
+    include_directories(join_paths(torch_include_base, 'ATen')),
+]
+
+# Create Torch dependency with includes and libraries
+torch_dep = declare_dependency(
+    include_directories: torch_inc_dirs,
+    link_args: ['-L' + torch_lib_dir, '-ltorch', '-ltorch_python', '-ltorch_cuda']
+)
+
+pybind_dep = dependency('pybind11', required: false)
+if not pybind_dep.found()
+    warning('pybind11 not found, skipping nixl_ep build')
+    subdir_done()
+endif
+
+nixl_dep = declare_dependency(link_with: nixl_lib, include_directories: nixl_inc_dirs)
+nixl_lib_dir = join_paths(meson.project_build_root(), 'src', 'core')
+
+ucx_build_deps = []
+if ucx_dep.found()
+    ucx_build_deps = [ucx_dep]
+endif
+
+doca_build_deps = []
+if doca_gpunetio_dep.found()
+    doca_build_deps = [doca_gpunetio_dep]
+endif
+
+nixl_ep_sources = [
+    'csrc/nixl_ep.cpp',
+    'csrc/kernels/nixl_ep.cu',
+]
+
+nixl_ep_inc_dirs = [
+    nixl_inc_dirs,
+    nixl_gpu_inc_dirs,
+    include_directories('csrc'),
+    include_directories('csrc/kernels'),
+    torch_inc_dirs,
+]
+
+nixl_ep_cpp_args = [
+    '-DHAVE_CUDA',
+    '-DTORCH_EXTENSION_NAME=nixl_ep_cpp',
+    '-Wno-deprecated-declarations',
+    '-Wno-unused-variable',
+    '-Wno-sign-compare',
+    '-Wno-reorder',
+    '-Wno-attributes',
+]
+
+nixl_ep_cuda_args = [
+    '-DHAVE_CUDA',
+    '-DTORCH_EXTENSION_NAME=nixl_ep_cpp',
+    '--expt-relaxed-constexpr',  # Allow calling constexpr __host__ functions from __device__ functions
+    '-arch=sm_90',  # Only compile for sm90 (overrides global -gencode flags)
+    '--ptxas-options=--register-usage-level=10',  # Allow more register usage (matches setup.py)
+    '-Xcompiler', '-Wno-deprecated-declarations',
+    '-Xcompiler', '-Wno-unused-variable',
+    '-Xcompiler', '-Wno-sign-compare',
+    '-Xcompiler', '-Wno-reorder',
+    '-Xcompiler', '-Wno-attributes',
+    '-DENABLE_FAST_DEBUG=1',
+]
+
+# NVCC 12.9 workaround for bug (https://nvbugspro.nvidia.com/bug/5595631)
+# This workaround is needed because NVCC 12.9 has a bug which fails UCX compilation
+nvcc = meson.get_compiler('cuda')
+if nvcc.version().version_compare('>=12.9') and nvcc.version().version_compare('<12.10')
+    nixl_ep_cpp_args += ['-D_LIBCUDACXX_ATOMIC_UNSAFE_AUTOMATIC_STORAGE']
+    nixl_ep_cuda_args += ['-D_LIBCUDACXX_ATOMIC_UNSAFE_AUTOMATIC_STORAGE']
+endif
+
+nixl_ep_rpath = torch_lib_dir
+nixl_ep_rpath += ':' + nixl_lib_dir
+nixl_ep_rpath += ':' + join_paths(nixl_lib_dir, 'core')
+nixl_ep_rpath += ':' + join_paths(nixl_lib_dir, 'plugins')
+
+# For now, nixl ep cannot be built with -G due to register usage limits
+nixl_ep_override_options = []
+if get_option('buildtype') == 'debug'
+    nixl_ep_override_options = ['optimization=3']
+    nixl_ep_cpp_args += ['-g']
+endif
+
+nixl_ep_install_rpath = join_paths(get_option('prefix'), get_option('libdir'))
+nixl_ep_install_rpath += ':' + join_paths(get_option('prefix'), get_option('libdir'), 'plugins')
+nixl_ep_install_rpath += ':' + torch_lib_dir
+
+nixl_ep_ext = py.extension_module('nixl_ep_cpp',
+    nixl_ep_sources,
+    subdir: 'nixl_ep',
+    dependencies: [
+        nixl_dep,
+        pybind_dep,
+        cuda_dep,
+        torch_dep,
+    ] + ucx_build_deps + doca_build_deps,
+    include_directories: nixl_ep_inc_dirs,
+    cpp_args: nixl_ep_cpp_args,
+    cuda_args: nixl_ep_cuda_args,
+    link_with: [nixl_lib],
+    build_rpath: nixl_ep_rpath,
+    install_rpath: nixl_ep_install_rpath,
+    override_options: nixl_ep_override_options,
+    install: true,
+)
+
+nixl_ep_py_files = files(
+    'nixl_ep/__init__.py',
+    'nixl_ep/buffer.py',
+    'nixl_ep/utils.py',
+)
+
+custom_target('nixl_ep_py_copy',
+    output: 'nixl_ep_py.stamp',
+    input: nixl_ep_py_files,
+    command: [
+        'bash', '-c',
+        'cp -r @0@/nixl_ep @1@/ && cp @2@ @1@/nixl_ep/ && touch @3@'.format(
+            meson.current_source_dir(),
+            meson.current_build_dir(),
+            nixl_ep_ext.full_path(),
+            join_paths(meson.current_build_dir(), 'nixl_ep_py.stamp')
+        )
+    ],
+    depends: nixl_ep_ext,
+    build_by_default: true,
+)
+
+py.install_sources(
+    'nixl_ep/__init__.py',
+    'nixl_ep/buffer.py',
+    'nixl_ep/utils.py',
+    subdir: 'nixl_ep',
+    pure: false,
+)
diff --git a/examples/device/ep/nixl_ep/__init__.py b/examples/device/ep/nixl_ep/__init__.py
new file mode 100644
index 0000000000..2e116870a2
--- /dev/null
+++ b/examples/device/ep/nixl_ep/__init__.py
@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This file incorporates material from the DeepSeek project, licensed under the MIT License.
+# The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+#
+# SPDX-License-Identifier: MIT AND Apache-2.0
+#
+# 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.
+
+from .buffer import Buffer
+from .utils import EventOverlap
+
+__all__ = ['Buffer', 'EventOverlap']
diff --git a/examples/device/ep/nixl_ep/buffer.py b/examples/device/ep/nixl_ep/buffer.py
new file mode 100644
index 0000000000..6fbf566849
--- /dev/null
+++ b/examples/device/ep/nixl_ep/buffer.py
@@ -0,0 +1,378 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This file incorporates material from the DeepSeek project, licensed under the MIT License.
+# The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+#
+# SPDX-License-Identifier: MIT AND Apache-2.0
+#
+# 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.
+
+import os
+from typing import TYPE_CHECKING, Callable, List, Literal, Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+
+# noinspection PyUnresolvedReferences
+from . import nixl_ep_cpp
+
+# noinspection PyUnresolvedReferences
+from .nixl_ep_cpp import EventHandle
+from .utils import EventOverlap
+
+if TYPE_CHECKING:
+    import mpi4py  # noqa: F401
+
+
+class Buffer:
+    """
+    The core expert-parallel (EP) communication buffers for Mixture of Experts (MoE) model, which supports dispatch and combine operations using NVLink and RDMA.
+
+    Attributes:
+        nvlink_backend: the backend for NVLink communication, you can choose from 'nixl', 'ipc' and 'none' (which disables NVLink entirely).
+        rank: the local rank number.
+        num_rdma_bytes: the buffer size for RDMA communication.
+        runtime: the C++ runtime.
+    """
+
+    num_sms: int = 20
+
+    def __init__(self, nvlink_backend: Literal['nixl', 'ipc', 'none'] = 'nixl',
+                 explicitly_destroy: bool = False, rank: int = 0, enable_shrink: bool = False,
+                 group: Optional[dist.ProcessGroup] = None, comm: Optional["mpi4py.MPI.Comm"] = None) -> None:
+        """
+        Initialize the nixl communication buffer.
+
+        Arguments:
+            nvlink_backend: nvlink implementation to use, you can choose from 'nixl', 'ipc' and 'none' (which disables NVLink entirely).
+            explicitly_destroy: If this flag is set to True, you need to explicitly call `destroy()` to release resources;
+                otherwise, the resources will be released by the destructor.
+                Note: Releasing resources in the destructor may cause Python's exception handling process to hang.
+            rank: the rank number.
+            group: the communication group (optional).
+            comm: the mpi4py.MPI.Comm communicator to use in case the group parameter is absent (optional).
+        """
+        self.rank = rank
+        self.group_size = 0  # Will be updated by `update_memory_buffers`
+        self.explicitly_destroy = explicitly_destroy
+        self.group = group
+        self.comm = comm
+        assert not (group and comm)
+
+        # Configure NVLINK backend
+        os.environ['NIXL_EP_NVLINK_BACKEND_IPC'] = '1' if nvlink_backend == 'ipc' else '0'
+        if nvlink_backend != 'nixl':
+            os.environ["UCX_TLS"] = "^cuda_ipc"
+
+        self.runtime = nixl_ep_cpp.Buffer(self.rank, explicitly_destroy, enable_shrink)
+
+    def destroy(self):
+        """
+        Destroy the cpp runtime and release resources.
+        """
+        assert self.explicitly_destroy, '`explicitly_destroy` flag must be set'
+
+        self.runtime.destroy()
+        self.runtime = None
+
+    @staticmethod
+    def is_sm90_compiled():
+        return nixl_ep_cpp.is_sm90_compiled()
+
+    @staticmethod
+    def set_num_sms(new_num_sms: int) -> None:
+        """
+        Set the number of SMs to use in high-throughput kernels.
+
+        Arguments:
+            new_num_sms: the new number to be set.
+        """
+
+        assert new_num_sms % 2 == 0, 'The SM count must be even'
+        Buffer.num_sms = new_num_sms
+
+    @staticmethod
+    def capture() -> EventOverlap:
+        """
+        Capture a CUDA event on the current stream, i.e. `torch.cuda.current_stream()`.
+
+        Returns:
+            event: the captured event.
+        """
+        return EventOverlap(EventHandle())
+
+    @staticmethod
+    def get_rdma_size_hint(num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int) -> int:
+        """
+        Get a minimum size requirement for the RDMA buffer. The size calculation will be done with BF16.
+
+        Arguments:
+            num_max_dispatch_tokens_per_rank: the maximum number of tokens to dispatch, all the ranks must hold the same value.
+            hidden: the hidden dimension of each token.
+            num_ranks: the number of EP group ranks.
+            num_experts: the number of all experts.
+
+        Returns:
+            size: the RDMA buffer size recommended.
+        """
+        return nixl_ep_cpp.get_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts)
+
+    def get_comm_stream(self) -> torch.Stream:
+        """
+        Get the communication stream.
+
+        Returns:
+            stream: the communication stream.
+        """
+        ts: torch.Stream = self.runtime.get_comm_stream()
+        return torch.cuda.Stream(stream_id=ts.stream_id, device_index=ts.device_index, device_type=ts.device_type)
+
+    def get_local_buffer_tensor(self, dtype: torch.dtype, size: Optional[torch.Size] = None,
+                                offset: int = 0) -> torch.Tensor:
+        """
+        Get the raw buffer (slice supported) as a PyTorch tensor.
+
+        Argument:
+            dtype: the data type (PyTorch `dtype`) for the tensor.
+            size: the slice size (by elements) to get from the buffer.
+            offset: the offset of the beginning element.
+        """
+        tensor = self.runtime.get_local_buffer_tensor(dtype, offset)
+        if size is None:
+            return tensor
+
+        assert tensor.numel() >= size.numel()
+        return tensor[:size.numel()].view(size)
+
+    @staticmethod
+    def _unpack_bias(bias: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]):
+        bias_0, bias_1 = None, None
+        if isinstance(bias, torch.Tensor):
+            bias_0 = bias
+        elif isinstance(bias, tuple):
+            assert len(bias) == 2
+            bias_0, bias_1 = bias
+        return bias_0, bias_1
+
+    def clean_buffer(self, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int) -> None:
+        """
+        As the kernels require part of the buffer to be zero-initialized, so it is vital to clean the buffer
+            if the buffer is dirty at some time.
+
+        Arguments:
+            num_max_dispatch_tokens_per_rank: the maximum number of tokens to dispatch, all the ranks must hold the same value.
+            hidden: the hidden dimension of each token.
+            num_experts: the number of all experts.
+        """
+        self.runtime.clean_buffer(num_max_dispatch_tokens_per_rank, hidden, num_experts)
+
+    # noinspection PyTypeChecker
+    def dispatch(self, x: torch.Tensor, topk_idx: torch.Tensor,
+                 num_max_dispatch_tokens_per_rank: int, num_experts: int,
+                 cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None,
+                 dispatch_wait_recv_cost_stats: Optional[torch.Tensor] = None,
+                 use_fp8: bool = True, round_scale: bool = False, use_ue8m0: bool = False,
+                 async_finish: bool = False, return_recv_hook: bool = False) -> \
+            Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor, Tuple, EventOverlap, Callable]:
+        """
+        A low-latency implementation for dispatching with NIXL device API.
+        This kernel requires all the ranks (no matter intranode or internode) should be visible via RDMA
+        Warning: as there are only two buffers, and the returned tensors reuse the buffer, you cannot hold more than 2
+            kernels' result tensors at a single moment.
+
+        Arguments:
+            x: `torch.Tensor` with `torch.bfloat16`, shaped as `[num_tokens, hidden]`, only several hidden shapes are
+                supported. The number of tokens to be dispatched must be less than `num_max_dispatch_tokens_per_rank`.
+            topk_idx: `torch.Tensor` with `torch.int64`, shaped as `[num_tokens, num_topk]`, only several top-k shapes
+                are supported. `-1` indices (not selecting any expert) are supported.
+            num_max_dispatch_tokens_per_rank: the maximum number of tokens to dispatch, all the ranks must hold the same value.
+            num_experts: the number of all experts.
+            cumulative_local_expert_recv_stats: a cumulative expert count tensor for statistics, which should have shape
+                `[num_local_experts]` and be typed as `torch.int`. This is useful for online service EP load balance
+                monitoring.
+            dispatch_wait_recv_cost_stats: a cumulative time spent waiting to receive each token tensor for statistics,
+                which should have shape `[num_ranks, num_ranks]` and be typed as `torch.int64`.
+                This is useful for detecting and pre-cisely localizing slow anomalies.
+            use_fp8: whether to enable FP8 casting, with this, the received data will be a tuple of FP8 tensor and scaling factors.
+            round_scale: whether round the scaling factors into power of 2.
+            use_ue8m0: whether use UE8M0 as scaling factor format (available only with `round_scale=True`).
+            async_finish: the current stream will not wait for the communication kernels to be finished if set.
+            return_recv_hook: return a receiving hook if set. If set, the kernel will just do the RDMA request issues,
+                but **without actually receiving the data**. You must call the received hook to make sure the data's arrival.
+                If you do not set this flag, the kernel will ensure the data's arrival.
+
+        Returns:
+            recv_x: a tensor or tuple with received tokens for each expert.
+                With `use_fp8=True`: the first element is a `torch.Tensor` shaped as
+                `[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden]` with `torch.float8_e4m3fn`.
+                The second tensor is the corresponding scales for the first element with shape
+                `[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden // 128]` with `torch.float`,
+                if `use_ue8m0=False`. With `use_ue8m0=True`, the second one is packed and shaped as
+                `[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden // 512]` with type `torch.int`.
+                Notice that, the last-two-dimension of the scaling tensors are in column-major for TMA compatibility.
+                With `use_fp8=False`, the result would be a tensor shaped as
+                `[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden]` with `torch.bfloat16`.
+                Moreover, not all tokens are valid, only some of the `num_max_dispatch_tokens_per_rank * num_ranks` are,
+                as we do not synchronize CPU received count with GPU (also not incompatible with CUDA graph if synced).
+            recv_count: a tensor shaped `[num_local_experts]` with type `torch.int`, indicating how many tokens each
+                expert receives. As mentioned before, not all tokens are valid in `recv_x`.
+            handle: the communication handle to be used in the `combine` function.
+            event: the event after executing the kernel (valid only if `async_finish` is set).
+            hook: the receiving hook function (valid only if `return_recv_hook` is set).
+        """
+        packed_recv_x, packed_recv_x_scales, packed_recv_count, packed_recv_src_info, packed_recv_layout_range, event, hook = \
+            self.runtime.dispatch(
+                x, topk_idx,
+                cumulative_local_expert_recv_stats,
+                dispatch_wait_recv_cost_stats,
+                num_max_dispatch_tokens_per_rank, num_experts,
+                use_fp8, round_scale, use_ue8m0,
+                async_finish, return_recv_hook)
+        handle = (packed_recv_src_info, packed_recv_layout_range, num_max_dispatch_tokens_per_rank, x.size(1), num_experts)
+        tensors_to_record = (x, topk_idx,
+                             packed_recv_x, packed_recv_x_scales, packed_recv_count,
+                             packed_recv_src_info, packed_recv_layout_range,
+                             cumulative_local_expert_recv_stats)
+        return (packed_recv_x, packed_recv_x_scales) if use_fp8 else packed_recv_x, packed_recv_count, handle, \
+            EventOverlap(event, tensors_to_record if async_finish else None), hook
+
+    # noinspection PyTypeChecker
+    def combine(self, x: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor,
+                handle: tuple, use_logfmt: bool = False, zero_copy: bool = False, async_finish: bool = False,
+                return_recv_hook: bool = False, out: Optional[torch.Tensor] = None,
+                combine_wait_recv_cost_stats: Optional[torch.Tensor] = None) -> \
+            Tuple[torch.Tensor, EventOverlap, Callable]:
+        """
+        A low-latency implementation for combining tokens (reduce **with weights**) with NIXL device API.
+        This kernel requires all the ranks (no matter intranode or internode) should be visible via RDMA
+        Warning: as there are only two buffers, and the returned tensors reuse the buffer, you cannot hold more than 2
+            kernels' result tensors at a single moment.
+
+        Arguments:
+            x: `[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden]` with `torch.bfloat16`,
+                the local calculated tokens to be sent to this original rank and reduced.
+            topk_idx: `[num_combined_tokens, num_topk]` with `torch.int64`, the expert indices selected by the dispatched
+                tokens. `-1` indices (not selecting any expert) are supported. Note that, `num_combined_tokens` equals
+                to the number of dispatched tokens.
+            topk_weights: `[num_combined_tokens, num_topk]` with `torch.float`, the expert weights selected by the dispatched
+                tokens. The received tokens will be reduced with the weights in this tensor.
+            handle: the communication handle given by the `dispatch` function.
+            use_logfmt: whether to use an internal "LogFMT with dynamic per-64-channel cast" format (10 bits).
+            zero_copy: whether the tensor is already copied into the RDMA buffer, should be cooperative
+                with `get_next_combine_buffer`.
+            async_finish: the current stream will not wait for the communication kernels to be finished if set.
+            return_recv_hook: return a receiving hook if set. If set, the kernel will just do the RDMA request issues,
+                but **without actually receiving the data**. You must call the received hook to make sure the data's arrival.
+                If you do not set this flag, the kernel will ensure the data's arrival.
+            out: the in-place output tensor, if set, the kernel will write the result to this tensor and return it directly.
+            combine_wait_recv_cost_stats: a cumulative time spent waiting to receive each token tensor for statistics,
+                which should have shape `[num_ranks, num_ranks]` and be typed as `torch.int64`.
+                This is useful for detecting and pre-cisely localizing slow anomalies.
+
+        Returns:
+            combined_x: the reduced token tensor, with shape `[num_combined_tokens, hidden]` and type `torch.bfloat16`.
+            event: the event after executing the kernel (valid only if `async_finish` is set).
+            hook: the receiving hook function (valid only if `return_recv_hook` is set).
+        """
+        src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
+        combined_x, event, hook = self.runtime.combine(
+            x, topk_idx, topk_weights, src_info, layout_range,
+            combine_wait_recv_cost_stats,
+            num_max_dispatch_tokens_per_rank, num_experts,
+            use_logfmt, zero_copy, async_finish, return_recv_hook,
+            out)
+        tensors_to_record = (x, topk_idx, topk_weights, src_info, layout_range, combined_x)
+        return combined_x, EventOverlap(event, tensors_to_record if async_finish else None), hook
+
+    def update_mask_buffer(self, rank_to_mask: int, mask: bool = False):
+        """
+        Mask (unmask) a rank during communication (dispatch, combine, and clean)
+
+        Arguments:
+            rank: the rank to mask (unmask).
+            mask: if True, will mask the rank (do not recvfrom/sendto the rank), otherwise will unmask the rank.
+        """
+        self.runtime.update_mask_buffer(rank_to_mask, mask)
+
+    def query_mask_buffer(self, mask_status: torch.Tensor):
+        """
+        Query the mask status of all ranks
+
+        Arguments:
+            mask_status: `[num_ranks]` with `torch.int`, the mask status of each rank. `1` means mask and `0` means unmasked.
+        """
+        self.runtime.query_mask_buffer(mask_status)
+
+    def clean_mask_buffer(self):
+        """
+        Clean the mask buffer
+
+        """
+        self.runtime.clean_mask_buffer()
+
+    def get_next_combine_buffer(self, handle: Tuple[torch.Tensor, torch.Tensor, int, int, int]):
+        """
+        Get the raw registered RDMA buffer tensor for next combine, so that the next combine kernel can skip the copying.
+
+        Arguments:
+            handle: the communication handle given by the `dispatch` function.
+
+        Returns:
+            buffer: the raw RDMA buffer as a BF16 PyTorch tensor with shape
+                `[num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank, hidden]`, you should fill this buffer
+                by yourself.
+        """
+        src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
+        return self.runtime.get_next_combine_buffer(num_max_dispatch_tokens_per_rank, hidden, num_experts)
+
+    def update_memory_buffers(self, num_ranks: int, num_experts_per_rank: int, num_rdma_bytes: int):
+        """
+        Allocate remote memory for the communication buffer.
+
+        Arguments:
+            num_ranks: the number of ranks.
+            num_experts_per_rank: the number of experts per rank.
+            num_rdma_bytes: the buffer size for RDMA communication.
+        """
+        self.group_size = num_ranks
+        self.num_rdma_bytes = num_rdma_bytes
+        os.environ['NIXL_EP_NUM_CHANNELS'] = str(num_experts_per_rank)
+        self.runtime.update_memory_buffers(num_ranks, num_rdma_bytes)
+
+    def connect_ranks(self, remote_ranks: List[int]) -> None:
+        """
+        Add connections to remote ranks.
+
+        Arguments:
+            remote_ranks: List of remote rank IDs to establish connections with.
+                         The current rank will be automatically filtered out.
+        """
+        self.runtime.connect_ranks(remote_ranks)
+
+    def disconnect_ranks(self, remote_ranks: List[int]) -> None:
+        """
+        Remove connections to remote ranks.
+
+        Arguments:
+            remote_ranks: List of remote rank IDs to remove connections from.
+                         These ranks must be at the end of the current remote_ranks list.
+        """
+        self.runtime.disconnect_ranks(remote_ranks)
+
+    def barrier(self) -> None:
+        """
+        barrier for all active ranks.
+        notice that this barrier does not flush the network QPs as it is currently doesn't have any use-case that requires it
+        """
+        self.runtime.barrier()
diff --git a/examples/device/ep/nixl_ep/utils.py b/examples/device/ep/nixl_ep/utils.py
new file mode 100644
index 0000000000..fc84155848
--- /dev/null
+++ b/examples/device/ep/nixl_ep/utils.py
@@ -0,0 +1,81 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This file incorporates material from the DeepSeek project, licensed under the MIT License.
+# The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+#
+# SPDX-License-Identifier: MIT AND Apache-2.0
+#
+# 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.
+
+from typing import Any, Optional, Tuple
+
+import torch
+
+# noinspection PyUnresolvedReferences
+from .nixl_ep_cpp import EventHandle
+
+
+class EventOverlap:
+    """
+    A wrapper class to manage CUDA events, also for better overlapping convenience.
+
+    Attributes:
+        event: the CUDA event captured.
+        extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph.
+    """
+
+    def __init__(self, event: Optional[EventHandle] = None,
+                 extra_tensors: Optional[Tuple[torch.Tensor, ...]] = None) -> None:
+        """
+        Initialize the class.
+
+        Arguments:
+            event: the CUDA event captured.
+            extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph.
+        """
+        self.event = event
+
+        # NOTES: we use extra tensors to achieve stream recording, otherwise,
+        # stream recording will be incompatible with CUDA graph.
+        self.extra_tensors = extra_tensors
+
+    def current_stream_wait(self) -> None:
+        """
+        The current stream `torch.cuda.current_stream()` waits for the event to be finished.
+        """
+        assert self.event is not None
+        self.event.current_stream_wait()
+
+    def __enter__(self) -> Any:
+        """
+        Utility for overlapping and Python `with` syntax.
+
+        You can overlap the kernels on the current stream with the following example:
+        ```python
+        event_overlap = event_after_all_to_all_kernels()
+        with event_overlap():
+            do_something_on_current_stream()
+        # After exiting the `with` scope, the current stream with wait the event to be finished.
+        ```
+        """
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """
+        Utility for overlapping and Python `with` syntax.
+
+        Please follow the example in the `__enter__` function.
+        """
+        if self.event is not None:
+            self.event.current_stream_wait()
diff --git a/examples/device/ep/scripts/reset_etcd.sh b/examples/device/ep/scripts/reset_etcd.sh
new file mode 100755
index 0000000000..a460b1ddfa
--- /dev/null
+++ b/examples/device/ep/scripts/reset_etcd.sh
@@ -0,0 +1,23 @@
+#!/bin/bash
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 +e
+pkill -9 -f etcd
+pkill -9 -f python
+rm -rf default.etcd
+HOST_IP=$(hostname -I | awk '{print $1}')
+etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://${HOST_IP}:2379 &
diff --git a/examples/device/ep/tests/elastic/README.md b/examples/device/ep/tests/elastic/README.md
new file mode 100644
index 0000000000..e02d864041
--- /dev/null
+++ b/examples/device/ep/tests/elastic/README.md
@@ -0,0 +1,34 @@
+### Elastic Test Suite
+
+#### Single Node (8 ranks, 4→8 expansion):
+```bash
+python3 tests/elastic/elastic.py \
+    --plan tests/elastic/single_expansion.json \
+    --num-processes 8 \
+    --etcd-server http://127.0.0.1:2379
+```
+
+#### Multi-Node Setup:
+
+**Node 1** (will launch the first phase with 4 ranks):
+```bash
+python3 tests/elastic/elastic.py \
+    --plan tests/elastic/single_expansion.json \
+    --num-processes 4 \
+```
+
+**Node 2** (will join the second phase with additional 4 ranks):
+```bash
+python3 tests/elastic/elastic.py \
+    --plan tests/elastic/single_expansion.json \
+    --num-processes 4 \
+    --rank-server $MASTER_IP \
+    --etcd-server http://$MASTER_IP:2379
+```
+
+### Available Test Plans
+
+- `no_expansion.json`: Static 4 ranks (baseline)
+- `single_expansion.json`: 4 → 8 ranks (single expansion)
+- `double_expansion.json`: 4 → 6 → 8 ranks (two expansions)
+- `expansion_contraction.json`: 4 → 8 → 6 ranks (scale up then down)
diff --git a/examples/device/ep/tests/elastic/double_expansion.json b/examples/device/ep/tests/elastic/double_expansion.json
new file mode 100644
index 0000000000..3966a37bdc
--- /dev/null
+++ b/examples/device/ep/tests/elastic/double_expansion.json
@@ -0,0 +1,5 @@
+[
+  [0, 1, 2, 3],
+  [0, 1, 2, 3, 4, 5],
+  [0, 1, 2, 3, 4, 5, 6, 7]
+]
diff --git a/examples/device/ep/tests/elastic/elastic.py b/examples/device/ep/tests/elastic/elastic.py
new file mode 100644
index 0000000000..a86ee65ae9
--- /dev/null
+++ b/examples/device/ep/tests/elastic/elastic.py
@@ -0,0 +1,399 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This file incorporates material from the DeepSeek project, licensed under the MIT License.
+# The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+#
+# SPDX-License-Identifier: MIT AND Apache-2.0
+#
+# 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.
+
+import argparse
+import os
+import random
+import signal
+import sys
+import threading
+import time
+from functools import partial
+from typing import cast
+
+import nixl_ep
+import rank_server
+import torch
+from plan import Plan
+
+# Add tests directory to path to import test utils
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from utils import (  # noqa: E402
+    bench,
+    bench_kineto,
+    calc_diff,
+    hash_tensor,
+    per_token_cast_back,
+)
+
+
+def handle_sigterm(signum, frame, buffer: nixl_ep.Buffer, plan: Plan, rank_client: rank_server.RankClient):
+    print(f"SIGTERM ({signum}) received for process {os.getpid()}! releasing rank and exiting...", flush=True)
+    if plan is not None:
+        rank_client.release_rank(user_context=plan.get_phase())
+    else:
+        rank_client.release_rank()
+    if buffer is not None and buffer.runtime is not None:
+        buffer.destroy()  # to invalidate local MD
+        del buffer
+    sys.exit(1)
+
+
+def self_kill():
+    os.kill(os.getpid(), signal.SIGTERM)
+
+
+def test_main(num_tokens: int, hidden: int, num_experts: int, num_topk: int,
+              rank: int, num_ranks: int, max_num_ranks: int, buffer: nixl_ep.Buffer,
+              use_logfmt: bool = False, seed: int = 0, kineto: bool = False, fault_tolerance_test: bool = False):
+    torch.manual_seed(seed + rank)
+    torch.cuda.manual_seed(seed + rank)
+    random.seed(seed + rank)
+
+    assert num_experts % num_ranks == 0
+    num_local_experts = num_experts // num_ranks
+
+    # NOTES: the integers greater than 256 exceed the BF16 precision limit
+    rank_offset = 128
+    assert num_ranks - rank_offset < 257, 'Too many ranks (exceeding test precision limit)'
+
+    # Track masked ranks (like shrink_test in elastic.py)
+    mask_status = torch.zeros((max_num_ranks, ), dtype=torch.int32, device='cuda')
+
+    x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') * (rank - rank_offset)
+    x[:, -128:] = torch.arange(num_tokens, device='cuda').to(torch.bfloat16).view(-1, 1)
+    x_list = [x]
+    for i in range(4 if use_logfmt else 0):
+        # NOTES: make more LogFMT casts and also with some BF16
+        x_list.append(torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') * 0.5 * random.random())
+    # NOTES: the last one is for performance testing
+    # Most of the values in the perf case is lower than the threshold, casting most channels
+    x_list.append(torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') * 0.1)
+
+    torch.manual_seed(seed + rank + 1000)
+    torch.cuda.manual_seed(seed + rank + 1000)
+    random.seed(seed + rank)
+
+    scores = torch.randn((num_tokens, num_experts), dtype=torch.float32, device='cuda').abs() + 1
+    topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=True)[1]
+    topk_weights = torch.randn((num_tokens, num_topk), dtype=torch.float32, device='cuda').abs()
+
+    # Randomly mask some positions
+    for i in range(10):
+        topk_idx[random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1)] = -1
+
+    all_topk_idx = torch.full((max_num_ranks, num_tokens, num_topk), -1, dtype=topk_idx.dtype, device='cuda')
+    for r in range(num_ranks):
+        # Use same deterministic reset as above (seed + r + 1000)
+        torch.manual_seed(seed + r + 1000)
+        torch.cuda.manual_seed(seed + r + 1000)
+        r_random = random.Random(seed + r)
+        r_scores = torch.randn((num_tokens, num_experts), dtype=torch.float32, device='cuda').abs() + 1
+        r_topk_idx = torch.topk(r_scores, num_topk, dim=-1, largest=True, sorted=True)[1]
+        # Apply same random masking
+        for i in range(10):
+            r_topk_idx[r_random.randint(0, num_tokens - 1), r_random.randint(0, num_topk - 1)] = -1
+        all_topk_idx[r] = r_topk_idx
+
+    # Check dispatch correctness
+    do_check = True
+    hash_value, num_times = 0, 0
+    timer = None
+    for current_x in x_list:
+        for return_recv_hook in (False, True):
+            for dispatch_use_fp8 in (False, True):
+                for round_scale in (False, True) if dispatch_use_fp8 else (False, ):
+                    for use_ue8m0 in (False, True) if round_scale else (False, ):
+                        num_times += 1
+                        for i in range((num_times % 2) + 1):
+                            # Kill this rank at the beginning of the first dispatch if marked to be killed
+                            if fault_tolerance_test and timer is None:
+                                print(f"[rank {rank}] Killing rank during dispatch/combine", flush=True)
+                                timer = threading.Timer(0.0001, self_kill)
+                                timer.start()
+
+                            cumulative_local_expert_recv_stats = torch.zeros((num_local_experts, ), dtype=torch.int, device='cuda')
+                            packed_recv_x, packed_recv_count, handle, event, hook = \
+                                buffer.dispatch(
+                                    current_x, topk_idx, num_tokens, num_experts,
+                                    use_fp8=dispatch_use_fp8, round_scale=round_scale, use_ue8m0=use_ue8m0,
+                                    cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats,
+                                    async_finish=not return_recv_hook, return_recv_hook=return_recv_hook)
+                            hook() if return_recv_hook else event.current_stream_wait()
+                        # Query mask buffer to get current failure status
+                        buffer.query_mask_buffer(mask_status)
+                        packed_recv_x = (packed_recv_x[0], packed_recv_x[1].contiguous()) if dispatch_use_fp8 else packed_recv_x
+                        simulated_gemm_x = per_token_cast_back(packed_recv_x[0].view(-1, hidden), packed_recv_x[1].view(-1, hidden // 128)).view(packed_recv_x[0].shape) \
+                            if dispatch_use_fp8 else cast(torch.Tensor, packed_recv_x).clone()
+
+                        for i in range(num_local_experts if do_check else 0):
+                            expert_id = rank * num_local_experts + i
+                            recv_x = per_token_cast_back(packed_recv_x[0][i], packed_recv_x[1][i]) if dispatch_use_fp8 else packed_recv_x[i]
+                            recv_count, recv_src_info, recv_layout_range = packed_recv_count[i], handle[0][i], handle[1][i]
+
+                            # Check expert indices
+                            int_mask = (2 ** 32) - 1
+                            num_valid_tokens = recv_count.item()
+                            assert cumulative_local_expert_recv_stats[i].item() == num_valid_tokens, f'{cumulative_local_expert_recv_stats[i].item()} != {num_valid_tokens}'
+                            assert num_valid_tokens == (recv_layout_range & int_mask).sum().item(), f'{num_valid_tokens} != {recv_layout_range & int_mask}.sum().item()'
+                            assert num_valid_tokens == (all_topk_idx == expert_id).sum(dim=[1, 2])[mask_status == 0].sum().item(), f'{num_valid_tokens} != {(all_topk_idx == expert_id).sum(dim=[1, 2])[mask_status == 0].sum().item()}'
+
+                            if num_valid_tokens == 0:
+                                continue
+                            # Check received data
+                            if current_x is x:
+                                recv_x = recv_x[:num_valid_tokens]
+                                recv_x_amin = recv_x[:, :-128].amin(dim=-1)
+                                recv_src_info = recv_src_info[:num_valid_tokens]
+                                assert torch.equal(recv_x_amin, recv_x[:, :-128].amax(dim=-1))
+                                if round_scale:
+                                    assert calc_diff(recv_x[:, -1], recv_src_info.view(-1)) < 0.007
+                                else:
+                                    assert (recv_x[:, -128:] - recv_src_info.view(-1, 1) % num_tokens).sum().item() == 0
+                                for j in range(num_ranks):
+                                    if mask_status[j]:
+                                        continue
+                                    begin_idx, count = (recv_layout_range[j] >> 32).item(), (recv_layout_range[j] & int_mask).item()
+                                    if not round_scale:
+                                        assert (recv_x_amin == j - rank_offset).sum().item() == (all_topk_idx[j] == expert_id).sum().item()
+                                        assert (recv_x[begin_idx:begin_idx + count, :-128] - j + rank_offset).sum().item() == 0
+                            if dispatch_use_fp8:
+                                hash_value ^= hash_tensor(packed_recv_x[0][i, :num_valid_tokens])
+                                hash_value ^= hash_tensor(packed_recv_x[1][i, :num_valid_tokens])
+                            else:
+                                hash_value ^= hash_tensor(cast(torch.Tensor, packed_recv_x)[i, :num_valid_tokens])
+
+                        # Check combine correctness
+                        for zero_copy in (False, ) if use_logfmt else (False, True):
+                            if zero_copy:
+                                buffer.get_next_combine_buffer(handle)[:, :, :] = simulated_gemm_x
+                            out = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
+                            combined_x, event, hook = buffer.combine(
+                                simulated_gemm_x, topk_idx, topk_weights, handle,
+                                use_logfmt=use_logfmt,
+                                async_finish=not return_recv_hook, zero_copy=zero_copy,
+                                return_recv_hook=return_recv_hook, out=out)
+                            hook() if return_recv_hook else event.current_stream_wait()
+                            # Query mask buffer again after combine
+                            buffer.query_mask_buffer(mask_status)
+                            if do_check:
+                                # Adjust topk_idx for validation: mark selections from masked ranks as -1
+                                owner_by_expert = (torch.arange(num_experts, device='cuda') // num_local_experts)
+                                fail_owner_mask = (mask_status != 0).index_select(0, owner_by_expert)
+                                valid_topk_idx = topk_idx >= 0
+                                failed_topk_idx = torch.zeros_like(topk_idx, device='cuda', dtype=torch.bool)
+                                failed_topk_idx[valid_topk_idx] = fail_owner_mask.index_select(0, topk_idx[valid_topk_idx])
+                                topk_idx[failed_topk_idx] = -1
+                                diff = calc_diff(current_x * topk_weights.masked_fill(topk_idx == -1, 0).sum(dim=1).view(-1, 1), combined_x)
+                                assert torch.isnan(combined_x).sum().item() == 0
+                                assert diff < (9e-4 if dispatch_use_fp8 else 1e-5), f'Error: {diff=}, {dispatch_use_fp8=}, {zero_copy=}'
+                                hash_value ^= hash_tensor(combined_x)
+
+    # noinspection PyShadowingNames
+    def large_gemm_with_hook(hook):
+        mat_0 = torch.randn((8192, 8192), dtype=torch.float)
+        mat_1 = torch.randn((8192, 8192), dtype=torch.float)
+        mat_0 @ mat_1
+        hook()
+
+    # noinspection PyShadowingNames
+    def test_func(return_recv_hook: bool):
+        recv_x, recv_count, handle, event, hook = buffer.dispatch(
+            current_x, topk_idx, num_tokens, num_experts,
+            cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats,
+            use_fp8=True, async_finish=False, return_recv_hook=return_recv_hook)
+        return_recv_hook and large_gemm_with_hook(hook)
+        combined_x, event, hook = buffer.combine(simulated_gemm_x, topk_idx, topk_weights, handle,
+                                                 use_logfmt=use_logfmt, return_recv_hook=return_recv_hook)
+        return_recv_hook and large_gemm_with_hook(hook)
+
+    def test_barrier():
+        buffer.barrier()
+
+    # Calculate bandwidth
+    num_fp8_bytes, num_bf16_bytes = (hidden + hidden / 128 * 4 + 16), hidden * 2
+    num_logfmt10_bytes = hidden * 10 / 8 + hidden / 128 * 4
+    num_dispatch_comm_bytes, num_combine_comm_bytes = 0, 0
+    for i in range(num_tokens):
+        num_selections = (topk_idx[i] != -1).sum().item()
+        num_dispatch_comm_bytes += num_fp8_bytes * num_selections
+        num_combine_comm_bytes += (num_logfmt10_bytes if use_logfmt else num_bf16_bytes) * num_selections
+
+    # Dispatch + combine testing
+    avg_t, min_t, max_t = bench(partial(test_func, return_recv_hook=False))
+    print(f'[rank {rank}] Dispatch + combine bandwidth: {(num_dispatch_comm_bytes + num_combine_comm_bytes) / 1e9 / avg_t:.2f} GB/s, '
+          f'avg_t={avg_t * 1e6:.2f} us, min_t={min_t * 1e6:.2f} us, max_t={max_t * 1e6:.2f} us', flush=True)
+
+    # Separate profiling
+    if not kineto:
+        return
+
+    for return_recv_hook in (False, True):
+        buffer.barrier()
+        dispatch_t, combine_t = bench_kineto(partial(test_func, return_recv_hook=return_recv_hook),
+                                             kernel_names=('dispatch', 'combine'), barrier_comm_profiling=True,
+                                             suppress_kineto_output=False, num_kernels_per_period=2 if return_recv_hook else 1,
+                                             barrier_fn=test_barrier)
+        if not return_recv_hook:
+            print(f'[rank {rank}] Dispatch bandwidth: {num_dispatch_comm_bytes / 1e9 / dispatch_t:.2f} GB/s, avg_t={dispatch_t * 1e6:.2f} us | '
+                  f'Combine bandwidth: {num_combine_comm_bytes / 1e9 / combine_t:.2f} GB/s, avg_t={combine_t * 1e6:.2f} us', flush=True)
+        else:
+            print(f'[rank {rank}] Dispatch send/recv time: {dispatch_t[0] * 1e6:.2f} + {dispatch_t[1] * 1e6:.2f} us | '
+                  f'Combine send/recv time: {combine_t[0] * 1e6:.2f} + {combine_t[1] * 1e6:.2f} us', flush=True)
+
+
+def worker(torch_rank: int, args: argparse.Namespace):
+    rank_client = rank_server.RankClient(args.rank_server if args.rank_server else "127.0.0.1")
+    local_rank, global_rank, last_active_phase = rank_client.get_rank()
+    plan = Plan(args.plan, global_rank, start_phase=last_active_phase if last_active_phase is not None else 0)
+    if plan.current_phase == -1:
+        print(f"Process {torch_rank} -> no plan phases were found for rank {global_rank} after phase {last_active_phase}, exiting", flush=True)
+        return
+
+    max_num_ranks = plan.get_max_rank() + 1
+    print(f"Process {torch_rank} -> global_rank={global_rank}, local_rank={local_rank}", flush=True)
+
+    # Initialize torch
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(local_rank % 8)
+    torch.set_default_dtype(torch.bfloat16)
+    torch.set_default_device('cuda')
+    torch.cuda.set_device(0)
+
+    # Initialize UCX
+    pxb_nics = ["mlx5_0", "mlx5_3", "mlx5_4", "mlx5_5", "mlx5_6", "mlx5_9", "mlx5_10", "mlx5_11"]
+    tcp_nics = ',ibp154s0,ibp192s0,ibp206s0,ibp220s0,ibp94s0'
+    os.environ['UCX_NET_DEVICES'] = f'cuda0-{pxb_nics[local_rank]}:1' + tcp_nics
+
+    # Initialize NIXL
+    os.environ['NIXL_ETCD_ENDPOINTS'] = args.etcd_server
+
+    # Initialize nixl_ep buffer
+    num_rdma_bytes = nixl_ep.Buffer.get_rdma_size_hint(args.num_tokens, args.hidden_dim, max_num_ranks, args.num_experts_per_rank * max_num_ranks)
+    if local_rank == 0:
+        print(f'Allocating buffer size: {num_rdma_bytes / 1e6} MB ...', flush=True)
+
+    buffer = nixl_ep.Buffer(rank=global_rank, nvlink_backend=args.nvlink_backend, explicitly_destroy=True, enable_shrink=True)
+    buffer.update_memory_buffers(num_ranks=max_num_ranks, num_experts_per_rank=args.num_experts_per_rank, num_rdma_bytes=num_rdma_bytes)
+    signal.signal(signal.SIGTERM, partial(handle_sigterm, buffer=buffer, plan=plan, rank_client=rank_client))
+    remote_ranks = set()
+    mask_status = torch.zeros((max_num_ranks, ), dtype=torch.int32, device='cuda')
+
+    while True:
+        print(f"global_rank={global_rank}, local_rank={local_rank} -> start phase {plan.get_phase()}", flush=True)
+
+        added_ranks = plan.get_new_ranks()
+        cleanly_removed = plan.get_removed_ranks()
+        ranks_to_kill = plan.get_killed_ranks()
+
+        # If this rank is being removed in this phase, exit gracefully
+        if global_rank in cleanly_removed:
+            print(f"global_rank={global_rank}, local_rank={local_rank} -> this rank is being removed in this phase, exiting", flush=True)
+            rank_client.release_rank(user_context=plan.get_phase())
+            break
+
+        if len(added_ranks) > 0:
+            print(f"global_rank={global_rank}, local_rank={local_rank} -> adding connections to {added_ranks}", flush=True)
+            buffer.connect_ranks(added_ranks)
+            remote_ranks.update(added_ranks)
+
+        # Check if this rank should be killed
+        kill_rank = global_rank in ranks_to_kill
+
+        if len(cleanly_removed) > 0:
+            print(f"global_rank={global_rank}, local_rank={local_rank} -> removing connections to {cleanly_removed}", flush=True)
+            buffer.disconnect_ranks(cleanly_removed)
+            remote_ranks.difference_update(cleanly_removed)
+            time.sleep(5)  # required to avoid race between MD invalidation and readdition of same ranks, if this is part of the test
+
+        # Use sparse num_ranks = max(active_ranks) + 1 for proper indexing
+        active_ranks_list = plan.get_active_ranks()
+        current_num_ranks = max(active_ranks_list) + 1  # Sparse indexing
+        current_num_experts = args.num_experts_per_rank * current_num_ranks
+
+        test_main(args.num_tokens, args.hidden_dim, current_num_experts, args.num_topk,
+                  global_rank, current_num_ranks, max_num_ranks, buffer, kineto=args.kineto, fault_tolerance_test=kill_rank)
+        # Query mask buffer to detect any unexpected rank failures and clean them up
+        buffer.query_mask_buffer(mask_status)
+        newly_failed_ranks = set()
+        for r in range(current_num_ranks):
+            if mask_status[r].item() != 0 and r in remote_ranks:
+                newly_failed_ranks.add(r)
+
+        if len(newly_failed_ranks) > 0:
+            print(f"global_rank={global_rank}, local_rank={local_rank} -> detected unexpected rank failures: {newly_failed_ranks}, cleaning up...", flush=True)
+            remote_ranks.difference_update(newly_failed_ranks)
+            buffer.disconnect_ranks(list(newly_failed_ranks))
+            time.sleep(5)
+
+        print(f"global_rank={global_rank}, local_rank={local_rank} -> end phase {plan.get_phase()}", flush=True)
+
+        if not plan.next():
+            break
+
+    buffer.destroy()
+
+    print(f"global_rank={global_rank}, local_rank={local_rank} -> done", flush=True)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Elastic EP Test")
+    parser.add_argument("--plan", type=str, default="plan.json", help="Path to plan file")
+    parser.add_argument("--num-processes", type=int, default=8, help="Number of worker processes to launch")
+    parser.add_argument("--num-tokens", type=int, default=128, help="Number of tokens")
+    parser.add_argument("--num-experts-per-rank", type=int, default=2, help="Number of experts per rank")
+    parser.add_argument("--hidden-dim", type=int, default=7168, help="Hidden dimension")
+    parser.add_argument("--num-topk", type=int, default=8, help="Number of topk")
+    parser.add_argument("--etcd-server", type=str, default="http://127.0.0.1:2379", help="ETCD server address for NIXL (default: http://127.0.0.1:2379)")
+    parser.add_argument("--rank-server", type=str, help="Rank server address. If not set, a rank server will be started locally and will be killed after all the workers launched in this run are finished.")
+    parser.add_argument("--kineto", action="store_true", help="Enable kineto profiling")
+    parser.add_argument('--nvlink-backend', choices=['nixl', 'ipc', 'none'], default='ipc', help='NVLink backend to use')
+
+    args = parser.parse_args()
+
+    # TODO_Itay: remove this once NIXL supports NVLink on multiple workers
+    assert args.nvlink_backend != 'nixl', "NIXL does not support NVLink on multiple workers yet"
+    rank_server_process = None
+    if not args.rank_server:
+        print("Starting rank server locally", flush=True)
+        rank_server_process = torch.multiprocessing.Process(target=rank_server.start_server, daemon=True)
+        rank_server_process.start()
+        time.sleep(0.5)
+    if args.num_processes == 1:
+        worker(0, args)
+        return
+
+    ctx = torch.multiprocessing.spawn(
+        worker,
+        args=(args,),
+        nprocs=args.num_processes,
+        join=False,
+        daemon=False,
+        start_method="spawn",)
+
+    for p in ctx.processes:
+        try:
+            p.join()
+        except Exception:
+            pass
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/device/ep/tests/elastic/expansion_contraction.json b/examples/device/ep/tests/elastic/expansion_contraction.json
new file mode 100644
index 0000000000..c5d14a708c
--- /dev/null
+++ b/examples/device/ep/tests/elastic/expansion_contraction.json
@@ -0,0 +1,6 @@
+[
+  [0, 1, 2, 3],
+  [0, 1, 2, 3, 4, 5, 6, 7],
+  [0, 1, 2, 3, 4, -6, 7],
+  [0, 1, 2, 3, 4, 5, 6, 7]
+]
diff --git a/examples/device/ep/tests/elastic/no_expansion.json b/examples/device/ep/tests/elastic/no_expansion.json
new file mode 100644
index 0000000000..3dbb524863
--- /dev/null
+++ b/examples/device/ep/tests/elastic/no_expansion.json
@@ -0,0 +1,3 @@
+[
+  [0, 1, 2, 3]
+]
diff --git a/examples/device/ep/tests/elastic/plan.py b/examples/device/ep/tests/elastic/plan.py
new file mode 100644
index 0000000000..23036bdbc1
--- /dev/null
+++ b/examples/device/ep/tests/elastic/plan.py
@@ -0,0 +1,95 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import json
+
+
+class Plan:
+    def __init__(self, plan_path: str, rank: int, start_phase: int = 0):
+        """Initialize plan for a specific rank."""
+        with open(plan_path, "r") as f:
+            self.phases = json.load(f)
+        self.rank = rank
+        # Auto-detect starting phase for this rank
+        self.current_phase = self._find_starting_phase(start_phase)
+        self.starting_phase = self.current_phase  # Store the starting phase
+
+    def _find_starting_phase(self, start_search_from_phase: int) -> int:
+        """Find the first phase where this rank appears after the start_search_from_phase"""
+
+        for i in range(start_search_from_phase, len(self.phases)):
+            if self.rank in self.phases[i]:
+                return i
+
+        return -1
+
+    def get_new_ranks(self):
+        """Get ranks to connect to at current phase (only the new ones)."""
+        if self.current_phase == self.starting_phase:
+            # First phase for this rank: connect to all other ranks in the phase (only positive)
+            return [r for r in self.phases[self.current_phase] if r >= 0 and r != self.rank]
+        else:
+            # Later phases: only connect to newly added ranks
+            prev_ranks = set(self.phases[self.current_phase - 1])
+            curr_ranks = set(self.phases[self.current_phase])
+            new_ranks = curr_ranks - prev_ranks
+            return [r for r in new_ranks if r >= 0 and r != self.rank]
+
+    def get_removed_ranks(self):
+        """Get ranks to remove at current phase."""
+        if self.current_phase == self.starting_phase:
+            # First phase: no ranks to remove
+            return []
+
+        # Use absolute values for comparison
+        prev_ranks_abs = set(abs(r) for r in self.phases[self.current_phase - 1])
+        curr_ranks_abs = set(abs(r) for r in self.phases[self.current_phase])
+
+        # Find negative ranks in current and previous phase
+        curr_negative_abs = set(abs(r) for r in self.phases[self.current_phase] if r < 0)
+        prev_negative_abs = set(abs(r) for r in self.phases[self.current_phase - 1] if r < 0)
+
+        # Ranks cleanly removed: in prev (absolute) but not in curr (absolute),
+        # excluding those being killed now and those that were killed before
+        cleanly_removed = list(prev_ranks_abs - curr_ranks_abs - curr_negative_abs - prev_negative_abs)
+
+        return cleanly_removed
+
+    def get_killed_ranks(self):
+        """Get ranks that were killed at current phase."""
+        return list(set(abs(r) for r in self.phases[self.current_phase] if r < 0))
+
+    def get_active_ranks(self):
+        """Get all active ranks in current phase."""
+        return [r for r in self.phases[self.current_phase] if r >= 0]
+
+    def get_max_rank(self):
+        """Get the maximum participating rank index"""
+        return max(max(phase) for phase in self.phases)
+
+    def get_min_active_ranks(self):
+        """Get the minimum number of active ranks in all phases."""
+        return min([len(phase) for phase in self.phases])
+
+    def next(self):
+        """Advance to next phase."""
+        if self.current_phase < len(self.phases) - 1:
+            self.current_phase += 1
+            return True
+        return False
+
+    def get_phase(self):
+        """Get current phase index."""
+        return self.current_phase
diff --git a/examples/device/ep/tests/elastic/rank_server.py b/examples/device/ep/tests/elastic/rank_server.py
new file mode 100644
index 0000000000..7b84cc0bda
--- /dev/null
+++ b/examples/device/ep/tests/elastic/rank_server.py
@@ -0,0 +1,141 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import os
+import socket
+import time
+from collections import defaultdict
+from socketserver import StreamRequestHandler, ThreadingTCPServer
+from threading import Lock
+
+import torch
+
+
+# --- Request handler ---
+class RankServerHandler(StreamRequestHandler):
+    _counts: defaultdict[str, list[int]] = defaultdict(list)  # List of used local ranks per host
+    _global: int = 0
+    _lock: Lock = Lock()
+    _rank_to_host: dict[int, tuple[str, int]] = {}  # Maps global rank to (host, local_rank)
+    _user_context: dict[str, str | None] = {}
+    _all_global_ranks: set[int] = set()
+    _removed_global_ranks: set[int] = set()
+
+    def handle(self):
+        with self._lock:
+            line = self.rfile.readline().strip().decode()
+
+            if line.startswith("RELEASE_RANK"):
+                # Handle rank release - remove from used ranks list
+                if len(self._all_global_ranks) > 0:
+                    # Find the highest assigned rank and release it
+                    rank_to_release = int(line.split()[1])
+                    self._all_global_ranks.remove(rank_to_release)
+                    self._removed_global_ranks.add(rank_to_release)
+                    self._user_context[str(rank_to_release)] = line.split()[2]
+
+                    if rank_to_release in self._rank_to_host:
+                        host, local_rank = self._rank_to_host[rank_to_release]
+                        # Remove this local rank from the used list
+                        if local_rank in self._counts[host]:
+                            self._counts[host].remove(local_rank)
+                        # Remove the mapping
+                        del self._rank_to_host[rank_to_release]
+                    self.wfile.write("OK\n".encode())
+                else:
+                    self.wfile.write("ERROR: No ranks to release\n".encode())
+            else:
+                # Handle rank assignment - find lowest unused local rank
+                host = line
+                # Find the lowest unused local rank for this host
+                used_ranks = set(self._counts[host])
+                local = 0
+                while local in used_ranks:
+                    local += 1
+
+                # Add this local rank to the used list
+                self._counts[host].append(local)
+                if len(self._removed_global_ranks) == 0:
+                    global_rank = len(self._all_global_ranks)
+                else:
+                    global_rank = min(self._removed_global_ranks)
+                    self._removed_global_ranks.remove(global_rank)
+
+                self._all_global_ranks.add(global_rank)
+                # Record which host and local rank this global rank maps to
+                self._rank_to_host[global_rank] = (host, local)
+                if str(global_rank) not in self._user_context:
+                    self._user_context[str(global_rank)] = None
+                self.wfile.write(f"{local} {global_rank} {self._user_context[str(global_rank)]}\n".encode()) if self._user_context[str(global_rank)] is not None else self.wfile.write(f"{local} {global_rank}\n".encode())
+
+
+# --- TCPServer subclass to reuse port immediately ---
+class ReusableTCPServer(ThreadingTCPServer):
+    allow_reuse_address = True
+
+
+# --- Lazy-start server ---
+def start_server(port: int = 9999) -> None:
+    try:
+        server = ReusableTCPServer(("0.0.0.0", port), RankServerHandler)
+        server.serve_forever()
+    except OSError:
+        pass  # another process already started the server
+
+
+def start_server_process(port: int = 9999):
+    server_process = torch.multiprocessing.Process(target=start_server, args=(port,), daemon=False)
+    server_process.start()
+    time.sleep(1)
+    return server_process
+
+
+# --- Client API ---
+class RankClient:
+
+    def __init__(self, server: str = "127.0.0.1", port: int = 9999):
+        self.self_global_rank = None
+        self.server = server
+        self.port = port
+
+    def get_rank(self):
+        if self.self_global_rank is not None:
+            print(f"WARNING: rank already assigned - returning existing rank {self.self_global_rank}", flush=True)
+            return self.self_global_rank
+        s = socket.create_connection((self.server, self.port))
+        s.sendall(f"{os.uname().nodename}\n".encode())
+        server_response = s.recv(1024).decode().split()
+        if len(server_response) == 2:
+            local_rank, global_rank = tuple(map(int, server_response))
+            user_context = None
+        else:
+            local_rank, global_rank, user_context = tuple(map(int, server_response))
+        s.close()
+        self.self_global_rank = global_rank
+        return local_rank, global_rank, user_context
+
+    def release_rank(self, user_context: str | None = None) -> bool:
+        """Release a rank (decrement the global counter by 1)"""
+        s = socket.create_connection((self.server, self.port))
+        s.sendall(f"RELEASE_RANK {self.self_global_rank if self.self_global_rank is not None else -1} {user_context}\n".encode())
+        response = s.recv(1024).decode().strip()
+        s.close()
+        self.self_global_rank = None
+        return response == "OK"
+
+
+# --- Example usage ---
+if __name__ == "__main__":
+    start_server()
diff --git a/examples/device/ep/tests/elastic/single_expansion.json b/examples/device/ep/tests/elastic/single_expansion.json
new file mode 100644
index 0000000000..fc0b73007a
--- /dev/null
+++ b/examples/device/ep/tests/elastic/single_expansion.json
@@ -0,0 +1,4 @@
+[
+  [0, 1, 2, 3],
+  [0, 1, 2, 3, 4, 5, 6, 7]
+]
diff --git a/examples/device/ep/tests/elastic/single_expansion_16_ranks.json b/examples/device/ep/tests/elastic/single_expansion_16_ranks.json
new file mode 100644
index 0000000000..2871773c0a
--- /dev/null
+++ b/examples/device/ep/tests/elastic/single_expansion_16_ranks.json
@@ -0,0 +1,4 @@
+[
+  [0, 1, 2, 3, 4, 5, 6, 7],
+  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
+]
diff --git a/examples/device/ep/tests/elastic/single_rank.json b/examples/device/ep/tests/elastic/single_rank.json
new file mode 100644
index 0000000000..903ce72321
--- /dev/null
+++ b/examples/device/ep/tests/elastic/single_rank.json
@@ -0,0 +1,3 @@
+[
+  [0]
+]
diff --git a/examples/device/ep/tests/utils.py b/examples/device/ep/tests/utils.py
new file mode 100644
index 0000000000..72f85823e5
--- /dev/null
+++ b/examples/device/ep/tests/utils.py
@@ -0,0 +1,247 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This file incorporates material from the DeepSeek project, licensed under the MIT License.
+# The modifications made by NVIDIA are licensed under the Apache License, Version 2.0.
+#
+# SPDX-License-Identifier: MIT AND Apache-2.0
+#
+# 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.
+
+import inspect
+import json
+import os
+import sys
+import tempfile
+from pathlib import Path
+from typing import Callable, Optional, Union
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+
+def init_dist(local_rank: int, num_local_ranks: int):
+    # NOTES: you may rewrite this function with your own cluster settings
+    ip = os.getenv('MASTER_ADDR', '127.0.0.1')
+    port = int(os.getenv('MASTER_PORT', '8361'))
+    num_nodes = int(os.getenv('WORLD_SIZE', 1))
+    node_rank = int(os.getenv('RANK', 0))
+
+    sig = inspect.signature(dist.init_process_group)
+    params = {
+        'backend': 'nccl',
+        'init_method': f'tcp://{ip}:{port}',
+        'world_size': num_nodes * num_local_ranks,
+        'rank': node_rank * num_local_ranks + local_rank,
+    }
+    if 'device_id' in sig.parameters:
+        # noinspection PyTypeChecker
+        params['device_id'] = torch.device('cuda:0')  # TODO: restore to local_rank once ucx fixes lane-selection
+    dist.init_process_group(**params)
+    torch.set_default_dtype(torch.bfloat16)
+    torch.set_default_device('cuda')
+    torch.cuda.set_device(0)  # TODO: restore to local_rank once ucx fixes lane-selection
+
+    return dist.get_rank(), dist.get_world_size(), dist.new_group(list(range(num_local_ranks * num_nodes)))
+
+
+def calc_diff(x: torch.Tensor, y: torch.Tensor):
+    x, y = x.double() + 1, y.double() + 1
+    denominator = (x * x + y * y).sum()
+    sim = 2 * (x * y).sum() / denominator
+    return (1 - sim).item()
+
+
+def per_token_cast_to_fp8(x: torch.Tensor):
+    assert x.dim() == 2 and x.size(1) % 128 == 0
+    m, n = x.shape
+    x_view = x.view(m, -1, 128)
+    x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
+    return (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view(m, n), (x_amax / 448.0).view(m, -1)
+
+
+def per_token_cast_back(x_fp8: torch.Tensor, x_scales: torch.Tensor):
+    if x_fp8.numel() == 0:
+        return x_fp8.to(torch.bfloat16)
+    if x_scales.dtype == torch.int:
+        x_scales = x_scales.view(dtype=torch.uint8).to(torch.int) << 23
+        x_scales = x_scales.view(dtype=torch.float)
+    x_fp32 = x_fp8.to(torch.float32).view(x_fp8.size(0), -1, 128)
+    x_scales = x_scales.view(x_fp8.size(0), -1, 1)
+    return (x_fp32 * x_scales).view(x_fp8.shape).to(torch.bfloat16)
+
+
+def inplace_unique(x: torch.Tensor, num_slots: int):
+    assert x.dim() == 2
+    mask = x < 0
+    x_padded = x.masked_fill(mask, num_slots)
+    bin_count = torch.zeros((x.size(0), num_slots + 1), dtype=x.dtype, device=x.device)
+    bin_count.scatter_add_(1, x_padded, torch.ones_like(x_padded))
+    bin_count = bin_count[:, :num_slots]
+    sorted_bin_count, sorted_bin_idx = torch.sort(bin_count, dim=-1, descending=True)
+    sorted_bin_idx.masked_fill_(sorted_bin_count == 0, -1)
+    sorted_bin_idx = torch.sort(sorted_bin_idx, descending=True, dim=-1).values
+    x[:, :].fill_(-1)
+    valid_len = min(num_slots, x.size(1))
+    x[:, :valid_len] = sorted_bin_idx[:, :valid_len]
+
+
+def create_grouped_scores(scores: torch.Tensor, group_idx: torch.Tensor, num_groups: int):
+    num_tokens, num_experts = scores.shape
+    scores = scores.view(num_tokens, num_groups, -1)
+    mask = torch.zeros((num_tokens, num_groups), dtype=torch.bool, device=scores.device)
+    mask = mask.scatter_(1, group_idx, True).unsqueeze(-1).expand_as(scores)
+    return (scores * mask).view(num_tokens, num_experts)
+
+
+def bench(fn, num_warmups: int = 50, num_tests: int = 50, post_fn=None):
+    # Flush L2 cache with 256 MB data
+    torch.cuda.synchronize()
+    cache = torch.empty(int(256e6 // 4), dtype=torch.int, device='cuda')
+
+    # Warmup
+    for _ in range(num_warmups):
+        fn()
+
+    # Flush L2
+    cache.zero_()
+
+    # Testing
+    start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)]
+    end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)]
+    for i in range(num_tests):
+        # Record
+        start_events[i].record()
+        fn()
+        end_events[i].record()
+        if post_fn is not None:
+            post_fn()
+    torch.cuda.synchronize()
+
+    times = np.array([s.elapsed_time(e) / 1e3 for s, e in zip(start_events, end_events)])[1:]
+    return np.average(times), np.min(times), np.max(times)
+
+
+class empty_suppress:
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_):
+        pass
+
+
+class suppress_stdout_stderr:
+    def __enter__(self):
+        self.outnull_file = open(os.devnull, 'w')
+        self.errnull_file = open(os.devnull, 'w')
+
+        self.old_stdout_fileno_undup = sys.stdout.fileno()
+        self.old_stderr_fileno_undup = sys.stderr.fileno()
+
+        self.old_stdout_fileno = os.dup(sys.stdout.fileno())
+        self.old_stderr_fileno = os.dup(sys.stderr.fileno())
+
+        self.old_stdout = sys.stdout
+        self.old_stderr = sys.stderr
+
+        os.dup2(self.outnull_file.fileno(), self.old_stdout_fileno_undup)
+        os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup)
+
+        sys.stdout = self.outnull_file
+        sys.stderr = self.errnull_file
+        return self
+
+    def __exit__(self, *_):
+        sys.stdout = self.old_stdout
+        sys.stderr = self.old_stderr
+
+        os.dup2(self.old_stdout_fileno, self.old_stdout_fileno_undup)
+        os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup)
+
+        os.close(self.old_stdout_fileno)
+        os.close(self.old_stderr_fileno)
+
+        self.outnull_file.close()
+        self.errnull_file.close()
+
+
+def bench_kineto(fn, kernel_names: Union[str, tuple], num_tests: int = 30, suppress_kineto_output: bool = False,
+                 trace_path: Optional[str] = None, barrier_comm_profiling: bool = False,
+                 num_kernels_per_period: int = 1, barrier_fn: Optional[Callable] = None):
+    # Profile
+    suppress = suppress_stdout_stderr if suppress_kineto_output else empty_suppress
+    with suppress():
+        schedule = torch.profiler.schedule(wait=0, warmup=1, active=1, repeat=1)
+        with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA], schedule=schedule) as prof:
+            for i in range(2):
+                # NOTES: use a large kernel and a barrier to eliminate the unbalanced CPU launch overhead
+                if barrier_comm_profiling:
+                    lhs = torch.randn((8192, 8192), dtype=torch.float, device='cuda')
+                    rhs = torch.randn((8192, 8192), dtype=torch.float, device='cuda')
+                    lhs @ rhs
+                    if barrier_fn is not None:
+                        barrier_fn()
+                    else:
+                        dist.all_reduce(torch.ones(1, dtype=torch.float, device='cuda'))
+                for _ in range(num_tests):
+                    fn()
+                torch.cuda.synchronize()
+                prof.step()
+
+    # Parse the profiling table
+    assert isinstance(kernel_names, str) or isinstance(kernel_names, tuple)
+    is_tuple = isinstance(kernel_names, tuple)
+    prof_lines = prof.key_averages().table(sort_by='cuda_time_total', max_name_column_width=100).split('\n')
+    kernel_names = (kernel_names, ) if isinstance(kernel_names, str) else kernel_names
+    assert all([isinstance(name, str) for name in kernel_names])
+    for name in kernel_names:
+        assert sum([name in line for line in prof_lines]) == 1, f'Errors of the kernel {name} in the profiling table'
+
+    # Save chrome traces
+    if trace_path is not None:
+        prof.export_chrome_trace(trace_path)
+
+    # Return average kernel durations
+    units = {'ms': 1e3, 'us': 1e6}
+    kernel_durations = []
+    for name in kernel_names:
+        for line in prof_lines:
+            if name in line:
+                time_str = line.split()[-2]
+                for unit, scale in units.items():
+                    if unit in time_str:
+                        kernel_durations.append(float(time_str.replace(unit, '')) / scale)
+                        break
+                break
+
+    # Expand the kernels by periods
+    if num_kernels_per_period > 1:
+        with tempfile.NamedTemporaryFile(suffix='.json') as tmp:
+            prof.export_chrome_trace(tmp.name)
+            profile_data = json.loads(Path(tmp.name).read_text())
+
+        for i, kernel_name in enumerate(kernel_names):
+            events = [event for event in profile_data['traceEvents'] if f'::{kernel_name}' in event['name']]
+            events = sorted(events, key=lambda event: event['ts'])
+            durations = [event['dur'] / 1e6 for event in events]
+            assert len(durations) % num_kernels_per_period == 0
+            num_kernel_patterns = len(durations) // num_kernels_per_period
+            kernel_durations[i] = [sum(durations[j::num_kernels_per_period]) / num_kernel_patterns for j in range(num_kernels_per_period)]  # type: ignore[call-overload]
+
+    # Return execution durations
+    return kernel_durations if is_tuple else kernel_durations[0]
+
+
+def hash_tensor(t: torch.Tensor):
+    return t.view(torch.int).sum().item()
diff --git a/examples/device/meson.build b/examples/device/meson.build
new file mode 100644
index 0000000000..c399921327
--- /dev/null
+++ b/examples/device/meson.build
@@ -0,0 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 get_option('build_nixl_ep')
+    subdir('ep')
+endif
diff --git a/examples/meson.build b/examples/meson.build
index 73bd4defff..95d7ed0164 100644
--- a/examples/meson.build
+++ b/examples/meson.build
@@ -14,3 +14,4 @@
 # limitations under the License.
 
 subdir('cpp')
+subdir('device')
diff --git a/meson.build b/meson.build
index d9f2a7ebe5..e0df99d763 100644
--- a/meson.build
+++ b/meson.build
@@ -91,7 +91,10 @@ if cuda_dep.found()
 
     # Please extend with your arch if not present in the list
     nvcc_flags = []
-    nvcc_flags += ['-gencode', 'arch=compute_80,code=sm_80']
+    # nixl ep cannot be built with sm_80
+    if not get_option('build_nixl_ep')
+        nvcc_flags += ['-gencode', 'arch=compute_80,code=sm_80']
+    endif
     nvcc_flags += ['-gencode', 'arch=compute_90,code=sm_90']
     add_project_arguments(nvcc_flags, language: 'cuda')
 
@@ -100,7 +103,10 @@ if cuda_dep.found()
     add_project_arguments('-rdc=true', language: 'cuda')
 
     nvcc_flags_link = []
-    nvcc_flags_link += ['-gencode=arch=compute_80,code=sm_80']
+    # nixl ep cannot be built with sm_80
+    if not get_option('build_nixl_ep')
+        nvcc_flags_link += ['-gencode=arch=compute_80,code=sm_80']
+    endif
     nvcc_flags_link += ['-gencode=arch=compute_90,code=sm_90']
     add_project_link_arguments(nvcc_flags_link, language: 'cuda')
     message('nvcc version: ' + nvcc.version())
diff --git a/meson_options.txt b/meson_options.txt
index fc28c08f2a..d4bf74eed0 100644
--- a/meson_options.txt
+++ b/meson_options.txt
@@ -32,4 +32,5 @@ option('rust', type: 'boolean', value: false, description: 'Build Rust bindings'
 # Tests
 option('build_tests', type: 'boolean', value: true, description: 'Build all tests')
 option('build_examples', type: 'boolean', value: true, description: 'Build all examples')
+option('build_nixl_ep', type: 'boolean', value: false, description: 'Build nixl_ep example (requires sm_90)')
 option('test_all_plugins', type: 'boolean', value: false, description: 'Testing all plugins in addition to the mocks..')
diff --git a/src/core/meson.build b/src/core/meson.build
index 54e6344d2a..c5f05da591 100644
--- a/src/core/meson.build
+++ b/src/core/meson.build
@@ -61,6 +61,7 @@ nixl_lib = library('nixl',
                    include_directories: [ nixl_inc_dirs, utils_inc_dirs ],
                    link_args: ['-lstdc++fs'],
                    dependencies: nixl_lib_deps,
-                   install: true)
+                   install: true,
+                   install_rpath: join_paths(get_option('prefix'), get_option('libdir')))
 
 nixl_dep = declare_dependency(link_with: nixl_lib, include_directories: nixl_inc_dirs)
diff --git a/src/plugins/ucx/meson.build b/src/plugins/ucx/meson.build
index aebfa92f91..c3debe9962 100644
--- a/src/plugins/ucx/meson.build
+++ b/src/plugins/ucx/meson.build
@@ -38,7 +38,8 @@ else
                cpp_args : compile_flags + ['-fPIC'],
                name_prefix: 'libplugin_',  # Custom prefix for plugin libraries
                install_dir: plugin_install_dir,
-               install_rpath: '$ORIGIN/..')
+               build_rpath: get_option('ucx_path') != '' ? get_option('ucx_path') + '/lib' : '',
+               install_rpath: get_option('ucx_path') != '' ? '$ORIGIN/..:' + get_option('ucx_path') + '/lib' : '$ORIGIN/..')
 
     if get_option('buildtype') == 'debug'
         run_command('sh', '-c',
diff --git a/src/utils/ucx/meson.build b/src/utils/ucx/meson.build
index cf70169441..60a6dfc93b 100644
--- a/src/utils/ucx/meson.build
+++ b/src/utils/ucx/meson.build
@@ -27,4 +27,6 @@ ucx_utils_lib = library('ucx_utils',
            ucx_utils_sources,
            dependencies: ucx_utils_dep,
            include_directories: [ nixl_inc_dirs, utils_inc_dirs ],
-           install: true)
+           install: true,
+           build_rpath: get_option('ucx_path') != '' ? get_option('ucx_path') + '/lib' : '',
+           install_rpath: get_option('ucx_path') != '' ? get_option('ucx_path') + '/lib' : '')