From f9ed2fae5caa8614992fe9ed0d0ccae2db31540e Mon Sep 17 00:00:00 2001 From: ovidiusm Date: Mon, 30 Mar 2026 17:19:57 +0200 Subject: [PATCH 01/59] Pin torch version to 2.11 (#1468) * NIXL EP: Add torch version check Signed-off-by: Ovidiu Mara * Pin torch version to 2.11 during wheel build Signed-off-by: Ovidiu Mara * Pin torch exactly to 2.11.* Signed-off-by: Ovidiu Mara * Add linter ignore Signed-off-by: Ovidiu Mara * Add copyright Signed-off-by: Ovidiu Mara --------- Signed-off-by: Ovidiu Mara --- .flake8 | 6 ++++++ contrib/Dockerfile.manylinux | 2 +- examples/device/ep/nixl_ep/__init__.py | 7 +++++++ pyproject.toml | 2 +- 4 files changed, 15 insertions(+), 2 deletions(-) create mode 100644 .flake8 diff --git a/.flake8 b/.flake8 new file mode 100644 index 00000000..2d2f3207 --- /dev/null +++ b/.flake8 @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[flake8] +per-file-ignores = + examples/device/ep/nixl_ep/__init__.py:E402 diff --git a/contrib/Dockerfile.manylinux b/contrib/Dockerfile.manylinux index 9dc8b0a2..b85ef634 100644 --- a/contrib/Dockerfile.manylinux +++ b/contrib/Dockerfile.manylinux @@ -317,7 +317,7 @@ ENV PATH="$VIRTUAL_ENV/bin:$PATH" RUN uv pip install --upgrade meson meson-python pybind11 patchelf pyYAML click setuptools tabulate auditwheel tomlkit # Install PyTorch RUN export UV_INDEX="https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d .)" && \ - uv pip install torch torchvision torchaudio + uv pip install 'torch==2.11.*' # Upgrade setuptools to latest version for compatibility with PEP 639 (license format) RUN uv pip install --upgrade 'setuptools>=80.9.0' diff --git a/examples/device/ep/nixl_ep/__init__.py b/examples/device/ep/nixl_ep/__init__.py index 719e149b..19f06cec 100644 --- a/examples/device/ep/nixl_ep/__init__.py +++ b/examples/device/ep/nixl_ep/__init__.py @@ -20,6 +20,13 @@ import torch +_torch_base_version = torch.__version__.split("+", 1)[0] +if not (_torch_base_version == "2.11" or _torch_base_version.startswith("2.11.")): + raise RuntimeError( + f"Unsupported torch version '{torch.__version__}'. " + "NIXL EP requires torch 2.11.*." + ) + from . import nixl_ep_cpp as _nixl_ep_cpp from .buffer import Buffer from .utils import EventOverlap diff --git a/pyproject.toml b/pyproject.toml index 29e5cda2..2c39306a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ # limitations under the License. [build-system] -requires = ["meson-python", "pybind11", "patchelf", "pyyaml", "types-PyYAML", "pytest", "build", "setuptools>=80.9.0", "torch"] +requires = ["meson-python", "pybind11", "patchelf", "pyyaml", "types-PyYAML", "pytest", "build", "setuptools>=80.9.0", "torch==2.11.*"] build-backend = "mesonpy" [project] From 48d4b35258ec6e21e5f7b5e15fafa74a8ea036df Mon Sep 17 00:00:00 2001 From: ovidiusm Date: Mon, 30 Mar 2026 17:23:24 +0200 Subject: [PATCH 02/59] Pin torch version to 2.11 (#1468) * NIXL EP: Add torch version check Signed-off-by: Ovidiu Mara * Pin torch version to 2.11 during wheel build Signed-off-by: Ovidiu Mara * Pin torch exactly to 2.11.* Signed-off-by: Ovidiu Mara * Add linter ignore Signed-off-by: Ovidiu Mara * Add copyright Signed-off-by: Ovidiu Mara --------- Signed-off-by: Ovidiu Mara From e0524ea3b5d9563c439a195b0b2e43aa2db97e99 Mon Sep 17 00:00:00 2001 From: Eric Raut <10216922+rauteric@users.noreply.github.com> Date: Mon, 30 Mar 2026 14:39:56 -0700 Subject: [PATCH 03/59] libfabric: Use lock for all access to endpoints (#1457) `FI_THREAD_COMPLETION` semantics require locked access to all objects bound to a completion queue, particularly endpoints. Expand the CQ mutex to cover posting operations to the endpoint as well. * In case we need to progress the CQ in retry while posting an operation, only do so if the progress thread is not enabled. Otherwise, let the progress thread handle CQ polling. Signed-off-by: Eric Raut --- src/plugins/libfabric/libfabric_backend.cpp | 4 + src/utils/libfabric/libfabric_rail.cpp | 104 ++++++++++++++------ src/utils/libfabric/libfabric_rail.h | 16 ++- 3 files changed, 92 insertions(+), 32 deletions(-) diff --git a/src/plugins/libfabric/libfabric_backend.cpp b/src/plugins/libfabric/libfabric_backend.cpp index 11c6fe79..5fed7b6e 100644 --- a/src/plugins/libfabric/libfabric_backend.cpp +++ b/src/plugins/libfabric/libfabric_backend.cpp @@ -388,6 +388,10 @@ nixlLibfabricEngine::nixlLibfabricEngine(const nixlBackendInitParams *init_param // Start Progress thread for rail completion processing if (progress_thread_enabled_) { + for (size_t i = 0; i < rail_manager.getNumRails(); ++i) { + rail_manager.getRail(i).setProgressThreadEnabled(true); + } + NIXL_DEBUG << "Starting Progress thread for rails with delay: " << progress_thread_delay_.count() << " microseconds"; progress_thread_stop_ = false; diff --git a/src/utils/libfabric/libfabric_rail.cpp b/src/utils/libfabric/libfabric_rail.cpp index edd7cc3d..278e3ea5 100644 --- a/src/utils/libfabric/libfabric_rail.cpp +++ b/src/utils/libfabric/libfabric_rail.cpp @@ -699,6 +699,11 @@ nixlLibfabricRail::setNotificationCallback(std::function callback) { xferIdCallback = callback; @@ -714,7 +719,7 @@ nixlLibfabricRail::progressCompletionQueue() const { // Only protect libfabric CQ hardware operations { - std::lock_guard cq_lock(cq_progress_mutex_); + std::lock_guard ep_lock(ep_mutex_); // Non-blocking read (used by progress thread or fallback) ret = fi_cq_read(cq, completions, NIXL_LIBFABRIC_CQ_BATCH_SIZE); @@ -997,7 +1002,11 @@ nixlLibfabricRail::postRecv(nixlLibfabricReq *req) const { NIXL_TRACE << "Posting receive on endpoint=" << endpoint << " buffer=" << req->buffer << " size=" << req->buffer_size << " context=" << &req->ctx; - int ret = fi_recvmsg(endpoint, &msg, 0); + int ret; + { + const std::lock_guard ep_lock(ep_mutex_); + ret = fi_recvmsg(endpoint, &msg, 0); + } if (ret) { NIXL_ERROR << "fi_recvmsg failed on rail " << rail_id << ": " << fi_strerror(-ret); return NIXL_ERR_BACKEND; @@ -1033,8 +1042,16 @@ nixlLibfabricRail::postSend(uint64_t immediate_data, while (true) { // Libfabric fi_senddata call - ret = fi_senddata( - endpoint, req->buffer, req->buffer_size, desc, immediate_data, dest_addr, &req->ctx); + { + const std::lock_guard ep_lock(ep_mutex_); + ret = fi_senddata(endpoint, + req->buffer, + req->buffer_size, + desc, + immediate_data, + dest_addr, + &req->ctx); + } if (ret == 0) { // Success @@ -1058,9 +1075,16 @@ nixlLibfabricRail::postSend(uint64_t immediate_data, } // Progress completion queue to drain pending completions before retry - nixl_status_t progress_status = progressCompletionQueue(); - if (progress_status == NIXL_SUCCESS) { - NIXL_TRACE << "Progressed completions on rail " << rail_id << " before retry"; + if (!progress_thread_enabled_) { + nixl_status_t progress_status = progressCompletionQueue(); + if (progress_status != NIXL_SUCCESS && progress_status != NIXL_IN_PROG) { + NIXL_ERROR << "progressCompletionQueue failed on rail " << rail_id + << " during fi_senddata retry"; + return progress_status; + } + if (progress_status == NIXL_SUCCESS) { + NIXL_TRACE << "Progressed completions on rail " << rail_id << " before retry"; + } } continue; @@ -1100,15 +1124,18 @@ nixlLibfabricRail::postWrite(const void *local_buffer, while (true) { // Libfabric fi_writedata call - ret = fi_writedata(endpoint, - local_buffer, - length, - local_desc, - immediate_data, - dest_addr, - remote_addr, - remote_key, - &req->ctx); + { + const std::lock_guard ep_lock(ep_mutex_); + ret = fi_writedata(endpoint, + local_buffer, + length, + local_desc, + immediate_data, + dest_addr, + remote_addr, + remote_key, + &req->ctx); + } if (ret == 0) { // Success @@ -1132,9 +1159,16 @@ nixlLibfabricRail::postWrite(const void *local_buffer, } // Progress completion queue to drain pending completions before retry - nixl_status_t progress_status = progressCompletionQueue(); - if (progress_status == NIXL_SUCCESS) { - NIXL_TRACE << "Progressed completions on rail " << rail_id << " before retry"; + if (!progress_thread_enabled_) { + nixl_status_t progress_status = progressCompletionQueue(); + if (progress_status != NIXL_SUCCESS && progress_status != NIXL_IN_PROG) { + NIXL_ERROR << "progressCompletionQueue failed on rail " << rail_id + << " during fi_writedata retry"; + return progress_status; + } + if (progress_status == NIXL_SUCCESS) { + NIXL_TRACE << "Progressed completions on rail " << rail_id << " before retry"; + } } continue; @@ -1173,14 +1207,17 @@ nixlLibfabricRail::postRead(void *local_buffer, while (true) { // Libfabric fi_read call - ret = fi_read(endpoint, - local_buffer, - length, - local_desc, - dest_addr, - remote_addr, - remote_key, - &req->ctx); + { + const std::lock_guard ep_lock(ep_mutex_); + ret = fi_read(endpoint, + local_buffer, + length, + local_desc, + dest_addr, + remote_addr, + remote_key, + &req->ctx); + } if (ret == 0) { // Success @@ -1204,9 +1241,16 @@ nixlLibfabricRail::postRead(void *local_buffer, } // Progress completion queue to drain pending completions before retry - nixl_status_t progress_status = progressCompletionQueue(); - if (progress_status == NIXL_SUCCESS) { - NIXL_TRACE << "Progressed completions on rail " << rail_id << " before retry"; + if (!progress_thread_enabled_) { + nixl_status_t progress_status = progressCompletionQueue(); + if (progress_status != NIXL_SUCCESS && progress_status != NIXL_IN_PROG) { + NIXL_ERROR << "progressCompletionQueue failed on rail " << rail_id + << " during fi_read retry"; + return progress_status; + } + if (progress_status == NIXL_SUCCESS) { + NIXL_TRACE << "Progressed completions on rail " << rail_id << " before retry"; + } } continue; diff --git a/src/utils/libfabric/libfabric_rail.h b/src/utils/libfabric/libfabric_rail.h index a6e0133a..e5df7541 100644 --- a/src/utils/libfabric/libfabric_rail.h +++ b/src/utils/libfabric/libfabric_rail.h @@ -335,6 +335,15 @@ class nixlLibfabricRail { void setNotificationCallback(std::function callback); + /** + * @brief Enable or disable a progress thread that handles CQ draining. + * + * @param enabled true to enable progress-thread CQ draining, + * false to disable it. + */ + void + setProgressThreadEnabled(bool enabled); + /** Set callback for XFER_ID tracking */ void setXferIdCallback(std::function callback); @@ -367,8 +376,11 @@ class nixlLibfabricRail { struct fid_cq *cq; // from rail_cqs[rail_id] struct fid_av *av; // from rail_avs[rail_id] - // CQ progress mutex to protect completion queue operations - mutable std::mutex cq_progress_mutex_; + // EP mutex to protect endpoint and CQ operations + mutable std::mutex ep_mutex_; + + // Whether a progress thread is handling CQ draining + bool progress_thread_enabled_ = false; // Callback functions std::function notificationCallback; From 88127130116e24ef6781cf281a789e527a1725c6 Mon Sep 17 00:00:00 2001 From: Colin Hirsch Date: Wed, 1 Apr 2026 15:26:02 +0200 Subject: [PATCH 04/59] TEST/GTEST: Ignore IB warning. (#1359) --- test/gtest/main.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/test/gtest/main.cpp b/test/gtest/main.cpp index 7a67a467..4cd83056 100644 --- a/test/gtest/main.cpp +++ b/test/gtest/main.cpp @@ -68,6 +68,8 @@ void ParseArguments(int argc, char **argv) { } namespace { + const std::regex + ib_regex("IB device\\(s\\) were detected, but accelerated IB support was not found"); const std::regex aws_regex("UCX version is less than 1.19, CUDA support is limited, including" " the lack of support for multi-GPU within a single process."); const std::regex non_gpu_regex("[0-9]+ NVIDIA GPU\\(s\\) were detected, but UCX CUDA support " @@ -80,6 +82,9 @@ RunAllTests() { LogProblemCounter lpc; std::list ligs; + // TODO: Remove after the CI issues spuriously triggering this message are fixed. + ligs.emplace_back(ib_regex); + if (std::getenv("AWS_BATCH_JOB_ID") != nullptr) { ligs.emplace_back(aws_regex); } From 1b83f63858062563e0c8fb9575caf61d73cbe049 Mon Sep 17 00:00:00 2001 From: ofarjon <34404710+ofirfarjun7@users.noreply.github.com> Date: Wed, 1 Apr 2026 18:16:04 +0300 Subject: [PATCH 05/59] Fix dependency issues (#1486) * use ucx_ref and ignore torch ver in module * Remove file * fix * fix --- .flake8 | 6 ------ contrib/Dockerfile | 15 ++++----------- contrib/build-container.sh | 3 +-- examples/device/ep/nixl_ep/__init__.py | 7 ------- 4 files changed, 5 insertions(+), 26 deletions(-) delete mode 100644 .flake8 diff --git a/.flake8 b/.flake8 deleted file mode 100644 index 2d2f3207..00000000 --- a/.flake8 +++ /dev/null @@ -1,6 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -[flake8] -per-file-ignores = - examples/device/ep/nixl_ep/__init__.py:E402 diff --git a/contrib/Dockerfile b/contrib/Dockerfile index a540ea20..591b4c02 100644 --- a/contrib/Dockerfile +++ b/contrib/Dockerfile @@ -179,17 +179,10 @@ RUN rm -rf /opt/hpcx/ucx RUN cd /usr/local/src && \ git clone https://github.com/openucx/ucx.git && \ - cd ucx && \ - if [ "$BUILD_NIXL_EP" = "true" ]; then \ - echo "=== BUILD_NIXL_EP=true: Using latest UCX master branch (ignoring UCX_REF=$UCX_REF) ===" && \ - UCX_COMMIT=$(git rev-parse --short HEAD) && \ - echo "Using UCX commit: $UCX_COMMIT" && \ - EXPERIMENTAL_API_PARAM="--enable-experimental-api"; \ - else \ - echo "=== Using UCX_REF=$UCX_REF ===" && \ - git checkout $UCX_REF && \ - EXPERIMENTAL_API_PARAM=""; \ - fi && \ + cd ucx && \ + echo "=== Using UCX_REF=$UCX_REF ===" && \ + git checkout $UCX_REF && \ + if [ "$BUILD_NIXL_EP" = "true" ]; then EXPERIMENTAL_API_PARAM="--enable-experimental-api"; else EXPERIMENTAL_API_PARAM=""; fi && \ ./autogen.sh && \ ./contrib/configure-release-mt \ --prefix=$UCX_PREFIX \ diff --git a/contrib/build-container.sh b/contrib/build-container.sh index 345a6e72..2e7cd14e 100755 --- a/contrib/build-container.sh +++ b/contrib/build-container.sh @@ -173,11 +173,10 @@ show_build_options() { echo "Container arch: ${ARCH}" echo "Python Versions for wheel build: ${WHL_PYTHON_VERSIONS}" echo "Wheel Platform: ${WHL_PLATFORM}" + echo "UCX Ref: ${UCX_REF}" if [ "$BUILD_NIXL_EP" = "true" ]; then - echo "UCX Ref: master (latest) - BUILD_NIXL_EP enabled" echo "NIXL EP: Enabled" else - echo "UCX Ref: ${UCX_REF}" echo "NIXL EP: Disabled" fi echo "Build Type: ${BUILD_TYPE}" diff --git a/examples/device/ep/nixl_ep/__init__.py b/examples/device/ep/nixl_ep/__init__.py index 19f06cec..719e149b 100644 --- a/examples/device/ep/nixl_ep/__init__.py +++ b/examples/device/ep/nixl_ep/__init__.py @@ -20,13 +20,6 @@ import torch -_torch_base_version = torch.__version__.split("+", 1)[0] -if not (_torch_base_version == "2.11" or _torch_base_version.startswith("2.11.")): - raise RuntimeError( - f"Unsupported torch version '{torch.__version__}'. " - "NIXL EP requires torch 2.11.*." - ) - from . import nixl_ep_cpp as _nixl_ep_cpp from .buffer import Buffer from .utils import EventOverlap From d0dcbbfac05fb01b9235d3ccb9f124dc60932559 Mon Sep 17 00:00:00 2001 From: Jason Goldschmidt <92800864+jgoldsch12@users.noreply.github.com> Date: Fri, 3 Apr 2026 18:14:10 -0400 Subject: [PATCH 06/59] Add Dell ObjectScale accelerated engine support to plugin_gtest for OBJ plugin. (#1441) * Add Dell ObjectScale accelerated engine support to plugin_gtest for OBJ plugin. Introduce NIXL_OBJ_ENDPOINT_OVERRIDE environment variable for setting the endpoint of the Dell ObjectScale cluster. Add support for VRAM_SEG tests. Signed-off-by: Jason Goldschmidt * Extract CUOBJ tests to their own plugin test file. Make endpoint override a universaly applied environment variable to allow all tests to run against a specific endpoint * Addressed code review comments * Address code review comment regarding guard naming in memory_handler.h Signed-off-by: Jason Goldschmidt * Address code review comment Signed-off-by: Jason Goldschmidt * Addressed Code review comments Signed-off-by: Jason Goldschmidt --------- Signed-off-by: Jason Goldschmidt Co-authored-by: Adit Ranadive --- test/gtest/plugins/memory_handler.h | 143 +++++++++++++++- test/gtest/plugins/meson.build | 17 +- test/gtest/plugins/obj_cuobj_plugin.cpp | 217 ++++++++++++++++++++++++ test/gtest/plugins/obj_plugin.cpp | 104 ++++-------- 4 files changed, 407 insertions(+), 74 deletions(-) create mode 100644 test/gtest/plugins/obj_cuobj_plugin.cpp diff --git a/test/gtest/plugins/memory_handler.h b/test/gtest/plugins/memory_handler.h index 632bf50c..4c4799ca 100644 --- a/test/gtest/plugins/memory_handler.h +++ b/test/gtest/plugins/memory_handler.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,15 +14,20 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#ifndef __MEMORY_HANDLER_H -#define __MEMORY_HANDLER_H +#ifndef NIXL_TEST_GTEST_PLUGINS_MEMORY_HANDLER_H +#define NIXL_TEST_GTEST_PLUGINS_MEMORY_HANDLER_H #include +#include #include "backend/backend_aux.h" #include "backend_engine.h" #include "common/nixl_log.h" #include "nixl.h" +#ifdef HAVE_CUDA +#include +#endif + namespace gtest::plugins { template class memoryHandler; @@ -127,5 +132,135 @@ template<> class memoryHandler { nixlBackendMD *md_; }; +#ifdef HAVE_CUDA +/** + * @brief Memory handler specialization for GPU (VRAM) memory segments. + * + * Manages CUDA device memory allocation and transfers. Uses host-side staging + * buffers for data initialization (setIncreasing) and verification (checkIncreasing), + * copying between host and device via cudaMemcpy. + */ +template<> class memoryHandler { +public: + memoryHandler(size_t len, int dev_id) : buf_(nullptr), len_(len), devId_(dev_id), md_(nullptr) { + cudaError_t err = cudaSetDevice(dev_id); + if (err != cudaSuccess) { + NIXL_ERROR << "cudaSetDevice(" << dev_id << ") failed: " << cudaGetErrorString(err); + throw std::runtime_error("cudaSetDevice failed"); + } + err = cudaMalloc(&buf_, len_); + if (err != cudaSuccess) { + buf_ = nullptr; + NIXL_ERROR << "cudaMalloc(" << len_ << ") failed: " << cudaGetErrorString(err); + throw std::runtime_error("cudaMalloc failed"); + } + } + + ~memoryHandler() { + if (buf_) { + cudaError_t err = cudaSetDevice(devId_); + if (err != cudaSuccess) { + NIXL_ERROR << "cudaSetDevice(" << devId_ + << ") failed in destructor: " << cudaGetErrorString(err); + } + err = cudaFree(buf_); + if (err != cudaSuccess) { + NIXL_ERROR << "cudaFree failed: " << cudaGetErrorString(err); + } + } + } + + memoryHandler(const memoryHandler &) = delete; + memoryHandler & + operator=(const memoryHandler &) = delete; + + void + setIncreasing(uint8_t start_byte) { + std::vector host(len_); + for (auto &entry : host) + entry = start_byte++; + cudaError_t err = cudaSetDevice(devId_); + if (err != cudaSuccess) { + NIXL_ERROR << "cudaSetDevice(" << devId_ << ") failed: " << cudaGetErrorString(err); + throw std::runtime_error("cudaSetDevice failed"); + } + err = cudaMemcpy(buf_, host.data(), len_, cudaMemcpyHostToDevice); + if (err != cudaSuccess) { + NIXL_ERROR << "cudaMemcpy H2D failed: " << cudaGetErrorString(err); + throw std::runtime_error("cudaMemcpy failed"); + } + } + + bool + checkIncreasing(uint8_t start_byte) { + std::vector host(len_); + cudaError_t err = cudaSetDevice(devId_); + if (err != cudaSuccess) { + NIXL_ERROR << "cudaSetDevice(" << devId_ << ") failed: " << cudaGetErrorString(err); + return false; + } + err = cudaMemcpy(host.data(), buf_, len_, cudaMemcpyDeviceToHost); + if (err != cudaSuccess) { + NIXL_ERROR << "cudaMemcpy D2H failed: " << cudaGetErrorString(err); + return false; + } + for (auto &entry : host) { + uint8_t expected_byte = start_byte++; + if (entry != expected_byte) { + NIXL_ERROR << "Verification failed! local: " << entry + << ", expected: " << expected_byte; + return false; + } + } + return true; + } + + void + reset() { + cudaError_t err = cudaSetDevice(devId_); + if (err != cudaSuccess) { + NIXL_ERROR << "cudaSetDevice(" << devId_ << ") failed: " << cudaGetErrorString(err); + throw std::runtime_error("cudaSetDevice failed"); + } + err = cudaMemset(buf_, 0x00, len_); + if (err != cudaSuccess) { + NIXL_ERROR << "cudaMemset failed: " << cudaGetErrorString(err); + throw std::runtime_error("cudaMemset failed"); + } + } + + void + populateBlobDesc(nixlBlobDesc *desc, int buf_index = 0) { + desc->addr = reinterpret_cast(buf_); + desc->len = len_; + desc->devId = devId_; + } + + void + populateMetaDesc(nixlMetaDesc *desc, int entry_index, size_t entry_size) { + desc->addr = reinterpret_cast(buf_) + entry_index * entry_size; + desc->len = entry_size; + desc->devId = devId_; + desc->metadataP = md_; + } + + void + setMD(nixlBackendMD *md) { + md_ = md; + } + + nixlBackendMD * + getMD() { + return md_; + } + +private: + void *buf_; + size_t len_; + int devId_; + nixlBackendMD *md_; +}; +#endif // HAVE_CUDA + } // namespace gtest::plugins -#endif // __MEMORY_HANDLER_H +#endif // NIXL_TEST_GTEST_PLUGINS_MEMORY_HANDLER_H diff --git a/test/gtest/plugins/meson.build b/test/gtest/plugins/meson.build index 5ef6476b..cb4573e2 100644 --- a/test/gtest/plugins/meson.build +++ b/test/gtest/plugins/meson.build @@ -37,11 +37,24 @@ if not aws_s3_crt.found() subdir_done() endif +plugins_gtest_cpp_flags = [] +plugins_gtest_cuda_deps = [] +if cuda_dep.found() + plugins_gtest_cpp_flags += '-DHAVE_CUDA' + plugins_gtest_cuda_deps = [cuda_dep] +endif + +plugins_gtest_sources = ['../main.cpp', '../common.cpp', 'obj_plugin.cpp'] +if cuobj_dep.found() + plugins_gtest_sources += 'obj_cuobj_plugin.cpp' +endif + plugins_test_exe = executable('plugins_gtest', - sources : ['../main.cpp', '../common.cpp', 'obj_plugin.cpp'], + sources : plugins_gtest_sources, include_directories: [nixl_inc_dirs, utils_inc_dirs, plugins_inc_dirs, '.'], + cpp_args : plugins_gtest_cpp_flags, dependencies : [nixl_dep, gtest_dep, absl_strings_dep, absl_time_dep, plugin_deps, - obj_backend_interface], + obj_backend_interface] + plugins_gtest_cuda_deps, link_with: [nixl_build_lib], install : true ) diff --git a/test/gtest/plugins/obj_cuobj_plugin.cpp b/test/gtest/plugins/obj_cuobj_plugin.cpp new file mode 100644 index 00000000..ddd378f5 --- /dev/null +++ b/test/gtest/plugins/obj_cuobj_plugin.cpp @@ -0,0 +1,217 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 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 defined HAVE_CUOBJ_CLIENT + +#include +#include + +#include "plugins_common.h" +#include "transfer_handler.h" +#include "obj/obj_backend.h" + +#ifdef HAVE_CUDA +#include +#endif + +namespace gtest::plugins::obj { + +nixl_b_params_t obj_accel_params = {{"accelerated", "true"}}; +nixl_b_params_t obj_dell_params = {{"accelerated", "true"}, + {"type", "dell"}, + {"req_checksum", "required"}, + {"scheme", "http"}}; +const std::string accel_agent_name = "Agent3-Accel"; +const std::string dell_agent_name = "Agent4-Dell"; + +const nixlBackendInitParams obj_accel_test_params = {.localAgent = accel_agent_name, + .type = "OBJ", + .customParams = &obj_accel_params, + .enableProgTh = false, + .pthrDelay = 0, + .syncMode = + nixl_thread_sync_t::NIXL_THREAD_SYNC_RW}; + +const nixlBackendInitParams obj_dell_test_params = {.localAgent = dell_agent_name, + .type = "OBJ", + .customParams = &obj_dell_params, + .enableProgTh = false, + .pthrDelay = 0, + .syncMode = + nixl_thread_sync_t::NIXL_THREAD_SYNC_RW}; + +// Separate test suite for S3 Accelerated client with accelerated=true +// Note: These tests require the cuobjclient library to be available at compile time +class setupObjAccelTestFixture : public setupBackendTestFixture { +protected: + nixl_b_params_t localParams_; + + setupObjAccelTestFixture() { + localParams_ = *GetParam().customParams; + const char *endpoint = std::getenv("NIXL_OBJ_ENDPOINT_OVERRIDE"); + if (endpoint && endpoint[0] != '\0') { + localParams_["endpoint_override"] = endpoint; + localParams_["req_checksum"] = "required"; + } + nixlBackendInitParams initParams = GetParam(); + initParams.customParams = &localParams_; + localBackendEngine_ = std::make_shared(&initParams); + } +}; + +TEST_P(setupObjAccelTestFixture, AccelXferTest) { + transferHandler transfer( + localBackendEngine_, localBackendEngine_, accel_agent_name, accel_agent_name, false, 1); + transfer.setLocalMem(); + transfer.testTransfer(NIXL_WRITE); + transfer.resetLocalMem(); + transfer.testTransfer(NIXL_READ); + transfer.checkLocalMem(); +} + +TEST_P(setupObjAccelTestFixture, AccelXferMultiBufsTest) { + transferHandler transfer( + localBackendEngine_, localBackendEngine_, accel_agent_name, accel_agent_name, false, 3); + transfer.setLocalMem(); + transfer.testTransfer(NIXL_WRITE); + transfer.resetLocalMem(); + transfer.testTransfer(NIXL_READ); + transfer.checkLocalMem(); +} + +TEST_P(setupObjAccelTestFixture, AccelQueryMemTest) { + transferHandler transfer( + localBackendEngine_, localBackendEngine_, accel_agent_name, accel_agent_name, false, 3); + transfer.setLocalMem(); + transfer.testTransfer(NIXL_WRITE); + + nixl_reg_dlist_t descs(OBJ_SEG); + descs.addDesc(nixlBlobDesc(nixlBasicDesc(), "test-obj-key-0")); + descs.addDesc(nixlBlobDesc(nixlBasicDesc(), "test-obj-key-1")); + descs.addDesc(nixlBlobDesc(nixlBasicDesc(), "test-obj-key-nonexistent")); + std::vector resp; + localBackendEngine_->queryMem(descs, resp); + + EXPECT_EQ(resp.size(), 3); + EXPECT_EQ(resp[0].has_value(), true); + EXPECT_EQ(resp[1].has_value(), true); + EXPECT_EQ(resp[2].has_value(), false); +} + +INSTANTIATE_TEST_SUITE_P(ObjAccelTests, + setupObjAccelTestFixture, + testing::Values(obj_accel_test_params)); + +/** + * @brief Test fixture for Dell ObjectScale S3 over RDMA engine. + * + * Reads the NIXL_OBJ_ENDPOINT_OVERRIDE environment variable to configure + * the Dell ObjectScale endpoint. Tests are skipped if the variable is not set. + * Requires cuobjclient library (HAVE_CUOBJ_CLIENT). + */ +class setupObjDellTestFixture : public setupBackendTestFixture { +protected: + nixl_b_params_t localParams_; + + setupObjDellTestFixture() { + localParams_ = *GetParam().customParams; + const char *endpoint = std::getenv("NIXL_OBJ_ENDPOINT_OVERRIDE"); + if (endpoint && endpoint[0] != '\0') { + localParams_["endpoint_override"] = endpoint; + nixlBackendInitParams initParams = GetParam(); + initParams.customParams = &localParams_; + localBackendEngine_ = std::make_shared(&initParams); + } + } + + void + SetUp() override { + const char *endpoint = std::getenv("NIXL_OBJ_ENDPOINT_OVERRIDE"); + if (!endpoint || endpoint[0] == '\0') { + GTEST_SKIP() << "NIXL_OBJ_ENDPOINT_OVERRIDE not set, skipping Dell tests"; + } + setupBackendTestFixture::SetUp(); + } +}; + +TEST_P(setupObjDellTestFixture, DellXferTest) { + transferHandler transfer( + localBackendEngine_, localBackendEngine_, dell_agent_name, dell_agent_name, false, 1); + transfer.setLocalMem(); + transfer.testTransfer(NIXL_WRITE); + transfer.resetLocalMem(); + transfer.testTransfer(NIXL_READ); + transfer.checkLocalMem(); +} + +TEST_P(setupObjDellTestFixture, DellXferMultiBufsTest) { + transferHandler transfer( + localBackendEngine_, localBackendEngine_, dell_agent_name, dell_agent_name, false, 3); + transfer.setLocalMem(); + transfer.testTransfer(NIXL_WRITE); + transfer.resetLocalMem(); + transfer.testTransfer(NIXL_READ); + transfer.checkLocalMem(); +} + +TEST_P(setupObjDellTestFixture, DellQueryMemTest) { + transferHandler transfer( + localBackendEngine_, localBackendEngine_, dell_agent_name, dell_agent_name, false, 3); + transfer.setLocalMem(); + transfer.testTransfer(NIXL_WRITE); + + nixl_reg_dlist_t descs(OBJ_SEG); + descs.addDesc(nixlBlobDesc(nixlBasicDesc(), "test-obj-key-0")); + descs.addDesc(nixlBlobDesc(nixlBasicDesc(), "test-obj-key-1")); + descs.addDesc(nixlBlobDesc(nixlBasicDesc(), "test-obj-key-nonexistent")); + std::vector resp; + localBackendEngine_->queryMem(descs, resp); + + EXPECT_EQ(resp.size(), 3); + EXPECT_EQ(resp[0].has_value(), true); + EXPECT_EQ(resp[1].has_value(), true); + EXPECT_EQ(resp[2].has_value(), false); +} + +#ifdef HAVE_CUDA +// GPU memory (VRAM_SEG) transfer test for Dell ObjectScale RDMA engine. +// Exercises the VRAM-specific code paths: cuMemObjGetDescriptor/PutDescriptor for +// RDMA descriptor registration, and putObjectRdmaAsync/getObjectRdmaAsync for +// GPU-direct RDMA transfers. +TEST_P(setupObjDellTestFixture, DellVramXferTest) { + int device_count = 0; + cudaError_t err = cudaGetDeviceCount(&device_count); + if (err != cudaSuccess || device_count == 0) { + GTEST_SKIP() << "No CUDA devices available, skipping VRAM test"; + } + transferHandler transfer( + localBackendEngine_, localBackendEngine_, dell_agent_name, dell_agent_name, false, 1); + transfer.setLocalMem(); + transfer.testTransfer(NIXL_WRITE); + transfer.resetLocalMem(); + transfer.testTransfer(NIXL_READ); + transfer.checkLocalMem(); +} +#endif // HAVE_CUDA + +INSTANTIATE_TEST_SUITE_P(ObjDellTests, + setupObjDellTestFixture, + testing::Values(obj_dell_test_params)); + +} // namespace gtest::plugins::obj + +#endif // HAVE_CUOBJ_CLIENT diff --git a/test/gtest/plugins/obj_plugin.cpp b/test/gtest/plugins/obj_plugin.cpp index de5590f5..f60fca5e 100644 --- a/test/gtest/plugins/obj_plugin.cpp +++ b/test/gtest/plugins/obj_plugin.cpp @@ -16,6 +16,7 @@ */ #include +#include #include "plugins_common.h" #include "transfer_handler.h" @@ -38,15 +39,25 @@ namespace gtest::plugins::obj { * crtMinLimit is set to the S3 minimum part size (5 MiB) so that * partSize is not clamped and MPU is exercised with multiple parts. * - ObjAccelTests: S3 Accelerated client tests (accelerated = true) - * Note: Only compiled if HAVE_CUOBJ_CLIENT is defined + * Note: Defined in obj_cuobj_plugin.cpp, only compiled if + * HAVE_CUOBJ_CLIENT is defined + * - ObjDellTests: Dell ObjectScale S3 over RDMA tests + * (accelerated = true, type = dell, req_checksum = required, scheme = http) + * Note: Defined in obj_cuobj_plugin.cpp, only compiled if + * HAVE_CUOBJ_CLIENT is defined. + * Skipped at runtime if NIXL_OBJ_ENDPOINT_OVERRIDE is not set. + * Includes DellVramXferTest (GPU memory) if HAVE_CUDA is also defined. + * + * Environment: + * - NIXL_OBJ_ENDPOINT_OVERRIDE (e.g. http://100.68.213.151:9020) + * Optional for all test suites. When set, overrides the S3 endpoint for + * Standard, CRT, Accel, and Dell tests alike. */ nixl_b_params_t obj_params = {{"crtMinLimit", "0"}}; nixl_b_params_t obj_crt_params = {{"crtMinLimit", "5242880"}}; // 5 MiB: S3 minimum part size -nixl_b_params_t obj_accel_params = {{"accelerated", "true"}}; const std::string local_agent_name = "Agent1"; const std::string crt_agent_name = "Agent2-CRT"; -const std::string accel_agent_name = "Agent3-Accel"; const nixlBackendInitParams obj_test_params = {.localAgent = local_agent_name, .type = "OBJ", .customParams = &obj_params, @@ -62,18 +73,20 @@ const nixlBackendInitParams obj_crt_test_params = {.localAgent = crt_agent_name, .syncMode = nixl_thread_sync_t::NIXL_THREAD_SYNC_RW}; -const nixlBackendInitParams obj_accel_test_params = {.localAgent = accel_agent_name, - .type = "OBJ", - .customParams = &obj_accel_params, - .enableProgTh = false, - .pthrDelay = 0, - .syncMode = - nixl_thread_sync_t::NIXL_THREAD_SYNC_RW}; - class setupObjTestFixture : public setupBackendTestFixture { protected: + nixl_b_params_t localParams_; + setupObjTestFixture() { - localBackendEngine_ = std::make_shared(&GetParam()); + localParams_ = *GetParam().customParams; + const char *endpoint = std::getenv("NIXL_OBJ_ENDPOINT_OVERRIDE"); + if (endpoint && endpoint[0] != '\0') { + localParams_["endpoint_override"] = endpoint; + localParams_["req_checksum"] = "required"; + } + nixlBackendInitParams initParams = GetParam(); + initParams.customParams = &localParams_; + localBackendEngine_ = std::make_shared(&initParams); } }; @@ -116,14 +129,23 @@ TEST_P(setupObjTestFixture, queryMemTest) { EXPECT_EQ(resp[2].has_value(), false); } - INSTANTIATE_TEST_SUITE_P(ObjTests, setupObjTestFixture, testing::Values(obj_test_params)); // Separate test suite for S3 CRT client with crtMinLimit enabled class setupObjCrtTestFixture : public setupBackendTestFixture { protected: + nixl_b_params_t localParams_; + setupObjCrtTestFixture() { - localBackendEngine_ = std::make_shared(&GetParam()); + localParams_ = *GetParam().customParams; + const char *endpoint = std::getenv("NIXL_OBJ_ENDPOINT_OVERRIDE"); + if (endpoint && endpoint[0] != '\0') { + localParams_["endpoint_override"] = endpoint; + localParams_["req_checksum"] = "required"; + } + nixlBackendInitParams initParams = GetParam(); + initParams.customParams = &localParams_; + localBackendEngine_ = std::make_shared(&initParams); } }; @@ -186,58 +208,4 @@ TEST_P(setupObjCrtTestFixture, CrtQueryMemTest) { INSTANTIATE_TEST_SUITE_P(ObjCrtTests, setupObjCrtTestFixture, testing::Values(obj_crt_test_params)); -#if defined HAVE_CUOBJ_CLIENT -// Separate test suite for S3 Accelerated client with accelerated=true -// Note: These tests require the cuobjclient library to be available at compile time -class setupObjAccelTestFixture : public setupBackendTestFixture { -protected: - setupObjAccelTestFixture() { - localBackendEngine_ = std::make_shared(&GetParam()); - } -}; - -TEST_P(setupObjAccelTestFixture, AccelXferTest) { - transferHandler transfer( - localBackendEngine_, localBackendEngine_, accel_agent_name, accel_agent_name, false, 1); - transfer.setLocalMem(); - transfer.testTransfer(NIXL_WRITE); - transfer.resetLocalMem(); - transfer.testTransfer(NIXL_READ); - transfer.checkLocalMem(); -} - -TEST_P(setupObjAccelTestFixture, AccelXferMultiBufsTest) { - transferHandler transfer( - localBackendEngine_, localBackendEngine_, accel_agent_name, accel_agent_name, false, 3); - transfer.setLocalMem(); - transfer.testTransfer(NIXL_WRITE); - transfer.resetLocalMem(); - transfer.testTransfer(NIXL_READ); - transfer.checkLocalMem(); -} - -TEST_P(setupObjAccelTestFixture, AccelQueryMemTest) { - transferHandler transfer( - localBackendEngine_, localBackendEngine_, accel_agent_name, accel_agent_name, false, 3); - transfer.setLocalMem(); - transfer.testTransfer(NIXL_WRITE); - - nixl_reg_dlist_t descs(OBJ_SEG); - descs.addDesc(nixlBlobDesc(nixlBasicDesc(), "test-obj-key-0")); - descs.addDesc(nixlBlobDesc(nixlBasicDesc(), "test-obj-key-1")); - descs.addDesc(nixlBlobDesc(nixlBasicDesc(), "test-obj-key-nonexistent")); - std::vector resp; - localBackendEngine_->queryMem(descs, resp); - - EXPECT_EQ(resp.size(), 3); - EXPECT_EQ(resp[0].has_value(), true); - EXPECT_EQ(resp[1].has_value(), true); - EXPECT_EQ(resp[2].has_value(), false); -} - -INSTANTIATE_TEST_SUITE_P(ObjAccelTests, - setupObjAccelTestFixture, - testing::Values(obj_accel_test_params)); -#endif // HAVE_CUOBJ_CLIENT - } // namespace gtest::plugins::obj From db55b35f4f091631ff53d71ccb8859e966b86d0e Mon Sep 17 00:00:00 2001 From: Daniel Pressler Date: Sun, 5 Apr 2026 13:26:43 +0300 Subject: [PATCH 07/59] CI: Run Artifactory cleanup at end of container build pipeline (#1384) Invoke artifactory-cleanup-by-time.sh in pipeline_stop (dry-run) to list Docker images that would be removed by age. Use REGISTRY_REPO_NAME, REGISTRY_REPO_PATH, ARTIFACTORY_CLEANUP_AGE (default 3m) and existing Artifactory credentials. Also: split REGISTRY_REPO into REGISTRY_REPO_NAME and REGISTRY_REPO_PATH, set ARTIFACTORY_REGISTRY_URL for push/echo, add containerSelector for podman on relevant steps, and simplify taskName (drop axis_index). Signed-off-by: Daniel Pressler --- .ci/jenkins/lib/build-container-matrix.yaml | 30 +++- .ci/scripts/artifactory-cleanup-by-time.sh | 175 ++++++++++++++++++++ 2 files changed, 196 insertions(+), 9 deletions(-) create mode 100755 .ci/scripts/artifactory-cleanup-by-time.sh diff --git a/.ci/jenkins/lib/build-container-matrix.yaml b/.ci/jenkins/lib/build-container-matrix.yaml index 09b0019f..42764d5e 100644 --- a/.ci/jenkins/lib/build-container-matrix.yaml +++ b/.ci/jenkins/lib/build-container-matrix.yaml @@ -16,7 +16,7 @@ kubernetes: requests: "{memory: 8Gi, cpu: 4000m}" runs_on_dockers: - - { name: "podman-v5.7.1", url: "quay.io/podman/stable:v5.7.1", privileged: true } + - { name: "podman", url: "quay.io/podman/stable:v5.7.1", privileged: true } # Build matrix matrix: @@ -27,12 +27,15 @@ matrix: # Configuration env: - REGISTRY_HOSTESS: "artifactory.nvidia.com" - REGISTRY_REPO: "sw-nbu-swx-nixl-docker-local/verification" + REGISTRY_HOST_NAME: "artifactory.nvidia.com" + REGISTRY_REPO_NAME: "sw-nbu-swx-nixl-docker-local" + REGISTRY_REPO_PATH: "verification" + ARTIFACTORY_REGISTRY_URL: "${REGISTRY_HOST_NAME}/${REGISTRY_REPO_NAME}/${REGISTRY_REPO_PATH}" + ARTIFACTORY_CLEANUP_AGE: "3m" # 3 months LOCAL_TAG_BASE: "nixl-ci:build-" MAIL_FROM: "jenkins@nvidia.com" -taskName: "${BUILD_TARGET}/${arch}/${axis_index}" +taskName: "${BUILD_TARGET}/${arch}" credentials: - credentialsId: 'svc-nixl-new-artifactory-token' @@ -56,6 +59,7 @@ pipeline_start: # Build pipeline steps: - name: Prepare + containerSelector: "{name: 'podman'}" run: | # Setup podman and dependencies set -x @@ -67,6 +71,7 @@ steps: - name: Build NIXLBench enable: ${ENABLE_NIXLBENCH_BUILD} + containerSelector: "{name: 'podman'}" run: | # Clone UCX source for nixlbench git clone https://github.com/openucx/ucx.git ucx-src @@ -83,6 +88,7 @@ steps: - name: Build NIXL enable: ${ENABLE_NIXL_BUILD} + containerSelector: "{name: 'podman'}" run: | export UCX_REF="${UCX_VERSION}" @@ -101,6 +107,7 @@ steps: ${NIXL_EP_FLAG} - name: Add Version Info + containerSelector: "{name: 'podman'}" run: | git config --global --add safe.directory '*' # Extract standardized 8-char commit hash for UCX version info: @@ -138,11 +145,12 @@ steps: docker rm -f tempcontainer || true - name: Push + containerSelector: "{name: 'podman'}" credentialsId: 'svc-nixl-new-artifactory-token' run: | source version-info - ARTIFACTORY_REGISTRY="${REGISTRY_HOSTESS}/${REGISTRY_REPO}/${BUILD_TARGET}" - ARTIFACTORY_API="https://${REGISTRY_HOSTESS}/artifactory/api/storage/${REGISTRY_REPO}/${BUILD_TARGET}" + ARTIFACTORY_REGISTRY="${ARTIFACTORY_REGISTRY_URL}/${BUILD_TARGET}" + ARTIFACTORY_API="https://${REGISTRY_HOST_NAME}/artifactory/api/storage/${REGISTRY_REPO_NAME}/${REGISTRY_REPO_PATH}/${BUILD_TARGET}" # Prepare image properties IMAGE_PROPERTIES="BUILD_TARGET=${BUILD_TARGET};NIXL_VERSION=${NIXL_VERSION};UCX_VERSION=${UCX_VERSION};arch=${arch};" @@ -150,7 +158,7 @@ steps: IMAGE_PROPERTIES+="BASE_IMAGE=${BASE_IMAGE};BASE_IMAGE_TAG=${BASE_IMAGE_TAG}" # Login to Artifactory - echo "$ARTIFACTORY_PASSWORD" | docker login "${REGISTRY_HOSTESS}" -u "$ARTIFACTORY_USERNAME" --password-stdin + echo "$ARTIFACTORY_PASSWORD" | docker login "${REGISTRY_HOST_NAME}" -u "$ARTIFACTORY_USERNAME" --password-stdin # Function to tag, push, and set properties tag_push_set_properties() { @@ -174,9 +182,9 @@ steps: run: | source version-info echo "Image type built: ${BUILD_TARGET} (${arch})" - echo "Image pushed to: ${REGISTRY_HOSTESS}/${REGISTRY_REPO}/${BUILD_TARGET}:${TAG_NAME}" + echo "Image pushed to: ${ARTIFACTORY_REGISTRY_URL}/${BUILD_TARGET}:${TAG_NAME}" if [[ "${UPDATE_LATEST}" == "true" ]]; then - echo "Latest tag updated: ${REGISTRY_HOSTESS}/${REGISTRY_REPO}/${BUILD_TARGET}:${BASE_IMAGE_TAG}-${arch}-latest" + echo "Latest tag updated: ${ARTIFACTORY_REGISTRY_URL}/${BUILD_TARGET}:${BASE_IMAGE_TAG}-${arch}-latest" fi echo -e "\nBuild config for manual repro:" @@ -195,6 +203,7 @@ steps: pipeline_stop: shell: action module: groovy + containerSelector: "{name: 'podman'}" run: | if (params.MAIL_TO) { def jobStatus = currentBuild.result ?: 'SUCCESS' @@ -223,3 +232,6 @@ pipeline_stop: """ ) } + withCredentials([usernamePassword(credentialsId: 'svc-nixl-new-artifactory-token', usernameVariable: 'ARTIFACTORY_USER', passwordVariable: 'ARTIFACTORY_TOKEN')]) { + sh "yum install -y jq > /dev/null && .ci/scripts/artifactory-cleanup-by-time.sh --path ${REGISTRY_REPO_PATH}/${BUILD_TARGET} ${REGISTRY_REPO_NAME} ${ARTIFACTORY_CLEANUP_AGE}" + } diff --git a/.ci/scripts/artifactory-cleanup-by-time.sh b/.ci/scripts/artifactory-cleanup-by-time.sh new file mode 100755 index 00000000..ede1524d --- /dev/null +++ b/.ci/scripts/artifactory-cleanup-by-time.sh @@ -0,0 +1,175 @@ +#!/bin/bash +# Artifactory cleanup by time (Docker images) +# Deletes Docker images in a JFrog Artifactory repository that are older than a given age. +# Only manifest.json (and list.manifest.json) are queried; image age is based on the +# manifest's last modified date. The entire image folder is then deleted. +# +# Usage: +# export ARTIFACTORY_URL="https://your-instance.jfrog.io/artifactory" +# export ARTIFACTORY_USER="user" +# export ARTIFACTORY_TOKEN="password-or-api-key" +# ./artifactory-cleanup-by-time.sh [OPTIONS] REPO_KEY [AGE] +# +# Arguments: +# REPO_KEY Repository key (e.g. sw-nbu-swx-nixl-docker-local) +# AGE Age threshold: delete images whose manifest was last modified before this (default: 4w) +# Examples: 4w (weeks), 30d (days), 2m (months) +# +# Options: +# --dry-run Only list images that would be deleted, do not delete +# --path PATH Only consider images under this path (AQL path match) +# +# Environment: +# ARTIFACTORY_URL Base Artifactory URL (required) +# ARTIFACTORY_USER Username for API auth (required unless using token only) +# ARTIFACTORY_TOKEN Password or API key (required) +# ARTIFACTORY_API_KEY Alternative to ARTIFACTORY_TOKEN +# +# Example: +# ./artifactory-cleanup-by-time.sh --dry-run sw-nbu-swx-nixl-docker-local 4w +# ./artifactory-cleanup-by-time.sh --path "verification/nixl/" sw-nbu-swx-nixl-docker-local 4w + +set -e + +DRY_RUN=false +PATH_FILTER="" + +while [[ $# -gt 0 ]]; do + case "$1" in + --dry-run) + DRY_RUN=true + shift + ;; + --path) + if [[ $# -lt 2 || "$2" == -* ]]; then + echo "Missing value for --path" >&2 + exit 1 + fi + PATH_FILTER="$2" + shift 2 + ;; + -*) + echo "Unknown option: $1" >&2 + exit 1 + ;; + *) + break + ;; + esac +done + +if [[ $# -lt 1 ]]; then + echo "Usage: $0 [--dry-run] [--path PATH] REPO_KEY [AGE]" >&2 + echo " AGE default: 4w (4 weeks)" >&2 + exit 1 +fi + +REPO_KEY="$1" +AGE="${2:-4w}" + +# Parse age (e.g. 4w, 30d, 2m) and compute cutoff via epoch (portable) +if [[ "$AGE" =~ ^([0-9]+)([wdm])$ ]]; then + NUM="${BASH_REMATCH[1]}" + UNIT="${BASH_REMATCH[2]}" +else + echo "Invalid AGE format: $AGE (e.g. 4w, 30d, 2m)" >&2 + exit 1 +fi + +NOW_EPOCH=$(date +%s) +case "$UNIT" in + w) SECS_AGO=$((NUM * 7 * 24 * 3600)) ;; + d) SECS_AGO=$((NUM * 24 * 3600)) ;; + m) SECS_AGO=$((NUM * 30 * 24 * 3600)) ;; + *) echo "Invalid age unit: $UNIT (use w, d, or m)" >&2; exit 1 ;; +esac +CUTOFF_EPOCH=$((NOW_EPOCH - SECS_AGO)) + +if date -u -d "@${CUTOFF_EPOCH}" +"%Y-%m-%dT%H:%M:%S.000Z" &>/dev/null; then + CUTOFF=$(date -u -d "@${CUTOFF_EPOCH}" +"%Y-%m-%dT%H:%M:%S.000Z") +elif date -u -r "${CUTOFF_EPOCH}" +"%Y-%m-%dT%H:%M:%S.000Z" &>/dev/null; then + CUTOFF=$(date -u -r "${CUTOFF_EPOCH}" +"%Y-%m-%dT%H:%M:%S.000Z") +else + echo "Cannot convert epoch to ISO date" >&2 + exit 1 +fi + +ARTIFACTORY_URL="${ARTIFACTORY_URL:-https://artifactory.nvidia.com/artifactory}" +TOKEN="${ARTIFACTORY_TOKEN:-${ARTIFACTORY_API_KEY:?Set ARTIFACTORY_TOKEN or ARTIFACTORY_API_KEY}}" +USER="${ARTIFACTORY_USER:-}" + +if [[ -n "$USER" ]]; then + CURL_AUTH=(-u "$USER:$TOKEN") +else + CURL_AUTH=(-H "X-JFrog-Art-Api: $TOKEN") +fi + +# Build AQL: only manifest files; filter by last modified date. Use $and so $or is valid. +if [[ -n "$PATH_FILTER" ]]; then + PATH_MATCH="${PATH_FILTER%/}" + AQL='items.find({"$and":[{"repo":"'"$REPO_KEY"'"},{"path":{"$match":"'"$PATH_MATCH"'/*"}},{"$or":[{"name":"manifest.json"},{"name":"list.manifest.json"}]},{"modified":{"$lt":"'"$CUTOFF"'"}}]}).include("repo","path","name")' +else + AQL='items.find({"$and":[{"repo":"'"$REPO_KEY"'"},{"$or":[{"name":"manifest.json"},{"name":"list.manifest.json"}]},{"modified":{"$lt":"'"$CUTOFF"'"}}]}).include("repo","path","name")' +fi + +echo "Current time (UTC): $(date -u +"%Y-%m-%dT%H:%M:%S.000Z")" +echo "Cutoff: images with manifest last modified before $CUTOFF will be deleted (repo: $REPO_KEY)" +if [[ -n "$PATH_FILTER" ]]; then + echo "Path filter: $PATH_FILTER" +fi +if [[ "$DRY_RUN" = true ]]; then echo "DRY RUN: no artifacts will be deleted"; fi + +RESPONSE=$(curl -s -w "\n%{http_code}" "${CURL_AUTH[@]}" -X POST \ + "${ARTIFACTORY_URL}/api/search/aql" \ + -H "Content-Type: text/plain" \ + -d "$AQL") || { echo "AQL request failed"; exit 1; } + +# Split body and HTTP code (last line) +HTTP_CODE=$(echo "$RESPONSE" | tail -n1) +RESPONSE_BODY=$(echo "$RESPONSE" | sed '$d') + +if ! echo "$RESPONSE_BODY" | jq -e . >/dev/null 2>&1; then + echo "Artifactory response is not valid JSON (HTTP $HTTP_CODE). First 500 chars:" + echo "$RESPONSE_BODY" | head -c 500 + echo "" + exit 1 +fi + +# Check for AQL errors in JSON body +if echo "$RESPONSE_BODY" | jq -e '.errors' >/dev/null 2>&1; then + echo "AQL error:" + echo "$RESPONSE_BODY" | jq -r '.errors[]? // .' + exit 1 +fi + +RESULTS_COUNT=$(echo "$RESPONSE_BODY" | jq -r '(.results // []) | length') +echo "AQL returned ${RESULTS_COUNT} manifest(s) matching criteria." + +# Each result is a manifest file; the image folder is repo/path (parent of manifest). Dedupe by folder. +COUNT=0 +FAILED=0 +while IFS= read -r image_path; do + [[ -z "$image_path" ]] && continue + if [[ "$DRY_RUN" = true ]]; then + echo " [dry-run] would delete image: $image_path" + COUNT=$((COUNT + 1)) + else + echo " Deleting image: $image_path" + if curl -sSf "${CURL_AUTH[@]}" -X DELETE "${ARTIFACTORY_URL}/${image_path}"; then + COUNT=$((COUNT + 1)) + else + echo " ERROR: delete failed for $image_path" >&2 + FAILED=$((FAILED + 1)) + fi + fi +done < <(echo "$RESPONSE_BODY" | jq -r '(.results // [])[] | "\(.repo)/\(.path)"' | sort -u) + +if [[ "$DRY_RUN" = true ]]; then + echo "Done. $COUNT image(s) would be deleted." +else + if [[ "$FAILED" -gt 0 ]]; then + echo "Done. $COUNT image(s) deleted, $FAILED failed." >&2 + exit 1 + fi + echo "Done. $COUNT image(s) deleted." +fi From 06491f8eb30c51c4558a611a602faa43fa95312e Mon Sep 17 00:00:00 2001 From: ofarjon <34404710+ofirfarjun7@users.noreply.github.com> Date: Sun, 5 Apr 2026 18:03:29 +0300 Subject: [PATCH 08/59] NIXL EP: Use VMM API for device memory allocation. (#1415) * NIXL/EP: Use vmm API instead of cudaMalloc * NIXL/EP: Use vmm API instead of cudaMalloc * NIXL/EP: revert * NIXL/EP: Support gdr copy with vmm. * NIXL/EP: Improve. * NIXL/EP: Format. * NIXL/EP: check return val. * NIXL/EP: Format. * NIXL/EP:Improve. * NIXL/EP: Format. * NIXL/EP: Improve. * NIXL/EP: Format. * NIXL/EP: fallback to cudaMalloc if fabric not supported * NIXL/EP: set default vals * NIXL/EP: Fix. * NIXL/EP: Fix comments. * NIXL/EP: Fix. * NIXL/EP: Fix. * NIXL/EP: not needed. * NIXL/EP: new files terms. * NIXL/EP: Fix comments. * NIXL/EP: Fix comment. * NIXL/EP: Fix. * NIXL/EP: fix. * NIXL/EP: Fix comment. * NIXL/EP: Fix comment. * NIXL/EP: Fix comment. * NIXL/EP: Fix comment. * NIXL/EP: Fix comment. --------- Co-authored-by: Ofir Farjon --- examples/device/ep/csrc/config.hpp | 6 +- examples/device/ep/csrc/nixl_ep.cpp | 40 ++++--- examples/device/ep/csrc/nixl_ep.hpp | 8 ++ examples/device/ep/csrc/vmm.cpp | 174 ++++++++++++++++++++++++++++ examples/device/ep/csrc/vmm.hpp | 55 +++++++++ examples/device/ep/meson.build | 1 + 6 files changed, 266 insertions(+), 18 deletions(-) create mode 100644 examples/device/ep/csrc/vmm.cpp create mode 100644 examples/device/ep/csrc/vmm.hpp diff --git a/examples/device/ep/csrc/config.hpp b/examples/device/ep/csrc/config.hpp index f17ac113..89b49cc4 100644 --- a/examples/device/ep/csrc/config.hpp +++ b/examples/device/ep/csrc/config.hpp @@ -122,7 +122,11 @@ struct EPLayout { } }; -size_t get_rdma_size_hint(int num_max_dispatch_tokens_per_rank, int hidden, int num_ranks, int num_experts) { +inline 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; } diff --git a/examples/device/ep/csrc/nixl_ep.cpp b/examples/device/ep/csrc/nixl_ep.cpp index ad0b6990..e8f2055d 100644 --- a/examples/device/ep/csrc/nixl_ep.cpp +++ b/examples/device/ep/csrc/nixl_ep.cpp @@ -103,25 +103,28 @@ void Buffer::init(int num_ranks, int num_experts_per_rank, int64_t num_rdma_byte EP_HOST_ASSERT(per_channel_bytes < std::numeric_limits::max()); // Create 32 MiB workspace - CUDA_CHECK(cudaMalloc(&workspace, NUM_WORKSPACE_BYTES)); + m_workspace_alloc = std::make_unique(NUM_WORKSPACE_BYTES); + workspace = m_workspace_alloc->ptr(); CUDA_CHECK(cudaMemsetAsync(workspace, 0, NUM_WORKSPACE_BYTES, comm_stream)); EP_HOST_ASSERT(max_experts_per_rank > 0); - CUDA_CHECK(cudaMalloc(&rdma_buffer_ptr, num_rdma_bytes)); + m_rdma_alloc = std::make_unique(static_cast(num_rdma_bytes)); + rdma_buffer_ptr = m_rdma_alloc->ptr(); CUDA_CHECK(cudaMemset(rdma_buffer_ptr, 0, num_rdma_bytes)); // Allocate and clean shrink buffer - mask_buffer_ptr = nullptr; - sync_buffer_ptr = nullptr; int num_mask_buffer_bytes = max_num_ranks * sizeof(int); - CUDA_CHECK(cudaMalloc(&mask_buffer_ptr, num_mask_buffer_bytes)); + m_mask_alloc = std::make_unique(static_cast(num_mask_buffer_bytes)); + mask_buffer_ptr = static_cast(m_mask_alloc->ptr()); 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)); + m_sync_alloc = std::make_unique(static_cast(num_sync_buffer_bytes)); + sync_buffer_ptr = static_cast(m_sync_alloc->ptr()); + m_sync_count_alloc = std::make_unique(static_cast(num_sync_buffer_bytes)); + sync_count_ptr = static_cast(m_sync_count_alloc->ptr()); CUDA_CHECK(cudaMemset(sync_buffer_ptr, 0, num_sync_buffer_bytes)); - CUDA_CHECK(cudaMalloc(&sync_count_ptr, num_sync_buffer_bytes)); CUDA_CHECK(cudaMemset(sync_count_ptr, 0, num_sync_buffer_bytes)); CUDA_CHECK(cudaDeviceSynchronize()); @@ -167,10 +170,10 @@ torch::Stream Buffer::get_comm_stream() const { } void Buffer::destroy() { - auto warn_cuda = [](cudaError_t status, const char* operation) noexcept { + auto warn_cuda = [](cudaError_t status, const char *operation) noexcept { if (status != cudaSuccess) { - std::cerr << "WARNING: destroy() failed to " << operation - << ": " << cudaGetErrorString(status) << '\n'; + std::cerr << "WARNING: destroy() failed to " << operation << ": " + << cudaGetErrorString(status) << '\n'; } }; @@ -195,7 +198,6 @@ void Buffer::destroy() { warn_nixl(nixl_agent_info->agent->invalidateLocalMD(), "invalidate local metadata"); } - warn_nixl(nixl_agent_info->agent->deregisterMem( nixl_agent_info->rdma_reg_descs, &nixl_agent_info->extra_params), @@ -212,13 +214,17 @@ void Buffer::destroy() { nixl_agent_info.reset(); } - warn_cuda(cudaFree(rdma_buffer_ptr), "free RDMA buffer"); - warn_cuda(cudaFree(mask_buffer_ptr), "free mask buffer"); - warn_cuda(cudaFree(sync_buffer_ptr), "free sync buffer"); - warn_cuda(cudaFree(sync_count_ptr), "free sync-count buffer"); + m_rdma_alloc.reset(); + rdma_buffer_ptr = nullptr; + m_mask_alloc.reset(); + mask_buffer_ptr = nullptr; + m_sync_alloc.reset(); + sync_buffer_ptr = nullptr; + m_sync_count_alloc.reset(); + sync_count_ptr = nullptr; - // Free workspace - warn_cuda(cudaFree(workspace), "free workspace"); + m_workspace_alloc.reset(); + workspace = nullptr; destroyed = true; available = false; diff --git a/examples/device/ep/csrc/nixl_ep.hpp b/examples/device/ep/csrc/nixl_ep.hpp index 8d0117a1..1053ab9d 100644 --- a/examples/device/ep/csrc/nixl_ep.hpp +++ b/examples/device/ep/csrc/nixl_ep.hpp @@ -40,6 +40,7 @@ #include "event.hpp" #include "kernels/configs.cuh" #include "kernels/exception.cuh" +#include "vmm.hpp" #include "nixl.h" @@ -86,6 +87,13 @@ struct Buffer { int *sync_buffer_ptr = nullptr; int *sync_count_ptr = nullptr; + /* Owning VMM allocations (keep raw ptrs above as aliases) */ + std::unique_ptr m_rdma_alloc; + std::unique_ptr m_mask_alloc; + std::unique_ptr m_sync_alloc; + std::unique_ptr m_sync_count_alloc; + std::unique_ptr m_workspace_alloc; + // Device info and communication int device_id; int num_device_sms; diff --git a/examples/device/ep/csrc/vmm.cpp b/examples/device/ep/csrc/vmm.cpp new file mode 100644 index 00000000..93c7f28b --- /dev/null +++ b/examples/device/ep/csrc/vmm.cpp @@ -0,0 +1,174 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 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. + */ + +#include +#include + +#include "config.hpp" +#include "vmm.hpp" + +namespace { + +constexpr const char *k_vmm_ctx = "vmm_region"; + +/** Log a non-fatal warning if a CUDA driver API call failed (e.g. during teardown). */ +void +warn_cu_api(CUresult status, const char *context, const char *operation) noexcept { + if (status != CUDA_SUCCESS) { + const char *msg = nullptr; + if (cuGetErrorString(status, &msg) != CUDA_SUCCESS || msg == nullptr) { + msg = "unknown CUDA driver error"; + } + std::cerr << "WARNING: " << context << " failed to " << operation << ": " << msg << '\n'; + } +} + +} // namespace + +namespace nixl_ep { + +void +vmm_region::release() noexcept { + if (is_cuda_malloc_) { + if (ptr_) { + warn_cu_api(cuMemFree(ptr_), k_vmm_ctx, "cuMemFree"); + } + ptr_ = 0; + return; + } + + if (vmm_mapped_) { + warn_cu_api(cuMemUnmap(ptr_, size_), k_vmm_ctx, "cuMemUnmap"); + vmm_mapped_ = false; + } + if (ptr_) { + warn_cu_api(cuMemAddressFree(ptr_, size_), k_vmm_ctx, "cuMemAddressFree"); + ptr_ = 0; + } + if (handle_) { + warn_cu_api(cuMemRelease(handle_), k_vmm_ctx, "cuMemRelease"); + handle_ = 0; + } +} + +vmm_region::~vmm_region() { + release(); +} + +vmm_region::vmm_region(size_t size) { + if (size == 0) { + throw std::invalid_argument("vmm_region: size must be non-zero"); + } + + struct cuda_alloc_ctx { + bool fabric_supported; + CUmemAllocationProp prop; + size_t granularity; + CUdevice device; + CUmemAccessDesc access_desc = {}; + + cuda_alloc_ctx() : fabric_supported(false), prop({}), granularity(0) { + int version; + + if (cuCtxGetDevice(&device) != CUDA_SUCCESS) { + throw std::runtime_error("CUDA device should be set before creating a vmm_region"); + } + + if (cuDriverGetVersion(&version) != CUDA_SUCCESS) { + throw std::runtime_error("Failed to get CUDA driver version"); + } + + if (version < 11000) { + return; /* too old — fall back to cudaMalloc */ + } + + int fab = 0; + if ((cuDeviceGetAttribute(&fab, + CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, + device) != CUDA_SUCCESS) || + (!fab)) { + return; /* no fabric — fall back to cudaMalloc */ + } + + int rdma_vmm_supported = 0; + if (cuDeviceGetAttribute(&rdma_vmm_supported, + CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED, + device) != CUDA_SUCCESS) { + throw std::runtime_error( + "Failed to query GPUDirect RDMA with VMM support attribute"); + } + + if (!rdma_vmm_supported) { + std::cerr << "DIAG: " << k_vmm_ctx + << " - GPUDirect RDMA with CUDA VMM not supported; falling back to " + "cuMemAlloc\n"; + return; + } + + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.location.id = device; + prop.allocFlags.gpuDirectRDMACapable = 1; + prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; + + if (cuMemGetAllocationGranularity( + &granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM) != CUDA_SUCCESS) { + throw std::runtime_error("Failed to get CUDA allocation granularity"); + } + + access_desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + access_desc.location.id = device; + access_desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + fabric_supported = true; + } + }; + + static cuda_alloc_ctx ctx; + + if (!ctx.fabric_supported) { + size_ = size; + is_cuda_malloc_ = true; + if (cuMemAlloc(&ptr_, size) != CUDA_SUCCESS) { + throw std::runtime_error("cuMemAlloc fallback failed"); + } + return; + } + + size_ = nixl_ep::align_up(size, ctx.granularity); + + if (cuMemCreate(&handle_, size_, &ctx.prop, 0) != CUDA_SUCCESS) { + throw std::runtime_error("Failed to create CUDA VMM allocation"); + } + + if (cuMemAddressReserve(&ptr_, size_, 0, 0, 0) != CUDA_SUCCESS) { + release(); + throw std::runtime_error("Failed to reserve CUDA virtual address"); + } + + if (cuMemMap(ptr_, size_, 0, handle_, 0) != CUDA_SUCCESS) { + release(); + throw std::runtime_error("Failed to map CUDA VMM memory"); + } + vmm_mapped_ = true; + + if (cuMemSetAccess(ptr_, size_, &ctx.access_desc, 1) != CUDA_SUCCESS) { + release(); + throw std::runtime_error("Failed to set CUDA memory access"); + } +} + +} // namespace nixl_ep diff --git a/examples/device/ep/csrc/vmm.hpp b/examples/device/ep/csrc/vmm.hpp new file mode 100644 index 00000000..9b007b12 --- /dev/null +++ b/examples/device/ep/csrc/vmm.hpp @@ -0,0 +1,55 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 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. + */ + +#pragma once + +#include +#include +#include + +namespace nixl_ep { + +class vmm_region { +public: + explicit vmm_region(size_t size); + + ~vmm_region(); + + vmm_region(const vmm_region &) = delete; + vmm_region & + operator=(const vmm_region &) = delete; + vmm_region(vmm_region &&) = delete; + vmm_region & + operator=(vmm_region &&) = delete; + + [[nodiscard]] void * + ptr() const noexcept { + return reinterpret_cast(static_cast(ptr_)); + } + +private: + void + release() noexcept; + + CUdeviceptr ptr_ = 0; + size_t size_ = 0; + CUmemGenericAllocationHandle handle_ = 0; + bool is_cuda_malloc_ = false; + bool vmm_mapped_ = false; +}; + +} // namespace nixl_ep diff --git a/examples/device/ep/meson.build b/examples/device/ep/meson.build index 123a90ba..52956f72 100644 --- a/examples/device/ep/meson.build +++ b/examples/device/ep/meson.build @@ -63,6 +63,7 @@ endif nixl_ep_sources = [ 'csrc/nixl_ep.cpp', + 'csrc/vmm.cpp', 'csrc/kernels/nixl_ep.cu', ] From 86b8c72093c33cc96af4b5c2268d866dfd3aaf18 Mon Sep 17 00:00:00 2001 From: NitayAmiel <95589670+NitayAmiel@users.noreply.github.com> Date: Mon, 6 Apr 2026 23:37:31 +0300 Subject: [PATCH 09/59] nixl_ep: Fix starting with non consecutive ranks (e.g [0,2]) (#1478) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move the p2p_ptr_get() call inside the is_rank_masked guard in three locations in nixl_ep.cu (two in dispatch, one in combine). In elastic scenarios, ranks can be set up in non-consecutive order — for example, ranks 0 and 2 may start before rank 1 has been initialized. During this early phase, rank 1's P2P memory mapping is not yet established. Previously, p2p_ptr_get() was called unconditionally for all ranks before the is_rank_masked check, so it would attempt to access the uninitialized P2P mapping for rank 1 and fail, even though rank 1 is correctly masked and should be skipped. In each of the three locations, the p2p_ptr_get(dst_ptr, dst_rank) call is moved to inside the rank mask guard, ensuring the P2P pointer lookup only occurs for ranks that are actually participating and initialized. --- examples/device/ep/csrc/kernels/nixl_ep.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/device/ep/csrc/kernels/nixl_ep.cu b/examples/device/ep/csrc/kernels/nixl_ep.cu index 1ddcbfaf..855d1c17 100644 --- a/examples/device/ep/csrc/kernels/nixl_ep.cu +++ b/examples/device/ep/csrc/kernels/nixl_ep.cu @@ -187,8 +187,8 @@ dispatch(void* packed_recv_x, void* packed_recv_x_scales, dst_expert_local_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_msg + rank * num_max_dispatch_tokens_per_rank * num_bytes_per_msg + slot_idx * num_bytes_per_msg; - void* dst_p2p_ptr = nixl_ctx.p2p_ptr_get(dst_ptr, dst_rank); if (not is_rank_masked(mask_buffer_ptr, dst_rank)) { + void* dst_p2p_ptr = nixl_ctx.p2p_ptr_get(dst_ptr, dst_rank); if (dst_p2p_ptr == 0) { nixlMemViewElem src_mdesc{nixl_ctx.local_mvh, 0, nixl_ctx.offset_get(src_ptr)}; nixlMemViewElem dst_mdesc{nixl_ctx.remote_mvh, (size_t) dst_rank, nixl_ctx.offset_get(dst_ptr)}; @@ -256,8 +256,8 @@ dispatch(void* packed_recv_x, void* packed_recv_x_scales, // Wait local sends issued and send expert counts while (ld_acquire_global(atomic_finish_counter_per_expert + responsible_expert_idx) != FINISHED_SUM_TAG * 2); auto dst_ptr = reinterpret_cast(rdma_recv_count + dst_expert_local_idx * num_ranks + rank); - void* dst_p2p_ptr = nixl_ctx.p2p_ptr_get(dst_ptr, dst_rank); if (not is_rank_masked(mask_buffer_ptr, dst_rank)) { + void* dst_p2p_ptr = nixl_ctx.p2p_ptr_get(dst_ptr, dst_rank); if (dst_p2p_ptr == 0) { nixlMemViewElem dst_mdesc{nixl_ctx.remote_mvh, static_cast(dst_rank), nixl_ctx.offset_get(dst_ptr)}; EP_DEVICE_ASSERT(nixlAtomicAdd(num_tokens_sent + 1, dst_mdesc, dst_expert_local_idx) == NIXL_IN_PROG); @@ -805,8 +805,8 @@ combine(void* combined_x, if (sub_warp_id == 1 and lane_id == 0) { while (ld_acquire_global(atomic_clean_flag) == 0); auto dst_ptr = reinterpret_cast(rdma_recv_flag + global_expert_idx); - void* dst_p2p_ptr = nixl_ctx.p2p_ptr_get(dst_ptr, dst_rank); if (not is_rank_masked(mask_buffer_ptr, dst_rank)) { + void* dst_p2p_ptr = nixl_ctx.p2p_ptr_get(dst_ptr, dst_rank); if (dst_p2p_ptr == 0) { nixlMemViewElem dst_mdesc{nixl_ctx.remote_mvh, (size_t) dst_rank, nixl_ctx.offset_get(dst_ptr)}; EP_DEVICE_ASSERT(nixlAtomicAdd(1, dst_mdesc, local_expert_idx) == NIXL_IN_PROG); From 3a812e8f3b6ee81f85d56d194c7fba2a25a5ab9b Mon Sep 17 00:00:00 2001 From: Yoray Zack Date: Tue, 7 Apr 2026 14:21:52 +0300 Subject: [PATCH 10/59] nixl_ep: Add high-throughput kernels (#1341) * nixl_ep: Add high-throughput kernels Port the DeepEP high-throughput dispatch and combine kernels to use the NIXL device API V2 (memory views, nixlPut, nixlAtomicAdd). Rename nixl_ep.cu -> nixl_ep_ll.cu Co-authored-by: Roey Azran Co-authored-by: Micha Dery Signed-off-by: Yoray Zack * fixups pre-commit * fix pre-commit 2 * fix pre-commit3 --------- Signed-off-by: Yoray Zack Co-authored-by: Roey Azran Co-authored-by: Micha Dery --- examples/device/ep/README.md | 2 +- examples/device/ep/csrc/config.hpp | 84 +- examples/device/ep/csrc/kernels/api.cuh | 175 +- examples/device/ep/csrc/kernels/buffer.cuh | 162 ++ examples/device/ep/csrc/kernels/configs.cuh | 3 + examples/device/ep/csrc/kernels/exception.cuh | 2 +- examples/device/ep/csrc/kernels/launch.cuh | 24 + examples/device/ep/csrc/kernels/layout.cu | 157 ++ examples/device/ep/csrc/kernels/nixl_ep_ht.cu | 2423 +++++++++++++++++ .../kernels/{nixl_ep.cu => nixl_ep_ll.cu} | 38 +- examples/device/ep/csrc/kernels/runtime.cu | 54 + examples/device/ep/csrc/kernels/utils.cuh | 101 +- examples/device/ep/csrc/nixl_ep.cpp | 665 ++++- examples/device/ep/csrc/nixl_ep.hpp | 81 +- examples/device/ep/meson.build | 5 +- examples/device/ep/nixl_ep/__init__.py | 3 +- examples/device/ep/nixl_ep/buffer.py | 393 ++- examples/device/ep/nixl_ep/utils.py | 2 +- examples/device/ep/tests/test_ht.py | 578 ++++ examples/device/ep/tests/utils.py | 10 +- 20 files changed, 4870 insertions(+), 92 deletions(-) create mode 100644 examples/device/ep/csrc/kernels/buffer.cuh create mode 100644 examples/device/ep/csrc/kernels/layout.cu create mode 100644 examples/device/ep/csrc/kernels/nixl_ep_ht.cu rename examples/device/ep/csrc/kernels/{nixl_ep.cu => nixl_ep_ll.cu} (97%) create mode 100644 examples/device/ep/csrc/kernels/runtime.cu create mode 100644 examples/device/ep/tests/test_ht.py diff --git a/examples/device/ep/README.md b/examples/device/ep/README.md index af7301c5..418e68b1 100644 --- a/examples/device/ep/README.md +++ b/examples/device/ep/README.md @@ -39,7 +39,7 @@ buffer.disconnect_ranks(ranks) ## Key APIs - `Buffer(rank_id, ...)`: 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 +- `update_memory_buffers(num_ranks, num_experts_per_rank, num_rdma_bytes, num_nvl_bytes=0)`: 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 diff --git a/examples/device/ep/csrc/config.hpp b/examples/device/ep/csrc/config.hpp index 89b49cc4..c8137ebe 100644 --- a/examples/device/ep/csrc/config.hpp +++ b/examples/device/ep/csrc/config.hpp @@ -37,6 +37,82 @@ dtype_t align_up(dtype_t a, dtype_t b) { return ceil_div(a, b) * b; } +struct Config { + int num_sms; + int num_max_nvl_chunked_send_tokens; + int num_max_nvl_chunked_recv_tokens; + int num_max_rdma_chunked_send_tokens; + int num_max_rdma_chunked_recv_tokens; + + Config(int num_sms, + int num_max_nvl_chunked_send_tokens, int num_max_nvl_chunked_recv_tokens, + int num_max_rdma_chunked_send_tokens, int num_max_rdma_chunked_recv_tokens) : + num_sms(num_sms), + num_max_nvl_chunked_send_tokens(num_max_nvl_chunked_send_tokens), + num_max_nvl_chunked_recv_tokens(num_max_nvl_chunked_recv_tokens), + num_max_rdma_chunked_send_tokens(num_max_rdma_chunked_send_tokens), + num_max_rdma_chunked_recv_tokens(num_max_rdma_chunked_recv_tokens) { + EP_HOST_ASSERT(num_sms >= 0); + EP_HOST_ASSERT(num_max_nvl_chunked_send_tokens > 0 and num_max_nvl_chunked_recv_tokens > 0); + EP_HOST_ASSERT(num_max_nvl_chunked_send_tokens < num_max_nvl_chunked_recv_tokens); + EP_HOST_ASSERT(num_max_rdma_chunked_send_tokens > 0 and num_max_rdma_chunked_recv_tokens > 0); + + // Ceil up RDMA buffer size + this->num_max_rdma_chunked_recv_tokens = align_up(num_max_rdma_chunked_recv_tokens, num_max_rdma_chunked_send_tokens); + EP_HOST_ASSERT(num_max_rdma_chunked_send_tokens < num_max_rdma_chunked_recv_tokens); + // NOTES: this assertion is related to RDMA lazy head update, we must ensure senders always have space to push + EP_HOST_ASSERT(num_max_rdma_chunked_send_tokens <= num_max_rdma_chunked_recv_tokens / 2); + } + + size_t get_nvl_buffer_size_hint(size_t hidden_bytes, int num_ranks) const { + // Below are some assumptions + // TODO: add assertions + constexpr int kNumMaxTopK = 128; + constexpr int kNumMaxScales = 128; + EP_HOST_ASSERT(num_ranks < NUM_MAX_NVL_PEERS or num_ranks % NUM_MAX_NVL_PEERS == 0); + EP_HOST_ASSERT(num_ranks <= NUM_MAX_NVL_PEERS or num_sms % 2 == 0); + const auto num_rdma_ranks = std::max(num_ranks / NUM_MAX_NVL_PEERS, 1); + const auto num_nvl_ranks = std::min(num_ranks, NUM_MAX_NVL_PEERS); + const int num_channels = num_sms / 2; + + size_t num_bytes = 0; + num_bytes += num_channels * num_nvl_ranks * (2 * num_rdma_ranks + 3) * sizeof(int); + num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * hidden_bytes; + num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * ht::get_source_meta_bytes(); + num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxTopK * sizeof(int64_t); + num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxTopK * sizeof(float); + num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxScales * sizeof(float); + num_bytes = ((num_bytes + 127) / 128) * 128; + return num_bytes; + } + + size_t get_rdma_buffer_size_hint(int64_t hidden_bytes, int num_ranks) const { + // Legacy mode + if (num_ranks <= NUM_MAX_NVL_PEERS) + return 0; + + // Below are some assumptions + // TODO: add assertions + constexpr int kNumMaxTopK = 128; + constexpr int kNumMaxScales = 128; + EP_HOST_ASSERT(num_ranks % NUM_MAX_NVL_PEERS == 0); + EP_HOST_ASSERT(num_sms % 2 == 0); + const int num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; + const int num_channels = num_sms / 2; + + size_t num_bytes = 0; + num_bytes += num_channels * num_rdma_ranks * (NUM_MAX_NVL_PEERS * 2 + 2) * 2 * sizeof(int); + num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * hidden_bytes * 2; + num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * ht::get_source_meta_bytes() * 2; + num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * kNumMaxTopK * sizeof(int64_t) * 2; + num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * kNumMaxTopK * sizeof(float) * 2; + num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * kNumMaxScales * sizeof(float) * 2; + num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * sizeof(int4) * 2; + num_bytes = ((num_bytes + 127) / 128) * 128; + return num_bytes; + } +}; + struct EPBuffer { int num_clean_int = 0; @@ -123,10 +199,10 @@ struct EPLayout { }; inline size_t -get_rdma_size_hint(int num_max_dispatch_tokens_per_rank, - int hidden, - int num_ranks, - int num_experts) { +get_low_latency_buffer_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; } diff --git a/examples/device/ep/csrc/kernels/api.cuh b/examples/device/ep/csrc/kernels/api.cuh index c72443f3..93f9ba48 100644 --- a/examples/device/ep/csrc/kernels/api.cuh +++ b/examples/device/ep/csrc/kernels/api.cuh @@ -30,22 +30,183 @@ namespace nixl_ep { -// EP kernels -namespace ep_kernels { +namespace intranode { + +void barrier(int** barrier_signal_ptrs, int rank, int num_nvl_ranks, cudaStream_t stream); + +} // namespace intranode + struct gpu_nixl_ctx { nixlMemViewH local_mvh; nixlMemViewH barrier_mvh; nixlMemViewH remote_mvh; + nixlMemViewH ht_barrier_mvh; int *sync_buffer_ptr; // [src_rank] int *sync_count_ptr; // [dst_rank] + uint64_t *last_ht_barrier_counter; + uint64_t *local_ht_barrier_counter_ptr; void *rdma_buffer_ptr; int max_num_ranks; + int num_rdma_ranks; int rank; - __device__ inline uint64_t offset_get(uint64_t ptr); - __device__ inline void* p2p_ptr_get(uint64_t dst_ptr, int dst_rank); + __device__ inline uint64_t offset_get(uint64_t ptr) { + return ptr - reinterpret_cast(rdma_buffer_ptr); + } }; + +// Layout kernels +namespace layout { + +void get_dispatch_layout(const topk_idx_t* topk_idx, + int* num_tokens_per_rank, + int* num_tokens_per_rdma_rank, + int* num_tokens_per_expert, + bool* is_token_in_rank, + int num_tokens, + int num_topk, + int num_ranks, + int num_experts, + cudaStream_t stream); + +} // namespace layout + +// High-throughput kernels +namespace ht { + +int get_source_meta_bytes(); + +void notify_dispatch(const int* num_tokens_per_rank, + int* moe_recv_counter_mapped, + int num_ranks, + const int* num_tokens_per_rdma_rank, + int* moe_recv_rdma_counter_mapped, + const int* num_tokens_per_expert, + int* moe_recv_expert_counter_mapped, + int num_experts, + const bool* is_token_in_rank, + int num_tokens, + int num_channels, + int hidden_int4, + int num_scales, + int num_topk, + int expert_alignment, + int* rdma_channel_prefix_matrix, + int* recv_rdma_rank_prefix_sum, + int* gbl_channel_prefix_matrix, + int* recv_gbl_rank_prefix_sum, + void* rdma_buffer_ptr, + int num_max_rdma_chunked_recv_tokens, + void** buffer_ptrs, + int num_max_nvl_chunked_recv_tokens, + int** barrier_signal_ptrs, + int rank, + cudaStream_t stream, + int64_t num_rdma_bytes, + int64_t num_nvl_bytes, + bool low_latency_mode, + gpu_nixl_ctx nixl_ctx); + +void dispatch(void* recv_x, + float* recv_x_scales, + topk_idx_t* recv_topk_idx, + float* recv_topk_weights, + void* recv_src_meta, + const void* x, + const float* x_scales, + const topk_idx_t* topk_idx, + const float* topk_weights, + int* send_rdma_head, + int* send_nvl_head, + int* recv_rdma_channel_prefix_matrix, + int* recv_gbl_channel_prefix_matrix, + const int* rdma_channel_prefix_matrix, + const int* recv_rdma_rank_prefix_sum, + const int* gbl_channel_prefix_matrix, + const int* recv_gbl_rank_prefix_sum, + const bool* is_token_in_rank, + int num_tokens, + int hidden_int4, + int num_scales, + int num_topk, + int num_experts, + int scale_token_stride, + int scale_hidden_stride, + void* rdma_buffer_ptr, + int num_max_rdma_chunked_send_tokens, + int num_max_rdma_chunked_recv_tokens, + void** buffer_ptrs, + int num_max_nvl_chunked_send_tokens, + int num_max_nvl_chunked_recv_tokens, + int rank, + int num_ranks, + bool is_cached_dispatch, + cudaStream_t stream, + int num_channels, + bool low_latency_mode, + gpu_nixl_ctx nixl_ctx); + +void cached_notify(int hidden_int4, + int num_scales, + int num_topk_idx, + int num_topk_weights, + int num_ranks, + int num_channels, + int num_combined_tokens, + int* combined_rdma_head, + const int* rdma_channel_prefix_matrix, + const int* rdma_rank_prefix_sum, + int* combined_nvl_head, + void* rdma_buffer_ptr, + int num_max_rdma_chunked_recv_tokens, + void** buffer_ptrs, + int num_max_nvl_chunked_recv_tokens, + int** barrier_signal_ptrs, + int rank, + cudaStream_t stream, + int64_t num_rdma_bytes, + int64_t num_nvl_bytes, + bool is_cached_dispatch, + bool low_latency_mode, + gpu_nixl_ctx nixl_ctx); + +void combine(cudaDataType_t type, + void* combined_x, + float* combined_topk_weights, + const bool* is_combined_token_in_rank, + const void* x, + const float* topk_weights, + const void* bias_0, + const void* bias_1, + const int* combined_rdma_head, + const int* combined_nvl_head, + const void* src_meta, + const int* rdma_channel_prefix_matrix, + const int* rdma_rank_prefix_sum, + const int* gbl_channel_prefix_matrix, + int num_tokens, + int num_combined_tokens, + int hidden, + int num_topk, + void* rdma_buffer_ptr, + int num_max_rdma_chunked_send_tokens, + int num_max_rdma_chunked_recv_tokens, + void** buffer_ptrs, + int num_max_nvl_chunked_send_tokens, + int num_max_nvl_chunked_recv_tokens, + int rank, + int num_ranks, + cudaStream_t stream, + int num_channels, + bool low_latency_mode, + gpu_nixl_ctx nixl_ctx); + +} // namespace ht + + +// EP kernels +namespace ep_kernels { 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, @@ -64,7 +225,7 @@ void dispatch(void* packed_recv_x, void* packed_recv_x_scales, 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); + cudaStream_t stream, int phases, nixl_ep::gpu_nixl_ctx nixl_ctx); void combine(void* combined_x, void* rdma_recv_x, uint64_t* rdma_recv_flag, void* rdma_send_x, @@ -77,9 +238,9 @@ void combine(void* combined_x, 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); + cudaStream_t stream, int phases, bool zero_copy, nixl_ep::gpu_nixl_ctx nixl_ctx); -void barrier(ep_kernels::gpu_nixl_ctx nixl_ctx, int* mask_buffer_ptr, cudaStream_t stream); +void barrier(gpu_nixl_ctx nixl_ctx, int* mask_buffer_ptr, cudaStream_t stream); void query_mask_buffer(int* mask_buffer_ptr, int num_ranks, int* output_mask_tensor, cudaStream_t stream); diff --git a/examples/device/ep/csrc/kernels/buffer.cuh b/examples/device/ep/csrc/kernels/buffer.cuh new file mode 100644 index 00000000..520ea0ed --- /dev/null +++ b/examples/device/ep/csrc/kernels/buffer.cuh @@ -0,0 +1,162 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek + * SPDX-FileCopyrightText: Copyright (c) 2026 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" + +namespace nixl_ep { + +template +struct Buffer { +private: + uint8_t* ptr; + +public: + int total_bytes; + + __device__ __forceinline__ Buffer() : ptr(nullptr), total_bytes(0) {} + + __device__ __forceinline__ Buffer(void* &gbl_ptr, int num_elems, int offset = 0) { + total_bytes = num_elems * sizeof(dtype_t); + ptr = reinterpret_cast(gbl_ptr) + offset * sizeof(dtype_t); + gbl_ptr = reinterpret_cast(gbl_ptr) + total_bytes; + } + + __device__ __forceinline__ Buffer advance_also(void* &gbl_ptr) { + gbl_ptr = reinterpret_cast(gbl_ptr) + total_bytes; + return *this; + } + + __device__ __forceinline__ dtype_t* buffer() { + return reinterpret_cast(ptr); + } + + __device__ __forceinline__ dtype_t& operator[](int idx) { + return buffer()[idx]; + } +}; + +template +struct AsymBuffer { +private: + uint8_t* ptrs[kNumRanks]; + int num_bytes; + +public: + int total_bytes; + + __device__ __forceinline__ AsymBuffer(void* &gbl_ptr, int num_elems, int num_ranks, + int sm_id = 0, int num_sms = 1, int offset = 0) { + EP_STATIC_ASSERT(kNumRanks == 1, ""); + num_bytes = num_elems * sizeof(dtype_t); + + int per_channel_bytes = num_bytes * num_ranks; + total_bytes = per_channel_bytes * num_sms; + ptrs[0] = reinterpret_cast(gbl_ptr) + per_channel_bytes * sm_id + num_bytes * offset; + gbl_ptr = reinterpret_cast(gbl_ptr) + total_bytes; + } + + __device__ __forceinline__ AsymBuffer(void** gbl_ptrs, int num_elems, int num_ranks, + int sm_id = 0, int num_sms = 1, int offset = 0) { + EP_STATIC_ASSERT(kNumRanks > 1, ""); + num_bytes = num_elems * sizeof(dtype_t); + + int per_channel_bytes = num_bytes * num_ranks; + total_bytes = per_channel_bytes * num_sms; + for (int i = 0; i < kNumRanks; ++ i) { + ptrs[i] = reinterpret_cast(gbl_ptrs[i]) + per_channel_bytes * sm_id + num_bytes * offset; + gbl_ptrs[i] = reinterpret_cast(gbl_ptrs[i]) + total_bytes; + } + } + + __device__ __forceinline__ void advance(int shift) { +#ifdef __CUDACC__ + #pragma unroll +#endif + for (int i = 0; i < kNumRanks; ++ i) + ptrs[i] = ptrs[i] + shift * sizeof(dtype_t); + } + + __device__ __forceinline__ AsymBuffer advance_also(void* &gbl_ptr) { + gbl_ptr = reinterpret_cast(gbl_ptr) + total_bytes; + return *this; + } + + template + __device__ __forceinline__ AsymBuffer advance_also(void** gbl_ptrs) { + for (int i = 0; i < kNumAlsoRanks; ++ i) + gbl_ptrs[i] = reinterpret_cast(gbl_ptrs[i]) + total_bytes; + return *this; + } + + __device__ __forceinline__ dtype_t* buffer(int idx = 0) { + EP_STATIC_ASSERT(kNumRanks == 1, "`buffer` is only available for single rank case"); + return reinterpret_cast(ptrs[0] + num_bytes * idx); + } + + __device__ __forceinline__ dtype_t* buffer_by(int rank_idx, int idx = 0) { + EP_STATIC_ASSERT(kNumRanks > 1, "`buffer` is only available for single rank case"); + return reinterpret_cast(ptrs[rank_idx] + num_bytes * idx); + } +}; + +template +struct SymBuffer { +private: + // NOTES: for non-decoupled case, `recv_ptr` is not used + uint8_t* send_ptr; + uint8_t* recv_ptr; + int num_bytes; + +public: + int total_bytes; + + __device__ __forceinline__ SymBuffer(void* &gbl_ptr, int num_elems, int num_ranks, + int sm_id = 0, int num_sms = 1) { + num_bytes = num_elems * sizeof(dtype_t); + + int per_channel_bytes = num_bytes * num_ranks; + total_bytes = per_channel_bytes * num_sms * (static_cast(kDecoupled) + 1); + send_ptr = reinterpret_cast(gbl_ptr) + per_channel_bytes * sm_id; + recv_ptr = reinterpret_cast(gbl_ptr) + per_channel_bytes * (sm_id + num_sms); + gbl_ptr = reinterpret_cast(gbl_ptr) + total_bytes; + } + + __device__ __forceinline__ dtype_t* send_buffer(int idx = 0) { + EP_STATIC_ASSERT(kDecoupled, "`send_buffer` is only available for non-decoupled case"); + return reinterpret_cast(send_ptr + num_bytes * idx); + } + + __device__ __forceinline__ dtype_t* recv_buffer(int idx = 0) { + EP_STATIC_ASSERT(kDecoupled, "`recv_buffer` is only available for non-decoupled case"); + return reinterpret_cast(recv_ptr + num_bytes * idx); + } + + __device__ __forceinline__ dtype_t* buffer(int idx = 0) { + EP_STATIC_ASSERT(not kDecoupled, "`buffer` is only available for decoupled case"); + return reinterpret_cast(send_ptr + num_bytes * idx); + } +}; + +} // namespace nixl_ep diff --git a/examples/device/ep/csrc/kernels/configs.cuh b/examples/device/ep/csrc/kernels/configs.cuh index cd3ce777..73b3c5da 100644 --- a/examples/device/ep/csrc/kernels/configs.cuh +++ b/examples/device/ep/csrc/kernels/configs.cuh @@ -22,11 +22,14 @@ #pragma once +#define NUM_MAX_NVL_PEERS 8 +#define NUM_MAX_RDMA_PEERS 20 #define NUM_WORKSPACE_BYTES (32 * 1024 * 1024) #define NUM_MAX_LOCAL_EXPERTS 1024 #define NUM_BUFFER_ALIGNMENT_BYTES 128 #define FINISHED_SUM_TAG 1024 +#define NUM_WAIT_NANOSECONDS 500 #ifndef ENABLE_FAST_DEBUG #define NUM_CPU_TIMEOUT_SECS 100 diff --git a/examples/device/ep/csrc/kernels/exception.cuh b/examples/device/ep/csrc/kernels/exception.cuh index 93bdc684..2de5c13d 100644 --- a/examples/device/ep/csrc/kernels/exception.cuh +++ b/examples/device/ep/csrc/kernels/exception.cuh @@ -1,6 +1,6 @@ /* * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 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. diff --git a/examples/device/ep/csrc/kernels/launch.cuh b/examples/device/ep/csrc/kernels/launch.cuh index 6c044a45..22090a40 100644 --- a/examples/device/ep/csrc/kernels/launch.cuh +++ b/examples/device/ep/csrc/kernels/launch.cuh @@ -73,6 +73,30 @@ cfg.dynamicSmemBytes = smem_size; #endif #endif +#define SWITCH_NVL_RANKS(case_macro) \ + switch (num_nvl_ranks) { \ + case 2: \ + case_macro(2); \ + case 4: \ + case_macro(4); \ + case 8: \ + case_macro(8); \ + default: \ + EP_HOST_ASSERT(false and "Unsupported NVL ranks"); \ + } \ + while (false) + +#define SWITCH_RDMA_RANKS(case_macro) \ + switch (num_ranks / NUM_MAX_NVL_PEERS) { \ + case 2: case_macro(2); \ + case 4: case_macro(4); \ + case 8: case_macro(8); \ + case 16: case_macro(16); \ + case 18: case_macro(18); \ + case 20: case_macro(20); \ + default: EP_HOST_ASSERT(false and "Unsupported RDMA ranks"); \ + } while (false) + #define SWITCH_HIDDEN(case_macro) \ switch (hidden) { \ case 2048: case_macro(2048); \ diff --git a/examples/device/ep/csrc/kernels/layout.cu b/examples/device/ep/csrc/kernels/layout.cu new file mode 100644 index 00000000..3fd3b719 --- /dev/null +++ b/examples/device/ep/csrc/kernels/layout.cu @@ -0,0 +1,157 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 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" + +namespace nixl_ep { + +namespace layout { + +template +__global__ void get_dispatch_layout(const topk_idx_t* topk_idx, + int* num_tokens_per_rank, int* num_tokens_per_rdma_rank, + int* num_tokens_per_expert, bool* is_token_in_rank, + int num_tokens, int num_topk, int num_ranks, int num_experts) { + auto sm_id = static_cast(blockIdx.x); + auto thread_id = static_cast(threadIdx.x); + + // Count expert statistics + __shared__ int num_tokens_per_expert_per_thread[kNumThreads][kNumExpertsPerSM]; + int expert_begin_idx = sm_id * kNumExpertsPerSM, expert_end_idx = min(expert_begin_idx + kNumExpertsPerSM, num_experts); + if (expert_begin_idx < expert_end_idx) { + // Per-thread count + #pragma unroll + for (int i = 0; i < kNumExpertsPerSM; ++ i) + num_tokens_per_expert_per_thread[thread_id][i] = 0; + #pragma unroll + for (int i = thread_id; i < num_tokens; i += kNumThreads) { + auto shifted_topk_idx = topk_idx + i * num_topk; + #pragma unroll + for (int j = 0, expert_idx; j < num_topk; ++ j) { + expert_idx = static_cast(shifted_topk_idx[j]); + if (expert_begin_idx <= expert_idx and expert_idx < expert_end_idx) + ++ num_tokens_per_expert_per_thread[thread_id][expert_idx - expert_begin_idx]; + } + } + __syncthreads(); + + // Sum up + EP_STATIC_ASSERT(kNumExpertsPerSM <= kNumThreads, "Too many experts per SM"); + if (expert_begin_idx + thread_id < expert_end_idx) { + int sum = 0; + #pragma unroll + for (int i = 0; i < kNumThreads; ++ i) + sum += num_tokens_per_expert_per_thread[i][thread_id]; + num_tokens_per_expert[expert_begin_idx + thread_id] = sum; + } + return; + } + + if (num_tokens_per_rdma_rank != nullptr) + EP_DEVICE_ASSERT(num_ranks % NUM_MAX_NVL_PEERS == 0 and num_ranks > NUM_MAX_NVL_PEERS); + + // Count rank statistics + constexpr int kNumRDMARanksPerSM = kNumRanksPerSM / NUM_MAX_NVL_PEERS; + __shared__ int num_tokens_per_rank_per_thread[kNumThreads][kNumRanksPerSM]; + __shared__ int num_tokens_per_rdma_rank_per_thread[kNumThreads][kNumRDMARanksPerSM]; + auto sm_begin = (num_experts + kNumExpertsPerSM - 1) / kNumExpertsPerSM; + int rank_begin_idx = (sm_id - sm_begin) * kNumRanksPerSM, rank_end_idx = min(rank_begin_idx + kNumRanksPerSM, num_ranks); + int rdma_rank_begin_idx = rank_begin_idx / NUM_MAX_NVL_PEERS, rdma_rank_end_idx = rank_end_idx / NUM_MAX_NVL_PEERS; + if (rank_begin_idx < rank_end_idx) { + const auto num_expert_per_rank = num_experts / num_ranks; + auto expert_begin = rank_begin_idx * num_expert_per_rank; + auto expert_end = rank_end_idx * num_expert_per_rank; + + // Per-thread count + #pragma unroll + for (int i = 0; i < kNumRanksPerSM; ++ i) + num_tokens_per_rank_per_thread[thread_id][i] = 0; + #pragma unroll + for (int i = 0; i < kNumRDMARanksPerSM; ++ i) + num_tokens_per_rdma_rank_per_thread[thread_id][i] = 0; + #pragma unroll + for (int i = thread_id; i < num_tokens; i += kNumThreads) { + auto shifted_topk_idx = topk_idx + i * num_topk; + int is_in_rank[kNumRanksPerSM] = {0}, is_in_rdma_rank[kNumRDMARanksPerSM] = {0}; + #pragma unroll + for (int j = 0, expert_idx, rank_idx; j < num_topk; ++j) { + expert_idx = static_cast(shifted_topk_idx[j]); + if (expert_begin <= expert_idx and expert_idx < expert_end) { + // Count single rank + rank_idx = expert_idx / num_expert_per_rank - rank_begin_idx; + is_in_rank[rank_idx] ++, is_in_rdma_rank[rank_idx / NUM_MAX_NVL_PEERS] ++; + } + } + + auto shifted_is_token_in_rank = is_token_in_rank + i * num_ranks; + #pragma unroll + for (int j = 0; j + rank_begin_idx < rank_end_idx; ++ j) { + shifted_is_token_in_rank[j + rank_begin_idx] = (is_in_rank[j] > 0); + num_tokens_per_rank_per_thread[thread_id][j] += (is_in_rank[j] > 0); + } + + #pragma unroll + for (int j = 0; j + rdma_rank_begin_idx < rdma_rank_end_idx; ++ j) + num_tokens_per_rdma_rank_per_thread[thread_id][j] += (is_in_rdma_rank[j] > 0); + } + __syncthreads(); + + // Sum up + EP_STATIC_ASSERT(kNumRanksPerSM <= kNumThreads, "Too many ranks per SM"); + if (rank_begin_idx + thread_id < rank_end_idx) { + int sum = 0; + #pragma unroll + for (int i = 0; i < kNumThreads; ++ i) + sum += num_tokens_per_rank_per_thread[i][thread_id]; + num_tokens_per_rank[rank_begin_idx + thread_id] = sum; + } + + if (num_tokens_per_rdma_rank != nullptr and rdma_rank_begin_idx + thread_id < rdma_rank_end_idx) { + int sum = 0; + #pragma unroll + for (int i = 0; i < kNumThreads; ++ i) + sum += num_tokens_per_rdma_rank_per_thread[i][thread_id]; + num_tokens_per_rdma_rank[rdma_rank_begin_idx + thread_id] = sum; + } + } +} + +void get_dispatch_layout(const topk_idx_t* topk_idx, + int* num_tokens_per_rank, int* num_tokens_per_rdma_rank, + int* num_tokens_per_expert, bool* is_token_in_rank, + int num_tokens, int num_topk, int num_ranks, int num_experts, + cudaStream_t stream) { + constexpr int kNumThreads = 256, kNumExpertsPerSM = 4, kNumRanksPerSM = 8; + int num_sms = ((num_experts + kNumExpertsPerSM - 1) / kNumExpertsPerSM) + (num_ranks + kNumRanksPerSM - 1) / kNumRanksPerSM; + EP_STATIC_ASSERT(kNumRanksPerSM % NUM_MAX_NVL_PEERS == 0, "Invalid number of ranks per SM"); + + SETUP_LAUNCH_CONFIG(num_sms, kNumThreads, stream); + LAUNCH_KERNEL(&cfg, (get_dispatch_layout), + topk_idx, num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, + num_tokens, num_topk, num_ranks, num_experts); +} + +} // namespace layout + +} // namespace nixl_ep diff --git a/examples/device/ep/csrc/kernels/nixl_ep_ht.cu b/examples/device/ep/csrc/kernels/nixl_ep_ht.cu new file mode 100644 index 00000000..140ba095 --- /dev/null +++ b/examples/device/ep/csrc/kernels/nixl_ep_ht.cu @@ -0,0 +1,2423 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek + * SPDX-FileCopyrightText: Copyright (c) 2026 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 +#include + +#include "buffer.cuh" +#include "configs.cuh" +#include "exception.cuh" +#include "launch.cuh" +#include "utils.cuh" +#include "api.cuh" +#include "nixl_device.cuh" +#include +#include +#include + +namespace nixl_ep { + +namespace ht { + +struct SourceMeta { + int src_rdma_rank, is_token_in_nvl_rank_bits; + + EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS == 8, "Invalid number of maximum NVL peers"); + + __forceinline__ SourceMeta() = default; + + // TODO: faster encoding + __device__ __forceinline__ SourceMeta(int rdma_rank, const bool* is_token_in_nvl_ranks) { + src_rdma_rank = rdma_rank; + is_token_in_nvl_rank_bits = is_token_in_nvl_ranks[0]; + #pragma unroll + for (int i = 1; i < NUM_MAX_NVL_PEERS; ++i) + is_token_in_nvl_rank_bits |= is_token_in_nvl_ranks[i] << i; + } + + __device__ __forceinline__ bool is_token_in_nvl_rank(int nvl_rank) const { return (is_token_in_nvl_rank_bits >> nvl_rank) & 1; } +}; + +EP_STATIC_ASSERT(sizeof(SourceMeta) % sizeof(int) == 0, "Invalid size of `SourceMeta`"); + +int get_source_meta_bytes() { + return sizeof(SourceMeta); +} + +__host__ __device__ __forceinline__ int get_num_bytes_per_token(int hidden_int4, int num_scales, int num_topk_idx, int num_topk_weights) { + return static_cast(align_up(hidden_int4 * sizeof(int4) + sizeof(SourceMeta) + num_scales * sizeof(float) + + num_topk_idx * sizeof(int) + num_topk_weights * sizeof(float), + sizeof(int4))); +} + +__host__ __device__ __forceinline__ std::pair get_rdma_clean_meta(int hidden_int4, + int num_scales, + int num_topk_idx, + int num_topk_weights, + int num_rdma_ranks, + int num_rdma_recv_buffer_tokens, + int num_channels) { + // Return `int32_t` offset and count to clean + return {(get_num_bytes_per_token(hidden_int4, num_scales, num_topk_idx, num_topk_weights) * num_rdma_recv_buffer_tokens * + num_rdma_ranks * 2 * num_channels) / + sizeof(int), + (NUM_MAX_NVL_PEERS * 2 + 4) * num_rdma_ranks * 2 * num_channels}; +} + +__host__ __device__ __forceinline__ std::pair get_nvl_clean_meta(int hidden_int4, + int num_scales, + int num_topk_idx, + int num_topk_weights, + int num_rdma_ranks, + int num_nvl_ranks, + int num_nvl_recv_buffer_tokens, + int num_channels, + bool is_dispatch) { + // Return `int32_t` offset and to clean + EP_STATIC_ASSERT(sizeof(SourceMeta) % sizeof(int) == 0, "Invalid size of `SourceMeta`"); + + return { + (num_nvl_recv_buffer_tokens * get_num_bytes_per_token(hidden_int4, num_scales, num_topk_idx, num_topk_weights) * num_nvl_ranks * + num_channels) / + sizeof(int), + num_nvl_ranks * (2 * num_rdma_ranks + 2) * num_channels, + }; +} + +template +__forceinline__ __device__ int translate_dst_rdma_rank(const int dst_rdma_rank, const int nvl_rank) { + return dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank; +} + +__forceinline__ __device__ void nixl_barrier_send_warp(nixl_ep::gpu_nixl_ctx nixl_ctx, int num_channels) { + int rdma_rank = nixl_ctx.rank / NUM_MAX_NVL_PEERS; + int nvl_rank = nixl_ctx.rank % NUM_MAX_NVL_PEERS; + int lane_id = get_lane_id(); + + for (int j = lane_id; j < num_channels; j += 32) { + for (int i = 0; i < nixl_ctx.num_rdma_ranks; i++) { + if (i == rdma_rank) continue; + int global_dst_rank = i * NUM_MAX_NVL_PEERS + nvl_rank; + nixlMemViewElem barrier_mdesc{nixl_ctx.ht_barrier_mvh, (size_t)global_dst_rank, 0}; + EP_DEVICE_ASSERT(nixlAtomicAdd( + 1, barrier_mdesc, j, 0) == NIXL_IN_PROG); + } + } +} + +__forceinline__ __device__ void nixl_barrier_wait(nixl_ep::gpu_nixl_ctx nixl_ctx, int num_channels) { + uint64_t epoch = ld_acquire_sys_global(nixl_ctx.last_ht_barrier_counter); + uint64_t expected_counter = (epoch + num_channels) * (nixl_ctx.num_rdma_ranks - 1); + while (ld_acquire_sys_global(nixl_ctx.local_ht_barrier_counter_ptr) < expected_counter); + st_release_sys_global(nixl_ctx.last_ht_barrier_counter, epoch + num_channels); +} + +template +__global__ void notify_dispatch(const int* num_tokens_per_rank, + int* moe_recv_counter_mapped, + int num_ranks, + const int* num_tokens_per_rdma_rank, + int* moe_recv_rdma_counter_mapped, + const int* num_tokens_per_expert, + int* moe_recv_expert_counter_mapped, + int num_experts, + const bool* is_token_in_rank, + int num_tokens, + int num_channels, + int expert_alignment, + const int rdma_clean_offset, + const int rdma_num_int_clean, + const int nvl_clean_offset, + const int nvl_num_int_clean, + int* rdma_channel_prefix_matrix, + int* recv_rdma_rank_prefix_sum, + int* gbl_channel_prefix_matrix, + int* recv_gbl_rank_prefix_sum, + void* rdma_buffer_ptr, + void** buffer_ptrs, + int** barrier_signal_ptrs, + int rank, + nixl_ep::gpu_nixl_ctx nixl_ctx) { + auto sm_id = static_cast(blockIdx.x); + auto thread_id = static_cast(threadIdx.x), warp_id = thread_id / 32, lane_id = get_lane_id(); + auto num_threads = static_cast(blockDim.x), num_warps = num_threads / 32; + + auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS; + auto num_rdma_experts = num_experts / kNumRDMARanks, num_nvl_experts = num_rdma_experts / NUM_MAX_NVL_PEERS; + + if (sm_id == 0) { + // Communication with others + // Global barrier: the first warp does intra-node sync, the second warp does internode sync + EP_DEVICE_ASSERT(num_warps > 1); + EP_DEVICE_ASSERT(kNumRDMARanks <= num_threads); + + __syncthreads(); + + if (warp_id == 1) { + nixl_barrier_send_warp(nixl_ctx, num_channels); + if (lane_id == 0) + nixl_barrier_wait(nixl_ctx, num_channels); + } + barrier_block(barrier_signal_ptrs, nvl_rank); + + // Send numbers of tokens per rank/expert to RDMA ranks + auto rdma_buffer_ptr_int = static_cast(rdma_buffer_ptr); + auto rdma_recv_num_tokens_mixed = SymBuffer(rdma_buffer_ptr, NUM_MAX_NVL_PEERS + num_rdma_experts + 1, kNumRDMARanks); + + // Clean up for later data dispatch + EP_DEVICE_ASSERT(rdma_recv_num_tokens_mixed.total_bytes <= rdma_clean_offset * sizeof(int)); + #pragma unroll + for (int i = thread_id; i < rdma_num_int_clean; i += num_threads) + rdma_buffer_ptr_int[rdma_clean_offset + i] = 0; + + // Copy to send buffer + #pragma unroll + for (int i = thread_id; i < num_ranks; i += num_threads) + rdma_recv_num_tokens_mixed.send_buffer(i / NUM_MAX_NVL_PEERS)[i % NUM_MAX_NVL_PEERS] = num_tokens_per_rank[i]; + #pragma unroll + for (int i = thread_id; i < num_experts; i += num_threads) + rdma_recv_num_tokens_mixed.send_buffer(i / num_rdma_experts)[NUM_MAX_NVL_PEERS + i % num_rdma_experts] = + num_tokens_per_expert[i]; + if (thread_id < kNumRDMARanks) + rdma_recv_num_tokens_mixed.send_buffer(thread_id)[NUM_MAX_NVL_PEERS + num_rdma_experts] = num_tokens_per_rdma_rank[thread_id]; + __syncthreads(); + + // Issue send + // TODO: more light fence or barrier or signaling + // TODO: overlap EP barrier and NVL cleaning + for (int i = warp_id; i < kNumRDMARanks; i += num_warps) { + if (i != rdma_rank) { + size_t src_offset = nixl_ctx.offset_get(reinterpret_cast(rdma_recv_num_tokens_mixed.send_buffer(i))); + size_t dst_offset = nixl_ctx.offset_get(reinterpret_cast(rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank))); + size_t msg_size = (NUM_MAX_NVL_PEERS + num_rdma_experts + 1) * sizeof(int); + int translated_dst = translate_dst_rdma_rank(i, nvl_rank); + nixlMemViewElem src_mdesc{nixl_ctx.local_mvh, 0, src_offset}; + nixlMemViewElem dst_mdesc{nixl_ctx.remote_mvh, (size_t)translated_dst, dst_offset}; + nixl_status_t status = nixlPut( + src_mdesc, dst_mdesc, msg_size, 0); + EP_DEVICE_ASSERT(status == NIXL_IN_PROG); + } else { + UNROLLED_WARP_COPY(1, + lane_id, + NUM_MAX_NVL_PEERS + num_rdma_experts + 1, + rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank), + rdma_recv_num_tokens_mixed.send_buffer(i), + ld_volatile_global, + st_na_global); + } + } + __syncthreads(); + + + // Barrier + if (warp_id == 0) { + nixl_barrier_send_warp(nixl_ctx, num_channels); + if (lane_id == 0) + nixl_barrier_wait(nixl_ctx, num_channels); + } + __syncthreads(); + + // NVL buffers + auto nvl_send_buffer = thread_id < NUM_MAX_NVL_PEERS ? buffer_ptrs[thread_id] : nullptr; + auto nvl_recv_buffer = buffer_ptrs[nvl_rank]; + auto nvl_reduced_num_tokens_per_expert = Buffer(nvl_recv_buffer, num_rdma_experts).advance_also(nvl_send_buffer); + auto nvl_send_num_tokens_per_rank = AsymBuffer(nvl_send_buffer, kNumRDMARanks, NUM_MAX_NVL_PEERS); + auto nvl_send_num_tokens_per_expert = AsymBuffer(nvl_send_buffer, num_nvl_experts, NUM_MAX_NVL_PEERS); + auto nvl_recv_num_tokens_per_rank = AsymBuffer(nvl_recv_buffer, kNumRDMARanks, NUM_MAX_NVL_PEERS); + auto nvl_recv_num_tokens_per_expert = AsymBuffer(nvl_recv_buffer, num_nvl_experts, NUM_MAX_NVL_PEERS); + + // Clean up for later data dispatch + auto nvl_buffer_ptr_int = static_cast(buffer_ptrs[nvl_rank]); + EP_DEVICE_ASSERT(nvl_reduced_num_tokens_per_expert.total_bytes + nvl_send_num_tokens_per_rank.total_bytes + + nvl_send_num_tokens_per_expert.total_bytes <= + nvl_clean_offset * sizeof(int)); + #pragma unroll + for (int i = thread_id; i < nvl_num_int_clean; i += num_threads) + nvl_buffer_ptr_int[nvl_clean_offset + i] = 0; + + // Reduce number of tokens per expert into the NVL send buffer + // TODO: may use NVSHMEM reduction + EP_DEVICE_ASSERT(num_rdma_experts <= num_threads); + if (thread_id < num_rdma_experts) { + int sum = 0; + #pragma unroll + for (int i = 0; i < kNumRDMARanks; ++i) + sum += rdma_recv_num_tokens_mixed.recv_buffer(i)[NUM_MAX_NVL_PEERS + thread_id]; + nvl_reduced_num_tokens_per_expert[thread_id] = sum; + } + __syncthreads(); + + // Reduce RDMA received tokens + if (thread_id == 0) { + int sum = 0; + #pragma unroll + for (int i = 0; i < kNumRDMARanks; ++i) { + sum += rdma_recv_num_tokens_mixed.recv_buffer(i)[NUM_MAX_NVL_PEERS + num_rdma_experts]; + recv_rdma_rank_prefix_sum[i] = sum; + } + while (ld_volatile_global(moe_recv_rdma_counter_mapped) != -1); + *moe_recv_rdma_counter_mapped = sum; + } + + // Send numbers of tokens per rank/expert to NVL ranks + EP_DEVICE_ASSERT(NUM_MAX_NVL_PEERS <= num_threads); + if (thread_id < NUM_MAX_NVL_PEERS) { + #pragma unroll + for (int i = 0; i < kNumRDMARanks; ++i) + nvl_send_num_tokens_per_rank.buffer(nvl_rank)[i] = rdma_recv_num_tokens_mixed.recv_buffer(i)[thread_id]; + #pragma unroll + for (int i = 0; i < num_nvl_experts; ++i) + nvl_send_num_tokens_per_expert.buffer(nvl_rank)[i] = nvl_reduced_num_tokens_per_expert[thread_id * num_nvl_experts + i]; + } + barrier_block(barrier_signal_ptrs, nvl_rank); + + // Reduce the number of tokens per rank/expert + EP_DEVICE_ASSERT(num_nvl_experts <= num_threads); + if (thread_id == 0) { + int sum = 0; + #pragma unroll + for (int i = 0; i < num_ranks; ++i) { + int src_rdma_rank = i / NUM_MAX_NVL_PEERS, src_nvl_rank = i % NUM_MAX_NVL_PEERS; + sum += nvl_recv_num_tokens_per_rank.buffer(src_nvl_rank)[src_rdma_rank]; + recv_gbl_rank_prefix_sum[i] = sum; + } + while (ld_volatile_global(moe_recv_counter_mapped) != -1); + *moe_recv_counter_mapped = sum; + } + if (thread_id < num_nvl_experts) { + int sum = 0; + #pragma unroll + for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) + sum += nvl_recv_num_tokens_per_expert.buffer(i)[thread_id]; + sum = (sum + expert_alignment - 1) / expert_alignment * expert_alignment; + while (ld_volatile_global(moe_recv_expert_counter_mapped + thread_id) != -1); + moe_recv_expert_counter_mapped[thread_id] = sum; + } + + // Finally barrier + if (warp_id == 1) { + nixl_barrier_send_warp(nixl_ctx, num_channels); + if (lane_id == 0) + nixl_barrier_wait(nixl_ctx, num_channels); + } + barrier_block(barrier_signal_ptrs, nvl_rank); + } else { + // Calculate meta data + int dst_rdma_rank = sm_id - 1; + for (int channel_id = warp_id; channel_id < num_channels; channel_id += num_warps) { + int token_start_idx, token_end_idx; + get_channel_task_range(num_tokens, num_channels, channel_id, token_start_idx, token_end_idx); + + // Iterate over tokens + int total_count = 0, per_nvl_rank_count[NUM_MAX_NVL_PEERS] = {0}; + for (int64_t i = token_start_idx + lane_id; i < token_end_idx; i += 32) { + EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS * sizeof(bool) == sizeof(uint64_t), "Invalid number of NVL peers"); + auto is_token_in_rank_uint64 = + *reinterpret_cast(is_token_in_rank + i * num_ranks + dst_rdma_rank * NUM_MAX_NVL_PEERS); + auto is_token_in_rank_values = reinterpret_cast(&is_token_in_rank_uint64); + #pragma unroll + for (int j = 0; j < NUM_MAX_NVL_PEERS; ++j) + per_nvl_rank_count[j] += is_token_in_rank_values[j]; + total_count += (is_token_in_rank_uint64 != 0); + } + + // Warp reduce + total_count = warp_reduce_sum(total_count); + #pragma unroll + for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) + per_nvl_rank_count[i] = warp_reduce_sum(per_nvl_rank_count[i]); + + // Write into channel matrix + if (elect_one_sync()) { + #pragma unroll + for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) + gbl_channel_prefix_matrix[(dst_rdma_rank * NUM_MAX_NVL_PEERS + i) * num_channels + channel_id] = per_nvl_rank_count[i]; + rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id] = total_count; + } + } + + // Calculate prefix sum + __syncthreads(); + if (thread_id == 0) { + auto prefix_row = rdma_channel_prefix_matrix + dst_rdma_rank * num_channels; + #pragma unroll + for (int i = 1; i < num_channels; ++i) + prefix_row[i] += prefix_row[i - 1]; + } + + EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); + if (thread_id < NUM_MAX_NVL_PEERS) { + auto prefix_row = gbl_channel_prefix_matrix + (dst_rdma_rank * NUM_MAX_NVL_PEERS + thread_id) * num_channels; + #pragma unroll + for (int i = 1; i < num_channels; ++i) + prefix_row[i] += prefix_row[i - 1]; + } + } +} + +void notify_dispatch(const int* num_tokens_per_rank, + int* moe_recv_counter_mapped, + int num_ranks, + const int* num_tokens_per_rdma_rank, + int* moe_recv_rdma_counter_mapped, + const int* num_tokens_per_expert, + int* moe_recv_expert_counter_mapped, + int num_experts, + const bool* is_token_in_rank, + int num_tokens, + int num_channels, + int hidden_int4, + int num_scales, + int num_topk, + int expert_alignment, + int* rdma_channel_prefix_matrix, + int* recv_rdma_rank_prefix_sum, + int* gbl_channel_prefix_matrix, + int* recv_gbl_rank_prefix_sum, + void* rdma_buffer_ptr, + int num_max_rdma_chunked_recv_tokens, + void** buffer_ptrs, + int num_max_nvl_chunked_recv_tokens, + int** barrier_signal_ptrs, + int rank, + cudaStream_t stream, + int64_t num_rdma_bytes, + int64_t num_nvl_bytes, + bool low_latency_mode, + nixl_ep::gpu_nixl_ctx nixl_ctx) { +#define NOTIFY_DISPATCH_LAUNCH_CASE(num_rdma_ranks) \ + { \ + auto notify_dispatch_func = low_latency_mode ? notify_dispatch : notify_dispatch; \ + LAUNCH_KERNEL(&cfg, \ + notify_dispatch_func, \ + num_tokens_per_rank, \ + moe_recv_counter_mapped, \ + num_ranks, \ + num_tokens_per_rdma_rank, \ + moe_recv_rdma_counter_mapped, \ + num_tokens_per_expert, \ + moe_recv_expert_counter_mapped, \ + num_experts, \ + is_token_in_rank, \ + num_tokens, \ + num_channels, \ + expert_alignment, \ + rdma_clean_meta.first, \ + rdma_clean_meta.second, \ + nvl_clean_meta.first, \ + nvl_clean_meta.second, \ + rdma_channel_prefix_matrix, \ + recv_rdma_rank_prefix_sum, \ + gbl_channel_prefix_matrix, \ + recv_gbl_rank_prefix_sum, \ + rdma_buffer_ptr, \ + buffer_ptrs, \ + barrier_signal_ptrs, \ + rank, \ + nixl_ctx); \ + } \ + break + + constexpr int kNumThreads = 512; + const auto num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; + + // Get clean meta + auto rdma_clean_meta = + get_rdma_clean_meta(hidden_int4, num_scales, num_topk, num_topk, num_rdma_ranks, num_max_rdma_chunked_recv_tokens, num_channels); + auto nvl_clean_meta = get_nvl_clean_meta(hidden_int4, + num_scales, + num_topk, + num_topk, + num_rdma_ranks, + NUM_MAX_NVL_PEERS, + num_max_nvl_chunked_recv_tokens, + num_channels, + true); + EP_HOST_ASSERT((rdma_clean_meta.first + rdma_clean_meta.second) * sizeof(int) <= num_rdma_bytes); + EP_HOST_ASSERT((nvl_clean_meta.first + nvl_clean_meta.second) * sizeof(int) <= num_nvl_bytes); + EP_HOST_ASSERT(num_rdma_bytes < std::numeric_limits::max()); + EP_HOST_ASSERT(num_nvl_bytes < std::numeric_limits::max()); + + // Launch kernel + SETUP_LAUNCH_CONFIG(1 + num_rdma_ranks, kNumThreads, stream); + SWITCH_RDMA_RANKS(NOTIFY_DISPATCH_LAUNCH_CASE); +#undef NOTIFY_DISPATCH_LAUNCH_CASE +} + +// At most 8 RDMA ranks to be sent +constexpr int get_num_topk_rdma_ranks(int num_rdma_ranks) { + return num_rdma_ranks < 8 ? num_rdma_ranks : 8; +} + +template +__global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NVL_PEERS) * 32), 1) + dispatch(int4* recv_x, + float* recv_x_scales, + topk_idx_t* recv_topk_idx, + float* recv_topk_weights, + SourceMeta* recv_src_meta, + const int4* x, + const float* x_scales, + const topk_idx_t* topk_idx, + const float* topk_weights, + int* send_rdma_head, + int* send_nvl_head, + int* recv_rdma_channel_prefix_matrix, + int* recv_gbl_channel_prefix_matrix, + const int* rdma_channel_prefix_matrix, + const int* recv_rdma_rank_prefix_sum, + const int* gbl_channel_prefix_matrix, + const int* recv_gbl_rank_prefix_sum, + const bool* is_token_in_rank, + int num_tokens, + int hidden_int4, + int num_scales, + int num_topk, + int num_experts, + int scale_token_stride, + int scale_hidden_stride, + void* rdma_buffer_ptr, + int num_max_rdma_chunked_send_tokens, + int num_max_rdma_chunked_recv_tokens, + void** buffer_ptrs, + int num_max_nvl_chunked_send_tokens, + int num_max_nvl_chunked_recv_tokens, + int rank, + int num_ranks, + nixl_ep::gpu_nixl_ctx nixl_ctx) { + enum class WarpRole { kRDMASender, kRDMASenderCoordinator, kRDMAAndNVLForwarder, kForwarderCoordinator, kNVLReceivers }; + + const auto num_sms = static_cast(gridDim.x); + const auto sm_id = static_cast(blockIdx.x); + const auto num_threads = static_cast(blockDim.x), num_warps = num_threads / 32; + const auto thread_id = static_cast(threadIdx.x), warp_id = thread_id / 32, lane_id = get_lane_id(); + const auto num_channels = num_sms / 2, channel_id = sm_id / 2; + const bool is_forwarder = sm_id % 2 == 0; + const auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS; + + + const auto role_meta = [=]() -> std::pair { + if (is_forwarder) { + if (warp_id < NUM_MAX_NVL_PEERS) { + return {WarpRole::kRDMAAndNVLForwarder, (warp_id + channel_id) % NUM_MAX_NVL_PEERS}; + } else { + return {WarpRole::kForwarderCoordinator, warp_id - NUM_MAX_NVL_PEERS}; + } + } else if (warp_id < kNumDispatchRDMASenderWarps) { + return {WarpRole::kRDMASender, -1}; + } else if (warp_id == kNumDispatchRDMASenderWarps) { + return {WarpRole::kRDMASenderCoordinator, -1}; + } else { + return {WarpRole::kNVLReceivers, (warp_id + channel_id - kNumDispatchRDMASenderWarps) % NUM_MAX_NVL_PEERS}; + } + }(); + auto warp_role = role_meta.first; + auto target_rank = role_meta.second; // Not applicable for RDMA senders + EP_DEVICE_ASSERT(num_warps == kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NVL_PEERS); + + // Data checks + EP_DEVICE_ASSERT(num_topk <= 32); + + // RDMA symmetric layout + EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS * sizeof(bool) == sizeof(uint64_t), "Invalid number of NVL peers"); + auto hidden_bytes = hidden_int4 * sizeof(int4); + auto scale_bytes = num_scales * sizeof(float); + auto num_bytes_per_token = get_num_bytes_per_token(hidden_int4, num_scales, num_topk, num_topk); + auto rdma_channel_data = SymBuffer( + rdma_buffer_ptr, num_max_rdma_chunked_recv_tokens * num_bytes_per_token, kNumRDMARanks, channel_id, num_channels); + auto rdma_channel_meta = SymBuffer(rdma_buffer_ptr, NUM_MAX_NVL_PEERS * 2 + 2, kNumRDMARanks, channel_id, num_channels); + auto rdma_channel_head = SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); + auto rdma_channel_tail = SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); + + // NVL buffer layouts + // NOTES: `rs_wr_buffer_ptr` means "Read for Senders, Write for Receivers", `ws_rr_buffer_ptr` means "Write for Senders, Read for + // Receivers" + void *rs_wr_buffer_ptr = nullptr, *ws_rr_buffer_ptr = nullptr; + int rs_wr_rank = 0, ws_rr_rank = 0; + if (warp_role == WarpRole::kRDMAAndNVLForwarder) + rs_wr_buffer_ptr = buffer_ptrs[nvl_rank], ws_rr_buffer_ptr = buffer_ptrs[target_rank], rs_wr_rank = nvl_rank, + ws_rr_rank = target_rank; + if (warp_role == WarpRole::kNVLReceivers) + rs_wr_buffer_ptr = buffer_ptrs[target_rank], ws_rr_buffer_ptr = buffer_ptrs[nvl_rank], rs_wr_rank = target_rank, + ws_rr_rank = nvl_rank; + + // Allocate buffers + auto nvl_channel_x = AsymBuffer(ws_rr_buffer_ptr, + num_max_nvl_chunked_recv_tokens * num_bytes_per_token, + NUM_MAX_NVL_PEERS, + channel_id, + num_channels, + rs_wr_rank) + .advance_also(rs_wr_buffer_ptr); + auto nvl_channel_prefix_start = + AsymBuffer(ws_rr_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) + .advance_also(rs_wr_buffer_ptr); + auto nvl_channel_prefix_end = AsymBuffer(ws_rr_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) + .advance_also(rs_wr_buffer_ptr); + auto nvl_channel_head = + AsymBuffer(rs_wr_buffer_ptr, 1, NUM_MAX_NVL_PEERS, channel_id, num_channels, ws_rr_rank).advance_also(ws_rr_buffer_ptr); + auto nvl_channel_tail = + AsymBuffer(ws_rr_buffer_ptr, 1, NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank).advance_also(rs_wr_buffer_ptr); + + // RDMA sender warp synchronization + // NOTES: `rdma_send_channel_tail` means the latest released tail + // NOTES: `rdma_send_channel_window` means the ongoing 32 transactions' status + __shared__ int rdma_send_channel_lock[kNumRDMARanks]; + __shared__ int rdma_send_channel_tail[kNumRDMARanks]; + __shared__ uint32_t rdma_send_channel_window[kNumRDMARanks]; + __shared__ nixlGpuXferStatusH rdma_put_xfer_status; + auto sync_rdma_sender_smem = []() { asm volatile("barrier.sync 0, %0;" ::"r"((kNumDispatchRDMASenderWarps + 1) * 32)); }; + + // TMA stuffs + extern __shared__ __align__(1024) uint8_t smem_tma_buffer[]; + auto tma_buffer = smem_tma_buffer + target_rank * kNumTMABytesPerWarp; + auto tma_mbarrier = reinterpret_cast(tma_buffer + num_bytes_per_token); + uint32_t tma_phase = 0; + if ((warp_role == WarpRole::kRDMAAndNVLForwarder or warp_role == WarpRole::kNVLReceivers) and elect_one_sync()) { + mbarrier_init(tma_mbarrier, 1); + fence_barrier_init(); + EP_DEVICE_ASSERT(num_bytes_per_token + sizeof(uint64_t) <= kNumTMABytesPerWarp); + } + __syncwarp(); + + // Forward warp synchronization + __shared__ volatile int forward_channel_head[NUM_MAX_NVL_PEERS][kNumRDMARanks]; + __shared__ volatile bool forward_channel_retired[NUM_MAX_NVL_PEERS]; + auto sync_forwarder_smem = []() { asm volatile("barrier.sync 1, %0;" ::"r"((NUM_MAX_NVL_PEERS + 1) * 32)); }; + + if (warp_role == WarpRole::kRDMASender) { + // Get tasks + int token_start_idx, token_end_idx; + get_channel_task_range(num_tokens, num_channels, channel_id, token_start_idx, token_end_idx); + + // Send number of tokens in this channel by `-value - 1` + EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS * 2 + 2 <= 32, "Invalid number of NVL peers"); + for (int dst_rdma_rank = warp_id; dst_rdma_rank < kNumRDMARanks; dst_rdma_rank += kNumDispatchRDMASenderWarps) { + auto dst_ptr = + dst_rdma_rank == rdma_rank ? rdma_channel_meta.recv_buffer(dst_rdma_rank) : rdma_channel_meta.send_buffer(dst_rdma_rank); + if (lane_id < NUM_MAX_NVL_PEERS) { + dst_ptr[lane_id] = + -(channel_id == 0 + ? 0 + : gbl_channel_prefix_matrix[(dst_rdma_rank * NUM_MAX_NVL_PEERS + lane_id) * num_channels + channel_id - 1]) - + 1; + } else if (lane_id < NUM_MAX_NVL_PEERS * 2) { + dst_ptr[lane_id] = + -gbl_channel_prefix_matrix[(dst_rdma_rank * NUM_MAX_NVL_PEERS + lane_id - NUM_MAX_NVL_PEERS) * num_channels + + channel_id] - + 1; + } else if (lane_id == NUM_MAX_NVL_PEERS * 2) { + dst_ptr[lane_id] = -(channel_id == 0 ? 0 : rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id - 1]) - 1; + } else if (lane_id == NUM_MAX_NVL_PEERS * 2 + 1) { + dst_ptr[lane_id] = -rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id] - 1; + } + __syncwarp(); + + // Issue RDMA for non-local ranks + if (dst_rdma_rank != rdma_rank) { + size_t src_offset = nixl_ctx.offset_get(reinterpret_cast(rdma_channel_meta.send_buffer(dst_rdma_rank))); + size_t dst_offset = nixl_ctx.offset_get(reinterpret_cast(rdma_channel_meta.recv_buffer(rdma_rank))); + size_t msg_size = sizeof(int) * (NUM_MAX_NVL_PEERS * 2 + 2); + int translated_rank = translate_dst_rdma_rank(dst_rdma_rank, nvl_rank); + nixlMemViewElem src_mdesc{nixl_ctx.local_mvh, 0, src_offset}; + nixlMemViewElem dst_mdesc{nixl_ctx.remote_mvh, (size_t)translated_rank, dst_offset}; + EP_DEVICE_ASSERT(nixlPut( + src_mdesc, dst_mdesc, msg_size, channel_id) == + NIXL_IN_PROG); + } + } + sync_rdma_sender_smem(); + + // Iterate over tokens and copy into buffer + int64_t token_idx; + int cached_rdma_channel_head = 0, global_rdma_tail_idx = 0; + auto send_buffer = lane_id == rdma_rank ? rdma_channel_data.recv_buffer(lane_id) : rdma_channel_data.send_buffer(lane_id); + for (token_idx = token_start_idx; token_idx < token_end_idx; ++token_idx) { + // Read RDMA rank existence + uint64_t is_token_in_rank_uint64 = 0; + if (lane_id < kNumRDMARanks) { + is_token_in_rank_uint64 = + __ldg(reinterpret_cast(is_token_in_rank + token_idx * num_ranks + lane_id * NUM_MAX_NVL_PEERS)); + global_rdma_tail_idx += (is_token_in_rank_uint64 != 0); + } + __syncwarp(); + + // Skip the token which does not belong to this warp + if ((token_idx - token_start_idx) % kNumDispatchRDMASenderWarps != warp_id) + continue; + auto rdma_tail_idx = is_token_in_rank_uint64 == 0 ? -1 : global_rdma_tail_idx - 1; + + // Wait the remote buffer to be released + auto start_time = clock64(); + while (is_token_in_rank_uint64 != 0 and rdma_tail_idx - cached_rdma_channel_head >= num_max_rdma_chunked_recv_tokens) { + cached_rdma_channel_head = static_cast(ld_volatile_global(rdma_channel_head.buffer(lane_id))); + + // Timeout check + if (clock64() - start_time >= NUM_TIMEOUT_CYCLES) { + printf("NixlEP dispatch RDMA sender timeout, channel: %d, RDMA: %d, nvl: %d, dst RDMA lane: %d, head: %d, tail: %d\n", + channel_id, + rdma_rank, + nvl_rank, + lane_id, + cached_rdma_channel_head, + rdma_tail_idx); + trap(); + } + } + __syncwarp(); + + // Store RDMA head for combine + if (lane_id < kNumRDMARanks and not kCachedMode) + send_rdma_head[token_idx * kNumRDMARanks + lane_id] = rdma_tail_idx; + + // Broadcast tails + SourceMeta src_meta; + int num_topk_ranks = 0, topk_ranks[kNumTopkRDMARanks]; + void* dst_send_buffers[kNumTopkRDMARanks]; + #pragma unroll + for (int i = 0, slot_idx; i < kNumRDMARanks; ++i) + if ((slot_idx = __shfl_sync(0xffffffff, rdma_tail_idx, i)) >= 0) { + slot_idx = slot_idx % num_max_rdma_chunked_recv_tokens; + topk_ranks[num_topk_ranks] = i; + auto recv_is_token_in_rank_uint64 = broadcast(is_token_in_rank_uint64, i); + auto recv_is_token_in_rank_values = reinterpret_cast(&recv_is_token_in_rank_uint64); + if (lane_id == num_topk_ranks) + src_meta = SourceMeta(rdma_rank, recv_is_token_in_rank_values); + dst_send_buffers[num_topk_ranks++] = + reinterpret_cast(broadcast(send_buffer, i)) + slot_idx * num_bytes_per_token; + } + EP_DEVICE_ASSERT(num_topk_ranks <= kNumTopkRDMARanks); + + // Copy `x` into symmetric send buffer + auto st_broadcast = [=](const int key, const int4& value) { + #pragma unroll + for (int j = 0; j < num_topk_ranks; ++j) + st_na_global(reinterpret_cast(dst_send_buffers[j]) + key, value); + }; + UNROLLED_WARP_COPY(5, lane_id, hidden_int4, 0, x + token_idx * hidden_int4, ld_nc_global, st_broadcast); + #pragma unroll + for (int i = 0; i < num_topk_ranks; ++i) + dst_send_buffers[i] = reinterpret_cast(dst_send_buffers[i]) + hidden_int4; + + // Copy `x_scales` into symmetric send buffer + #pragma unroll + for (int i = lane_id; i < num_scales; i += 32) { + auto offset = token_idx * scale_token_stride + i * scale_hidden_stride; + auto value = ld_nc_global(x_scales + offset); + #pragma unroll + for (int j = 0; j < num_topk_ranks; ++j) + st_na_global(reinterpret_cast(dst_send_buffers[j]) + i, value); + } + #pragma unroll + for (int i = 0; i < num_topk_ranks; ++i) + dst_send_buffers[i] = reinterpret_cast(dst_send_buffers[i]) + num_scales; + + // Copy source metadata into symmetric send buffer + if (lane_id < num_topk_ranks) + st_na_global(reinterpret_cast(dst_send_buffers[lane_id]), src_meta); + #pragma unroll + for (int i = 0; i < num_topk_ranks; ++i) + dst_send_buffers[i] = reinterpret_cast(dst_send_buffers[i]) + 1; + + // Copy `topk_idx` and `topk_weights` into symmetric send buffer + #pragma unroll + for (int i = lane_id; i < num_topk * num_topk_ranks; i += 32) { + auto rank_idx = i / num_topk, copy_idx = i % num_topk; + auto idx_value = static_cast(ld_nc_global(topk_idx + token_idx * num_topk + copy_idx)); + auto weight_value = ld_nc_global(topk_weights + token_idx * num_topk + copy_idx); + st_na_global(reinterpret_cast(dst_send_buffers[rank_idx]) + copy_idx, idx_value); + st_na_global(reinterpret_cast(dst_send_buffers[rank_idx]) + num_topk + copy_idx, weight_value); + } + __syncwarp(); + + // Release the transaction in the window + if (is_token_in_rank_uint64 != 0) { + // Acquire lock first + acquire_lock(rdma_send_channel_lock + lane_id); + auto latest_tail = rdma_send_channel_tail[lane_id]; + auto offset = rdma_tail_idx - latest_tail; + while (offset >= 32) { + release_lock(rdma_send_channel_lock + lane_id); + acquire_lock(rdma_send_channel_lock + lane_id); + latest_tail = rdma_send_channel_tail[lane_id]; + offset = rdma_tail_idx - latest_tail; + } + + // Release the transaction slot + // Add the bit and move the ones if possible + auto window = rdma_send_channel_window[lane_id] | (1u << offset); + if (offset == 0) { + auto num_empty_slots = (~window) == 0 ? 32 : __ffs(~window) - 1; + st_release_cta(rdma_send_channel_tail + lane_id, latest_tail + num_empty_slots); + window >>= num_empty_slots; + } + rdma_send_channel_window[lane_id] = window; + + // Release lock + release_lock(rdma_send_channel_lock + lane_id); + } + __syncwarp(); + } + } else if (warp_role == WarpRole::kRDMASenderCoordinator) { + // NOTES: in case of splitting, the issued put at the end of the buffer + EP_DEVICE_ASSERT(num_max_rdma_chunked_recv_tokens % num_max_rdma_chunked_send_tokens == 0); + + // Clean shared memory + EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA ranks"); + (lane_id < kNumRDMARanks) ? (rdma_send_channel_lock[lane_id] = 0) : 0; + (lane_id < kNumRDMARanks) ? (rdma_send_channel_tail[lane_id] = 0) : 0; + (lane_id < kNumRDMARanks) ? (rdma_send_channel_window[lane_id] = 0) : 0; + + // Synchronize shared memory + sync_rdma_sender_smem(); + + // Get number of tokens to send for each RDMA rank + int num_tokens_to_send = 0; + if (lane_id < kNumRDMARanks) { + num_tokens_to_send = rdma_channel_prefix_matrix[lane_id * num_channels + channel_id]; + if (channel_id > 0) + num_tokens_to_send -= rdma_channel_prefix_matrix[lane_id * num_channels + channel_id - 1]; + } + + // Iterate all RDMA ranks + int last_issued_tail = 0; + auto start_time = clock64(); + while (__any_sync(0xffffffff, num_tokens_to_send > 0)) { + // Timeout check + if (clock64() - start_time > NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) { + printf("NixlEP RDMA sender coordinator timeout, channel: %d, IB: %d, nvl %d, dst IB: %d, tail: %d, remaining: %d\n", + channel_id, + rdma_rank, + nvl_rank, + lane_id, + last_issued_tail, + num_tokens_to_send); + trap(); + } + + // TODO: try thread-level `put_nbi`? + for (int i = 0, synced_num_tokens_to_send; i < kNumRDMARanks; ++i) { + // To mitigate incast congestion, shuffle the starting index of target rank for different ranks and channels + int dst_rdma_rank = (i + channel_id + rdma_rank) % kNumRDMARanks; + synced_num_tokens_to_send = __shfl_sync(0xffffffff, num_tokens_to_send, dst_rdma_rank); + if (synced_num_tokens_to_send == 0) + continue; + + // Read the latest progress + // NOTES: `rdma_send_channel_tail` does not need to be protected by lock + auto processed_tail = + __shfl_sync(0xffffffff, ld_acquire_cta(const_cast(rdma_send_channel_tail + dst_rdma_rank)), 0); + auto synced_last_issued_tail = __shfl_sync(0xffffffff, last_issued_tail, dst_rdma_rank); + auto num_tokens_processed = processed_tail - synced_last_issued_tail; + if (num_tokens_processed != synced_num_tokens_to_send and num_tokens_processed < num_max_rdma_chunked_send_tokens) + continue; + + // Issue RDMA send + auto num_tokens_to_issue = min(num_tokens_processed, num_max_rdma_chunked_send_tokens); + EP_DEVICE_ASSERT(num_tokens_to_issue >= 0 and num_tokens_to_issue <= synced_num_tokens_to_send); + if (dst_rdma_rank != rdma_rank) { + auto dst_slot_idx = synced_last_issued_tail % num_max_rdma_chunked_recv_tokens; + EP_DEVICE_ASSERT(dst_slot_idx + num_tokens_to_issue <= num_max_rdma_chunked_recv_tokens); + const size_t num_bytes_per_msg = num_bytes_per_token * num_tokens_to_issue; + const auto dst_ptr = + reinterpret_cast(rdma_channel_data.recv_buffer(rdma_rank) + dst_slot_idx * num_bytes_per_token); + const auto src_ptr = + reinterpret_cast(rdma_channel_data.send_buffer(dst_rdma_rank) + dst_slot_idx * num_bytes_per_token); + size_t src_offset = nixl_ctx.offset_get(src_ptr); + size_t dst_offset = nixl_ctx.offset_get(dst_ptr); + + int translated_dst = translate_dst_rdma_rank(dst_rdma_rank, nvl_rank); + nixlMemViewElem src_mdesc{nixl_ctx.local_mvh, 0, src_offset}; + nixlMemViewElem dst_mdesc{nixl_ctx.remote_mvh, (size_t)translated_dst, dst_offset}; + EP_DEVICE_ASSERT(nixlPut( + src_mdesc, dst_mdesc, num_bytes_per_msg, channel_id) == + NIXL_IN_PROG); + } else { + // Lighter fence for local RDMA rank + memory_fence(); + } + __syncwarp(); + + // Update tails + if (lane_id == dst_rdma_rank) { + last_issued_tail += num_tokens_to_issue; + num_tokens_to_send -= num_tokens_to_issue; + if (dst_rdma_rank == rdma_rank) { + atomicAdd(reinterpret_cast(rdma_channel_tail.buffer(dst_rdma_rank)), static_cast(num_tokens_to_issue)); + } else { + size_t tail_counter_offset = nixl_ctx.offset_get(reinterpret_cast(rdma_channel_tail.buffer(rdma_rank))); + int translated_dst_tail = translate_dst_rdma_rank(dst_rdma_rank, nvl_rank); + + nixlMemViewElem tail_mdesc{nixl_ctx.remote_mvh, (size_t)translated_dst_tail, tail_counter_offset}; + EP_DEVICE_ASSERT(nixlAtomicAdd( + num_tokens_to_issue, tail_mdesc, channel_id, 0) == + NIXL_IN_PROG); + } + } + __syncwarp(); + } + } + } else if (warp_role == WarpRole::kRDMAAndNVLForwarder) { + // RDMA consumers and NVL producers + const auto dst_nvl_rank = target_rank; + + // Wait counters to arrive + int num_tokens_to_recv_from_rdma = 0, src_rdma_channel_prefix = 0; + EP_DEVICE_ASSERT(kNumRDMARanks <= 32); + auto start_time = clock64(); + if (lane_id < kNumRDMARanks) { + while (true) { + auto meta_0 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + dst_nvl_rank); + auto meta_1 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + NUM_MAX_NVL_PEERS + dst_nvl_rank); + auto meta_2 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + NUM_MAX_NVL_PEERS * 2); + auto meta_3 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + NUM_MAX_NVL_PEERS * 2 + 1); + if (meta_0 < 0 and meta_1 < 0 and meta_2 < 0 and meta_3 < 0) { + // Notify NVL ranks + int start_sum = -meta_0 - 1, end_sum = -meta_1 - 1; + EP_DEVICE_ASSERT(start_sum >= 0 and end_sum >= 0 and end_sum >= start_sum); + st_relaxed_sys_global(nvl_channel_prefix_start.buffer() + lane_id, -start_sum - 1); + st_relaxed_sys_global(nvl_channel_prefix_end.buffer() + lane_id, -end_sum - 1); + + // Save RDMA channel received token count + src_rdma_channel_prefix = -meta_2 - 1; + auto src_rdma_channel_prefix_1 = -meta_3 - 1; + num_tokens_to_recv_from_rdma = src_rdma_channel_prefix_1 - src_rdma_channel_prefix; + if (not kCachedMode) + recv_rdma_channel_prefix_matrix[lane_id * num_channels + channel_id] = src_rdma_channel_prefix_1; + src_rdma_channel_prefix += lane_id == 0 ? 0 : recv_rdma_rank_prefix_sum[lane_id - 1]; + EP_DEVICE_ASSERT(num_tokens_to_recv_from_rdma >= 0); + break; + } + + // Timeout check + if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { + printf( + "NixlEP dispatch forwarder timeout (RDMA meta), channel: %d, RDMA: %d, nvl: %d, src RDMA lane: %d, dst NVL: %d, meta: %d, %d, %d, %d\n", + channel_id, + rdma_rank, + nvl_rank, + lane_id, + dst_nvl_rank, + meta_0, + meta_1, + meta_2, + meta_3); + trap(); + } + } + } + __syncwarp(); + + // Shift cached head + send_nvl_head += src_rdma_channel_prefix * NUM_MAX_NVL_PEERS + dst_nvl_rank; + + // Wait shared memory to be cleaned + sync_forwarder_smem(); + + // Forward tokens from RDMA buffer + // NOTES: always start from the local rank + int src_rdma_rank = sm_id % kNumRDMARanks; + int cached_rdma_channel_head = 0, cached_rdma_channel_tail = 0; + int cached_nvl_channel_head = 0, cached_nvl_channel_tail = 0, rdma_nvl_token_idx = 0; + while (__any_sync(0xffffffff, num_tokens_to_recv_from_rdma > 0)) { + // Check destination queue emptiness, or wait a buffer to be released + start_time = clock64(); + while (true) { + const int num_used_slots = cached_nvl_channel_tail - cached_nvl_channel_head; + if (num_max_nvl_chunked_recv_tokens - num_used_slots >= num_max_nvl_chunked_send_tokens) + break; + cached_nvl_channel_head = __shfl_sync(0xffffffffu, ld_volatile_global(nvl_channel_head.buffer()), 0); + + // Timeout check + if (elect_one_sync() and clock64() - start_time > NUM_TIMEOUT_CYCLES) { + printf( + "NixlEP dispatch forwarder timeout (NVL check), channel: %d, RDMA: %d, nvl: %d, dst NVL: %d, head: %d, tail: %d\n", + channel_id, + rdma_rank, + nvl_rank, + dst_nvl_rank, + ld_volatile_global(nvl_channel_head.buffer()), + cached_nvl_channel_tail); + trap(); + } + } + + // Find next source RDMA rank (round-robin) + start_time = clock64(); + while (true) { + src_rdma_rank = (src_rdma_rank + 1) % kNumRDMARanks; + if (__shfl_sync(0xffffffff, num_tokens_to_recv_from_rdma, src_rdma_rank) > 0) { + if (lane_id == src_rdma_rank and cached_rdma_channel_head == cached_rdma_channel_tail) + { + cached_rdma_channel_tail = static_cast(ld_acquire_sys_global(rdma_channel_tail.buffer(src_rdma_rank))); + } + if (__shfl_sync(0xffffffff, cached_rdma_channel_tail > cached_rdma_channel_head, src_rdma_rank)) + { + break; + } + } + + // Timeout check + if (clock64() - start_time > NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) { + printf( + "NixlEP dispatch forwarder timeout (RDMA check), channel: %d, RDMA: %d, nvl: %d, dst NVL: %d, src RDMA lane: %d, " + "head: %d, tail: %d, expected: %d\n", + channel_id, + rdma_rank, + nvl_rank, + dst_nvl_rank, + lane_id, + cached_rdma_channel_head, + cached_rdma_channel_tail, + num_tokens_to_recv_from_rdma); + trap(); + } + } + auto src_rdma_head = __shfl_sync(0xffffffff, cached_rdma_channel_head, src_rdma_rank); + auto src_rdma_tail = __shfl_sync(0xffffffff, cached_rdma_channel_tail, src_rdma_rank); + + // Iterate over every token from the RDMA buffer + for (int i = src_rdma_head, num_tokens_sent = 0; i < src_rdma_tail; ++i) { + auto rdma_slot_idx = i % num_max_rdma_chunked_recv_tokens; + auto shifted = rdma_channel_data.recv_buffer(src_rdma_rank) + rdma_slot_idx * num_bytes_per_token; + auto src_meta = ld_nc_global(reinterpret_cast(shifted + hidden_bytes + scale_bytes)); + + lane_id == src_rdma_rank ? (num_tokens_to_recv_from_rdma -= 1) : 0; + bool is_in_dst_nvl_rank = src_meta.is_token_in_nvl_rank(dst_nvl_rank); + if (lane_id == src_rdma_rank) { + auto cached_head = is_in_dst_nvl_rank ? rdma_nvl_token_idx : -1; + rdma_nvl_token_idx += is_in_dst_nvl_rank; + if (not kCachedMode) + send_nvl_head[i * NUM_MAX_NVL_PEERS] = cached_head; + } + if (not is_in_dst_nvl_rank) + continue; + + // Get an empty slot + int dst_slot_idx = (cached_nvl_channel_tail++) % num_max_nvl_chunked_recv_tokens; + auto dst_shifted = nvl_channel_x.buffer() + dst_slot_idx * num_bytes_per_token; + + // Copy data + if (elect_one_sync()) { + tma_load_1d(tma_buffer, shifted, tma_mbarrier, num_bytes_per_token, false); + mbarrier_arrive_and_expect_tx(tma_mbarrier, num_bytes_per_token); + } + __syncwarp(); + mbarrier_wait(tma_mbarrier, tma_phase); + if (elect_one_sync()) + tma_store_1d(tma_buffer, dst_shifted, num_bytes_per_token); + __syncwarp(); + + // In case of insufficient NVL buffers, early stopping + if ((++num_tokens_sent) == num_max_nvl_chunked_send_tokens) + src_rdma_tail = i + 1; + + // Wait TMA to be finished + tma_store_wait<0>(); + __syncwarp(); + } + + // Sync head index + if (lane_id == src_rdma_rank) + { + forward_channel_head[dst_nvl_rank][src_rdma_rank] = (cached_rdma_channel_head = src_rdma_tail); + } + + // Move tail index + __syncwarp(); + if (elect_one_sync()) + st_release_sys_global(nvl_channel_tail.buffer(), cached_nvl_channel_tail); + } + + // Retired + __syncwarp(); + if (elect_one_sync()) + { + forward_channel_retired[dst_nvl_rank] = true; + } + } else if (warp_role == WarpRole::kForwarderCoordinator) { + // Extra warps for forwarder coordinator should exit directly + if (target_rank > 0) + return; + + // Forward warp coordinator + EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); + + // Clean shared memory + EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); + #pragma unroll + for (int i = lane_id; i < kNumRDMARanks * NUM_MAX_NVL_PEERS; i += 32) + forward_channel_head[i % NUM_MAX_NVL_PEERS][i / NUM_MAX_NVL_PEERS] = 0; + if (lane_id < NUM_MAX_NVL_PEERS) + forward_channel_retired[lane_id] = false; + sync_forwarder_smem(); + + int last_head = 0, target_rdma = lane_id < kNumRDMARanks ? lane_id : 0; + while (true) { + // Find minimum head + int min_head = std::numeric_limits::max(); + #pragma unroll + for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) + if (not forward_channel_retired[i]) + min_head = min(min_head, forward_channel_head[i][target_rdma]); + if (__all_sync(0xffffffff, min_head == std::numeric_limits::max())) + break; + + // Update remote head + if (min_head != std::numeric_limits::max() and min_head >= last_head + num_max_rdma_chunked_send_tokens and + lane_id < kNumRDMARanks) { + if(lane_id == rdma_rank){ + atomicAdd(reinterpret_cast(rdma_channel_head.buffer(rdma_rank)), static_cast(min_head - last_head)); + } else { + size_t head_counter_offset = nixl_ctx.offset_get(reinterpret_cast(rdma_channel_head.buffer(rdma_rank))); + int translated_dst_head = translate_dst_rdma_rank(lane_id, nvl_rank); + nixlMemViewElem head_mdesc{nixl_ctx.remote_mvh, (size_t)translated_dst_head, head_counter_offset}; + EP_DEVICE_ASSERT(nixlAtomicAdd( + min_head - last_head, head_mdesc, channel_id, 0) == NIXL_IN_PROG); + } + last_head = min_head; + } + + // Nanosleep and let other warps work + __nanosleep(NUM_WAIT_NANOSECONDS); + } + } else { + // NVL consumers + // Retrieve rank offset from barrier results (each lane's register stores an RDMA rank) + int src_nvl_rank = target_rank, total_offset = 0; + const int local_expert_begin = rank * (num_experts / num_ranks); + const int local_expert_end = local_expert_begin + (num_experts / num_ranks); + + EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); + if (lane_id < kNumRDMARanks and lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank > 0) + total_offset = recv_gbl_rank_prefix_sum[lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank - 1]; + + // Receive channel offsets + int start_offset = 0, end_offset = 0, num_tokens_to_recv; + auto start_time = clock64(); + while (lane_id < kNumRDMARanks) { + start_offset = ld_volatile_global(nvl_channel_prefix_start.buffer() + lane_id); + end_offset = ld_volatile_global(nvl_channel_prefix_end.buffer() + lane_id); + if (start_offset < 0 and end_offset < 0) { + start_offset = -start_offset - 1, end_offset = -end_offset - 1; + total_offset += start_offset; + break; + } + + // Timeout check + if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { + printf( + "NixlEP dispatch NVL receiver timeout, channel: %d, RDMA: %d, nvl: %d, src RDMA: %d, src nvl: %d, start: %d, end: %d\n", + channel_id, + rdma_rank, + nvl_rank, + lane_id, + src_nvl_rank, + start_offset, + end_offset); + trap(); + } + } + num_tokens_to_recv = warp_reduce_sum(end_offset - start_offset); + + // Save for combine usage + if (lane_id < kNumRDMARanks and not kCachedMode) + recv_gbl_channel_prefix_matrix[(lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank) * num_channels + channel_id] = total_offset; + __syncwarp(); + + int cached_channel_head_idx = 0, cached_channel_tail_idx = 0; + while (num_tokens_to_recv > 0) { + // Check channel status by lane 0 + start_time = clock64(); + while (true) { + // Ready to copy + if (cached_channel_head_idx != cached_channel_tail_idx) + break; + cached_channel_tail_idx = __shfl_sync(0xffffffff, ld_acquire_sys_global(nvl_channel_tail.buffer()), 0); + + // Timeout check + if (elect_one_sync() and clock64() - start_time > NUM_TIMEOUT_CYCLES) { + printf("NixlEP dispatch NVL receiver timeout, channel: %d, RDMA: %d, nvl: %d, src NVL: %d, head: %d, tail: %d\n", + channel_id, + rdma_rank, + nvl_rank, + src_nvl_rank, + cached_channel_head_idx, + cached_channel_tail_idx); + trap(); + } + } + + // Copy data + int num_recv_tokens = cached_channel_tail_idx - cached_channel_head_idx; + for (int chunk_idx = 0; chunk_idx < num_recv_tokens; ++chunk_idx, --num_tokens_to_recv) { + int token_idx_in_buffer = (cached_channel_head_idx++) % num_max_nvl_chunked_recv_tokens; + auto shifted = nvl_channel_x.buffer() + token_idx_in_buffer * num_bytes_per_token; + auto meta = ld_nc_global(reinterpret_cast(shifted + hidden_bytes + scale_bytes)); + + int64_t recv_token_idx = __shfl_sync(0xffffffff, total_offset, meta.src_rdma_rank); + (lane_id == meta.src_rdma_rank) ? (total_offset += 1) : 0; + + bool scale_aligned = (scale_bytes % 16 == 0); + auto tma_load_bytes = hidden_bytes + (scale_aligned ? scale_bytes : 0); + + // Copy data + if (elect_one_sync()) { + tma_load_1d(tma_buffer, shifted, tma_mbarrier, tma_load_bytes); + mbarrier_arrive_and_expect_tx(tma_mbarrier, tma_load_bytes); + } + __syncwarp(); + mbarrier_wait(tma_mbarrier, tma_phase); + if (elect_one_sync()) { + tma_store_1d(tma_buffer, recv_x + recv_token_idx * hidden_int4, hidden_bytes, false); + if (scale_aligned) + tma_store_1d(tma_buffer + hidden_bytes, recv_x_scales + recv_token_idx * num_scales, scale_bytes, false); + } + __syncwarp(); + shifted += hidden_bytes; + + // Copy scales + // TODO: make it as templated + if (not scale_aligned) { + UNROLLED_WARP_COPY(1, + lane_id, + num_scales, + recv_x_scales + recv_token_idx * num_scales, + reinterpret_cast(shifted), + ld_nc_global, + st_na_global); + } + shifted += scale_bytes; + + // Copy source meta + if (not kCachedMode and elect_one_sync()) + st_na_global(recv_src_meta + recv_token_idx, meta); + shifted += sizeof(SourceMeta); + + // Copy `topk_idx` and `topk_weights` + if (lane_id < num_topk) { + // Read + auto idx_value = static_cast(ld_nc_global(reinterpret_cast(shifted) + lane_id)); + auto weight_value = ld_nc_global(reinterpret_cast(shifted + sizeof(int) * num_topk) + lane_id); + auto recv_idx = recv_token_idx * num_topk + lane_id; + + // Transform and write + idx_value = (idx_value >= local_expert_begin and idx_value < local_expert_end) ? idx_value - local_expert_begin : -1; + weight_value = idx_value >= 0 ? weight_value : 0.0f; + st_na_global(recv_topk_idx + recv_idx, idx_value); + st_na_global(recv_topk_weights + recv_idx, weight_value); + } + + // Wait TMA to be finished + tma_store_wait<0>(); + __syncwarp(); + } + + // Move queue + if (elect_one_sync()) + st_relaxed_sys_global(nvl_channel_head.buffer(), cached_channel_head_idx); + } + } +} +void dispatch(void* recv_x, + float* recv_x_scales, + topk_idx_t* recv_topk_idx, + float* recv_topk_weights, + void* recv_src_meta, + const void* x, + const float* x_scales, + const topk_idx_t* topk_idx, + const float* topk_weights, + int* send_rdma_head, + int* send_nvl_head, + int* recv_rdma_channel_prefix_matrix, + int* recv_gbl_channel_prefix_matrix, + const int* rdma_channel_prefix_matrix, + const int* recv_rdma_rank_prefix_sum, + const int* gbl_channel_prefix_matrix, + const int* recv_gbl_rank_prefix_sum, + const bool* is_token_in_rank, + int num_tokens, + int hidden_int4, + int num_scales, + int num_topk, + int num_experts, + int scale_token_stride, + int scale_hidden_stride, + void* rdma_buffer_ptr, + int num_max_rdma_chunked_send_tokens, + int num_max_rdma_chunked_recv_tokens, + void** buffer_ptrs, + int num_max_nvl_chunked_send_tokens, + int num_max_nvl_chunked_recv_tokens, + int rank, + int num_ranks, + bool is_cached_dispatch, + cudaStream_t stream, + int num_channels, + bool low_latency_mode, + gpu_nixl_ctx nixl_ctx) { + constexpr int kNumDispatchRDMASenderWarps = 7; + constexpr int kNumTMABytesPerWarp = 16384; + constexpr int smem_size = kNumTMABytesPerWarp * NUM_MAX_NVL_PEERS; + + // Make sure never OOB + EP_HOST_ASSERT(static_cast(num_scales) * scale_hidden_stride < std::numeric_limits::max()); + +#define DISPATCH_LAUNCH_CASE(num_rdma_ranks) \ + { \ + auto dispatch_func = low_latency_mode \ + ? (is_cached_dispatch ? dispatch \ + : dispatch) \ + : (is_cached_dispatch ? dispatch \ + : dispatch); \ + SET_SHARED_MEMORY_FOR_TMA(dispatch_func); \ + LAUNCH_KERNEL(&cfg, \ + dispatch_func, \ + reinterpret_cast(recv_x), \ + recv_x_scales, \ + recv_topk_idx, \ + recv_topk_weights, \ + reinterpret_cast(recv_src_meta), \ + reinterpret_cast(x), \ + x_scales, \ + topk_idx, \ + topk_weights, \ + send_rdma_head, \ + send_nvl_head, \ + recv_rdma_channel_prefix_matrix, \ + recv_gbl_channel_prefix_matrix, \ + rdma_channel_prefix_matrix, \ + recv_rdma_rank_prefix_sum, \ + gbl_channel_prefix_matrix, \ + recv_gbl_rank_prefix_sum, \ + is_token_in_rank, \ + num_tokens, \ + hidden_int4, \ + num_scales, \ + num_topk, \ + num_experts, \ + scale_token_stride, \ + scale_hidden_stride, \ + rdma_buffer_ptr, \ + num_max_rdma_chunked_send_tokens, \ + num_max_rdma_chunked_recv_tokens, \ + buffer_ptrs, \ + num_max_nvl_chunked_send_tokens, \ + num_max_nvl_chunked_recv_tokens, \ + rank, \ + num_ranks, \ + nixl_ctx); \ + } \ + break + + EP_HOST_ASSERT((topk_idx == nullptr) == (topk_weights == nullptr)); + EP_HOST_ASSERT((recv_topk_idx == nullptr) == (recv_topk_weights == nullptr)); + + SETUP_LAUNCH_CONFIG(num_channels * 2, (kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NVL_PEERS) * 32, stream); + SWITCH_RDMA_RANKS(DISPATCH_LAUNCH_CASE); +#undef DISPATCH_LAUNCH_CASE +} + +template +__global__ void cached_notify(const int rdma_clean_offset, + const int rdma_num_int_clean, + const int nvl_clean_offset, + const int nvl_num_int_clean, + int* combined_rdma_head, + int num_combined_tokens, + int num_channels, + const int* rdma_channel_prefix_matrix, + const int* rdma_rank_prefix_sum, + int* combined_nvl_head, + void* rdma_buffer_ptr, + void** buffer_ptrs, + int** barrier_signal_ptrs, + int rank, + int num_ranks, + bool is_cached_dispatch, + gpu_nixl_ctx nixl_ctx) { + auto sm_id = static_cast(blockIdx.x); + auto thread_id = static_cast(threadIdx.x); + auto num_threads = static_cast(blockDim.x); + auto num_warps = num_threads / 32; + auto warp_id = thread_id / 32; + auto lane_id = get_lane_id(); + + auto nvl_rank = rank % NUM_MAX_NVL_PEERS; + auto num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; + + // Using two SMs, which clean the RDMA/NVL buffer respectively + if (sm_id == 0) { + __syncthreads(); + + // Barrier for RDMA + if (warp_id == 1) { + nixl_barrier_send_warp(nixl_ctx, num_channels); + if (lane_id == 0) + nixl_barrier_wait(nixl_ctx, num_channels); + } + + // Barrier for NVL + barrier_block(barrier_signal_ptrs, nvl_rank); + + // Clean RDMA buffer + auto rdma_buffer_ptr_int = static_cast(rdma_buffer_ptr); + #pragma unroll + for (int i = thread_id; i < rdma_num_int_clean; i += num_threads) + rdma_buffer_ptr_int[rdma_clean_offset + i] = 0; + + // Clean NVL buffer + auto nvl_buffer_ptr_int = static_cast(buffer_ptrs[nvl_rank]); + #pragma unroll + for (int i = thread_id; i < nvl_num_int_clean; i += num_threads) + nvl_buffer_ptr_int[nvl_clean_offset + i] = 0; + __syncthreads(); + + // Barrier again + if (warp_id == 1) { + nixl_barrier_send_warp(nixl_ctx, num_channels); + if (lane_id == 0) + nixl_barrier_wait(nixl_ctx, num_channels); + } + barrier_block(barrier_signal_ptrs, nvl_rank); + } else if (sm_id == 1) { + if (is_cached_dispatch) + return; + + EP_DEVICE_ASSERT(num_warps >= num_channels); + EP_DEVICE_ASSERT(num_rdma_ranks <= 32); + + // Iterate in reverse order + if (lane_id < num_rdma_ranks and warp_id < num_channels) { + int token_start_idx, token_end_idx; + get_channel_task_range(num_combined_tokens, num_channels, warp_id, token_start_idx, token_end_idx); + + // NOTES: `1 << 25` is a heuristic large number + int last_head = 1 << 25; + for (int token_idx = token_end_idx - 1; token_idx >= token_start_idx; --token_idx) { + auto current_head = __ldg(combined_rdma_head + token_idx * num_rdma_ranks + lane_id); + if (current_head < 0) { + combined_rdma_head[token_idx * num_rdma_ranks + lane_id] = -last_head - 1; + } else { + last_head = current_head; + } + } + } + } else { + if (is_cached_dispatch) + return; + + EP_DEVICE_ASSERT(num_warps >= num_channels); + EP_DEVICE_ASSERT(rdma_channel_prefix_matrix != nullptr and rdma_rank_prefix_sum != nullptr); + EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 32, "Too many NVL peers"); + + if (warp_id < num_channels) { + constexpr int tma_batch_size = kNumTMABytesPerWarp - sizeof(uint64_t); + constexpr int num_bytes_per_token = sizeof(int) * NUM_MAX_NVL_PEERS; + constexpr int num_tokens_per_batch = tma_batch_size / num_bytes_per_token; + EP_STATIC_ASSERT(num_bytes_per_token % 16 == 0, "num_bytes_per_token should be divisible by 16"); + + // TMA stuffs + extern __shared__ __align__(1024) uint8_t smem_tma_buffer[]; + auto tma_buffer = smem_tma_buffer + warp_id * kNumTMABytesPerWarp; + auto tma_mbarrier = reinterpret_cast(tma_buffer + tma_batch_size); + uint32_t tma_phase = 0; + if (elect_one_sync()) { + mbarrier_init(tma_mbarrier, 1); + fence_barrier_init(); + } + __syncwarp(); + + for (int dst_rdma_rank = sm_id - 2; dst_rdma_rank < num_rdma_ranks; dst_rdma_rank += num_channels * 2 - 2) { + // Iterate in reverse order + int token_start_idx = warp_id == 0 ? 0 : rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + warp_id - 1]; + int token_end_idx = rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + warp_id]; + int shift = dst_rdma_rank == 0 ? 0 : rdma_rank_prefix_sum[dst_rdma_rank - 1]; + token_start_idx += shift, token_end_idx += shift; + + // NOTES: `1 << 25` is a heuristic large number + int last_head = 1 << 25; + for (int batch_end_idx = token_end_idx; batch_end_idx > token_start_idx; batch_end_idx -= num_tokens_per_batch) { + auto batch_start_idx = max(token_start_idx, batch_end_idx - num_tokens_per_batch); + + if (elect_one_sync()) { + tma_load_1d(tma_buffer, + combined_nvl_head + batch_start_idx * NUM_MAX_NVL_PEERS, + tma_mbarrier, + (batch_end_idx - batch_start_idx) * num_bytes_per_token); + mbarrier_arrive_and_expect_tx(tma_mbarrier, (batch_end_idx - batch_start_idx) * num_bytes_per_token); + } + mbarrier_wait(tma_mbarrier, tma_phase); + __syncwarp(); + + for (int token_idx = batch_end_idx - 1; token_idx >= batch_start_idx; --token_idx) { + if (lane_id < NUM_MAX_NVL_PEERS) { + auto current_head = + reinterpret_cast(tma_buffer)[(token_idx - batch_start_idx) * NUM_MAX_NVL_PEERS + lane_id]; + if (current_head < 0) { + reinterpret_cast(tma_buffer)[(token_idx - batch_start_idx) * NUM_MAX_NVL_PEERS + lane_id] = + -last_head - 1; + } else { + last_head = current_head; + } + } + } + tma_store_fence(); + __syncwarp(); + + if (elect_one_sync()) + tma_store_1d(tma_buffer, + combined_nvl_head + batch_start_idx * NUM_MAX_NVL_PEERS, + (batch_end_idx - batch_start_idx) * num_bytes_per_token); + tma_store_wait<0>(); + __syncwarp(); + } + } + } + } +} + +void cached_notify(int hidden_int4, + int num_scales, + int num_topk_idx, + int num_topk_weights, + int num_ranks, + int num_channels, + int num_combined_tokens, + int* combined_rdma_head, + const int* rdma_channel_prefix_matrix, + const int* rdma_rank_prefix_sum, + int* combined_nvl_head, + void* rdma_buffer_ptr, + int num_max_rdma_chunked_recv_tokens, + void** buffer_ptrs, + int num_max_nvl_chunked_recv_tokens, + int** barrier_signal_ptrs, + int rank, + cudaStream_t stream, + int64_t num_rdma_bytes, + int64_t num_nvl_bytes, + bool is_cached_dispatch, + bool low_latency_mode, + gpu_nixl_ctx nixl_ctx) { + const int num_threads = std::max(128, 32 * num_channels); + const int num_warps = num_threads / 32; + const auto num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; + const int kNumTMABytesPerWarp = 8192; + const int smem_size = kNumTMABytesPerWarp * num_warps; + + // Get clean meta + auto rdma_clean_meta = get_rdma_clean_meta( + hidden_int4, num_scales, num_topk_idx, num_topk_weights, num_rdma_ranks, num_max_rdma_chunked_recv_tokens, num_channels); + auto nvl_clean_meta = get_nvl_clean_meta(hidden_int4, + num_scales, + num_topk_idx, + num_topk_weights, + num_rdma_ranks, + NUM_MAX_NVL_PEERS, + num_max_nvl_chunked_recv_tokens, + num_channels, + is_cached_dispatch); + EP_HOST_ASSERT((rdma_clean_meta.first + rdma_clean_meta.second) * sizeof(int) <= num_rdma_bytes); + EP_HOST_ASSERT((nvl_clean_meta.first + nvl_clean_meta.second) * sizeof(int) <= num_nvl_bytes); + EP_HOST_ASSERT(num_rdma_bytes < std::numeric_limits::max()); + EP_HOST_ASSERT(num_nvl_bytes < std::numeric_limits::max()); + EP_HOST_ASSERT(num_channels * 2 > 3); + + // Launch kernel + auto cached_notify_func = low_latency_mode ? cached_notify : cached_notify; + SETUP_LAUNCH_CONFIG(num_channels * 2, num_threads, stream); + SET_SHARED_MEMORY_FOR_TMA(cached_notify_func); + LAUNCH_KERNEL(&cfg, + cached_notify_func, + rdma_clean_meta.first, + rdma_clean_meta.second, + nvl_clean_meta.first, + nvl_clean_meta.second, + combined_rdma_head, + num_combined_tokens, + num_channels, + rdma_channel_prefix_matrix, + rdma_rank_prefix_sum, + combined_nvl_head, + rdma_buffer_ptr, + buffer_ptrs, + barrier_signal_ptrs, + rank, + num_ranks, + is_cached_dispatch, + nixl_ctx); +} + +template +__device__ int combine_token(bool is_token_in_rank, + int head_idx, + int lane_id, + int hidden_int4, + int num_topk, + int4* combined_row, + float* combined_topk_weights, + const int4* bias_0_int4, + const int4* bias_1_int4, + int num_max_recv_tokens, + const GetAddrFn& get_addr_fn, + const ReceiveTWFn& recv_tw_fn, + uint8_t* smem_ptr, + uint32_t (&tma_phase)[kNumStages]) { + constexpr auto kDtypePerInt4 = sizeof(int4) / sizeof(dtype_t); + + // Broadcast current heads + // Lane `i` holds the head of rank `i` and `is_token_in_rank` + EP_STATIC_ASSERT(kMaxNumRanks <= 32, "Too many ranks"); + int num_topk_ranks = 0, topk_ranks[kMaxNumRanks], slot_indices[kMaxNumRanks]; + #pragma unroll + for (int i = 0; i < kNumRanks; ++i) + if (__shfl_sync(0xffffffff, is_token_in_rank, i)) { + slot_indices[num_topk_ranks] = __shfl_sync(0xffffffff, head_idx, i) % num_max_recv_tokens; + topk_ranks[num_topk_ranks++] = i; + } + EP_DEVICE_ASSERT(num_topk_ranks <= kMaxNumRanks); + EP_STATIC_ASSERT(not(kUseTMA and kMaybeWithBias), "TMA cannot be used by receiver warps"); + EP_STATIC_ASSERT(kNumStages == 2, "Only support 2 stages now"); + + // Reduce data + if constexpr (kUseTMA) { + constexpr int kNumTMABufferBytesPerStage = kNumTMALoadBytes * (NUM_MAX_NVL_PEERS + 1) + 16; + EP_DEVICE_ASSERT(hidden_int4 % 32 == 0); + + auto tma_load_buffer = [=](const int& i, const int& j) -> int4* { + return reinterpret_cast(smem_ptr + i * kNumTMABufferBytesPerStage + j * kNumTMALoadBytes); + }; + auto tma_store_buffer = [=](const int& i) -> int4* { + return reinterpret_cast(smem_ptr + i * kNumTMABufferBytesPerStage + NUM_MAX_NVL_PEERS * kNumTMALoadBytes); + }; + auto tma_mbarrier = [=](const int& i) -> uint64_t* { + return reinterpret_cast(smem_ptr + i * kNumTMABufferBytesPerStage + (NUM_MAX_NVL_PEERS + 1) * kNumTMALoadBytes); + }; + + // Prefetch + if (lane_id < num_topk_ranks) + tma_load_1d( + tma_load_buffer(0, lane_id), get_addr_fn(topk_ranks[lane_id], slot_indices[lane_id], 0), tma_mbarrier(0), kNumTMALoadBytes); + mbarrier_arrive_and_expect_tx(tma_mbarrier(0), lane_id < num_topk_ranks ? kNumTMALoadBytes : 0); + __syncwarp(); + + for (int shifted = 0, iter = 0; shifted < hidden_int4; shifted += 32, iter += 1) { + const int stage_idx = iter % kNumStages; + const int next_stage_idx = (iter + 1) % kNumStages; + + // Prefetch next stage + if (shifted + 32 < hidden_int4) { + if (lane_id < num_topk_ranks) + tma_load_1d(tma_load_buffer(next_stage_idx, lane_id), + get_addr_fn(topk_ranks[lane_id], slot_indices[lane_id], shifted + 32), + tma_mbarrier(next_stage_idx), + kNumTMALoadBytes); + mbarrier_arrive_and_expect_tx(tma_mbarrier(next_stage_idx), lane_id < num_topk_ranks ? kNumTMALoadBytes : 0); + __syncwarp(); + } + + mbarrier_wait(tma_mbarrier(stage_idx), tma_phase[stage_idx]); + float values[kDtypePerInt4] = {0}; + #pragma unroll + for (int j = 0; j < num_topk_ranks; ++j) { + auto recv_value_dtypes = reinterpret_cast(tma_load_buffer(stage_idx, j) + lane_id); + #pragma unroll + for (int k = 0; k < kDtypePerInt4; ++k) + values[k] += static_cast(recv_value_dtypes[k]); + } + + // Wait shared memory to be released + tma_store_wait(); + + // Copy into shared and issue TMA + auto out_dtypes = reinterpret_cast(tma_store_buffer(stage_idx) + lane_id); + #pragma unroll + for (int j = 0; j < kDtypePerInt4; ++j) + out_dtypes[j] = static_cast(values[j]); + tma_store_fence(); + __syncwarp(); + + if (elect_one_sync()) + tma_store_1d(tma_store_buffer(stage_idx), combined_row + shifted, kNumTMALoadBytes); + __syncwarp(); + } + + // Flush all writes + tma_store_wait<0>(); + } else { + #pragma unroll + for (int i = lane_id; i < hidden_int4; i += 32) { + // Read bias + // TODO: make it as a finer-grained template + int4 bias_0_value_int4, bias_1_value_int4; + if constexpr (kMaybeWithBias) { + bias_0_value_int4 = bias_0_int4 != nullptr ? ld_nc_global(bias_0_int4 + i) : make_int4(0, 0, 0, 0); + bias_1_value_int4 = bias_1_int4 != nullptr ? ld_nc_global(bias_1_int4 + i) : make_int4(0, 0, 0, 0); + } + + // Read buffers + // TODO: maybe too many registers here + int4 recv_value_int4[kMaxNumRanks]; + #pragma unroll + for (int j = 0; j < num_topk_ranks; ++j) + recv_value_int4[j] = ld_nc_global(get_addr_fn(topk_ranks[j], slot_indices[j], i)); + + // Clean + // Reduce bias + float values[kDtypePerInt4] = {0}; + if constexpr (kMaybeWithBias) { + auto bias_0_values = reinterpret_cast(&bias_0_value_int4); + auto bias_1_values = reinterpret_cast(&bias_1_value_int4); + #pragma unroll + for (int j = 0; j < kDtypePerInt4; ++j) + values[j] = static_cast(bias_0_values[j]) + static_cast(bias_1_values[j]); + } + + // Reduce all-to-all results + #pragma unroll + for (int j = 0; j < num_topk_ranks; ++j) { + auto recv_value_dtypes = reinterpret_cast(&recv_value_int4[j]); + #pragma unroll + for (int k = 0; k < kDtypePerInt4; ++k) + values[k] += static_cast(recv_value_dtypes[k]); + } + + // Cast back to `dtype_t` and write + int4 out_int4; + auto out_dtypes = reinterpret_cast(&out_int4); + #pragma unroll + for (int j = 0; j < kDtypePerInt4; ++j) + out_dtypes[j] = static_cast(values[j]); + st_na_global(combined_row + i, out_int4); + } + } + + // Reduce `topk_weights` + if (lane_id < num_topk) { + float value = 0; + #pragma unroll + for (int i = 0; i < num_topk_ranks; ++i) + value += recv_tw_fn(topk_ranks[i], slot_indices[i], lane_id); + st_na_global(combined_topk_weights + lane_id, value); + } + + // Return the minimum top-k rank + return topk_ranks[0]; +} + +template 0) ? kNumCombineForwarderWarps / kNumRDMARanks : 1, + int kNumForwarders = kNumRDMARanks * kNumWarpsPerForwarder, + int kNumRDMAReceivers = kNumForwarders - NUM_MAX_NVL_PEERS> +__global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* combined_x, + float* combined_topk_weights, + const bool* is_combined_token_in_rank, + const int4* x, + const float* topk_weights, + const int4* bias_0, + const int4* bias_1, + const int* combined_rdma_head, + const int* combined_nvl_head, + const SourceMeta* src_meta, + const int* rdma_channel_prefix_matrix, + const int* rdma_rank_prefix_sum, + const int* gbl_channel_prefix_matrix, + int num_tokens, + int num_combined_tokens, + int hidden, + int num_topk, + void* rdma_buffer_ptr, + int num_max_rdma_chunked_send_tokens, + int num_max_rdma_chunked_recv_tokens, + void** buffer_ptrs, + int num_max_nvl_chunked_send_tokens, + int num_max_nvl_chunked_recv_tokens, + int rank, + int num_ranks, + gpu_nixl_ctx nixl_ctx) { + enum class WarpRole { kNVLSender, kNVLAndRDMAForwarder, kRDMAReceiver, kCoordinator }; + + const auto sm_id = static_cast(blockIdx.x); + const auto num_threads = static_cast(blockDim.x), num_warps = num_threads / 32; + const auto thread_id = static_cast(threadIdx.x), lane_id = get_lane_id(); + const auto num_channels = static_cast(gridDim.x) / 2, channel_id = sm_id / 2; + const bool is_forwarder_sm = sm_id % 2 == 1; + + EP_DEVICE_ASSERT(num_topk <= 32); + EP_DEVICE_ASSERT(hidden % (sizeof(int4) / sizeof(dtype_t)) == 0); + const auto hidden_int4 = hidden / (sizeof(int4) / sizeof(dtype_t)); + const auto hidden_bytes = hidden_int4 * sizeof(int4); + const auto num_bytes_per_token = get_num_bytes_per_token(hidden_int4, 0, 0, num_topk); + + // NOTES: we decouple a channel into 2 SMs + const auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS; + auto role_meta = [=]() -> std::pair { + auto warp_id = thread_id / 32; + if (not is_forwarder_sm) { + if (warp_id < NUM_MAX_NVL_PEERS) { + auto shuffled_warp_id = warp_id; + shuffled_warp_id = (shuffled_warp_id + channel_id) % NUM_MAX_NVL_PEERS; + return {WarpRole::kNVLSender, shuffled_warp_id}; + } else if (warp_id < kNumForwarders) { + return {WarpRole::kRDMAReceiver, warp_id - NUM_MAX_NVL_PEERS}; + } else { + return {WarpRole::kCoordinator, 0}; + } + } else { + if (warp_id < kNumForwarders) { + auto shuffled_warp_id = (warp_id + channel_id) % kNumForwarders; + return {WarpRole::kNVLAndRDMAForwarder, shuffled_warp_id}; + } else { + return {WarpRole::kCoordinator, 0}; + } + } + }(); + auto warp_role = role_meta.first; + auto warp_id = role_meta.second; + + EP_DEVICE_ASSERT(num_warps == kNumForwarders + 1); + auto num_max_nvl_chunked_recv_tokens_per_rdma = num_max_nvl_chunked_recv_tokens / kNumRDMARanks; + + if (warp_role == WarpRole::kNVLSender) { + // NVL producers + const auto dst_nvl_rank = warp_id; + + // NVL layouts + // NOTES: to avoid deadlocks, we use separate NVL buffers for different RDMA sources + auto dst_buffer_ptr = buffer_ptrs[dst_nvl_rank], local_buffer_ptr = buffer_ptrs[nvl_rank]; + auto nvl_channel_x = AsymBuffer(dst_buffer_ptr, + num_max_nvl_chunked_recv_tokens * num_bytes_per_token, + NUM_MAX_NVL_PEERS, + channel_id, + num_channels, + nvl_rank) + .advance_also(local_buffer_ptr); + auto nvl_channel_head = AsymBuffer(local_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, dst_nvl_rank) + .advance_also(dst_buffer_ptr); + auto nvl_channel_tail = AsymBuffer(dst_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank) + .advance_also(local_buffer_ptr); + + // TMA stuffs + extern __shared__ __align__(1024) uint8_t smem_tma_buffer[]; + auto tma_buffer = smem_tma_buffer + dst_nvl_rank * kNumTMABytesPerSenderWarp; + auto tma_mbarrier = reinterpret_cast(tma_buffer + num_bytes_per_token); + uint32_t tma_phase = 0; + if (elect_one_sync()) { + mbarrier_init(tma_mbarrier, 1); + fence_barrier_init(); + EP_DEVICE_ASSERT(num_bytes_per_token + sizeof(uint64_t) <= kNumTMABytesPerSenderWarp); + } + __syncwarp(); + + // Get tasks for each RDMA lane + int token_start_idx = 0, token_end_idx = 0; + if (lane_id < kNumRDMARanks) { + int prefix_idx = (lane_id * NUM_MAX_NVL_PEERS + dst_nvl_rank) * num_channels + channel_id; + token_start_idx = gbl_channel_prefix_matrix[prefix_idx]; + token_end_idx = (prefix_idx == num_channels * num_ranks - 1) ? num_tokens : gbl_channel_prefix_matrix[prefix_idx + 1]; + } + __syncwarp(); + + // NOTES: here the cached value of each lane is only responsible for a single RDMA buffer + int cached_channel_head_idx = 0, cached_channel_tail_idx = 0; + EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); + + // Iterate over all tokens and send by chunks + int current_rdma_idx = channel_id % kNumRDMARanks; + while (true) { + // Exit if possible + if (__all_sync(0xffffffff, token_start_idx >= token_end_idx)) + break; + + // Decide the next RDMA buffer to send + bool is_lane_ready = false; + auto start_time = clock64(); + while (true) { + int num_used_slots = cached_channel_tail_idx - cached_channel_head_idx; + is_lane_ready = lane_id < kNumRDMARanks and token_start_idx < token_end_idx and + num_max_nvl_chunked_recv_tokens_per_rdma - num_used_slots >= num_max_nvl_chunked_send_tokens; + if (__any_sync(0xffffffff, is_lane_ready)) + break; + + // Retry + if (lane_id < kNumRDMARanks and token_start_idx < token_end_idx) + cached_channel_head_idx = ld_volatile_global(nvl_channel_head.buffer() + lane_id); + + // Timeout check + if (clock64() - start_time > NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) { + printf( + "NixlEP combine NVL sender timeout, channel: %d, RDMA: %d, nvl: %d, dst NVL: %d, RDMA lane: %d, head: %d, tail: " + "%d, start: %d, end: %d\n", + channel_id, + rdma_rank, + nvl_rank, + dst_nvl_rank, + lane_id, + ld_volatile_global(nvl_channel_head.buffer() + lane_id), + cached_channel_tail_idx, + token_start_idx, + token_end_idx); + trap(); + } + } + + // Sync token start index and count + for (int i = 0; i < kNumRDMARanks; ++i) { + current_rdma_idx = (current_rdma_idx + 1) % kNumRDMARanks; + if (__shfl_sync(0xffffffff, (token_start_idx >= token_end_idx) or (not is_lane_ready), current_rdma_idx)) + continue; + + // Sync token start index + auto token_idx = static_cast(__shfl_sync(0xffffffff, token_start_idx, current_rdma_idx)); + int num_tokens_in_chunk = + __shfl_sync(0xffffffff, min(num_max_nvl_chunked_send_tokens, token_end_idx - token_start_idx), current_rdma_idx); + + // Send by chunk + for (int chunk_idx = 0; chunk_idx < num_tokens_in_chunk; ++chunk_idx, ++token_idx) { + // Get an empty slot + int dst_slot_idx = 0; + if (lane_id == current_rdma_idx) { + dst_slot_idx = (cached_channel_tail_idx++) % num_max_nvl_chunked_recv_tokens_per_rdma; + dst_slot_idx = current_rdma_idx * num_max_nvl_chunked_recv_tokens_per_rdma + dst_slot_idx; + } + dst_slot_idx = __shfl_sync(0xffffffff, dst_slot_idx, current_rdma_idx); + + // Load data + auto shifted_x_buffers = nvl_channel_x.buffer() + dst_slot_idx * num_bytes_per_token; + auto shifted_x = x + token_idx * hidden_int4; + tma_store_wait<0>(); + if (elect_one_sync()) { + tma_load_1d(tma_buffer, shifted_x, tma_mbarrier, hidden_bytes); + mbarrier_arrive_and_expect_tx(tma_mbarrier, hidden_bytes); + } + __syncwarp(); + mbarrier_wait(tma_mbarrier, tma_phase); + + // Load source meta + if (lane_id == num_topk) + *reinterpret_cast(tma_buffer + hidden_bytes) = ld_nc_global(src_meta + token_idx); + + // Load `topk_weights` + if (lane_id < num_topk) + *reinterpret_cast(tma_buffer + hidden_bytes + sizeof(SourceMeta) + lane_id * sizeof(float)) = + ld_nc_global(topk_weights + token_idx * num_topk + lane_id); + + // Issue TMA store + tma_store_fence(); + __syncwarp(); + if (elect_one_sync()) + tma_store_1d(tma_buffer, shifted_x_buffers, num_bytes_per_token, false); + } + lane_id == current_rdma_idx ? (token_start_idx = static_cast(token_idx)) : 0; + } + + // Move queue tail + tma_store_wait<0>(); + __syncwarp(); + if (lane_id < kNumRDMARanks and is_lane_ready) + st_release_sys_global(nvl_channel_tail.buffer() + lane_id, cached_channel_tail_idx); + } + } else { + // Combiners and coordinators + // RDMA symmetric layout + auto rdma_channel_data = SymBuffer( + rdma_buffer_ptr, num_max_rdma_chunked_recv_tokens * num_bytes_per_token, kNumRDMARanks, channel_id, num_channels); + auto rdma_channel_head = SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); + auto rdma_channel_tail = SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); + + // NVL layouts + void* local_nvl_buffer = buffer_ptrs[nvl_rank]; + void* nvl_buffers[NUM_MAX_NVL_PEERS]; + #pragma unroll + for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) + nvl_buffers[i] = buffer_ptrs[i]; + auto nvl_channel_x = + AsymBuffer( + local_nvl_buffer, num_max_nvl_chunked_recv_tokens * num_bytes_per_token, NUM_MAX_NVL_PEERS, channel_id, num_channels) + .advance_also(nvl_buffers); + auto nvl_channel_head = + AsymBuffer(nvl_buffers, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank) + .advance_also(local_nvl_buffer); + auto nvl_channel_tail = AsymBuffer(local_nvl_buffer, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels) + .advance_also(nvl_buffers); + + // Combiner warp synchronization + __shared__ volatile int forwarder_nvl_head[kNumForwarders][NUM_MAX_NVL_PEERS]; + __shared__ volatile bool forwarder_retired[kNumForwarders]; + __shared__ volatile int rdma_receiver_rdma_head[kNumRDMAReceivers][kNumRDMARanks]; + __shared__ volatile bool rdma_receiver_retired[kNumRDMAReceivers]; + auto sync_forwarder_smem = [=]() { asm volatile("barrier.sync 0, %0;" ::"r"((kNumForwarders + 1) * 32)); }; + auto sync_rdma_receiver_smem = [=]() { asm volatile("barrier.sync 1, %0;" ::"r"((kNumRDMAReceivers + 1) * 32)); }; + + if (warp_role == WarpRole::kNVLAndRDMAForwarder) { + // Receive from NVL ranks and forward to RDMA ranks + // NOTES: this part is using "large warps" for each RDMA ranks + const auto dst_rdma_rank = warp_id / kNumWarpsPerForwarder; + const auto sub_warp_id = warp_id % kNumWarpsPerForwarder; + auto send_buffer = + dst_rdma_rank == rdma_rank ? rdma_channel_data.recv_buffer(dst_rdma_rank) : rdma_channel_data.send_buffer(dst_rdma_rank); + auto sync_large_warp = [=]() { + if (kNumWarpsPerForwarder == 1) { + __syncwarp(); + } else { + asm volatile("bar.sync %0, %1;" ::"r"(dst_rdma_rank + 2), "r"(kNumWarpsPerForwarder * 32)); + } + }; + EP_STATIC_ASSERT(kNumWarpsPerForwarder == 1 or kNumRDMARanks + 2 <= 16, "Barriers are not enough"); + + // TMA stuffs + constexpr int kNumStages = 2; + constexpr int kNumTMALoadBytes = sizeof(int4) * 32; + constexpr int kNumTMABufferBytesPerStage = kNumTMALoadBytes * (NUM_MAX_NVL_PEERS + 1) + 16; + EP_STATIC_ASSERT(kNumTMABufferBytesPerStage * kNumStages <= kNumTMABytesPerForwarderWarp, "TMA buffer is not larger enough"); + + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + auto smem_ptr = smem_buffer + warp_id * kNumStages * kNumTMABufferBytesPerStage; + auto tma_mbarrier = [=](const int& i) { + return reinterpret_cast(smem_ptr + i * kNumTMABufferBytesPerStage + kNumTMALoadBytes * (NUM_MAX_NVL_PEERS + 1)); + }; + uint32_t tma_phase[kNumStages] = {0}; + if (lane_id < kNumStages) { + mbarrier_init(tma_mbarrier(lane_id), 32); + fence_barrier_init(); + } + __syncwarp(); + + // Advance to the corresponding NVL buffer + nvl_channel_x.advance(dst_rdma_rank * num_max_nvl_chunked_recv_tokens_per_rdma * num_bytes_per_token); + nvl_channel_head.advance(dst_rdma_rank); + nvl_channel_tail.advance(dst_rdma_rank); + + // Clean shared memory and sync + EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); + lane_id < NUM_MAX_NVL_PEERS ? (forwarder_nvl_head[warp_id][lane_id] = 0) : 0; + lane_id == 0 ? (forwarder_retired[warp_id] = false) : false; + sync_forwarder_smem(); + + // Get count and cached head + int cached_nvl_channel_tail_idx = 0; + int num_tokens_to_combine = rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id]; + int num_tokens_prefix = channel_id == 0 ? 0 : rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id - 1]; + num_tokens_to_combine -= num_tokens_prefix; + num_tokens_prefix += dst_rdma_rank == 0 ? 0 : rdma_rank_prefix_sum[dst_rdma_rank - 1]; + combined_nvl_head += num_tokens_prefix * NUM_MAX_NVL_PEERS; + + // Iterate over all tokens and combine by chunks + for (int token_start_idx = 0; token_start_idx < num_tokens_to_combine; token_start_idx += num_max_rdma_chunked_send_tokens) { + // Check destination queue emptiness, or wait a buffer to be released + auto token_end_idx = min(token_start_idx + num_max_rdma_chunked_send_tokens, num_tokens_to_combine); + auto num_chunked_tokens = token_end_idx - token_start_idx; + auto start_time = clock64(); + while (sub_warp_id == 0 and lane_id == 0) { + // Inequality: `num_max_rdma_chunked_recv_tokens - (tail - head) >= num_chunked_tokens` + // Here, `token_start_idx` is the actual tail + int num_used_slots = token_start_idx - ld_volatile_global(rdma_channel_head.buffer(dst_rdma_rank)); + if (num_max_rdma_chunked_recv_tokens - num_used_slots >= num_chunked_tokens) + break; + + // Timeout check + if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { + printf( + "NixlEP combine forwarder (RDMA check) timeout, channel: %d, RDMA: %d, nvl: %d, dst RDMA: %d, head: %ld, tail: " + "%d, chunked: %d\n", + channel_id, + rdma_rank, + nvl_rank, + dst_rdma_rank, + ld_volatile_global(rdma_channel_head.buffer(dst_rdma_rank)), + token_start_idx, + num_chunked_tokens); + trap(); + } + } + sync_large_warp(); + + // Combine and write to the RDMA buffer + for (int token_idx = token_start_idx + sub_warp_id; token_idx < token_end_idx; token_idx += kNumWarpsPerForwarder) { + // Read expected head + EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); + int expected_head = -1; + if (lane_id < NUM_MAX_NVL_PEERS) { + expected_head = ld_nc_global(combined_nvl_head + token_idx * NUM_MAX_NVL_PEERS + lane_id); + expected_head < 0 ? (forwarder_nvl_head[warp_id][lane_id] = -expected_head - 1) + : (forwarder_nvl_head[warp_id][lane_id] = expected_head); + } + + // Wait lanes to be ready + start_time = clock64(); + while (cached_nvl_channel_tail_idx <= expected_head) { + cached_nvl_channel_tail_idx = ld_acquire_sys_global(nvl_channel_tail.buffer(lane_id)); + + // Timeout check + if (clock64() - start_time > NUM_TIMEOUT_CYCLES and lane_id < NUM_MAX_NVL_PEERS) { + printf( + "NixlEP combine forwarder (NVL check) timeout, channel: %d, RDMA: %d, nvl: %d, src NVL: %d, dst RDMA: %d, " + "tail: %d, waiting: %d, total: %d, sub: %d, large: %d, expected: %d\n", + channel_id, + rdma_rank, + nvl_rank, + lane_id, + dst_rdma_rank, + cached_nvl_channel_tail_idx, + token_idx, + num_tokens_to_combine, + sub_warp_id, + kNumWarpsPerForwarder, + expected_head); + trap(); + } + } + + // Combine current token + auto rdma_slot_idx = token_idx % num_max_rdma_chunked_recv_tokens; + void* shifted = send_buffer + rdma_slot_idx * num_bytes_per_token; + auto get_addr_fn = [&](int src_nvl_rank, int slot_idx, int hidden_int4_idx) -> int4* { + return reinterpret_cast(nvl_channel_x.buffer(src_nvl_rank) + slot_idx * num_bytes_per_token) + + hidden_int4_idx; + }; + auto recv_tw_fn = [&](int src_nvl_rank, int slot_idx, int topk_idx) -> float { + return ld_nc_global(reinterpret_cast(nvl_channel_x.buffer(src_nvl_rank) + slot_idx * num_bytes_per_token + + hidden_bytes + sizeof(SourceMeta)) + + topk_idx); + }; + combine_token( + expected_head >= 0, + expected_head, + lane_id, + hidden_int4, + num_topk, + static_cast(shifted), + reinterpret_cast(static_cast(shifted) + hidden_bytes + sizeof(SourceMeta)), + nullptr, + nullptr, + num_max_nvl_chunked_recv_tokens_per_rdma, + get_addr_fn, + recv_tw_fn, + smem_ptr, + tma_phase); + + // Update head + if (lane_id < NUM_MAX_NVL_PEERS) + expected_head < 0 ? (forwarder_nvl_head[warp_id][lane_id] = -expected_head - 1) + : (forwarder_nvl_head[warp_id][lane_id] = expected_head + 1); + } + sync_large_warp(); + + // Issue RDMA send + if (sub_warp_id == kNumWarpsPerForwarder - 1) { + if (dst_rdma_rank != rdma_rank) { + auto rdma_slot_idx = token_start_idx % num_max_rdma_chunked_recv_tokens; + const size_t num_bytes_per_msg = num_chunked_tokens * num_bytes_per_token; + const auto dst_ptr = + reinterpret_cast(rdma_channel_data.recv_buffer(rdma_rank) + rdma_slot_idx * num_bytes_per_token); + const auto src_ptr = + reinterpret_cast(rdma_channel_data.send_buffer(dst_rdma_rank) + rdma_slot_idx * num_bytes_per_token); + int translated_dst_comb = translate_dst_rdma_rank(dst_rdma_rank, nvl_rank); + nixlMemViewElem src_mdesc_comb{nixl_ctx.local_mvh, 0, nixl_ctx.offset_get(src_ptr)}; + nixlMemViewElem dst_mdesc_comb{nixl_ctx.remote_mvh, (size_t)translated_dst_comb, nixl_ctx.offset_get(dst_ptr)}; + EP_DEVICE_ASSERT(nixlPut( + src_mdesc_comb, dst_mdesc_comb, num_bytes_per_msg, channel_id) == + NIXL_IN_PROG); + } else { + memory_fence(); + } + + // Write new RDMA tail + __syncwarp(); + if (elect_one_sync()) { + auto tail_ptr = reinterpret_cast(rdma_channel_tail.buffer(rdma_rank)); + if(dst_rdma_rank == rdma_rank){ + atomicAdd(reinterpret_cast(tail_ptr), static_cast(num_chunked_tokens)); + } else { + int translated_dst_ct = translate_dst_rdma_rank(dst_rdma_rank, nvl_rank); + nixlMemViewElem tail_mdesc_ct{nixl_ctx.remote_mvh, (size_t)translated_dst_ct, nixl_ctx.offset_get(tail_ptr)}; + EP_DEVICE_ASSERT(nixlAtomicAdd( + num_chunked_tokens, tail_mdesc_ct, channel_id, 0) == + NIXL_IN_PROG); + } + } + } + } + + // Retired + __syncwarp(); + if (elect_one_sync()) + forwarder_retired[warp_id] = true; + } else if (warp_role == WarpRole::kRDMAReceiver) { + // Receive from RDMA ranks and write to the output tensor + // Clean shared memory and sync + EP_DEVICE_ASSERT(kNumRDMARanks <= 32); + lane_id < kNumRDMARanks ? (rdma_receiver_rdma_head[warp_id][lane_id] = 0) : 0; + lane_id == 0 ? (rdma_receiver_retired[warp_id] = false) : 0; + sync_rdma_receiver_smem(); + + // The same tokens as the dispatch process + int token_start_idx, token_end_idx; + get_channel_task_range(num_combined_tokens, num_channels, channel_id, token_start_idx, token_end_idx); + + // Iterate over all tokens and combine + int cached_channel_tail_idx = 0; + for (int64_t token_idx = token_start_idx + warp_id; token_idx < token_end_idx; token_idx += kNumRDMAReceivers) { + // Read expected head + EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); + int expected_head = -1; + if (lane_id < kNumRDMARanks) { + expected_head = ld_nc_global(combined_rdma_head + token_idx * kNumRDMARanks + lane_id); + (expected_head < 0) ? (rdma_receiver_rdma_head[warp_id][lane_id] = -expected_head - 1) + : (rdma_receiver_rdma_head[warp_id][lane_id] = expected_head); + } + + // Wait lanes to be ready + auto start_time = clock64(); + while (cached_channel_tail_idx <= expected_head) { + cached_channel_tail_idx = static_cast(ld_acquire_sys_global(rdma_channel_tail.buffer(lane_id))); + + // Timeout check + if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { + printf( + "NixlEP combine RDMA receiver timeout, channel: %d, RDMA: %d, nvl: %d, src RDMA: %d, tail: %d, waiting: %ld, " + "expect: %d\n", + channel_id, + rdma_rank, + nvl_rank, + lane_id, + cached_channel_tail_idx, + token_idx, + expected_head); + trap(); + } + } + __syncwarp(); + + // Combine current token + auto get_addr_fn = [&](int src_rdma_rank, int slot_idx, int hidden_int4_idx) -> int4* { + return reinterpret_cast(rdma_channel_data.recv_buffer(src_rdma_rank) + slot_idx * num_bytes_per_token) + + hidden_int4_idx; + }; + auto recv_tw_fn = [&](int src_rdma_rank, int slot_idx, int topk_idx) -> float { + return ld_nc_global(reinterpret_cast(rdma_channel_data.recv_buffer(src_rdma_rank) + + slot_idx * num_bytes_per_token + hidden_bytes + sizeof(SourceMeta)) + + topk_idx); + }; + uint32_t dummy_tma_phases[2]; + combine_token( + expected_head >= 0, + expected_head, + lane_id, + hidden_int4, + num_topk, + combined_x + token_idx * hidden_int4, + combined_topk_weights + token_idx * num_topk, + bias_0 == nullptr ? nullptr : bias_0 + token_idx * hidden_int4, + bias_1 == nullptr ? nullptr : bias_1 + token_idx * hidden_int4, + num_max_rdma_chunked_recv_tokens, + get_addr_fn, + recv_tw_fn, + nullptr, + dummy_tma_phases); + } + + // Retired + __syncwarp(); + if (elect_one_sync()) + rdma_receiver_retired[warp_id] = true; + } else { + // Coordinator + // Sync shared memory status + is_forwarder_sm ? sync_forwarder_smem() : sync_rdma_receiver_smem(); + const auto num_warps_per_rdma_rank = kNumForwarders / kNumRDMARanks; + + int last_rdma_head = 0; + int last_nvl_head[kNumRDMARanks] = {0}; + int dst_rdma_rank = lane_id < kNumRDMARanks ? lane_id : 0; + int dst_nvl_rank = lane_id < NUM_MAX_NVL_PEERS ? lane_id : 0; + EP_STATIC_ASSERT(kNumCombineForwarderWarps <= 32, "Invalid number of forwarder warps"); + while (true) { + // Retired + if (not is_forwarder_sm and __all_sync(0xffffffff, lane_id >= kNumRDMAReceivers or rdma_receiver_retired[lane_id])) + break; + if (is_forwarder_sm and __all_sync(0xffffffff, lane_id >= kNumForwarders or forwarder_retired[lane_id])) + break; + + // Find minimum head for RDMA ranks + if (not is_forwarder_sm) { + int min_head = std::numeric_limits::max(); + #pragma unroll + for (int i = 0; i < kNumRDMAReceivers; ++i) + if (not rdma_receiver_retired[i]) + min_head = min(min_head, rdma_receiver_rdma_head[i][dst_rdma_rank]); + if (min_head != std::numeric_limits::max() and min_head >= last_rdma_head + num_max_rdma_chunked_send_tokens and + lane_id < kNumRDMARanks) { + if (dst_rdma_rank == rdma_rank) { + atomicAdd(reinterpret_cast(rdma_channel_head.buffer(rdma_rank)), static_cast(min_head - last_rdma_head)); + } else { + size_t head_counter_offset = nixl_ctx.offset_get(reinterpret_cast(rdma_channel_head.buffer(rdma_rank))); + int translated_dst_ch = translate_dst_rdma_rank(dst_rdma_rank, nvl_rank); + nixlMemViewElem head_mdesc_ch{nixl_ctx.remote_mvh, (size_t)translated_dst_ch, head_counter_offset}; + EP_DEVICE_ASSERT(nixlAtomicAdd(min_head - last_rdma_head, head_mdesc_ch, channel_id, 0) == NIXL_IN_PROG); + } + last_rdma_head = min_head; + } + } else { + // Find minimum head for NVL ranks + #pragma unroll + for (int i = 0; i < kNumRDMARanks; ++i) { + int min_head = std::numeric_limits::max(); + #pragma unroll + for (int j = 0; j < num_warps_per_rdma_rank; ++j) + if (not forwarder_retired[i * num_warps_per_rdma_rank + j]) + min_head = min(min_head, forwarder_nvl_head[i * num_warps_per_rdma_rank + j][dst_nvl_rank]); + if (min_head != std::numeric_limits::max() and min_head > last_nvl_head[i] and lane_id < NUM_MAX_NVL_PEERS) + st_relaxed_sys_global(nvl_channel_head.buffer_by(dst_nvl_rank) + i, last_nvl_head[i] = min_head); + } + } + + // Nanosleep and let other warps work + __nanosleep(NUM_WAIT_NANOSECONDS); + } + } + } +} + +void combine(cudaDataType_t type, + void* combined_x, + float* combined_topk_weights, + const bool* is_combined_token_in_rank, + const void* x, + const float* topk_weights, + const void* bias_0, + const void* bias_1, + const int* combined_rdma_head, + const int* combined_nvl_head, + const void* src_meta, + const int* rdma_channel_prefix_matrix, + const int* rdma_rank_prefix_sum, + const int* gbl_channel_prefix_matrix, + int num_tokens, + int num_combined_tokens, + int hidden, + int num_topk, + void* rdma_buffer_ptr, + int num_max_rdma_chunked_send_tokens, + int num_max_rdma_chunked_recv_tokens, + void** buffer_ptrs, + int num_max_nvl_chunked_send_tokens, + int num_max_nvl_chunked_recv_tokens, + int rank, + int num_ranks, + cudaStream_t stream, + int num_channels, + bool low_latency_mode, + gpu_nixl_ctx nixl_ctx) { + constexpr int kNumCombineForwarderWarps = 24; + constexpr int kNumTMABytesPerSenderWarp = 16384; + constexpr int kNumTMABytesPerForwarderWarp = 9248; + constexpr int smem_size = + std::max(kNumTMABytesPerSenderWarp * NUM_MAX_NVL_PEERS, kNumTMABytesPerForwarderWarp * kNumCombineForwarderWarps); + +#define COMBINE_LAUNCH_CASE(num_rdma_ranks) \ + { \ + auto combine_func = low_latency_mode ? combine \ + : combine; \ + SET_SHARED_MEMORY_FOR_TMA(combine_func); \ + LAUNCH_KERNEL(&cfg, \ + combine_func, \ + reinterpret_cast(combined_x), \ + combined_topk_weights, \ + is_combined_token_in_rank, \ + reinterpret_cast(x), \ + topk_weights, \ + reinterpret_cast(bias_0), \ + reinterpret_cast(bias_1), \ + combined_rdma_head, \ + combined_nvl_head, \ + reinterpret_cast(src_meta), \ + rdma_channel_prefix_matrix, \ + rdma_rank_prefix_sum, \ + gbl_channel_prefix_matrix, \ + num_tokens, \ + num_combined_tokens, \ + hidden, \ + num_topk, \ + rdma_buffer_ptr, \ + num_max_rdma_chunked_send_tokens, \ + num_max_rdma_chunked_recv_tokens, \ + buffer_ptrs, \ + num_max_nvl_chunked_send_tokens, \ + num_max_nvl_chunked_recv_tokens, \ + rank, \ + num_ranks, \ + nixl_ctx); \ + } \ + break + + int num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; + auto num_warps_per_forwarder = std::max(kNumCombineForwarderWarps / num_rdma_ranks, 1); + int num_forwarder_warps = num_rdma_ranks * num_warps_per_forwarder; + EP_HOST_ASSERT(num_rdma_ranks <= kNumCombineForwarderWarps); + EP_HOST_ASSERT(num_forwarder_warps > NUM_MAX_NVL_PEERS and num_forwarder_warps % num_rdma_ranks == 0); + EP_HOST_ASSERT(num_max_nvl_chunked_recv_tokens % num_rdma_ranks == 0); + EP_HOST_ASSERT(num_max_nvl_chunked_recv_tokens / num_rdma_ranks > + std::max(num_max_rdma_chunked_send_tokens, num_max_nvl_chunked_send_tokens)); + EP_HOST_ASSERT(num_max_nvl_chunked_recv_tokens / num_rdma_ranks - num_warps_per_forwarder >= num_max_nvl_chunked_send_tokens); + EP_HOST_ASSERT(num_max_rdma_chunked_send_tokens >= num_warps_per_forwarder); + EP_HOST_ASSERT(type == CUDA_R_16BF); + + SETUP_LAUNCH_CONFIG(num_channels * 2, (num_forwarder_warps + 1) * 32, stream); + SWITCH_RDMA_RANKS(COMBINE_LAUNCH_CASE); +#undef COMBINE_LAUNCH_CASE +} + +} // namespace ht + +} // namespace nixl_ep diff --git a/examples/device/ep/csrc/kernels/nixl_ep.cu b/examples/device/ep/csrc/kernels/nixl_ep_ll.cu similarity index 97% rename from examples/device/ep/csrc/kernels/nixl_ep.cu rename to examples/device/ep/csrc/kernels/nixl_ep_ll.cu index 855d1c17..b03539f7 100644 --- a/examples/device/ep/csrc/kernels/nixl_ep.cu +++ b/examples/device/ep/csrc/kernels/nixl_ep_ll.cu @@ -32,21 +32,17 @@ namespace cg = cooperative_groups; namespace nixl_ep { -namespace ep_kernels { - -__device__ inline uint64_t gpu_nixl_ctx::offset_get(uint64_t ptr) { - return ptr - reinterpret_cast(rdma_buffer_ptr); -} +__device__ inline void* p2p_ptr_get(gpu_nixl_ctx& ctx, uint64_t dst_ptr, int dst_rank) { + if (dst_rank == ctx.rank) return (void*) dst_ptr; -__device__ inline void* gpu_nixl_ctx::p2p_ptr_get(uint64_t dst_ptr, int dst_rank) { - if (dst_rank == rank) return (void*) dst_ptr; - - void *remote_ptr = nixlGetPtr(remote_mvh, dst_rank); + void *remote_ptr = nixlGetPtr(ctx.remote_mvh, dst_rank); if (remote_ptr == nullptr) return nullptr; - return (void*) ((uint64_t) remote_ptr + offset_get(dst_ptr)); + return (void*) ((uint64_t) remote_ptr + ctx.offset_get(dst_ptr)); } +namespace ep_kernels { + template __forceinline__ __device__ bool is_rank_masked(int* mask_buffer_ptr, int rank) { if (mask_buffer_ptr == nullptr) { @@ -78,7 +74,7 @@ dispatch(void* packed_recv_x, void* packed_recv_x_scales, 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) { + bool round_scale, int phases, nixl_ep::gpu_nixl_ctx nixl_ctx) { const auto sm_id = static_cast(blockIdx.x); const auto thread_id = static_cast(threadIdx.x); const auto warp_id = thread_id / 32, lane_id = get_lane_id(); @@ -187,8 +183,8 @@ dispatch(void* packed_recv_x, void* packed_recv_x_scales, dst_expert_local_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_msg + rank * num_max_dispatch_tokens_per_rank * num_bytes_per_msg + slot_idx * num_bytes_per_msg; + void* dst_p2p_ptr = p2p_ptr_get(nixl_ctx, dst_ptr, dst_rank); if (not is_rank_masked(mask_buffer_ptr, dst_rank)) { - void* dst_p2p_ptr = nixl_ctx.p2p_ptr_get(dst_ptr, dst_rank); if (dst_p2p_ptr == 0) { nixlMemViewElem src_mdesc{nixl_ctx.local_mvh, 0, nixl_ctx.offset_get(src_ptr)}; nixlMemViewElem dst_mdesc{nixl_ctx.remote_mvh, (size_t) dst_rank, nixl_ctx.offset_get(dst_ptr)}; @@ -256,8 +252,8 @@ dispatch(void* packed_recv_x, void* packed_recv_x_scales, // Wait local sends issued and send expert counts while (ld_acquire_global(atomic_finish_counter_per_expert + responsible_expert_idx) != FINISHED_SUM_TAG * 2); auto dst_ptr = reinterpret_cast(rdma_recv_count + dst_expert_local_idx * num_ranks + rank); + void* dst_p2p_ptr = p2p_ptr_get(nixl_ctx, dst_ptr, dst_rank); if (not is_rank_masked(mask_buffer_ptr, dst_rank)) { - void* dst_p2p_ptr = nixl_ctx.p2p_ptr_get(dst_ptr, dst_rank); if (dst_p2p_ptr == 0) { nixlMemViewElem dst_mdesc{nixl_ctx.remote_mvh, static_cast(dst_rank), nixl_ctx.offset_get(dst_ptr)}; EP_DEVICE_ASSERT(nixlAtomicAdd(num_tokens_sent + 1, dst_mdesc, dst_expert_local_idx) == NIXL_IN_PROG); @@ -400,7 +396,7 @@ void dispatch(void* packed_recv_x, void* packed_recv_x_scales, 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) { + cudaStream_t stream, int phases, nixl_ep::gpu_nixl_ctx nixl_ctx) { constexpr int kNumMaxTopK = 11; const int num_warp_groups = ceil_div(num_experts, num_device_sms); const int num_warps_per_group = 32 / num_warp_groups; @@ -620,7 +616,7 @@ combine(void* combined_x, 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) { + int phases, bool zero_copy, nixl_ep::gpu_nixl_ctx nixl_ctx) { const auto sm_id = __shfl_sync(0xffffffff, static_cast(blockIdx.x), 0); const auto num_sms = __shfl_sync(0xffffffff, static_cast(gridDim.x), 0); const auto thread_id = static_cast(threadIdx.x); @@ -726,7 +722,7 @@ combine(void* combined_x, const auto buf_ptr = reinterpret_cast(rdma_send_x_vec_row); const auto dst_ptr = reinterpret_cast(rdma_recv_x) + (global_expert_idx * num_max_dispatch_tokens_per_rank + src_idx) * num_bytes_per_slot; - void* dst_p2p_ptr = nixl_ctx.p2p_ptr_get(dst_ptr, dst_rank); + void* dst_p2p_ptr = p2p_ptr_get(nixl_ctx, dst_ptr, dst_rank); int num_send_bytes = hidden * sizeof(nv_bfloat16); if (not zero_copy or dst_p2p_ptr != 0) { @@ -805,8 +801,8 @@ combine(void* combined_x, if (sub_warp_id == 1 and lane_id == 0) { while (ld_acquire_global(atomic_clean_flag) == 0); auto dst_ptr = reinterpret_cast(rdma_recv_flag + global_expert_idx); + void* dst_p2p_ptr = p2p_ptr_get(nixl_ctx, dst_ptr, dst_rank); if (not is_rank_masked(mask_buffer_ptr, dst_rank)) { - void* dst_p2p_ptr = nixl_ctx.p2p_ptr_get(dst_ptr, dst_rank); if (dst_p2p_ptr == 0) { nixlMemViewElem dst_mdesc{nixl_ctx.remote_mvh, (size_t) dst_rank, nixl_ctx.offset_get(dst_ptr)}; EP_DEVICE_ASSERT(nixlAtomicAdd(1, dst_mdesc, local_expert_idx) == NIXL_IN_PROG); @@ -1013,7 +1009,7 @@ void combine(void* combined_x, 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) { + cudaStream_t stream, int phases, bool zero_copy, nixl_ep::gpu_nixl_ctx nixl_ctx) { constexpr int kNumMaxTopk = 11; const int num_warp_groups = ceil_div(num_experts, num_device_sms); const int num_warps_per_group = 32 / num_warp_groups; @@ -1124,7 +1120,7 @@ void clean_mask_buffer(int* mask_buffer_ptr, int num_ranks, cudaStream_t stream) } template -__forceinline__ __device__ void barrier(ep_kernels::gpu_nixl_ctx nixl_ctx, int* mask_buffer_ptr, int thread_id) { +__forceinline__ __device__ void barrier(nixl_ep::gpu_nixl_ctx nixl_ctx, int* mask_buffer_ptr, int thread_id) { EP_DEVICE_ASSERT(kNumThreads >= nixl_ctx.max_num_ranks); if (thread_id < nixl_ctx.max_num_ranks && thread_id != nixl_ctx.rank) { @@ -1154,12 +1150,12 @@ __forceinline__ __device__ void barrier(ep_kernels::gpu_nixl_ctx nixl_ctx, int* } template -__global__ void barrier_kernel(ep_kernels::gpu_nixl_ctx nixl_ctx, int* mask_buffer_ptr) { +__global__ void barrier_kernel(nixl_ep::gpu_nixl_ctx nixl_ctx, int* mask_buffer_ptr) { const auto thread_id = static_cast(threadIdx.x); barrier(nixl_ctx, mask_buffer_ptr, thread_id); } -void barrier(ep_kernels::gpu_nixl_ctx nixl_ctx, int* mask_buffer_ptr, cudaStream_t stream) { +void barrier(nixl_ep::gpu_nixl_ctx nixl_ctx, int* mask_buffer_ptr, cudaStream_t stream) { constexpr int kNumThreads = 32; SETUP_LAUNCH_CONFIG(1, kNumThreads, stream); LAUNCH_KERNEL(&cfg, barrier_kernel, nixl_ctx, mask_buffer_ptr); diff --git a/examples/device/ep/csrc/kernels/runtime.cu b/examples/device/ep/csrc/kernels/runtime.cu new file mode 100644 index 00000000..c4d10e46 --- /dev/null +++ b/examples/device/ep/csrc/kernels/runtime.cu @@ -0,0 +1,54 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 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 +#include + +#include "configs.cuh" +#include "exception.cuh" +#include "launch.cuh" +#include "utils.cuh" + +#include + +namespace nixl_ep { + +namespace intranode { + +template +__global__ void barrier(int** barrier_signal_ptrs, int rank) { + barrier_block(barrier_signal_ptrs, rank); +} + +void barrier(int** barrier_signal_ptrs, int rank, int num_nvl_ranks, cudaStream_t stream) { +#define BARRIER_LAUNCH_CASE(ranks) \ + LAUNCH_KERNEL(&cfg, barrier, barrier_signal_ptrs, rank); \ + break + + SETUP_LAUNCH_CONFIG(1, 32, stream); + SWITCH_NVL_RANKS(BARRIER_LAUNCH_CASE); +#undef BARRIER_LAUNCH_CASE +} + +} // namespace intranode + +} // namespace nixl_ep diff --git a/examples/device/ep/csrc/kernels/utils.cuh b/examples/device/ep/csrc/kernels/utils.cuh index 7649326b..10bce4ca 100644 --- a/examples/device/ep/csrc/kernels/utils.cuh +++ b/examples/device/ep/csrc/kernels/utils.cuh @@ -85,6 +85,9 @@ __device__ __forceinline__ void trap() { __device__ __forceinline__ void memory_fence() { asm volatile("fence.acq_rel.sys;":: : "memory"); } +__device__ __forceinline__ void st_relaxed_sys_global(const int *ptr, int val) { + asm volatile("st.relaxed.sys.global.s32 [%0], %1;"::"l"(ptr), "r"(val) : "memory"); +} __device__ __forceinline__ void st_release_sys_global(const int *ptr, int val) { asm volatile("st.release.sys.global.s32 [%0], %1;"::"l"(ptr), "r"(val) : "memory"); @@ -137,6 +140,41 @@ __device__ __forceinline__ int atomic_add_release_global(const uint64_t* ptr, ui asm volatile("atom.add.release.gpu.global.u64 %0, [%1], %2;" : "=l"(ret) : "l"(ptr), "l"(value)); return ret; } +__device__ __forceinline__ int ld_acquire_cta(const int *ptr) { + int ret; + asm volatile("ld.acquire.cta.s32 %0, [%1];" : "=r"(ret) : "l"(ptr)); + return ret; +} + +__device__ __forceinline__ int ld_acquire_cta(const volatile int *ptr) { + int ret; + asm volatile("ld.acquire.cta.s32 %0, [%1];" : "=r"(ret) : "l"(ptr)); + return ret; +} + +__device__ __forceinline__ int ld_volatile_global(const int *ptr) { + int ret; + asm volatile("ld.volatile.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr)); + return ret; +} + +__device__ __forceinline__ float ld_volatile_global(const float *ptr) { + float ret; + asm volatile("ld.volatile.global.f32 %0, [%1];" : "=f"(ret) : "l"(ptr)); + return ret; +} + +__device__ __forceinline__ int64_t ld_volatile_global(const int64_t *ptr) { + int64_t ret; + asm volatile("ld.volatile.global.s64 %0, [%1];" : "=l"(ret) : "l"(ptr)); + return ret; +} + +__device__ __forceinline__ int64_t ld_volatile_global(const uint64_t *ptr) { + int64_t ret; + asm volatile("ld.volatile.global.u64 %0, [%1];" : "=l"(ret) : "l"(ptr)); + return ret; +} #ifndef DISABLE_AGGRESSIVE_PTX_INSTRS #define LD_NC_FUNC "ld.global.nc.L1::no_allocate.L2::256B" @@ -266,6 +304,11 @@ __device__ __forceinline__ uint32_t elect_one_sync() { #endif } +__device__ __forceinline__ void fence_view_async_shared() { + asm volatile("fence.proxy.async.shared::cta; \n" :: ); +} + + // TMA PTX instructions #ifndef DISABLE_SM90_FEATURES __device__ __forceinline__ void fence_barrier_init() { @@ -332,7 +375,7 @@ __device__ __forceinline__ void tma_store_1d(const void* smem_ptr, const void* g asm volatile("cp.async.bulk.commit_group;"); } -template +template __device__ __forceinline__ void tma_store_wait() { asm volatile("cp.async.bulk.wait_group.read %0;" :: "n"(N) : "memory"); } @@ -421,6 +464,62 @@ __forceinline__ __device__ out_dtype_t extract_required_scale_format(float value } } +template +__forceinline__ __device__ void +barrier_block(int** barrier_signal_ptrs, int rank) { + auto thread_id = static_cast(threadIdx.x); + + // For non-sync-only cases, the memory operations by other threads in the block must be visible to the `sys` scope + if constexpr (not kSyncOnly) { + memory_fence(); + __syncthreads(); + } + + // Add self-ranks, sub other ranks + if (thread_id < kNumRanks) { + atomicAdd_system(barrier_signal_ptrs[rank] + thread_id, FINISHED_SUM_TAG); + atomicSub_system(barrier_signal_ptrs[thread_id] + rank, FINISHED_SUM_TAG); + } + EP_DEVICE_ASSERT(kNumRanks <= blockDim.x); + + // Check timeout + auto start_time = clock64(); + while (true) { + auto value = thread_id < kNumRanks ? ld_volatile_global(barrier_signal_ptrs[rank] + thread_id) : 0; + if (__all_sync(0xffffffff, value <= 0)) + break; + + if (clock64() - start_time > NUM_TIMEOUT_CYCLES and thread_id < kNumRanks) { + printf("NixlEP timeout check failed: rank = %d, thread = %d, value = %d)\n", rank, thread_id, value); + trap(); + } + } + __syncthreads(); +} + +__forceinline__ __device__ int atomic_cas_cta_acquire(int* addr, int x, int y) { + int ret; + asm volatile("atom.acquire.cta.shared::cta.cas.b32 %0, [%1], %2, %3;" : "=r"(ret) : "l"(addr), "r"(x), "r"(y) : "memory"); + return ret; +} + +__forceinline__ __device__ int atomic_exch_cta_release(int* addr, int x) { + int ret; + asm volatile("atom.release.cta.shared::cta.exch.b32 %0, [%1], %2;" : "=r"(ret) : "l"(addr), "r"(x) : "memory"); + return ret; +} + +__forceinline__ __device__ void acquire_lock(int* mutex) { + // To make later memory operations valid, we must use `acquire` for memory semantics + while (atomic_cas_cta_acquire(mutex, 0, 1) != 0); +} + +__forceinline__ __device__ void release_lock(int* mutex) { + // To make previous memory operations visible to other threads, we must use `release` for memory semantics + atomic_exch_cta_release(mutex, 0); +} + + // Operation functors template struct ReduceSum { __device__ T operator()(T a, T b) const { return a + b; } }; template struct ReduceMax { __device__ T operator()(T a, T b) const { return a > b ? a : b; } }; diff --git a/examples/device/ep/csrc/nixl_ep.cpp b/examples/device/ep/csrc/nixl_ep.cpp index e8f2055d..a5a33ca9 100644 --- a/examples/device/ep/csrc/nixl_ep.cpp +++ b/examples/device/ep/csrc/nixl_ep.cpp @@ -52,48 +52,55 @@ #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)); } -void Buffer::update_memory_buffers(int num_ranks, int num_experts_per_rank, int64_t num_rdma_bytes) +void Buffer::update_memory_buffers(int num_ranks, int num_experts_per_rank, int64_t num_rdma_bytes, int64_t num_nvl_bytes) { if (!available) { - init(num_ranks, num_experts_per_rank, num_rdma_bytes); + init(num_ranks, num_experts_per_rank, num_nvl_bytes, 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): +Buffer::Buffer(int rank, bool explicitly_destroy, bool low_latency_mode): + low_latency_mode(low_latency_mode), rank(rank), num_ranks(1), explicitly_destroy(explicitly_destroy), comm_stream(at::cuda::getStreamFromPool(true)) {} -void Buffer::init(int num_ranks, int num_experts_per_rank, int64_t num_rdma_bytes) +void Buffer::init(int num_ranks, int num_experts_per_rank, int64_t num_nvl_bytes, int64_t num_rdma_bytes) { // Update buffer attributes this->max_num_ranks = num_ranks; this->max_experts_per_rank = num_experts_per_rank; + this->num_nvl_bytes = num_nvl_bytes; this->num_rdma_bytes = num_rdma_bytes; + // Metadata memory + int64_t barrier_signal_bytes = NUM_MAX_NVL_PEERS * sizeof(int); + int64_t buffer_ptr_bytes = NUM_MAX_NVL_PEERS * sizeof(void*); + int64_t barrier_signal_ptr_bytes = NUM_MAX_NVL_PEERS * sizeof(int*); + // Common checks EP_STATIC_ASSERT(NUM_BUFFER_ALIGNMENT_BYTES % sizeof(int4) == 0, "Invalid alignment"); - EP_HOST_ASSERT(num_rdma_bytes % NUM_BUFFER_ALIGNMENT_BYTES == 0); - EP_HOST_ASSERT(num_rdma_bytes / sizeof(int4) < std::numeric_limits::max()); - EP_HOST_ASSERT(0 <= rank and rank < num_ranks); + EP_HOST_ASSERT(num_nvl_bytes % NUM_BUFFER_ALIGNMENT_BYTES == 0 and (num_nvl_bytes <= std::numeric_limits::max() or num_rdma_bytes == 0)); + EP_HOST_ASSERT(num_rdma_bytes % NUM_BUFFER_ALIGNMENT_BYTES == 0 and (low_latency_mode or num_rdma_bytes <= std::numeric_limits::max())); + EP_HOST_ASSERT(0 <= rank and rank < num_ranks and (num_ranks <= NUM_MAX_NVL_PEERS * NUM_MAX_RDMA_PEERS or low_latency_mode)); + EP_HOST_ASSERT(num_ranks < NUM_MAX_NVL_PEERS or num_ranks % NUM_MAX_NVL_PEERS == 0); + if (num_rdma_bytes > 0) + EP_HOST_ASSERT(num_ranks > NUM_MAX_NVL_PEERS or low_latency_mode); // Get ranks CUDA_CHECK(cudaGetDevice(&device_id)); + rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS; + num_rdma_ranks = std::max(1, num_ranks / NUM_MAX_NVL_PEERS), num_nvl_ranks = std::min(num_ranks, NUM_MAX_NVL_PEERS); + // Get device info cudaDeviceProp device_prop = {}; CUDA_CHECK(cudaGetDeviceProperties(&device_prop, device_id)); @@ -102,11 +109,40 @@ void Buffer::init(int num_ranks, int num_experts_per_rank, int64_t num_rdma_byte auto per_channel_bytes = ceil_div(num_rdma_bytes, denom_sms); EP_HOST_ASSERT(per_channel_bytes < std::numeric_limits::max()); + if (num_nvl_bytes > 0) { + // Local IPC: alloc local memory and set local IPC handles + CUDA_CHECK(cudaMalloc(&buffer_ptrs[nvl_rank], num_nvl_bytes + barrier_signal_bytes + buffer_ptr_bytes + barrier_signal_ptr_bytes)); + CUDA_CHECK(cudaIpcGetMemHandle(&ipc_handles[nvl_rank], buffer_ptrs[nvl_rank])); + buffer_ptrs_gpu = reinterpret_cast(static_cast(buffer_ptrs[nvl_rank]) + num_nvl_bytes + barrier_signal_bytes); + + // Set barrier signals + barrier_signal_ptrs[nvl_rank] = reinterpret_cast(static_cast(buffer_ptrs[nvl_rank]) + num_nvl_bytes); + barrier_signal_ptrs_gpu = reinterpret_cast(static_cast(buffer_ptrs[nvl_rank]) + num_nvl_bytes + barrier_signal_bytes + buffer_ptr_bytes); + + // No need to synchronize, will do a full device sync during `sync` + CUDA_CHECK(cudaMemsetAsync(barrier_signal_ptrs[nvl_rank], 0, barrier_signal_bytes, comm_stream)); + } + // Create 32 MiB workspace m_workspace_alloc = std::make_unique(NUM_WORKSPACE_BYTES); workspace = m_workspace_alloc->ptr(); CUDA_CHECK(cudaMemsetAsync(workspace, 0, NUM_WORKSPACE_BYTES, comm_stream)); + // MoE counter + CUDA_CHECK(cudaMallocHost(&moe_recv_counter, sizeof(int64_t), cudaHostAllocMapped)); + CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_counter_mapped, const_cast(moe_recv_counter), 0)); + *moe_recv_counter = -1; + + // MoE expert-level counter + CUDA_CHECK(cudaMallocHost(&moe_recv_expert_counter, sizeof(int) * NUM_MAX_LOCAL_EXPERTS, cudaHostAllocMapped)); + CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_expert_counter_mapped, const_cast(moe_recv_expert_counter), 0)); + for (int i = 0; i < NUM_MAX_LOCAL_EXPERTS; ++ i) + moe_recv_expert_counter[i] = -1; + + // MoE RDMA-level counter + CUDA_CHECK(cudaMallocHost(&moe_recv_rdma_counter, sizeof(int), cudaHostAllocMapped)); + CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_rdma_counter_mapped, const_cast(moe_recv_rdma_counter), 0)); + *moe_recv_rdma_counter = -1; EP_HOST_ASSERT(max_experts_per_rank > 0); m_rdma_alloc = std::make_unique(static_cast(num_rdma_bytes)); rdma_buffer_ptr = m_rdma_alloc->ptr(); @@ -126,11 +162,20 @@ void Buffer::init(int num_ranks, int num_experts_per_rank, int64_t num_rdma_byte sync_count_ptr = static_cast(m_sync_count_alloc->ptr()); CUDA_CHECK(cudaMemset(sync_buffer_ptr, 0, num_sync_buffer_bytes)); CUDA_CHECK(cudaMemset(sync_count_ptr, 0, num_sync_buffer_bytes)); + CUDA_CHECK(cudaMalloc(&local_barrier_cnt_ptr, num_sync_buffer_bytes)); + CUDA_CHECK(cudaMemset(local_barrier_cnt_ptr, 0, num_sync_buffer_bytes)); + + // Allocate barrier counters for high-throughput mode + CUDA_CHECK(cudaMalloc(&local_ht_barrier_counter, sizeof(uint64_t))); + CUDA_CHECK(cudaMemset(local_ht_barrier_counter, 0, sizeof(uint64_t))); + CUDA_CHECK(cudaMalloc(&last_ht_barrier_counter, sizeof(uint64_t))); + CUDA_CHECK(cudaMemset(last_ht_barrier_counter, 0, sizeof(uint64_t))); CUDA_CHECK(cudaDeviceSynchronize()); my_peer_info.rdma_buffer_ptr = rdma_buffer_ptr; my_peer_info.device_id = get_local_device_id(); my_peer_info.sync_buffer_ptr = sync_buffer_ptr; + my_peer_info.ht_barrier_ptr = local_ht_barrier_counter; my_peer_info.rank = rank; nixl_peer_info.resize(max_num_ranks); @@ -154,15 +199,36 @@ bool Buffer::is_available() const { return available; } +bool Buffer::is_ht_available() const { + return is_available() and num_ranks > NUM_MAX_NVL_PEERS; +} + +int Buffer::get_num_rdma_ranks() const { + return num_rdma_ranks; +} + +int Buffer::get_rdma_rank() const { + return rdma_rank; +} + +int Buffer::get_root_rdma_rank(bool global) const { + return global ? nvl_rank : 0; +} + 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 { +pybind11::bytearray Buffer::get_local_ipc_handle() const { + return {ipc_handles[nvl_rank].reserved, CUDA_IPC_HANDLE_SIZE}; +} + +torch::Tensor Buffer::get_local_buffer_tensor(const pybind11::object& dtype, int64_t offset, bool use_rdma_buffer) const { torch::ScalarType casted_dtype = torch::python::detail::py_object_to_dtype(dtype); auto element_bytes = static_cast(elementSize(casted_dtype)); - auto base_ptr = static_cast(rdma_buffer_ptr) + offset; - return torch::from_blob(base_ptr, num_rdma_bytes / element_bytes, torch::TensorOptions().dtype(casted_dtype).device(at::kCUDA)); + auto base_ptr = static_cast(use_rdma_buffer ? rdma_buffer_ptr : buffer_ptrs[nvl_rank]) + offset; + auto num_bytes = use_rdma_buffer ? num_rdma_bytes : num_nvl_bytes; + return torch::from_blob(base_ptr, num_bytes / element_bytes, torch::TensorOptions().dtype(casted_dtype).device(at::kCUDA)); } torch::Stream Buffer::get_comm_stream() const { @@ -193,6 +259,20 @@ void Buffer::destroy() { _nixl_ep_destroy(); + if (num_nvl_bytes > 0) { + intranode::barrier(barrier_signal_ptrs_gpu, nvl_rank, num_nvl_ranks, comm_stream); + CUDA_CHECK(cudaDeviceSynchronize()); + + // Close remote IPC + if (is_available()) { + for (int i = 0; i < num_nvl_ranks; ++ i) if (i != nvl_rank) + CUDA_CHECK(cudaIpcCloseMemHandle(buffer_ptrs[i])); + } + + // Free local buffer + CUDA_CHECK(cudaFree(buffer_ptrs[nvl_rank])); + } + if (nixl_agent_info and nixl_agent_info->agent != nullptr) { if (getenv("NIXL_ETCD_ENDPOINTS")) { warn_nixl(nixl_agent_info->agent->invalidateLocalMD(), @@ -223,9 +303,17 @@ void Buffer::destroy() { m_sync_count_alloc.reset(); sync_count_ptr = nullptr; + warn_cuda(cudaFree(local_barrier_cnt_ptr), "free local barrier count"); + warn_cuda(cudaFree(local_ht_barrier_counter), "free local ht barrier counter"); + warn_cuda(cudaFree(last_ht_barrier_counter), "free last ht barrier counter"); + m_workspace_alloc.reset(); workspace = nullptr; + CUDA_CHECK(cudaFreeHost(const_cast(moe_recv_counter))); + CUDA_CHECK(cudaFreeHost(const_cast(moe_recv_expert_counter))); + CUDA_CHECK(cudaFreeHost(const_cast(moe_recv_rdma_counter))); + destroyed = true; available = false; } @@ -301,15 +389,46 @@ void Buffer::_nixl_agents_peer_info_gather(std::vector& ranks) { } } -void Buffer::connect_ranks(const std::vector& remote_ranks_list, const std::optional>& remote_mds) { +void Buffer::_ipc_handles_sync(const std::vector> &all_gathered_handles = {}) { + if (num_nvl_bytes > 0) { + EP_HOST_ASSERT(all_gathered_handles.size() == max_num_ranks); + for (int i = 0, offset = rdma_rank * num_nvl_ranks; i < num_nvl_ranks; ++ i) { + EP_HOST_ASSERT(all_gathered_handles[offset + i].has_value()); + auto handle_str = std::string(all_gathered_handles[offset + i].value()); + EP_HOST_ASSERT(handle_str.size() == CUDA_IPC_HANDLE_SIZE); + if (offset + i != rank) { + std::memcpy(ipc_handles[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE); + CUDA_CHECK(cudaIpcOpenMemHandle(&buffer_ptrs[i], ipc_handles[i], cudaIpcMemLazyEnablePeerAccess)); + barrier_signal_ptrs[i] = reinterpret_cast(static_cast(buffer_ptrs[i]) + num_nvl_bytes); + } else { + EP_HOST_ASSERT(std::memcmp(ipc_handles[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE) == 0); + } + } + + // Copy all buffer and barrier signal pointers to GPU + CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs, sizeof(void*) * NUM_MAX_NVL_PEERS, cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(barrier_signal_ptrs_gpu, barrier_signal_ptrs, sizeof(int*) * NUM_MAX_NVL_PEERS, cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaDeviceSynchronize()); + } +} + +void Buffer::connect_ranks(const std::vector& remote_ranks_list, const std::optional>& remote_mds, + const std::vector> &all_gathered_handles) { EP_HOST_ASSERT(!remote_ranks_list.empty()); EP_HOST_ASSERT(!remote_mds.has_value() || remote_mds->size() == remote_ranks_list.size()); + if (!low_latency_mode && num_nvl_bytes > 0) { + EP_HOST_ASSERT(remote_ranks.empty() && "connect_ranks called more than once in high-throughput mode; elasticity is not yet supported"); + } + std::vector new_ranks; std::vector new_ranks_mds; 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); + if (all_gathered_handles.size() > 0) + _ipc_handles_sync(all_gathered_handles); + for (size_t i = 0; i < remote_ranks_list.size(); i++) { int remote_rank = remote_ranks_list[i]; // Skip self and ranks we are already connected to @@ -319,6 +438,7 @@ void Buffer::connect_ranks(const std::vector& remote_ranks_list, const std: new_ranks.push_back(remote_rank); CUDA_CHECK(cudaMemset(mask_buffer_ptr + remote_rank, 0, sizeof(int))); CUDA_CHECK(cudaMemset(sync_count_ptr + remote_rank, 0, sizeof(int))); + CUDA_CHECK(cudaMemset(local_barrier_cnt_ptr + remote_rank, 0, sizeof(int))); CUDA_CHECK(cudaMemset(sync_buffer_ptr + remote_rank, 0, sizeof(int))); if (remote_mds.has_value()) @@ -377,6 +497,469 @@ void Buffer::disconnect_ranks(const std::vector& remote_ranks_list) { _nixl_ep_memory_views_create(); } +std::tuple, torch::Tensor, torch::Tensor, std::optional> +Buffer::get_dispatch_layout(const torch::Tensor& topk_idx, int num_experts, + std::optional& previous_event, bool async, bool allocate_on_comm_stream) { + EP_HOST_ASSERT(topk_idx.dim() == 2); + EP_HOST_ASSERT(topk_idx.is_contiguous()); + EP_HOST_ASSERT(num_experts > 0); + + // Allocate all tensors on comm stream if set + // NOTES: do not allocate tensors upfront! + auto compute_stream = at::cuda::getCurrentCUDAStream(); + if (allocate_on_comm_stream) { + EP_HOST_ASSERT(previous_event.has_value() and async); + at::cuda::setCurrentCUDAStream(comm_stream); + } + + // Wait previous tasks to be finished + if (previous_event.has_value()) { + stream_wait(comm_stream, previous_event.value()); + } else { + stream_wait(comm_stream, compute_stream); + } + + auto num_tokens = static_cast(topk_idx.size(0)), num_topk = static_cast(topk_idx.size(1)); + auto num_tokens_per_rank = torch::empty({num_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); + auto num_tokens_per_rdma_rank = std::optional(); + auto num_tokens_per_expert = torch::empty({num_experts}, dtype(torch::kInt32).device(torch::kCUDA)); + auto is_token_in_rank = torch::empty({num_tokens, num_ranks}, dtype(torch::kBool).device(torch::kCUDA)); + if (is_ht_available()) + num_tokens_per_rdma_rank = torch::empty({num_rdma_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); + + layout::get_dispatch_layout(topk_idx.data_ptr(), + num_tokens_per_rank.data_ptr(), + num_tokens_per_rdma_rank.has_value() ? num_tokens_per_rdma_rank.value().data_ptr() : nullptr, + num_tokens_per_expert.data_ptr(), + is_token_in_rank.data_ptr(), + num_tokens, num_topk, num_ranks, num_experts, + comm_stream); + + // Wait streams + std::optional event; + if (async) { + event = EventHandle(comm_stream); + for (auto& t: {topk_idx, num_tokens_per_rank, num_tokens_per_expert, is_token_in_rank}) { + t.record_stream(comm_stream); + if (allocate_on_comm_stream) + t.record_stream(compute_stream); + } + for (auto& to: {num_tokens_per_rdma_rank}) { + to.has_value() ? to->record_stream(comm_stream) : void(); + if (allocate_on_comm_stream) + to.has_value() ? to->record_stream(compute_stream) : void(); + } + } else { + stream_wait(compute_stream, comm_stream); + } + + // Switch back compute stream + if (allocate_on_comm_stream) + at::cuda::setCurrentCUDAStream(compute_stream); + + return {num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, event}; +} + +std::tuple, std::optional, std::optional, std::vector, torch::Tensor, torch::Tensor, std::optional, torch::Tensor, std::optional, torch::Tensor, std::optional, std::optional, std::optional, std::optional> +Buffer::ht_dispatch(const torch::Tensor& x, const std::optional& x_scales, + const std::optional& topk_idx, const std::optional& topk_weights, + const std::optional& num_tokens_per_rank, const std::optional& num_tokens_per_rdma_rank, + const torch::Tensor& is_token_in_rank, const std::optional& num_tokens_per_expert, + int cached_num_recv_tokens, int cached_num_rdma_recv_tokens, + const std::optional& cached_rdma_channel_prefix_matrix, const std::optional& cached_recv_rdma_rank_prefix_sum, + const std::optional& cached_gbl_channel_prefix_matrix, const std::optional& cached_recv_gbl_rank_prefix_sum, + int expert_alignment, const Config& config, std::optional& previous_event, bool async, bool allocate_on_comm_stream) { + // In dispatch, CPU will busy-wait until GPU receive tensor size metadata from other ranks, which can be quite long. + // If users of DeepEP need to execute other Python code on other threads, such as KV transfer, their code will get stuck due to GIL + // unless we release GIL here. + pybind11::gil_scoped_release release; + const int num_channels = config.num_sms / 2; + EP_HOST_ASSERT(config.num_sms % 2 == 0); + EP_HOST_ASSERT(0 < get_num_rdma_ranks() and get_num_rdma_ranks() <= NUM_MAX_RDMA_PEERS); + + bool cached_mode = cached_rdma_channel_prefix_matrix.has_value(); + if (cached_mode) { + EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix.has_value()); + EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum.has_value()); + EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix.has_value()); + EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum.has_value()); + } else { + EP_HOST_ASSERT(num_tokens_per_rank.has_value()); + EP_HOST_ASSERT(num_tokens_per_rdma_rank.has_value()); + EP_HOST_ASSERT(num_tokens_per_expert.has_value()); + } + + // Type checks + if (cached_mode) { + EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum->scalar_type() == torch::kInt32); + } else { + EP_HOST_ASSERT(num_tokens_per_rank->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(num_tokens_per_rdma_rank->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(num_tokens_per_expert->scalar_type() == torch::kInt32); + } + + // Shape and contiguous checks + EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous()); + EP_HOST_ASSERT((x.size(1) * x.element_size()) % sizeof(int4) == 0); + if (cached_mode) { + EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix->dim() == 2 and cached_rdma_channel_prefix_matrix->is_contiguous()); + EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix->size(0) == num_rdma_ranks and cached_rdma_channel_prefix_matrix->size(1) == num_channels); + EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum->dim() == 1 and cached_recv_rdma_rank_prefix_sum->is_contiguous()); + EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum->size(0) == num_rdma_ranks); + EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix->dim() == 2 and cached_gbl_channel_prefix_matrix->is_contiguous()); + EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix->size(0) == num_ranks and cached_gbl_channel_prefix_matrix->size(1) == num_channels); + EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum->dim() == 1 and cached_recv_gbl_rank_prefix_sum->is_contiguous()); + EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum->size(0) == num_ranks); + } else { + EP_HOST_ASSERT(num_tokens_per_rank->dim() == 1 and num_tokens_per_rank->is_contiguous()); + EP_HOST_ASSERT(num_tokens_per_rdma_rank->dim() == 1 and num_tokens_per_rdma_rank->is_contiguous()); + EP_HOST_ASSERT(num_tokens_per_expert->dim() == 1 and num_tokens_per_expert->is_contiguous()); + EP_HOST_ASSERT(num_tokens_per_rank->size(0) == num_ranks); + EP_HOST_ASSERT(num_tokens_per_rdma_rank->size(0) == num_rdma_ranks); + EP_HOST_ASSERT(num_tokens_per_expert->size(0) % num_ranks == 0); + EP_HOST_ASSERT(num_tokens_per_expert->size(0) / num_ranks <= NUM_MAX_LOCAL_EXPERTS); + } + + auto num_tokens = static_cast(x.size(0)), hidden = static_cast(x.size(1)), hidden_int4 = static_cast(x.size(1) * x.element_size() / sizeof(int4)); + auto num_experts = cached_mode ? 0 : static_cast(num_tokens_per_expert->size(0)), num_local_experts = num_experts / num_ranks; + + // Top-k checks + int num_topk = 0; + topk_idx_t* topk_idx_ptr = nullptr; + float* topk_weights_ptr = nullptr; + EP_HOST_ASSERT(topk_idx.has_value() == topk_weights.has_value()); + if (topk_idx.has_value()) { + num_topk = static_cast(topk_idx->size(1)); + EP_HOST_ASSERT(num_experts > 0); + EP_HOST_ASSERT(topk_idx->dim() == 2 and topk_idx->is_contiguous()); + EP_HOST_ASSERT(topk_weights->dim() == 2 and topk_weights->is_contiguous()); + EP_HOST_ASSERT(num_tokens == topk_idx->size(0) and num_tokens == topk_weights->size(0)); + EP_HOST_ASSERT(num_topk == topk_weights->size(1)); + EP_HOST_ASSERT(topk_weights->scalar_type() == torch::kFloat32); + topk_idx_ptr = topk_idx->data_ptr(); + topk_weights_ptr = topk_weights->data_ptr(); + } + + // FP8 scales checks + float* x_scales_ptr = nullptr; + int num_scales = 0, scale_token_stride = 0, scale_hidden_stride = 0; + if (x_scales.has_value()) { + EP_HOST_ASSERT(x.element_size() == 1); + EP_HOST_ASSERT(x_scales->scalar_type() == torch::kFloat32 or x_scales->scalar_type() == torch::kInt); + EP_HOST_ASSERT(x_scales->dim() == 2); + EP_HOST_ASSERT(x_scales->size(0) == num_tokens); + num_scales = static_cast(x_scales->size(1)); + x_scales_ptr = static_cast(x_scales->data_ptr()); + scale_token_stride = static_cast(x_scales->stride(0)); + scale_hidden_stride = static_cast(x_scales->stride(1)); + } + + // Allocate all tensors on comm stream if set + // NOTES: do not allocate tensors upfront! + auto compute_stream = at::cuda::getCurrentCUDAStream(); + if (allocate_on_comm_stream) { + EP_HOST_ASSERT(previous_event.has_value() and async); + at::cuda::setCurrentCUDAStream(comm_stream); + } + + // Wait previous tasks to be finished + if (previous_event.has_value()) { + stream_wait(comm_stream, previous_event.value()); + } else { + stream_wait(comm_stream, compute_stream); + } + + // Create handles (only return for non-cached mode) + int num_recv_tokens = -1, num_rdma_recv_tokens = -1; + auto rdma_channel_prefix_matrix = torch::Tensor(); + auto recv_rdma_rank_prefix_sum = torch::Tensor(); + auto gbl_channel_prefix_matrix = torch::Tensor(); + auto recv_gbl_rank_prefix_sum = torch::Tensor(); + std::vector num_recv_tokens_per_expert_list; + + // Barrier or send sizes + if (cached_mode) { + num_recv_tokens = cached_num_recv_tokens; + num_rdma_recv_tokens = cached_num_rdma_recv_tokens; + rdma_channel_prefix_matrix = cached_rdma_channel_prefix_matrix.value(); + recv_rdma_rank_prefix_sum = cached_recv_rdma_rank_prefix_sum.value(); + gbl_channel_prefix_matrix = cached_gbl_channel_prefix_matrix.value(); + recv_gbl_rank_prefix_sum = cached_recv_gbl_rank_prefix_sum.value(); + + // Just a barrier and clean flags + ht::cached_notify(hidden_int4, num_scales, num_topk, num_topk, + num_ranks, num_channels, 0, nullptr, + nullptr, nullptr, nullptr, + rdma_buffer_ptr, config.num_max_rdma_chunked_recv_tokens, + buffer_ptrs_gpu, config.num_max_nvl_chunked_recv_tokens, + barrier_signal_ptrs_gpu, rank, comm_stream, + config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks), + num_nvl_bytes, true, low_latency_mode, gpu_ctx); + } else { + rdma_channel_prefix_matrix = torch::empty({num_rdma_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); + recv_rdma_rank_prefix_sum = torch::empty({num_rdma_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); + gbl_channel_prefix_matrix = torch::empty({num_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); + recv_gbl_rank_prefix_sum = torch::empty({num_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); + + // Send sizes + *moe_recv_counter = -1, *moe_recv_rdma_counter = -1; + for (int i = 0; i < num_local_experts; ++ i) + moe_recv_expert_counter[i] = -1; + ht::notify_dispatch(num_tokens_per_rank->data_ptr(), moe_recv_counter_mapped, num_ranks, + num_tokens_per_rdma_rank->data_ptr(), moe_recv_rdma_counter_mapped, + num_tokens_per_expert->data_ptr(), moe_recv_expert_counter_mapped, num_experts, + is_token_in_rank.data_ptr(), num_tokens, num_channels, + hidden_int4, num_scales, num_topk, expert_alignment, + rdma_channel_prefix_matrix.data_ptr(), recv_rdma_rank_prefix_sum.data_ptr(), + gbl_channel_prefix_matrix.data_ptr(), recv_gbl_rank_prefix_sum.data_ptr(), + rdma_buffer_ptr, config.num_max_rdma_chunked_recv_tokens, + buffer_ptrs_gpu, config.num_max_nvl_chunked_recv_tokens, + barrier_signal_ptrs_gpu, rank, comm_stream, + config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks), + num_nvl_bytes, low_latency_mode, gpu_ctx); + + // Synchronize total received tokens and tokens per expert + auto start_time = std::chrono::high_resolution_clock::now(); + while (true) { + // Read total count + num_recv_tokens = static_cast(*moe_recv_counter); + num_rdma_recv_tokens = static_cast(*moe_recv_rdma_counter); + + // Read per-expert count + bool ready = (num_recv_tokens >= 0) and (num_rdma_recv_tokens >= 0); + for (int i = 0; i < num_local_experts and ready; ++ i) + ready &= moe_recv_expert_counter[i] >= 0; + + if (ready) + break; + + // Timeout check + if (std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - start_time).count() > NUM_CPU_TIMEOUT_SECS) { + for (int i = 0; i < num_local_experts; ++ i) + printf("moe_recv_expert_counter[%d]: %d\n", i, moe_recv_expert_counter[i]); + throw std::runtime_error("NixlEP error: timeout (dispatch CPU)"); + } + } + num_recv_tokens_per_expert_list = std::vector(moe_recv_expert_counter, moe_recv_expert_counter + num_local_experts); + } + + // Allocate new tensors + auto recv_x = torch::empty({num_recv_tokens, hidden}, x.options()); + auto recv_topk_idx = std::optional(), recv_topk_weights = std::optional(), recv_x_scales = std::optional(); + auto recv_src_meta = std::optional(); + auto recv_rdma_channel_prefix_matrix = std::optional(); + auto recv_gbl_channel_prefix_matrix = std::optional(); + auto send_rdma_head = std::optional(); + auto send_nvl_head = std::optional(); + if (not cached_mode) { + recv_src_meta = torch::empty({num_recv_tokens, ht::get_source_meta_bytes()}, dtype(torch::kByte).device(torch::kCUDA)); + recv_rdma_channel_prefix_matrix = torch::empty({num_rdma_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); + recv_gbl_channel_prefix_matrix = torch::empty({num_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); + send_rdma_head = torch::empty({num_tokens, num_rdma_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); + send_nvl_head = torch::empty({num_rdma_recv_tokens, NUM_MAX_NVL_PEERS}, dtype(torch::kInt32).device(torch::kCUDA)); + } + + // Assign pointers + topk_idx_t* recv_topk_idx_ptr = nullptr; + float* recv_topk_weights_ptr = nullptr; + float* recv_x_scales_ptr = nullptr; + if (topk_idx.has_value()) { + recv_topk_idx = torch::empty({num_recv_tokens, num_topk}, topk_idx->options()); + recv_topk_weights = torch::empty({num_recv_tokens, num_topk}, topk_weights->options()); + recv_topk_idx_ptr = recv_topk_idx->data_ptr(); + recv_topk_weights_ptr = recv_topk_weights->data_ptr(); + } + if (x_scales.has_value()) { + recv_x_scales = torch::empty({num_recv_tokens, num_scales}, x_scales->options()); + recv_x_scales_ptr = static_cast(recv_x_scales->data_ptr()); + } + + // Launch data dispatch + // NOTES: the buffer size checks are moved into the `.cu` file + ht::dispatch(recv_x.data_ptr(), recv_x_scales_ptr, recv_topk_idx_ptr, recv_topk_weights_ptr, + cached_mode ? nullptr : recv_src_meta->data_ptr(), + x.data_ptr(), x_scales_ptr, topk_idx_ptr, topk_weights_ptr, + cached_mode ? nullptr : send_rdma_head->data_ptr(), cached_mode ? nullptr : send_nvl_head->data_ptr(), + cached_mode ? nullptr : recv_rdma_channel_prefix_matrix->data_ptr(), + cached_mode ? nullptr : recv_gbl_channel_prefix_matrix->data_ptr(), + rdma_channel_prefix_matrix.data_ptr(), recv_rdma_rank_prefix_sum.data_ptr(), + gbl_channel_prefix_matrix.data_ptr(), recv_gbl_rank_prefix_sum.data_ptr(), + is_token_in_rank.data_ptr(), + num_tokens, hidden_int4, num_scales, num_topk, num_experts, + scale_token_stride, scale_hidden_stride, + rdma_buffer_ptr, config.num_max_rdma_chunked_send_tokens, config.num_max_rdma_chunked_recv_tokens, + buffer_ptrs_gpu, config.num_max_nvl_chunked_send_tokens, config.num_max_nvl_chunked_recv_tokens, + rank, num_ranks, cached_mode, + comm_stream, num_channels, low_latency_mode, gpu_ctx); + + // Wait streams + std::optional event; + if (async) { + event = EventHandle(comm_stream); + for (auto& t: {x, is_token_in_rank, recv_x, + rdma_channel_prefix_matrix, recv_rdma_rank_prefix_sum, gbl_channel_prefix_matrix, recv_gbl_rank_prefix_sum}) { + t.record_stream(comm_stream); + if (allocate_on_comm_stream) + t.record_stream(compute_stream); + } + for (auto& to: {x_scales, topk_idx, topk_weights, + num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, + cached_rdma_channel_prefix_matrix, cached_recv_rdma_rank_prefix_sum, + cached_gbl_channel_prefix_matrix, cached_recv_gbl_rank_prefix_sum, + recv_topk_idx, recv_topk_weights, recv_x_scales, + recv_rdma_channel_prefix_matrix, recv_gbl_channel_prefix_matrix, send_rdma_head, send_nvl_head, + recv_src_meta}) { + to.has_value() ? to->record_stream(comm_stream) : void(); + if (allocate_on_comm_stream) + to.has_value() ? to->record_stream(compute_stream) : void(); + } + } else { + stream_wait(compute_stream, comm_stream); + } + + // Switch back compute stream + if (allocate_on_comm_stream) + at::cuda::setCurrentCUDAStream(compute_stream); + + // Return values + return {recv_x, recv_x_scales, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, + rdma_channel_prefix_matrix, gbl_channel_prefix_matrix, + recv_rdma_channel_prefix_matrix, recv_rdma_rank_prefix_sum, + recv_gbl_channel_prefix_matrix, recv_gbl_rank_prefix_sum, + recv_src_meta, send_rdma_head, send_nvl_head, event}; +} + +std::tuple, std::optional> +Buffer::ht_combine(const torch::Tensor& x, const std::optional& topk_weights, + const std::optional& bias_0, const std::optional& bias_1, + const torch::Tensor& src_meta, const torch::Tensor& is_combined_token_in_rank, + const torch::Tensor& rdma_channel_prefix_matrix, const torch::Tensor& rdma_rank_prefix_sum, const torch::Tensor& gbl_channel_prefix_matrix, + const torch::Tensor& combined_rdma_head, const torch::Tensor& combined_nvl_head, + const Config& config, std::optional& previous_event, bool async, bool allocate_on_comm_stream) { + const int num_channels = config.num_sms / 2; + EP_HOST_ASSERT(config.num_sms % 2 == 0); + + // Shape and contiguous checks + EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous()); + EP_HOST_ASSERT(src_meta.dim() == 2 and src_meta.is_contiguous() and src_meta.scalar_type() == torch::kByte); + EP_HOST_ASSERT(is_combined_token_in_rank.dim() == 2 and is_combined_token_in_rank.is_contiguous() and is_combined_token_in_rank.scalar_type() == torch::kBool); + EP_HOST_ASSERT(rdma_channel_prefix_matrix.dim() == 2 and rdma_channel_prefix_matrix.is_contiguous() and rdma_channel_prefix_matrix.scalar_type() == torch::kInt32); + EP_HOST_ASSERT(rdma_rank_prefix_sum.dim() == 1 and rdma_rank_prefix_sum.is_contiguous() and rdma_rank_prefix_sum.scalar_type() == torch::kInt32); + EP_HOST_ASSERT(gbl_channel_prefix_matrix.dim() == 2 and gbl_channel_prefix_matrix.is_contiguous() and gbl_channel_prefix_matrix.scalar_type() == torch::kInt32); + EP_HOST_ASSERT(combined_rdma_head.dim() == 2 and combined_rdma_head.is_contiguous() and combined_rdma_head.scalar_type() == torch::kInt32); + EP_HOST_ASSERT(combined_nvl_head.dim() == 2 and combined_nvl_head.is_contiguous() and combined_nvl_head.scalar_type() == torch::kInt32); + + auto num_tokens = static_cast(x.size(0)), hidden = static_cast(x.size(1)), hidden_int4 = static_cast(x.size(1) * x.element_size() / sizeof(int4)); + auto num_combined_tokens = static_cast(is_combined_token_in_rank.size(0)); + EP_HOST_ASSERT((hidden * x.element_size()) % sizeof(int4) == 0); + EP_HOST_ASSERT(src_meta.size(1) == ht::get_source_meta_bytes()); + EP_HOST_ASSERT(is_combined_token_in_rank.size(1) == num_ranks); + EP_HOST_ASSERT(rdma_channel_prefix_matrix.size(0) == num_rdma_ranks and rdma_channel_prefix_matrix.size(1) == num_channels); + EP_HOST_ASSERT(rdma_rank_prefix_sum.size(0) == num_rdma_ranks); + EP_HOST_ASSERT(gbl_channel_prefix_matrix.size(0) == num_ranks and gbl_channel_prefix_matrix.size(1) == num_channels); + EP_HOST_ASSERT(combined_rdma_head.dim() == 2 and combined_rdma_head.size(0) == num_combined_tokens and combined_rdma_head.size(1) == num_rdma_ranks); + EP_HOST_ASSERT(combined_nvl_head.dim() == 2 and combined_nvl_head.size(1) == NUM_MAX_NVL_PEERS); + + // Allocate all tensors on comm stream if set + // NOTES: do not allocate tensors upfront! + auto compute_stream = at::cuda::getCurrentCUDAStream(); + if (allocate_on_comm_stream) { + EP_HOST_ASSERT(previous_event.has_value() and async); + at::cuda::setCurrentCUDAStream(comm_stream); + } + + // Wait previous tasks to be finished + if (previous_event.has_value()) { + stream_wait(comm_stream, previous_event.value()); + } else { + stream_wait(comm_stream, compute_stream); + } + + // Top-k checks + int num_topk = 0; + auto combined_topk_weights = std::optional(); + float* topk_weights_ptr = nullptr; + float* combined_topk_weights_ptr = nullptr; + if (topk_weights.has_value()) { + EP_HOST_ASSERT(topk_weights->dim() == 2 and topk_weights->is_contiguous()); + EP_HOST_ASSERT(topk_weights->size(0) == num_tokens); + EP_HOST_ASSERT(topk_weights->scalar_type() == torch::kFloat32); + num_topk = static_cast(topk_weights->size(1)); + topk_weights_ptr = topk_weights->data_ptr(); + combined_topk_weights = torch::empty({num_combined_tokens, num_topk}, topk_weights->options()); + combined_topk_weights_ptr = combined_topk_weights->data_ptr(); + } + + // Extra check for avoid-dead-lock design + EP_HOST_ASSERT(config.num_max_nvl_chunked_recv_tokens % num_rdma_ranks == 0); + EP_HOST_ASSERT(config.num_max_nvl_chunked_send_tokens <= config.num_max_nvl_chunked_recv_tokens / num_rdma_ranks); + + // Launch barrier and reset queue head and tail + ht::cached_notify(hidden_int4, 0, 0, num_topk, + num_ranks, num_channels, + num_combined_tokens, combined_rdma_head.data_ptr(), + rdma_channel_prefix_matrix.data_ptr(), rdma_rank_prefix_sum.data_ptr(), combined_nvl_head.data_ptr(), + rdma_buffer_ptr, config.num_max_rdma_chunked_recv_tokens, + buffer_ptrs_gpu, config.num_max_nvl_chunked_recv_tokens, + barrier_signal_ptrs_gpu, rank, comm_stream, + config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks), + num_nvl_bytes, false, low_latency_mode, gpu_ctx); + + // Assign bias pointers + auto bias_opts = std::vector>({bias_0, bias_1}); + void* bias_ptrs[2] = {nullptr, nullptr}; + for (int i = 0; i < 2; ++ i) if (bias_opts[i].has_value()) { + auto bias = bias_opts[i].value(); + EP_HOST_ASSERT(bias.dim() == 2 and bias.is_contiguous()); + EP_HOST_ASSERT(bias.scalar_type() == x.scalar_type()); + EP_HOST_ASSERT(bias.size(0) == num_combined_tokens and bias.size(1) == hidden); + bias_ptrs[i] = bias.data_ptr(); + } + + // Launch data combine + auto combined_x = torch::empty({num_combined_tokens, hidden}, x.options()); + ht::combine(at::cuda::ScalarTypeToCudaDataType(x.scalar_type()), + combined_x.data_ptr(), combined_topk_weights_ptr, + is_combined_token_in_rank.data_ptr(), + x.data_ptr(), topk_weights_ptr, bias_ptrs[0], bias_ptrs[1], + combined_rdma_head.data_ptr(), combined_nvl_head.data_ptr(), + src_meta.data_ptr(), rdma_channel_prefix_matrix.data_ptr(), rdma_rank_prefix_sum.data_ptr(), gbl_channel_prefix_matrix.data_ptr(), + num_tokens, num_combined_tokens, hidden, num_topk, + rdma_buffer_ptr, config.num_max_rdma_chunked_send_tokens, config.num_max_rdma_chunked_recv_tokens, + buffer_ptrs_gpu, config.num_max_nvl_chunked_send_tokens, config.num_max_nvl_chunked_recv_tokens, + rank, num_ranks, comm_stream, num_channels, low_latency_mode, gpu_ctx); + + // Wait streams + std::optional event; + if (async) { + event = EventHandle(comm_stream); + for (auto& t: {x, src_meta, + is_combined_token_in_rank, rdma_channel_prefix_matrix, rdma_rank_prefix_sum, gbl_channel_prefix_matrix, + combined_x, combined_rdma_head, combined_nvl_head}) { + t.record_stream(comm_stream); + if (allocate_on_comm_stream) + t.record_stream(compute_stream); + } + for (auto& to: {topk_weights, combined_topk_weights, bias_0, bias_1}) { + to.has_value() ? to->record_stream(comm_stream) : void(); + if (allocate_on_comm_stream) + to.has_value() ? to->record_stream(compute_stream) : void(); + } + } else { + stream_wait(compute_stream, comm_stream); + } + + // Switch back compute stream + if (allocate_on_comm_stream) + at::cuda::setCurrentCUDAStream(compute_stream); + + // Return values + return {combined_x, combined_topk_weights, event}; +} + std::tuple, torch::Tensor, torch::Tensor, torch::Tensor, std::optional, std::optional>> Buffer::dispatch(const torch::Tensor& x, const torch::Tensor& topk_idx, const std::optional& cumulative_local_expert_recv_stats, @@ -646,7 +1229,7 @@ std::string Buffer::get_local_metadata() const { void Buffer::_nixl_ep_memory_views_create(void) { nixl_remote_dlist_t remote_descs(VRAM_SEG); nixl_remote_dlist_t barrier_descs(VRAM_SEG); - nixl_xfer_dlist_t local_descs(VRAM_SEG); + nixl_local_dlist_t local_descs(VRAM_SEG); local_descs.addDesc(nixlBlobDesc((uintptr_t)(rdma_buffer_ptr), num_rdma_bytes, get_local_device_id(), "")); local_descs.addDesc(nixlBlobDesc((uintptr_t)(sync_count_ptr), max_num_ranks * sizeof(int), get_local_device_id(), "")); @@ -662,6 +1245,15 @@ void Buffer::_nixl_ep_memory_views_create(void) { if (!remote_ranks.empty()) { EP_HOST_ASSERT(nixl_agent_info->agent->prepMemView(remote_descs, gpu_ctx.remote_mvh, &nixl_agent_info->extra_params) == NIXL_SUCCESS); EP_HOST_ASSERT(nixl_agent_info->agent->prepMemView(barrier_descs, gpu_ctx.barrier_mvh, &nixl_agent_info->extra_params) == NIXL_SUCCESS); + + if (!low_latency_mode && num_ranks > NUM_MAX_NVL_PEERS) { + nixl_remote_dlist_t ht_barrier_descs(VRAM_SEG); + for (int r = 0; r < max_num_ranks; r++) { + std::string remote_agent_name = remote_set.count(r) ? nixl_agent_info->remote_agent_names[r] : nixl_null_agent; + ht_barrier_descs.addDesc(nixlRemoteDesc((uintptr_t)nixl_peer_info[r].ht_barrier_ptr, sizeof(uint64_t), nixl_peer_info[r].device_id, remote_agent_name)); + } + EP_HOST_ASSERT(nixl_agent_info->agent->prepMemView(ht_barrier_descs, gpu_ctx.ht_barrier_mvh, &nixl_agent_info->extra_params) == NIXL_SUCCESS); + } } } @@ -669,17 +1261,22 @@ void Buffer::_nixl_ep_memory_views_destroy(void) { if (gpu_ctx.local_mvh) nixl_agent_info->agent->releaseMemView(gpu_ctx.local_mvh); if (gpu_ctx.remote_mvh) nixl_agent_info->agent->releaseMemView(gpu_ctx.remote_mvh); if (gpu_ctx.barrier_mvh) nixl_agent_info->agent->releaseMemView(gpu_ctx.barrier_mvh); + if (gpu_ctx.ht_barrier_mvh) nixl_agent_info->agent->releaseMemView(gpu_ctx.ht_barrier_mvh); gpu_ctx.local_mvh = nullptr; gpu_ctx.remote_mvh = nullptr; gpu_ctx.barrier_mvh = nullptr; + gpu_ctx.ht_barrier_mvh = nullptr; } void Buffer::_nixl_ep_init(void) { gpu_ctx = { .sync_buffer_ptr = sync_buffer_ptr, .sync_count_ptr = sync_count_ptr, + .last_ht_barrier_counter = last_ht_barrier_counter, + .local_ht_barrier_counter_ptr = local_ht_barrier_counter, .rdma_buffer_ptr = rdma_buffer_ptr, .max_num_ranks = max_num_ranks, + .num_rdma_ranks = num_rdma_ranks, .rank = rank, }; } @@ -706,7 +1303,6 @@ void Buffer::_nixl_agent_init() { ", status: " + std::to_string(status)); } - // Set UCX-specific parameters const char* num_channels_env = std::getenv("NIXL_EP_NUM_CHANNELS"); init_params["ucx_num_device_channels"] = num_channels_env ? num_channels_env : "4"; init_params["ucx_error_handling_mode"] = "none"; @@ -739,6 +1335,12 @@ void Buffer::_nixl_agent_init() { EP_HOST_ASSERT(agent->registerMem(nixl_agent_info->sync_reg_descs, &nixl_agent_info->extra_params) == NIXL_SUCCESS); EP_HOST_ASSERT(agent->registerMem(nixl_agent_info->sync_count_reg_descs, &nixl_agent_info->extra_params) == NIXL_SUCCESS); + if (!low_latency_mode && local_ht_barrier_counter) { + nixl_reg_dlist_t ht_barrier_dlist(VRAM_SEG); + ht_barrier_dlist.addDesc(nixlBlobDesc((uintptr_t)(local_ht_barrier_counter), sizeof(uint64_t), get_local_device_id(), "")); + EP_HOST_ASSERT(agent->registerMem(ht_barrier_dlist) == NIXL_SUCCESS); + } + if (getenv("NIXL_ETCD_ENDPOINTS")) { status = nixl_agent_info->agent->sendLocalMD(); if (status != NIXL_SUCCESS) { @@ -788,27 +1390,42 @@ static std::optional> convert_mds(const std::optional(m, "Config") + .def(pybind11::init(), + py::arg("num_sms") = 20, + py::arg("num_max_nvl_chunked_send_tokens") = 6, py::arg("num_max_nvl_chunked_recv_tokens") = 256, + py::arg("num_max_rdma_chunked_send_tokens") = 6, py::arg("num_max_rdma_chunked_recv_tokens") = 256) + .def("get_nvl_buffer_size_hint", &nixl_ep::Config::get_nvl_buffer_size_hint) + .def("get_rdma_buffer_size_hint", &nixl_ep::Config::get_rdma_buffer_size_hint); pybind11::class_(m, "EventHandle") .def(pybind11::init<>()) .def("current_stream_wait", &nixl_ep::EventHandle::current_stream_wait); pybind11::class_(m, "Buffer") - .def(pybind11::init()) + .def(pybind11::init()) .def("update_memory_buffers", &nixl_ep::Buffer::update_memory_buffers) .def("barrier", &nixl_ep::Buffer::barrier) - .def("connect_ranks", [](nixl_ep::Buffer &buffer, const std::vector& remote_ranks, const std::optional>& remote_mds) { - buffer.connect_ranks(remote_ranks, nixl_ep::convert_mds(remote_mds)); - }, py::arg("remote_ranks"), py::arg("remote_mds") = std::nullopt) + .def("connect_ranks", [](nixl_ep::Buffer &buffer, const std::vector& remote_ranks, const std::optional>& remote_mds, const std::vector> &all_gathered_handles) { + buffer.connect_ranks(remote_ranks, nixl_ep::convert_mds(remote_mds), all_gathered_handles); + }, py::arg("remote_ranks"), py::arg("remote_mds") = std::nullopt, py::arg("ipc_handles") = std::vector>{}) .def("disconnect_ranks", &nixl_ep::Buffer::disconnect_ranks) .def("is_available", &nixl_ep::Buffer::is_available) + .def("get_num_rdma_ranks", &nixl_ep::Buffer::get_num_rdma_ranks) + .def("get_rdma_rank", &nixl_ep::Buffer::get_rdma_rank) + .def("get_root_rdma_rank", &nixl_ep::Buffer::get_root_rdma_rank) .def("get_local_device_id", &nixl_ep::Buffer::get_local_device_id) + .def("get_local_ipc_handle", &nixl_ep::Buffer::get_local_ipc_handle) .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("get_dispatch_layout", &nixl_ep::Buffer::get_dispatch_layout) .def("dispatch", &nixl_ep::Buffer::dispatch) .def("combine", &nixl_ep::Buffer::combine) + .def("ht_dispatch", &nixl_ep::Buffer::ht_dispatch) + .def("ht_combine", &nixl_ep::Buffer::ht_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) diff --git a/examples/device/ep/csrc/nixl_ep.hpp b/examples/device/ep/csrc/nixl_ep.hpp index 1053ab9d..0c58cde8 100644 --- a/examples/device/ep/csrc/nixl_ep.hpp +++ b/examples/device/ep/csrc/nixl_ep.hpp @@ -53,6 +53,7 @@ namespace nixl_ep { struct NixlPeerInfo { void* rdma_buffer_ptr; int* sync_buffer_ptr; + uint64_t* ht_barrier_ptr; int device_id; int rank; }; @@ -76,8 +77,16 @@ struct NixlAgentInfo }; struct Buffer { + EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS == 8, "The number of maximum NVLink peers must be 8"); + private: int buffer_idx = 0; // Double buffering index + bool low_latency_mode = false; + + // NVLink Buffer + int64_t num_nvl_bytes; + void* buffer_ptrs[NUM_MAX_NVL_PEERS] = {nullptr}; + void** buffer_ptrs_gpu = nullptr; // RDMA Buffer int64_t num_rdma_bytes; @@ -86,6 +95,7 @@ struct Buffer { int *mask_buffer_ptr = nullptr; int *sync_buffer_ptr = nullptr; int *sync_count_ptr = nullptr; + int *local_barrier_cnt_ptr = nullptr; /* Owning VMM allocations (keep raw ptrs above as aliases) */ std::unique_ptr m_rdma_alloc; @@ -97,9 +107,10 @@ struct Buffer { // Device info and communication int device_id; int num_device_sms; - int rank; - int num_ranks; + int rank, rdma_rank, nvl_rank; + int num_ranks, num_rdma_ranks, num_nvl_ranks; std::vector remote_ranks; /* global ranks */ + cudaIpcMemHandle_t ipc_handles[NUM_MAX_NVL_PEERS]; // Stream for communication at::cuda::CUDAStream comm_stream; @@ -112,15 +123,33 @@ struct Buffer { // After `destroy()` be called, this flag will be true bool destroyed = false; + // Barrier signals + int* barrier_signal_ptrs[NUM_MAX_NVL_PEERS] = {nullptr}; + int** barrier_signal_ptrs_gpu = nullptr; + // Workspace void* workspace = nullptr; + // Host-side MoE info + volatile int* moe_recv_counter = nullptr; + int* moe_recv_counter_mapped = nullptr; + + // Host-side expert-level MoE info + volatile int* moe_recv_expert_counter = nullptr; + int* moe_recv_expert_counter_mapped = nullptr; + + // Host-side RDMA-level MoE info + volatile int* moe_recv_rdma_counter = nullptr; + int* moe_recv_rdma_counter_mapped = nullptr; + std::unique_ptr nixl_agent_info; std::vector nixl_peer_info; NixlPeerInfo my_peer_info; int max_num_ranks; int max_experts_per_rank; - ep_kernels::gpu_nixl_ctx gpu_ctx; + nixl_ep::gpu_nixl_ctx gpu_ctx; + uint64_t* last_ht_barrier_counter = nullptr; + uint64_t* local_ht_barrier_counter = nullptr; /* Common private funcs */ void _nixl_agent_init(); @@ -134,29 +163,65 @@ struct Buffer { void _nixl_ep_memory_views_destroy(void); void _nixl_ep_destroy(void); + /* high-throughput mode private funcs */ + void _ipc_handles_sync(const std::vector> &all_gathered_handles); + public: - Buffer(int rank, bool explicitly_destroy); + Buffer(int rank, bool explicitly_destroy, bool low_latency_mode = true); - void update_memory_buffers(int num_ranks, int max_experts_per_rank, int64_t num_rdma_bytes); + void update_memory_buffers(int num_ranks, int max_experts_per_rank, int64_t num_rdma_bytes, int64_t num_nvl_bytes = 0); - void connect_ranks(const std::vector& remote_ranks_list, const std::optional>& remote_mds = std::nullopt); + void connect_ranks(const std::vector& remote_ranks_list, const std::optional>& remote_mds = std::nullopt, const std::vector>& all_gathered_handles = {}); void disconnect_ranks(const std::vector& remote_ranks_list); - void init(int num_ranks, int max_experts_per_rank, int64_t num_rdma_bytes); + void init(int num_ranks, int max_experts_per_rank, int64_t num_nvl_bytes, int64_t num_rdma_bytes); ~Buffer() noexcept; bool is_available() const; + bool is_ht_available() const; + + int get_num_rdma_ranks() const; + + int get_rdma_rank() const; + + int get_root_rdma_rank(bool global) const; + int get_local_device_id() const; - torch::Tensor get_local_buffer_tensor(const pybind11::object& dtype, int64_t offset) const; + pybind11::bytearray get_local_ipc_handle() const; + + torch::Tensor get_local_buffer_tensor(const pybind11::object& dtype, int64_t offset, bool use_rdma_buffer = false) const; torch::Stream get_comm_stream() const; + void destroy(); + std::tuple, torch::Tensor, torch::Tensor, std::optional> + get_dispatch_layout(const torch::Tensor& topk_idx, int num_experts, std::optional& previous_event, + bool async, bool allocate_on_comm_stream); + + std::tuple, std::optional, std::optional, std::vector, torch::Tensor, torch::Tensor, std::optional, torch::Tensor, std::optional, torch::Tensor, std::optional, std::optional, std::optional, std::optional> + ht_dispatch(const torch::Tensor& x, const std::optional& x_scales, + const std::optional& topk_idx, const std::optional& topk_weights, + const std::optional& num_tokens_per_rank, const std::optional& num_tokens_per_rdma_rank, + const torch::Tensor& is_token_in_rank, const std::optional& num_tokens_per_expert, + int cached_num_recv_tokens, int cached_num_rdma_recv_tokens, + const std::optional& cached_rdma_channel_prefix_matrix, const std::optional& cached_recv_rdma_rank_prefix_sum, + const std::optional& cached_gbl_channel_prefix_matrix, const std::optional& cached_recv_gbl_rank_prefix_sum, + int expert_alignment, const Config& config, std::optional& previous_event, bool async, bool allocate_on_comm_stream); + + std::tuple, std::optional> + ht_combine(const torch::Tensor& x, const std::optional& topk_weights, + const std::optional& bias_0, const std::optional& bias_1, + const torch::Tensor& src_meta, const torch::Tensor& is_combined_token_in_rank, + const torch::Tensor& rdma_channel_prefix_matrix, const torch::Tensor& rdma_rank_prefix_sum, const torch::Tensor& gbl_channel_prefix_matrix, + const torch::Tensor& combined_rdma_head, const torch::Tensor& combined_nvl_head, + const Config& config, std::optional& previous_event, bool async, bool allocate_on_comm_stream); + void clean_buffer(int num_max_dispatch_tokens_per_rank, int hidden, int num_experts); std::tuple, torch::Tensor, torch::Tensor, torch::Tensor, std::optional, std::optional>> diff --git a/examples/device/ep/meson.build b/examples/device/ep/meson.build index 52956f72..a9ba19fb 100644 --- a/examples/device/ep/meson.build +++ b/examples/device/ep/meson.build @@ -64,7 +64,10 @@ endif nixl_ep_sources = [ 'csrc/nixl_ep.cpp', 'csrc/vmm.cpp', - 'csrc/kernels/nixl_ep.cu', + 'csrc/kernels/nixl_ep_ll.cu', + 'csrc/kernels/nixl_ep_ht.cu', + 'csrc/kernels/layout.cu', + 'csrc/kernels/runtime.cu', ] nixl_ep_inc_dirs = [ diff --git a/examples/device/ep/nixl_ep/__init__.py b/examples/device/ep/nixl_ep/__init__.py index 719e149b..488ad2fc 100644 --- a/examples/device/ep/nixl_ep/__init__.py +++ b/examples/device/ep/nixl_ep/__init__.py @@ -25,5 +25,6 @@ from .utils import EventOverlap topk_idx_t = getattr(_nixl_ep_cpp, "topk_idx_t", torch.int64) +Config = _nixl_ep_cpp.Config -__all__ = ["Buffer", "EventOverlap"] +__all__ = ["Buffer", "EventOverlap", "Config"] diff --git a/examples/device/ep/nixl_ep/buffer.py b/examples/device/ep/nixl_ep/buffer.py index d3610a00..b7dfb278 100644 --- a/examples/device/ep/nixl_ep/buffer.py +++ b/examples/device/ep/nixl_ep/buffer.py @@ -30,7 +30,7 @@ from . import nixl_ep_cpp # noinspection PyUnresolvedReferences -from .nixl_ep_cpp import EventHandle +from .nixl_ep_cpp import Config, EventHandle from .utils import EventOverlap if TYPE_CHECKING: @@ -55,6 +55,7 @@ def __init__( disable_ll_nvlink: bool = False, explicitly_destroy: bool = False, rank: int = 0, + low_latency_mode: bool = True, group: Optional[dist.ProcessGroup] = None, comm: Optional["mpi4py.MPI.Comm"] = None, tcp_store_group: Optional[dist.TCPStore] = None, @@ -68,12 +69,15 @@ def __init__( 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. + low_latency_mode: whether to enable low-latency mode. group: the communication group (optional). comm: the mpi4py.MPI.Comm communicator to use in case the group parameter is absent (optional). tcp_store_group: TCPStore for metadata exchange (optional). """ self.rank = rank self.group_size = 0 # Will be updated by `update_memory_buffers` + self.low_latency_mode = low_latency_mode + self.explicitly_destroy = explicitly_destroy self.group = group self.comm = comm @@ -83,7 +87,9 @@ def __init__( if disable_ll_nvlink: os.environ["UCX_TLS"] = "^cuda_ipc" - self.runtime = nixl_ep_cpp.Buffer(self.rank, explicitly_destroy) + self.runtime = nixl_ep_cpp.Buffer( + self.rank, explicitly_destroy, low_latency_mode + ) def destroy(self): """ @@ -121,14 +127,14 @@ def capture() -> EventOverlap: return EventOverlap(EventHandle()) @staticmethod - def get_rdma_size_hint( + def get_low_latency_buffer_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. + Get a minimum buffer size requirement for low-latency mode. 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. @@ -139,7 +145,19 @@ def get_rdma_size_hint( Returns: size: the RDMA buffer size recommended. """ - return nixl_ep_cpp.get_rdma_size_hint( + return nixl_ep_cpp.get_low_latency_buffer_size_hint( + num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts + ) + + @staticmethod + def get_rdma_size_hint( + num_max_dispatch_tokens_per_rank: int, + hidden: int, + num_ranks: int, + num_experts: int, + ) -> int: + """Backward-compatible alias for get_low_latency_buffer_size_hint.""" + return Buffer.get_low_latency_buffer_size_hint( num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts ) @@ -158,7 +176,11 @@ def get_comm_stream(self) -> torch.Stream: ) def get_local_buffer_tensor( - self, dtype: torch.dtype, size: Optional[torch.Size] = None, offset: int = 0 + self, + dtype: torch.dtype, + size: Optional[torch.Size] = None, + offset: int = 0, + use_rdma_buffer: bool = False, ) -> torch.Tensor: """ Get the raw buffer (slice supported) as a PyTorch tensor. @@ -167,8 +189,9 @@ def get_local_buffer_tensor( 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. + use_rdma_buffer: whether to return the RDMA buffer. """ - tensor = self.runtime.get_local_buffer_tensor(dtype, offset) + tensor = self.runtime.get_local_buffer_tensor(dtype, offset, use_rdma_buffer) if size is None: return tensor @@ -185,6 +208,113 @@ def _unpack_bias(bias: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]): bias_0, bias_1 = bias return bias_0, bias_1 + @staticmethod + def get_dispatch_config(num_ranks: int) -> Config: + """ + Get a recommended dispatch config. + + Argument: + num_ranks: the number of ranks. + + Returns: + config: the recommended config. + """ + + # TODO: automatically tune + config_map = { + 2: Config(Buffer.num_sms, 24, 256, 6, 128), + 4: Config(Buffer.num_sms, 6, 256, 6, 128), + 8: Config(Buffer.num_sms, 6, 256, 6, 128), + 16: Config(Buffer.num_sms, 36, 288, 20, 128), + 24: Config(Buffer.num_sms, 8, 288, 32, 128), + 32: Config(Buffer.num_sms, 32, 288, 32, 128), + 64: Config(Buffer.num_sms, 20, 288, 28, 128), + 128: Config(Buffer.num_sms, 20, 560, 32, 128), + 144: Config(Buffer.num_sms, 32, 720, 12, 128), + 160: Config(Buffer.num_sms, 28, 720, 12, 128), + } + assert num_ranks in config_map, f"Unsupported number of EP ranks: {num_ranks}" + return config_map[num_ranks] + + @staticmethod + def get_combine_config(num_ranks: int) -> Config: + """ + Get a recommended combine config. + + Argument: + num_ranks: the number of ranks. + + Returns: + config: the recommended config. + """ + + # TODO: automatically tune + config_map = { + 2: Config(Buffer.num_sms, 10, 256, 6, 128), + 4: Config(Buffer.num_sms, 9, 256, 6, 128), + 8: Config(Buffer.num_sms, 4, 256, 6, 128), + 16: Config(Buffer.num_sms, 4, 288, 12, 128), + 24: Config(Buffer.num_sms, 1, 288, 8, 128), + 32: Config(Buffer.num_sms, 1, 288, 8, 128), + 64: Config(Buffer.num_sms, 1, 288, 20, 128), + 128: Config(Buffer.num_sms, 1, 560, 12, 128), + 144: Config(Buffer.num_sms, 2, 720, 8, 128), + 160: Config(Buffer.num_sms, 2, 720, 8, 128), + } + assert num_ranks in config_map, f"Unsupported number of EP ranks: {num_ranks}" + return config_map[num_ranks] + + # noinspection PyTypeChecker + def get_dispatch_layout( + self, + topk_idx: torch.Tensor, + num_experts: int, + previous_event: Optional[EventOverlap] = None, + async_finish: bool = False, + allocate_on_comm_stream: bool = False, + ) -> Tuple[ + torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor, EventOverlap + ]: + """ + Calculate the layout required for later communication. + + Arguments: + topk_idx: `[num_tokens, num_topk]`, dtype must be `torch.int64`, the expert indices selected by each token, + `-1` means no selections. + num_experts: the number of experts. + previous_event: the event to wait before actually executing the kernel. + async_finish: the current stream will not wait for the communication kernels to be finished if set. + allocate_on_comm_stream: control whether all the allocated tensors' ownership to be on the communication stream. + + Returns: + num_tokens_per_rank: `[num_ranks]` with `torch.int`, the number of tokens to be sent to each rank. + num_tokens_per_rdma_rank: `[num_rdma_ranks]` with `torch.int`, the number of tokens to be sent to each RDMA + rank (with the same GPU index), return `None` for intranode settings. + num_tokens_per_expert: `[num_experts]` with `torch.int`, the number of tokens to be sent to each expert. + is_token_in_rank: `[num_tokens, num_ranks]` with `torch.bool`, whether a token be sent to a rank. + event: the event after executing the kernel (valid only if `async_finish` is set). + """ + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + event, + ) = self.runtime.get_dispatch_layout( + topk_idx, + num_experts, + getattr(previous_event, "event", None), + async_finish, + allocate_on_comm_stream, + ) + return ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + EventOverlap(event), + ) + def clean_buffer( self, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int ) -> None: @@ -200,7 +330,7 @@ def clean_buffer( self.runtime.clean_buffer(num_max_dispatch_tokens_per_rank, hidden, num_experts) # noinspection PyTypeChecker - def dispatch( + def low_latency_dispatch( self, x: torch.Tensor, topk_idx: torch.Tensor, @@ -308,8 +438,11 @@ def dispatch( hook, ) + def dispatch(self, *args, **kwargs): + return self.low_latency_dispatch(*args, **kwargs) + # noinspection PyTypeChecker - def combine( + def low_latency_combine( self, x: torch.Tensor, topk_idx: torch.Tensor, @@ -390,6 +523,200 @@ def combine( hook, ) + def combine(self, *args, **kwargs): + return self.low_latency_combine(*args, **kwargs) + + # noinspection PyTypeChecker + def ht_dispatch( + self, + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + handle: Optional[Tuple] = None, + num_tokens_per_rank: Optional[torch.Tensor] = None, + num_tokens_per_rdma_rank: Optional[torch.Tensor] = None, + is_token_in_rank: Optional[torch.Tensor] = None, + num_tokens_per_expert: Optional[torch.Tensor] = None, + topk_idx: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + expert_alignment: int = 1, + config: Optional[Config] = None, + previous_event: Optional[EventOverlap] = None, + async_finish: bool = False, + allocate_on_comm_stream: bool = False, + ) -> Tuple[ + Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + List[int], + Tuple, + EventOverlap, + ]: + """ + Internode dispatch implementation, for more details, please refer to the `dispatch` docs. + Normally, you should not directly call this function. + """ + config = self.get_dispatch_config(self.group_size) if config is None else config + assert config is not None + + # Launch the kernel with cached or non-cached mode + x, x_scales = x if isinstance(x, tuple) else (x, None) + if handle is not None: + assert topk_idx is None and topk_weights is None + ( + is_token_in_rank, + rdma_channel_prefix_matrix, + gbl_channel_prefix_matrix, + recv_rdma_channel_prefix_matrix, + recv_rdma_rank_prefix_sum, + recv_gbl_channel_prefix_matrix, + recv_gbl_rank_prefix_sum, + recv_src_meta, + send_rdma_head, + send_nvl_head, + ) = handle + num_recv_tokens = recv_src_meta.size(0) + num_rdma_recv_tokens = send_nvl_head.size(0) + recv_x, recv_x_scales, _, _, _, _, _, _, _, _, _, _, _, _, event = ( + self.runtime.ht_dispatch( + x, + x_scales, + topk_idx, + topk_weights, + None, + None, + is_token_in_rank, + None, + num_recv_tokens, + num_rdma_recv_tokens, + rdma_channel_prefix_matrix, + recv_rdma_rank_prefix_sum, + gbl_channel_prefix_matrix, + recv_gbl_rank_prefix_sum, + expert_alignment, + config, + getattr(previous_event, "event", None), + async_finish, + allocate_on_comm_stream, + ) + ) + return (recv_x, recv_x_scales) if x_scales is not None else recv_x, None, None, None, None, EventOverlap(event) # type: ignore[return-value] + else: + assert ( + num_tokens_per_rank is not None + and is_token_in_rank is not None + and num_tokens_per_expert is not None + ) + ( + recv_x, + recv_x_scales, + recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert_list, + rdma_channel_prefix_matrix, + gbl_channel_prefix_matrix, + recv_rdma_channel_prefix_matrix, + recv_rdma_rank_prefix_sum, + recv_gbl_channel_prefix_matrix, + recv_gbl_rank_prefix_sum, + recv_src_meta, + send_rdma_head, + send_nvl_head, + event, + ) = self.runtime.ht_dispatch( + x, + x_scales, + topk_idx, + topk_weights, + num_tokens_per_rank, + num_tokens_per_rdma_rank, + is_token_in_rank, + num_tokens_per_expert, + 0, + 0, + None, + None, + None, + None, + expert_alignment, + config, + getattr(previous_event, "event", None), + async_finish, + allocate_on_comm_stream, + ) + handle = ( + is_token_in_rank, + rdma_channel_prefix_matrix, + gbl_channel_prefix_matrix, + recv_rdma_channel_prefix_matrix, + recv_rdma_rank_prefix_sum, + recv_gbl_channel_prefix_matrix, + recv_gbl_rank_prefix_sum, + recv_src_meta, + send_rdma_head, + send_nvl_head, + ) + return ( + (recv_x, recv_x_scales) if x_scales is not None else recv_x, + recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert_list, + handle, + EventOverlap(event), + ) + + # noinspection PyTypeChecker + def ht_combine( + self, + x: torch.Tensor, + handle: Union[tuple, list], + topk_weights: Optional[torch.Tensor] = None, + bias: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] = None, + config: Optional[Config] = None, + previous_event: Optional[EventOverlap] = None, + async_finish: bool = False, + allocate_on_comm_stream: bool = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], EventOverlap]: + """ + Internode combine implementation, for more details, please refer to the `combine` docs. + Normally, you should not directly call this function. + """ + config = self.get_combine_config(self.group_size) if config is None else config + assert config is not None + + # Unpack handle and bias + ( + is_combined_token_in_rank, + _, + _, + rdma_channel_prefix_matrix, + rdma_rank_prefix_sum, + gbl_channel_prefix_matrix, + gbl_rank_prefix_sum, + src_meta, + send_rdma_head, + send_nvl_head, + ) = handle + bias_0, bias_1 = Buffer._unpack_bias(bias) + + # Launch the kernel + combined_x, combined_topk_weights, event = self.runtime.ht_combine( + x, + topk_weights, + bias_0, + bias_1, + src_meta, + is_combined_token_in_rank, + rdma_channel_prefix_matrix, + rdma_rank_prefix_sum, + gbl_channel_prefix_matrix, + send_rdma_head, + send_nvl_head, + config, + getattr(previous_event, "event", None), + async_finish, + allocate_on_comm_stream, + ) + return combined_x, combined_topk_weights, EventOverlap(event) + def update_mask_buffer(self, rank_to_mask: int, mask: bool = False): """ Mask (unmask) a rank during communication (dispatch, combine, and clean) @@ -442,7 +769,11 @@ def get_next_combine_buffer( ) def update_memory_buffers( - self, num_ranks: int, num_experts_per_rank: int, num_rdma_bytes: int + self, + num_ranks: int, + num_experts_per_rank: int, + num_rdma_bytes: int, + num_nvl_bytes: int = 0, ): """ Allocate remote memory for the communication buffer. @@ -451,11 +782,13 @@ def update_memory_buffers( 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. + num_nvl_bytes: the buffer size for intranode NVLink communication (default: 0). """ self.group_size = num_ranks + self.num_nvl_bytes = num_nvl_bytes self.num_rdma_bytes = num_rdma_bytes self.runtime.update_memory_buffers( - num_ranks, num_experts_per_rank, num_rdma_bytes + num_ranks, num_experts_per_rank, num_rdma_bytes, num_nvl_bytes ) def set_tcp_store_group(self, tcp_store_group: Optional[dist.TCPStore]) -> None: @@ -486,6 +819,31 @@ def _fetch_remote_metadata_from_tcp_store(self, remote_ranks: List[int]): finally: self.tcp_store_group.delete_key(md_key) + def _ht_connect_ranks(self, remote_ranks: List[int]) -> None: + if self.group is not None: + + def all_gather_object(obj): + object_list = [None] * self.group_size + dist.all_gather_object(object_list, obj, self.group) + return object_list + + elif self.comm is not None: + + def all_gather_object(obj): + return self.comm.allgather(obj) + + else: + raise ValueError("Either 'group' or 'comm' must be configured.") + + local_ipc_handle = self.runtime.get_local_ipc_handle() + ipc_handles = all_gather_object(local_ipc_handle) + + if self.tcp_store_group is not None: + with self._fetch_remote_metadata_from_tcp_store(remote_ranks) as remote_mds: + self.runtime.connect_ranks(remote_ranks, remote_mds, ipc_handles) + else: + self.runtime.connect_ranks(remote_ranks, None, ipc_handles) + def connect_ranks(self, remote_ranks: List[int]) -> None: """ Add connections to remote ranks. @@ -494,11 +852,16 @@ def connect_ranks(self, remote_ranks: List[int]) -> None: remote_ranks: List of remote rank IDs to establish connections with. The current rank will be automatically filtered out. """ - if self.tcp_store_group is not None: - with self._fetch_remote_metadata_from_tcp_store(remote_ranks) as remote_mds: - self.runtime.connect_ranks(remote_ranks, remote_mds) + if self.low_latency_mode: + if self.tcp_store_group is not None: + with self._fetch_remote_metadata_from_tcp_store( + remote_ranks + ) as remote_mds: + self.runtime.connect_ranks(remote_ranks, remote_mds) + else: + self.runtime.connect_ranks(remote_ranks) else: - self.runtime.connect_ranks(remote_ranks) + self._ht_connect_ranks(remote_ranks) def disconnect_ranks(self, remote_ranks: List[int]) -> None: """ diff --git a/examples/device/ep/nixl_ep/utils.py b/examples/device/ep/nixl_ep/utils.py index 69d68208..4aad3449 100644 --- a/examples/device/ep/nixl_ep/utils.py +++ b/examples/device/ep/nixl_ep/utils.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 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. diff --git a/examples/device/ep/tests/test_ht.py b/examples/device/ep/tests/test_ht.py new file mode 100644 index 00000000..45c752d7 --- /dev/null +++ b/examples/device/ep/tests/test_ht.py @@ -0,0 +1,578 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 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 sys +import time + +# Add elastic subdirectory to path for store_group import +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "elastic")) +# noinspection PyUnresolvedReferences +import nixl_ep # noqa: E402 +import store_group # noqa: E402 +import torch # noqa: E402 +import torch.distributed as dist # noqa: E402 + +from utils import ( # noqa: E402 + bench, + bench_kineto, + calc_diff, + create_grouped_scores, + init_dist, + inplace_unique, + per_token_cast_back, + per_token_cast_to_fp8, +) + +TCP_STORE_PORT = 9999 + + +# noinspection PyShadowingNames +def test_main( + args: argparse.Namespace, + num_sms: int, + local_rank: int, + num_local_ranks: int, + num_ranks: int, + num_nodes: int, + rank: int, + buffer: nixl_ep.Buffer, + group: dist.ProcessGroup, +): + # Settings + num_tokens, hidden = args.num_tokens, args.hidden + num_topk_groups, num_topk, num_experts = ( + args.num_topk_groups, + args.num_topk, + args.num_experts, + ) + + assert num_experts % num_ranks == 0 and num_local_ranks == 8 + if local_rank == 0: + print( + f"[config] num_tokens={num_tokens}, hidden={hidden}, num_topk_groups={num_topk_groups}, num_topk={num_topk}", + flush=True, + ) + + # Random data + x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * rank + x_pure_rand = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") + x_e4m3 = per_token_cast_to_fp8(x) + x_e4m3 = (x_e4m3[0], x_e4m3[1].T.contiguous().T) + scores = ( + torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs() + + 1 + ) + group_scores = scores.view(num_tokens, num_nodes, -1).amax(dim=-1) + group_idx = torch.topk( + group_scores, k=num_topk_groups, dim=-1, sorted=False + ).indices + masked_scores = create_grouped_scores(scores, group_idx, num_nodes) + topk_idx = torch.topk(masked_scores, num_topk, dim=-1, largest=True, sorted=False)[ + 1 + ] + topk_idx = topk_idx.to(nixl_ep.topk_idx_t) + topk_weights = ( + torch.ones((num_tokens, num_topk), dtype=torch.float32, device="cuda") * rank + ) + topk_weights_pure_rand = torch.randn( + (num_tokens, num_topk), dtype=torch.float32, device="cuda" + ) + rank_idx = topk_idx // (num_experts // num_ranks) + rank_idx = rank_idx.to(torch.int64) + rank_idx.masked_fill_(topk_idx == -1, -1) + inplace_unique(rank_idx, num_ranks) + rdma_rank_idx = rank_idx // num_local_ranks + rdma_rank_idx.masked_fill_(rank_idx == -1, -1) + inplace_unique(rdma_rank_idx, num_nodes) + + # RDMA dispatch counts + rdma_idx = topk_idx // (num_experts // num_nodes) + rdma_idx.masked_fill_(topk_idx == -1, -1) + inplace_unique(rdma_idx, num_nodes) + num_rdma_token_sent = rdma_idx.ne(-1).sum().item() + + # Expert meta + num_tokens_per_expert = torch.zeros((num_experts,), dtype=torch.int, device="cuda") + for i in range(num_experts): + num_tokens_per_expert[i] = (topk_idx == i).sum() + gbl_num_tokens_per_expert = num_tokens_per_expert.clone() + dist.all_reduce(gbl_num_tokens_per_expert, group=group) + + # Rank layout meta + num_tokens_per_rank = torch.empty((num_ranks,), dtype=torch.int, device="cuda") + num_tokens_per_rdma_rank = torch.empty((num_nodes,), dtype=torch.int, device="cuda") + token_idx_in_rank = torch.full( + (num_ranks, num_tokens), -1, dtype=torch.long, device="cuda" + ) + for i in range(num_ranks): + num_tokens_per_rank[i] = (rank_idx == i).sum() + token_sel = (rank_idx == i).max(dim=-1)[0] + count = token_sel.sum().item() + tokens = torch.sort(token_sel.to(torch.int), descending=True)[1] + tokens[:count] = torch.sort(tokens[:count])[0] + token_idx_in_rank[i][tokens[:count]] = torch.arange( + count, dtype=torch.long, device="cuda" + ) + for i in range(num_nodes): + num_tokens_per_rdma_rank[i] = (rdma_rank_idx == i).sum() + token_idx_in_rank = token_idx_in_rank.T.contiguous().to(torch.int) + is_token_in_rank = token_idx_in_rank >= 0 + gbl_num_tokens_per_rank = num_tokens_per_rank.clone() + dist.all_reduce(gbl_num_tokens_per_rank, group=group) + + ( + ref_num_tokens_per_rank, + ref_num_tokens_per_rdma_rank, + ref_num_tokens_per_expert, + ref_is_token_in_rank, + _, + ) = buffer.get_dispatch_layout(topk_idx, num_experts) + assert torch.allclose(ref_num_tokens_per_rank, num_tokens_per_rank) + assert torch.allclose(ref_num_tokens_per_rdma_rank, num_tokens_per_rdma_rank) + assert torch.allclose(ref_num_tokens_per_expert, num_tokens_per_expert) + assert torch.allclose(ref_is_token_in_rank, is_token_in_rank) + t = bench(lambda: buffer.get_dispatch_layout(topk_idx, num_experts))[0] + if local_rank == 0: + print(f"[layout] Kernel performance: {t * 1000:.3f} ms", flush=True) + print("", flush=True) + group.barrier() + time.sleep(1) + + # Config + rdma_buffer_size, nvl_buffer_size = 128, (720 if num_ranks in (144, 160) else 512) + config = nixl_ep.Config(num_sms, 8, nvl_buffer_size, 16, rdma_buffer_size) + + # Test dispatch + # noinspection PyShadowingNames + def check_data(check_x, recv_gbl_rank_prefix_sum): + assert torch.allclose(check_x.amin(dim=1), check_x.amax(dim=1)) + check_start = 0 + for i in range(num_ranks): + check_end = recv_gbl_rank_prefix_sum[i].item() + assert (check_x[check_start:check_end, :].int() - i).sum().item() == 0 + check_start = check_end + + for previous_mode in (False, True): + for async_mode in (False, True): + for current_x in (x_pure_rand, x, x_e4m3): + for with_topk in (False, True): + if local_rank == 0: + print( + f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...', + flush=True, + end="", + ) + dispatch_args = { + "x": current_x, + "num_tokens_per_rank": num_tokens_per_rank, + "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, + "is_token_in_rank": is_token_in_rank, + "num_tokens_per_expert": num_tokens_per_expert, + "config": config, + "async_finish": async_mode, + } + if with_topk: + dispatch_args.update( + { + "topk_idx": topk_idx, + "topk_weights": ( + topk_weights_pure_rand + if current_x is x_pure_rand + else topk_weights + ), + } + ) + if previous_mode: + dispatch_args.update({"previous_event": buffer.capture()}) + ( + recv_x, + recv_topk_idx, + recv_topk_weights, + recv_num_tokens_per_expert_list, + handle, + event, + ) = buffer.ht_dispatch(**dispatch_args) + event.current_stream_wait() if async_mode else () + recv_x = ( + per_token_cast_back(*recv_x) + if isinstance(recv_x, tuple) + else recv_x + ) + + # Checks + recv_gbl_rank_prefix_sum = handle[-4] + assert gbl_num_tokens_per_rank[rank].item() == recv_x.size( + 0 + ), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}" + assert ( + gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist() + == recv_num_tokens_per_expert_list + ) + if current_x is not x_pure_rand: + check_data(recv_x, recv_gbl_rank_prefix_sum) + if with_topk: + # Check `topk_idx` + assert recv_topk_idx is not None + assert recv_topk_weights is not None + assert ( + recv_topk_idx.eq(-1) + | ( + (recv_topk_idx >= 0) + & (recv_topk_idx < (num_experts // num_ranks)) + ) + ).sum().item() == recv_topk_idx.numel() + for i, count in enumerate(recv_num_tokens_per_expert_list): + assert recv_topk_idx.eq(i).sum().item() == count + + # Check `topk_weights` + if current_x is not x_pure_rand: + recv_topk_weights[recv_topk_idx.eq(-1)] = ( + recv_topk_weights.amax(dim=1, keepdim=True).expand_as( + recv_topk_weights + )[recv_topk_idx.eq(-1)] + ) + check_data(recv_topk_weights, recv_gbl_rank_prefix_sum) + + # Test cached dispatch (must without top-k staffs) + if not with_topk: + dispatch_args = { + "x": current_x, + "handle": handle, + "config": config, + "async_finish": async_mode, + } + if previous_mode: + dispatch_args.update({"previous_event": buffer.capture()}) + recv_x_cached, _, _, _, _, event = buffer.ht_dispatch( + **dispatch_args + ) + event.current_stream_wait() if async_mode else () + recv_x_cached = ( + per_token_cast_back(*recv_x_cached) + if isinstance(recv_x_cached, tuple) + else recv_x_cached + ) + + if current_x is not x_pure_rand: + check_data(recv_x_cached, recv_gbl_rank_prefix_sum) + + # Use cached result for combine + recv_x = recv_x_cached + + # Test combine + bias_0 = torch.ones( + (num_tokens, hidden), dtype=torch.bfloat16, device="cuda" + ) + bias_1 = torch.randn( + (num_tokens, hidden), dtype=torch.bfloat16, device="cuda" + ) + combine_args = { + "x": recv_x, + "bias": (bias_0, bias_1), + "handle": handle, + "config": config, + "async_finish": async_mode, + } + if with_topk: + combine_args.update({"topk_weights": recv_topk_weights}) + if previous_mode: + combine_args.update({"previous_event": buffer.capture()}) + combined_x, combined_topk_weights, event = buffer.ht_combine( + **combine_args + ) + event.current_stream_wait() if async_mode else () + + check_x = ( + combined_x.float() - bias_0.float() - bias_1.float() + ) / is_token_in_rank.sum(dim=1).unsqueeze(1) + ref_x = x_pure_rand if current_x is x_pure_rand else x + assert calc_diff(check_x, ref_x) < 5e-6 + if with_topk: + check_topk_weights = ( + combined_topk_weights + if (current_x is x_pure_rand) + else ( + combined_topk_weights + / is_token_in_rank.sum(dim=1).unsqueeze(1) + ) + ) + ref_topk_weights = ( + topk_weights_pure_rand + if current_x is x_pure_rand + else topk_weights + ) + assert calc_diff(check_topk_weights, ref_topk_weights) < 1e-9 + + # For later tuning + dispatch_bf16_rdma_send_bytes = num_rdma_token_sent * hidden * 2 + dispatch_bf16_nvl_recv_bytes = recv_x.numel() * 2 + combine_bf16_nvl_send_bytes = dispatch_bf16_nvl_recv_bytes + combine_bf16_rdma_recv_bytes = dispatch_bf16_rdma_send_bytes + + # Sync all ranks before printing passed + group.barrier() + if local_rank == 0: + print(" passed", flush=True) + group.barrier() + if local_rank == 0: + print("", flush=True) + + # Tune dispatch performance + best_dispatch_results = None + fp8_factor = (1 + 4 / 128) / 2 + for current_x in (x_e4m3, x): + best_time, best_results = 1e10, None + rdma_send_bytes = ( + (dispatch_bf16_rdma_send_bytes * fp8_factor) + if isinstance(current_x, tuple) + else dispatch_bf16_rdma_send_bytes + ) + nvl_recv_bytes = ( + (dispatch_bf16_nvl_recv_bytes * fp8_factor) + if isinstance(current_x, tuple) + else dispatch_bf16_nvl_recv_bytes + ) + for nvl_chunk_size in range(4, 45, 4): + for rdma_chunk_size in range(4, 33, 4): + config = nixl_ep.Config( + num_sms, + nvl_chunk_size, + nvl_buffer_size, + rdma_chunk_size, + rdma_buffer_size, + ) + tune_args = {"x": current_x, "handle": handle, "config": config} + t, notify_t = bench_kineto( + lambda: buffer.ht_dispatch(**tune_args), ("dispatch", "notify") + ) + if t < best_time: + best_time, best_results = t, ( + num_sms, + nvl_chunk_size, + rdma_chunk_size, + notify_t, + ) + if local_rank == 0: + print( + f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}, RDMA chunk {rdma_chunk_size}, transmit: {t * 1e6:.2f} us, notify: {notify_t * 1e6:.2f} us, BW: {rdma_send_bytes / 1e9 / t:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / t:.2f} GB/s (NVL) ", + flush=True, + ) + if local_rank == 0: + print( + f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}, transmit: {best_time * 1e6:.2f} us, notify: {best_results[3] * 1e6:.2f} us, BW: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)', # type: ignore[index] + flush=True, + ) + print("", flush=True) + + if isinstance(current_x, tuple): + # Gather FP8 the best config from rank 0 + best_dispatch_results = torch.tensor([best_results[0], best_results[1], best_results[2]], dtype=torch.int32, device="cuda") # type: ignore[index] + all_best_fp8_results_list = [ + torch.zeros_like(best_dispatch_results) + for _ in range(torch.distributed.get_world_size()) + ] + dist.all_gather( + all_best_fp8_results_list, best_dispatch_results, group=group + ) + best_dispatch_results = all_best_fp8_results_list[0].tolist() + dispatch_config = nixl_ep.Config(best_dispatch_results[0], best_dispatch_results[1], nvl_buffer_size, best_dispatch_results[2], rdma_buffer_size) # type: ignore[index] + + dispatch_args = { + "x": x, + "num_tokens_per_rank": num_tokens_per_rank, + "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, + "is_token_in_rank": is_token_in_rank, + "num_tokens_per_expert": num_tokens_per_expert, + "config": dispatch_config if dispatch_config is not None else config, + } + recv_x, _, _, _, handle, _ = buffer.ht_dispatch(**dispatch_args) + + # Tune combine performance + best_time, best_results = 1e10, None + for nvl_chunk_size in range(1, 8, 1): + for rdma_chunk_size in range(12 if num_nodes == 2 else 8, 33, 4): + config = nixl_ep.Config( + num_sms, + nvl_chunk_size, + nvl_buffer_size, + rdma_chunk_size, + rdma_buffer_size, + ) + tune_args = {"x": recv_x, "handle": handle, "config": config} + t, notify_t = bench_kineto( + lambda: buffer.ht_combine(**tune_args), ("combine", "notify") + ) + if local_rank == 0: + print( + f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}, RDMA chunk {rdma_chunk_size}, transmit: {t * 1e6:.2f} us, notify: {notify_t * 1e6:.2f} us, BW: {combine_bf16_rdma_recv_bytes / 1e9 / t:.2f} GB/s (RDMA), {combine_bf16_nvl_send_bytes / 1e9 / t:.2f} GB/s (NVL) ", + flush=True, + ) + if t < best_time: + best_time, best_results = t, ( + num_sms, + nvl_chunk_size, + rdma_chunk_size, + notify_t, + ) + + if local_rank == 0: + print(f"[tuning] Best combine: SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}, transmit: {best_time * 1e6:.2f} us, notify: {best_results[3] * 1e6:.2f} us, BW: {combine_bf16_rdma_recv_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {combine_bf16_nvl_send_bytes / 1e9 / best_time:.2f} GB/s (NVL)", flush=True) # type: ignore[index] + print("", flush=True) + + +# noinspection PyUnboundLocalVariable,PyShadowingNames +def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + # Pin each process to a distinct GPU so NCCL does not see duplicate devices. + # Use local_rank so NCCL gets correct device_id; avoid CUDA_VISIBLE_DEVICES + # so that UCX/DOCA can see all GPUs for GPU-initiated RDMA when needed. + torch.set_default_dtype(torch.bfloat16) + torch.set_default_device("cuda") + torch.cuda.set_device(local_rank % 8) + + num_nodes = int(os.getenv("WORLD_SIZE", 1)) + + rank, num_ranks, group = init_dist(local_rank, num_local_ranks) + print( + f"pid: {os.getpid()}, rank: {rank}, num_ranks: {num_ranks} ,local_rank: {local_rank}", + flush=True, + ) + if args.test_ll_compatibility: + ll_num_experts = 256 + + num_sms = 24 + num_qps_per_rank = max( + num_sms // 2, ll_num_experts // num_ranks if args.test_ll_compatibility else 0 + ) + + # Create TCPStore client for NIXL metadata exchange + tcp_server = args.tcp_server if args.tcp_server else "127.0.0.1" + tcp_store = store_group.create_client_store( + master_addr=tcp_server, + port=TCP_STORE_PORT, + ) + + # Initialize NIXL buffer with group (for IPC handles) and TCPStore (for NIXL metadata) + print( + f"pid: {os.getpid()}, rank: {rank}, num_ranks: {num_ranks}, initializing buffer", + flush=True, + ) + buffer = nixl_ep.Buffer( + rank=rank, + low_latency_mode=False, + explicitly_destroy=True, + group=group, + tcp_store_group=tcp_store, + ) + buffer.update_memory_buffers( + num_ranks=num_ranks, + num_experts_per_rank=num_qps_per_rank, + num_nvl_bytes=int(2e9), + num_rdma_bytes=int(1e9), + ) + buffer.connect_ranks([i for i in range(num_ranks) if i != rank]) + + assert num_local_ranks == 8 and num_ranks > 8 + torch.manual_seed(rank) + + for i in (num_sms,): + test_main( + args, + i, + local_rank, + num_local_ranks, + num_ranks, + num_nodes, + rank, + buffer, + group, + ) + if local_rank == 0: + print("", flush=True) + + # Destroy the buffer runtime and communication group + buffer.destroy() + dist.barrier() + dist.destroy_process_group() + + +def run_server(): + _store = store_group.create_master_store(port=TCP_STORE_PORT) # noqa: F841 + # Keep the server process alive while TCPStore serves requests + while True: + time.sleep(1) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Test high-throughput EP kernels") + parser.add_argument( + "--num-processes", + type=int, + default=8, + help="Number of processes to spawn (default: 8)", + ) + parser.add_argument( + "--num-tokens", type=int, default=4096, help="Number of tokens (default: 4096)" + ) + parser.add_argument( + "--hidden", type=int, default=7168, help="Hidden dimension size (default: 7168)" + ) + parser.add_argument( + "--num-topk-groups", + type=int, + default=None, + help="Number of top-k groups (default: `min(num_nodes, 4)`)", + ) + parser.add_argument( + "--num-topk", type=int, default=8, help="Number of top-k experts (default: 8)" + ) + parser.add_argument( + "--num-experts", type=int, default=256, help="Number of experts (default: 256)" + ) + parser.add_argument( + "--test-ll-compatibility", + action="store_true", + help="whether to test compatibility with low-latency kernels", + ) + parser.add_argument( + "--tcp-server", + type=str, + help="TCP server address (for both TCPStore and rank server). If not set, both will be started locally.", + ) + args = parser.parse_args() + + if not args.tcp_server: + print("Starting TCPStore and rank server locally", flush=True) + server_process = torch.multiprocessing.Process(target=run_server, daemon=True) + server_process.start() + time.sleep(0.5) + + # Set default `num_topk_groups` if not provided + if args.num_topk_groups is None: + num_nodes = int(os.getenv("WORLD_SIZE", 1)) + args.num_topk_groups = min(num_nodes, 4) + + num_processes = args.num_processes + # 2-node run (WORLD_SIZE=2): run on both nodes with same MASTER_ADDR/MASTER_PORT; node1 needs --tcp-server . + # NVL/RDMA timeouts across nodes usually mean RDMA/IB/UCX between nodes is broken or slow (e.g. "accelerated IB support was not found" on one node). + torch.multiprocessing.spawn( + test_loop, args=(num_processes, args), nprocs=num_processes + ) diff --git a/examples/device/ep/tests/utils.py b/examples/device/ep/tests/utils.py index 70fc210f..ec1edb11 100644 --- a/examples/device/ep/tests/utils.py +++ b/examples/device/ep/tests/utils.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 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. @@ -47,15 +47,11 @@ def init_dist(local_rank: int, num_local_ranks: int): } 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 + params["device_id"] = torch.device(f"cuda:{local_rank}") 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 + torch.cuda.set_device(local_rank) return ( dist.get_rank(), From 0bae9b542dba9ffefffb64477c48d56e59f8ff12 Mon Sep 17 00:00:00 2001 From: Nate Mailhot Date: Thu, 9 Apr 2026 10:35:44 -0700 Subject: [PATCH 11/59] chore: update attributions for the latest release (#1480) * chore: update attributions for the latest release * fix missing license * update attr * fix dup * more changes * add missing license * add new pckage --- ATTRIBUTIONS-CPP.md | 47 ++++++++++++++-- ATTRIBUTIONS-Python.md | 118 ++++++++++++++++++++++++++--------------- ATTRIBUTIONS-Rust.md | 50 +++++++---------- 3 files changed, 138 insertions(+), 77 deletions(-) diff --git a/ATTRIBUTIONS-CPP.md b/ATTRIBUTIONS-CPP.md index edf0494d..8522ad8b 100644 --- a/ATTRIBUTIONS-CPP.md +++ b/ATTRIBUTIONS-CPP.md @@ -479,7 +479,7 @@ AND/OR REDISTRIBUTE THE LIBRARY AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR ``` -## prometheus-cpp +## prometheus-cpp - v1.3.0 - **Repository URL**: https://github.com/jupp0r/prometheus-cpp - **License URL**: https://github.com/jupp0r/prometheus-cpp/blob/master/LICENSE @@ -514,7 +514,7 @@ FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ``` -## uccl +## uccl - 1.4.0 - **Repository URL**: https://github.com/uccl-project/uccl - **License URL**: https://github.com/uccl-project/uccl/blob/main/LICENSE @@ -700,7 +700,44 @@ DEALINGS IN THE SOFTWARE. END OF TERMS AND CONDITIONS ``` -## libxml2 +## LibFabric - 2.3.0 + +- **Repository URL**: https://github.com/ofiwg/libfabric +- **License URL**: https://github.com/ofiwg/libfabric/blob/main/COPYING +- **License name**: BSD-3-Clause +### License Text: +``` +Copyright (c) 2015-2021 Intel Corporation. All rights reserved. +Copyright (c) 2015-2021 Cisco Systems, Inc. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in + the documentation and/or other materials provided with the + distribution. + * Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +``` + +## libxml2 - 2.9.14 - **Repository URL**: https://gitlab.gnome.org/GNOME/libxml2 - **License URL**: https://gitlab.gnome.org/GNOME/libxml2/-/blob/master/Copyright @@ -5274,7 +5311,7 @@ Random Hacker. That's all there is to it! ``` -## Mooncake - v0.3.9 +## Mooncake - v0.3.2.post1 - **Repository URL**: https://github.com/kvcache-ai/Mooncake - **License URL**: https://github.com/kvcache-ai/Mooncake/blob/main/LICENSE @@ -6141,7 +6178,7 @@ OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ``` -## RE2 +## RE2 - 2023-03-01 - **Repository URL**: https://github.com/google/re2 - **License URL**: https://github.com/google/re2/blob/main/LICENSE diff --git a/ATTRIBUTIONS-Python.md b/ATTRIBUTIONS-Python.md index d01bfc43..48b6f1d7 100644 --- a/ATTRIBUTIONS-Python.md +++ b/ATTRIBUTIONS-Python.md @@ -21,7 +21,7 @@ This project uses the following third-party libraries. Each library is open-sour This file is automatically generated. Please do not edit it directly. -## black (26.1.0) +## black (26.3.1) ### Licenses License: `MIT` @@ -90,7 +90,7 @@ THE SOFTWARE. - `Homepage`: https://github.com/asottile/cfgv -## click (8.1.7) +## click (8.3.1) ### Licenses License: `BSD-3-Clause` @@ -137,7 +137,7 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - `Source Code`: https://github.com/pallets/click/ -## cuda-bindings (12.9.4) +## cuda-bindings (13.2.0) ### Licenses License: `NVIDIA Proprietary Software` @@ -199,7 +199,7 @@ g. You agree to defend, indemnify and hold harmless NVIDIA and its affiliates, - `Repository`: https://github.com/NVIDIA/cuda-python -## cuda-pathfinder (1.3.3) +## cuda-pathfinder (1.5.0) ### Licenses License: `Apache-2.0` @@ -879,10 +879,10 @@ License: `Apache Software License 2.0` - `Homepage`: https://github.com/kragniz/python-etcd3 -## filelock (3.20.3) +## filelock (3.25.2) ### Licenses -License: `Unlicense` +License: `MIT` - `licenses/LICENSE`: ``` @@ -954,7 +954,7 @@ SOFTWARE. - `Homepage`: https://github.com/pycqa/flake8 -## fsspec (2026.2.0) +## fsspec (2026.3.0) ### Licenses License: `BSD-3-Clause` @@ -998,7 +998,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - `Homepage`: https://github.com/fsspec/filesystem_spec -## grpcio (1.78.0) +## grpcio (1.80.0) ### Licenses License: `Apache-2.0` @@ -1624,7 +1624,7 @@ Exhibit B - "Incompatible With Secondary Licenses" Notice - `Source Code`: https://github.com/grpc/grpc -## identify (2.6.16) +## identify (2.6.18) ### Licenses License: `MIT` @@ -8418,7 +8418,7 @@ SOFTWARE. - `Homepage`: https://github.com/pytest-dev/iniconfig -## isort (7.0.0) +## isort (8.0.1) ### Licenses License: `MIT` @@ -8500,7 +8500,7 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - `Source`: https://github.com/pallets/jinja/ -## librt (0.7.8) +## librt (0.8.1) ### Licenses License: `MIT` @@ -9160,7 +9160,7 @@ DEALINGS IN THE SOFTWARE. ## mypy-extensions (1.1.0) ### Licenses -License: `None` +License: `Apache-2.0` - `licenses/LICENSE`: ``` @@ -9250,10 +9250,10 @@ NetworkX is distributed with the 3-clause BSD license. - `Source Code`: https://github.com/networkx/networkx -## nixl (0.10.0) +## nixl (1.0.1) ### Licenses -License: `None` +License: `Apache-2.0` - `licenses/LICENSE`: ``` @@ -9471,10 +9471,10 @@ License: `None` -## nixl-cu12 (0.10.0) +## nixl-cu13 (1.0.1) ### Licenses -License: `None` +License: `Apache-2.0` - `licenses/LICENSE`: ``` @@ -9737,7 +9737,7 @@ DAMAGE. - `Homepage`: https://github.com/ekalinin/nodeenv -## numpy (2.4.2) +## numpy (2.4.4) ### Licenses License: `BSD-3-Clause AND 0BSD AND MIT AND Zlib AND CC0-1.0` @@ -11588,7 +11588,7 @@ PERFORMANCE OF THIS SOFTWARE. - `tracker`: https://github.com/numpy/numpy/issues -## nvidia-cublas-cu12 (12.8.4.1) +## nvidia-cublas (13.1.0.3) ### Licenses License: `NVIDIA Proprietary Software` @@ -13169,7 +13169,7 @@ conditions: - `Homepage`: https://developer.nvidia.com/cuda-zone -## nvidia-cuda-cupti-cu12 (12.8.90) +## nvidia-cuda-cupti (13.0.85) ### Licenses License: `NVIDIA Proprietary Software` @@ -14750,7 +14750,7 @@ conditions: - `Homepage`: https://developer.nvidia.com/cuda-zone -## nvidia-cuda-nvrtc-cu12 (12.8.93) +## nvidia-cuda-nvrtc (13.0.88) ### Licenses License: `NVIDIA Proprietary Software` @@ -16331,7 +16331,7 @@ conditions: - `Homepage`: https://developer.nvidia.com/cuda-zone -## nvidia-cuda-runtime-cu12 (12.8.90) +## nvidia-cuda-runtime (13.0.96) ### Licenses License: `NVIDIA Proprietary Software` @@ -17912,7 +17912,7 @@ conditions: - `Homepage`: https://developer.nvidia.com/cuda-zone -## nvidia-cudnn-cu12 (9.10.2.21) +## nvidia-cudnn-cu13 (9.19.0.56) ### Licenses License: `LicenseRef-NVIDIA-Proprietary` @@ -18079,7 +18079,7 @@ In addition to the rights above, for parties that are developing software intend - `Homepage`: https://developer.nvidia.com/cuda-zone -## nvidia-cufft-cu12 (11.3.3.83) +## nvidia-cufft (12.0.0.61) ### Licenses License: `NVIDIA Proprietary Software` @@ -19660,7 +19660,7 @@ conditions: - `Homepage`: https://developer.nvidia.com/cuda-zone -## nvidia-cufile-cu12 (1.13.1.3) +## nvidia-cufile (1.15.1.6) ### Licenses License: `NVIDIA Proprietary Software` @@ -21241,7 +21241,7 @@ conditions: - `Homepage`: https://developer.nvidia.com/cuda-zone -## nvidia-curand-cu12 (10.3.9.90) +## nvidia-curand (10.4.0.35) ### Licenses License: `NVIDIA Proprietary Software` @@ -22822,7 +22822,7 @@ conditions: - `Homepage`: https://developer.nvidia.com/cuda-zone -## nvidia-cusolver-cu12 (11.7.3.90) +## nvidia-cusolver (12.0.4.66) ### Licenses License: `NVIDIA Proprietary Software` @@ -24403,7 +24403,7 @@ conditions: - `Homepage`: https://developer.nvidia.com/cuda-zone -## nvidia-cusparse-cu12 (12.5.8.93) +## nvidia-cusparse (12.6.3.3) ### Licenses License: `NVIDIA Proprietary Software` @@ -25984,7 +25984,7 @@ conditions: - `Homepage`: https://developer.nvidia.com/cuda-zone -## nvidia-cusparselt-cu12 (0.7.1) +## nvidia-cusparselt-cu13 (0.8.0) ### Licenses License: `NVIDIA Proprietary Software` @@ -25993,7 +25993,7 @@ License: `NVIDIA Proprietary Software` - `Homepage`: https://developer.nvidia.com/cusparselt -## nvidia-nccl-cu12 (2.27.5) +## nvidia-nccl-cu13 (2.28.9) ### Licenses License: `BSD-3-Clause` @@ -26045,7 +26045,7 @@ for more information and license details. - `Homepage`: https://developer.nvidia.com/cuda-zone -## nvidia-nvjitlink-cu12 (12.8.93) +## nvidia-nvjitlink (13.0.88) ### Licenses License: `NVIDIA Proprietary Software` @@ -27626,7 +27626,7 @@ conditions: - `Homepage`: https://developer.nvidia.com/cuda-zone -## nvidia-nvshmem-cu12 (3.4.5) +## nvidia-nvshmem-cu13 (3.4.5) ### Licenses License: `NVIDIA Proprietary Software` @@ -29207,7 +29207,7 @@ conditions: - `Homepage`: https://developer.nvidia.com/cuda-zone -## nvidia-nvtx-cu12 (12.8.90) +## nvidia-nvtx (13.0.85) ### Licenses License: `Apache 2.0` @@ -30051,7 +30051,7 @@ Exhibit B - "Incompatible With Secondary Licenses" Notice - `Source Code`: https://github.com/cpburnz/python-pathspec -## platformdirs (4.5.1) +## platformdirs (4.9.4) ### Licenses License: `MIT` @@ -30153,7 +30153,7 @@ THE SOFTWARE. - `Homepage`: https://github.com/pre-commit/pre-commit -## protobuf (6.33.5) +## protobuf (7.34.1) ### Licenses License: `3-Clause BSD License` @@ -30271,7 +30271,7 @@ WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. - `Homepage`: https://github.com/PyCQA/pyflakes -## pygments (2.19.2) +## pygments (2.20.0) ### Licenses License: `BSD-2-Clause` @@ -30313,7 +30313,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - `Source`: https://github.com/pygments/pygments -## pytest (7.4.4) +## pytest (9.0.2) ### Licenses License: `MIT` @@ -30424,7 +30424,7 @@ SOFTWARE. - `Source Code`: https://github.com/yaml/pyyaml -## setuptools (82.0.0) +## setuptools (81.0.0) ### Licenses License: `MIT` @@ -31301,7 +31301,7 @@ SOFTWARE. - `Source`: https://github.com/sympy/sympy -## tabulate (0.9.0) +## tabulate (0.10.0) ### Licenses License: `MIT` @@ -31587,7 +31587,7 @@ THE SOFTWARE.``` - `Homepage`: https://github.com/uiri/toml -## tomli (2.4.0) +## tomli (2.4.1) ### Licenses License: `MIT` @@ -31622,7 +31622,7 @@ SOFTWARE. - `Homepage`: https://github.com/hukkin/tomli -## torch (2.10.0) +## torch (2.11.0) ### Licenses License: `BSD-3-Clause` @@ -40600,7 +40600,7 @@ If the Library as you received it specifies that a proxy can decide whether futu - `Repository`: https://github.com/pytorch/pytorch -## tqdm (4.66.5) +## tqdm (4.67.3) ### Licenses License: `MPL-2.0 AND MIT` @@ -41252,7 +41252,7 @@ PERFORMANCE OF THIS SOFTWARE. - `Repository`: https://github.com/python/typing_extensions -## virtualenv (20.36.1) +## virtualenv (21.2.0) ### Licenses License: `MIT` @@ -41287,4 +41287,38 @@ WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. - `Source`: https://github.com/pypa/virtualenv - `Tracker`: https://github.com/pypa/virtualenv/issues +## cuda-toolkit (13.0.2) + +- PyPI: https://pypi.org/project/cuda-toolkit/13.0.2/ +- License: `NVIDIA Proprietary Software` + +NVIDIA Software License Agreement. +See https://docs.nvidia.com/cuda/eula/ for full terms. + +## python-discovery (1.2.1) + +- PyPI: https://pypi.org/project/python-discovery/1.2.1/ +- `repository`: https://github.com/tox-dev/python-discovery +- License: `MIT` +``` +MIT License + +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/ATTRIBUTIONS-Rust.md b/ATTRIBUTIONS-Rust.md index a1449577..892cd65d 100644 --- a/ATTRIBUTIONS-Rust.md +++ b/ATTRIBUTIONS-Rust.md @@ -4856,22 +4856,17 @@ END OF TERMS AND CONDITIONS **License Type(s)**: Apache-2.0 OR LGPL-2.1-or-later OR MIT ### License: https://github.com/r-efi/r-efi ``` +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 - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - 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 + 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. +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. ``` ## regex - 1.11.1 @@ -5085,7 +5080,7 @@ limitations under the License. ## regex-automata - 0.4.9 **Repository URL**: https://github.com/rust-lang/regex/tree/master/regex-automata **License Type(s)**: Apache-2.0 OR MIT -### License: https://github.com/rust-lang/regex/tree/master/regex-automata/blob/HEAD/LICENSE-APACHE +### License: https://github.com/rust-lang/regex/blob/HEAD/LICENSE-APACHE ``` Apache License Version 2.0, January 2004 @@ -5293,7 +5288,7 @@ limitations under the License. ## regex-syntax - 0.8.5 **Repository URL**: https://github.com/rust-lang/regex/tree/master/regex-syntax **License Type(s)**: Apache-2.0 OR MIT -### License: https://github.com/rust-lang/regex/tree/master/regex-syntax/blob/HEAD/LICENSE-APACHE +### License: https://github.com/rust-lang/regex/blob/HEAD/LICENSE-APACHE ``` Apache License Version 2.0, January 2004 @@ -11476,21 +11471,16 @@ limitations under the License. **License Type(s)**: Apache-2.0 OR Apache-2.0 WITH LLVM-exception OR MIT ### License: https://github.com/bytecodealliance/wit-bindgen/blob/main/LICENSE-APACHE ``` +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 - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - 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 + 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. +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 ae5ae82dfa45228d9b395bd16427a6801a8a2f7c Mon Sep 17 00:00:00 2001 From: Ye Xiang Date: Thu, 9 Apr 2026 10:48:20 -0700 Subject: [PATCH 12/59] libfabric: fix multi-GPU memory registration (#1506) PR #1258 changed two independent if-statements into if/else in registerMem(), breaking the fallthrough that calls cudaSetDevice() when the CUDA address workaround is disabled for device N>0, in multi-GPU per process case. Without cudaSetDevice(), cuMemGetAddressRange() fails inside EFA's dmabuf path for large buffers, causing fi_mr_key to return FI_KEY_NOTAVAIL. Blackwell (B200) requires the correct CUDA context for cuMemGetAddressRange; Hopper (H200) continues to work based on tests. Restore the original two-if pattern and downgrade the expected multi-GPU detection log from WARN to INFO. Validated with nixlbench --scheme tp --mode MG --num-initiator-dev 4 --num-target-dev 8 --check-consistency on p6 B200 instances. Signed-off-by: Ye Xiang --- src/plugins/libfabric/libfabric_backend.cpp | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/plugins/libfabric/libfabric_backend.cpp b/src/plugins/libfabric/libfabric_backend.cpp index 5fed7b6e..a66dac67 100644 --- a/src/plugins/libfabric/libfabric_backend.cpp +++ b/src/plugins/libfabric/libfabric_backend.cpp @@ -710,23 +710,23 @@ nixlLibfabricEngine::registerMem(const nixlBlobDesc &mem, if (cuda_addr_wa_) { bool need_restart; if (vramUpdateCtx((void *)mem.addr, mem.devId, need_restart)) { - NIXL_WARN << "CUDA address workaround failed for device " << mem.devId - << ", disabling workaround for multi-GPU support"; - cuda_addr_wa_ = false; // Disable workaround for subsequent registrations + NIXL_INFO << "Multi-GPU detected (device " << mem.devId + << "), using cudaSetDevice fallback"; + cuda_addr_wa_ = false; } else if (need_restart) { - // Restart progress thread if needed NIXL_DEBUG << "CUDA context updated, restarting progress thread"; vramApplyCtx(); } - } else { - // Set CUDA device context directly for multi-GPU support + } + // Fallback: set device via runtime API (uses primary context) + if (!cuda_addr_wa_) { cudaError_t cuda_ret = cudaSetDevice(mem.devId); if (cuda_ret != cudaSuccess) { NIXL_ERROR << "Failed to set CUDA device " << mem.devId << ": " << cudaGetErrorString(cuda_ret); return NIXL_ERR_NOT_SUPPORTED; } - NIXL_DEBUG << "Set CUDA device context to GPU " << mem.devId; + NIXL_INFO << "Set CUDA device context to GPU " << mem.devId; } // Query PCI bus ID from memory address (AFTER setting context) From 27589401584e361cadf99319a5f2890c6ebb99e3 Mon Sep 17 00:00:00 2001 From: Raul Akhmetshin <74596089+rakhmets@users.noreply.github.com> Date: Fri, 10 Apr 2026 11:08:38 +0200 Subject: [PATCH 13/59] EXAMPLES/DEVICE/EP/CSRS/KERNELS: Removed unused variable. (#1508) Signed-off-by: Raul Akhmetshin --- examples/device/ep/csrc/kernels/nixl_ep_ht.cu | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/device/ep/csrc/kernels/nixl_ep_ht.cu b/examples/device/ep/csrc/kernels/nixl_ep_ht.cu index 140ba095..4ab48f38 100644 --- a/examples/device/ep/csrc/kernels/nixl_ep_ht.cu +++ b/examples/device/ep/csrc/kernels/nixl_ep_ht.cu @@ -589,7 +589,6 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV __shared__ int rdma_send_channel_lock[kNumRDMARanks]; __shared__ int rdma_send_channel_tail[kNumRDMARanks]; __shared__ uint32_t rdma_send_channel_window[kNumRDMARanks]; - __shared__ nixlGpuXferStatusH rdma_put_xfer_status; auto sync_rdma_sender_smem = []() { asm volatile("barrier.sync 0, %0;" ::"r"((kNumDispatchRDMASenderWarps + 1) * 32)); }; // TMA stuffs From c8d39f7ff4ca7480867ab580e7f0de6d45ad25c2 Mon Sep 17 00:00:00 2001 From: Mikhail Brinskiy Date: Sun, 12 Apr 2026 11:27:53 +0200 Subject: [PATCH 14/59] Update codeowners for NIXL EP and rust (#1489) --- CODEOWNERS | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CODEOWNERS b/CODEOWNERS index 3ba85f24..5b752c09 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -15,7 +15,7 @@ CODEOWNERS @ai-dynamo/Devops @ai-dynamo/nixl-maintainers # Bindings /src/bindings/python @ovidiusm @mkhazraee @roiedanino -/src/bindings/rust @roiedanino @gleon99 @mkhazraee +/src/bindings/rust @roiedanino @gleon99 @tomerg-nvidia @mkhazraee # UCX Plugin /src/plugins/ucx @brminich @yosefe @gleon99 @@ -36,4 +36,4 @@ CODEOWNERS @ai-dynamo/Devops @ai-dynamo/nixl-maintainers /test/unit/utils/libfabric @amitrad-aws @yexiang-aws @akkart-aws @fengjica # NIXL EP example -/examples/device/ep @itayalroy @ebarilanM @ai-dynamo/nixl-maintainers +/examples/device/ep @itayalroy @ebarilanM @ofirfarjun7 @ai-dynamo/nixl-maintainers From b908b41675b90fe4acf72a9a68f73fad8cf85238 Mon Sep 17 00:00:00 2001 From: youkaichao Date: Sun, 12 Apr 2026 21:05:58 +0800 Subject: [PATCH 15/59] BINDINGS/PYTHON/NIXL_META/NIXL: fix IDE/linters/static-analysis name resolution (#1402) * BINDINGS/PYTHON/NIXL_META/NIXL: fix static analysis issues for nixl_agent and nixl_logger --------- Signed-off-by: youkaichao Signed-off-by: Roie Danino Co-authored-by: Roie Danino Co-authored-by: Mikhail Brinskiy --- examples/python/basic_two_peers.py | 4 +-- .../python/nixl-meta/nixl/__init__.py | 12 +++++-- src/bindings/python/nixl-meta/nixl/_api.py | 31 +++++++++++++++++++ src/bindings/python/nixl-meta/nixl/logging.py | 25 +++++++++++++++ .../python/nixl-meta/nixl/meson.build | 4 ++- 5 files changed, 71 insertions(+), 5 deletions(-) create mode 100644 src/bindings/python/nixl-meta/nixl/_api.py create mode 100644 src/bindings/python/nixl-meta/nixl/logging.py diff --git a/examples/python/basic_two_peers.py b/examples/python/basic_two_peers.py index 82b7aa8f..02fa422a 100755 --- a/examples/python/basic_two_peers.py +++ b/examples/python/basic_two_peers.py @@ -1,13 +1,13 @@ #!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import argparse import torch -from nixl._api import nixl_agent, nixl_agent_config +from nixl import nixl_agent, nixl_agent_config from nixl.logging import get_logger logger = get_logger(__name__) diff --git a/src/bindings/python/nixl-meta/nixl/__init__.py b/src/bindings/python/nixl-meta/nixl/__init__.py index f39b5d60..d0b7ef63 100644 --- a/src/bindings/python/nixl-meta/nixl/__init__.py +++ b/src/bindings/python/nixl-meta/nixl/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -15,6 +15,7 @@ import importlib import sys +from typing import TYPE_CHECKING # Try packages in order candidates = ["nixl_cu13", "nixl_cu12"] @@ -34,7 +35,7 @@ submodules = ["_api", "_bindings", "_utils", "logging"] for sub_name in submodules: # Import submodule from actual wheel - module = importlib.import_module(f"{pkg}.{sub_name}") + module = importlib.import_module(f"{_pkg.__name__}.{sub_name}") # Make it accessible as nixl._api, nixl._utils, nixl.logging sys.modules[f"nixl.{sub_name}"] = module # Also add the submodule itself to the nixl namespace @@ -44,3 +45,10 @@ for attr in dir(module): if not attr.startswith("_"): setattr(sys.modules[__name__], attr, getattr(module, attr)) + +if TYPE_CHECKING: + from nixl import logging # noqa: F401 + from nixl._api import ( # type: ignore[attr-defined] # noqa: F401 + nixl_agent, + nixl_agent_config, + ) diff --git a/src/bindings/python/nixl-meta/nixl/_api.py b/src/bindings/python/nixl-meta/nixl/_api.py new file mode 100644 index 00000000..953ce2e7 --- /dev/null +++ b/src/bindings/python/nixl-meta/nixl/_api.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +# This file is a type stub for static analysis tools (pyright, mypy, IDEs). +# At runtime it is shadowed by the actual nixl_cu12._api or nixl_cu13._api +# module, which __init__.py injects into sys.modules["nixl._api"]. +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + try: + from nixl_cu13._api import ( # type: ignore[import] # noqa: F401 + nixl_agent, + nixl_agent_config, + ) + except ImportError: + from nixl_cu12._api import ( # type: ignore[import] # noqa: F401 + nixl_agent, + nixl_agent_config, + ) diff --git a/src/bindings/python/nixl-meta/nixl/logging.py b/src/bindings/python/nixl-meta/nixl/logging.py new file mode 100644 index 00000000..224a2056 --- /dev/null +++ b/src/bindings/python/nixl-meta/nixl/logging.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +# This file is a type stub for static analysis tools (pyright, mypy, IDEs). +# At runtime it is shadowed by the actual nixl_cu12.logging or nixl_cu13.logging +# module, which __init__.py injects into sys.modules["nixl.logging"]. +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + try: + from nixl_cu13.logging import get_logger # type: ignore[import] # noqa: F401 + except ImportError: + from nixl_cu12.logging import get_logger # type: ignore[import] # noqa: F401 diff --git a/src/bindings/python/nixl-meta/nixl/meson.build b/src/bindings/python/nixl-meta/nixl/meson.build index 5e681089..d71f5518 100644 --- a/src/bindings/python/nixl-meta/nixl/meson.build +++ b/src/bindings/python/nixl-meta/nixl/meson.build @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -15,3 +15,5 @@ fs = import('fs') fs.copyfile('__init__.py') +fs.copyfile('_api.py') +fs.copyfile('logging.py') From 3a15bc5c302395d4a41dab3ed1af8a47492f19a4 Mon Sep 17 00:00:00 2001 From: Itay Alroy <75032521+itayalroy@users.noreply.github.com> Date: Mon, 13 Apr 2026 15:07:12 +0300 Subject: [PATCH 16/59] nixl_ep: Fixes to high-throughput commit (#1503) * Remove redundant count buffer * nixl_ep: Fix internode destruction flows * Re-guard p2p_ptr_get with is_rank_masked * Merge two !low_latency_mode blocks * Remove redundant condition --- examples/device/ep/csrc/kernels/nixl_ep_ll.cu | 6 +- examples/device/ep/csrc/nixl_ep.cpp | 89 +++++++++++-------- examples/device/ep/csrc/nixl_ep.hpp | 2 +- 3 files changed, 55 insertions(+), 42 deletions(-) diff --git a/examples/device/ep/csrc/kernels/nixl_ep_ll.cu b/examples/device/ep/csrc/kernels/nixl_ep_ll.cu index b03539f7..41ff55c8 100644 --- a/examples/device/ep/csrc/kernels/nixl_ep_ll.cu +++ b/examples/device/ep/csrc/kernels/nixl_ep_ll.cu @@ -183,8 +183,8 @@ dispatch(void* packed_recv_x, void* packed_recv_x_scales, dst_expert_local_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_msg + rank * num_max_dispatch_tokens_per_rank * num_bytes_per_msg + slot_idx * num_bytes_per_msg; - void* dst_p2p_ptr = p2p_ptr_get(nixl_ctx, dst_ptr, dst_rank); if (not is_rank_masked(mask_buffer_ptr, dst_rank)) { + void* dst_p2p_ptr = p2p_ptr_get(nixl_ctx, dst_ptr, dst_rank); if (dst_p2p_ptr == 0) { nixlMemViewElem src_mdesc{nixl_ctx.local_mvh, 0, nixl_ctx.offset_get(src_ptr)}; nixlMemViewElem dst_mdesc{nixl_ctx.remote_mvh, (size_t) dst_rank, nixl_ctx.offset_get(dst_ptr)}; @@ -252,8 +252,8 @@ dispatch(void* packed_recv_x, void* packed_recv_x_scales, // Wait local sends issued and send expert counts while (ld_acquire_global(atomic_finish_counter_per_expert + responsible_expert_idx) != FINISHED_SUM_TAG * 2); auto dst_ptr = reinterpret_cast(rdma_recv_count + dst_expert_local_idx * num_ranks + rank); - void* dst_p2p_ptr = p2p_ptr_get(nixl_ctx, dst_ptr, dst_rank); if (not is_rank_masked(mask_buffer_ptr, dst_rank)) { + void* dst_p2p_ptr = p2p_ptr_get(nixl_ctx, dst_ptr, dst_rank); if (dst_p2p_ptr == 0) { nixlMemViewElem dst_mdesc{nixl_ctx.remote_mvh, static_cast(dst_rank), nixl_ctx.offset_get(dst_ptr)}; EP_DEVICE_ASSERT(nixlAtomicAdd(num_tokens_sent + 1, dst_mdesc, dst_expert_local_idx) == NIXL_IN_PROG); @@ -801,8 +801,8 @@ combine(void* combined_x, if (sub_warp_id == 1 and lane_id == 0) { while (ld_acquire_global(atomic_clean_flag) == 0); auto dst_ptr = reinterpret_cast(rdma_recv_flag + global_expert_idx); - void* dst_p2p_ptr = p2p_ptr_get(nixl_ctx, dst_ptr, dst_rank); if (not is_rank_masked(mask_buffer_ptr, dst_rank)) { + void* dst_p2p_ptr = p2p_ptr_get(nixl_ctx, dst_ptr, dst_rank); if (dst_p2p_ptr == 0) { nixlMemViewElem dst_mdesc{nixl_ctx.remote_mvh, (size_t) dst_rank, nixl_ctx.offset_get(dst_ptr)}; EP_DEVICE_ASSERT(nixlAtomicAdd(1, dst_mdesc, local_expert_idx) == NIXL_IN_PROG); diff --git a/examples/device/ep/csrc/nixl_ep.cpp b/examples/device/ep/csrc/nixl_ep.cpp index a5a33ca9..2fb3b14d 100644 --- a/examples/device/ep/csrc/nixl_ep.cpp +++ b/examples/device/ep/csrc/nixl_ep.cpp @@ -128,21 +128,24 @@ void Buffer::init(int num_ranks, int num_experts_per_rank, int64_t num_nvl_bytes workspace = m_workspace_alloc->ptr(); CUDA_CHECK(cudaMemsetAsync(workspace, 0, NUM_WORKSPACE_BYTES, comm_stream)); - // MoE counter - CUDA_CHECK(cudaMallocHost(&moe_recv_counter, sizeof(int64_t), cudaHostAllocMapped)); - CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_counter_mapped, const_cast(moe_recv_counter), 0)); - *moe_recv_counter = -1; - - // MoE expert-level counter - CUDA_CHECK(cudaMallocHost(&moe_recv_expert_counter, sizeof(int) * NUM_MAX_LOCAL_EXPERTS, cudaHostAllocMapped)); - CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_expert_counter_mapped, const_cast(moe_recv_expert_counter), 0)); - for (int i = 0; i < NUM_MAX_LOCAL_EXPERTS; ++ i) - moe_recv_expert_counter[i] = -1; - - // MoE RDMA-level counter - CUDA_CHECK(cudaMallocHost(&moe_recv_rdma_counter, sizeof(int), cudaHostAllocMapped)); - CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_rdma_counter_mapped, const_cast(moe_recv_rdma_counter), 0)); - *moe_recv_rdma_counter = -1; + if (!low_latency_mode) { + // MoE counter + CUDA_CHECK(cudaMallocHost(&moe_recv_counter, sizeof(int), cudaHostAllocMapped)); + CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_counter_mapped, const_cast(moe_recv_counter), 0)); + *moe_recv_counter = -1; + + // MoE expert-level counter + CUDA_CHECK(cudaMallocHost(&moe_recv_expert_counter, sizeof(int) * NUM_MAX_LOCAL_EXPERTS, cudaHostAllocMapped)); + CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_expert_counter_mapped, const_cast(moe_recv_expert_counter), 0)); + for (int i = 0; i < NUM_MAX_LOCAL_EXPERTS; ++ i) + moe_recv_expert_counter[i] = -1; + + // MoE RDMA-level counter + CUDA_CHECK(cudaMallocHost(&moe_recv_rdma_counter, sizeof(int), cudaHostAllocMapped)); + CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_rdma_counter_mapped, const_cast(moe_recv_rdma_counter), 0)); + *moe_recv_rdma_counter = -1; + } + EP_HOST_ASSERT(max_experts_per_rank > 0); m_rdma_alloc = std::make_unique(static_cast(num_rdma_bytes)); rdma_buffer_ptr = m_rdma_alloc->ptr(); @@ -162,14 +165,14 @@ void Buffer::init(int num_ranks, int num_experts_per_rank, int64_t num_nvl_bytes sync_count_ptr = static_cast(m_sync_count_alloc->ptr()); CUDA_CHECK(cudaMemset(sync_buffer_ptr, 0, num_sync_buffer_bytes)); CUDA_CHECK(cudaMemset(sync_count_ptr, 0, num_sync_buffer_bytes)); - CUDA_CHECK(cudaMalloc(&local_barrier_cnt_ptr, num_sync_buffer_bytes)); - CUDA_CHECK(cudaMemset(local_barrier_cnt_ptr, 0, num_sync_buffer_bytes)); - - // Allocate barrier counters for high-throughput mode - CUDA_CHECK(cudaMalloc(&local_ht_barrier_counter, sizeof(uint64_t))); - CUDA_CHECK(cudaMemset(local_ht_barrier_counter, 0, sizeof(uint64_t))); - CUDA_CHECK(cudaMalloc(&last_ht_barrier_counter, sizeof(uint64_t))); - CUDA_CHECK(cudaMemset(last_ht_barrier_counter, 0, sizeof(uint64_t))); + + if (!low_latency_mode) { + CUDA_CHECK(cudaMalloc(&local_ht_barrier_counter, sizeof(uint64_t))); + CUDA_CHECK(cudaMemset(local_ht_barrier_counter, 0, sizeof(uint64_t))); + CUDA_CHECK(cudaMalloc(&last_ht_barrier_counter, sizeof(uint64_t))); + CUDA_CHECK(cudaMemset(last_ht_barrier_counter, 0, sizeof(uint64_t))); + } + CUDA_CHECK(cudaDeviceSynchronize()); my_peer_info.rdma_buffer_ptr = rdma_buffer_ptr; @@ -261,16 +264,16 @@ void Buffer::destroy() { if (num_nvl_bytes > 0) { intranode::barrier(barrier_signal_ptrs_gpu, nvl_rank, num_nvl_ranks, comm_stream); - CUDA_CHECK(cudaDeviceSynchronize()); + warn_cuda(cudaDeviceSynchronize(), "synchronize device after intranode barrier"); // Close remote IPC if (is_available()) { for (int i = 0; i < num_nvl_ranks; ++ i) if (i != nvl_rank) - CUDA_CHECK(cudaIpcCloseMemHandle(buffer_ptrs[i])); + warn_cuda(cudaIpcCloseMemHandle(buffer_ptrs[i]), "close remote IPC handle"); } // Free local buffer - CUDA_CHECK(cudaFree(buffer_ptrs[nvl_rank])); + warn_cuda(cudaFree(buffer_ptrs[nvl_rank]), "free local NVL buffer"); } if (nixl_agent_info and nixl_agent_info->agent != nullptr) { @@ -290,6 +293,11 @@ void Buffer::destroy() { nixl_agent_info->sync_count_reg_descs, &nixl_agent_info->extra_params), "deregister sync-count memory"); + if (local_ht_barrier_counter != nullptr) { + warn_nixl(nixl_agent_info->agent->deregisterMem( + nixl_agent_info->ht_barrier_reg_descs), + "deregister ht barrier memory"); + } nixl_agent_info.reset(); } @@ -303,17 +311,22 @@ void Buffer::destroy() { m_sync_count_alloc.reset(); sync_count_ptr = nullptr; - warn_cuda(cudaFree(local_barrier_cnt_ptr), "free local barrier count"); - warn_cuda(cudaFree(local_ht_barrier_counter), "free local ht barrier counter"); - warn_cuda(cudaFree(last_ht_barrier_counter), "free last ht barrier counter"); + if (!low_latency_mode) { + warn_cuda(cudaFree(local_ht_barrier_counter), "free local ht barrier counter"); + local_ht_barrier_counter = nullptr; + warn_cuda(cudaFree(last_ht_barrier_counter), "free last ht barrier counter"); + last_ht_barrier_counter = nullptr; + warn_cuda(cudaFreeHost(const_cast(moe_recv_counter)), "free moe receive counter"); + moe_recv_counter = nullptr; + warn_cuda(cudaFreeHost(const_cast(moe_recv_expert_counter)), "free moe receive expert counter"); + moe_recv_expert_counter = nullptr; + warn_cuda(cudaFreeHost(const_cast(moe_recv_rdma_counter)), "free moe receive rdma counter"); + moe_recv_rdma_counter = nullptr; + } m_workspace_alloc.reset(); workspace = nullptr; - CUDA_CHECK(cudaFreeHost(const_cast(moe_recv_counter))); - CUDA_CHECK(cudaFreeHost(const_cast(moe_recv_expert_counter))); - CUDA_CHECK(cudaFreeHost(const_cast(moe_recv_rdma_counter))); - destroyed = true; available = false; } @@ -438,7 +451,6 @@ void Buffer::connect_ranks(const std::vector& remote_ranks_list, const std: new_ranks.push_back(remote_rank); CUDA_CHECK(cudaMemset(mask_buffer_ptr + remote_rank, 0, sizeof(int))); CUDA_CHECK(cudaMemset(sync_count_ptr + remote_rank, 0, sizeof(int))); - CUDA_CHECK(cudaMemset(local_barrier_cnt_ptr + remote_rank, 0, sizeof(int))); CUDA_CHECK(cudaMemset(sync_buffer_ptr + remote_rank, 0, sizeof(int))); if (remote_mds.has_value()) @@ -1330,15 +1342,16 @@ void Buffer::_nixl_agent_init() { nixl_agent_info->sync_count_reg_descs.clear(); nixl_agent_info->sync_count_reg_descs.addDesc( nixlBlobDesc(reinterpret_cast(sync_count_ptr), max_num_ranks * sizeof(int), device_id, "")); + nixl_agent_info->ht_barrier_reg_descs.clear(); EP_HOST_ASSERT(agent->registerMem(nixl_agent_info->rdma_reg_descs, &nixl_agent_info->extra_params) == NIXL_SUCCESS); EP_HOST_ASSERT(agent->registerMem(nixl_agent_info->sync_reg_descs, &nixl_agent_info->extra_params) == NIXL_SUCCESS); EP_HOST_ASSERT(agent->registerMem(nixl_agent_info->sync_count_reg_descs, &nixl_agent_info->extra_params) == NIXL_SUCCESS); - if (!low_latency_mode && local_ht_barrier_counter) { - nixl_reg_dlist_t ht_barrier_dlist(VRAM_SEG); - ht_barrier_dlist.addDesc(nixlBlobDesc((uintptr_t)(local_ht_barrier_counter), sizeof(uint64_t), get_local_device_id(), "")); - EP_HOST_ASSERT(agent->registerMem(ht_barrier_dlist) == NIXL_SUCCESS); + if (local_ht_barrier_counter) { + nixl_agent_info->ht_barrier_reg_descs.addDesc( + nixlBlobDesc((uintptr_t)(local_ht_barrier_counter), sizeof(uint64_t), get_local_device_id(), "")); + EP_HOST_ASSERT(agent->registerMem(nixl_agent_info->ht_barrier_reg_descs) == NIXL_SUCCESS); } if (getenv("NIXL_ETCD_ENDPOINTS")) { diff --git a/examples/device/ep/csrc/nixl_ep.hpp b/examples/device/ep/csrc/nixl_ep.hpp index 0c58cde8..03a89816 100644 --- a/examples/device/ep/csrc/nixl_ep.hpp +++ b/examples/device/ep/csrc/nixl_ep.hpp @@ -73,6 +73,7 @@ struct NixlAgentInfo nixl_reg_dlist_t rdma_reg_descs{VRAM_SEG}; nixl_reg_dlist_t sync_reg_descs{VRAM_SEG}; nixl_reg_dlist_t sync_count_reg_descs{VRAM_SEG}; + nixl_reg_dlist_t ht_barrier_reg_descs{VRAM_SEG}; std::vector wire_up_done; // [num_peers] }; @@ -95,7 +96,6 @@ struct Buffer { int *mask_buffer_ptr = nullptr; int *sync_buffer_ptr = nullptr; int *sync_count_ptr = nullptr; - int *local_barrier_cnt_ptr = nullptr; /* Owning VMM allocations (keep raw ptrs above as aliases) */ std::unique_ptr m_rdma_alloc; From e6dee7f0f92b68aa58de41dda8aecdc451908a8c Mon Sep 17 00:00:00 2001 From: BatshevaBlack <132911331+BatshevaBlack@users.noreply.github.com> Date: Mon, 13 Apr 2026 15:24:57 +0300 Subject: [PATCH 17/59] Fix unbounded growth of in-memory telemetry event buffer (#1516) * Fix unbounded growth of in-memory telemetry event buffer The events_ vector grew without limit between periodic flushes. Enforce the configured buffer size (NIXL_TELEMETRY_BUFFER_SIZE) by dropping new events when the buffer is full (drop-newest policy). Signed-off-by: Batsheva Black * Reduce naming redundancy for telemetry buffer size Remove the buffer_size local variable in initializeTelemetry() and assign directly to maxBufferedEvents_. Use maxBufferedEvents_ instead of exporter_->getMaxEventsBuffered() in writeEventHelper() since they hold the same value. --------- Signed-off-by: Batsheva Black --- src/core/telemetry/telemetry.cpp | 17 +++++++++++------ src/core/telemetry/telemetry.h | 1 + 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/src/core/telemetry/telemetry.cpp b/src/core/telemetry/telemetry.cpp index 022725c7..5af5ed3c 100644 --- a/src/core/telemetry/telemetry.cpp +++ b/src/core/telemetry/telemetry.cpp @@ -91,10 +91,10 @@ getExporterName() { void nixlTelemetry::initializeTelemetry() { - const auto buffer_size = nixl::config::getValueDefaulted(TELEMETRY_BUFFER_SIZE_VAR, - DEFAULT_TELEMETRY_BUFFER_SIZE); + maxBufferedEvents_ = nixl::config::getValueDefaulted(TELEMETRY_BUFFER_SIZE_VAR, + DEFAULT_TELEMETRY_BUFFER_SIZE); - if (buffer_size == 0) { + if (maxBufferedEvents_ == 0) { throw std::invalid_argument("Telemetry buffer size cannot be 0"); } @@ -112,7 +112,7 @@ nixlTelemetry::initializeTelemetry() { throw std::runtime_error("Failed to load telemetry plugin: " + *exporter_name); } - const nixlTelemetryExporterInitParams init_params{agentName_, buffer_size}; + const nixlTelemetryExporterInitParams init_params{agentName_, maxBufferedEvents_}; exporter_ = plugin_handle->createExporter(init_params); if (!exporter_) { NIXL_ERROR << "Failed to create telemetry exporter: " << *exporter_name; @@ -134,8 +134,7 @@ nixlTelemetry::initializeTelemetry() { bool nixlTelemetry::writeEventHelper() { std::vector next_queue; - // assume next buffer will be the same size as the current one - next_queue.reserve(exporter_->getMaxEventsBuffered()); + next_queue.reserve(maxBufferedEvents_); { std::lock_guard lock(mutex_); events_.swap(next_queue); @@ -172,6 +171,9 @@ nixlTelemetry::updateData(const std::string &event_name, uint64_t value) { // agent can be multi-threaded std::lock_guard lock(mutex_); + if (events_.size() >= maxBufferedEvents_) { + return; + } events_.emplace_back(std::chrono::duration_cast( std::chrono::system_clock::now().time_since_epoch()) .count(), @@ -235,6 +237,9 @@ nixlTelemetry::addXferTime(std::chrono::microseconds xfer_time, bool is_write, u .count(); const std::lock_guard lock(mutex_); + if (events_.size() + 3 > maxBufferedEvents_) { + return; + } events_.emplace_back(time, nixl_telemetry_category_t::NIXL_TELEMETRY_PERFORMANCE, "agent_xfer_time", diff --git a/src/core/telemetry/telemetry.h b/src/core/telemetry/telemetry.h index ebfa93a3..dc0bae91 100644 --- a/src/core/telemetry/telemetry.h +++ b/src/core/telemetry/telemetry.h @@ -85,6 +85,7 @@ class nixlTelemetry { std::unique_ptr exporter_; std::unique_ptr> buffer_; std::vector events_; + size_t maxBufferedEvents_; std::mutex mutex_; asio::thread_pool pool_; periodicTask writeTask_; From 2bdcdb5fea25c5f99aba035aae3c6cc332c16ee3 Mon Sep 17 00:00:00 2001 From: Itay Alroy <75032521+itayalroy@users.noreply.github.com> Date: Mon, 13 Apr 2026 17:00:52 +0300 Subject: [PATCH 18/59] nixl_ep: Handle planned SIGTERM in elastic test (#1500) elastic.py intentionally sends SIGTERM to simulate rank failure during the elastic test, but PR #1427 treats those expected exits as worker errors. Ignore planned SIGTERM exits when checking worker exit codes after the test completes. --- examples/device/ep/tests/elastic/elastic.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/examples/device/ep/tests/elastic/elastic.py b/examples/device/ep/tests/elastic/elastic.py index 642d340b..f68a2e0d 100644 --- a/examples/device/ep/tests/elastic/elastic.py +++ b/examples/device/ep/tests/elastic/elastic.py @@ -67,7 +67,10 @@ def handle_sigterm( if buffer is not None and buffer.runtime is not None: buffer.destroy() # to invalidate local MD del buffer - sys.exit(1) + + # Continue with default signal handler + signal.signal(signum, signal.SIG_DFL) + signal.raise_signal(signum) def self_kill(): @@ -645,11 +648,11 @@ def main(): daemon=False, start_method="spawn", ) - failed = [] for i, p in enumerate(ctx.processes): p.join() - if p.exitcode != 0: + # Ignore expected fault-tolerance SIGTERM exits. + if p.exitcode not in (0, -signal.SIGTERM): failed.append((i, p.exitcode)) if failed: raise RuntimeError( From 44af005498846a55c64cd516e05aa92748070115 Mon Sep 17 00:00:00 2001 From: ovidiusm Date: Mon, 13 Apr 2026 16:18:11 +0200 Subject: [PATCH 19/59] Add dockerfiles for inference images (#1477) * Integration build scripts for vllm Signed-off-by: Ovidiu Mara * Add copyright Signed-off-by: Ovidiu Mara * Remove CUDA version parameter, we can use CUDA_VERSION from base image instead Signed-off-by: Ovidiu Mara * Address AI comments Signed-off-by: Ovidiu Mara --------- Signed-off-by: Ovidiu Mara --- contrib/Dockerfile.sglang | 29 +++++++++++++++++++++++++++++ contrib/Dockerfile.vllm | 30 ++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+) create mode 100644 contrib/Dockerfile.sglang create mode 100644 contrib/Dockerfile.vllm diff --git a/contrib/Dockerfile.sglang b/contrib/Dockerfile.sglang new file mode 100644 index 00000000..5716e9af --- /dev/null +++ b/contrib/Dockerfile.sglang @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +ARG BASE_IMAGE=lmsysorg/sglang:latest +FROM --platform=$TARGETPLATFORM ${BASE_IMAGE} + +# Build context is expected to contain only wheel files for one architecture. +COPY . /tmp/nixl-wheels + +RUN set -eux; \ + ls -lah /tmp/nixl-wheels/*.whl ; \ + pip install --no-cache-dir --no-deps --force-reinstall /tmp/nixl-wheels/nixl-*.whl /tmp/nixl-wheels/nixl_cu12*cp312*.whl; \ + CUDA_MAJOR="${CUDA_VERSION%%.*}"; \ + if [ "${CUDA_MAJOR}" = "13" ]; then \ + pip install --no-cache-dir --no-deps --force-reinstall /tmp/nixl-wheels/nixl_cu13*cp312*.whl; \ + fi; \ + rm -rf /tmp/nixl-wheels diff --git a/contrib/Dockerfile.vllm b/contrib/Dockerfile.vllm new file mode 100644 index 00000000..09bbf9ed --- /dev/null +++ b/contrib/Dockerfile.vllm @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +ARG BASE_IMAGE=vllm/vllm-openai:latest +FROM --platform=$TARGETPLATFORM ${BASE_IMAGE} + +# Build context is expected to contain only wheel files for one architecture. +COPY . /tmp/nixl-wheels + +RUN set -eux; \ + ls /tmp/nixl-wheels/*.whl >/dev/null; \ + uv pip install --system --no-deps --force-reinstall /tmp/nixl-wheels/nixl-*.whl /tmp/nixl-wheels/nixl_cu12*cp312*.whl; \ + uv pip install --system quart; \ + rm -rf /tmp/nixl-wheels + +RUN curl -fL --retry 3 --retry-delay 2 \ + https://raw.githubusercontent.com/vllm-project/vllm/b73b5b06290c0d1439b09db71eef15fe59bc1fbb/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \ + -o toy_proxy_server.py From 8daa52184bbd7c4fe499c73515fb5e810ba721ed Mon Sep 17 00:00:00 2001 From: Guy Ealey Morag Date: Mon, 13 Apr 2026 16:18:52 +0200 Subject: [PATCH 20/59] Fix IB warning test failure in CI (#1521) * Count all matches in LogProblemCounter * Fix version check for WarnWhenIbPresentButRdmaNotSupported * Use constexpr variable for ucp version * Exit after the loop if it matched any regex --- src/plugins/ucx/ucx_utils.cpp | 3 +-- src/plugins/ucx/ucx_utils.h | 4 ++++ test/gtest/common.cpp | 8 +++++++- test/gtest/hw_warning_test.cpp | 6 ++---- 4 files changed, 14 insertions(+), 7 deletions(-) diff --git a/src/plugins/ucx/ucx_utils.cpp b/src/plugins/ucx/ucx_utils.cpp index 1e15aafd..e5edabd0 100644 --- a/src/plugins/ucx/ucx_utils.cpp +++ b/src/plugins/ucx/ucx_utils.cpp @@ -640,8 +640,7 @@ nixlUcxContext::warnAboutHardwareSupportMismatch() const { << "GPU memory is not supported."; } - if (ucpVersion_ >= UCP_VERSION(1, 22)) { - // `UCS_MEMORY_TYPE_RDMA` to be checked explicitly only from UCX 1.22 + if (ucpVersion_ >= ucp_version_mem_type_rdma) { if (hw_info.numIbDevices > 0 && !UCS_BIT_GET(attr.memory_types, UCS_MEMORY_TYPE_RDMA)) { NIXL_WARN << hw_info.numIbDevices << " IB device(s) were detected, but accelerated IB support was not found! " diff --git a/src/plugins/ucx/ucx_utils.h b/src/plugins/ucx/ucx_utils.h index 00c38b17..9eafdfaf 100644 --- a/src/plugins/ucx/ucx_utils.h +++ b/src/plugins/ucx/ucx_utils.h @@ -34,6 +34,10 @@ enum class nixl_ucx_mt_t { SINGLE, CTX, WORKER }; inline constexpr std::string_view nixl_ucx_err_handling_param_name = "ucx_error_handling_mode"; +// The API `ucp_context_query(ctx, &attr)` sets `UCS_MEMORY_TYPE_RDMA` in `attr.memory_types` +// field only from UCX 1.22 +inline constexpr unsigned ucp_version_mem_type_rdma = UCP_VERSION(1, 22); + template [[nodiscard]] constexpr auto enumToInteger(const Enum e) noexcept { diff --git a/test/gtest/common.cpp b/test/gtest/common.cpp index 650ad176..c715fd69 100644 --- a/test/gtest/common.cpp +++ b/test/gtest/common.cpp @@ -182,15 +182,21 @@ LogProblemCounter::Send(const absl::LogEntry &entry) { return; } + bool matched = false; const std::string msg(entry.text_message()); { const std::lock_guard lock(log_problem_mutex); for (auto &[rx, count] : log_problem_ignore) { if (std::regex_search(msg, rx)) { + matched = true; ++count; - return; } } + + if (matched) { + return; + } + ++global_problem_count; } diff --git a/test/gtest/hw_warning_test.cpp b/test/gtest/hw_warning_test.cpp index 346b215a..83228b3b 100644 --- a/test/gtest/hw_warning_test.cpp +++ b/test/gtest/hw_warning_test.cpp @@ -73,12 +73,10 @@ TEST_F(HardwareWarningTest, WarnWhenGpuPresentButCudaNotSupported) { /** * Test that a warning is logged when IB devices are present but UCX * RDMA support is not available. - * - * Note: This warning only triggers for UCX >= 1.21. */ TEST_F(HardwareWarningTest, WarnWhenIbPresentButRdmaNotSupported) { - if (ucpVersion_ < UCP_VERSION(1, 21)) { - GTEST_SKIP() << "UCX version is less than 1.21, skipping test"; + if (ucpVersion_ < ucp_version_mem_type_rdma) { + GTEST_SKIP() << "UCX version too old for RDMA memory type check, skipping test"; } const auto &hw_info = nixl::hwInfo::instance(); From 4b998db09ea2c52df7ee48698b4e64bb7b668c2e Mon Sep 17 00:00:00 2001 From: Stary Date: Mon, 13 Apr 2026 22:38:39 +0800 Subject: [PATCH 21/59] build: use Ninja + memory-aware parallelism for Mooncake build (#1485) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * build: use Ninja + memory-aware parallelism for Mooncake build Upgrade Mooncake from v0.3.9 to v0.3.10.post1 and switch from Make to Ninja generator. v0.3.10.post1 includes limit_jobs.cmake (Mooncake PR#1718) which auto-detects available memory and creates Ninja job pools to cap compile and link parallelism separately. Previously the build hardcoded `make -j4` as a workaround for OOM during linking. With Ninja job pools, compilation can use all cores while memory-heavy link steps are automatically throttled — no manual -j tuning needed. Signed-off-by: staryxchen * docs: update Mooncake version in ATTRIBUTIONS-CPP.md Bump Mooncake attribution from v0.3.9 to v0.3.10.post1 to match the version upgrade in .gitlab/build.sh. Signed-off-by: staryxchen * build: disable Mooncake store module (nixl only needs transfer engine) Add -DWITH_STORE=OFF to skip building mooncake-store, which nixl does not use (the Mooncake plugin only links libtransfer_engine). This also works around a Pimpl compilation error in S3SnapshotObjectStore introduced in v0.3.10.post1 (see https://github.com/kvcache-ai/Mooncake/pull/1796). Signed-off-by: staryxchen * ci: bump CI_IMAGE_TAG to 20260402-1 Mooncake version and build flags changed (v0.3.10.post1, Ninja, -DWITH_STORE=OFF), so the CI image must be rebuilt. Signed-off-by: staryxchen --------- Signed-off-by: staryxchen Signed-off-by: ovidiusm Co-authored-by: Mikhail Brinskiy Co-authored-by: ovidiusm --- .ci/jenkins/lib/build-matrix.yaml | 2 +- .ci/jenkins/lib/test-matrix.yaml | 2 +- .gitlab/build.sh | 8 ++++---- ATTRIBUTIONS-CPP.md | 2 +- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.ci/jenkins/lib/build-matrix.yaml b/.ci/jenkins/lib/build-matrix.yaml index 380dcb3f..f33e6db0 100644 --- a/.ci/jenkins/lib/build-matrix.yaml +++ b/.ci/jenkins/lib/build-matrix.yaml @@ -43,7 +43,7 @@ env: TEST_TIMEOUT: 30 UCX_TLS: "^shm" STORAGE_DRIVER: 'overlay' - CI_IMAGE_TAG: "20260323-1" + CI_IMAGE_TAG: "20260402-1" runs_on_dockers: diff --git a/.ci/jenkins/lib/test-matrix.yaml b/.ci/jenkins/lib/test-matrix.yaml index c425aa60..54621f97 100644 --- a/.ci/jenkins/lib/test-matrix.yaml +++ b/.ci/jenkins/lib/test-matrix.yaml @@ -49,7 +49,7 @@ env: SLURM_JOB_TIMEOUT: '02:20:00' SLURM_IMMEDIATE_TIMEOUT: "3600" STORAGE_DRIVER: overlay - CI_IMAGE_TAG: "20260323-1" + CI_IMAGE_TAG: "20260402-1" empty_volumes: - {mountPath: /var/lib/containers/storage, memory: false} diff --git a/.gitlab/build.sh b/.gitlab/build.sh index 5d2fe021..a66818b0 100755 --- a/.gitlab/build.sh +++ b/.gitlab/build.sh @@ -282,15 +282,15 @@ else ( \ cd ${TMPDIR} && \ - MOONCAKE_VERSION="${MOONCAKE_VERSION:-v0.3.9}" && \ + MOONCAKE_VERSION="${MOONCAKE_VERSION:-v0.3.10.post1}" && \ echo "MOONCAKE_VERSION: ${MOONCAKE_VERSION}" && \ git clone --depth 1 --branch "${MOONCAKE_VERSION}" https://github.com/kvcache-ai/Mooncake.git && \ cd Mooncake && \ $SUDO bash dependencies.sh -y && \ mkdir build && cd build && \ - cmake .. -DBUILD_SHARED_LIBS=ON && \ - make -j4 && \ - $SUDO make install && \ + cmake .. -DBUILD_SHARED_LIBS=ON -DWITH_STORE=OFF -G Ninja && \ + ninja && \ + $SUDO ninja install && \ $SUDO ldconfig && \ cd .. && \ rm -rf Mooncake diff --git a/ATTRIBUTIONS-CPP.md b/ATTRIBUTIONS-CPP.md index 8522ad8b..078fb352 100644 --- a/ATTRIBUTIONS-CPP.md +++ b/ATTRIBUTIONS-CPP.md @@ -5311,7 +5311,7 @@ Random Hacker. That's all there is to it! ``` -## Mooncake - v0.3.2.post1 +## Mooncake - v0.3.10.post1 - **Repository URL**: https://github.com/kvcache-ai/Mooncake - **License URL**: https://github.com/kvcache-ai/Mooncake/blob/main/LICENSE From dd9af4ce402c2ebb3a59ea566e5d219e4937c610 Mon Sep 17 00:00:00 2001 From: Raul Akhmetshin <74596089+rakhmets@users.noreply.github.com> Date: Mon, 13 Apr 2026 17:38:30 +0200 Subject: [PATCH 22/59] PLUGINS/UCX: Removed unused class field. (#1512) Signed-off-by: Raul Akhmetshin --- src/plugins/ucx/ucx_backend.cpp | 9 +-------- src/plugins/ucx/ucx_backend.h | 1 - 2 files changed, 1 insertion(+), 9 deletions(-) diff --git a/src/plugins/ucx/ucx_backend.cpp b/src/plugins/ucx/ucx_backend.cpp index 81626770..552e3810 100644 --- a/src/plugins/ucx/ucx_backend.cpp +++ b/src/plugins/ucx/ucx_backend.cpp @@ -932,21 +932,14 @@ nixl_status_t nixlUcxEngine::loadRemoteConnInfo (const std::string &remote_agent nixlSerDes::_stringToBytes(addr.data(), remote_conn_info, size); std::shared_ptr conn = std::make_shared(); - bool error = false; for (auto &uw: uws) { auto result = uw->connect(addr.data(), size); if (!result.ok()) { - error = true; - break; + return NIXL_ERR_BACKEND; } conn->eps.push_back(std::move(*result)); } - if (error) - return NIXL_ERR_BACKEND; - - conn->remoteAgent = remote_agent; - remoteConnMap.insert({remote_agent, conn}); return NIXL_SUCCESS; diff --git a/src/plugins/ucx/ucx_backend.h b/src/plugins/ucx/ucx_backend.h index 50662936..19031184 100644 --- a/src/plugins/ucx/ucx_backend.h +++ b/src/plugins/ucx/ucx_backend.h @@ -42,7 +42,6 @@ enum ucx_cb_op_t { NOTIF_STR }; class nixlUcxConnection : public nixlBackendConnMD { private: - std::string remoteAgent; std::vector> eps; public: From 02691a5209888e9cec40d7eb2daa998096a91ad4 Mon Sep 17 00:00:00 2001 From: lishapira Date: Mon, 13 Apr 2026 20:26:04 +0300 Subject: [PATCH 23/59] nixl_ep: fix GCC maybe-uninitialized warning in ht_dispatch (#1525) Use raw pointers instead of dereferencing std::optional in ternary expressions, following the same pattern as recv_topk_idx_ptr and recv_topk_weights_ptr. This prevents GCC's -Wmaybe-uninitialized warning, which becomes a build error with -Werror. Made-with: Cursor --- examples/device/ep/csrc/nixl_ep.cpp | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/examples/device/ep/csrc/nixl_ep.cpp b/examples/device/ep/csrc/nixl_ep.cpp index 2fb3b14d..3c8ed6fe 100644 --- a/examples/device/ep/csrc/nixl_ep.cpp +++ b/examples/device/ep/csrc/nixl_ep.cpp @@ -766,12 +766,22 @@ Buffer::ht_dispatch(const torch::Tensor& x, const std::optional& auto recv_gbl_channel_prefix_matrix = std::optional(); auto send_rdma_head = std::optional(); auto send_nvl_head = std::optional(); + void* recv_src_meta_ptr = nullptr; + int* recv_rdma_channel_prefix_matrix_ptr = nullptr; + int* recv_gbl_channel_prefix_matrix_ptr = nullptr; + int* send_rdma_head_ptr = nullptr; + int* send_nvl_head_ptr = nullptr; if (not cached_mode) { recv_src_meta = torch::empty({num_recv_tokens, ht::get_source_meta_bytes()}, dtype(torch::kByte).device(torch::kCUDA)); recv_rdma_channel_prefix_matrix = torch::empty({num_rdma_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); recv_gbl_channel_prefix_matrix = torch::empty({num_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); send_rdma_head = torch::empty({num_tokens, num_rdma_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); send_nvl_head = torch::empty({num_rdma_recv_tokens, NUM_MAX_NVL_PEERS}, dtype(torch::kInt32).device(torch::kCUDA)); + recv_src_meta_ptr = recv_src_meta->data_ptr(); + recv_rdma_channel_prefix_matrix_ptr = recv_rdma_channel_prefix_matrix->data_ptr(); + recv_gbl_channel_prefix_matrix_ptr = recv_gbl_channel_prefix_matrix->data_ptr(); + send_rdma_head_ptr = send_rdma_head->data_ptr(); + send_nvl_head_ptr = send_nvl_head->data_ptr(); } // Assign pointers @@ -792,11 +802,11 @@ Buffer::ht_dispatch(const torch::Tensor& x, const std::optional& // Launch data dispatch // NOTES: the buffer size checks are moved into the `.cu` file ht::dispatch(recv_x.data_ptr(), recv_x_scales_ptr, recv_topk_idx_ptr, recv_topk_weights_ptr, - cached_mode ? nullptr : recv_src_meta->data_ptr(), + recv_src_meta_ptr, x.data_ptr(), x_scales_ptr, topk_idx_ptr, topk_weights_ptr, - cached_mode ? nullptr : send_rdma_head->data_ptr(), cached_mode ? nullptr : send_nvl_head->data_ptr(), - cached_mode ? nullptr : recv_rdma_channel_prefix_matrix->data_ptr(), - cached_mode ? nullptr : recv_gbl_channel_prefix_matrix->data_ptr(), + send_rdma_head_ptr, send_nvl_head_ptr, + recv_rdma_channel_prefix_matrix_ptr, + recv_gbl_channel_prefix_matrix_ptr, rdma_channel_prefix_matrix.data_ptr(), recv_rdma_rank_prefix_sum.data_ptr(), gbl_channel_prefix_matrix.data_ptr(), recv_gbl_rank_prefix_sum.data_ptr(), is_token_in_rank.data_ptr(), From f48000a2172eece0f79b5d6425e8f2ed114139d2 Mon Sep 17 00:00:00 2001 From: Adit Ranadive Date: Mon, 13 Apr 2026 07:36:35 -1000 Subject: [PATCH 24/59] Fix/nixlbench poll shutdown (#1421) * nixlbench: fix poll() infinite loop, add SIGTERM handling - Register SIGTERM handler in main.cpp alongside existing SIGINT handler - Add !signaled() check to both do-while conditions in poll() so the loop exits immediately when terminate is set - Add checkLiveness() lambda in poll() that queries arePeersAlive() every 5s; sets terminate=1 if peer's etcd key is gone, triggering the !signaled() break - Thread &terminate through execTransfer() and execTransferIterations() so the hot transfer loop breaks early on signal instead of running all N iterations Signed-off-by: Adit Ranadive * nixlbench: add etcd lease-based peer liveness detection Each rank attaches its presence key to an etcd v3 lease (TTL=15s). A standalone etcd::KeepAlive object (independent gRPC channel) sends periodic renewals. If the process dies (any signal including SIGKILL), the KeepAlive thread dies, the gRPC connection closes, and etcd auto-deletes the key after the TTL. - runtime.h: add virtual arePeersAlive() default (returns true) to base class - etcd_rt.h: add KeepAlive member + arePeersAlive() override declaration - etcd_rt.cpp: - setup(): attach rank key to lease via standalone KeepAlive constructor (avoids sharing the main SyncClient gRPC channel) - ~xferBenchEtcdRT(): Cancel()+reset() keepalive before rmdir() on clean exit - arePeersAlive(): query each peer's rank key; detect absence via resp.value().key().empty() (is_ok() returns true even for missing keys) - Store etcd endpoints string for standalone KeepAlive constructor - worker.cpp: use std::_Exit() instead of exit() in synchronize() to bypass gRPC atexit handlers that deadlock with the KeepAlive background thread Signed-off-by: Adit Ranadive * Fix clang Signed-off-by: Adit Ranadive * Fix copyright Signed-off-by: Adit Ranadive * nixlbench: address style review comments - arePeersAlive(): add [[nodiscard]] attribute - keepalive: change shared_ptr to unique_ptr (ownership is exclusive) - LEASE_TTL_S: move from class header to anonymous namespace in .cpp, rename to lease_ttl_s (ALL_CAPS reserved for macros) - checkLiveness: use chrono_literals (5s) and direct duration comparison instead of duration_cast + count(); add const to local 'now' Signed-off-by: Adit Ranadive * nixlbench: address safety review comments - terminate flag: change to std::atomic to prevent data races between the signal handler, poll() loop, and OpenMP transfer threads that all read or write it concurrently - destructor null guard: check client \!= nullptr before calling rmdir() in ~xferBenchEtcdRT(); setup() may fail after partial init, leaving client null - cleanupForExit(): add best-effort namespace cleanup before _Exit() so the non-leased 'size' key (and any others) are removed; a stale size key would corrupt rank allocation in the next run of the same benchmark_group Signed-off-by: Adit Ranadive * nixlbench: fix python_bindings.cpp for atomic terminate Signed-off-by: Adit Ranadive * nixlbench: remove redundant keepalive.reset() from destructor unique_ptr destroys the held object automatically when the destructor returns; the explicit reset() was a no-op. cleanupForExit() retains reset() because it runs before _Exit() which skips destructors. Update header guard. Signed-off-by: Adit Ranadive * nixlbench: handle keepalive and rank-key registration failures in setup() Wrap KeepAlive construction in try/catch and check the put() result. Both paths unlock the etcd lock and return -1 on failure, preventing setup() from reporting success with the rank unregistered. Signed-off-by: Adit Ranadive * nixlbench: address coderabbit review comments - python_bindings: replace atomic* constructor binding with a 3-arg lambda factory; use raw pointer return so it works with all pybind11 versions (unique_ptr return requires pybind11 >= 2.6) - nixl_worker: add terminate check in recreate_per_iteration loop to match the fast-exit behaviour already present in the standard path - worker: add static_assert for lock-free atomic and use explicit fetch_add(1, relaxed) in signalHandler instead of prefix ++ Signed-off-by: Adit Ranadive * nixlbench: register leased rank key before bumping size counter A crash between put("size") and put("rank/N", lease) left a stale non-leased size key that poisoned subsequent runs. Register the leased rank key first so a mid-setup crash leaves size unchanged; only bump size once rank/ is successfully attached to the lease. Also add terminate_ptr check inside execSingleTransfer() polling loop so threads don't spin in getXferStatus() after SIGTERM is received. Signed-off-by: Adit Ranadive * nixlbench: rename arePeersAlive and use consistent atomic ordering - Rename arePeersAlive() -> areAllPeersAlive() for clarity - Use load(memory_order_relaxed) consistently at all terminate_ptr read sites instead of implicit seq_cst via operator*; use store(memory_order_relaxed) for the liveness-check write Signed-off-by: Adit Ranadive * nixlbench: use default memory ordering for terminate flag Drop explicit memory_order_relaxed from all load()/store()/fetch_add() calls on the terminate flag. The default seq_cst is fine for a flag checked a few times per second; cleaner and more readable. Signed-off-by: Adit Ranadive --------- Signed-off-by: Adit Ranadive --- benchmark/nixlbench/src/main.cpp | 1 + .../nixlbench/src/runtime/etcd/etcd_rt.cpp | 73 ++++++++++++++++++- .../nixlbench/src/runtime/etcd/etcd_rt.h | 29 ++++++-- .../src/runtime/etcd/python_bindings.cpp | 11 ++- benchmark/nixlbench/src/runtime/runtime.h | 11 +++ .../nixlbench/src/worker/nixl/nixl_worker.cpp | 73 ++++++++++++++++--- benchmark/nixlbench/src/worker/worker.cpp | 22 ++++-- benchmark/nixlbench/src/worker/worker.h | 3 +- 8 files changed, 189 insertions(+), 34 deletions(-) diff --git a/benchmark/nixlbench/src/main.cpp b/benchmark/nixlbench/src/main.cpp index bcb00942..aec1fdcd 100644 --- a/benchmark/nixlbench/src/main.cpp +++ b/benchmark/nixlbench/src/main.cpp @@ -191,6 +191,7 @@ int main(int argc, char *argv[]) { } std::signal(SIGINT, worker_ptr->signalHandler); + std::signal(SIGTERM, worker_ptr->signalHandler); // Ensure all processes are ready before exchanging metadata ret = worker_ptr->synchronizeStart(); diff --git a/benchmark/nixlbench/src/runtime/etcd/etcd_rt.cpp b/benchmark/nixlbench/src/runtime/etcd/etcd_rt.cpp index 394e3491..c749f62f 100644 --- a/benchmark/nixlbench/src/runtime/etcd/etcd_rt.cpp +++ b/benchmark/nixlbench/src/runtime/etcd/etcd_rt.cpp @@ -22,16 +22,22 @@ #include #include #include +#include #include #include "etcd_rt.h" #define ETCD_EP_DEFAULT "http://localhost:2379" +namespace { +// Lease TTL in seconds: rank key auto-expires this long after the last keepalive. +constexpr int lease_ttl_s = 15; +} // namespace + // ETCD Runtime implementation xferBenchEtcdRT::xferBenchEtcdRT(const std::string &benchmark_group, const std::string &etcd_endpoints, const int size, - int *terminate_input) { + std::atomic *terminate_input) { // Store parameters for later use in setup stored_etcd_endpoints = etcd_endpoints.empty() ? ETCD_EP_DEFAULT : etcd_endpoints; global_size = size; @@ -75,9 +81,30 @@ xferBenchEtcdRT::setup() { my_rank = 0; } - // Update registration information + // Register the rank key with a lease before bumping "size"; a crash before + // registration leaves size unchanged. Standalone KeepAlive uses its own + // gRPC channel to avoid races with the main client. + try { + keepalive = std::make_unique(stored_etcd_endpoints, lease_ttl_s); + } + catch (const std::exception &e) { + client->unlock(lock_response.lock_key()); + std::cerr << "Failed to start etcd keepalive: " << e.what() << std::endl; + return -1; + } + + const auto rank_response = client->put(makeKey("rank", my_rank), "active", keepalive->Lease()); + if (!rank_response.is_ok()) { + client->unlock(lock_response.lock_key()); + keepalive->Cancel(); + keepalive.reset(); + std::cerr << "Failed to register rank key: " << rank_response.error_message() << std::endl; + return -1; + } + + // Bump the non-leased "size" counter only after the rank key is successfully + // registered with a lease. client->put(makeKey("size"), std::to_string(my_rank + 1)); - client->put(makeKey("rank", my_rank), "active"); // Release the lock client->unlock(lock_response.lock_key()); @@ -99,8 +126,46 @@ xferBenchEtcdRT::setup() { } xferBenchEtcdRT::~xferBenchEtcdRT() { + // Stop lease keepalive before cleanup so the rank key is removed by rmdir + // rather than expiring after the TTL + if (keepalive) { + keepalive->Cancel(); + } // All ranks delete, as some could be missing if ETCD state is confused - client->rmdir(makeKey(""), true); + if (client) { + client->rmdir(makeKey(""), true); + } +} + +void +xferBenchEtcdRT::cleanupForExit() { + // Cancel keepalive first so the gRPC stream closes before we touch etcd. + if (keepalive) { + keepalive->Cancel(); + keepalive.reset(); + } + // Remove the whole namespace so non-leased keys (e.g. "size") don't + // linger and poison the next run in the same benchmark_group. + if (client) { + client->rmdir(makeKey(""), true); + } +} + +bool +xferBenchEtcdRT::areAllPeersAlive() { + for (int r = 0; r < global_size; r++) { + if (r == my_rank) continue; + auto resp = client->get(makeKey("rank", r)); + // For a single-key get(), is_ok() reflects gRPC success, NOT key existence: + // a missing key returns is_ok()=true with an empty value().key(). + // resp.value().key() is non-empty only when the key actually exists in etcd. + if (!resp.is_ok() || resp.value().key().empty()) { + std::cerr << "nixlbench: peer rank " << r + << " key missing (or etcd error) -- peer may have died" << std::endl; + return false; + } + } + return true; } int diff --git a/benchmark/nixlbench/src/runtime/etcd/etcd_rt.h b/benchmark/nixlbench/src/runtime/etcd/etcd_rt.h index d93b2f91..5d051342 100644 --- a/benchmark/nixlbench/src/runtime/etcd/etcd_rt.h +++ b/benchmark/nixlbench/src/runtime/etcd/etcd_rt.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -15,9 +15,10 @@ * limitations under the License. */ -#ifndef _ETCD_RT_H -#define _ETCD_RT_H +#ifndef NIXL_BENCHMARK_NIXLBENCH_SRC_RUNTIME_ETCD_ETCD_RT_H +#define NIXL_BENCHMARK_NIXLBENCH_SRC_RUNTIME_ETCD_ETCD_RT_H +#include #include #include #include @@ -27,9 +28,10 @@ #include #include "runtime/runtime.h" -// Forward declaration for etcd client +// Forward declarations for etcd client namespace etcd { class SyncClient; +class KeepAlive; } enum xferBenchEtcdMsgType { XFER_BENCH_ETCD_MSG_TYPE_INT = 1, XFER_BENCH_ETCD_MSG_TYPE_CHAR = 2 }; @@ -46,10 +48,15 @@ class xferBenchEtcdRT: public xferBenchRT { std::string benchmark_group; std::unique_ptr client; + // Lease-based peer liveness: each rank's key is attached to a lease that + // auto-expires if the process dies (including SIGKILL). areAllPeersAlive() + // checks whether all peer rank keys are still present in etcd. + std::unique_ptr keepalive; + int my_rank; // Rank information int global_size; uint64_t barrier_gen; - int *terminate; + std::atomic *terminate; bool error() const { return terminate != nullptr && *terminate; }; bool should_retry(int value, int max = 60) const { @@ -73,7 +80,7 @@ class xferBenchEtcdRT: public xferBenchRT { xferBenchEtcdRT(const std::string &benchmark_group, const std::string &etcd_endpoints, const int size, - int *terminate = nullptr); + std::atomic *terminate = nullptr); ~xferBenchEtcdRT(); // Setup function to initialize connection and perform registration @@ -93,6 +100,14 @@ class xferBenchEtcdRT: public xferBenchRT { // Barrier synchronization int barrier(const std::string& barrier_id) override; + + // Check if all peer rank keys are still present in etcd + bool + areAllPeersAlive() override; + + // Cancel keepalive and remove namespace keys before a forced _Exit() + void + cleanupForExit() override; }; -#endif // _ETCD_RT_H +#endif // NIXL_BENCHMARK_NIXLBENCH_SRC_RUNTIME_ETCD_ETCD_RT_H diff --git a/benchmark/nixlbench/src/runtime/etcd/python_bindings.cpp b/benchmark/nixlbench/src/runtime/etcd/python_bindings.cpp index d1653aca..d7a25836 100644 --- a/benchmark/nixlbench/src/runtime/etcd/python_bindings.cpp +++ b/benchmark/nixlbench/src/runtime/etcd/python_bindings.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -26,11 +26,14 @@ PYBIND11_MODULE (etcd_runtime, m) { m.doc() = "Python bindings for ETCD runtime"; py::class_(m, "EtcdRuntime") - .def(py::init(), + .def(py::init([](const std::string &benchmark_group, + const std::string &etcd_endpoints, + const int size) { + return new xferBenchEtcdRT(benchmark_group, etcd_endpoints, size, nullptr); + }), py::arg("benchmark_group"), py::arg("etcd_endpoints"), - py::arg("size"), - py::arg("terminate") = nullptr) + py::arg("size")) .def("setup", &xferBenchEtcdRT::setup) .def("get_rank", &xferBenchEtcdRT::getRank) .def("get_size", &xferBenchEtcdRT::getSize) diff --git a/benchmark/nixlbench/src/runtime/runtime.h b/benchmark/nixlbench/src/runtime/runtime.h index e3c40023..dd309e3e 100644 --- a/benchmark/nixlbench/src/runtime/runtime.h +++ b/benchmark/nixlbench/src/runtime/runtime.h @@ -41,6 +41,17 @@ class xferBenchRT { // Add a barrier function to synchronize all processes virtual int barrier(const std::string& barrier_id) = 0; + + // Check if all peer processes are still alive; returns true by default + [[nodiscard]] virtual bool + areAllPeersAlive() { + return true; + } + + // Best-effort cleanup of runtime state (e.g. etcd keys) before a + // forced exit that bypasses normal destructors. + virtual void + cleanupForExit() {} }; #endif // NIXL_BENCHMARK_NIXLBENCH_SRC_RUNTIME_RUNTIME_H diff --git a/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp b/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp index c18b6fcc..23a36711 100644 --- a/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp +++ b/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp @@ -18,6 +18,7 @@ #include "worker/nixl/nixl_worker.h" #include #include +#include #include #if HAVE_CUDA #include @@ -1203,10 +1204,14 @@ static inline nixl_status_t execSingleTransfer(nixlAgent *agent, nixlXferReqH *req, xferBenchTimer &timer, - xferBenchStats &thread_stats) { + xferBenchStats &thread_stats, + const std::atomic *terminate_ptr = nullptr) { nixl_status_t rc = agent->postXferReq(req); thread_stats.post_duration.add(timer.lap()); while (NIXL_IN_PROG == rc) { + if (__builtin_expect(terminate_ptr && terminate_ptr->load(), 0)) { + break; + } rc = agent->getXferStatus(req); } return rc; @@ -1243,7 +1248,8 @@ execTransferIterations(nixlAgent *agent, const int num_iter, xferBenchTimer &timer, xferBenchStats &thread_stats, - const bool recreate_per_iteration) { + const bool recreate_per_iteration, + const std::atomic *terminate_ptr = nullptr) { nixlXferReqH *req = nullptr; nixlTime::us_t total_prepare_duration = 0; @@ -1264,6 +1270,10 @@ execTransferIterations(nixlAgent *agent, if (__builtin_expect(recreate_per_iteration, 0)) { // GUSLI path: Create/execute/release per iteration for (int i = 0; i < num_iter; ++i) { + // Check for signal (SIGTERM/SIGINT) to allow fast exit on peer death + if (__builtin_expect(terminate_ptr && terminate_ptr->load(), 0)) { + return -1; + } nixl_status_t create_rc = agent->createXferReq(op, local_desc, remote_desc, target, req, ¶ms); if (__builtin_expect(create_rc != NIXL_SUCCESS, 0)) { @@ -1273,7 +1283,7 @@ execTransferIterations(nixlAgent *agent, } total_prepare_duration += timer.lap(); - nixl_status_t rc = execSingleTransfer(agent, req, timer, thread_stats); + nixl_status_t rc = execSingleTransfer(agent, req, timer, thread_stats, terminate_ptr); if (__builtin_expect(rc != NIXL_SUCCESS, 0)) { std::cout << "NIXL Xfer failed with status: " << nixlEnumStrings::statusStr(rc) @@ -1293,7 +1303,12 @@ execTransferIterations(nixlAgent *agent, } else { // Standard path: Single request for all iterations for (int i = 0; i < num_iter; ++i) { - nixl_status_t rc = execSingleTransfer(agent, req, timer, thread_stats); + // Check for signal (SIGTERM/SIGINT) to allow fast exit on peer death + if (__builtin_expect(terminate_ptr && terminate_ptr->load(), 0)) { + agent->releaseXferReq(req); + return -1; + } + nixl_status_t rc = execSingleTransfer(agent, req, timer, thread_stats, terminate_ptr); if (__builtin_expect(rc != NIXL_SUCCESS, 0)) { std::cout << "NIXL Xfer failed with status: " << nixlEnumStrings::statusStr(rc) @@ -1321,7 +1336,8 @@ execTransfer(nixlAgent *agent, const nixl_xfer_op_t op, const int num_iter, const int num_threads, - xferBenchStats &stats) { + xferBenchStats &stats, + const std::atomic *terminate_ptr = nullptr) { int ret = 0; stats.clear(); @@ -1357,7 +1373,8 @@ execTransfer(nixlAgent *agent, num_iter, timer, thread_stats, - xferBenchConfig::recreate_xfer); + xferBenchConfig::recreate_xfer, + terminate_ptr); if (__builtin_expect(result != 0, 0)) { ret = result; @@ -1390,8 +1407,14 @@ xferBenchNixlWorker::transfer(size_t block_size, } if (skip > 0) { - ret = execTransfer( - agent, local_iovs, remote_iovs, xfer_op, skip, xferBenchConfig::num_threads, stats); + ret = execTransfer(agent, + local_iovs, + remote_iovs, + xfer_op, + skip, + xferBenchConfig::num_threads, + stats, + &terminate); if (ret < 0) { return std::variant(ret); } @@ -1402,8 +1425,14 @@ xferBenchNixlWorker::transfer(size_t block_size, stats.clear(); - ret = execTransfer( - agent, local_iovs, remote_iovs, xfer_op, num_iter, xferBenchConfig::num_threads, stats); + ret = execTransfer(agent, + local_iovs, + remote_iovs, + xfer_op, + num_iter, + xferBenchConfig::num_threads, + stats, + &terminate); if (ret < 0) { return std::variant(ret); } @@ -1427,16 +1456,36 @@ xferBenchNixlWorker::poll(size_t block_size) { } total_iter = skip + num_iter; + // Periodically check if all peers are still alive via etcd lease keys. + // Fires at most once every liveness_check_interval to avoid + // saturating etcd with get() calls during tight polling loops. + using namespace std::chrono_literals; + auto last_liveness_check = std::chrono::steady_clock::now(); + constexpr auto liveness_check_interval = 5s; + auto checkLiveness = [&]() { + const auto now = std::chrono::steady_clock::now(); + if (now - last_liveness_check >= liveness_check_interval) { + last_liveness_check = now; + if (rt && !rt->areAllPeersAlive()) { + std::cerr << "nixlbench: peer liveness check failed — aborting poll" << std::endl; + terminate.store(1); + } + } + }; + /* Ensure warmup is done*/ do { status = agent->getNotifs(notifs); - } while (status == NIXL_SUCCESS && skip != int(notifs["initiator"].size())); + checkLiveness(); + } while (!signaled() && status == NIXL_SUCCESS && skip != int(notifs["initiator"].size())); synchronize(); /* Polling for actual iterations*/ do { status = agent->getNotifs(notifs); - } while (status == NIXL_SUCCESS && total_iter != int(notifs["initiator"].size())); + checkLiveness(); + } while (!signaled() && status == NIXL_SUCCESS && + total_iter != int(notifs["initiator"].size())); synchronize(); } diff --git a/benchmark/nixlbench/src/worker/worker.cpp b/benchmark/nixlbench/src/worker/worker.cpp index 95d28f02..8ed0c756 100644 --- a/benchmark/nixlbench/src/worker/worker.cpp +++ b/benchmark/nixlbench/src/worker/worker.cpp @@ -68,7 +68,8 @@ class xferBenchNullRT : public xferBenchRT { } }; -static xferBenchRT *createRT(int *terminate) { +static xferBenchRT * +createRT(std::atomic *terminate) { // For storage backends without ETCD endpoints, use null runtime if (xferBenchConfig::isStorageBackend() && xferBenchConfig::etcd_endpoints.empty()) { std::cout << "Using null runtime for storage backend without ETCD" << std::endl; @@ -108,9 +109,15 @@ int xferBenchWorker::synchronize() { if (rt->barrier("sync") != 0) { std::cerr << "Failed to synchronize" << std::endl; - // assuming this is a fatal error, continue benchmarking after synchronization failure does - // not make sense - exit(EXIT_FAILURE); + // Best-effort cleanup of non-leased etcd keys (e.g. the "size" key) + // so they don't poison subsequent runs in the same benchmark_group. + rt->cleanupForExit(); + // Use _Exit() instead of exit() to bypass atexit handlers (e.g. gRPC shutdown). + // exit() would deadlock with the etcd KeepAlive background thread: gRPC shutdown + // waits for open streams to close, but the KeepAlive thread keeps renewing the + // lease stream indefinitely. _Exit() kills all threads immediately, which closes + // the gRPC stream and lets the lease expire on the etcd server side. + std::_Exit(EXIT_FAILURE); } return 0; @@ -168,7 +175,10 @@ bool xferBenchWorker::isTarget() { return ("target" == name); } -int xferBenchWorker::terminate = 0; +static_assert(std::atomic::is_always_lock_free, + "xferBenchWorker::terminate must be lock-free for safe use in signal handlers"); + +std::atomic xferBenchWorker::terminate = 0; void xferBenchWorker::signalHandler(int signal) { static const char msg[] = "Ctrl-C received, exiting...\n"; @@ -177,7 +187,7 @@ void xferBenchWorker::signalHandler(int signal) { auto size = write(stdout_fd, msg, sizeof(msg) - 1); (void)size; - if (++terminate > max_count) { + if (terminate.fetch_add(1) >= max_count) { std::_Exit(EXIT_FAILURE); } } diff --git a/benchmark/nixlbench/src/worker/worker.h b/benchmark/nixlbench/src/worker/worker.h index 8af54eaa..55ce34a0 100644 --- a/benchmark/nixlbench/src/worker/worker.h +++ b/benchmark/nixlbench/src/worker/worker.h @@ -20,6 +20,7 @@ #include "runtime/runtime.h" #include "utils/utils.h" +#include #include #include #include @@ -29,7 +30,7 @@ class xferBenchWorker { protected: std::string name; xferBenchRT *rt; - static int terminate; + static std::atomic terminate; public: xferBenchWorker(); From 6ba5277fb657d173d0330f7c003f666b4939cd7f Mon Sep 17 00:00:00 2001 From: Ye Xiang Date: Mon, 13 Apr 2026 23:20:55 -0700 Subject: [PATCH 25/59] Support Neuron in NIXLBench (#1454) * libfabric: Fix Neuron mem type support regression Signed-off-by: Ye Xiang * feat(nixlbench): add Neuron VRAM support for non-CUDA builds Enable nixlbench VRAM segment allocation, deallocation, and consistency checks on Neuron (Trainium/Inferentia) devices when CUDA is not available. Signed-off-by: Ye Xiang * address review: do-while(0) wrap macro, add #else guard in cleanupBasicDescVram * address review: refactor nixlbench VRAM helpers --------- Signed-off-by: Ye Xiang --- benchmark/nixlbench/src/utils/utils.cpp | 57 ++++-------- .../nixlbench/src/worker/nixl/nixl_worker.cpp | 90 ++++++++++--------- .../nixlbench/src/worker/nixl/nixl_worker.h | 8 +- src/plugins/libfabric/libfabric_backend.cpp | 8 +- 4 files changed, 76 insertions(+), 87 deletions(-) diff --git a/benchmark/nixlbench/src/utils/utils.cpp b/benchmark/nixlbench/src/utils/utils.cpp index ef954e20..8ef9de11 100644 --- a/benchmark/nixlbench/src/utils/utils.cpp +++ b/benchmark/nixlbench/src/utils/utils.cpp @@ -739,6 +739,23 @@ xferBenchUtils::getDevToUse() { return dev_to_use; } +static void +copyVramToHost(void *host_addr, const void *device_addr, size_t len) { + if (neuronCoreCount() > 0) { + CHECK_NEURON_ERROR( + neuronMemcpy(host_addr, (void *)device_addr, len, neuronMemcpyDeviceToHost), + "nrt_tensor_read failed"); + return; + } +#if HAVE_CUDA + CHECK_CUDA_ERROR(cudaMemcpy(host_addr, (void *)device_addr, len, cudaMemcpyDeviceToHost), + "cudaMemcpy failed"); +#else + std::cerr << "VRAM not supported without CUDA or Neuron" << std::endl; + exit(EXIT_FAILURE); +#endif +} + static bool allBytesAre(void *buffer, size_t size, uint8_t value) { uint8_t *byte_buffer = static_cast(buffer); @@ -872,31 +889,13 @@ xferBenchUtils::checkConsistency(std::vector> &iov_lis xferBenchConfig::backend == XFERBENCH_BACKEND_GPUNETIO) { if (xferBenchConfig::op_type == XFERBENCH_OP_READ) { if (xferBenchConfig::initiator_seg_type == XFERBENCH_SEG_TYPE_VRAM) { -#if HAVE_CUDA if (posix_memalign(&addr, xferBenchConfig::page_size, len) != 0) { std::cerr << "Failed to allocate aligned buffer of size: " << len << std::endl; exit(EXIT_FAILURE); } is_allocated = true; - // Assume no CUDA cores exist if Neuron cores are found. - // There are no AWS instance types with both NVIDIA GPUs and Neuron - // accelerators. - if (neuronCoreCount() > 0) { - CHECK_NEURON_ERROR( - neuronMemcpy(addr, (void *)iov.addr, len, neuronMemcpyDeviceToHost), - "nrt_tensor_read failed"); - } else { - CHECK_CUDA_ERROR( - cudaMemcpy(addr, (void *)iov.addr, len, cudaMemcpyDeviceToHost), - "cudaMemcpy failed"); - } -#else - std::cerr << "Failure in consistency check: VRAM segment type not " - "supported without CUDA" - << std::endl; - exit(EXIT_FAILURE); -#endif + copyVramToHost(addr, (void *)iov.addr, len); } else { addr = (void *)iov.addr; } @@ -974,27 +973,9 @@ xferBenchUtils::checkConsistency(std::vector> &iov_lis xferBenchConfig::target_seg_type == XFERBENCH_SEG_TYPE_VRAM) || (xferBenchConfig::op_type == XFERBENCH_OP_READ && xferBenchConfig::initiator_seg_type == XFERBENCH_SEG_TYPE_VRAM)) { -#if HAVE_CUDA addr = calloc(1, len); is_allocated = true; - // Assume no CUDA cores exist if Neuron cores are found. - // There are no AWS instance types with both NVIDIA GPUs and Neuron - // accelerators. - if (neuronCoreCount() > 0) { - CHECK_NEURON_ERROR( - neuronMemcpy(addr, (void *)iov.addr, len, neuronMemcpyDeviceToHost), - "nrt_tensor_read failed"); - } else { - CHECK_CUDA_ERROR( - cudaMemcpy(addr, (void *)iov.addr, len, cudaMemcpyDeviceToHost), - "cudaMemcpy failed"); - } -#else - std::cerr << "Failure in consistency check: VRAM segment type not supported " - "without CUDA" - << std::endl; - exit(EXIT_FAILURE); -#endif + copyVramToHost(addr, (void *)iov.addr, len); } else if ((xferBenchConfig::op_type == XFERBENCH_OP_WRITE && xferBenchConfig::target_seg_type == XFERBENCH_SEG_TYPE_DRAM) || (xferBenchConfig::op_type == XFERBENCH_OP_READ && diff --git a/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp b/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp index 23a36711..9fcf8f9d 100644 --- a/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp +++ b/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp @@ -48,13 +48,16 @@ } \ } while (0) +static nixl_mem_t +resolveVramSegment() { #if HAVE_CUDA -#define HANDLE_VRAM_SEGMENT(_seg_type) _seg_type = VRAM_SEG; + return VRAM_SEG; #else -#define HANDLE_VRAM_SEGMENT(_seg_type) \ - std::cerr << "VRAM segment type not supported without CUDA" << std::endl; \ + if (neuronCoreCount() > 0) return VRAM_SEG; + std::cerr << "VRAM not supported without CUDA or Neuron" << std::endl; std::exit(EXIT_FAILURE); #endif +} #define GET_SEG_TYPE(is_initiator) \ ({ \ @@ -64,7 +67,7 @@ if (0 == _seg_type_str.compare("DRAM")) { \ _seg_type = DRAM_SEG; \ } else if (0 == _seg_type_str.compare("VRAM")) { \ - HANDLE_VRAM_SEGMENT(_seg_type); \ + _seg_type = resolveVramSegment(); \ } else { \ std::cerr << "Invalid segment type: " << _seg_type_str << std::endl; \ exit(EXIT_FAILURE); \ @@ -396,6 +399,21 @@ xferBenchNixlWorker::initBasicDescDram(size_t buffer_size, int mem_dev_id) { return std::optional(std::in_place, (uintptr_t)addr, buffer_size, mem_dev_id); } +static std::optional +getVramDescNeuron(int devid, size_t buffer_size, uint8_t memset_value) { + void *addr; + CHECK_NEURON_ERROR(neuronMalloc(&addr, buffer_size, devid), "Failed to allocate nrt tensor"); + CHECK_NEURON_ERROR(neuronMemset(addr, memset_value, buffer_size), + "Failed to set device memory"); + + return std::optional(std::in_place, (uintptr_t)addr, buffer_size, devid); +} + +static void +cleanupVramNeuron(xferBenchIOV &iov) { + CHECK_NEURON_ERROR(neuronFree((void *)iov.addr), "Failed to free nrt tensor"); +} + #if HAVE_CUDA static std::optional getVramDescCuda(int devid, size_t buffer_size, uint8_t memset_value) { @@ -460,33 +478,42 @@ getVramDescCudaVmm(int devid, size_t buffer_size, uint8_t memset_value) { #endif /* HAVE_CUDA_FABRIC */ } -static std::optional -getVramDescNeuron(int devid, size_t buffer_size, uint8_t memset_value) { - void *addr; - CHECK_NEURON_ERROR(neuronMalloc(&addr, buffer_size, devid), "Failed to allocate nrt tensor"); - CHECK_NEURON_ERROR(neuronMemset(addr, memset_value, buffer_size), - "Failed to set device memory"); - - return std::optional(std::in_place, (uintptr_t)addr, buffer_size, devid); +static void +cleanupVramCuda(xferBenchIOV &iov) { + CHECK_CUDA_ERROR(cudaSetDevice(iov.devId), "Failed to set device"); + if (xferBenchConfig::enable_vmm) { + CHECK_CUDA_DRIVER_ERROR(cuMemUnmap(iov.addr, iov.len), "Failed to unmap memory"); + CHECK_CUDA_DRIVER_ERROR(cuMemRelease(iov.handle), "Failed to release memory"); + CHECK_CUDA_DRIVER_ERROR(cuMemAddressFree(iov.addr, iov.padded_size), + "Failed to free reserved address"); + } else { + CHECK_CUDA_ERROR(cudaFreeAsync((void *)iov.addr, 0), "Failed to deallocate CUDA buffer"); + CHECK_CUDA_ERROR(cudaStreamSynchronize(0), "Failed to synchronize stream 0"); + } } +#endif /* HAVE_CUDA */ + static std::optional getVramDesc(int devid, size_t buffer_size, bool isInit) { uint8_t memset_value = isInit ? XFERBENCH_INITIATOR_BUFFER_ELEMENT : XFERBENCH_TARGET_BUFFER_ELEMENT; - // Assume no CUDA cores exist if Neuron cores are found. - // There are no AWS instance types with both NVIDIA GPUs and Neuron accelerators. if (neuronCoreCount() > 0) { return getVramDescNeuron(devid, buffer_size, memset_value); } +#if HAVE_CUDA CHECK_CUDA_ERROR(cudaSetDevice(devid), "Failed to set device"); if (xferBenchConfig::enable_vmm) { return getVramDescCudaVmm(devid, buffer_size, memset_value); } else { return getVramDescCuda(devid, buffer_size, memset_value); } +#else + std::cerr << "VRAM not supported without CUDA or Neuron" << std::endl; + return std::nullopt; +#endif } std::optional @@ -505,7 +532,6 @@ xferBenchNixlWorker::initBasicDescVram(size_t buffer_size, int mem_dev_id) { return getVramDesc(mem_dev_id, buffer_size, isInitiator()); } -#endif /* HAVE_CUDA */ // Helper to open a single file with appropriate flags static std::optional @@ -650,37 +676,19 @@ xferBenchNixlWorker::cleanupBasicDescDram(xferBenchIOV &iov) { free((void *)iov.addr); } -#if HAVE_CUDA void xferBenchNixlWorker::cleanupBasicDescVram(xferBenchIOV &iov) { - // Assume no CUDA cores exist if Neuron cores are found. - // There are no AWS instance types with both NVIDIA GPUs and Neuron accelerators. if (neuronCoreCount() > 0) { - CHECK_NEURON_ERROR(neuronFree((void *)iov.addr), "Failed to free nrt tensor"); + cleanupVramNeuron(iov); return; } - CHECK_CUDA_ERROR(cudaSetDevice(iov.devId), "Failed to set device"); - if (xferBenchConfig::enable_vmm) { - CHECK_CUDA_DRIVER_ERROR(cuMemUnmap(iov.addr, iov.len), "Failed to unmap memory"); - CHECK_CUDA_DRIVER_ERROR(cuMemRelease(iov.handle), "Failed to release memory"); - CHECK_CUDA_DRIVER_ERROR(cuMemAddressFree(iov.addr, iov.padded_size), - "Failed to free reserved address"); - } else { - /* - * CUDA streams allow for concurrent execution of kernels and memory operations. However, - * memory management functions like cudaFree are implicitly synchronized with all streams to - * guarantee safety. This means cudaFree will wait for all kernels (in any stream) that - * might use the memory to finish before actually freeing it. - * If the application hangs on cudaFree due to kernels running in other streams, switching - * to cudaFreeAsync can allow the host to proceed without waiting for the entire device - * synchronization. - */ - CHECK_CUDA_ERROR(cudaFreeAsync((void *)iov.addr, 0), "Failed to deallocate CUDA buffer"); - CHECK_CUDA_ERROR(cudaStreamSynchronize(0), "Failed to synchronize stream 0"); - } +#if HAVE_CUDA + cleanupVramCuda(iov); +#else + std::cerr << "VRAM not supported without CUDA or Neuron" << std::endl; +#endif } -#endif /* HAVE_CUDA */ void xferBenchNixlWorker::cleanupBasicDescFile(xferBenchIOV &iov) { @@ -936,11 +944,9 @@ xferBenchNixlWorker::allocateMemory(int num_threads) { basic_desc = initBasicDescDram(buffer_size, mem_dev_id); break; } -#if HAVE_CUDA case VRAM_SEG: basic_desc = initBasicDescVram(buffer_size, i); break; -#endif default: std::cerr << "Unsupported mem type: " << seg_type << std::endl; exit(EXIT_FAILURE); @@ -994,11 +1000,9 @@ xferBenchNixlWorker::deallocateMemory(std::vector> &io case DRAM_SEG: cleanupBasicDescDram(iov); break; -#if HAVE_CUDA case VRAM_SEG: cleanupBasicDescVram(iov); break; -#endif default: std::cerr << "Unsupported mem type: " << seg_type << std::endl; exit(EXIT_FAILURE); diff --git a/benchmark/nixlbench/src/worker/nixl/nixl_worker.h b/benchmark/nixlbench/src/worker/nixl/nixl_worker.h index d293ff8c..a1878af6 100644 --- a/benchmark/nixlbench/src/worker/nixl/nixl_worker.h +++ b/benchmark/nixlbench/src/worker/nixl/nixl_worker.h @@ -73,11 +73,11 @@ class xferBenchNixlWorker: public xferBenchWorker { private: std::optional initBasicDescDram(size_t buffer_size, int mem_dev_id); - void cleanupBasicDescDram(xferBenchIOV &basic_desc); -#if HAVE_CUDA + void + cleanupBasicDescDram(xferBenchIOV &basic_desc); std::optional initBasicDescVram(size_t buffer_size, int mem_dev_id); - void cleanupBasicDescVram(xferBenchIOV &basic_desc); -#endif + void + cleanupBasicDescVram(xferBenchIOV &basic_desc); std::optional initBasicDescFile(size_t buffer_size, xferFileState &fstate, int mem_dev_id); void cleanupBasicDescFile(xferBenchIOV &basic_desc); diff --git a/src/plugins/libfabric/libfabric_backend.cpp b/src/plugins/libfabric/libfabric_backend.cpp index a66dac67..cdd0cbc1 100644 --- a/src/plugins/libfabric/libfabric_backend.cpp +++ b/src/plugins/libfabric/libfabric_backend.cpp @@ -676,10 +676,14 @@ nixlLibfabricEngine::getSupportedMems() const { if (runtime_ == FI_HMEM_CUDA) { NIXL_DEBUG << "CUDA runtime detected, adding VRAM support"; mems.push_back(VRAM_SEG); + } else +#endif + if (runtime_ == FI_HMEM_NEURON) { + NIXL_DEBUG << "Neuron runtime detected, adding VRAM support"; + mems.push_back(VRAM_SEG); } else { - NIXL_DEBUG << "Non-CUDA runtime, skipping VRAM support"; + NIXL_DEBUG << "No accelerator runtime, skipping VRAM support"; } -#endif return mems; } From 26ee1691398eb43f8f7eb431fd4d79ebcc4bafb3 Mon Sep 17 00:00:00 2001 From: Ben Walker Date: Tue, 14 Apr 2026 00:43:21 -0700 Subject: [PATCH 26/59] nixl: fix test build when UCX plugin is not enabled (#1491) The gtest subdir was gated on ucx_dep.found() which checks for the system UCX library, not whether the UCX plugin is actually being built. On machines with UCX installed, this enters the gtest subdir even when UCX is not in enable_plugins, hitting an undefined ucx_backend_lib variable in the mock backend build. The mock backend has no UCX source dependency -- the link_with was unnecessary. Remove it and gate the gtest subdir on the UCX plugin being enabled rather than the system library being present. Signed-off-by: Ben Walker Co-authored-by: Adit Ranadive --- test/gtest/mocks/meson.build | 3 +-- test/meson.build | 4 ++-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/test/gtest/mocks/meson.build b/test/gtest/mocks/meson.build index 03cbc160..6281960d 100644 --- a/test/gtest/mocks/meson.build +++ b/test/gtest/mocks/meson.build @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -19,7 +19,6 @@ mock_backend_sources = ['gmock_engine.cpp', 'mock_backend_plugin.cpp', 'mock_bac mock_backend_plugin = shared_library('MOCK_BACKEND', mock_backend_sources, dependencies: [nixl_infra, nixl_common_dep, gmock_dep], include_directories: [nixl_inc_dirs, utils_inc_dirs, gtest_inc_dirs], - link_with : [ucx_backend_lib], name_prefix: 'libplugin_', install: true, install_dir: plugin_install_dir) diff --git a/test/meson.build b/test/meson.build index c0b0b5d5..09f21a56 100644 --- a/test/meson.build +++ b/test/meson.build @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -16,6 +16,6 @@ subdir('nixl') subdir('unit') -if ucx_dep.found() +if enabled_plugins.get('UCX') and ucx_dep.found() subdir('gtest') endif From 5115f8c8baf07e5c98d3b9cbc92e3563ffc8411d Mon Sep 17 00:00:00 2001 From: Guy Ealey Morag Date: Tue, 14 Apr 2026 10:02:21 +0200 Subject: [PATCH 27/59] Add batch insertion for sorted descriptor lists (#1479) --- src/core/nixl_agent.cpp | 2 +- src/infra/mem_section.h | 28 ++- src/infra/nixl_descriptors.cpp | 109 +++++++-- src/infra/nixl_memory_section.cpp | 59 ++--- test/gtest/unit/descriptors/meson.build | 22 ++ test/gtest/unit/descriptors/sec_desc_list.cpp | 215 ++++++++++++++++++ test/gtest/unit/meson.build | 3 + 7 files changed, 386 insertions(+), 52 deletions(-) create mode 100644 test/gtest/unit/descriptors/meson.build create mode 100644 test/gtest/unit/descriptors/sec_desc_list.cpp diff --git a/src/core/nixl_agent.cpp b/src/core/nixl_agent.cpp index 899583db..94ceed46 100644 --- a/src/core/nixl_agent.cpp +++ b/src/core/nixl_agent.cpp @@ -454,7 +454,7 @@ nixlAgent::registerMem(const nixl_reg_dlist_t &descs, const auto [it, inserted] = data->remoteSections_.try_emplace(data->name_, data->name_); - ret = it->second.loadLocalData(sec_descs, backend); + ret = it->second.loadLocalData(std::move(sec_descs), backend); if (ret == NIXL_SUCCESS) { count++; } else { diff --git a/src/infra/mem_section.h b/src/infra/mem_section.h index 8afc3439..28b183a6 100644 --- a/src/infra/mem_section.h +++ b/src/infra/mem_section.h @@ -59,7 +59,7 @@ class nixlSectionDesc : public nixlMetaDesc { } inline friend bool operator==(const nixlSectionDesc &lhs, const nixlSectionDesc &rhs) { - return (static_cast(lhs) == static_cast(rhs)); + return static_cast(lhs) == static_cast(rhs); } inline void print(const std::string &suffix) const { @@ -69,6 +69,8 @@ class nixlSectionDesc : public nixlMetaDesc { class nixlSecDescList : public nixlDescList { public: + enum class order : bool { UNSORTED, SORTED }; + explicit nixlSecDescList(const nixl_mem_t &type) : nixlDescList(type, 0) {} using nixlDescList::operator[]; // bring in const overload @@ -76,12 +78,19 @@ class nixlSecDescList : public nixlDescList { void addDesc(const nixlSectionDesc &desc) override; - bool - verifySorted() const; + void + addDesc(nixlSectionDesc &&desc); - nixlSectionDesc & + void + addDescs(std::vector batch, order ord = order::UNSORTED); + + void + addDescs(nixlSecDescList &&other); + + // Shadow the parent's non-const operator[] to return a const ref, + // this prevents mutation of descriptor fields after insertion + const nixlSectionDesc & operator[](size_t index) { - assert(verifySorted()); return descs[index]; } @@ -98,6 +107,13 @@ class nixlSecDescList : public nixlDescList { nixlSecDescList(const nixlSecDescList &) = default; nixlSecDescList & operator=(const nixlSecDescList &) = default; + nixlSecDescList(nixlSecDescList &&) = default; + nixlSecDescList & + operator=(nixlSecDescList &&) = default; + +private: + void + addSortedDescs(std::vector batch); }; using nixl_sec_dlist_t = nixlSecDescList; @@ -168,7 +184,7 @@ class nixlRemoteSection : public nixlMemSection { // When adding self as a remote agent for local operations nixl_status_t - loadLocalData(const nixlSecDescList &mem_elms, nixlBackendEngine *backend); + loadLocalData(nixlSecDescList mem_elms, nixlBackendEngine *backend); ~nixlRemoteSection(); }; diff --git a/src/infra/nixl_descriptors.cpp b/src/infra/nixl_descriptors.cpp index a24fa51a..1a16f059 100644 --- a/src/infra/nixl_descriptors.cpp +++ b/src/infra/nixl_descriptors.cpp @@ -15,10 +15,11 @@ * limitations under the License. */ #include -#include +#include +#include #include #include -#include "nixl.h" + #include "nixl_descriptors.h" #include "mem_section.h" #include "backend/backend_aux.h" @@ -353,30 +354,102 @@ void nixlSecDescList::addDesc(const nixlSectionDesc &desc) { auto &vec = this->descs; auto itr = std::upper_bound(vec.begin(), vec.end(), desc); - if (itr == vec.end()) - vec.push_back(desc); - else - vec.insert(itr, desc); + vec.insert(itr, desc); } -bool -nixlSecDescList::verifySorted() const { - const auto &vec = this->descs; - int size = (int)vec.size(); - if (size <= 1) return (size == 1); - for (int i = 0; i < size - 1; ++i) { - if (vec[i + 1] < vec[i]) return false; +void +nixlSecDescList::addDesc(nixlSectionDesc &&desc) { + auto &vec = this->descs; + auto itr = std::upper_bound(vec.begin(), vec.end(), desc); + vec.insert(itr, std::move(desc)); +} + +namespace { +void +appendAll(std::vector &dst, std::vector &src) { + dst.reserve(dst.size() + src.size()); + for (auto &d : src) { + dst.emplace_back(std::move(d)); } - return true; } +} // namespace + +void +nixlSecDescList::addSortedDescs(std::vector batch) { + auto &vec = this->descs; + if (vec.empty()) { + vec = std::move(batch); + return; + } + + // Check if the batch comes after the existing elements + if (!(batch.front() < vec.back())) { + appendAll(vec, batch); + return; + } + + // Check if the batch comes before the existing elements + if (!(vec.front() < batch.back())) { + appendAll(batch, vec); + vec = std::move(batch); + return; + } + + // Merge the batch into the existing vector in-place in reverse order to reduce copies + const size_t old_size = vec.size(); + const size_t batch_size = batch.size(); + const size_t new_size = old_size + batch_size; + + vec.resize(new_size); + + auto dst = vec.rbegin(); + auto a = std::make_reverse_iterator(vec.begin() + old_size); + auto a_end = vec.rend(); + auto b = batch.rbegin(); + auto b_end = batch.rend(); + + while (a != a_end && b != b_end) { + auto &src = (*b < *a) ? a : b; + *dst++ = std::move(*src++); + } + while (b != b_end) { + *dst++ = std::move(*b++); + } +} + +void +nixlSecDescList::addDescs(std::vector batch, order ord) { + if (batch.empty()) return; + + if (batch.size() == 1) { + // It's more efficient to insert a single element directly + addDesc(std::move(batch[0])); + return; + } + + if (ord == order::SORTED) { + NIXL_ASSERT(std::is_sorted(batch.begin(), batch.end())); + } else { + std::sort(batch.begin(), batch.end()); + } + + addSortedDescs(std::move(batch)); +} + +void +nixlSecDescList::addDescs(nixlSecDescList &&other) { + NIXL_ASSERT(type == other.type) << "Memory type mismatch: " << static_cast(type) + << " != " << static_cast(other.type); + addDescs(std::move(other.descs), order::SORTED); +} + int nixlSecDescList::getIndex(const nixlBasicDesc &query) const { auto itr = std::lower_bound(this->descs.begin(), this->descs.end(), query); - if (itr == this->descs.end()) return NIXL_ERR_NOT_FOUND; - if (static_cast(*itr) == query) - return static_cast(itr - this->descs.begin()); - return NIXL_ERR_NOT_FOUND; + if (itr == this->descs.end() || static_cast(*itr) != query) + return NIXL_ERR_NOT_FOUND; + return static_cast(itr - this->descs.begin()); } int diff --git a/src/infra/nixl_memory_section.cpp b/src/infra/nixl_memory_section.cpp index 12f479ec..4987fd44 100644 --- a/src/infra/nixl_memory_section.cpp +++ b/src/infra/nixl_memory_section.cpp @@ -142,17 +142,23 @@ nixlLocalSection::addDescList(const nixl_reg_dlist_t &mem_elms, nixlSecDescList &target = emplace(nixl_mem, backend); - // Add entries to the target list nixlSectionDesc local_sec, self_sec; nixlBasicDesc *lp = &local_sec; nixlBasicDesc *rp = &self_sec; nixl_status_t ret = NIXL_SUCCESS; - int i; - for (i = 0; i < mem_elms.descCount(); ++i) { + // Accumulate entries into batches, then merge on success + std::vector local_batch; + std::vector self_batch; + local_batch.reserve(mem_elms.descCount()); + if (backend->supportsLocal()) { + self_batch.reserve(mem_elms.descCount()); + } + + for (const auto &mem : mem_elms) { // TODO: For now trusting the user, but there can be a more checks mode // where we find overlaps and split the memories or warn the user - ret = backend->registerMem(mem_elms[i], nixl_mem, local_sec.metadataP); + ret = backend->registerMem(mem, nixl_mem, local_sec.metadataP); if (ret != NIXL_SUCCESS) break; @@ -175,34 +181,32 @@ nixlLocalSection::addDescList(const nixl_reg_dlist_t &mem_elms, } } - *lp = mem_elms[i]; // Copy the basic desc part + *lp = mem; // Copy the basic desc part if (((nixl_mem == BLK_SEG) || (nixl_mem == OBJ_SEG) || (nixl_mem == FILE_SEG)) && (lp->len==0)) lp->len = SIZE_MAX; // File has no range limit - target.addDesc(local_sec); + local_batch.push_back(local_sec); if (backend->supportsLocal()) { *rp = *lp; - remote_self.addDesc(self_sec); + self_batch.push_back(self_sec); } } - // Abort in case of error - if (ret != NIXL_SUCCESS) { - for (int j = 0; j < i; ++j) { - int index = target.getIndex(mem_elms[j]); - + if (ret == NIXL_SUCCESS) { + target.addDescs(std::move(local_batch)); + if (backend->supportsLocal()) { + remote_self.addDescs(std::move(self_batch)); + } + } else { + for (size_t j = 0; j < local_batch.size(); ++j) { if (backend->supportsLocal()) { - int self_index = remote_self.getIndex(mem_elms[j]); - // Should never be negative, as we just added it in previous loop - if (self_index >= 0 && remote_self[self_index].metadataP != target[index].metadataP) - backend->unloadMD(remote_self[self_index].metadataP); + if (self_batch[j].metadataP != local_batch[j].metadataP) + backend->unloadMD(self_batch[j].metadataP); } - backend->deregisterMem(target[index].metadataP); - target.remDesc(index); + backend->deregisterMem(local_batch[j].metadataP); } - remote_self.clear(); } return ret; } @@ -307,18 +311,21 @@ nixl_status_t nixlLocalSection::serializePartial(nixlSerDes* serializer, } const nixlSecDescList &base = it->second; - nixlSecDescList resp(nixl_mem); + std::vector descs; + descs.reserve(mem_elms.descCount()); for (const auto &desc : mem_elms) { - int index = base.getIndex(desc); + const int index = base.getIndex(desc); if (index < 0) { ret = NIXL_ERR_NOT_FOUND; break; } - resp.addDesc(base[index]); + descs.push_back(base[index]); } if (ret != NIXL_SUCCESS) { break; } + nixlSecDescList resp(nixl_mem); + resp.addDescs(std::move(descs)); mem_elms_to_serialize.try_emplace(sec_key, std::move(resp)); } @@ -418,8 +425,7 @@ nixlRemoteSection::loadRemoteData(nixlSerDes *deserializer, backend_map_t &backe } nixl_status_t -nixlRemoteSection::loadLocalData(const nixlSecDescList &mem_elms, nixlBackendEngine *backend) { - +nixlRemoteSection::loadLocalData(nixlSecDescList mem_elms, nixlBackendEngine *backend) { if (mem_elms.isEmpty()) { // Shouldn't happen return NIXL_ERR_UNKNOWN; } @@ -428,9 +434,8 @@ nixlRemoteSection::loadLocalData(const nixlSecDescList &mem_elms, nixlBackendEng nixlSecDescList &target = emplace(nixl_mem, backend); - for (auto &elm : mem_elms) { - target.addDesc(elm); - } + target.addDescs(std::move(mem_elms)); + return NIXL_SUCCESS; } diff --git a/test/gtest/unit/descriptors/meson.build b/test/gtest/unit/descriptors/meson.build new file mode 100644 index 00000000..823a41ce --- /dev/null +++ b/test/gtest/unit/descriptors/meson.build @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +descriptors_unit_test_dep = declare_dependency( + sources: [ + 'sec_desc_list.cpp', + ], + include_directories: [nixl_inc_dirs, gtest_inc_dirs], + dependencies: [nixl_common_dep], +) diff --git a/test/gtest/unit/descriptors/sec_desc_list.cpp b/test/gtest/unit/descriptors/sec_desc_list.cpp new file mode 100644 index 00000000..2792cc63 --- /dev/null +++ b/test/gtest/unit/descriptors/sec_desc_list.cpp @@ -0,0 +1,215 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 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. + */ + +#include +#include +#include +#include + +#include "mem_section.h" + +namespace descriptors { + +class secDescListTest : public ::testing::Test { +protected: + static constexpr uint64_t defaultDevId = 1; + static constexpr size_t defaultLen = 64; + + static nixlSectionDesc + makeDesc(uintptr_t addr, uint64_t dev_id = defaultDevId, size_t len = defaultLen) { + return nixlSectionDesc(addr, len, dev_id); + } + + nixlSecDescList + makeList() { + return nixlSecDescList(DRAM_SEG); + } + + nixlSecDescList + makeList(std::initializer_list addrs) { + auto list = makeList(); + std::vector batch; + batch.reserve(addrs.size()); + for (auto a : addrs) { + batch.push_back(makeDesc(a)); + } + list.addDescs(std::move(batch)); + assertSorted(list); + return list; + } + + static void + assertSorted(const nixlSecDescList &list) { + ASSERT_TRUE(std::is_sorted(list.begin(), list.end())); + } + + static void + expectAddrs(const nixlSecDescList &list, const std::vector &expected) { + ASSERT_EQ(list.descCount(), static_cast(expected.size())); + ASSERT_TRUE(std::is_sorted(list.begin(), list.end())); + + for (size_t i = 0; i < expected.size(); ++i) { + EXPECT_EQ(list[i].addr, expected[i]) << "mismatch at index " << i; + EXPECT_EQ(list[i].devId, defaultDevId) << "devId mismatch at index " << i; + EXPECT_EQ(list[i].len, defaultLen) << "len mismatch at index " << i; + } + } + + static void + expectAddrsDevIds(const nixlSecDescList &list, + const std::vector> &expected) { + ASSERT_EQ(list.descCount(), static_cast(expected.size())); + ASSERT_TRUE(std::is_sorted(list.begin(), list.end())); + + for (size_t i = 0; i < expected.size(); ++i) { + EXPECT_EQ(list[i].addr, expected[i].first) << "addr mismatch at index " << i; + EXPECT_EQ(list[i].devId, expected[i].second) << "devId mismatch at index " << i; + EXPECT_EQ(list[i].len, defaultLen) << "len mismatch at index " << i; + } + } +}; + +TEST_F(secDescListTest, EmptyBatchOnEmptyList) { + auto list = makeList(); + list.addDescs({}); + ASSERT_TRUE(list.isEmpty()); +} + +TEST_F(secDescListTest, EmptyBatchOnNonEmptyList) { + auto list = makeList({100, 200}); + list.addDescs({}); + expectAddrs(list, {100, 200}); +} + +TEST_F(secDescListTest, SingleElementBatch) { + auto list = makeList({20, 40}); + + list.addDescs({makeDesc(10)}); + expectAddrs(list, {10, 20, 40}); + + list.addDescs({makeDesc(30)}); + expectAddrs(list, {10, 20, 30, 40}); + + list.addDescs({makeDesc(50)}); + expectAddrs(list, {10, 20, 30, 40, 50}); +} + +TEST_F(secDescListTest, AllAfterAppend) { + auto list = makeList({10, 20}); + list.addDescs({makeDesc(30), makeDesc(40)}); + expectAddrs(list, {10, 20, 30, 40}); +} + +TEST_F(secDescListTest, AllBeforePrepend) { + auto list = makeList({30, 40}); + list.addDescs({makeDesc(10), makeDesc(20)}); + expectAddrs(list, {10, 20, 30, 40}); +} + +TEST_F(secDescListTest, InterleavedMerge) { + auto list = makeList({10, 30, 50}); + list.addDescs({makeDesc(20), makeDesc(40)}); + expectAddrs(list, {10, 20, 30, 40, 50}); +} + +TEST_F(secDescListTest, UnsortedInput) { + auto list = makeList({40, 10, 30, 20}); + expectAddrs(list, {10, 20, 30, 40}); +} + +TEST_F(secDescListTest, SortedInput) { + auto list = makeList(); + list.addDescs({makeDesc(10), makeDesc(20), makeDesc(30)}, nixlSecDescList::order::SORTED); + expectAddrs(list, {10, 20, 30}); +} + +TEST_F(secDescListTest, SortedHintWithUnsortedInputDies) { + auto list = makeList(); + EXPECT_DEBUG_DEATH( + list.addDescs({makeDesc(30), makeDesc(10), makeDesc(20)}, nixlSecDescList::order::SORTED), + ""); +} + +TEST_F(secDescListTest, OtherListOverload) { + auto list = makeList({20, 40}); + auto other = makeList({10, 30, 50}); + + list.addDescs(std::move(other)); + expectAddrs(list, {10, 20, 30, 40, 50}); +} + +TEST_F(secDescListTest, DuplicateDescriptors) { + auto list = makeList({10, 20}); + list.addDescs({makeDesc(10), makeDesc(20)}); + expectAddrs(list, {10, 10, 20, 20}); +} + +TEST_F(secDescListTest, MultipleDevIds) { + auto list = makeList(); + + // empty batch + list.addDescs({}); + expectAddrsDevIds(list, {}); + + // single element + list.addDesc(makeDesc(50, 1)); + expectAddrsDevIds(list, {{50, 1}}); + + // all after existing + list.addDescs({makeDesc(100, 1), makeDesc(10, 2)}); + expectAddrsDevIds(list, {{50, 1}, {100, 1}, {10, 2}}); + + // all before existing + list.addDescs({makeDesc(10, 1), makeDesc(50, 0)}); + expectAddrsDevIds(list, {{50, 0}, {10, 1}, {50, 1}, {100, 1}, {10, 2}}); + + // interleaved + list.addDescs({makeDesc(10, 0), makeDesc(75, 1), makeDesc(200, 1), makeDesc(20, 2)}); + expectAddrsDevIds( + list, {{10, 0}, {50, 0}, {10, 1}, {50, 1}, {75, 1}, {100, 1}, {200, 1}, {10, 2}, {20, 2}}); +} + +TEST_F(secDescListTest, AddRandomBatches) { + constexpr size_t total_elements = 4096; + + auto list = makeList(); + + std::mt19937 rng(42); + std::uniform_int_distribution addr_dist(1, 1000000); + std::uniform_int_distribution dev_dist(0, 8); + std::uniform_int_distribution len_dist(1, 64); + std::uniform_int_distribution batch_dist(1, 256); + + size_t num_added = 0; + while (num_added < total_elements) { + size_t batch_size = std::min(batch_dist(rng), total_elements - num_added); + + std::vector batch; + batch.reserve(batch_size); + for (size_t j = 0; j < batch_size; ++j) { + batch.push_back(makeDesc(addr_dist(rng), dev_dist(rng), len_dist(rng))); + } + + list.addDescs(std::move(batch)); + + num_added += batch_size; + ASSERT_EQ(list.descCount(), num_added); + assertSorted(list); + } +} + +} // namespace descriptors diff --git a/test/gtest/unit/meson.build b/test/gtest/unit/meson.build index 64f8e69d..4f9aee42 100644 --- a/test/gtest/unit/meson.build +++ b/test/gtest/unit/meson.build @@ -26,6 +26,9 @@ endif subdir('agent') unit_test_deps += [agent_unit_test_dep] +subdir('descriptors') +unit_test_deps += [descriptors_unit_test_dep] + cpp = meson.get_compiler('cpp') azure_storage_blobs = cpp.find_library('azure-storage-blobs', required: false) if enabled_plugins.get('AZURE_BLOB') and azure_storage_blobs.found() From a7bb070fd773a263908aaf2280cd4d8cca60a7c1 Mon Sep 17 00:00:00 2001 From: ovidiusm Date: Tue, 14 Apr 2026 13:34:19 +0200 Subject: [PATCH 28/59] Refactor manylinux dockerfile to allow faster UCX & NIXL bisection (#1523) * Build UCX and NIXL at the end to allow easy bisection with cached builds Signed-off-by: Ovidiu Mara * Remove pip install step Signed-off-by: Ovidiu Mara * Cleanup Signed-off-by: Ovidiu Mara * Add git log and clean whitespace Signed-off-by: Ovidiu Mara * Update help string Signed-off-by: Ovidiu Mara * Quote UCX_REF Signed-off-by: Ovidiu Mara --------- Signed-off-by: Ovidiu Mara Co-authored-by: Adit Ranadive --- contrib/Dockerfile.manylinux | 49 ++++++++++++++++++------------------ contrib/build-container.sh | 14 +++++++---- 2 files changed, 33 insertions(+), 30 deletions(-) diff --git a/contrib/Dockerfile.manylinux b/contrib/Dockerfile.manylinux index b85ef634..128b7d57 100644 --- a/contrib/Dockerfile.manylinux +++ b/contrib/Dockerfile.manylinux @@ -20,7 +20,6 @@ FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} ARG DEFAULT_PYTHON_VERSION="3.14" ARG ARCH="x86_64" -ARG UCX_REF="v1.21.x" ARG LIBFABRIC_VERSION="v1.21.0" ARG HWLOC_VERSION="2.12.2" ARG BUILD_TYPE="release" @@ -258,28 +257,6 @@ RUN cd /workspace && \ rpm -Uvh gdrcopy-*.el8.$ARCH.rpm && \ rpm -Uvh gdrcopy-devel-*.el8.noarch.rpm -RUN cd /usr/local/src && \ - git clone https://github.com/openucx/ucx.git && \ - cd ucx && \ - git checkout $UCX_REF && \ - ./autogen.sh && \ - ./contrib/configure-release-mt \ - --enable-shared \ - --disable-static \ - --disable-doxygen-doc \ - --enable-experimental-api \ - --enable-optimizations \ - --enable-cma \ - --enable-devel-headers \ - --with-cuda=/usr/local/cuda \ - --with-verbs \ - --with-dm \ - --without-gdrcopy \ - --with-efa && \ - make -j && \ - make -j install-strip && \ - ldconfig - # Build libfabric from source RUN wget --tries=3 --waitretry=5 --timeout=30 --read-timeout=60 \ "https://github.com/ofiwg/libfabric/releases/download/${LIBFABRIC_VERSION}/libfabric-${LIBFABRIC_VERSION#v}.tar.bz2" -O libfabric.tar.bz2 && \ @@ -321,6 +298,30 @@ RUN export UV_INDEX="https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | c # Upgrade setuptools to latest version for compatibility with PEP 639 (license format) RUN uv pip install --upgrade 'setuptools>=80.9.0' +ARG UCX_REF="v1.21.x" +RUN cd /usr/local/src && \ + git clone https://github.com/openucx/ucx.git && \ + cd ucx && \ + git checkout "${UCX_REF}" && \ + git log -1 && \ + ./autogen.sh && \ + ./contrib/configure-release-mt \ + --enable-shared \ + --disable-static \ + --disable-doxygen-doc \ + --enable-experimental-api \ + --enable-optimizations \ + --enable-cma \ + --enable-devel-headers \ + --with-cuda=/usr/local/cuda \ + --with-verbs \ + --with-dm \ + --without-gdrcopy \ + --with-efa && \ + make -j && \ + make -j install-strip && \ + ldconfig + COPY . /workspace/nixl WORKDIR /workspace/nixl @@ -372,5 +373,3 @@ RUN IFS=',' read -ra PYTHON_VERSIONS <<< "$WHL_PYTHON_VERSIONS" && \ RUN if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \ cp build/src/bindings/python/nixl-meta/nixl*.whl dist/; \ fi - -RUN uv pip install dist/nixl*cp${DEFAULT_PYTHON_VERSION//./}*.whl build/src/bindings/python/nixl-meta/nixl*.whl diff --git a/contrib/build-container.sh b/contrib/build-container.sh index 2e7cd14e..02b7216c 100755 --- a/contrib/build-container.sh +++ b/contrib/build-container.sh @@ -119,9 +119,13 @@ get_options() { missing_requirement $1 fi ;; - --ucx-upstream) - # Master branch (v1.20) also containing EFA SRD support - UCX_REF=9d2b88a1f67faf9876f267658bd077b379b8bb76 + --ucx-ref) + if [ "$2" ]; then + UCX_REF=$2 + shift + else + missing_requirement $1 + fi ;; --build-nixl-ep) BUILD_NIXL_EP=true @@ -192,8 +196,8 @@ show_help() { echo " [--build-type [debug|release] to select build type (default: release)]" echo " [--tag tag for image]" echo " [--python-versions python versions to build for, comma separated]" - echo " [--ucx-upstream use ucx master branch]" - echo " [--build-nixl-ep build NIXL with NIXL EP support (uses latest UCX master)]" + echo " [--ucx-ref ucx git reference (branch, tag, or sha)]" + echo " [--build-nixl-ep build NIXL with NIXL EP support (requires UCX >= 1.21)]" echo " [--arch [x86_64|aarch64] to select target architecture]" echo " [--dockerfile path to a dockerfile to use]" exit 0 From d6456361eb48aa7b111aa8741eb07e5709d2cdbf Mon Sep 17 00:00:00 2001 From: Ben Walker Date: Tue, 14 Apr 2026 10:11:46 -0700 Subject: [PATCH 29/59] test: fix maybe-uninitialized warning in metadata_exchange test (#1473) Initialize status = NIXL_ERR_NOT_FOUND before the backend creation loop to silence g++ -Werror=maybe-uninitialized. Signed-off-by: Ben Walker --- test/gtest/metadata_exchange.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/gtest/metadata_exchange.cpp b/test/gtest/metadata_exchange.cpp index 7986a4fc..c3811011 100644 --- a/test/gtest/metadata_exchange.cpp +++ b/test/gtest/metadata_exchange.cpp @@ -550,8 +550,8 @@ TEST_F(MetadataExchangeTestFixture, LocalNonLocalMDExchange) { auto &src = agents_[0]; auto &dst = agents_[1]; + nixl_status_t status = NIXL_ERR_NOT_FOUND; nixlBackendH *backend; - nixl_status_t status; std::string backend_name; for (const auto& name : std::set{"GDS", "POSIX"}) { const LogIgnoreGuard lig1("Error initializing GPU Direct Storage driver"); From 3702ccecd0a4ed68cb9f975cf30ef16c680b5ec2 Mon Sep 17 00:00:00 2001 From: oa-aws Date: Tue, 14 Apr 2026 20:54:55 +0300 Subject: [PATCH 30/59] libfabric: Fix warnings for NUMA-aware rail selection policy (support more instance types) (#1461) --- .../libfabric/libfabric_rail_manager.cpp | 80 ++++++-- src/utils/libfabric/libfabric_rail_manager.h | 5 +- src/utils/libfabric/libfabric_topology.cpp | 174 +++++++++++------- src/utils/libfabric/libfabric_topology.h | 12 ++ test/gtest/hw_warning_test.cpp | 6 - .../libfabric/libfabric_topology_test.cpp | 22 ++- .../utils/libfabric/topo/c5n.18xl-topo.xml | 69 +++++++ 7 files changed, 277 insertions(+), 91 deletions(-) create mode 100644 test/unit/utils/libfabric/topo/c5n.18xl-topo.xml diff --git a/src/utils/libfabric/libfabric_rail_manager.cpp b/src/utils/libfabric/libfabric_rail_manager.cpp index 7c86c340..0e2e78b8 100644 --- a/src/utils/libfabric/libfabric_rail_manager.cpp +++ b/src/utils/libfabric/libfabric_rail_manager.cpp @@ -253,13 +253,25 @@ nixlLibfabricRailManager::init(const nixl_b_params_t &custom_params) { // get bandwidth/rail limit from user or compute it, then select policy size_t max_bw = 0; size_t max_rails = 0; + size_t recommended_rails = 0; size_t nic_speed = topology->getAvgNicBandwidth(); - if (!getDramRailLimit(custom_params, max_bw, max_rails) || max_rails == 0) { + if (!getDramRailLimit(custom_params, max_bw, max_rails, recommended_rails) || max_rails == 0) { // had some error in deducing rail count, so just use default policy NIXL_WARN << "Using default (all) rail selection policy for DRAM memory type due to " "previous errors"; dram_rail_selection_policy_ = std::make_unique(); - } else if (max_rails < topology->getTotalNicCount()) { + } else if (topology->getAllDevices().size() == 1) { + // system has only 1 EFA, so use default policy (don't issue warning) + dram_rail_selection_policy_ = std::make_unique(); + } + // NOTE: from this point onward we regard only rails that are attached to a PCIe switch because + // these are the only ones that matter for avoiding PCIe congestion (the instance types we deal + // with have either all rails connected to PCIe switches, or have only one EFA that is either + // not attached to a NUMA node, as in C-series, or not attached to a PCIe switch, as in + // G-series), so we test rail count against topology->getTotalNicCount() (which represents the + // total number of "usable" rails, i.e. that are attached to a PCIe switch) and not against + // topology->getAllDevices().size() + else if (max_rails < topology->getTotalNicCount()) { // bandwidth does not exceed total machine capacity, so use NUMA-aware rail selection policy NIXL_TRACE << "Using NUMA-aware rail selection policy for DRAM memory type"; size_t numa_rail_count = topology->getNumaRailCount(); // NOTE: averaged if non-uniform @@ -268,7 +280,18 @@ nixlLibfabricRailManager::init(const nixl_b_params_t &custom_params) { NIXL_WARN << "User-provided configuration value for max_bw_per_dram_seg (" << max_bw << " Gbps) exceeds single NUMA node capacity of " << numa_speed << " Gbps, and will spill over to other NUMA nodes"; + } else if (max_rails > recommended_rails) { + // configured rail count does not spill over to other nodes, but still exceeds PCIe + // switch capacity, and is expected to cause congestion, so warn user + size_t recommended_bw = recommended_rails * nic_speed; + NIXL_WARN << "User-provided configuration value for max_bw_per_dram_seg (" << max_bw + << " Gbps), which results in " << max_rails + << " rails, exceeds the congestion-free recommendation of " << recommended_bw + << " Gbps (" << recommended_rails + << " rails per NUMA node), and is expected to cause PCIe congestion"; } + NIXL_TRACE << "Using " << max_rails + << " rails in NUMA-aware rail selection policy for DRAM_SEG"; dram_rail_selection_policy_ = std::make_unique(max_rails); } else { @@ -524,10 +547,29 @@ nixlLibfabricRailManager::prepareAndSubmitTransfer( bool nixlLibfabricRailManager::getDramRailLimit(const nixl_b_params_t &custom_params, size_t &max_bw, - size_t &max_rails) { - // first make sure there are any EFA devices - if (topology->getTotalNicCount() == 0) { + size_t &max_rails, + size_t &recommended_rails) { + // trivial case: no EFA devices + if (topology->getAllDevices().empty()) { + max_rails = 0; NIXL_WARN << "Could not find EFA devices, rail selection for DRAM memory type aborted"; + // this is unexpected, we except to find at least one EFA + return false; + } + + // trivial case: only one EFA device + if (topology->getAllDevices().size() == 1) { + max_rails = 1; + NIXL_INFO << "NUMA-aware rail selection disabled when system has only one EFA device"; + // this could happen (e.g. C-series instance types) + return true; + } + + // first make sure there are any usable EFA devices + if (topology->getTotalNicCount() == 0) { + NIXL_WARN + << "Could not find EFA devices connected to any PCIe switch, NUMA-aware rail selection " + "for DRAM memory type aborted"; return false; } @@ -545,7 +587,20 @@ nixlLibfabricRailManager::getDramRailLimit(const nixl_b_params_t &custom_params, return false; } - // now compute rail limit based on bandwidth limit + // compute recommended number of rails and corresponding bandwidth limit + size_t recommended_bw = topology->getAvgNumaNodeBandwidth(); + if (recommended_bw == 0) { + NIXL_WARN << "Could not deduce average bandwidth limit per NUMA node, NUMA-aware rail " + "selection for DRAM type aborted"; + return false; + } + // NOTE: when rail limit is computed, we divide switch capacity by NIC upstream link + recommended_rails = recommended_bw / nic_upstream_speed; + NIXL_TRACE << "Computed (average) NUMA node combined PCIe switch bandwidth limit is " + << recommended_bw << " Gbps"; + NIXL_TRACE << "Computed (average) rails per NUMA node is " << recommended_rails; + + // now compute rail limit based on configured or computed bandwidth limit max_rails = 0; max_bw = 0; @@ -559,17 +614,8 @@ nixlLibfabricRailManager::getDramRailLimit(const nixl_b_params_t &custom_params, // bandwidth limit could not be obtained from user (either user did not specify, or // configuration was malformed) so compute bandwidth limit from topology, and then deduce // rail count - max_bw = topology->getAvgNumaNodeBandwidth(); - if (max_bw == 0) { - NIXL_WARN << "Could not deduce average bandwidth limit per NUMA node, NUMA-aware rail " - "selection for DRAM type aborted"; - return false; - } - // NOTE: when rail limit is computed, we divide switch capacity by NIC upstream link - max_rails = max_bw / nic_upstream_speed; - NIXL_TRACE << "Computed (average) NUMA node combined PCIe switch bandwidth limit is " - << max_bw << " Gbps"; - NIXL_TRACE << "Computed (average) rails per NUMA node is " << max_rails; + max_rails = recommended_rails; + max_bw = recommended_bw; } else { // print warning if bandwidth limit provided by user is less than the speed of a single NIC if (max_bw < nic_speed) { diff --git a/src/utils/libfabric/libfabric_rail_manager.h b/src/utils/libfabric/libfabric_rail_manager.h index 8b44de85..1d4b424f 100644 --- a/src/utils/libfabric/libfabric_rail_manager.h +++ b/src/utils/libfabric/libfabric_rail_manager.h @@ -360,7 +360,10 @@ class nixlLibfabricRailManager { // get rail count limit for DRAM memory type, either computed or from user bool - getDramRailLimit(const nixl_b_params_t &custom_params, size_t &max_bw, size_t &max_rails); + getDramRailLimit(const nixl_b_params_t &custom_params, + size_t &max_bw, + size_t &max_rails, + size_t &recommended_rails); // Internal rail selection method std::vector diff --git a/src/utils/libfabric/libfabric_topology.cpp b/src/utils/libfabric/libfabric_topology.cpp index c2accffa..b6de3c6e 100644 --- a/src/utils/libfabric/libfabric_topology.cpp +++ b/src/utils/libfabric/libfabric_topology.cpp @@ -90,6 +90,10 @@ nixlLibfabricTopology::discoverTopology() { NIXL_ERROR << "Failed to discover hwloc topology"; return status; } + + // build nic info map regardless of accelerator to EFA mapping + buildNicInfoMap(); + // Build nVidia accelerator to EFA mapping based on PCIe topology if (num_nvidia_accel > 0) { status = buildAccelToEfaMapping(); @@ -200,7 +204,7 @@ nixlLibfabricTopology::getDeviceNumaNode(const std::string &efa_device) const { } else { device_numa_node = itr->second.numa_node_id; if (device_numa_node == INVALID_NUMA_NODE_ID) { - NIXL_WARN << "EFA device " << efa_device << " is not associated with a NUMA node"; + NIXL_INFO << "EFA device " << efa_device << " is not associated with a NUMA node"; } else { NIXL_DEBUG << "EFA device " << efa_device << " is on NUMA node " << device_numa_node; } @@ -550,64 +554,11 @@ nixlLibfabricTopology::buildTopologyAwareGrouping() { // Step 1: Build NIC info structures by correlating libfabric with hwloc std::vector discovered_nics; std::vector discovered_accel; - // Discover NICs by correlating libfabric devices with hwloc objects - for (const auto &pair : pcie_to_libfabric_map) { - const std::string &pcie_addr = pair.first; - const std::string &libfabric_name = pair.second; - - // Parse PCIe address - uint16_t domain_id; - uint8_t bus_id, device_id, function_id; - if (sscanf(pcie_addr.c_str(), - "%hx:%hhx:%hhx.%hhx", - &domain_id, - &bus_id, - &device_id, - &function_id) != 4) { - NIXL_WARN << "Failed to parse PCIe address: " << pcie_addr; - continue; - } - - // Find corresponding hwloc object - hwloc_obj_t hwloc_node = - hwloc_get_pcidev_by_busid(hwloc_topology, domain_id, bus_id, device_id, function_id); - - if (hwloc_node) { - NicInfo nic; - nic.libfabric_name = libfabric_name; - nic.hwloc_node = hwloc_node; - nic.line_speed = getPcieDevSpeed(pcie_addr); - // NOTE: upstream link speed is given in decimal GB/s (i.e. multiple of 10^9) as float, - // so we convert it to size_t decimal Gbps - nic.upstream_link_speed = (size_t)(hwloc_node->attr->pcidev.linkspeed * 8.0f); - nic.numa_node_id = getPcieDevNumaNodeId(hwloc_node, pcie_addr); - nic.domain_id = domain_id; - nic.bus_id = bus_id; - nic.device_id = device_id; - nic.function_id = function_id; - if (!getPcieDevParentSwitchData(hwloc_node, - pcie_addr, - nic.parent_switch_domain, - nic.parent_switch_bus_id, - nic.parent_switch_link_speed)) { - NIXL_TRACE << "Could not locate parent PCIe bridge/switch of NIC " - << libfabric_name; - nic.parent_switch_domain = UINT16_MAX; - nic.parent_switch_bus_id = UINT8_MAX; - } - nic_info_map.insert(NicInfoMap::value_type(libfabric_name, nic)); - NIXL_DEBUG << "EFA device " << libfabric_name << " mapped to NUMA node " - << nic.numa_node_id << " (PCIe address " << pcie_addr - << ", speed: " << nic.line_speed - << " Gbps, upstream link speed: " << nic.upstream_link_speed - << " Gbps, parent switch domain/bus-id: " << nic.parent_switch_domain << "/" - << nic.parent_switch_bus_id << ")"; - discovered_nics.push_back(nic); - NIXL_TRACE << "Correlated NIC: " << pcie_addr << " → " << libfabric_name; - } else { - NIXL_WARN << "Could not find hwloc object for PCIe address: " << pcie_addr; - } + // the nic info map is already populated, so just use its mapped values + for (const auto &entry : nic_info_map) { + discovered_nics.push_back(entry.second); } + // Step 2: Discover accelerators hwloc_obj_t pci_obj = nullptr; while ((pci_obj = hwloc_get_next_pcidev(hwloc_topology, pci_obj)) != nullptr) { @@ -671,12 +622,56 @@ nixlLibfabricTopology::buildTopologyAwareGrouping() { } } } - // step 5: compute the capacity limit of each NUMA node and some other topology metrics + return NIXL_SUCCESS; +} + +void +nixlLibfabricTopology::buildNicInfoMap() { + // Build NIC info structures by correlating libfabric with hwloc + // Discover NICs by correlating libfabric devices with hwloc objects + for (const auto &pair : pcie_to_libfabric_map) { + // parse PCIe address from string + const std::string &pcie_addr = pair.first; + const std::string &libfabric_name = pair.second; + + // Parse PCIe address + uint16_t domain_id; + uint8_t bus_id, device_id, function_id; + if (sscanf(pcie_addr.c_str(), + "%hx:%hhx:%hhx.%hhx", + &domain_id, + &bus_id, + &device_id, + &function_id) != 4) { + NIXL_WARN << "Failed to parse PCIe address: " << pcie_addr; + continue; + } + + // Find corresponding hwloc object + hwloc_obj_t hwloc_node = + hwloc_get_pcidev_by_busid(hwloc_topology, domain_id, bus_id, device_id, function_id); + + if (hwloc_node) { + NicInfo nic; + collectNicInfo(nic, + libfabric_name, + pcie_addr, + hwloc_node, + domain_id, + bus_id, + device_id, + function_id); + NIXL_TRACE << "Correlated NIC: " << pcie_addr << " → " << libfabric_name; + } else { + NIXL_WARN << "Could not find hwloc object for PCIe address: " << pcie_addr; + } + } + + // now compute the capacity limit of each NUMA node and some other topology metrics buildNumaSpeedMap(); calcAvgNumaNodeBandwidth(); calcAvgNicBandwidth(); calcAvgNicUpstreamBandwidth(); - return NIXL_SUCCESS; } nixl_status_t @@ -796,9 +791,9 @@ nixlLibfabricTopology::getPcieDevNumaNodeId(hwloc_obj_t obj, const std::string & return INVALID_NUMA_NODE_ID; } if (node_count == 0) { - // this is possible in some instance types (e.g. g5.48xl), so we issue only a warning - NIXL_WARN << "Failed to identify the NUMA node closest to NIC PCIe device at address " - << pcie_addr << ": no node found"; + // this is possible in some instance types (e.g. g5.48xl), so we issue info here + NIXL_DEBUG << "NIC PCIe device at address " << pcie_addr + << " is not associated with a NUMA node"; return INVALID_NUMA_NODE_ID; } if (node_count > 1) { @@ -874,6 +869,44 @@ nixlLibfabricTopology::getPcieDevParentSwitchData(hwloc_obj_t obj, return found; } +void +nixlLibfabricTopology::collectNicInfo(NicInfo &nic, + const std::string &name, + const std::string &pcie_addr, + hwloc_obj_t hwloc_node, + uint16_t domain_id, + uint8_t bus_id, + uint8_t device_id, + uint8_t function_id) { + + nic.libfabric_name = name; + nic.hwloc_node = hwloc_node; + nic.line_speed = getPcieDevSpeed(pcie_addr); + // NOTE: upstream link speed is given in decimal GB/s (i.e. multiple of 10^9) as float, + // so we convert it to size_t decimal Gbps + nic.upstream_link_speed = (size_t)(hwloc_node->attr->pcidev.linkspeed * 8.0f); + nic.numa_node_id = getPcieDevNumaNodeId(hwloc_node, pcie_addr); + nic.domain_id = domain_id; + nic.bus_id = bus_id; + nic.device_id = device_id; + nic.function_id = function_id; + if (!getPcieDevParentSwitchData(hwloc_node, + pcie_addr, + nic.parent_switch_domain, + nic.parent_switch_bus_id, + nic.parent_switch_link_speed)) { + NIXL_TRACE << "Could not locate parent PCIe bridge/switch of NIC " << name; + nic.parent_switch_domain = UINT16_MAX; + nic.parent_switch_bus_id = UINT8_MAX; + } + nic_info_map.insert(NicInfoMap::value_type(name, nic)); + NIXL_DEBUG << "EFA device " << name << " mapped to NUMA node " << nic.numa_node_id + << " (PCIe address " << pcie_addr << ", speed: " << nic.line_speed + << " Gbps, upstream link speed: " << nic.upstream_link_speed + << " Gbps, parent switch domain/bus-id: " << nic.parent_switch_domain << "/" + << nic.parent_switch_bus_id << ")"; +} + void nixlLibfabricTopology::buildNumaSpeedMap() { // build the NUMA bandwidth limit map @@ -954,13 +987,14 @@ nixlLibfabricTopology::buildNumaSpeedMap() { // (otherwise we are probably completely off here, or there is HW issue) if (pairib.first->second.second != numa_node_id) { NIXL_WARN << "Invalid NUMA node id " << numa_node_id << " for link bus " - << topmost_bus_id << ", expecting instead " - << pairib.first->second.second + << std::hex << static_cast(topmost_bus_id) << std::dec + << ", expecting instead " << pairib.first->second.second << ", entry will be ignored (sub-optimal DRAM performance may be " "observed)"; } } else { - NIXL_DEBUG << "Recording link bus " << topmost_bus_id << " speed " + NIXL_DEBUG << "Recording link bus " << std::hex + << static_cast(topmost_bus_id) << std::dec << " speed " << topmost_speed << " GB/s to capacity of NUMA node " << numa_node_id; } @@ -968,6 +1002,14 @@ nixlLibfabricTopology::buildNumaSpeedMap() { } } + // if nothing usable was found then stop + // (usable means a NIC with parent switch having non-zero link speed) + if (link_speed_map.empty()) { + NIXL_TRACE + << "Could not find any NIC connected to a PCIe switch, NIC speed map building aborted"; + return; + } + // now we process the link speed map and accumulate per NUMA node numa_speed_map.resize(max_node_id + 1, 0); for (const auto &entry : link_speed_map) { diff --git a/src/utils/libfabric/libfabric_topology.h b/src/utils/libfabric/libfabric_topology.h index 8f3b6381..cd846dea 100644 --- a/src/utils/libfabric/libfabric_topology.h +++ b/src/utils/libfabric/libfabric_topology.h @@ -84,6 +84,8 @@ class nixlLibfabricTopology { nixl_status_t buildAccelToEfaMapping(); void + buildNicInfoMap(); + void cleanupHwlocTopology(); // Data structures for NIXL topology-aware grouping algorithm @@ -162,6 +164,16 @@ class nixlLibfabricTopology { uint8_t &bus_id, size_t &link_speed); + void + collectNicInfo(NicInfo &nic, + const std::string &name, + const std::string &pcie_addr, + hwloc_obj_t hwloc_node, + uint16_t domain_id, + uint8_t bus_id, + uint8_t device_id, + uint8_t function_id); + // finds out the PCIe bandwidth limit of all NUMA nodes (determined by sum of connected PCIe // switches/bridges) void diff --git a/test/gtest/hw_warning_test.cpp b/test/gtest/hw_warning_test.cpp index 83228b3b..eef9b259 100644 --- a/test/gtest/hw_warning_test.cpp +++ b/test/gtest/hw_warning_test.cpp @@ -181,12 +181,6 @@ TEST_F(HardwareWarningTest, EfaHardwareMismatchNoWarning) { nixlAgent agent("EfaTestAgent", nixlAgentConfig(true)); - const gtest::LogIgnoreGuard lig_no_efa( - "Could not find EFA devices, rail selection for DRAM memory type aborted"); - const gtest::LogIgnoreGuard lig_rail_fallback( - "Using default \\(all\\) rail selection policy for DRAM memory type due to " - "previous errors"); - for (const auto &name : backends) { nixlBackendH *backend; EXPECT_EQ(agent.createBackend(name, {}, backend), NIXL_SUCCESS); diff --git a/test/unit/utils/libfabric/libfabric_topology_test.cpp b/test/unit/utils/libfabric/libfabric_topology_test.cpp index ac6bb605..76776c31 100644 --- a/test/unit/utils/libfabric/libfabric_topology_test.cpp +++ b/test/unit/utils/libfabric/libfabric_topology_test.cpp @@ -182,6 +182,24 @@ static TopologyInfo topologies[] = { .numa_capacity = 32, // topmost switch on g6 report link speed 32 GB/s .numa_rail_count = 1, // should default to all rails, which is 1 .test_scenarios = {{false, 0}, {true, 50}, {true, 100}, {true, 200}}, + .rail_partition = {{{0}}, {{0}}}}, // pretending single switch in each node, both using rail 0 + + // + // C-series + // + + // c5.18xl instance type (EFA connected directly to host bridge, no PCIe switch involved) + {.enable = true, + .instance_type = "c5n.18xl", + .topo_file = "c5n.18xl-topo.xml", + .numa_node_count = 0, // no NIC is attached to NUMA node, the only NIC is attached to machine + .nic_count = 1, + .nic_line_speed = 0, // neither fi-info nor hwloc report speed for the NIC + .nic_upstream_link_speed = 0, // c5 reports upstream link speed 0 + .switch_count = 0, // no bridge/switch spec other than host bridge with link speed 0 + .numa_capacity = 0, // topmost switch/bridge on c5 report link speed 0 + .numa_rail_count = 1, // should default to all rails, which is 1 + .test_scenarios = {{false, 0}, {true, 50}, {true, 100}, {true, 200}}, .rail_partition = {{{0}}, {{0}}}} // pretending single switch in each node, both using rail 0 // end of list @@ -1153,7 +1171,9 @@ validateRailSelection(const std::vector &selected_rails, size_t bandwidth = override_bandwidth ? bandwidth_limit : curr_topology->numa_rail_count * curr_topology->nic_line_speed; - size_t expected_rail_count = bandwidth / curr_topology->nic_line_speed; + // on C series the nic line speed is reported as zero, so be careful here + size_t expected_rail_count = + curr_topology->nic_line_speed ? bandwidth / curr_topology->nic_line_speed : 1; expected_rail_count = std::max(1ul, expected_rail_count); expected_rail_count = std::min(curr_topology->nic_count, expected_rail_count); if (curr_topology->numa_capacity == 0) { diff --git a/test/unit/utils/libfabric/topo/c5n.18xl-topo.xml b/test/unit/utils/libfabric/topo/c5n.18xl-topo.xml new file mode 100644 index 00000000..c87c312c --- /dev/null +++ b/test/unit/utils/libfabric/topo/c5n.18xl-topo.xml @@ -0,0 +1,69 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file From 7385451477acde0c000a209fd7aedb5e729143b2 Mon Sep 17 00:00:00 2001 From: omer-b1 Date: Tue, 14 Apr 2026 23:19:02 +0300 Subject: [PATCH 31/59] nixlbench: Fix pairwise SG stats scaling for multi-rank (#1495) In pairwise SG mode with a single-process storage backend (getSize()==1), one process handles all devices under a shared timer. Multiplying total_data_transferred by num_initiator_dev correctly accounts for all devices' data in that shared measurement. However, commit 7c4f144 (#1407) applies this scaling to all pairwise SG configurations. In multi-rank VRAM setups (getSize() > 1), each rank handles a single device with its own independent timer -- the per-device data is already correct. The multiplication inflates B/W by num_initiator_dev (8x for 8 devices). Fix: add rt->getSize() == 1 guard so multi-rank SG skips the scaling while single-process SG and MG mode remain unchanged. Fixes: #1463 Signed-off-by: Omer Brezel Co-authored-by: Adit Ranadive --- benchmark/nixlbench/src/utils/utils.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/benchmark/nixlbench/src/utils/utils.cpp b/benchmark/nixlbench/src/utils/utils.cpp index 8ef9de11..72be26fd 100644 --- a/benchmark/nixlbench/src/utils/utils.cpp +++ b/benchmark/nixlbench/src/utils/utils.cpp @@ -1096,7 +1096,8 @@ xferBenchUtils::printStats(bool is_target, total_data_transferred = ((block_size * batch_size) * num_iter); // In Bytes avg_latency = (total_duration / (num_iter * batch_size)); // In microsec - if (IS_PAIRWISE_AND_MG() || (IS_PAIRWISE_AND_SG() && xferBenchConfig::num_initiator_dev > 1)) { + if (IS_PAIRWISE_AND_MG() || + (IS_PAIRWISE_AND_SG() && xferBenchConfig::num_initiator_dev > 1 && rt->getSize() == 1)) { total_data_transferred *= xferBenchConfig::num_initiator_dev; // In Bytes avg_latency /= xferBenchConfig::num_initiator_dev; // In microsec } From 2c717f9093b5ea180fc0294bbeaa467cbbdc4f31 Mon Sep 17 00:00:00 2001 From: fengjica <227908179+fengjica@users.noreply.github.com> Date: Tue, 14 Apr 2026 15:02:13 -0700 Subject: [PATCH 32/59] libfabric: fix active rail tracking (#1510) * libfabric: fix active rail tracking Today there is no reference counting to correctly keep track of the status of rails. In addition to fixing that, I add a unit test with some code refactoring to avoid duplicating the common test code. * addressed review feedback: - Only decrement rail refcount when deregisterMemory succeeds, in both the normal deregister path and the registerMemory cleanup paths - Add defensive check in decRailActive to prevent size_t underflow when refcount is already zero - Remove unbalanced incRailActive(0) from postControlMessage since rail 0 is already unconditionally polled in progressActiveRails() - Add unit test for failed deregister preserving active rail state Signed-off-by: Feng Ji Co-authored-by: Adit Ranadive --- .../libfabric/libfabric_rail_manager.cpp | 59 ++-- src/utils/libfabric/libfabric_rail_manager.h | 12 +- .../utils/libfabric/libfabric_mock_stubs.h | 273 +++++++++++++++++ .../libfabric/libfabric_topology_test.cpp | 208 +------------ test/unit/utils/libfabric/meson.build | 17 + .../libfabric/rail_active_refcount_test.cpp | 290 ++++++++++++++++++ 6 files changed, 624 insertions(+), 235 deletions(-) create mode 100644 test/unit/utils/libfabric/libfabric_mock_stubs.h create mode 100644 test/unit/utils/libfabric/rail_active_refcount_test.cpp diff --git a/src/utils/libfabric/libfabric_rail_manager.cpp b/src/utils/libfabric/libfabric_rail_manager.cpp index 0e2e78b8..5b906dad 100644 --- a/src/utils/libfabric/libfabric_rail_manager.cpp +++ b/src/utils/libfabric/libfabric_rail_manager.cpp @@ -758,7 +758,10 @@ nixlLibfabricRailManager::registerMemory(void *buffer, for (size_t j = 0; j < i; ++j) { const size_t cleanup_idx = selected_rails[j]; if (mr_list_out[cleanup_idx]) { - rails_[cleanup_idx]->deregisterMemory(mr_list_out[cleanup_idx]); + if (rails_[cleanup_idx]->deregisterMemory(mr_list_out[cleanup_idx]) == + NIXL_SUCCESS) { + decRailActive(cleanup_idx); + } mr_list_out[cleanup_idx] = nullptr; } } @@ -776,7 +779,10 @@ nixlLibfabricRailManager::registerMemory(void *buffer, for (size_t j = 0; j < i; ++j) { const size_t cleanup_idx = selected_rails[j]; if (mr_list_out[cleanup_idx]) { - rails_[cleanup_idx]->deregisterMemory(mr_list_out[cleanup_idx]); + if (rails_[cleanup_idx]->deregisterMemory(mr_list_out[cleanup_idx]) == + NIXL_SUCCESS) { + decRailActive(cleanup_idx); + } mr_list_out[cleanup_idx] = nullptr; } } @@ -787,7 +793,7 @@ nixlLibfabricRailManager::registerMemory(void *buffer, key_list_out[rail_idx] = key; // Mark rail as active for progress tracking optimization - markRailActive(rail_idx); + incRailActive(rail_idx); NIXL_DEBUG << "Registered memory on rail " << rail_idx << " (mr=" << static_cast(mr) << ", key=" << key << ")"; @@ -816,11 +822,12 @@ nixlLibfabricRailManager::deregisterMemory(const std::vector &selected_r if (mr_list[rail_idx]) { nixl_status_t status = rails_[rail_idx]->deregisterMemory(mr_list[rail_idx]); - if (status != NIXL_SUCCESS) { + if (status == NIXL_SUCCESS) { + decRailActive(rail_idx); + } else { NIXL_ERROR << "Failed to deregister memory on rail " << rail_idx; overall_status = status; } - markRailInactive(rail_idx); } } @@ -934,9 +941,6 @@ nixlLibfabricRailManager::postControlMessage(ControlMessageType msg_type, NIXL_DEBUG << "Sending control message type " << msg_type_value << " agent_idx=" << agent_idx << " XFER_ID=" << xfer_id << " imm_data=" << imm_data << " on rail " << rail_id; - // Mark rail 0 as active so its CQ gets progressed - markRailActive(rail_id); - // Use rail 0 for notifications nixl_status_t status = rails_[rail_id]->postSend(imm_data, dest_addr, req); @@ -959,7 +963,9 @@ nixlLibfabricRailManager::progressActiveRails() { std::lock_guard lock(active_rails_mutex_); // Always progress rail 0 for notifications (SEND/RECV) rails_to_process.insert(0); - rails_to_process.insert(active_rails_.begin(), active_rails_.end()); + for (const auto &[rail_id, refcount] : active_rails_) { + rails_to_process.insert(rail_id); + } } // Process rails without holding the lock @@ -1177,30 +1183,35 @@ nixlLibfabricRailManager::deserializeRailEndpoints( } void -nixlLibfabricRailManager::markRailActive(size_t rail_id) { +nixlLibfabricRailManager::incRailActive(size_t rail_id) { if (rail_id >= rails_.size()) { - NIXL_ERROR << "Invalid rail ID for markRailActive: " << rail_id; + NIXL_ERROR << "Invalid rail ID for incRailActive: " << rail_id; return; } std::lock_guard lock(active_rails_mutex_); - bool was_inserted = active_rails_.insert(rail_id).second; - - if (was_inserted) { - NIXL_DEBUG << "Marked rail " << rail_id - << " as active (total active: " << active_rails_.size() << ")"; - } else { - NIXL_TRACE << "Rail " << rail_id << " was already active"; - } + active_rails_[rail_id]++; + NIXL_DEBUG << "incRailActive rail " << rail_id << " (refcount: " << active_rails_[rail_id] + << ", total active: " << active_rails_.size() << ")"; } void -nixlLibfabricRailManager::markRailInactive(size_t rail_id) { +nixlLibfabricRailManager::decRailActive(size_t rail_id) { std::lock_guard lock(active_rails_mutex_); - size_t erased = active_rails_.erase(rail_id); - if (erased > 0) { - NIXL_DEBUG << "Marked rail " << rail_id - << " as inactive (total active: " << active_rails_.size() << ")"; + auto it = active_rails_.find(rail_id); + if (it != active_rails_.end()) { + // The reference count is always non-zero under normal situation. + // Here is a defensive check, assuming a bug somewhere else set it to zero. + if (it->second > 0) { + it->second--; + } + if (it->second == 0) { + active_rails_.erase(it); + NIXL_DEBUG << "decRailActive rail " << rail_id + << " now inactive (total active: " << active_rails_.size() << ")"; + } else { + NIXL_DEBUG << "decRailActive rail " << rail_id << " (refcount: " << it->second << ")"; + } } else { NIXL_TRACE << "Rail " << rail_id << " was not in active set"; } diff --git a/src/utils/libfabric/libfabric_rail_manager.h b/src/utils/libfabric/libfabric_rail_manager.h index 1d4b424f..3e816edc 100644 --- a/src/utils/libfabric/libfabric_rail_manager.h +++ b/src/utils/libfabric/libfabric_rail_manager.h @@ -253,13 +253,13 @@ class nixlLibfabricRailManager { validateAllRailsInitialized(); // Active Rail Management APIs - /** Mark rail as active for progress tracking optimization */ + /** Increment active reference count on a rail */ void - markRailActive(size_t rail_id); + incRailActive(size_t rail_id); - /** Mark rail as inactive for progress tracking optimization */ + /** Decrement active reference count on a rail; rail becomes inactive when count reaches zero */ void - markRailInactive(size_t rail_id); + decRailActive(size_t rail_id); /** Clear all active rail markings */ void @@ -351,8 +351,8 @@ class nixlLibfabricRailManager { // EFA device to rail mapping std::unordered_map efa_device_to_rail_map; - // Active Rail Tracking System - std::unordered_set active_rails_; + // active rails with reference counting (always positive) + std::unordered_map active_rails_; mutable std::mutex active_rails_mutex_; // rail selection policy for DRAM memory type diff --git a/test/unit/utils/libfabric/libfabric_mock_stubs.h b/test/unit/utils/libfabric/libfabric_mock_stubs.h new file mode 100644 index 00000000..b48b1236 --- /dev/null +++ b/test/unit/utils/libfabric/libfabric_mock_stubs.h @@ -0,0 +1,273 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 Amazon.com, Inc. and affiliates. + * 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. + */ + +/* + * Shared libfabric stub ops for unit tests that need to construct a + * nixlLibfabricRailManager without real hardware. Include this header + * in exactly one translation unit per test executable (it contains + * definitions, not just declarations). + */ +#ifndef NIXL_TEST_UNIT_UTILS_LIBFABRIC_LIBFABRIC_MOCK_STUBS_H +#define NIXL_TEST_UNIT_UTILS_LIBFABRIC_LIBFABRIC_MOCK_STUBS_H + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +// helper template for allocating zeroed C structs (free()-able) +template +T * +malloc_zero() { + T *res = (T *)calloc(1, sizeof(T)); + if (res == nullptr) { + std::stringstream ss; + ss << "Failed to allocate " << sizeof(T) << " bytes"; + throw std::runtime_error(ss.str()); + } + return res; +} + +// --- AV stubs --- + +static int +fi_av_close_stub(struct fid *fid) { + free(fid); + return 0; +} + +static struct fi_ops fi_av_fid_ops_stub{ + .close = fi_av_close_stub, +}; + +static fi_ops_av av_ops_stub = {}; + +static int +fi_av_open_stub(struct fid_domain *domain, + struct fi_av_attr *attr, + struct fid_av **av, + void *context) { + *av = malloc_zero(); + (*av)->fid.ops = &fi_av_fid_ops_stub; + (*av)->ops = &av_ops_stub; + return 0; +} + +// --- CQ stubs --- + +static int +fi_cq_close_stub(struct fid *fid) { + free(fid); + return 0; +} + +static struct fi_ops fi_cq_fid_ops_stub{ + .close = fi_cq_close_stub, +}; + +static fi_ops_cq cq_ops_stub = {}; + +static int +fi_cq_open_stub(struct fid_domain *domain, + struct fi_cq_attr *attr, + struct fid_cq **cq, + void *context) { + *cq = malloc_zero(); + (*cq)->fid.ops = &fi_cq_fid_ops_stub; + (*cq)->ops = &cq_ops_stub; + return 0; +} + +// --- EP stubs --- + +static int +fi_ep_setopt_stub(fid_t fid, int level, int optname, const void *optval, size_t optlen) { + return 0; +} + +static fi_ops_ep fi_ep_ops_stub = { + .setopt = fi_ep_setopt_stub, +}; + +static int +fi_ep_bind_stub(struct fid *fid, struct fid *bfid, uint64_t flags) { + return 0; +} + +static int +fi_ep_control_stub(struct fid *fid, int command, void *arg) { + return 0; +} + +static int +fi_ep_cm_getname_stub(fid_t fid, void *addr, size_t *addrlen) { + return 0; +} + +static int +fi_ep_close_stub(struct fid *fid) { + free(fid); + return 0; +} + +static struct fi_ops fi_ep_fid_ops_stub{ + .close = fi_ep_close_stub, + .bind = fi_ep_bind_stub, + .control = fi_ep_control_stub, +}; + +static struct fi_ops_cm fi_ep_cm_ops_stub{ + .getname = fi_ep_cm_getname_stub, +}; + +static ssize_t +fi_ep_recvmsg_stub(struct fid_ep *ep, const struct fi_msg *msg, uint64_t flags) { + return 0; +} + +static fi_ops_msg fi_ep_msg_ops_stub{ + .recvmsg = fi_ep_recvmsg_stub, +}; + +static int +fi_endpoint_stub(struct fid_domain *domain, + struct fi_info *info, + struct fid_ep **ep, + void *context) { + *ep = malloc_zero(); + (*ep)->ops = &fi_ep_ops_stub; + (*ep)->fid.ops = &fi_ep_fid_ops_stub; + (*ep)->cm = &fi_ep_cm_ops_stub; + (*ep)->msg = &fi_ep_msg_ops_stub; + return 0; +} + +// --- MR stubs --- + +static int +fi_mr_close_stub(struct fid *fid) { + free(fid); + return 0; +} + +static fi_ops fi_mr_self_ops_stub = { + .close = fi_mr_close_stub, +}; + +static int +fi_mr_reg_stub(struct fid *fid, + const void *buf, + size_t len, + uint64_t access, + uint64_t offset, + uint64_t requested_key, + uint64_t flags, + struct fid_mr **mr, + void *context) { + *mr = malloc_zero(); + (*mr)->fid.ops = &fi_mr_self_ops_stub; + return 0; +} + +static int +fi_mr_regattr_stub(struct fid *fid, + const struct fi_mr_attr *attr, + uint64_t flags, + struct fid_mr **mr) { + *mr = malloc_zero(); + (*mr)->fid.ops = &fi_mr_self_ops_stub; + return 0; +} + +static fi_ops_mr fi_mr_ops_stub{ + .reg = fi_mr_reg_stub, + .regv = nullptr, + .regattr = fi_mr_regattr_stub, +}; + +// --- Domain stubs --- + +static fi_ops_domain domain_ops_stub = {.av_open = fi_av_open_stub, + .cq_open = fi_cq_open_stub, + .endpoint = fi_endpoint_stub, + .scalable_ep = nullptr, + .cntr_open = nullptr, + .poll_open = nullptr, + .stx_ctx = nullptr, + .srx_ctx = nullptr, + .query_atomic = nullptr, + .query_collective = nullptr, + .endpoint2 = nullptr}; + +static int +fi_domain_close_stub(struct fid *fid) { + free(fid); + return 0; +} + +static struct fi_ops fi_domain_ops_stub{ + .close = fi_domain_close_stub, +}; + +static int +fi_domain_stub(struct fid_fabric *fabric, + struct fi_info *info, + struct fid_domain **domain, + void *context) { + *domain = malloc_zero(); + (*domain)->fid.ops = &fi_domain_ops_stub; + (*domain)->ops = &domain_ops_stub; + (*domain)->mr = &fi_mr_ops_stub; + return 0; +} + +// --- Fabric stubs --- + +static fi_ops_fabric fabric_ops_stub{.domain = fi_domain_stub, + .passive_ep = nullptr, + .eq_open = nullptr, + .wait_open = nullptr, + .trywait = nullptr, + .domain2 = nullptr}; + +static int +fi_fabric_close_stub(struct fid *fid) { + free(fid); + return 0; +} + +static struct fi_ops fi_fabric_ops_stub{ + .close = fi_fabric_close_stub, +}; + +// Helper: build a mock fid_fabric using the stubs above +static struct fid_fabric * +mock_fabric_create() { + fid_fabric *fabric = malloc_zero(); + fabric->fid.ops = &fi_fabric_ops_stub; + fabric->ops = &fabric_ops_stub; + return fabric; +} + +#endif // NIXL_TEST_UNIT_UTILS_LIBFABRIC_LIBFABRIC_MOCK_STUBS_H diff --git a/test/unit/utils/libfabric/libfabric_topology_test.cpp b/test/unit/utils/libfabric/libfabric_topology_test.cpp index 76776c31..c6cc98d9 100644 --- a/test/unit/utils/libfabric/libfabric_topology_test.cpp +++ b/test/unit/utils/libfabric/libfabric_topology_test.cpp @@ -566,19 +566,8 @@ __wrap_get_mempolicy(int *mode, return 0; } -// helper template function for allocating flat type zeroed memory with malloc but keeping new -// operator semantics of throwing bad_alloc when allocation fails -template -T * -malloc_zero() { - T *res = (T *)calloc(1, sizeof(T)); - if (res == nullptr) { - std::stringstream ss; - ss << "Failed to allocate " << sizeof(T) << " bytes"; - throw std::runtime_error(ss.str()); - } - return res; -} +// shared stub ops for mocking libfabric objects +#include "libfabric_mock_stubs.h" // mock fi_getinfo by current topology in use extern "C" int @@ -649,195 +638,6 @@ __wrap_fi_getinfo(uint32_t version, return 0; // or FI_SUCCESS } -// stubs -static int -fi_av_close_stub(struct fid *fid) { - free(fid); - return 0; -} - -struct fi_ops fi_av_fid_ops_stub{ - .close = fi_av_close_stub, -}; - -static fi_ops_av av_ops_stub = {}; - -static int -fi_av_open_stub(struct fid_domain *domain, - struct fi_av_attr *attr, - struct fid_av **av, - void *context) { - *av = malloc_zero(); - (*av)->fid.ops = &fi_av_fid_ops_stub; - (*av)->ops = &av_ops_stub; - return 0; -} - -static int -fi_cq_close_stub(struct fid *fid) { - free(fid); - return 0; -} - -struct fi_ops fi_cq_fid_ops_stub{ - .close = fi_cq_close_stub, -}; - -static fi_ops_cq cq_ops_stub = {}; - -static int -fi_cq_open_stub(struct fid_domain *domain, - struct fi_cq_attr *attr, - struct fid_cq **cq, - void *context) { - *cq = malloc_zero(); - (*cq)->fid.ops = &fi_cq_fid_ops_stub; - (*cq)->ops = &cq_ops_stub; - return 0; -} - -static int -fi_ep_setopt_stub(fid_t fid, int level, int optname, const void *optval, size_t optlen) { - return 0; -} - -static fi_ops_ep fi_ep_ops_stub = { - .setopt = fi_ep_setopt_stub, -}; - -static int -fi_ep_bind_stub(struct fid *fid, struct fid *bfid, uint64_t flags) { - return 0; -} - -static int -fi_ep_control_stub(struct fid *fid, int command, void *arg) { - return 0; -} - -static int -fi_ep_cm_getname_stub(fid_t fid, void *addr, size_t *addrlen) { - return 0; -} - -static int -fi_ep_close_stub(struct fid *fid) { - free(fid); - return 0; -} - -struct fi_ops fi_ep_fid_ops_stub{ - .close = fi_ep_close_stub, - .bind = fi_ep_bind_stub, - .control = fi_ep_control_stub, -}; - -struct fi_ops_cm fi_ep_cm_ops_stub{ - .getname = fi_ep_cm_getname_stub, -}; - -static ssize_t -fi_ep_recvmsg_stub(struct fid_ep *ep, const struct fi_msg *msg, uint64_t flags) { - return 0; -} - -static fi_ops_msg fi_ep_msg_ops_stub{ - .recvmsg = fi_ep_recvmsg_stub, -}; - -static int -fi_endpoint_stub(struct fid_domain *domain, - struct fi_info *info, - struct fid_ep **ep, - void *context) { - *ep = malloc_zero(); - (*ep)->ops = &fi_ep_ops_stub; - (*ep)->fid.ops = &fi_ep_fid_ops_stub; - (*ep)->cm = &fi_ep_cm_ops_stub; - (*ep)->msg = &fi_ep_msg_ops_stub; - return 0; -} - -static fi_ops_domain domain_ops_stub = {.av_open = fi_av_open_stub, - .cq_open = fi_cq_open_stub, - .endpoint = fi_endpoint_stub, - .scalable_ep = nullptr, - .cntr_open = nullptr, - .poll_open = nullptr, - .stx_ctx = nullptr, - .srx_ctx = nullptr, - .query_atomic = nullptr, - .query_collective = nullptr, - .endpoint2 = nullptr}; - -static int -fi_mr_close_stub(struct fid *fid) { - free(fid); - return 0; -} - -static fi_ops fi_mr_self_ops_stub = { - .close = fi_mr_close_stub, -}; - -static int -fi_mr_reg_stub(struct fid *fid, - const void *buf, - size_t len, - uint64_t access, - uint64_t offset, - uint64_t requested_key, - uint64_t flags, - struct fid_mr **mr, - void *context) { - *mr = malloc_zero(); - (*mr)->fid.ops = &fi_mr_self_ops_stub; - return 0; -} - -static fi_ops_mr fi_mr_ops_stub{ - .reg = fi_mr_reg_stub, -}; - -static int -fi_domain_close_stub(struct fid *fid) { - free(fid); - return 0; -} - -struct fi_ops fi_domain_ops_stub{ - .close = fi_domain_close_stub, -}; - -static int -fi_domain_stub(struct fid_fabric *fabric, - struct fi_info *info, - struct fid_domain **domain, - void *context) { - *domain = malloc_zero(); - (*domain)->fid.ops = &fi_domain_ops_stub; - (*domain)->ops = &domain_ops_stub; - (*domain)->mr = &fi_mr_ops_stub; - return 0; -} - -static fi_ops_fabric fabric_ops_stub{.domain = fi_domain_stub, - .passive_ep = nullptr, - .eq_open = nullptr, - .wait_open = nullptr, - .trywait = nullptr, - .domain2 = nullptr}; - -static int -fi_fabric_close_stub(struct fid *fid) { - free(fid); - return 0; -} - -struct fi_ops fi_fabric_ops_stub{ - .close = fi_fabric_close_stub, -}; - // mock fi_fabric to get past rail manager constructor extern "C" int __wrap_fi_fabric(struct fi_fabric_attr *attr, struct fid_fabric **fabric, void *context) { @@ -847,9 +647,7 @@ __wrap_fi_fabric(struct fi_fabric_attr *attr, struct fid_fabric **fabric, void * return __real_fi_fabric(attr, fabric, context); } - *fabric = malloc_zero(); - (*fabric)->fid.ops = &fi_fabric_ops_stub; - (*fabric)->ops = &fabric_ops_stub; + *fabric = mock_fabric_create(); return 0; } diff --git a/test/unit/utils/libfabric/meson.build b/test/unit/utils/libfabric/meson.build index 4975abee..2ce31cd7 100644 --- a/test/unit/utils/libfabric/meson.build +++ b/test/unit/utils/libfabric/meson.build @@ -45,4 +45,21 @@ if get_option('buildtype') != 'release' # install topology test files install_subdir('topo', strip_directory : true, install_dir : get_option('bindir')) + rail_active_refcount_test_link_args = [ + '-Wl,--wrap=fi_getinfo', + '-Wl,--wrap=numa_max_node', + '-Wl,--wrap=numa_num_configured_nodes', + '-Wl,--wrap=fi_fabric' ] + + rail_active_refcount_test_bin = executable('rail_active_refcount_test', + 'rail_active_refcount_test.cpp', + dependencies: libfabric_utils_dep, + include_directories: [nixl_inc_dirs, utils_inc_dirs], + link_with: libfabric_utils_lib, + cpp_args: libfabric_test_cpp_args, + link_args: rail_active_refcount_test_link_args, + install: true) + + test('rail_active_refcount_test', rail_active_refcount_test_bin) + endif diff --git a/test/unit/utils/libfabric/rail_active_refcount_test.cpp b/test/unit/utils/libfabric/rail_active_refcount_test.cpp new file mode 100644 index 00000000..09860727 --- /dev/null +++ b/test/unit/utils/libfabric/rail_active_refcount_test.cpp @@ -0,0 +1,290 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 Amazon.com, Inc. and affiliates. + * 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. + */ + +/* + * Unit test for active rail reference counting in nixlLibfabricRailManager. + */ + +#include "libfabric/libfabric_rail_manager.h" +#include "libfabric/libfabric_common.h" +#include "common/nixl_log.h" +#include "libfabric_mock_stubs.h" + +#include +#include + +// Number of fake EFA devices to create +static const size_t NUM_FAKE_RAILS = 4; + +// --- Unconditional __wrap_* functions (no isTesting gate needed) --- + +extern "C" int +__wrap_numa_max_node() { + return 1; +} + +extern "C" int +__wrap_numa_num_configured_nodes() { + return 2; +} + +extern "C" int +__wrap_fi_getinfo(uint32_t /*version*/, + const char * /*node*/, + const char * /*service*/, + uint64_t /*flags*/, + const struct fi_info * /*hints*/, + struct fi_info **info) { + // Build a linked list of NUM_FAKE_RAILS fake EFA devices + fi_info *head = nullptr; + fi_info *prev = nullptr; + for (size_t i = 0; i < NUM_FAKE_RAILS; ++i) { + fi_info *fi = malloc_zero(); + + fi->domain_attr = malloc_zero(); + std::string name = "efa_" + std::to_string(i); + fi->domain_attr->name = strdup(name.c_str()); + + fi->fabric_attr = malloc_zero(); + fi->fabric_attr->prov_name = strdup("efa"); + fi->fabric_attr->name = strdup("efa"); + + fi->ep_attr = malloc_zero(); + fi->ep_attr->type = FI_EP_RDM; + + fi->nic = malloc_zero(); + fi->nic->bus_attr = malloc_zero(); + fi->nic->bus_attr->bus_type = FI_BUS_PCI; + fi->nic->bus_attr->attr.pci.domain_id = 0; + fi->nic->bus_attr->attr.pci.bus_id = static_cast(i); + fi->nic->bus_attr->attr.pci.device_id = 0; + fi->nic->bus_attr->attr.pci.function_id = 0; + + fi->nic->link_attr = malloc_zero(); + fi->nic->link_attr->speed = 100ull * NIXL_LIBFABRIC_GIGA; + + if (prev) { + prev->next = fi; + } else { + head = fi; + } + prev = fi; + } + *info = head; + return 0; +} + +extern "C" int +__wrap_fi_fabric(struct fi_fabric_attr * /*attr*/, struct fid_fabric **fabric, void * /*context*/) { + *fabric = mock_fabric_create(); + return 0; +} + +// --- Test helpers --- + +#define TEST_ASSERT(cond, msg) \ + do { \ + if (!(cond)) { \ + std::cerr << "FAIL: " << (msg) << " [" << __FILE__ << ":" << __LINE__ << "]" \ + << std::endl; \ + return 1; \ + } \ + } while (0) + +// --- Tests --- + +static int +testIncDecBasic(nixlLibfabricRailManager &mgr) { + NIXL_INFO << " testIncDecBasic"; + + // inc rail 0 twice + mgr.incRailActive(0); + mgr.incRailActive(0); + TEST_ASSERT(mgr.getActiveRailCount() == 1, "one rail active after two incs on same rail"); + + // dec once — refcount drops to 1, rail still active + mgr.decRailActive(0); + TEST_ASSERT(mgr.getActiveRailCount() == 1, "rail still active after one dec"); + + // dec again — refcount drops to 0, rail removed + mgr.decRailActive(0); + TEST_ASSERT(mgr.getActiveRailCount() == 0, "rail removed after second dec"); + + return 0; +} + +static int +testDecNonExistent(nixlLibfabricRailManager &mgr) { + NIXL_INFO << " testDecNonExistent"; + + // dec on a rail that was never inc'd — should not crash + mgr.decRailActive(1); + TEST_ASSERT(mgr.getActiveRailCount() == 0, "count unchanged after dec on non-existent rail"); + + return 0; +} + +static int +testIncInvalidRailId(nixlLibfabricRailManager &mgr) { + NIXL_INFO << " testIncInvalidRailId"; + + // inc with rail_id >= rails_.size() — should be a no-op + mgr.incRailActive(NUM_FAKE_RAILS + 10); + TEST_ASSERT(mgr.getActiveRailCount() == 0, "no-op for invalid rail id"); + + return 0; +} + +static int +testClearActiveRails(nixlLibfabricRailManager &mgr) { + NIXL_INFO << " testClearActiveRails"; + + mgr.incRailActive(0); + mgr.incRailActive(1); + mgr.incRailActive(2); + TEST_ASSERT(mgr.getActiveRailCount() == 3, "three rails active"); + + mgr.clearActiveRails(); + TEST_ASSERT(mgr.getActiveRailCount() == 0, "all rails cleared"); + + return 0; +} + +static int +testMultipleRailsIndependent(nixlLibfabricRailManager &mgr) { + NIXL_INFO << " testMultipleRailsIndependent"; + + // rail 0: refcount 3, rail 1: refcount 1, rail 2: refcount 2 + mgr.incRailActive(0); + mgr.incRailActive(0); + mgr.incRailActive(0); + mgr.incRailActive(1); + mgr.incRailActive(2); + mgr.incRailActive(2); + TEST_ASSERT(mgr.getActiveRailCount() == 3, "three distinct rails active"); + + // remove rail 1 entirely + mgr.decRailActive(1); + TEST_ASSERT(mgr.getActiveRailCount() == 2, "rail 1 removed, two rails remain"); + + // dec rail 0 once — still active (refcount 2) + mgr.decRailActive(0); + TEST_ASSERT(mgr.getActiveRailCount() == 2, "rail 0 still active after one dec"); + + // clean up + mgr.clearActiveRails(); + return 0; +} + +static int +testRegisterDeregisterRefcount(nixlLibfabricRailManager &mgr) { + NIXL_INFO << " testRegisterDeregisterRefcount"; + + // Allocate a DRAM buffer to register + char buf[4096] = {}; + std::vector mr_list; + std::vector key_list; + std::vector selected_rails; + + TEST_ASSERT(mgr.getActiveRailCount() == 0, "no active rails before register"); + + nixl_status_t st = + mgr.registerMemory(buf, sizeof(buf), DRAM_SEG, 0, "", mr_list, key_list, selected_rails); + TEST_ASSERT(st == NIXL_SUCCESS, "registerMemory succeeded"); + TEST_ASSERT(mgr.getActiveRailCount() == NUM_FAKE_RAILS, "all rails active after DRAM register"); + + // Register a second time — refcounts should increase but active count stays the same + std::vector mr_list2; + std::vector key_list2; + std::vector selected_rails2; + st = + mgr.registerMemory(buf, sizeof(buf), DRAM_SEG, 0, "", mr_list2, key_list2, selected_rails2); + TEST_ASSERT(st == NIXL_SUCCESS, "second registerMemory succeeded"); + TEST_ASSERT(mgr.getActiveRailCount() == NUM_FAKE_RAILS, + "active count unchanged after second register"); + + // Deregister first — rails still active (refcount > 0) + st = mgr.deregisterMemory(selected_rails, mr_list); + TEST_ASSERT(st == NIXL_SUCCESS, "first deregisterMemory succeeded"); + TEST_ASSERT(mgr.getActiveRailCount() == NUM_FAKE_RAILS, + "all rails still active after first deregister"); + + // Deregister second — refcounts drop to 0, rails removed + st = mgr.deregisterMemory(selected_rails2, mr_list2); + TEST_ASSERT(st == NIXL_SUCCESS, "second deregisterMemory succeeded"); + TEST_ASSERT(mgr.getActiveRailCount() == 0, "no active rails after full deregister"); + + return 0; +} + +static int +fi_mr_close_fail(struct fid * /*fid*/) { + return -FI_EIO; +} + +static int +testDeregisterFailureKeepsRailActive(nixlLibfabricRailManager &mgr) { + NIXL_INFO << " testDeregisterFailureKeepsRailActive"; + + char buf[4096] = {}; + std::vector mr_list; + std::vector key_list; + std::vector selected_rails; + + nixl_status_t st = + mgr.registerMemory(buf, sizeof(buf), DRAM_SEG, 0, "", mr_list, key_list, selected_rails); + TEST_ASSERT(st == NIXL_SUCCESS, "registerMemory succeeded"); + TEST_ASSERT(mgr.getActiveRailCount() == NUM_FAKE_RAILS, "all rails active after register"); + + // Make fi_close fail so deregisterMemory fails + fi_mr_self_ops_stub.close = fi_mr_close_fail; + st = mgr.deregisterMemory(selected_rails, mr_list); + fi_mr_self_ops_stub.close = fi_mr_close_stub; + + TEST_ASSERT(st != NIXL_SUCCESS, "deregisterMemory should fail"); + TEST_ASSERT(mgr.getActiveRailCount() == NUM_FAKE_RAILS, + "all rails still active after failed deregister"); + + // Clean up: successful deregister to reset state + st = mgr.deregisterMemory(selected_rails, mr_list); + mgr.clearActiveRails(); + return 0; +} + +int +main() { + NIXL_INFO << "=== Rail Active Refcount Test ==="; + NIXL_INFO << "Using mock stubs (__wrap_fi_getinfo, __wrap_fi_fabric, etc.)"; + + // Construct rail manager with mocked hardware (NUM_FAKE_RAILS rails) + nixlLibfabricRailManager mgr(0); + TEST_ASSERT(mgr.getNumRails() == NUM_FAKE_RAILS, + "expected " + std::to_string(NUM_FAKE_RAILS) + " rails"); + + int res; + if ((res = testIncDecBasic(mgr)) != 0) return res; + if ((res = testDecNonExistent(mgr)) != 0) return res; + if ((res = testIncInvalidRailId(mgr)) != 0) return res; + if ((res = testClearActiveRails(mgr)) != 0) return res; + if ((res = testMultipleRailsIndependent(mgr)) != 0) return res; + if ((res = testRegisterDeregisterRefcount(mgr)) != 0) return res; + if ((res = testDeregisterFailureKeepsRailActive(mgr)) != 0) return res; + + NIXL_INFO << "=== All rail active refcount tests PASSED ==="; + return 0; +} From 1169d091de37b7bc60c8f82b000efac3716d7ae2 Mon Sep 17 00:00:00 2001 From: Anshuman Goswami Date: Tue, 14 Apr 2026 18:16:00 -0700 Subject: [PATCH 33/59] fix: improve log levels in libfabric/EFA code (#1462) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: improve log levels in libfabric/EFA code Audit and correct log levels across libfabric and EFA-related files per the logging guidelines: - Error cases → NIXL_ERROR - Unexpected things causing unlikely paths → NIXL_WARN - Connection establishments and control plane info → NIXL_INFO - Per-transfer debug information → NIXL_DEBUG Changes by file: libfabric_backend.cpp: - DEBUG→INFO: backend init, rail manager creation, self-connection, progress thread lifecycle, connection establishment/state changes, stored multirail connection, created agent connection - INFO→DEBUG: rail count mismatch detail (per-connection) - ERROR→WARN: disconnect of non-existent connection (non-fatal) libfabric_rail.cpp: - INFO→DEBUG: pool init/expansion, buffer chunk creation, buffer initialization, recv pre-posting, EAGAIN retry logging - INFO→TRACE: freeing info structure in cleanup libfabric_rail_manager.cpp: - WARN→ERROR: fallback to default rail policy due to previous errors libfabric_topology.cpp: - WARN→ERROR: no network devices found (returns error status) libfabric_common.cpp: - WARN→ERROR: no network devices found with any provider - DEBUG→TRACE: per-parameter env var check (very verbose) - TRACE→DEBUG: env var config override (notable decision) - ERROR→WARN: non-fatal config parse errors with default fallback Resolves: NIXL-46 * fix: add missing INFO logs for customer diagnostics in libfabric/EFA Add INFO-level logs at key configuration and establishment points so that an INFO-level log from a customer is sufficient to diagnose basic configuration and connection issues without reproducing. New INFO logs added: libfabric_backend.cpp: - CUDA address workaround state (enabled/disabled) - Striping threshold value and source (custom/default) - Successful disconnect from agent libfabric_rail_manager.cpp: - Rails creation summary (count + provider) - DRAM rail selection policy chosen (numa-aware with max_rails, or all) - Memory registration summary (rail count, mem_type, device) - Memory deregistration summary (rail count) libfabric_topology.cpp: - Topology summary (NUMA nodes, NICs, GPUs, accelerators) - Average NIC bandwidth - Accelerator-PCI to EFA device mapping (GPU-to-NIC affinity) * style: fix clang-format alignment in log statements * fix: address coderabbitai review comments - Accept SIZE_MAX as valid size_t value (>= to > in boundary check) - Add 'Accelerator-PCI' prefix to topology mapping INFO logs for context --- src/plugins/libfabric/libfabric_backend.cpp | 46 +++++++-------- src/utils/libfabric/libfabric_common.cpp | 16 +++--- src/utils/libfabric/libfabric_rail.cpp | 57 ++++++++++--------- .../libfabric/libfabric_rail_manager.cpp | 10 +++- src/utils/libfabric/libfabric_topology.cpp | 11 +++- 5 files changed, 75 insertions(+), 65 deletions(-) diff --git a/src/plugins/libfabric/libfabric_backend.cpp b/src/plugins/libfabric/libfabric_backend.cpp index cdd0cbc1..aaad7c3c 100644 --- a/src/plugins/libfabric/libfabric_backend.cpp +++ b/src/plugins/libfabric/libfabric_backend.cpp @@ -295,7 +295,7 @@ nixlLibfabricEngine::nixlLibfabricEngine(const nixlBackendInitParams *init_param rail_manager(NIXL_LIBFABRIC_DEFAULT_STRIPING_THRESHOLD), runtime_(FI_HMEM_SYSTEM) { - NIXL_DEBUG << "Initializing Libfabric Backend"; + NIXL_INFO << "Initializing Libfabric Backend"; // this is required for loading rail selection policy by configuration if (rail_manager.init(getCustomParams()) != NIXL_SUCCESS) { @@ -316,11 +316,11 @@ nixlLibfabricEngine::nixlLibfabricEngine(const nixlBackendInitParams *init_param vramInitCtx(); // CUDA address workaround if (nixl::config::checkExistence("NIXL_DISABLE_CUDA_ADDR_WA")) { - NIXL_DEBUG << "Disabling CUDA address workaround"; + NIXL_INFO << "CUDA address workaround: disabled"; cuda_addr_wa_ = false; } else { cuda_addr_wa_ = true; - NIXL_DEBUG << "CUDA address workaround enabled"; + NIXL_INFO << "CUDA address workaround: enabled"; } } #endif @@ -332,19 +332,19 @@ nixlLibfabricEngine::nixlLibfabricEngine(const nixlBackendInitParams *init_param if (getInitParam("striping_threshold", threshold_str) == NIXL_SUCCESS) { try { striping_threshold_ = std::stoull(threshold_str); - NIXL_DEBUG << "Using custom striping threshold: " << striping_threshold_ << " bytes"; + NIXL_INFO << "Striping threshold: " << striping_threshold_ << " bytes (custom)"; } catch (const std::exception &e) { NIXL_WARN << "Invalid striping_threshold value '" << threshold_str << "', using default: " << striping_threshold_ << " bytes"; } } else { - NIXL_DEBUG << "Using default striping threshold: " << striping_threshold_ << " bytes"; + NIXL_INFO << "Striping threshold: " << striping_threshold_ << " bytes (default)"; } // Initialize Rail Manager which will discover the topology and create all rails. try { - NIXL_DEBUG << "Rail Manager created with " << rail_manager.getNumRails() << " rails"; + NIXL_INFO << "Rail Manager created with " << rail_manager.getNumRails() << " rails"; // Set up notification callback on rail 0 size_t notification_rail_id = 0; @@ -383,8 +383,8 @@ nixlLibfabricEngine::nixlLibfabricEngine(const nixlBackendInitParams *init_param std::to_string(conn_status)); } - NIXL_DEBUG << "Created self-connection for agent: " << localAgent << " on " - << rail_manager.getNumRails() << " rails"; + NIXL_INFO << "Created self-connection for agent: " << localAgent << " on " + << rail_manager.getNumRails() << " rails"; // Start Progress thread for rail completion processing if (progress_thread_enabled_) { @@ -392,8 +392,8 @@ nixlLibfabricEngine::nixlLibfabricEngine(const nixlBackendInitParams *init_param rail_manager.getRail(i).setProgressThreadEnabled(true); } - NIXL_DEBUG << "Starting Progress thread for rails with delay: " - << progress_thread_delay_.count() << " microseconds"; + NIXL_INFO << "Starting Progress thread for rails with delay: " + << progress_thread_delay_.count() << " microseconds"; progress_thread_stop_ = false; progress_thread_ = std::thread(&nixlLibfabricEngine::progressThread, this); @@ -401,7 +401,7 @@ nixlLibfabricEngine::nixlLibfabricEngine(const nixlBackendInitParams *init_param NIXL_ERROR << "Failed to start Progress thread"; throw std::runtime_error("Failed to start Progress thread"); } - NIXL_DEBUG << "Progress thread started successfully"; + NIXL_INFO << "Progress thread started successfully"; } else { NIXL_DEBUG << "Progress thread disabled, using manual progress in checkXfer/getNotifs"; } @@ -495,8 +495,8 @@ nixlLibfabricEngine::loadRemoteConnInfo(const std::string &remote_agent, return conn_status; } - NIXL_DEBUG << "Successfully stored multirail connection for " << remote_agent << " on " - << rail_manager.getNumRails() << " rails"; + NIXL_INFO << "Successfully stored multirail connection for " << remote_agent << " on " + << rail_manager.getNumRails() << " rails"; return NIXL_SUCCESS; } @@ -510,8 +510,8 @@ nixlLibfabricEngine::connect(const std::string &remote_agent) { // Check if connection is already established auto it = connections_.find(remote_agent); if (it != connections_.end() && it->second->overall_state_ == ConnectionState::CONNECTED) { - NIXL_DEBUG << "Connection already established for " << remote_agent - << ", fi_addr=" << it->second->rail_remote_addr_list_[0][0]; + NIXL_INFO << "Connection already established for " << remote_agent + << ", fi_addr=" << it->second->rail_remote_addr_list_[0][0]; return NIXL_SUCCESS; } @@ -534,7 +534,7 @@ nixlLibfabricEngine::connect(const std::string &remote_agent) { return NIXL_ERR_NOT_FOUND; } - NIXL_DEBUG << "Successfully established connection for " << remote_agent; + NIXL_INFO << "Successfully established connection for " << remote_agent; return NIXL_SUCCESS; } @@ -543,7 +543,7 @@ nixlLibfabricEngine::disconnect(const std::string &remote_agent) { std::lock_guard lock(connection_state_mutex_); auto it = connections_.find(remote_agent); if (it == connections_.end()) { - NIXL_ERROR << "Disconnect failed. No metadata connection info for " << remote_agent; + NIXL_WARN << "Disconnect failed. No metadata connection info for " << remote_agent; return NIXL_ERR_NOT_FOUND; } // Connection exists - check if already disconnected @@ -560,7 +560,7 @@ nixlLibfabricEngine::disconnect(const std::string &remote_agent) { // Remove connection from map connections_.erase(remote_agent); - NIXL_DEBUG << "Connection erased from the connection map for agent: " << remote_agent; + NIXL_INFO << "Disconnected from agent: " << remote_agent; return NIXL_SUCCESS; } @@ -577,7 +577,7 @@ nixlLibfabricEngine::createAgentConnection( } if (data_rail_endpoints.size() != rail_manager.getNumRails()) { - NIXL_INFO << "Rail count (local: " << rail_manager.getNumRails() + NIXL_WARN << "Rail count mismatch (local: " << rail_manager.getNumRails() << ", remote: " << data_rail_endpoints.size() << ")"; } @@ -611,8 +611,8 @@ nixlLibfabricEngine::createAgentConnection( // Store connection connections_[agent_name] = conn; - NIXL_DEBUG << "Successfully created connection for agent: " << agent_name << " on " - << rail_manager.getNumRails() << " rails"; + NIXL_INFO << "Successfully created connection for agent: " << agent_name << " on " + << rail_manager.getNumRails() << " rails"; return NIXL_SUCCESS; } @@ -658,8 +658,8 @@ nixlLibfabricEngine::establishConnection(const std::string &remote_agent) const NIXL_DEBUG << "Agent index: " << it->second->agent_index_; conn_info->overall_state_ = ConnectionState::CONNECTED; - NIXL_DEBUG << "Connection state for agent " << remote_agent << " is now " - << conn_info->overall_state_; + NIXL_INFO << "Connection state for agent " << remote_agent << " is now " + << conn_info->overall_state_; return NIXL_SUCCESS; } diff --git a/src/utils/libfabric/libfabric_common.cpp b/src/utils/libfabric/libfabric_common.cpp index c1127cc9..99ab88ea 100644 --- a/src/utils/libfabric/libfabric_common.cpp +++ b/src/utils/libfabric/libfabric_common.cpp @@ -108,7 +108,7 @@ getAvailableNetworkDevices() { return {"sockets", {provider_device_map["sockets"][0]}}; } - NIXL_WARN << "No network devices found with any provider"; + NIXL_ERROR << "No network devices found with any provider"; return {"none", {}}; } @@ -274,15 +274,15 @@ getCustomIntParam(const nixl_b_params_t &custom_params, const std::string &key, std::size_t pos = 0; uint64_t parsed_value = std::stoull(value_str, &pos, 10); if (pos != value_str.size()) { - NIXL_ERROR << "Invalid " << key << " configuration value '" << value_str - << "': excess non-digit characters from position " << pos << "('" - << value_str.substr(pos) << "'), using default: " << value; + NIXL_WARN << "Invalid " << key << " configuration value '" << value_str + << "': excess non-digit characters from position " << pos << "('" + << value_str.substr(pos) << "'), using default: " << value; return NIXL_ERR_INVALID_PARAM; } - if (parsed_value >= SIZE_MAX) { - NIXL_ERROR << "Invalid " << key << " configuration value '" << parsed_value - << "': exceeding maximum allowed " << SIZE_MAX - << ", using default: " << value; + if (parsed_value > SIZE_MAX) { + NIXL_WARN << "Invalid " << key << " configuration value '" << parsed_value + << "': exceeding maximum allowed " << SIZE_MAX + << ", using default: " << value; return NIXL_ERR_INVALID_PARAM; } diff --git a/src/utils/libfabric/libfabric_rail.cpp b/src/utils/libfabric/libfabric_rail.cpp index 278e3ea5..e6080e5a 100644 --- a/src/utils/libfabric/libfabric_rail.cpp +++ b/src/utils/libfabric/libfabric_rail.cpp @@ -45,9 +45,9 @@ RequestPool::initializeBasePool(size_t pool_size) { free_indices_.push(i); } - NIXL_INFO << "InitializeBasePool - Rail " << rail_id_ - << " completed. Total requests: " << requests_.size() - << " Free requests: " << free_indices_.size(); + NIXL_DEBUG << "InitializeBasePool - Rail " << rail_id_ + << " completed. Total requests: " << requests_.size() + << " Free requests: " << free_indices_.size(); } void @@ -139,8 +139,8 @@ RequestPool::allocateReq(uint32_t req_id) { return nullptr; } - NIXL_INFO << "AllocateReq on Rail " << rail_id_ << " successfully expanded pool from " - << old_size << " to " << requests_.size() << " requests"; + NIXL_DEBUG << "AllocateReq on Rail " << rail_id_ << " successfully expanded pool from " + << old_size << " to " << requests_.size() << " requests"; } size_t idx = free_indices_.top(); @@ -207,9 +207,9 @@ ControlRequestPool::createBufferChunk(size_t chunk_size, BufferChunk &chunk) { return NIXL_ERR_BACKEND; } - NIXL_INFO << "CreateBufferChunk on Rail " << rail_id_ << " successfully created buffer chunk:" - << " buffer=" << chunk.buffer << " size=" << chunk.size << " mr=" << chunk.mr - << " mr_key=" << fi_mr_key(chunk.mr); + NIXL_DEBUG << "CreateBufferChunk on Rail " << rail_id_ << " successfully created buffer chunk:" + << " buffer=" << chunk.buffer << " size=" << chunk.size << " mr=" << chunk.mr + << " mr_key=" << fi_mr_key(chunk.mr); return NIXL_SUCCESS; } @@ -241,16 +241,16 @@ ControlRequestPool::initialize(struct fid_domain *domain) { requests_[i].operation_type = nixlLibfabricReq::SEND; // Default for control } - NIXL_INFO << "InitializeWithBuffers on Rail " << rail_id_ << " successfully initialized with " - << buffer_chunks_.size() << " buffer chunks"; + NIXL_DEBUG << "InitializeWithBuffers on Rail " << rail_id_ << " successfully initialized with " + << buffer_chunks_.size() << " buffer chunks"; return NIXL_SUCCESS; } nixl_status_t ControlRequestPool::expandPool() { - NIXL_INFO << "Expanding control request pool on rail " << rail_id_ << " from " - << requests_.size() << " to " << (requests_.size() * 2) << " requests"; + NIXL_DEBUG << "Expanding control request pool on rail " << rail_id_ << " from " + << requests_.size() << " to " << (requests_.size() * 2) << " requests"; size_t current_size = requests_.size(); size_t expansion_size = initial_pool_size_; // Add same amount as initial size @@ -292,8 +292,9 @@ ControlRequestPool::expandPool() { requests_[i].operation_type = nixlLibfabricReq::SEND; } - NIXL_INFO << "Successfully expanded control request pool on rail " << rail_id_ << " to " - << requests_.size() << " requests with " << buffer_chunks_.size() << " buffer chunks"; + NIXL_DEBUG << "Successfully expanded control request pool on rail " << rail_id_ << " to " + << requests_.size() << " requests with " << buffer_chunks_.size() + << " buffer chunks"; return NIXL_SUCCESS; } @@ -346,8 +347,8 @@ DataRequestPool::initialize() { nixl_status_t DataRequestPool::expandPool() { - NIXL_INFO << "Expanding data request pool on rail " << rail_id_ << " from " << requests_.size() - << " to " << (requests_.size() * 2) << " requests"; + NIXL_DEBUG << "Expanding data request pool on rail " << rail_id_ << " from " << requests_.size() + << " to " << (requests_.size() * 2) << " requests"; size_t current_size = requests_.size(); size_t expansion_size = initial_pool_size_; // Add same amount as initial size @@ -363,8 +364,8 @@ DataRequestPool::expandPool() { requests_[i].operation_type = nixlLibfabricReq::WRITE; // Default for data } - NIXL_INFO << "Successfully expanded data request pool on rail " << rail_id_ << " to " - << requests_.size() << " requests"; + NIXL_DEBUG << "Successfully expanded data request pool on rail " << rail_id_ << " to " + << requests_.size() << " requests"; return NIXL_SUCCESS; } @@ -589,8 +590,8 @@ nixlLibfabricRail::nixlLibfabricRail(const std::string &device, << " data requests for rail " << rail_id; // Post initial pool of receives using new resource management system - NIXL_INFO << "Pre-posting " << NIXL_LIBFABRIC_RECV_POOL_SIZE << " recv requests for rail " - << rail_id; + NIXL_DEBUG << "Pre-posting " << NIXL_LIBFABRIC_RECV_POOL_SIZE << " recv requests for rail " + << rail_id; for (size_t i = 0; i < NIXL_LIBFABRIC_RECV_POOL_SIZE; ++i) { nixlLibfabricReq *recv_req = allocateControlRequest( @@ -609,8 +610,8 @@ nixlLibfabricRail::nixlLibfabricRail(const std::string &device, } } - NIXL_INFO << "Successfully pre-posted " << NIXL_LIBFABRIC_RECV_POOL_SIZE - << " recv requests for rail " << rail_id; + NIXL_DEBUG << "Successfully pre-posted " << NIXL_LIBFABRIC_RECV_POOL_SIZE + << " recv requests for rail " << rail_id; NIXL_TRACE << "Successfully initialized rail " << rail_id; } catch (...) { @@ -1067,8 +1068,8 @@ nixlLibfabricRail::postSend(uint64_t immediate_data, // Log every N attempts to avoid log spam if (attempt % NIXL_LIBFABRIC_LOG_INTERVAL_ATTEMPTS == 0) { - NIXL_INFO << "fi_senddata still retrying EAGAIN on rail " << rail_id << " after " - << attempt << " attempts"; + NIXL_DEBUG << "fi_senddata still retrying EAGAIN on rail " << rail_id << " after " + << attempt << " attempts"; } else { NIXL_TRACE << "fi_senddata returned EAGAIN on rail " << rail_id << ", retrying (attempt " << attempt << ")"; @@ -1151,8 +1152,8 @@ nixlLibfabricRail::postWrite(const void *local_buffer, // Log every N attempts to avoid log spam if (attempt % NIXL_LIBFABRIC_LOG_INTERVAL_ATTEMPTS == 0) { - NIXL_INFO << "fi_writedata still retrying EAGAIN on rail " << rail_id << " after " - << attempt << " attempts"; + NIXL_DEBUG << "fi_writedata still retrying EAGAIN on rail " << rail_id << " after " + << attempt << " attempts"; } else { NIXL_TRACE << "fi_writedata returned EAGAIN on rail " << rail_id << ", retrying (attempt " << attempt << ")"; @@ -1233,8 +1234,8 @@ nixlLibfabricRail::postRead(void *local_buffer, // Log every N attempts to avoid log spam if (attempt % NIXL_LIBFABRIC_LOG_INTERVAL_ATTEMPTS == 0) { - NIXL_INFO << "fi_read still retrying EAGAIN on rail " << rail_id << " after " - << attempt << " attempts"; + NIXL_DEBUG << "fi_read still retrying EAGAIN on rail " << rail_id << " after " + << attempt << " attempts"; } else { NIXL_TRACE << "fi_read returned EAGAIN on rail " << rail_id << ", retrying (attempt " << attempt << ")"; diff --git a/src/utils/libfabric/libfabric_rail_manager.cpp b/src/utils/libfabric/libfabric_rail_manager.cpp index 5b906dad..d1e3e016 100644 --- a/src/utils/libfabric/libfabric_rail_manager.cpp +++ b/src/utils/libfabric/libfabric_rail_manager.cpp @@ -229,8 +229,7 @@ nixlLibfabricRailManager::nixlLibfabricRailManager(size_t striping_threshold) throw std::runtime_error("Failed to create rails for libfabric rail manager"); } - NIXL_DEBUG << "Successfully created " << rails_.size() - << " rails using provider=" << selected_provider_name; + NIXL_INFO << "Created " << rails_.size() << " rails using provider=" << selected_provider_name; } nixlLibfabricRailManager::~nixlLibfabricRailManager() { @@ -273,7 +272,7 @@ nixlLibfabricRailManager::init(const nixl_b_params_t &custom_params) { // topology->getAllDevices().size() else if (max_rails < topology->getTotalNicCount()) { // bandwidth does not exceed total machine capacity, so use NUMA-aware rail selection policy - NIXL_TRACE << "Using NUMA-aware rail selection policy for DRAM memory type"; + NIXL_INFO << "DRAM rail selection policy: numa-aware (max_rails=" << max_rails << ")"; size_t numa_rail_count = topology->getNumaRailCount(); // NOTE: averaged if non-uniform if (max_rails > numa_rail_count) { size_t numa_speed = numa_rail_count * nic_speed; @@ -799,6 +798,9 @@ nixlLibfabricRailManager::registerMemory(void *buffer, << " (mr=" << static_cast(mr) << ", key=" << key << ")"; } + NIXL_INFO << "Registered memory on " << selected_rails.size() << " rails, mem_type=" << mem_type + << " device=" << device_id; + return NIXL_SUCCESS; } @@ -831,6 +833,8 @@ nixlLibfabricRailManager::deregisterMemory(const std::vector &selected_r } } + NIXL_INFO << "Deregistered memory from " << selected_rails.size() << " rails"; + return overall_status; } diff --git a/src/utils/libfabric/libfabric_topology.cpp b/src/utils/libfabric/libfabric_topology.cpp index b6de3c6e..b54fc6c8 100644 --- a/src/utils/libfabric/libfabric_topology.cpp +++ b/src/utils/libfabric/libfabric_topology.cpp @@ -136,7 +136,7 @@ nixlLibfabricTopology::discoverProviderWithDevices() { NIXL_INFO << "Discovered " << num_devices << " " << provider_name << " devices (TCP fallback)"; } else if (provider_name == "none" || all_devices.empty()) { - NIXL_WARN << "No network devices found"; + NIXL_ERROR << "No network devices found"; return NIXL_ERR_BACKEND; } @@ -251,6 +251,11 @@ nixlLibfabricTopology::getNumaRailCount() const { void nixlLibfabricTopology::printTopologyInfo() const { + NIXL_INFO << "Topology: " << num_numa_nodes << " NUMA nodes, " << num_devices << " NICs, " + << num_nvidia_accel << " NVIDIA GPUs, " << num_aws_accel << " AWS accelerators"; + if (avg_nic_speed > 0) { + NIXL_INFO << "Avg NIC bandwidth: " << avg_nic_speed << " Gbps"; + } NIXL_TRACE << "=== Libfabric Topology Information ==="; NIXL_TRACE << "Topology discovered: " << (topology_discovered ? "Yes" : "No"); NIXL_TRACE << "Number of AWS accelerators: " << num_aws_accel; @@ -263,13 +268,13 @@ nixlLibfabricTopology::printTopologyInfo() const { NIXL_TRACE << "Accelerator-PCI → EFA mapping:"; for (const auto &pair : pci_to_efa_devices) { std::stringstream ss; - ss << " Accelerator-PCI " << pair.first << " → ["; + ss << "Accelerator-PCI " << pair.first << " → ["; for (size_t i = 0; i < pair.second.size(); ++i) { if (i > 0) ss << ", "; ss << pair.second[i]; } ss << "]"; - NIXL_TRACE << ss.str(); + NIXL_INFO << ss.str(); } NIXL_TRACE << "Host memory (DRAM) will limit number of EFA devices used per-NUMA node " "according to maximum PCIe switch bandwidth"; From 018584c9d2c421a16a1b7499a7a47ccee32985d6 Mon Sep 17 00:00:00 2001 From: BatshevaBlack <132911331+BatshevaBlack@users.noreply.github.com> Date: Wed, 15 Apr 2026 14:58:22 +0300 Subject: [PATCH 34/59] Remove timestamp field from telemetry events (#1522) * Remove timestamp field from telemetry events The per-event timestamp (timestampUs_) added overhead on every event without being consumed by any exporter. Remove it from the struct, constructors, producers, readers, and bump TELEMETRY_VERSION to 2. Signed-off-by: Batsheva Black --- docs/telemetry.md | 5 +---- examples/cpp/telemetry_reader.cpp | 16 --------------- examples/python/telemetry_reader.py | 19 +++++------------- src/api/cpp/backend/backend_engine.h | 8 ++------ src/api/cpp/telemetry/telemetry_plugin.h | 10 +++++----- src/core/nixl_plugin_manager.cpp | 7 +++---- src/core/telemetry/buffer_plugin.cpp | 5 +++-- src/core/telemetry/telemetry.cpp | 20 ++++--------------- src/core/telemetry/telemetry_event.h | 14 +++++-------- src/plugins/telemetry/README.md | 9 ++++----- .../prometheus/prometheus_plugin.cpp | 4 ++-- test/gtest/telemetry_test.cpp | 4 +--- 12 files changed, 35 insertions(+), 86 deletions(-) diff --git a/docs/telemetry.md b/docs/telemetry.md index f231b807..668a9dda 100644 --- a/docs/telemetry.md +++ b/docs/telemetry.md @@ -18,7 +18,6 @@ Custom telemetry exporter plug-ins can be created according to [src/plugins/tele ### Event Structure Each telemetry event contains: -- **Timestamp**: Microsecond precision timestamp - **Category**: Event category for filtering and aggregation - **Event Name**: Descriptive name/identifier for the event - **Value**: Numeric value associated with the event @@ -61,7 +60,7 @@ The **Shared Memory Buffer** plug-in, contains the data per transaction event, w ### Telemetry Details -- Timestamp is done after the operation completion. +- Producers append events under a mutex, consumers read them in **ring insertion order**. - Current design allows silent telemetry loss. - Current design does not support selective telemetry(e.g per category). All the telemetry events could be either ON or OFF. @@ -128,14 +127,12 @@ Both readers produce similar formatted output: ``` === NIXL Telemetry Event === -Timestamp: 2025-01-15 14:30:25.123456 Category: TRANSFER Event: agent_tx_bytes Value: 1048576 =========================== === NIXL Telemetry Event === -Timestamp: 2025-01-15 14:30:25.124567 Category: MEMORY Event: agent_memory_registered Value: 4096 diff --git a/examples/cpp/telemetry_reader.cpp b/examples/cpp/telemetry_reader.cpp index 2274f2f3..ba745c31 100644 --- a/examples/cpp/telemetry_reader.cpp +++ b/examples/cpp/telemetry_reader.cpp @@ -44,21 +44,6 @@ signal_handler(int signal) { } } -std::string -format_timestamp(uint64_t timestamp_us) { - auto time_point = - std::chrono::system_clock::time_point(std::chrono::microseconds(timestamp_us)); - auto time_t = std::chrono::system_clock::to_time_t(time_point); - - std::stringstream ss; - ss << std::put_time(std::localtime(&time_t), "%Y-%m-%d %H:%M:%S"); - - auto microseconds = timestamp_us % 1000000; - ss << "." << std::setfill('0') << std::setw(6) << microseconds; - - return ss.str(); -} - std::string format_bytes(uint64_t bytes) { const char *units[] = {"B", "KB", "MB", "GB", "TB"}; @@ -80,7 +65,6 @@ print_telemetry_event(const nixlTelemetryEvent &event) { // Can be extended to more general ostream if needed // friend std::ostream &operator<<(std::ostream &os, const nixlTelemetryEvent &event) std::cout << "\n=== NIXL Telemetry Event ===" << std::endl; - std::cout << "Timestamp: " << format_timestamp(event.timestampUs_) << std::endl; std::cout << "Category: " << nixlEnumStrings::telemetryCategoryStr(event.category_) << std::endl; std::cout << "Event name: " << event.eventName_ << std::endl; diff --git a/examples/python/telemetry_reader.py b/examples/python/telemetry_reader.py index 4eb7e0fd..4733db1c 100755 --- a/examples/python/telemetry_reader.py +++ b/examples/python/telemetry_reader.py @@ -23,7 +23,6 @@ import signal import sys import time -from datetime import datetime # Set up logging logging.basicConfig( @@ -34,7 +33,7 @@ logger = logging.getLogger(__name__) # Constants from telemetry_event.h -TELEMETRY_VERSION = 1 +TELEMETRY_VERSION = 2 MAX_EVENT_NAME_LEN = 32 # NIXL telemetry categories @@ -65,7 +64,6 @@ class NixlTelemetryEvent(ctypes.Structure): _pack_ = 1 _fields_ = [ - ("timestamp_us", ctypes.c_uint64), ("category", ctypes.c_int), ("event_name", ctypes.c_char * MAX_EVENT_NAME_LEN), ("_padding", ctypes.c_uint32), @@ -90,7 +88,7 @@ class BufferHeader(ctypes.Structure): class SharedRingBuffer: """Python wrapper for the C++ SharedRingBuffer class""" - def __init__(self, file_path, version=1): + def __init__(self, file_path, version=TELEMETRY_VERSION): self.file_path = file_path self.version = version self.file_fd = -1 @@ -122,11 +120,12 @@ def _map_header_only(self): temp_header = BufferHeader.from_buffer(header_mmap) - if temp_header.version != self.version: + actual_version = temp_header.version + if actual_version != self.version: del temp_header header_mmap.close() raise RuntimeError( - f"Version mismatch: expected {self.version}, got {temp_header.version}" + f"Version mismatch: expected {self.version}, got {actual_version}" ) self.buffer_size = temp_header.capacity @@ -197,13 +196,6 @@ def __del__(self): os.close(self.file_fd) -def format_timestamp(timestamp_us): - """Format timestamp in microseconds to readable format""" - dt = datetime.fromtimestamp(timestamp_us / 1_000_000) - microseconds = timestamp_us % 1_000_000 - return f"{dt.strftime('%Y-%m-%d %H:%M:%S')}.{microseconds:06d}" - - def format_bytes(bytes_val): """Format bytes to human readable format""" units = ["B", "KB", "MB", "GB", "TB"] @@ -235,7 +227,6 @@ def get_telemetry_category_string(category): def print_telemetry_event(event): """Print telemetry event in a formatted way""" logger.info("\n=== NIXL Telemetry Event ===") - logger.info("Timestamp: %s", format_timestamp(event.timestamp_us)) # Decode event name event_name = event.event_name.decode("utf-8").rstrip("\x00") diff --git a/src/api/cpp/backend/backend_engine.h b/src/api/cpp/backend/backend_engine.h index d7a1cb72..b62d284d 100644 --- a/src/api/cpp/backend/backend_engine.h +++ b/src/api/cpp/backend/backend_engine.h @@ -66,12 +66,8 @@ class nixlBackendEngine { if (!enableTelemetry_) return; if (telemetryEvents_.size() >= MAX_TELEMETRY_QUEUE_SIZE) return; std::lock_guard lock(telemetryEventsMutex_); - telemetryEvents_.emplace_back(std::chrono::duration_cast( - std::chrono::system_clock::now().time_since_epoch()) - .count(), - nixl_telemetry_category_t::NIXL_TELEMETRY_BACKEND, - event_name, - value); + telemetryEvents_.emplace_back( + nixl_telemetry_category_t::NIXL_TELEMETRY_BACKEND, event_name, value); } public: diff --git a/src/api/cpp/telemetry/telemetry_plugin.h b/src/api/cpp/telemetry/telemetry_plugin.h index fa038d97..bffa1761 100644 --- a/src/api/cpp/telemetry/telemetry_plugin.h +++ b/src/api/cpp/telemetry/telemetry_plugin.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -24,10 +24,10 @@ #include #include -enum class nixl_telemetry_plugin_api_version : unsigned int { V1 = 1 }; - -inline constexpr nixl_telemetry_plugin_api_version nixlTelemetryPluginApiVersionV1 = - nixl_telemetry_plugin_api_version::V1; +enum class nixl_telemetry_plugin_api_version : unsigned int { + V1 = 1, + V2 = 2, +}; // Type alias for exporter creation function using exporter_creator_fn_t = diff --git a/src/core/nixl_plugin_manager.cpp b/src/core/nixl_plugin_manager.cpp index fc351d14..96c1ff69 100644 --- a/src/core/nixl_plugin_manager.cpp +++ b/src/core/nixl_plugin_manager.cpp @@ -186,11 +186,10 @@ telemetryLoader(void *handle, const std::string &plugin_path) { return nullptr; } - // Check API version - if (plugin->api_version != nixlTelemetryPluginApiVersionV1) { + if (plugin->api_version != nixl_telemetry_plugin_api_version::V2) { NIXL_ERROR << "Plugin API version mismatch for " << plugin_path << ": expected " - << static_cast(nixlTelemetryPluginApiVersionV1) << ", got " - << static_cast(plugin->api_version); + << static_cast(nixl_telemetry_plugin_api_version::V2) << ", got " + << static_cast(plugin->api_version); dlclose(handle); return nullptr; } diff --git a/src/core/telemetry/buffer_plugin.cpp b/src/core/telemetry/buffer_plugin.cpp index 76491c9d..2359d018 100644 --- a/src/core/telemetry/buffer_plugin.cpp +++ b/src/core/telemetry/buffer_plugin.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -24,5 +24,6 @@ using buffer_exporter_plugin_t = nixlTelemetryPluginCreator= maxBufferedEvents_) { return; } - events_.emplace_back(std::chrono::duration_cast( - std::chrono::system_clock::now().time_since_epoch()) - .count(), - category, - event_name, - value); + events_.emplace_back(category, event_name, value); } // The next 4 methods might be removed, as addXferTime covers them. @@ -232,22 +227,15 @@ nixlTelemetry::addXferTime(std::chrono::microseconds xfer_time, bool is_write, u const char *bytes_name = is_write ? "agent_tx_bytes" : "agent_rx_bytes"; const char *requests_name = is_write ? "agent_tx_requests_num" : "agent_rx_requests_num"; - const auto time = std::chrono::duration_cast( - std::chrono::system_clock::now().time_since_epoch()) - .count(); - const std::lock_guard lock(mutex_); if (events_.size() + 3 > maxBufferedEvents_) { return; } - events_.emplace_back(time, - nixl_telemetry_category_t::NIXL_TELEMETRY_PERFORMANCE, + events_.emplace_back(nixl_telemetry_category_t::NIXL_TELEMETRY_PERFORMANCE, "agent_xfer_time", xfer_time.count()); - events_.emplace_back( - time, nixl_telemetry_category_t::NIXL_TELEMETRY_TRANSFER, bytes_name, bytes); - events_.emplace_back( - time, nixl_telemetry_category_t::NIXL_TELEMETRY_TRANSFER, requests_name, 1); + events_.emplace_back(nixl_telemetry_category_t::NIXL_TELEMETRY_TRANSFER, bytes_name, bytes); + events_.emplace_back(nixl_telemetry_category_t::NIXL_TELEMETRY_TRANSFER, requests_name, 1); } void diff --git a/src/core/telemetry/telemetry_event.h b/src/core/telemetry/telemetry_event.h index 308a2679..e5a26c9c 100644 --- a/src/core/telemetry/telemetry_event.h +++ b/src/core/telemetry/telemetry_event.h @@ -27,7 +27,7 @@ constexpr char TELEMETRY_BUFFER_SIZE_VAR[] = "NIXL_TELEMETRY_BUFFER_SIZE"; constexpr char TELEMETRY_RUN_INTERVAL_VAR[] = "NIXL_TELEMETRY_RUN_INTERVAL"; -constexpr inline int TELEMETRY_VERSION = 1; +constexpr inline int TELEMETRY_VERSION = 2; constexpr inline size_t MAX_EVENT_NAME_LEN = 32; /** @@ -55,29 +55,25 @@ telemetryCategoryStr(const nixl_telemetry_category_t &category); * @brief A structure to hold individual telemetry event data for cyclic buffer storage */ struct nixlTelemetryEvent { - uint64_t timestampUs_; nixl_telemetry_category_t category_; // Main event category for filtering char eventName_[MAX_EVENT_NAME_LEN]; // Detailed event name/identifier uint64_t value_; // Numeric value associated with the event nixlTelemetryEvent() noexcept = default; - nixlTelemetryEvent(uint64_t timestamp_us, - nixl_telemetry_category_t category, + nixlTelemetryEvent(nixl_telemetry_category_t category, const char *event_name, uint64_t value) noexcept - : timestampUs_(timestamp_us), - category_(category), + : category_(category), value_(value) { strncpy(eventName_, event_name, MAX_EVENT_NAME_LEN - 1); eventName_[MAX_EVENT_NAME_LEN - 1] = '\0'; } - nixlTelemetryEvent(uint64_t timestamp_us, - nixl_telemetry_category_t category, + nixlTelemetryEvent(nixl_telemetry_category_t category, const std::string &event_name, uint64_t value) noexcept - : nixlTelemetryEvent(timestamp_us, category, event_name.c_str(), value) {} + : nixlTelemetryEvent(category, event_name.c_str(), value) {} }; #endif diff --git a/src/plugins/telemetry/README.md b/src/plugins/telemetry/README.md index d65c2761..f25ef6ec 100644 --- a/src/plugins/telemetry/README.md +++ b/src/plugins/telemetry/README.md @@ -1,5 +1,5 @@