From b5ecf06f6595c94b8874d1cea468238de4232dba Mon Sep 17 00:00:00 2001 From: wehub-resource-sync Date: Mon, 13 Jul 2026 12:24:32 +0800 Subject: [PATCH] chore: import upstream snapshot with attribution --- .clang-format | 21 + .editorconfig | 58 + .github/workflows/format.yml | 28 + .gitignore | 28 + .gitmodules | 3 + CMakeLists.txt | 52 + LICENSE | 21 + README.md | 449 ++++ README.wehub.md | 7 + build.sh | 12 + csrc/CMakeLists.txt | 1 + csrc/elastic/buffer.hpp | 1382 ++++++++++ csrc/elastic/utils.hpp | 28 + csrc/indexing/main.cu | 21 + csrc/jit/api.hpp | 20 + csrc/jit/cache.hpp | 34 + csrc/jit/compiler.hpp | 269 ++ csrc/jit/device_runtime.hpp | 67 + csrc/jit/handle.hpp | 155 ++ csrc/jit/include_parser.hpp | 80 + csrc/jit/kernel_runtime.hpp | 87 + csrc/jit/launch_runtime.hpp | 74 + csrc/kernels/CMakeLists.txt | 15 + csrc/kernels/backend/CMakeLists.txt | 6 + csrc/kernels/backend/api.cuh | 104 + csrc/kernels/backend/cuda_driver.cu | 54 + csrc/kernels/backend/nccl.cu | 156 ++ csrc/kernels/backend/nvshmem.cu | 87 + csrc/kernels/backend/symmetric.hpp | 319 +++ csrc/kernels/elastic/api.hpp | 9 + csrc/kernels/elastic/barrier.hpp | 86 + csrc/kernels/elastic/combine.hpp | 289 ++ csrc/kernels/elastic/dispatch.hpp | 342 +++ csrc/kernels/elastic/engram.hpp | 191 ++ csrc/kernels/elastic/pp_send_recv.hpp | 182 ++ csrc/kernels/legacy/CMakeLists.txt | 7 + csrc/kernels/legacy/api.cuh | 326 +++ csrc/kernels/legacy/buffer.cuh | 132 + csrc/kernels/legacy/compiled.cuh | 18 + csrc/kernels/legacy/ibgda_device.cuh | 496 ++++ csrc/kernels/legacy/internode.cu | 2384 +++++++++++++++++ csrc/kernels/legacy/internode_ll.cu | 1289 +++++++++ csrc/kernels/legacy/intranode.cu | 1117 ++++++++ csrc/kernels/legacy/launch.cuh | 134 + csrc/kernels/legacy/layout.cu | 154 ++ csrc/kernels/legacy/utils.cuh | 650 +++++ csrc/legacy/buffer.hpp | 1794 +++++++++++++ csrc/legacy/config.hpp | 190 ++ csrc/python_api.cpp | 39 + csrc/utils/event.hpp | 44 + csrc/utils/format.hpp | 6 + csrc/utils/hash.hpp | 40 + csrc/utils/lazy_driver.hpp | 57 + csrc/utils/lazy_init.hpp | 27 + csrc/utils/shared_memory.hpp | 127 + csrc/utils/system.hpp | 111 + deep_ep/__init__.py | 97 + deep_ep/buffers/__init__.py | 0 deep_ep/buffers/elastic.py | 1107 ++++++++ deep_ep/buffers/legacy.py | 713 +++++ deep_ep/include/deep_ep/common/comm.cuh | 266 ++ deep_ep/include/deep_ep/common/compiled.cuh | 87 + deep_ep/include/deep_ep/common/exception.cuh | 94 + deep_ep/include/deep_ep/common/handle.cuh | 230 ++ deep_ep/include/deep_ep/common/layout.cuh | 313 +++ deep_ep/include/deep_ep/common/math.cuh | 68 + deep_ep/include/deep_ep/common/ptx.cuh | 458 ++++ deep_ep/include/deep_ep/impls/barrier.cuh | 42 + deep_ep/include/deep_ep/impls/combine.cuh | 245 ++ .../deep_ep/impls/combine_reduce_epilogue.cuh | 145 + .../include/deep_ep/impls/combine_utils.cuh | 172 ++ deep_ep/include/deep_ep/impls/dispatch.cuh | 411 +++ .../deep_ep/impls/dispatch_copy_epilogue.cuh | 325 +++ .../include/deep_ep/impls/engram_fetch.cuh | 108 + .../deep_ep/impls/engram_fetch_wait.cuh | 31 + .../include/deep_ep/impls/hybrid_combine.cuh | 626 +++++ .../include/deep_ep/impls/hybrid_dispatch.cuh | 677 +++++ .../include/deep_ep/impls/pp_send_recv.cuh | 214 ++ deep_ep/utils/__init__.py | 3 + deep_ep/utils/comm.py | 83 + deep_ep/utils/envs.py | 268 ++ deep_ep/utils/event.py | 96 + deep_ep/utils/find_pkgs.py | 82 + deep_ep/utils/gate.py | 180 ++ deep_ep/utils/math.py | 103 + deep_ep/utils/refs.py | 243 ++ deep_ep/utils/semantic.py | 27 + deep_ep/utils/testing.py | 219 ++ develop.sh | 21 + docs/legacy.md | 320 +++ docs/nvshmem.md | 77 + figures/low-latency.png | Bin 0 -> 688315 bytes figures/normal.png | Bin 0 -> 518233 bytes format.sh | 194 ++ install.sh | 12 + pyproject.toml | 39 + requirements-lint.txt | 3 + setup.py | 218 ++ tests/elastic/test_agrs.py | 202 ++ tests/elastic/test_barrier.py | 62 + tests/elastic/test_engram.py | 124 + tests/elastic/test_ep.py | 609 +++++ tests/elastic/test_pp.py | 139 + tests/legacy/test_internode.py | 395 +++ tests/legacy/test_intranode.py | 311 +++ tests/legacy/test_low_latency.py | 332 +++ tests/utils/test_gate.py | 57 + 107 files changed, 24727 insertions(+) create mode 100644 .clang-format create mode 100644 .editorconfig create mode 100644 .github/workflows/format.yml create mode 100644 .gitignore create mode 100644 .gitmodules create mode 100644 CMakeLists.txt create mode 100644 LICENSE create mode 100644 README.md create mode 100644 README.wehub.md create mode 100755 build.sh create mode 100644 csrc/CMakeLists.txt create mode 100644 csrc/elastic/buffer.hpp create mode 100644 csrc/elastic/utils.hpp create mode 100644 csrc/indexing/main.cu create mode 100644 csrc/jit/api.hpp create mode 100644 csrc/jit/cache.hpp create mode 100644 csrc/jit/compiler.hpp create mode 100644 csrc/jit/device_runtime.hpp create mode 100644 csrc/jit/handle.hpp create mode 100644 csrc/jit/include_parser.hpp create mode 100644 csrc/jit/kernel_runtime.hpp create mode 100644 csrc/jit/launch_runtime.hpp create mode 100644 csrc/kernels/CMakeLists.txt create mode 100644 csrc/kernels/backend/CMakeLists.txt create mode 100644 csrc/kernels/backend/api.cuh create mode 100644 csrc/kernels/backend/cuda_driver.cu create mode 100644 csrc/kernels/backend/nccl.cu create mode 100644 csrc/kernels/backend/nvshmem.cu create mode 100644 csrc/kernels/backend/symmetric.hpp create mode 100644 csrc/kernels/elastic/api.hpp create mode 100644 csrc/kernels/elastic/barrier.hpp create mode 100644 csrc/kernels/elastic/combine.hpp create mode 100644 csrc/kernels/elastic/dispatch.hpp create mode 100644 csrc/kernels/elastic/engram.hpp create mode 100644 csrc/kernels/elastic/pp_send_recv.hpp create mode 100644 csrc/kernels/legacy/CMakeLists.txt create mode 100644 csrc/kernels/legacy/api.cuh create mode 100644 csrc/kernels/legacy/buffer.cuh create mode 100644 csrc/kernels/legacy/compiled.cuh create mode 100644 csrc/kernels/legacy/ibgda_device.cuh create mode 100644 csrc/kernels/legacy/internode.cu create mode 100644 csrc/kernels/legacy/internode_ll.cu create mode 100644 csrc/kernels/legacy/intranode.cu create mode 100644 csrc/kernels/legacy/launch.cuh create mode 100644 csrc/kernels/legacy/layout.cu create mode 100644 csrc/kernels/legacy/utils.cuh create mode 100644 csrc/legacy/buffer.hpp create mode 100644 csrc/legacy/config.hpp create mode 100644 csrc/python_api.cpp create mode 100644 csrc/utils/event.hpp create mode 100644 csrc/utils/format.hpp create mode 100644 csrc/utils/hash.hpp create mode 100644 csrc/utils/lazy_driver.hpp create mode 100644 csrc/utils/lazy_init.hpp create mode 100644 csrc/utils/shared_memory.hpp create mode 100644 csrc/utils/system.hpp create mode 100644 deep_ep/__init__.py create mode 100644 deep_ep/buffers/__init__.py create mode 100644 deep_ep/buffers/elastic.py create mode 100644 deep_ep/buffers/legacy.py create mode 100644 deep_ep/include/deep_ep/common/comm.cuh create mode 100644 deep_ep/include/deep_ep/common/compiled.cuh create mode 100644 deep_ep/include/deep_ep/common/exception.cuh create mode 100644 deep_ep/include/deep_ep/common/handle.cuh create mode 100644 deep_ep/include/deep_ep/common/layout.cuh create mode 100644 deep_ep/include/deep_ep/common/math.cuh create mode 100644 deep_ep/include/deep_ep/common/ptx.cuh create mode 100644 deep_ep/include/deep_ep/impls/barrier.cuh create mode 100644 deep_ep/include/deep_ep/impls/combine.cuh create mode 100644 deep_ep/include/deep_ep/impls/combine_reduce_epilogue.cuh create mode 100644 deep_ep/include/deep_ep/impls/combine_utils.cuh create mode 100644 deep_ep/include/deep_ep/impls/dispatch.cuh create mode 100644 deep_ep/include/deep_ep/impls/dispatch_copy_epilogue.cuh create mode 100644 deep_ep/include/deep_ep/impls/engram_fetch.cuh create mode 100644 deep_ep/include/deep_ep/impls/engram_fetch_wait.cuh create mode 100644 deep_ep/include/deep_ep/impls/hybrid_combine.cuh create mode 100644 deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh create mode 100644 deep_ep/include/deep_ep/impls/pp_send_recv.cuh create mode 100644 deep_ep/utils/__init__.py create mode 100644 deep_ep/utils/comm.py create mode 100644 deep_ep/utils/envs.py create mode 100644 deep_ep/utils/event.py create mode 100644 deep_ep/utils/find_pkgs.py create mode 100644 deep_ep/utils/gate.py create mode 100644 deep_ep/utils/math.py create mode 100644 deep_ep/utils/refs.py create mode 100644 deep_ep/utils/semantic.py create mode 100644 deep_ep/utils/testing.py create mode 100755 develop.sh create mode 100644 docs/legacy.md create mode 100644 docs/nvshmem.md create mode 100644 figures/low-latency.png create mode 100644 figures/normal.png create mode 100755 format.sh create mode 100755 install.sh create mode 100644 pyproject.toml create mode 100644 requirements-lint.txt create mode 100644 setup.py create mode 100644 tests/elastic/test_agrs.py create mode 100644 tests/elastic/test_barrier.py create mode 100644 tests/elastic/test_engram.py create mode 100644 tests/elastic/test_ep.py create mode 100644 tests/elastic/test_pp.py create mode 100644 tests/legacy/test_internode.py create mode 100644 tests/legacy/test_intranode.py create mode 100644 tests/legacy/test_low_latency.py create mode 100644 tests/utils/test_gate.py diff --git a/.clang-format b/.clang-format new file mode 100644 index 0000000..d6b4037 --- /dev/null +++ b/.clang-format @@ -0,0 +1,21 @@ +BasedOnStyle: Google +UseTab: Never +IndentWidth: 4 +ColumnLimit: 140 +AccessModifierOffset: -4 + +# Force pointers to the type for C++. +DerivePointerAlignment: false +PointerAlignment: Left +ReferenceAlignment: Left +AllowShortFunctionsOnASingleLine: Inline +AllowShortIfStatementsOnASingleLine: false +AllowShortLoopsOnASingleLine: false +AlignOperands: false +BreakBeforeBinaryOperators: None +Cpp11BracedListStyle: true +ContinuationIndentWidth: 4 + +BinPackArguments: false +BinPackParameters: false + diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..7b75c68 --- /dev/null +++ b/.editorconfig @@ -0,0 +1,58 @@ +# https://editorconfig.org/ + +root = true + +[*] +charset = utf-8 +end_of_line = lf +indent_style = space +indent_size = 4 +trim_trailing_whitespace = true +insert_final_newline = true + +[*.{py,pyi}] +indent_size = 4 + +[*.{cpp,hpp,cxx,cc,c,h,cu,cuh}] +indent_size = 4 + +[*.rs] +indent_size = 4 + +[*.go] +indent_style = tab + +[*.{yaml,yml}] +indent_size = 2 + +[.clang-{format,tidy}] +indent_size = 2 + +[Makefile] +indent_style = tab + +[*.sh] +indent_size = 4 + +[*.bat] +indent_size = 4 +end_of_line = crlf + +[*.md] +indent_size = 2 +x-soft-wrap-text = true + +[*.rst] +indent_size = 4 +x-soft-wrap-text = true + +[*.{html,xml,css,scss,js,jsx,ts,tsx,vue}] +indent_size = 2 + +[**/test{,s,ing}/**/*.txt] +trim_trailing_whitespace = false +insert_final_newline = false + +[**/example{,s}/**/*.txt] +trim_trailing_whitespace = false +insert_final_newline = false diff --git a/.github/workflows/format.yml b/.github/workflows/format.yml new file mode 100644 index 0000000..95dc015 --- /dev/null +++ b/.github/workflows/format.yml @@ -0,0 +1,28 @@ +name: Code Format Check + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + +jobs: + format-check: + runs-on: ubuntu-latest + + steps: + - name: Checkout source + uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - name: Setup environment + run: | + sudo apt-get update + sudo apt-get install -y bash + + - name: Run format.sh + run: | + bash ./format.sh + + # If format.sh return non-zero, GitHub Actions will mark it as failure. \ No newline at end of file diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9f2fd31 --- /dev/null +++ b/.gitignore @@ -0,0 +1,28 @@ +# VS Code +compile_commands.json +.clangd +.cache/ +.vscode/ + +# macOS +.DS_Store + +# JetBrains +.idea +*/cmake-build-*/ + +# Python +*.pyc + +# Build +build/ +*.so +deep_ep.egg-info/ +dist/ + +# Coredump +coredump/ + +# OpenCode +.sisyphus/ +AGENTS.md diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..4c15fff --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "third-party/fmt"] + path = third-party/fmt + url = https://github.com/fmtlib/fmt.git diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..02e1ed3 --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,52 @@ +# NOTES: this CMake is only for debugging; for setup, please use Torch extension +cmake_minimum_required(VERSION 3.10) +project(deep_ep LANGUAGES CUDA CXX) +set(CMAKE_VERBOSE_MAKEFILE ON) + +set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -O3 -fPIC") +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC") +set(CUDA_SEPARABLE_COMPILATION ON) +list(APPEND CUDA_NVCC_FLAGS "-DENABLE_FAST_DEBUG") +list(APPEND CUDA_NVCC_FLAGS "-O3") +list(APPEND CUDA_NVCC_FLAGS "--ptxas-options=--verbose,--register-usage-level=10,--warn-on-local-memory-usage") +list(APPEND CUDA_NVCC_FLAGS "-Xcompiler=-rdynamic") +list(APPEND CUDA_NVCC_FLAGS "-lineinfo") +# Suppress warnings for the `fmt` library +list(APPEND CUDA_NVCC_FLAGS "--diag-suppress=128,2417") + +set(USE_SYSTEM_NVTX on) +set(CUDA_ARCH_LIST "9.0" CACHE STRING "List of CUDA architectures to compile") +set(TORCH_CUDA_ARCH_LIST "${CUDA_ARCH_LIST}") + +find_package(CUDAToolkit REQUIRED) +find_package(pybind11 REQUIRED) +find_package(Torch REQUIRED) + +# NVSHMEM +find_package(NVSHMEM REQUIRED HINTS ${NVSHMEM_ROOT_DIR}/lib/cmake/nvshmem) +add_library(nvshmem ALIAS nvshmem::nvshmem) +add_library(nvshmem_host ALIAS nvshmem::nvshmem_host) +add_library(nvshmem_device ALIAS nvshmem::nvshmem_device) + +# NCCL +# TODO: use `find_package` instead of manual checks +if (NOT NCCL_ROOT_DIR) + message(FATAL_ERROR "NCCL_ROOT_DIR is not set.") +endif() + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CUDA_STANDARD 20) + +include_directories(deep_ep/include third-party/fmt/include) +include_directories(${CUDA_TOOLKIT_ROOT_DIR}/include ${CUDA_TOOLKIT_ROOT_DIR}/include/cccl ${TORCH_INCLUDE_DIRS} ${PYTHON_INCLUDE_DIRS} ${NVSHMEM_INCLUDE_DIR} ${NCCL_ROOT_DIR}/include .) +link_directories(${TORCH_INSTALL_PREFIX}/lib ${CUDA_TOOLKIT_ROOT_DIR}/lib ${NVSHMEM_LIB_DIR} ${NCCL_ROOT_DIR}/lib) + +# Add kernels +add_subdirectory(csrc) + +# Link CPP and CUDA together +pybind11_add_module(_C csrc/python_api.cpp) +target_link_libraries(_C PRIVATE nccl ${RUNTIME_CUDA_LIBRARIES} ${LEGACY_CUDA_LIBRARIES} ${TORCH_LIBRARIES} torch_python) + +# Enable kernel code indexing with CMake-based IDEs +cuda_add_library(deep_ep_indexing_cuda STATIC csrc/indexing/main.cu) diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..5c48bdc --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 DeepSeek + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..078a151 --- /dev/null +++ b/README.md @@ -0,0 +1,449 @@ +# DeepEP + +DeepEP (DeepEveryParallel) is a high-performance communication library for modern machine learning training and inference. The library currently focuses on expert parallelism (EP) — providing high-throughput and low-latency all-to-all GPU kernels (MoE dispatch and combine) with low-precision support including FP8 — while also offering experimental primitives for pipeline parallelism (PP), context parallelism (CP), and remote memory access (Engram), all designed for zero or minimal SM occupation. All kernels are compiled at runtime via a lightweight Just-In-Time (JIT) module, requiring no CUDA compilation during installation. + +Despite its lightweight design, DeepEP's performance matches or exceeds hardware bandwidth limits across various configurations. + +## News + +- **V2 release**: A complete refactoring of Expert Parallelism — achieving extreme performance with several times fewer SM resources compared to V1, while supporting significantly larger scale-up and scale-out domains. V2 has also switched from the NVSHMEM backend to the more lightweight **NCCL Gin backend**. + +### New features + +- **Fully JIT** (Just-In-Time compilation) +- **NCCL Gin backend** + - Header-only & lightweight + - Able to reuse existing NCCL communicators +- **EPv2** + - High-throughput and low-latency APIs unified into a single `ElasticBuffer` interface, with a new GEMM layout + - Larger scale-up & scale-out domain support (up to EP2048) + - Analytical SM & QP count calculation — no more auto-tuning needed + - Both hybrid & direct modes remain supported + - For V3-like legacy training, SM usage reduced from 24 to 4 - 6 while maintaining equivalent or better performance +- **0 SM Engram** (with RDMA) +- **0 SM PP** (with RDMA) +- **0 SM CP** (with Copy Engine) + +### Notes + +- Buffer size consumption is larger than V1 +- 0 SM RDMA low-latency EP is no longer supported +- Engram, PP, and CP are experimental features + +### Still on-going features + +- **Elastic GPU & CPU buffers**: A contiguous virtual address space that maps to a hybrid of GPU and CPU physical memory under the hood, enabling fully automatic and transparent Engram or imbalanced EP +- Reducing intermediate buffer sizes by leveraging EP replay to handle load imbalance +- All-gather updates and reduce-scatter implementations for DP & TP + +For the legacy V1 documentation (NVSHMEM-based), see [docs/legacy.md](docs/legacy.md). + +## Performance + +Following V3's configuration, we tested with 8K tokens per batch, 7168 hidden dimensions, top 8 experts, FP8 dispatching, and BF16 combining, and obtained the following results: + +| Arch | NIC type | Topo | Dispatch Bottleneck Bandwidth | Combine Bottleneck Bandwidth | #SMs | +|--|--|--|--|--|--| +| SM90 | CX7 | EP 8 x 2 | 90 GB/s (RDMA) | 81 GB/s (RDMA) | 12 | +| SM90 | CX7 | EP 8 x 4 | 61 GB/s (RDMA) | 61 GB/s (RDMA) | 6 | +| SM100 | CX7 | EP 8 x 2 | 90 GB/s (RDMA) | 91 GB/s (RDMA) | 12 | +| SM100 | N/A | EP 8 | 726 GB/s (NVLink) | 740 GB/s (NVLink) | 64 (Max perf) | +| SM100 | N/A | EP 8 | 643 GB/s (NVLink) | 675 GB/s (NVLink) | 24 (Min #SM) | + +Notes: the results are logical bandwidth. For example, under the `EP 8 x 2` case, 90 GB/s actually contains local rank traffic. + +Comparing with V1, **V2 achieves up to 1.3x peak performance, while saving up to 4x SM count**. + +We omit results for larger EP configurations for the time being, but encourage interested users to benchmark them directly. Based on our internal experience, we expect the kernel to continue saturating hardware bandwidth at scale. + +For V1 performance data, see [docs/legacy.md](docs/legacy.md#performance). + +## Quick start + +### Requirements + +- Hopper (SM90) GPUs, or other architectures with SM90 PTX ISA support +- Python 3.8 and above +- CUDA version + - CUDA 12.3 and above for SM90 GPUs +- PyTorch 2.10 and above +- NCCL 2.30.4 and above +- NVLink for intranode communication +- RDMA network for internode communication + +### Install NCCL dependency + +We recommend using pip to install NCCL so that DeepEP can automatically locate it within the Python environment. You can install it using the following command: + +```bash +pip install "nvidia-nccl-cu13>=2.30.4" --no-deps +``` + +### Install NVSHMEM dependency + +DeepEP also depends on NVSHMEM to provide support for legacy methods. Please refer to our [NVSHMEM Installation Guide](docs/nvshmem.md) for instructions. + +### Development + +```bash +# Build and make symbolic links for SO files +python setup.py build +# You may modify the specific SO names according to your own platform +ln -s build/lib.linux-x86_64-cpython-38/deep_ep_cpp.cpython-38-x86_64-linux-gnu.so + +# Run test cases +# NOTES: you may modify the `init_dist` function in `tests/utils/envs.py` +# according to your own cluster settings, and launch into multiple nodes +python tests/elastic/test_ep.py +python tests/elastic/test_agrs.py +python tests/elastic/test_engram.py +python tests/elastic/test_pp.py +``` + +### Installation + +```bash +python setup.py install +``` + +Then, import `deep_ep` in your Python project, and enjoy! + +## Interfaces and examples + +### Buffer initialization + +In V2, all EP operations — high-throughput and low-latency — are unified under a single `ElasticBuffer` interface. The buffer can be initialized by specifying MoE settings directly, and the optimal SM and QP counts are calculated analytically. + +```python +import torch +import torch.distributed as dist +from typing import Optional + +from deep_ep import ElasticBuffer + +# Communication buffer (will allocate at runtime) +_buffer: Optional[ElasticBuffer] = None + +# Number of SMs to use for communication kernels (will be set at buffer creation) +_num_comm_sms: int = 0 + + +def get_buffer(group: dist.ProcessGroup, + num_max_tokens_per_rank: int, + hidden: int, + num_topk: int, + num_experts: int, + use_fp8_dispatch: bool = False) -> ElasticBuffer: + """Initialize or retrieve the ElasticBuffer for EP communication.""" + global _buffer, _num_comm_sms + + # Check if we can reuse the existing buffer + required_bytes = ElasticBuffer.get_buffer_size_hint( + group, num_max_tokens_per_rank, hidden, + num_topk=num_topk, use_fp8_dispatch=use_fp8_dispatch, + ) + if _buffer is not None and _buffer.group == group and _buffer.num_bytes >= required_bytes: + return _buffer + + # Allocate a new buffer with MoE settings + # NOTES: V2 buffer size consumption is larger than V1 + _buffer = ElasticBuffer( + group, + num_max_tokens_per_rank=num_max_tokens_per_rank, + hidden=hidden, + num_topk=num_topk, + use_fp8_dispatch=use_fp8_dispatch, + ) + + # V2 analytically calculates the optimal SM count — no more auto-tuning needed + # You may also specify `num_sms` manually in dispatch/combine calls to override + _num_comm_sms = _buffer.get_theoretical_num_sms(num_experts, num_topk) + + return _buffer +``` + +### Example use in model training or inference prefilling + +V2 unifies the dispatch and combine APIs into a single `ElasticBuffer` interface. The example below shows how to use them for training (with backward passes) or inference prefilling. + +```python +import torch +import torch.distributed as dist +from typing import Tuple, Union + +from deep_ep import ElasticBuffer, EPHandle, EventOverlap + + +def dispatch_forward(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + topk_idx: torch.Tensor, topk_weights: torch.Tensor, + num_experts: int, + num_max_tokens_per_rank: int, + expert_alignment: int = 1) -> \ + Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + torch.Tensor, torch.Tensor, EPHandle, EventOverlap]: + """ + MoE dispatch: route tokens to the corresponding experts across all ranks. + Supports both BF16 and FP8 (x as a tuple of [data, scale_factors]) inputs. + """ + global _buffer, _num_comm_sms + + recv_x, recv_topk_idx, recv_topk_weights, handle, event = _buffer.dispatch( + x, + topk_idx=topk_idx, + topk_weights=topk_weights, + num_experts=num_experts, + num_max_tokens_per_rank=num_max_tokens_per_rank, + expert_alignment=expert_alignment, + num_sms=_num_comm_sms, + async_with_compute_stream=True, + ) + + # `handle` contains routing metadata for the subsequent combine call + # `handle.num_recv_tokens_per_expert_list` provides per-expert token counts for GEMM + # Use `event.current_stream_wait()` to synchronize the compute stream before using results + return recv_x, recv_topk_idx, recv_topk_weights, handle, event + + +def dispatch_backward(grad_recv_x: torch.Tensor, + grad_recv_topk_weights: torch.Tensor, + handle: EPHandle) -> Tuple[torch.Tensor, torch.Tensor, EventOverlap]: + """The backward pass of MoE dispatch is actually a combine.""" + global _buffer, _num_comm_sms + + combined_grad_x, combined_grad_topk_weights, event = _buffer.combine( + grad_recv_x, + handle=handle, + topk_weights=grad_recv_topk_weights, + num_sms=_num_comm_sms, + async_with_compute_stream=True, + ) + + return combined_grad_x, combined_grad_topk_weights, event + + +def combine_forward(x: torch.Tensor, + handle: EPHandle) -> Tuple[torch.Tensor, EventOverlap]: + """MoE combine: reduce expert outputs back to their original ranks.""" + global _buffer, _num_comm_sms + + combined_x, _, event = _buffer.combine( + x, + handle=handle, + num_sms=_num_comm_sms, + async_with_compute_stream=True, + ) + + return combined_x, event + + +def combine_backward(grad_combined_x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + handle: EPHandle) -> \ + Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], EventOverlap]: + """The backward pass of MoE combine is actually a dispatch.""" + global _buffer, _num_comm_sms + + grad_x, _, _, _, event = _buffer.dispatch( + grad_combined_x, + handle=handle, + num_sms=_num_comm_sms, + async_with_compute_stream=True, + ) + + return grad_x, event +``` + +For communication-computation overlap, use the `EventOverlap` interface to manage dependencies between the communication stream and the compute stream: + +```python +# After dispatch, overlap computation while communication is in-flight +recv_x, recv_topk_idx, recv_topk_weights, handle, event = dispatch_forward(...) + +# ... do some independent computation here ... + +# Wait for communication to finish before using results +event.current_stream_wait() + +# Now safe to use recv_x, recv_topk_idx, recv_topk_weights +``` + +### Example use in inference decoding + +For inference decoding, the same `ElasticBuffer` is used. The handle-caching pattern allows reusing routing metadata across iterations when the gating decisions remain unchanged, avoiding redundant CPU synchronization. + +```python +import torch +from typing import Tuple, Optional, Union + +from deep_ep import ElasticBuffer, EPHandle, EventOverlap + + +def decode_dispatch(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + topk_idx: torch.Tensor, topk_weights: torch.Tensor, + num_experts: int, + num_max_tokens_per_rank: int, + cached_handle: Optional[EPHandle] = None) -> \ + Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + torch.Tensor, torch.Tensor, EPHandle, EventOverlap]: + """ + MoE dispatch for inference decoding. + If `cached_handle` is provided, the layout is reused without CPU synchronization. + """ + global _buffer, _num_comm_sms + + if cached_handle is not None: + # Reuse cached handle: skip layout recomputation and CPU sync + recv_x, _, _, handle, event = _buffer.dispatch( + x, + handle=cached_handle, + num_sms=_num_comm_sms, + async_with_compute_stream=True, + ) + return recv_x, cached_handle.topk_idx, None, handle, event + + recv_x, recv_topk_idx, recv_topk_weights, handle, event = _buffer.dispatch( + x, + topk_idx=topk_idx, + topk_weights=topk_weights, + num_experts=num_experts, + num_max_tokens_per_rank=num_max_tokens_per_rank, + num_sms=_num_comm_sms, + async_with_compute_stream=True, + ) + + return recv_x, recv_topk_idx, recv_topk_weights, handle, event + + +def decode_combine(x: torch.Tensor, + handle: EPHandle) -> Tuple[torch.Tensor, EventOverlap]: + """MoE combine for inference decoding.""" + global _buffer, _num_comm_sms + + combined_x, _, event = _buffer.combine( + x, + handle=handle, + num_sms=_num_comm_sms, + async_with_compute_stream=True, + ) + + return combined_x, event +``` + +### Environment variables + +The library provides some environment variables, which may be useful: + +- General + - `EP_BUFFER_DEBUG`: `0` or `1`, print buffer initialization, SM approximation, and backend debugging information, `0` by default + - `EP_SUPPRESS_NCCL_CHECK`: `0` or `1`, suppress NCCL version mismatch checking, `0` by default + - `EP_AVOID_RECORD_STREAM`: `0` or `1`, avoid `record_stream` on output tensors, `0` by default + - `EP_NUM_TOPK_IDX_BITS`: integer, override the number of bits for top-k index encoding, `0` (auto) by default +- Networking + - `EP_NIC_NAME`: string, the default NIC name used to query NIC properties, `mlx5_0` by default + - `EP_OVERRIDE_RDMA_SL`: integer, override the RDMA service level index for traffic isolation + - `EP_DISABLE_GIN`: `0` or `1`, disable the NCCL Gin backend (fall back to non-Gin path), `0` by default +- JIT + - `EP_JIT_DEBUG`: `0` or `1`, print JIT debugging information, `0` by default + - `EP_JIT_CACHE_DIR`: string, cache directory for compiled kernels, `$HOME/.deep_ep` by default + - `EP_JIT_NVCC_COMPILER`: string, NVCC compiler path; defaults to `torch.utils.cpp_extension.CUDA_HOME` + - `EP_JIT_CPP_STANDARD`: integer, C++ standard version, `20` by default + - `EP_JIT_PRINT_COMPILER_COMMAND`: `0` or `1`, print compilation commands, `0` by default + - `EP_JIT_PTXAS_VERBOSE`: `0` or `1`, show detailed PTXAS output, `0` by default + - `EP_JIT_PTXAS_CHECK`: `0` or `1`, assert no local memory usage in compiled kernels, `0` by default + - `EP_JIT_WITH_LINEINFO`: `0` or `1`, embed source line info for profiling tools, `0` by default + - `EP_JIT_DUMP_ASM`: `0` or `1`, dump both PTX and SASS, `0` by default + - `EP_JIT_DUMP_PTX`: `0` or `1`, dump PTX output, `0` by default + - `EP_JIT_DUMP_SASS`: `0` or `1`, dump SASS output, `0` by default +- Debug and profiling + - `EP_GIN_GDAKI_DEBUG`: `0` or `1`, enable NCCL Gin GDAKI debugging output, `0` by default + - `EP_USE_NVIDIA_TOOLS`: `0` or `1`, skip internal profiling when running under external NVIDIA tools, `0` by default + - `EP_DISABLE_BARRIER_PROFILING`: `0` or `1`, disable barrier-based communication profiling in benchmarks, `0` by default +- Build + - `EP_NCCL_ROOT_DIR`: string, path to the NCCL installation directory; auto-detected from the Python environment if not set + - `EP_NVSHMEM_ROOT_DIR`: string, path to the NVSHMEM installation directory; auto-detected from the Python environment if not set + - `TORCH_CUDA_ARCH_LIST`: string, list of target CUDA architectures, e.g. `"9.0"` + - `DISABLE_SM90_FEATURES`: `0` or `1`, disable SM90 features for legacy methods, `0` by default + - `DISABLE_AGGRESSIVE_PTX_INSTRS`: `0` or `1`, disable aggressive load/store instructions in legacy methods, `0` by default + +Some environment variables are **persistent**: they are captured at build time and baked into the installed package as default values. At import time, these defaults are applied automatically unless overridden by current environment variables. The persistent variables are: `EP_JIT_CACHE_DIR`, `EP_JIT_PRINT_COMPILER_COMMAND`, `EP_NUM_TOPK_IDX_BITS`, `EP_NCCL_ROOT_DIR`. + +For additional details, please refer to [the test code](tests/elastic/test_ep.py) or review the corresponding Python documentation. + +## Network configurations + +DeepEP is fully tested with InfiniBand networks. However, it is theoretically compatible with RDMA over Converged Ethernet (RoCE) as well. + +### Traffic isolation + +Traffic isolation is supported by InfiniBand through Virtual Lanes (VL). + +To prevent interference between different types of traffic, we recommend segregating workloads across different virtual lanes as follows: + +- expert-parallel workloads +- other workloads + +For DeepEP V2, you can control the virtual lane assignment by setting the `sl_idx` argument or the `EP_OVERRIDE_RDMA_SL` environment variable. + +### Adaptive routing + +Adaptive routing is an advanced routing feature provided by InfiniBand switches that can evenly distribute traffic across multiple paths. Even though adaptive routing introduces additional latency, we still recommend enabling it under all network load conditions. + +### Congestion control + +Congestion control is disabled because it hurts maximum bandwidth. If congestion is unavoidable in some scenarios, we recommend assigning those workloads to low-priority virtual lanes. + +### PCI atomic mode + +If the hardware supports it, we recommend using the following command to set the NIC's `PCI_ATOMIC_MODE` to improve RDMA atomic operation performance: + +```bash +sudo mlxconfig -y -d mlx5_$i set PCI_ATOMIC_MODE=4 +``` + +## Experimental branches + +- [Zero-copy](https://github.com/deepseek-ai/DeepEP/pull/453) + - Removing the copy between PyTorch tensors and communication buffers, which reduces the SM usages significantly for normal kernels + - This PR is authored by **Tencent Network Platform Department** +- [Eager](https://github.com/deepseek-ai/DeepEP/pull/437) + - Using a low-latency protocol removes the extra RTT latency introduced by RDMA atomic OPs +- [Hybrid-EP](https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep) + - A new backend implementation using TMA instructions for minimal SM usage and larger NVLink domain support + - Fine-grained communication-computation overlap for single-batch scenarios + - PCIe kernel support for non-NVLink environments + - NVFP4 data type support +- [AntGroup-Opt](https://github.com/deepseek-ai/DeepEP/tree/antgroup-opt) + - This optimization series is authored by **AntGroup Network Platform Department** + - [Normal-SMFree](https://github.com/deepseek-ai/DeepEP/pull/347) Eliminating SM from RDMA path by decoupling comm-kernel execution from NIC token transfer, freeing SMs for compute + - [LL-SBO](https://github.com/deepseek-ai/DeepEP/pull/483) Overlapping Down GEMM computation with Combine Send communication via signaling mechanism to reduce end-to-end latency + - [LL-Layered](https://github.com/deepseek-ai/DeepEP/pull/500) Optimizing cross-node LL operator communication using rail-optimized forwarding and data merging to reduce latency +- [Mori-EP](https://github.com/deepseek-ai/DeepEP/tree/mori-ep) + - ROCm/AMD GPU support powered by [MORI](https://github.com/ROCm/mori) backend (low-latency mode) +- [nvDev](https://github.com/deepseek-ai/DeepEP/tree/nvDev) + - V2-based branch with the latest CUDA features, such as Compute Fabric Transport (CFT) that brings better latency on small token sizes. + +## Community forks + +- [uccl/uccl-ep](https://github.com/uccl-project/uccl/tree/main/ep) - Enables running DeepEP on heterogeneous GPUs (e.g., Nvidia, AMD) and NICs (e.g., EFA, Broadcom, CX7) +- [Infrawaves/DeepEP_ibrc_dual-ports_multiQP](https://github.com/Infrawaves/DeepEP_ibrc_dual-ports_multiQP) - Adds multi-QP solution and dual-port NIC support in IBRC transport +- [antgroup/DeepXTrace](https://github.com/antgroup/DeepXTrace) - A diagnostic analyzer for efficient and precise localization of slow ranks +- [ROCm/mori](https://github.com/ROCm/mori) - AMD's next-generation communication library for performance-critical AI workloads (e.g., Wide EP, KVCache transfer, Collectives) + +## Acknowledgement + +DeepEP V2 is built on top of the [NCCL](https://github.com/nvidia/nccl) Gin backend. Thanks to @sjeaugey, @pakmarkthub, @sb17v, @xiaofanl-nvidia, and the NCCL team for their support! + +## License + +This code repository is released under [the MIT License](LICENSE). + +## Citation + +```bibtex +@misc{deepep2025, + title={DeepEP: an efficient expert-parallel communication library}, + author={Chenggang Zhao and Shangyan Zhou and Liyue Zhang and Chengqi Deng and Zhean Xu and Yuxuan Liu and Kuai Yu and Jiashi Li and Liang Zhao}, + year={2025}, + publisher = {GitHub}, + howpublished = {\url{https://github.com/deepseek-ai/DeepEP}}, +} +``` diff --git a/README.wehub.md b/README.wehub.md new file mode 100644 index 0000000..ca8c3a8 --- /dev/null +++ b/README.wehub.md @@ -0,0 +1,7 @@ +# WeHub 来源说明 + +- 原始项目：`deepseek-ai/DeepEP` +- 原始仓库：https://github.com/deepseek-ai/DeepEP +- 导入方式：上游默认分支的最新快照 +- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准 +- 本文件仅用于记录来源，不代表 WeHub 是原项目作者 diff --git a/build.sh b/build.sh new file mode 100755 index 0000000..abdfc40 --- /dev/null +++ b/build.sh @@ -0,0 +1,12 @@ +# Change current directory into project root +original_dir=$(pwd) +script_dir=$(realpath "$(dirname "$0")") +cd "$script_dir" + +# Remove old dist file, build files, and install +rm -rf build dist +rm -rf *.egg-info +python setup.py bdist_wheel + +# Open users' original directory +cd "$original_dir" diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt new file mode 100644 index 0000000..94e1eba --- /dev/null +++ b/csrc/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(kernels) diff --git a/csrc/elastic/buffer.hpp b/csrc/elastic/buffer.hpp new file mode 100644 index 0000000..0700646 --- /dev/null +++ b/csrc/elastic/buffer.hpp @@ -0,0 +1,1382 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include +#include + +#include "../kernels/backend/api.cuh" +#include "../kernels/elastic/api.hpp" +#include "../utils/event.hpp" +#include "utils.hpp" + +namespace deep_ep::elastic { + +class ElasticBuffer { + // Buffer bytes = GPU buffer + CPU buffer (excludes workspace) + // Memory layout: [[[Workspace] GPU buffer] CPU buffer] + int64_t num_buffer_bytes; + int64_t num_gpu_buffer_bytes; + int64_t num_cpu_buffer_bytes; + void* buffer; + + // Destructor settings + bool explicitly_destroy; + bool destroyed = false; + + // Workspace + // NOTES: for all workspace, we must keep them as zeros + void *workspace; + void *host_workspace, *mapped_host_workspace; + std::shared_ptr workspace_layout_wo_expert; + + // CUDA streams + at::cuda::CUDAStream comm_stream; + + // Whether to use hybrid mode (scale-out with scale-up) + bool allow_hybrid_mode; + + // Whether to allow multiple reductions + bool allow_multiple_reduction; + + // Whether to prefer overlapping communication with compute (use more SMs and channels if false) + bool prefer_overlap_with_compute; + + // Timeout settings + int num_cpu_timeout_secs; + int64_t num_gpu_timeout_cycles; + + // NCCL context + std::shared_ptr nccl_context; + + // Some EP hybrid mode settings + static constexpr int kNumMaxChannelsPerSM = 8; + static constexpr int kNumMaxSMs = 160; + static constexpr int kNumMaxChannels = kNumMaxChannelsPerSM * kNumMaxSMs; + + // Some Engram storage settings + int num_engram_entries = 0, engram_hidden = 0; + std::optional engram_sf; + + // PP settings + int prev_rank_idx = 0, next_rank_idx = 0; + int64_t num_max_pp_tensor_bytes = 0; + int num_max_pp_inflight_tensors = 0; + + // AGRS session settings + int64_t num_max_agrs_session_bytes = 0; + int num_max_agrs_per_session = 0; + int agrs_session_idx = 0; + bool agrs_in_session = false; + + // AGRS in-session settings + int64_t agrs_buffer_offset = 0; + int agrs_buffer_slot_idx = 0; + +public: + ElasticBuffer(const int& rank_idx, const int& num_ranks, + const int64_t& nccl_comm, const symmetric::cpu_comm_t& cpu_comm, + const int64_t& num_buffer_bytes, const int64_t& num_cpu_buffer_bytes, + const bool& allow_hybrid_mode, + const bool& allow_multiple_reduction, + const bool& prefer_overlap_with_compute, + const int& sl_idx, const int& num_allocated_qps, + const int& num_cpu_timeout_secs, const int& num_gpu_timeout_secs, + const bool& explicitly_destroy): + num_buffer_bytes(num_buffer_bytes), + num_cpu_buffer_bytes(num_cpu_buffer_bytes), + explicitly_destroy(explicitly_destroy), + comm_stream(get_global_comm_stream()), + allow_hybrid_mode(allow_hybrid_mode), + allow_multiple_reduction(allow_multiple_reduction), + prefer_overlap_with_compute(prefer_overlap_with_compute) { + // Check buffer bytes alignment (2 MB) + EP_HOST_ASSERT(num_buffer_bytes > 0 and num_buffer_bytes % symmetric::kNumAlignmentBytes == 0); + EP_HOST_ASSERT(num_cpu_buffer_bytes >= 0 and num_cpu_buffer_bytes % symmetric::kNumAlignmentBytes == 0); + EP_HOST_ASSERT(num_cpu_buffer_bytes <= num_buffer_bytes); + num_gpu_buffer_bytes = num_buffer_bytes - num_cpu_buffer_bytes; + + // Workspace is aligned to 2 MB so that it sits cleanly at the front of the GPU segment + const auto num_workspace_bytes = math::align( + layout::WorkspaceLayout::get_num_bytes(), symmetric::kNumAlignmentBytes); + + // Create NCCL symmetric memory context + // Symmetric memory layout: [[[Workspace] GPU buffer] CPU buffer] + // sym.num_bytes = workspace + buffer, sym.num_cpu_bytes = CPU buffer + const auto num_sym_bytes = num_workspace_bytes + num_buffer_bytes; + this->nccl_context = std::make_shared( + nccl_comm, cpu_comm, num_ranks, rank_idx, + num_sym_bytes, num_cpu_buffer_bytes, + allow_hybrid_mode, sl_idx, num_allocated_qps); + + // Verify the symmetric memory layout matches our expectations + EP_HOST_ASSERT(num_workspace_bytes + num_gpu_buffer_bytes == nccl_context->num_gpu_bytes); + EP_HOST_ASSERT(num_cpu_buffer_bytes == nccl_context->num_cpu_bytes); + + // Timeout + this->num_cpu_timeout_secs = num_cpu_timeout_secs; + this->num_gpu_timeout_cycles = static_cast(num_gpu_timeout_secs); + this->num_gpu_timeout_cycles *= jit::device_runtime->get_clock_rate(); + + // Assign workspaces and buffers + workspace = this->nccl_context->mapped_window_ptr; + workspace_layout_wo_expert = std::make_shared( + workspace, nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks, 0); + buffer = static_cast(workspace) + num_workspace_bytes; + CUDA_RUNTIME_CHECK(cudaMemset(workspace, 0, num_workspace_bytes)); + + // Allocate host workspaces + CUDA_RUNTIME_CHECK(cudaMallocHost(&host_workspace, layout::WorkspaceLayout::get_num_bytes(), cudaHostAllocMapped)); + CUDA_RUNTIME_CHECK(cudaHostGetDevicePointer(&mapped_host_workspace, host_workspace, 0)); + std::memset(host_workspace, 0, layout::WorkspaceLayout::get_num_bytes()); + + // We should call a barrier at the end + // The barrier should be called by Python `dist.barrier` + // NOTES: do not call our barrier, as the workspace is not ready yet + } + + ~ElasticBuffer() noexcept(false) { + if (not explicitly_destroy) + destroy(); + + if (not destroyed) { + printf("`destroy()` is not called before DeepEP elastic buffer destruction, which can leak resources.\n"); + fflush(stdout); + } + } + + void destroy() { + EP_HOST_ASSERT(not destroyed); + + // Finish all works on all GPUs + barrier(true, true); + + // Deallocate host workspaces + CUDA_RUNTIME_CHECK(cudaFreeHost(host_workspace)); + + // Destroy NCCL context + nccl_context->finalize(); + + // Cannot use anymore + destroyed = true; + } + + torch::Stream get_comm_stream() const { + return comm_stream; + } + + std::tuple get_physical_domain_size() const { + return {nccl_context->num_rdma_ranks, nccl_context->num_nvl_ranks}; + } + + std::tuple get_logical_domain_size() const { + return {nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks}; + } + + // ReSharper disable once CppMemberFunctionMayBeStatic + void barrier(const bool& use_comm_stream, const bool& with_cpu_sync, const bool& sequential = true) const { + const auto compute_stream = at::cuda::getCurrentCUDAStream(); + const auto stream = use_comm_stream ? comm_stream : compute_stream; + if (use_comm_stream) + stream_wait(comm_stream, compute_stream); + + // Wait all streams to finish on this GPU + if (with_cpu_sync) + CUDA_RUNTIME_CHECK(cudaDeviceSynchronize()); + + // Launch GPU barrier + launch_barrier(nccl_context->dev_comm, nccl_context->window, + workspace, + nccl_context->scaleout_rank_idx, nccl_context->scaleup_rank_idx, + nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks, + num_gpu_timeout_cycles, + nccl_context->is_scaleup_nvlink, + sequential, + stream); + + // Let CPU wait + if (with_cpu_sync) + CUDA_RUNTIME_CHECK(cudaDeviceSynchronize()); + + // Compute stream should also wait for the barrier + if (use_comm_stream) + stream_wait(compute_stream, comm_stream); + } + + void engram_write(const torch::Tensor& storage, const std::optional& sf) { + // Ensure previous fetch are finished + barrier(false, true); + + const auto compute_stream = at::cuda::getCurrentCUDAStream(); + + // Check storage + const auto [num_entries, hidden] = get_shape<2>(storage); + EP_HOST_ASSERT(storage.scalar_type() == torch::kBFloat16 or + storage.scalar_type() == torch::kFloat8_e4m3fn); + EP_HOST_ASSERT(storage.is_cuda() and storage.is_contiguous()); + num_engram_entries = num_entries, engram_hidden = hidden; + + // Store globally-replicated FP8 scaling factors for local gather during fetch. + if (sf.has_value()) + EP_HOST_ASSERT(sf->dim() == 2 and sf->is_cuda() and sf->is_contiguous() and + sf->element_size() == sizeof(sf_pack_t)); + engram_sf = sf; + + // Write storage to CPU segment at back of buffer + EP_HOST_ASSERT(storage.nbytes() <= num_cpu_buffer_bytes and "Engram storage exceeds CPU buffer size"); + const auto cpu_write_offset = allow_hybrid_mode + ? static_cast(nccl_context->scaleup_rank_idx) * num_cpu_buffer_bytes : 0; + CUDA_RUNTIME_CHECK(cudaMemcpyAsync( + math::advance_ptr(buffer, num_gpu_buffer_bytes + cpu_write_offset), + storage.data_ptr(), storage.nbytes(), + cudaMemcpyDeviceToDevice, compute_stream)); + + // Ensure data is visible for all ranks + barrier(false, true); + } + + std::function>()> + engram_fetch(const torch::Tensor& indices, int num_qps, const bool& use_tma_aligned_col_major_sf) const { + const auto use_fp8 = engram_sf.has_value(); + const auto fetched_dtype = use_fp8 ? torch::kFloat8_e4m3fn : torch::kBFloat16; + const int elem_size = use_fp8 ? sizeof(__nv_fp8_e4m3) : sizeof(nv_bfloat16); + const auto [num_tokens, num_entries_per_token] = get_shape<2>(indices); + EP_HOST_ASSERT(indices.scalar_type() == torch::kInt); + EP_HOST_ASSERT(indices.is_cuda() and indices.is_contiguous()); + EP_HOST_ASSERT(num_tokens * num_entries_per_token * engram_hidden * elem_size <= num_gpu_buffer_bytes); + + // Calculate a QP count + if (num_qps == 0) + num_qps = nccl_context->num_allocated_qps; + + // Return tensor from the raw buffer: each token's entries are concatenated along hidden + EP_HOST_ASSERT(num_engram_entries > 0); + const auto fetched = torch::from_blob( + buffer, + {num_tokens, engram_hidden * num_entries_per_token}, + torch::TensorOptions().dtype(fetched_dtype).device(torch::kCUDA) + ); + + auto fetched_sf = std::optional(); + void* sf_table_ptr = nullptr; + void* fetched_sf_ptr = nullptr; + int num_sf_packs = 0, sf_token_stride = 0, sf_hidden_stride = 0; + if (use_fp8) { + num_sf_packs = static_cast(engram_sf->size(1)); + if (use_tma_aligned_col_major_sf) { + // TMA-aligned column-major layout for the next GEMM input + sf_token_stride = 1, sf_hidden_stride = math::align(num_tokens, kNumAlignedSFPacks); + } else { + sf_token_stride = num_entries_per_token * num_sf_packs, sf_hidden_stride = 1; + } + fetched_sf = torch::empty_strided({num_tokens, num_entries_per_token * num_sf_packs}, + {sf_token_stride, sf_hidden_stride}, engram_sf->options()); + sf_table_ptr = engram_sf->data_ptr(); + fetched_sf_ptr = fetched_sf->data_ptr(); + } + + // Last issued Gin requests + const auto num_gin_ranks = allow_hybrid_mode + ? nccl_context->num_scaleout_ranks + : nccl_context->num_ranks; + const auto last_gin_requests = torch::empty( + {num_gin_ranks * num_qps, sizeof(ncclGinRequest_t)}, + torch::TensorOptions().dtype(torch::kByte).device(torch::kCUDA) + ); + + // Launch the fetch kernel + launch_engram_fetch( + nccl_context->dev_comm, nccl_context->window, + math::advance_ptr(buffer, num_gpu_buffer_bytes), + fetched.data_ptr(), + indices.data_ptr(), + static_cast(last_gin_requests.data_ptr()), + sf_table_ptr, fetched_sf_ptr, sf_token_stride, sf_hidden_stride, + num_engram_entries, + engram_hidden, elem_size, num_sf_packs, + num_entries_per_token, + num_tokens, + nccl_context->num_scaleout_ranks, + nccl_context->num_scaleup_ranks, + num_cpu_buffer_bytes, + num_qps, + allow_hybrid_mode, + at::cuda::getCurrentCUDAStream() + ); + + return [=, this]() -> std::tuple> { + // Wait for all RDMA gets to complete + launch_engram_fetch_wait( + static_cast(last_gin_requests.data_ptr()), + nccl_context->dev_comm, + nccl_context->window, + nccl_context->num_scaleout_ranks, + nccl_context->num_scaleup_ranks, + num_qps, + allow_hybrid_mode, + at::cuda::getCurrentCUDAStream() + ); + return {fetched, fetched_sf}; + }; + } + + void pp_set_config(const int64_t& num_max_tensor_bytes, const int& num_max_inflight_tensors) { + // Flush previous operations + barrier(false, true); + + EP_HOST_ASSERT(num_max_tensor_bytes > 0 and num_max_inflight_tensors > 0); + EP_HOST_ASSERT(num_max_tensor_bytes * num_max_inflight_tensors * 2 * 2 <= num_buffer_bytes); + this->prev_rank_idx = (nccl_context->rank_idx + nccl_context->num_ranks - 1) % nccl_context->num_ranks; + this->next_rank_idx = (nccl_context->rank_idx + 1) % nccl_context->num_ranks; + this->num_max_pp_tensor_bytes = math::align(num_max_tensor_bytes, 32); + this->num_max_pp_inflight_tensors = num_max_inflight_tensors; + } + + void pp_send(const torch::Tensor& x, const int& dst_rank_idx, const int& num_sms) const { + EP_HOST_ASSERT(num_max_pp_tensor_bytes > 0 and num_max_pp_inflight_tensors > 0); + EP_HOST_ASSERT(x.is_cuda() and x.is_contiguous() and x.nbytes() <= num_max_pp_tensor_bytes); + EP_HOST_ASSERT(dst_rank_idx == prev_rank_idx or dst_rank_idx == next_rank_idx); + + launch_pp_send( + nccl_context->dev_comm, nccl_context->window, + x.data_ptr(), x.nbytes(), + buffer, workspace, + nccl_context->rank_idx, dst_rank_idx, nccl_context->num_ranks, + num_max_pp_tensor_bytes, + num_max_pp_inflight_tensors, + num_sms == 0 ? jit::device_runtime->get_num_sms() : num_sms, + num_gpu_timeout_cycles, + jit::device_runtime->get_num_smem_bytes(), + at::cuda::getCurrentCUDAStream() + ); + } + + void pp_recv(const torch::Tensor& x, const int& src_rank_idx, const int& num_sms) const { + EP_HOST_ASSERT(num_max_pp_tensor_bytes > 0 and num_max_pp_inflight_tensors > 0); + EP_HOST_ASSERT(x.is_cuda() and x.is_contiguous() and x.nbytes() <= num_max_pp_tensor_bytes); + EP_HOST_ASSERT(src_rank_idx == prev_rank_idx or src_rank_idx == next_rank_idx); + + launch_pp_recv( + nccl_context->dev_comm, nccl_context->window, + x.data_ptr(), x.nbytes(), + buffer, workspace, + nccl_context->rank_idx, src_rank_idx, nccl_context->num_ranks, + num_max_pp_tensor_bytes, + num_max_pp_inflight_tensors, + num_sms == 0 ? jit::device_runtime->get_num_sms() : num_sms, + num_gpu_timeout_cycles, + jit::device_runtime->get_num_smem_bytes(), + at::cuda::getCurrentCUDAStream() + ); + } + + void agrs_set_config(const int64_t& num_max_session_bytes, + const int& new_num_max_agrs_per_session) { + // Flush previous operations + barrier(true, true); + + EP_HOST_ASSERT(nccl_context->num_ranks > 1); + EP_HOST_ASSERT(num_max_session_bytes > 0 and new_num_max_agrs_per_session > 0); + EP_HOST_ASSERT(num_max_session_bytes <= num_buffer_bytes); + EP_HOST_ASSERT(new_num_max_agrs_per_session <= layout::WorkspaceLayout::kNumMaxInflightAGRS); + EP_HOST_ASSERT(nccl_context->num_nvl_ranks == nccl_context->num_ranks); + this->num_max_agrs_session_bytes = math::align(num_max_session_bytes, 32); + this->num_max_agrs_per_session = new_num_max_agrs_per_session; + } + + void create_agrs_session() { + EP_HOST_ASSERT(not agrs_in_session); + agrs_in_session = true; + agrs_buffer_offset = 0; + agrs_buffer_slot_idx = 0; + agrs_session_idx += 1; + } + + void destroy_agrs_session() { + // Must be in a session + EP_HOST_ASSERT(agrs_in_session); + agrs_in_session = false; + + // Wait compute stream + stream_wait(comm_stream, at::cuda::getCurrentCUDAStream()); + + // Notify that the buffer is now available & Wait for the buffer to be ready + // NOTES: self-wait is guaranteed by in-stream order + std::vector write_ptrs(nccl_context->num_ranks - 1); + std::vector wait_ptrs(nccl_context->num_ranks - 1); + for (int i = 0; i < nccl_context->num_ranks - 1; ++ i) { + const auto dst_rank_idx = (nccl_context->rank_idx + i + 1) % nccl_context->num_ranks; + write_ptrs[i] = static_cast( + nccl_context->get_sym_ptr(workspace_layout_wo_expert->get_agrs_session_signal_ptr(nccl_context->rank_idx), dst_rank_idx)); + wait_ptrs[i] = workspace_layout_wo_expert->get_agrs_session_signal_ptr(dst_rank_idx); + } + cuda_driver::batched_write_and_wait(comm_stream, write_ptrs, wait_ptrs, agrs_session_idx); + } + + std::vector agrs_get_inplace_tensor(const std::vector& num_bytes_list) const { + EP_HOST_ASSERT(num_bytes_list.size() >= 1); + EP_HOST_ASSERT(num_max_agrs_session_bytes > 0 and num_max_agrs_per_session > 0 and agrs_in_session); + + std::vector out; + out.reserve(num_bytes_list.size()); + int64_t offset = agrs_buffer_offset; + for (const auto& num_bytes: num_bytes_list) { + EP_HOST_ASSERT(offset + num_bytes * nccl_context->num_ranks <= num_max_agrs_session_bytes and + agrs_buffer_slot_idx < num_max_agrs_per_session and + "Not enough session buffer size. Did you forget to flush session?"); + out.push_back(torch::from_blob(math::advance_ptr(buffer, offset + num_bytes * nccl_context->rank_idx), + {num_bytes}, torch::TensorOptions().dtype(torch::kByte).device(torch::kCUDA))); + offset += math::align(num_bytes * nccl_context->num_ranks, 32); + } + return out; + } + + std::pair, std::function> + all_gather(const std::vector& tensors) { + const int num_tensors = tensors.size(); + EP_HOST_ASSERT(num_max_agrs_session_bytes > 0 and num_max_agrs_per_session > 0 and agrs_in_session); + EP_HOST_ASSERT(num_tensors >= 1); + + int num_copies = 0; + std::vector offset(num_tensors); + for (int i = 0; i < num_tensors; ++ i) { + const auto& x = tensors[i]; + EP_HOST_ASSERT(x.is_contiguous()); + + const auto x_offset = math::ptr_diff(x.data_ptr(), buffer); + const bool is_inplace = 0 <= x_offset and x_offset < num_max_agrs_session_bytes; + offset[i] = agrs_buffer_offset; + num_copies += nccl_context->num_ranks - is_inplace; + agrs_buffer_offset += math::align(x.nbytes() * nccl_context->num_ranks, 32); + EP_HOST_ASSERT(not is_inplace or x.data_ptr() == math::advance_ptr(buffer, offset[i] + x.nbytes() * nccl_context->rank_idx)); + } + EP_HOST_ASSERT(agrs_buffer_offset <= num_max_agrs_session_bytes and + agrs_buffer_slot_idx < num_max_agrs_per_session and + "Not enough session buffer size. Did you forget to flush session?"); + + // Wait compute stream + const auto compute_stream = at::cuda::getCurrentCUDAStream(); + stream_wait(comm_stream, compute_stream); + + // Send data to all ranks + std::vector sizes(num_copies); + std::vector dst_ptrs(num_copies), src_ptrs(num_copies); + int count = 0; + for (int i = 0; i < nccl_context->num_ranks; ++ i) { + for (int j = 0; j < num_tensors; ++ j) { + const auto& x = tensors[j]; + const auto dst_rank_idx = (nccl_context->rank_idx + i) % nccl_context->num_ranks; + void* src_ptr = x.data_ptr(); + void* dst_ptr = + nccl_context->get_sym_ptr(math::advance_ptr(buffer, offset[j] + x.nbytes() * nccl_context->rank_idx), dst_rank_idx); + if (src_ptr != dst_ptr) { + src_ptrs[count] = src_ptr; + dst_ptrs[count] = dst_ptr; + sizes[count] = x.nbytes(); + count += 1; + } + } + } + cudaMemcpyAttributes attrs = { + .srcAccessOrder = cudaMemcpySrcAccessOrderStream, + .flags = cudaMemcpyFlagPreferOverlapWithCompute + }; +#if defined(CUDART_VERSION) and CUDART_VERSION >= 13000 + CUDA_RUNTIME_CHECK(cudaMemcpyBatchAsync(dst_ptrs.data(), src_ptrs.data(), sizes.data(), num_copies, attrs, comm_stream)); +#else + CUDA_RUNTIME_CHECK(cudaMemcpyBatchAsync(dst_ptrs.data(), src_ptrs.data(), sizes.data(), num_copies, attrs, nullptr, comm_stream)); +#endif + + // Wait for data from other ranks + const int current_session = agrs_session_idx; + const int slot_idx = agrs_buffer_slot_idx; + agrs_buffer_slot_idx += 1; + std::vector write_ptrs(nccl_context->num_ranks - 1); + std::vector wait_ptrs(nccl_context->num_ranks - 1); + for (int i = 0; i < nccl_context->num_ranks - 1; ++ i) { + const auto dst_rank_idx = (nccl_context->rank_idx + i + 1) % nccl_context->num_ranks; + write_ptrs[i] = nccl_context->get_sym_ptr( + workspace_layout_wo_expert->get_agrs_recv_signal_ptr(slot_idx, nccl_context->rank_idx), dst_rank_idx); + wait_ptrs[i] = workspace_layout_wo_expert->get_agrs_recv_signal_ptr(slot_idx, dst_rank_idx); + } + cuda_driver::batched_write_and_wait(comm_stream, write_ptrs, wait_ptrs, current_session); + + // Build output tensors eagerly + std::vector out(num_tensors); + for (int i = 0; i < num_tensors; ++ i) { + auto shape = tensors[i].sizes().vec(); + shape.insert(shape.begin(), nccl_context->num_ranks); + out[i] = torch::from_blob(math::advance_ptr(buffer, offset[i]), shape, tensors[i].options()); + } + + // Return tensors and a handle to wait for data arrival + const auto event = EventHandle(comm_stream); + auto handle = [=, this]() { + EP_HOST_ASSERT(compute_stream == at::cuda::getCurrentCUDAStream()); + EP_HOST_ASSERT(agrs_in_session and current_session == this->agrs_session_idx); + stream_wait(compute_stream, event); + }; + return {std::move(out), std::move(handle)}; + } + + torch::cuda::CUDAStream stream_control_prologue(const std::optional& previous_event, + const bool& allocate_on_comm_stream, + const bool& async_with_compute_stream) const { + // Allocate all tensors on communication stream if set + // NOTES: do not allocate tensors upfront! + const auto compute_stream = at::cuda::getCurrentCUDAStream(); + if (allocate_on_comm_stream) + at::cuda::setCurrentCUDAStream(comm_stream); + + // Assertion for safety + // `previous_event` implicitly means the overlapping computation kernels are launched first, + // in order not to use the memory on the compute stream, we must allocate on the communication stream. + // If you launch the communication kernels firstly, then `previous_event` must be unnecessary. + if (previous_event.has_value()) + EP_HOST_ASSERT(allocate_on_comm_stream); + + // Wait previous tasks to finish + if (previous_event.has_value()) { + stream_wait(comm_stream, previous_event.value()); + } else { + stream_wait(comm_stream, compute_stream); + } + return compute_stream; + } + + void stream_control_before_epilogue(const std::optional& previous_event_before_epilogue) const { + if (previous_event_before_epilogue.has_value()) + stream_wait(comm_stream, previous_event_before_epilogue.value()); + } + + std::optional stream_control_epilogue(const std::vector>& tensors, + const at::cuda::CUDAStream& compute_stream, + const bool& allocate_on_comm_stream, + const bool& async_with_compute_stream) const { + // Ensure memory access safety between two streams + std::optional event; + if (async_with_compute_stream) { + event = EventHandle(comm_stream); + + // NOTES: this environment only applies to V2 APIs + if (get_env("EP_AVOID_RECORD_STREAM", 0)) { + event->tensors_to_record = tensors; + } else { + for (auto& t: tensors) if (t.has_value()) { + t->record_stream(compute_stream); + t->record_stream(comm_stream); + } + } + } else { + stream_wait(compute_stream, comm_stream); + } + + // Switch back compute stream + if (allocate_on_comm_stream) + at::cuda::setCurrentCUDAStream(compute_stream); + + // The CUDA event marking the finishing + return event; + } + + static int64_t get_dispatch_buffer_size(const int& num_max_tokens_per_rank, + const int& hidden, const int& num_sf_packs, const int& num_topk, + const int& elem_size, + const int& num_scaleout_ranks, const int& num_scaleup_ranks, + const bool& is_scaleup_nvlink) { + const auto num_ranks = num_scaleup_ranks * num_scaleout_ranks; + const auto token_layout = get_dispatch_token_layout(hidden, elem_size, num_sf_packs, num_topk); + + if (num_scaleout_ranks == 1) { + // Direct dispatch + const auto send_buffer_layout = layout::BufferLayout( + token_layout, is_scaleup_nvlink ? 0 : 1, num_max_tokens_per_rank); + const auto recv_buffer_layout = layout::BufferLayout( + token_layout, num_ranks, num_max_tokens_per_rank); + return send_buffer_layout.get_num_bytes() + recv_buffer_layout.get_num_bytes(); + } else { + // Hybrid dispatch + const auto scaleup_recv_buffer = layout::BufferLayout( + token_layout, num_scaleup_ranks, num_scaleout_ranks * num_max_tokens_per_rank); + const auto scaleout_send_buffer = layout::BufferLayout( + token_layout, 1, num_max_tokens_per_rank); + const auto scaleout_recv_buffer = layout::BufferLayout( + token_layout, num_scaleout_ranks, + /* kNumChannels * kNumMaxTokensPerChannel */ num_max_tokens_per_rank + kNumMaxChannels); + return scaleup_recv_buffer.get_num_bytes() + + scaleout_send_buffer.get_num_bytes() + + scaleout_recv_buffer.get_num_bytes(); + } + } + + static int64_t get_combine_buffer_size(const int& num_max_tokens_per_rank, const int& hidden, const int& num_topk, + const int& num_scaleout_ranks, const int& num_scaleup_ranks, + const bool& is_scaleup_nvlink, + const bool& allow_multiple_reduction) { + const auto num_ranks = num_scaleup_ranks * num_scaleout_ranks; + const auto token_layout = get_combine_token_layout(hidden, sizeof(nv_bfloat16), num_topk); + + if (num_scaleout_ranks == 1) { + // Direct combine + const auto num_tokens_in_layout = allow_multiple_reduction ? std::min(num_ranks, num_topk) : num_topk; + const auto send_buffer_layout = layout::BufferLayout( + token_layout, is_scaleup_nvlink ? 0 : num_ranks, + // For single reduction cases, the maximum number of received tokens is + // `num_ranks * num_topk * num_max_tokens_per_rank` (we assume the bad case of `do_expand=True`) + num_max_tokens_per_rank * (allow_multiple_reduction ? 1 : num_topk)); + const auto recv_buffer_layout = layout::BufferLayout( + token_layout, num_tokens_in_layout, num_max_tokens_per_rank); + return send_buffer_layout.get_num_bytes() + recv_buffer_layout.get_num_bytes(); + } else { + // Hybrid combine + const int num_tokens_in_scaleup_layout = allow_multiple_reduction ? std::min(num_scaleup_ranks, num_topk) : num_topk; + const int num_tokens_in_scaleout_layout = allow_multiple_reduction ? std::min(num_scaleout_ranks, num_topk) : num_topk; + const auto scaleup_recv_buffer = layout::BufferLayout( + token_layout, num_tokens_in_scaleup_layout, num_scaleout_ranks * num_max_tokens_per_rank); + const auto scaleout_recv_buffer = layout::BufferLayout( + token_layout, num_tokens_in_scaleout_layout, num_max_tokens_per_rank); + const auto scaleout_send_buffer = layout::BufferLayout( + token_layout, allow_multiple_reduction ? 1 : num_topk, + /* kNumChannels * num_scaleout_ranks * kNumMaxTokensPerChannel */ + num_scaleout_ranks * (num_max_tokens_per_rank + kNumMaxChannels)); + return scaleup_recv_buffer.get_num_bytes() + + scaleout_send_buffer.get_num_bytes() + + scaleout_recv_buffer.get_num_bytes(); + } + } + + static int64_t calculate_buffer_size(const int64_t& nccl_comm, + const int& num_max_tokens_per_rank, const int& hidden, + int num_topk, const bool& use_fp8_dispatch, + const bool& allow_hybrid_mode, + const bool& allow_multiple_reduction) { + EP_HOST_ASSERT(num_max_tokens_per_rank > 0 and hidden > 0); + + // The worst case SF bytes must be less than the main part + EP_HOST_ASSERT(math::ceil_div(hidden, 32) * sizeof(float) <= hidden); + + // NOTES: there are lots of `kNumTopk <= 32` restrictions, so we use 32 to calculate token size + num_topk = num_topk == 0 ? 32 : num_topk; + + // Topology + const auto [num_rdma_ranks, num_nvl_ranks] = nccl::get_physical_domain_size(nccl_comm); + const auto [num_scaleout_ranks, num_scaleup_ranks] = nccl::get_logical_domain_size(nccl_comm, allow_hybrid_mode); + const auto is_scaleup_nvlink = num_scaleup_ranks == num_nvl_ranks; + + // Dispatch size + const auto elem_size = use_fp8_dispatch ? sizeof(__nv_fp8_e4m3) : sizeof(nv_bfloat16); + const auto num_sf_packs = use_fp8_dispatch ? math::ceil_div(hidden, 32) : 0; // An approximation for number of SF packs + const auto num_dispatch_bytes = get_dispatch_buffer_size( + num_max_tokens_per_rank, hidden, num_sf_packs, num_topk, elem_size, + num_scaleout_ranks, num_scaleup_ranks, + is_scaleup_nvlink); + + // Combine layout + const auto num_combine_bytes = get_combine_buffer_size( + num_max_tokens_per_rank, hidden, num_topk, + num_scaleout_ranks, num_scaleup_ranks, + is_scaleup_nvlink, allow_multiple_reduction); + + // Return the maximum of those layouts, aligned to 2 MB + return math::align(std::max(num_dispatch_bytes, num_combine_bytes), symmetric::kNumAlignmentBytes); + } + + static symmetric::cpu_handle_t create_cpu_handle(const int64_t& num_cpu_bytes) { + EP_HOST_ASSERT(num_cpu_bytes > 0 and num_cpu_bytes % symmetric::kNumAlignmentBytes == 0); + return symmetric::HybridElasticSymmetricMemory::create_cpu_handle(num_cpu_bytes); + } + + std::tuple, + std::optional, std::optional, + std::optional, + int, int, + std::vector, + torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, + std::optional, std::optional, + std::optional> + dispatch(const torch::Tensor& x, + const std::optional& sf, + const torch::Tensor& topk_idx, + const std::optional& topk_weights, + const std::optional& cumulative_local_expert_recv_stats, + const std::optional& cached_num_recv_tokens, + const std::optional& cached_num_expanded_tokens, + const std::optional>& cached_num_recv_tokens_per_expert_list, + const std::optional& cached_psum_num_recv_tokens_per_scaleup_rank, + const std::optional& cached_psum_num_recv_tokens_per_expert, + const std::optional& cached_num_unaligned_recv_tokens_per_expert, + const std::optional& cached_dst_buffer_slot_idx, + const std::optional& cached_token_metadata_at_forward, + const std::optional& cached_recv_src_metadata, + const std::optional& cached_channel_linked_list, + const int& num_max_tokens_per_rank, + const int& num_experts, const int& expert_alignment, + const int& num_sms, const int& num_qps, + const std::optional& previous_event, + const std::optional& previous_event_before_epilogue, + const bool& async_with_compute_stream, + const bool& allocate_on_comm_stream, + const bool& do_handle_copy, const bool& do_cpu_sync, + const bool& do_expand, const bool& do_zero_padding, + const bool& use_tma_aligned_col_major_sf) const { + // Check SM count + EP_HOST_ASSERT(num_sms > 0); + + // Zero padding only makes sense with expand mode + EP_HOST_ASSERT(not do_zero_padding or do_expand); + + // Cached mode must have responding handles + const bool cached_mode = cached_num_recv_tokens.has_value(); + if (cached_mode) { + EP_HOST_ASSERT(cached_num_recv_tokens.has_value()); + EP_HOST_ASSERT(cached_num_recv_tokens_per_expert_list.has_value()); + EP_HOST_ASSERT(cached_num_expanded_tokens.has_value()); + EP_HOST_ASSERT(cached_psum_num_recv_tokens_per_scaleup_rank.has_value()); + EP_HOST_ASSERT(cached_psum_num_recv_tokens_per_expert.has_value()); + EP_HOST_ASSERT(cached_dst_buffer_slot_idx.has_value()); + + // Hybrid kernels require more + if (nccl_context->num_scaleout_ranks > 1) { + EP_HOST_ASSERT(cached_token_metadata_at_forward.has_value()); + EP_HOST_ASSERT(cached_channel_linked_list.has_value()); + } + } + + // Check data tensor + const auto [num_tokens, hidden] = get_shape<2>(x); + const auto num_hidden_bytes = hidden * static_cast(x.element_size()); + const auto num_local_experts = num_experts / nccl_context->num_ranks; + EP_HOST_ASSERT(x.is_cuda() and x.is_contiguous()); + EP_HOST_ASSERT((x.size(1) * x.element_size()) % sizeof(int4) == 0); + EP_HOST_ASSERT(num_tokens <= num_max_tokens_per_rank); + + // Check SF stuffs + int num_sf_packs = 0; + void* sf_ptr = nullptr; + int sf_token_stride = 0, sf_hidden_stride = 0; + if (sf.has_value()) { + // SF must be FP32 or packed UE8M0x4 + const auto [num_tokens_, num_sf_packs_] = get_shape<2>(sf.value()); + EP_HOST_ASSERT(num_tokens == num_tokens_); + EP_HOST_ASSERT(sf->is_cuda()); + EP_HOST_ASSERT(sf->element_size() == sizeof(sf_pack_t)); + num_sf_packs = num_sf_packs_; + sf_ptr = sf->data_ptr(); + sf_token_stride = sf->stride(0); + sf_hidden_stride = sf->stride(1); + } + + // Check top-k stuffs + const auto [num_tokens_, num_topk] = get_shape<2>(topk_idx); + EP_HOST_ASSERT(num_tokens == num_tokens_); + EP_HOST_ASSERT(topk_idx.scalar_type() == c10::CppTypeToScalarType::value); + EP_HOST_ASSERT(topk_idx.is_cuda() and topk_idx.is_contiguous()); + + // Weights are optional for training backward + float* topk_weights_ptr = nullptr; + if (topk_weights.has_value()) { + const auto [num_tokens__, num_topk_] = get_shape<2>(topk_weights.value()); + EP_HOST_ASSERT(num_tokens == num_tokens__); + EP_HOST_ASSERT(topk_weights->is_cuda() and topk_weights->is_contiguous()); + topk_weights_ptr = topk_weights->data_ptr(); + } + + // Expert receiving counter + int* cumulative_local_expert_recv_stats_ptr = nullptr; + if (cumulative_local_expert_recv_stats.has_value()) { + const auto [num_local_experts_] = get_shape<1>(cumulative_local_expert_recv_stats.value()); + EP_HOST_ASSERT(cumulative_local_expert_recv_stats->is_cuda() and + cumulative_local_expert_recv_stats->is_contiguous()); + EP_HOST_ASSERT(num_local_experts == num_local_experts_); + cumulative_local_expert_recv_stats_ptr = cumulative_local_expert_recv_stats->data_ptr(); + } + + // Stream control + // All new tensor allocations should happen after this + const auto compute_stream = stream_control_prologue(previous_event, allocate_on_comm_stream, async_with_compute_stream); + + // The number of received tokens per expert + // This is useful for expanding mode + EP_HOST_ASSERT(num_experts % nccl_context->num_ranks == 0); + auto psum_num_recv_tokens_per_expert = cached_psum_num_recv_tokens_per_expert.value_or(torch::Tensor()); + if (cached_mode) { + const auto& [num_local_experts_] = get_shape<1>(psum_num_recv_tokens_per_expert); + EP_HOST_ASSERT(num_local_experts == num_local_experts_); + EP_HOST_ASSERT(psum_num_recv_tokens_per_expert.is_cuda() and psum_num_recv_tokens_per_expert.is_contiguous()); + EP_HOST_ASSERT(psum_num_recv_tokens_per_expert.scalar_type() == torch::kInt); + } else { + // NOTES: for expand mode, the input is exclusive prefix sum, while for non-expand, it is inclusive + psum_num_recv_tokens_per_expert = torch::empty( + {num_local_experts + 1}, at::TensorOptions(torch::kCUDA).dtype(torch::kInt)); + } + + // The unaligned (actual) number of received tokens per expert + // Written by the dispatch kernel's notify warps, used by the epilogue for zero padding + auto num_unaligned_recv_tokens_per_expert = cached_num_unaligned_recv_tokens_per_expert.value_or(torch::Tensor()); + int* num_unaligned_recv_tokens_per_expert_ptr = nullptr; + if (cached_mode) { + const auto& [num_local_experts_] = get_shape<1>(num_unaligned_recv_tokens_per_expert); + EP_HOST_ASSERT(num_local_experts == num_local_experts_); + EP_HOST_ASSERT(num_unaligned_recv_tokens_per_expert.is_cuda() and num_unaligned_recv_tokens_per_expert.is_contiguous()); + EP_HOST_ASSERT(num_unaligned_recv_tokens_per_expert.scalar_type() == torch::kInt); + } else { + num_unaligned_recv_tokens_per_expert = torch::empty( + {num_local_experts}, at::TensorOptions(torch::kCUDA).dtype(torch::kInt)); + } + num_unaligned_recv_tokens_per_expert_ptr = num_unaligned_recv_tokens_per_expert.data_ptr(); + + // The prefix sum tensor of number of received tokens from each rank + // Will also be used in combine as the dispatch handle + auto psum_num_recv_tokens_per_scaleup_rank = cached_psum_num_recv_tokens_per_scaleup_rank.value_or(torch::Tensor()); + if (cached_mode) { + const auto [num_scaleup_ranks] = get_shape<1>(psum_num_recv_tokens_per_scaleup_rank); + EP_HOST_ASSERT(num_scaleup_ranks == nccl_context->num_scaleup_ranks); + EP_HOST_ASSERT(psum_num_recv_tokens_per_scaleup_rank.is_cuda() and psum_num_recv_tokens_per_scaleup_rank.is_contiguous()); + EP_HOST_ASSERT(psum_num_recv_tokens_per_scaleup_rank.scalar_type() == torch::kInt); + } else { + psum_num_recv_tokens_per_scaleup_rank = torch::empty( + {nccl_context->num_scaleup_ranks}, at::TensorOptions(torch::kCUDA).dtype(torch::kInt)); + } + + // Decide number of channels by shared memory consumption + // Only for hybrid version + int num_channels_per_sm = 1, num_channels = 1; + const int num_smem_bytes = jit::device_runtime->get_num_smem_bytes(); + if (nccl_context->num_scaleout_ranks > 1) { + const auto dispatch_token_layout = get_dispatch_token_layout(hidden, x.element_size(), num_sf_packs, num_topk); + const auto combine_token_layout = get_combine_token_layout(hidden, sizeof(nv_bfloat16), num_topk); + EP_HOST_ASSERT(num_sms <= kNumMaxSMs); + num_channels_per_sm = std::min( + (num_smem_bytes - get_num_notify_smem_bytes(nccl_context->num_ranks, num_experts)) / dispatch_token_layout.get_num_bytes(), + 32 - kNumNotifyWarps); + num_channels_per_sm = std::min( + num_smem_bytes / combine_token_layout.get_num_bytes(), + num_channels_per_sm); + num_channels_per_sm = std::min( + /* 2 kinds of warps */ num_channels_per_sm / 2, kNumMaxChannelsPerSM); + if (not prefer_overlap_with_compute) + num_channels_per_sm = std::min(num_channels_per_sm, 4); + num_channels = num_sms * num_channels_per_sm; + if (get_env("EP_BUFFER_DEBUG")) + printf("Elastic buffer uses %d channels per SM\n", num_channels_per_sm); + } + + // Non-hybrid mode handles + auto dst_buffer_slot_idx = cached_dst_buffer_slot_idx.value_or(torch::Tensor()); + if (nccl_context->num_scaleout_ranks == 1) { + if (cached_mode) { + const auto [num_tokens__, num_topk_] = get_shape<2>(dst_buffer_slot_idx); + EP_HOST_ASSERT(num_tokens == num_tokens__ and num_topk == num_topk_); + EP_HOST_ASSERT(dst_buffer_slot_idx.is_cuda() and dst_buffer_slot_idx.is_contiguous()); + EP_HOST_ASSERT(dst_buffer_slot_idx.scalar_type() == torch::kInt); + } else { + // Allocate a new tensor + dst_buffer_slot_idx = torch::empty( + {num_tokens, num_topk}, torch::TensorOptions(torch::kCUDA).dtype(torch::kInt)); + } + } + + // Hybrid mode handles + std::optional token_metadata_at_forward, channel_linked_list; + int *token_metadata_at_forward_ptr = nullptr, *channel_linked_list_ptr = nullptr; + if (nccl_context->num_scaleout_ranks > 1) { + // The token destination slot idx during forward + // `[i, j, k, l]` means: from channel i from scale-out peer k, the j-th token's index in the l-th rank buffer + // NOTES: Used primarily for cached mode + // TODO: May make it a linked list to remove the redundant info in `token_metadata_at_forward` + const auto num_max_tokens_per_channel = math::ceil_div(num_max_tokens_per_rank, num_channels); + if (cached_mode) { + const auto [num_channels_, num_scaleout_ranks_, num_max_tokens_per_channel_, num_topk_] = + get_shape<4>(dst_buffer_slot_idx); + EP_HOST_ASSERT(num_channels == num_channels_ and nccl_context->num_scaleout_ranks == num_scaleout_ranks_ and + num_max_tokens_per_channel == num_max_tokens_per_channel_ and num_topk == num_topk_); + EP_HOST_ASSERT(dst_buffer_slot_idx.is_cuda() and dst_buffer_slot_idx.is_contiguous()); + EP_HOST_ASSERT(dst_buffer_slot_idx.scalar_type() == torch::kInt); + } else { + dst_buffer_slot_idx = torch::empty( + {num_channels, nccl_context->num_scaleout_ranks, num_max_tokens_per_channel, num_topk}, + torch::TensorOptions().device(torch::kCUDA).dtype(torch::kInt) + ); + } + + // The token metadata during forward + // `[i, j]` means: in channel i, the j-th forwarded token's metadata + // Info contains: + // - Scaleout rank index and source token index in the original rank (0) + // - Whether the token is the last one in the chunk (1) + // - cached top-k scaleup peer indices (top-k) + // - each selections' destination slot indices (top-k) + const auto num_max_forwarded_tokens = nccl_context->num_scaleout_ranks * num_max_tokens_per_channel + 1; + const auto num_forward_metadata_dims = 2 + num_topk * 2; + if (cached_mode) { + token_metadata_at_forward = cached_token_metadata_at_forward; + const auto [num_channels_, num_max_forwarded_tokens_, num_forward_metadata_dims_] = get_shape<3>(token_metadata_at_forward.value()); + EP_HOST_ASSERT(num_channels == num_channels_ and num_max_forwarded_tokens == num_max_forwarded_tokens_ + and num_forward_metadata_dims == num_forward_metadata_dims_); + EP_HOST_ASSERT(token_metadata_at_forward->is_cuda() and token_metadata_at_forward->is_contiguous()); + EP_HOST_ASSERT(token_metadata_at_forward->scalar_type() == torch::kInt); + } else { + token_metadata_at_forward = torch::empty( + {num_channels, num_max_forwarded_tokens, num_forward_metadata_dims}, + torch::TensorOptions().device(torch::kCUDA).dtype(torch::kInt) + ); + } + token_metadata_at_forward_ptr = token_metadata_at_forward->data_ptr(); + + // Per-scaleup-peer-per-channel linked list + // `[i, j, k]` means: from channel i from scaleup peer k, the j-th token's index in the combine's input + if (cached_mode) { + channel_linked_list = cached_channel_linked_list; + const auto [num_channels__, d1_, d2_] = get_shape<3>(channel_linked_list.value()); + channel_linked_list_ptr = channel_linked_list->data_ptr(); + EP_HOST_ASSERT(num_channels == num_channels__); + EP_HOST_ASSERT(d1_ == nccl_context->num_scaleout_ranks * num_max_tokens_per_channel + 1); + EP_HOST_ASSERT(d2_ == nccl_context->num_scaleup_ranks); + EP_HOST_ASSERT(channel_linked_list->is_cuda() and channel_linked_list->is_contiguous()); + EP_HOST_ASSERT(channel_linked_list->scalar_type() == torch::kInt); + } else { + channel_linked_list = torch::empty( + // Index 0 of the list means the starting item + {num_channels, + nccl_context->num_scaleout_ranks * num_max_tokens_per_channel + 1, + nccl_context->num_scaleup_ranks}, + torch::TensorOptions().device(torch::kCUDA).dtype(torch::kInt) + ); + } + channel_linked_list_ptr = channel_linked_list->data_ptr(); + } + + // Clone `topk_idx` for saving in the handle (to prevent users' modification) + auto copied_topk_idx = std::optional(); + topk_idx_t* copied_topk_idx_ptr = nullptr; + if (do_handle_copy and not cached_mode) { + copied_topk_idx = torch::empty_like(topk_idx); + copied_topk_idx_ptr = copied_topk_idx->data_ptr(); + } + + // Check buffer size + EP_HOST_ASSERT(get_dispatch_buffer_size( + num_max_tokens_per_rank, hidden, num_sf_packs, num_topk, x.element_size(), + nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks, + nccl_context->is_scaleup_nvlink) <= num_buffer_bytes); + + // Ready and clean host workspace for this round + const auto host_workspace_layout = layout::WorkspaceLayout( + host_workspace, + nccl_context->num_scaleout_ranks, + nccl_context->num_scaleup_ranks, + num_experts); + std::fill_n(host_workspace_layout.get_scaleup_rank_count_ptr(), nccl_context->num_scaleup_ranks, 0); + std::fill_n(host_workspace_layout.get_scaleup_expert_count_ptr(), num_local_experts, 0); + std::atomic_thread_fence(std::memory_order_seq_cst); + + // Do dispatch into the buffers (with SM limitation) + EP_HOST_ASSERT(num_sms <= jit::device_runtime->get_num_sms()); + launch_dispatch(x.data_ptr(), sf_ptr, + topk_idx.data_ptr(), topk_weights_ptr, + copied_topk_idx_ptr, + cumulative_local_expert_recv_stats_ptr, + psum_num_recv_tokens_per_scaleup_rank.data_ptr(), + psum_num_recv_tokens_per_expert.data_ptr(), + num_unaligned_recv_tokens_per_expert_ptr, + dst_buffer_slot_idx.data_ptr(), + token_metadata_at_forward_ptr, + num_tokens, num_max_tokens_per_rank, + hidden, x.element_size(), + num_sf_packs, sf_token_stride, sf_hidden_stride, + num_experts, num_topk, expert_alignment, + nccl_context->dev_comm, nccl_context->window, + buffer, + workspace, mapped_host_workspace, + nccl_context->scaleout_rank_idx, nccl_context->scaleup_rank_idx, + nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks, + nccl_context->is_scaleup_nvlink, + num_sms, num_channels_per_sm, + num_smem_bytes, + num_qps, num_gpu_timeout_cycles, + cached_mode, do_cpu_sync, + comm_stream); + + // Received token counters + int num_recv_tokens = 0, num_expanded_tokens = 0; + int counter_scaleup_rank_idx = 0, counter_local_expert_idx = 0; + std::vector num_recv_tokens_per_expert_list; + + // Assign these values according to modes + if (cached_mode) { + // Cached mode + EP_HOST_ASSERT(not do_cpu_sync and "Cannot do CPU sync with cached mode"); + num_recv_tokens = cached_num_recv_tokens.value(); + num_recv_tokens_per_expert_list = cached_num_recv_tokens_per_expert_list.value(); + num_expanded_tokens = cached_num_expanded_tokens.value(); + } else if (do_cpu_sync) { + // Non-cached mode with sync + const auto start_cpu_time = std::chrono::high_resolution_clock::now(); + while (true) { + bool ready = true; + + // Read number of received tokens from each scaleup rank + while (counter_scaleup_rank_idx < nccl_context->num_scaleup_ranks and ready) { + const auto count = math::encode_decode_positive( + host_workspace_layout.get_scaleup_rank_count_ptr()[counter_scaleup_rank_idx]); + if ((ready = math::is_decoded_positive_ready(count))) { + num_recv_tokens += count; + ++ counter_scaleup_rank_idx; + } + } + + // Read expert counts + while (counter_local_expert_idx < num_local_experts and ready) { + const auto count = math::encode_decode_positive( + host_workspace_layout.get_scaleup_expert_count_ptr()[counter_local_expert_idx]); + if ((ready = math::is_decoded_positive_ready(count))) { + num_recv_tokens_per_expert_list.push_back(count); + num_expanded_tokens += count; + ++ counter_local_expert_idx; + } + } + + // Ready and do next steps + const auto get_buffer_info = [&]() { + std::stringstream ss; + ss << "CPU side received count (scaleup: " << nccl_context->scaleup_rank_idx << "): "; + for (int i = 0; i < nccl_context->num_scaleup_ranks + num_local_experts; ++ i) { + ss << host_workspace_layout.get_scaleup_rank_expert_count_ptr()[i]; + ss << (i == nccl_context->num_scaleup_ranks - 1 ? " # ": " "); + } + return ss.str(); + }; + if (ready) { + if (get_env("EP_BUFFER_DEBUG")) + printf("%s\n", get_buffer_info().c_str()); + break; + } + + // Timeout checks + const auto now = std::chrono::high_resolution_clock::now(); + if (std::chrono::duration_cast(now - start_cpu_time).count() > num_cpu_timeout_secs) + throw EPExceptionWithLineInfo("Dispatch CPU wait", get_buffer_info()); + } + } else { + // Non-cached mode without CPU sync, allocate with the worst case + num_recv_tokens = num_max_tokens_per_rank * nccl_context->num_ranks; + num_expanded_tokens = nccl_context->num_ranks * num_max_tokens_per_rank * std::min(num_topk, num_local_experts); + num_expanded_tokens += (expert_alignment - 1) * num_local_experts; + num_expanded_tokens = math::align(num_expanded_tokens, expert_alignment); + } + + // Allocate received tensors + // `recv_src_metadata` includes source token indices and buffer slot indices + const auto num_allocated_tokens = do_expand ? num_expanded_tokens : num_recv_tokens; + auto recv_x = torch::empty({num_allocated_tokens, hidden}, x.options()); + auto recv_sf = std::optional(); + auto recv_topk_idx = std::optional(); + auto recv_topk_weights = std::optional(); + auto recv_src_metadata = cached_mode ? + cached_recv_src_metadata.value() : + torch::empty({num_recv_tokens, num_topk + 2}, + torch::TensorOptions(torch::kCUDA).dtype(torch::kInt)); + + // Optional tensors + void* recv_sf_ptr = nullptr; + topk_idx_t* recv_topk_idx_ptr = nullptr; + float* recv_topk_weights_ptr = nullptr; + int recv_sf_token_stride = 0, recv_sf_hidden_stride = 0; + if (sf.has_value()) { + if (not use_tma_aligned_col_major_sf) { + recv_sf_token_stride = num_sf_packs, recv_sf_hidden_stride = 1; + } else { + // TMA-aligned layout for the next GEMM input + recv_sf_token_stride = 1, recv_sf_hidden_stride = math::align(num_allocated_tokens, kNumAlignedSFPacks); + } + recv_sf = torch::empty_strided({num_allocated_tokens, num_sf_packs}, + {recv_sf_token_stride, recv_sf_hidden_stride}, + sf->options()); + recv_sf_ptr = recv_sf->data_ptr(); + } + if (not do_expand) { + recv_topk_idx = torch::empty({num_allocated_tokens, num_topk}, topk_idx.options()); + recv_topk_idx_ptr = recv_topk_idx->data_ptr(); + } + if (topk_weights.has_value()) { + recv_topk_weights = do_expand ? + torch::empty({num_allocated_tokens}, topk_weights->options()) : + torch::empty({num_allocated_tokens, num_topk}, topk_weights->options()); + recv_topk_weights_ptr = recv_topk_weights->data_ptr(); + } + + // Process prefix sum, in expanding mode, it is also atomic counters + if (not cached_mode) { + if (do_expand) { + // Slice the exclusive part and do atomic additions into inclusive + psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert.slice(0, 0, num_local_experts); + } else { + // Slice the inclusive part (and will not be used in the epilogue) + psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert.slice(0, 1, num_local_experts + 1); + } + } + EP_HOST_ASSERT(psum_num_recv_tokens_per_expert.size(0) == num_local_experts); + + // Launch copy kernels with full SMs + stream_control_before_epilogue(previous_event_before_epilogue); + launch_dispatch_copy_epilogue(buffer, workspace, + psum_num_recv_tokens_per_scaleup_rank.data_ptr(), + psum_num_recv_tokens_per_expert.data_ptr(), + recv_x.data_ptr(), recv_sf_ptr, + recv_topk_idx_ptr, recv_topk_weights_ptr, + recv_src_metadata.data_ptr(), + channel_linked_list_ptr, + num_unaligned_recv_tokens_per_expert_ptr, + num_recv_tokens, num_max_tokens_per_rank, + num_hidden_bytes, + num_sf_packs, recv_sf_token_stride, recv_sf_hidden_stride, + num_experts, num_topk, expert_alignment, + nccl_context->scaleout_rank_idx, nccl_context->scaleup_rank_idx, + nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks, + jit::device_runtime->get_num_sms(), + jit::device_runtime->get_num_smem_bytes(), + num_channels, + do_expand, cached_mode, + do_zero_padding, + comm_stream); + + // Stream control + const auto event = stream_control_epilogue( + {x, sf, topk_idx, topk_weights, + recv_x, recv_sf, recv_topk_idx, recv_topk_weights, + cumulative_local_expert_recv_stats, + copied_topk_idx, + psum_num_recv_tokens_per_scaleup_rank, + psum_num_recv_tokens_per_expert, + num_unaligned_recv_tokens_per_expert, + recv_src_metadata, + dst_buffer_slot_idx, + token_metadata_at_forward, + channel_linked_list}, + compute_stream, + allocate_on_comm_stream, async_with_compute_stream); + + return {recv_x, recv_sf, + recv_topk_idx, recv_topk_weights, + copied_topk_idx, + num_recv_tokens, num_expanded_tokens, + num_recv_tokens_per_expert_list, + psum_num_recv_tokens_per_scaleup_rank, + psum_num_recv_tokens_per_expert, + num_unaligned_recv_tokens_per_expert, + recv_src_metadata, + dst_buffer_slot_idx, + token_metadata_at_forward, + channel_linked_list, + event}; + } + + std::tuple, std::optional> + combine(const torch::Tensor& x, + const std::optional& topk_weights, + const std::optional& bias_0, + const std::optional& bias_1, + const torch::Tensor& src_metadata, + const torch::Tensor& combined_topk_idx, + const torch::Tensor& psum_num_recv_tokens_per_scaleup_rank, + const std::optional& token_metadata_at_forward, + const std::optional& channel_linked_list, + const int& num_experts, + const int& num_max_tokens_per_rank, + const int& num_sms, const int& num_qps, + const std::optional& previous_event, + const std::optional& previous_event_before_epilogue, + const bool& async_with_compute_stream, + const bool& allocate_on_comm_stream, + const bool& use_expanded_layout) const { + // Check SM count + EP_HOST_ASSERT(num_sms > 0); + + // Check data + const auto [num_tokens, hidden] = get_shape<2>(x); + EP_HOST_ASSERT(x.is_cuda() and x.is_contiguous()); + EP_HOST_ASSERT(x.scalar_type() == torch::kBFloat16); + EP_HOST_ASSERT((x.size(1) * x.element_size()) % sizeof(int4) == 0); + + // Check tensors at dispatch + const auto [num_combined_tokens, num_topk] = get_shape<2>(combined_topk_idx); + const auto [num_scaleup_ranks] = get_shape<1>(psum_num_recv_tokens_per_scaleup_rank); + EP_HOST_ASSERT(combined_topk_idx.is_cuda() and combined_topk_idx.is_contiguous()); + EP_HOST_ASSERT(combined_topk_idx.scalar_type() == c10::CppTypeToScalarType::value); + EP_HOST_ASSERT(num_scaleup_ranks == nccl_context->num_scaleup_ranks); + EP_HOST_ASSERT(psum_num_recv_tokens_per_scaleup_rank.is_cuda() and psum_num_recv_tokens_per_scaleup_rank.is_contiguous()); + EP_HOST_ASSERT(psum_num_recv_tokens_per_scaleup_rank.scalar_type() == torch::kInt); + EP_HOST_ASSERT(num_combined_tokens <= num_max_tokens_per_rank); + + // Check metadata + // For reduction mode, `num_tokens_` means the number of unexpanded tokens + const auto [num_reduced_tokens, num_topk_p2] = get_shape<2>(src_metadata); + EP_HOST_ASSERT(num_reduced_tokens == (use_expanded_layout ? num_reduced_tokens : num_tokens)); + EP_HOST_ASSERT(num_topk_p2 == num_topk + 2); + EP_HOST_ASSERT(src_metadata.is_cuda() and src_metadata.is_contiguous()); + EP_HOST_ASSERT(src_metadata.scalar_type() == torch::kInt); + + // Check optional tensors + if (topk_weights.has_value()) { + if (use_expanded_layout) { + const auto [num_tokens__] = get_shape<1>(topk_weights.value()); + EP_HOST_ASSERT(num_tokens == num_tokens__); + } else { + const auto [num_tokens__, num_topk__] = get_shape<2>(topk_weights.value()); + EP_HOST_ASSERT(num_tokens == num_tokens__ and num_topk == num_topk__); + } + EP_HOST_ASSERT(topk_weights->is_cuda() and topk_weights->is_contiguous()); + EP_HOST_ASSERT(topk_weights->scalar_type() == torch::kFloat); + } + + const 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(); + } + } + + // Stream control + // All new tensor allocations should happen after this + const auto compute_stream = stream_control_prologue(previous_event, allocate_on_comm_stream, async_with_compute_stream); + + // Check buffer size + EP_HOST_ASSERT(get_combine_buffer_size(num_max_tokens_per_rank, hidden, num_topk, + nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks, + nccl_context->is_scaleup_nvlink, allow_multiple_reduction) <= num_buffer_bytes); + + // Optional configs and metadata for hybrid combine + int num_channels = 1; + int* token_metadata_at_forward_ptr = nullptr; + int* channel_linked_list_ptr = nullptr; + if (nccl_context->num_scaleout_ranks > 1) { + // The token metadata during forward + const auto [num_channels_, d1, d2] = get_shape<3>(token_metadata_at_forward.value()); + const auto num_max_tokens_per_channel = math::ceil_div(num_max_tokens_per_rank, num_channels_); + num_channels = num_channels_; + token_metadata_at_forward_ptr = token_metadata_at_forward->data_ptr(); + EP_HOST_ASSERT(d1 == nccl_context->num_scaleout_ranks * num_max_tokens_per_channel + 1); + EP_HOST_ASSERT(d2 == 2 + num_topk * 2); + EP_HOST_ASSERT(token_metadata_at_forward->is_cuda() and token_metadata_at_forward->is_contiguous()); + EP_HOST_ASSERT(token_metadata_at_forward->scalar_type() == torch::kInt); + + // Per-scaleup-peer-per-channel linked list + const auto [num_channels__, d1_, d2_] = get_shape<3>(channel_linked_list.value()); + channel_linked_list_ptr = channel_linked_list->data_ptr(); + EP_HOST_ASSERT(num_channels == num_channels__); + EP_HOST_ASSERT(d1_ == nccl_context->num_scaleout_ranks * num_max_tokens_per_channel + 1); + EP_HOST_ASSERT(d2_ == nccl_context->num_scaleup_ranks); + EP_HOST_ASSERT(channel_linked_list->is_cuda() and channel_linked_list->is_contiguous()); + EP_HOST_ASSERT(channel_linked_list->scalar_type() == torch::kInt); + } + + // Push data into remote buffers + // NOTES: we don't use `num_hidden_bytes` due to enable later quantization possibility + const auto reduce_buffer = launch_combine( + x.data_ptr(), + topk_weights.has_value() ? topk_weights->data_ptr() : nullptr, + src_metadata.data_ptr(), + psum_num_recv_tokens_per_scaleup_rank.data_ptr(), + token_metadata_at_forward_ptr, + channel_linked_list_ptr, + nccl_context->dev_comm, nccl_context->window, + buffer, workspace, + num_reduced_tokens, num_max_tokens_per_rank, + hidden, num_experts, num_topk, + num_qps, num_gpu_timeout_cycles, + nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks, + nccl_context->scaleout_rank_idx, nccl_context->scaleup_rank_idx, + nccl_context->is_scaleup_nvlink, + num_sms, jit::device_runtime->get_num_smem_bytes(), + num_channels, + use_expanded_layout, allow_multiple_reduction, + comm_stream); + + // Allocate output tensors + auto combined_x = torch::empty({num_combined_tokens, hidden}, x.options()); + auto combined_topk_weights = std::optional(); + float* combined_topk_weights_ptr = nullptr; + if (topk_weights.has_value()) { + combined_topk_weights = torch::empty({num_combined_tokens, num_topk}, topk_weights->options()); + combined_topk_weights_ptr = combined_topk_weights->data_ptr(); + } + + // Combine pushed data + stream_control_before_epilogue(previous_event_before_epilogue); + launch_combine_reduce_epilogue(combined_x.data_ptr(), + combined_topk_weights_ptr, + combined_topk_idx.data_ptr(), + num_combined_tokens, num_max_tokens_per_rank, + hidden, + num_experts, num_topk, + reduce_buffer, + bias_ptrs[0], bias_ptrs[1], + nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks, + nccl_context->scaleout_rank_idx, nccl_context->scaleup_rank_idx, + jit::device_runtime->get_num_sms(), + jit::device_runtime->get_num_smem_bytes(), + use_expanded_layout, allow_multiple_reduction, + comm_stream); + + // Stream control + const auto event = stream_control_epilogue( + {x, topk_weights, bias_0, bias_1, + src_metadata, + combined_topk_idx, + combined_x, combined_topk_weights, + psum_num_recv_tokens_per_scaleup_rank, + token_metadata_at_forward, + channel_linked_list}, + compute_stream, + allocate_on_comm_stream, async_with_compute_stream); + return {combined_x, combined_topk_weights, event}; + } +}; + +static void register_apis(pybind11::module_& m) { + pybind11::class_(m, "ElasticBuffer") + .def(pybind11::init()) + .def("destroy", &ElasticBuffer::destroy) + .def("get_comm_stream", &ElasticBuffer::get_comm_stream) + .def("get_physical_domain_size", &ElasticBuffer::get_physical_domain_size) + .def("get_logical_domain_size", &ElasticBuffer::get_logical_domain_size) + .def("barrier", &ElasticBuffer::barrier) + .def("engram_write", &ElasticBuffer::engram_write) + .def("engram_fetch", &ElasticBuffer::engram_fetch) + .def("pp_set_config", &ElasticBuffer::pp_set_config) + .def("pp_send", &ElasticBuffer::pp_send) + .def("pp_recv", &ElasticBuffer::pp_recv) + .def("create_agrs_session", &ElasticBuffer::create_agrs_session) + .def("destroy_agrs_session", &ElasticBuffer::destroy_agrs_session) + .def("agrs_set_config", &ElasticBuffer::agrs_set_config) + .def("agrs_get_inplace_tensor", &ElasticBuffer::agrs_get_inplace_tensor) + .def("all_gather", &ElasticBuffer::all_gather) + .def("dispatch", &ElasticBuffer::dispatch) + .def("combine", &ElasticBuffer::combine); + m.def("create_cpu_handle", &ElasticBuffer::create_cpu_handle); + m.def("calculate_elastic_buffer_size", &ElasticBuffer::calculate_buffer_size); + m.def("get_elastic_buffer_alignment", [=]() { + return symmetric::kNumAlignmentBytes; + }); + + // NCCL communicator handle + m.def("get_local_nccl_unique_id", &nccl::get_local_unique_id); + m.def("create_nccl_comm", &nccl::create_nccl_comm); + m.def("destroy_nccl_comm", &nccl::destroy_nccl_comm); + + // Communication domain utilities + m.def("get_physical_domain_size", &nccl::get_physical_domain_size); + m.def("get_logical_domain_size", &nccl::get_logical_domain_size); +} + +} // namespace deep_ep diff --git a/csrc/elastic/utils.hpp b/csrc/elastic/utils.hpp new file mode 100644 index 0000000..f919b57 --- /dev/null +++ b/csrc/elastic/utils.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include +#include + +namespace deep_ep::elastic { + +static at::cuda::CUDAStream get_global_comm_stream() { + static std::optional comm_stream = std::nullopt; + if (not comm_stream.has_value()) + comm_stream = at::cuda::getStreamFromPool(true); + return comm_stream.value(); +} + +template +static auto get_shape(const torch::Tensor& t) { + EP_HOST_ASSERT(t.dim() == kNumDims); + return [&t] (std::index_sequence) { + return std::make_tuple(static_cast(t.sizes()[Is])...); + }(std::make_index_sequence()); +} + +template +static dtype_t* get_data_ptr(const std::optional& t) { + return t.has_value() ? t->data_ptr() : nullptr; +} + +} // deep_ep::elastic diff --git a/csrc/indexing/main.cu b/csrc/indexing/main.cu new file mode 100644 index 0000000..092c05e --- /dev/null +++ b/csrc/indexing/main.cu @@ -0,0 +1,21 @@ +// Utils +#include + +// EP +#include +#include +#include +#include +#include +#include + +// Engram +#include +#include + +// PP +#include + +int main() { + return 0; +} diff --git a/csrc/jit/api.hpp b/csrc/jit/api.hpp new file mode 100644 index 0000000..8332d91 --- /dev/null +++ b/csrc/jit/api.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "compiler.hpp" +#include "include_parser.hpp" +#include "kernel_runtime.hpp" + +namespace deep_ep::jit { + +static void init(const std::string& library_root_path, + const std::string& cuda_home_path_by_python, const std::string& nccl_root_path_by_python) { + Compiler::prepare_init(library_root_path, cuda_home_path_by_python, nccl_root_path_by_python); + KernelRuntime::prepare_init(cuda_home_path_by_python); + IncludeParser::prepare_init(library_root_path); +} + +static void register_apis(pybind11::module_& m) { + m.def("init_jit", &init); +} + +} // namespace deep_ep::jit diff --git a/csrc/jit/cache.hpp b/csrc/jit/cache.hpp new file mode 100644 index 0000000..96d89e5 --- /dev/null +++ b/csrc/jit/cache.hpp @@ -0,0 +1,34 @@ +#pragma once + +#include +#include +#include + +#include "kernel_runtime.hpp" + +namespace deep_ep::jit { + +class KernelRuntimeCache { + std::unordered_map> cache; + +public: + KernelRuntimeCache() = default; + + void clear() { + cache.clear(); + } + + std::shared_ptr get(const std::filesystem::path& dir_path) { + // Hit the runtime cache + if (const auto iterator = cache.find(dir_path); iterator != cache.end()) + return iterator->second; + + if (KernelRuntime::check_validity(dir_path)) + return cache[dir_path] = std::make_shared(dir_path); + return nullptr; + } +}; + +static auto kernel_runtime_cache = std::make_shared(); + +} // namespace deep_ep::jit diff --git a/csrc/jit/compiler.hpp b/csrc/jit/compiler.hpp new file mode 100644 index 0000000..ad01b3c --- /dev/null +++ b/csrc/jit/compiler.hpp @@ -0,0 +1,269 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "../utils/format.hpp" +#include "../utils/hash.hpp" +#include "../utils/lazy_init.hpp" +#include "../utils/system.hpp" +#include "cache.hpp" +#include "device_runtime.hpp" + +namespace deep_ep::jit { + +class Compiler { +public: + static std::filesystem::path library_root_path; + static std::filesystem::path library_include_path; + static std::filesystem::path cuda_home; + static std::filesystem::path nccl_root; + static std::filesystem::path cuobjdump_path; + + static void prepare_init(const std::string& library_root_path, + const std::string& cuda_home_path_by_python, + const std::string& nccl_root_path_by_python) { + // NOTES: if you are adding some third-party includes for kernels, please add its hash value + Compiler::library_root_path = library_root_path; + Compiler::library_include_path = Compiler::library_root_path / "include"; + Compiler::cuda_home = cuda_home_path_by_python; + Compiler::nccl_root = nccl_root_path_by_python; + Compiler::cuobjdump_path = Compiler::cuda_home / "bin" / "cuobjdump"; + } + + std::string signature, flags; + std::filesystem::path cache_dir_path; + + Compiler() { + EP_HOST_ASSERT(not library_root_path.empty()); + EP_HOST_ASSERT(not library_include_path.empty()); + EP_HOST_ASSERT(not cuda_home.empty()); + EP_HOST_ASSERT(not nccl_root.empty()); + EP_HOST_ASSERT(not cuobjdump_path.empty()); + + // Cache settings + cache_dir_path = std::filesystem::path(get_env("HOME")) / ".deep_ep"; + if (const auto env_cache_dir_path = get_env("EP_JIT_CACHE_DIR"); not env_cache_dir_path.empty()) + cache_dir_path = env_cache_dir_path; + + // The compiler flags applied to all derived compilers + signature = "unknown-compiler"; + flags = fmt::format("-std=c++{} --diag-suppress=39,161,174,177,186,940,3012 " + "--ptxas-options=--register-usage-level=10", + get_env("EP_JIT_CPP_STANDARD", 20)); + if (get_env("EP_JIT_DEBUG", 0) or get_env("EP_JIT_PTXAS_VERBOSE", 0)) + flags += " --ptxas-options=--verbose"; + if (get_env("EP_JIT_DEBUG", 0) or get_env("EP_JIT_WITH_LINEINFO", 0)) + flags += " -Xcompiler -rdynamic -lineinfo"; + if (get_env("EP_GIN_GDAKI_DEBUG", 0)) + flags += " -DNCCL_DEVICE_GIN_GDAKI_ENABLE_DEBUG=1"; + flags += fmt::format(" -I {}/include", nccl_root.c_str()); + + // Some special flags for EP + // TODO: make it more general, e.g. `EP_JIT_EXTRA_FLAGS` + if (int num_topk_idx_bits = get_env("EP_NUM_TOPK_IDX_BITS", 0); num_topk_idx_bits != 0) + flags += fmt::format(" -DEP_NUM_TOPK_IDX_BITS={}", num_topk_idx_bits); + } + + virtual ~Compiler() = default; + + std::filesystem::path make_tmp_dir() const { + return make_dirs(cache_dir_path / "tmp"); + } + + static void fsync_path(const std::filesystem::path& path) { + const auto fd = ::open(path.c_str(), O_RDONLY); + if (fd >= 0) { + ::fsync(fd); + ::close(fd); + } + } + + // Recursively fsync a directory: files and subdirectories first (bottom-up), then the directory itself + // NOTES: ensures data and directory entries are visible on other nodes in distributed filesystems + static void fsync_dir(const std::filesystem::path& dir_path) { // NOLINT(*-no-recursion) + for (const auto& entry: std::filesystem::directory_iterator(dir_path)) { + if (entry.is_directory()) + fsync_dir(entry.path()); + else if (entry.is_regular_file()) + fsync_path(entry.path()); + } + fsync_path(dir_path); + } + + static void put(const std::filesystem::path& path, const std::string& data) { + std::ofstream out(path, std::ios::binary); + EP_HOST_ASSERT(out.write(data.data(), data.size())); + out.close(); + + // NOTES: fsync to ensure the data is visible to other processes (e.g., NVCC) + // on distributed filesystems, where `close()` alone does not guarantee persistence + fsync_path(path); + } + + std::shared_ptr build(const std::string& name, const std::string& code) const { + const auto kernel_signature = fmt::format("{}$${}$${}$${}", name, signature, flags, code); + const auto dir_path = cache_dir_path / "cache" / fmt::format("kernel.{}.{}", name, get_hex_digest(kernel_signature)); + + // Hit the runtime cache + if (const auto runtime = kernel_runtime_cache->get(dir_path); runtime != nullptr) + return runtime; + + // Compile into a temporary directory, then atomically rename the whole directory + // NOTES: renaming a directory is atomic on both local and distributed filesystems, + // avoiding the stale inode issue that occurs when renaming individual files + const auto tmp_dir_path = make_tmp_dir() / get_uuid(); + make_dirs(tmp_dir_path); + + // Compile into the temporary directory + const auto tmp_cubin_path = tmp_dir_path / "kernel.cubin"; + if (get_env("EP_JIT_DUMP_ASM") or get_env("EP_JIT_DUMP_PTX")) { + const auto tmp_ptx_path = tmp_dir_path / "kernel.ptx"; + compile(code, tmp_dir_path, tmp_cubin_path, tmp_ptx_path); + } else { + compile(code, tmp_dir_path, tmp_cubin_path); + } + + // Disassemble if needed + if (get_env("EP_JIT_DUMP_ASM") or get_env("EP_JIT_DUMP_SASS")) { + const auto tmp_sass_path = tmp_dir_path / "kernel.sass"; + disassemble(tmp_cubin_path, tmp_sass_path); + } + + // Fsync before rename to ensure visibility on distributed filesystems + fsync_dir(tmp_dir_path); + + // Atomically rename the temporary directory to the final cache path + // NOTES: if another rank already created dir_path, rename will fail — that's fine + make_dirs(dir_path.parent_path()); + std::error_code error_code; + std::filesystem::rename(tmp_dir_path, dir_path, error_code); + if (error_code) { + // Another rank beat us, then clean up our dir and use the existing one + // NOTES: avoid `std::filesystem::remove_all` here — it can segfault on + // distributed filesystems, when concurrent processes operate + // on the same parent directory, causing stale directory entries + safe_remove_all(tmp_dir_path); + } + + // Put into the runtime cache + const auto runtime = kernel_runtime_cache->get(dir_path); + EP_HOST_ASSERT(runtime != nullptr); + return runtime; + } + + static void disassemble(const std::filesystem::path &cubin_path, const std::filesystem::path &sass_path) { + // Disassemble the CUBIN file to SASS + const auto command = fmt::format("{} --dump-sass {} > {}", cuobjdump_path.c_str(), cubin_path.c_str(), sass_path.c_str()); + if (get_env("EP_JIT_DEBUG", 0) or get_env("EP_JIT_PRINT_COMPILER_COMMAND", 0)) + printf("Running cuobjdump command: %s\n", command.c_str()); + const auto [return_code, output] = call_external_command(command); + if (return_code != 0) { + printf("cuobjdump failed: %s\n", output.c_str()); + EP_HOST_ASSERT(false and "cuobjdump failed"); + } + } + + virtual void compile(const std::string &code, const std::filesystem::path& dir_path, const std::filesystem::path &cubin_path, const std::optional &ptx_path = std::nullopt) const = 0; +}; + +EP_DECLARE_STATIC_VAR_IN_CLASS(Compiler, library_root_path); +EP_DECLARE_STATIC_VAR_IN_CLASS(Compiler, library_include_path); +EP_DECLARE_STATIC_VAR_IN_CLASS(Compiler, cuda_home); +EP_DECLARE_STATIC_VAR_IN_CLASS(Compiler, nccl_root); +EP_DECLARE_STATIC_VAR_IN_CLASS(Compiler, cuobjdump_path); + +class NVCCCompiler final: public Compiler { + std::filesystem::path nvcc_path; + + std::pair get_nvcc_version() const { + EP_HOST_ASSERT(std::filesystem::exists(nvcc_path)); + + // Call the version command + const auto command = std::string(nvcc_path) + " --version"; + const auto [return_code, output] = call_external_command(command); + EP_HOST_ASSERT(return_code == 0); + + // The version should be at least 12.3 + int major, minor; + std::smatch match; + EP_HOST_ASSERT(std::regex_search(output, match, std::regex(R"(release (\d+\.\d+))"))); + std::sscanf(match[1].str().c_str(), "%d.%d", &major, &minor); + EP_HOST_ASSERT((major > 12 or (major == 12 and minor >= 3)) and "NVCC version should be >= 12.3"); + return {major, minor}; + } + +public: + NVCCCompiler() { + // Override the compiler signature + nvcc_path = cuda_home / "bin" / "nvcc"; + cuobjdump_path = cuda_home / "bin" / "cuobjdump"; + if (const auto env_nvcc_path = get_env("EP_JIT_NVCC_COMPILER"); not env_nvcc_path.empty()) + nvcc_path = env_nvcc_path; + const auto [nvcc_major, nvcc_minor] = get_nvcc_version(); + signature = fmt::format("NVCC{}.{}", nvcc_major, nvcc_minor); + + // The override the compiler flags + // Only NVCC >= 12.9 supports arch-specific family suffix + const auto arch = device_runtime->get_arch(false, nvcc_major > 12 or nvcc_minor >= 9); + flags = fmt::format("{} -I{} --gpu-architecture=sm_{} " + "--compiler-options=-fPIC,-O3,-fconcepts,-Wno-deprecated-declarations,-Wno-abi " + "-O3 --expt-relaxed-constexpr --expt-extended-lambda", + flags, library_include_path.c_str(), arch); + } + + void compile(const std::string &code, const std::filesystem::path& dir_path, + const std::filesystem::path &cubin_path, + const std::optional &ptx_path) const override { + // Write the code into the cache directory + const auto code_path = dir_path / "kernel.cu"; + put(code_path, code); + + // Compile to CUBIN + // Avoid cwd files shadowing C++ standard library headers + const auto compile_dir = make_tmp_dir(); + const auto command = fmt::format("cd {} && {} {} -cubin -o {} {}", + compile_dir.c_str(), nvcc_path.c_str(), code_path.c_str(), cubin_path.c_str(), flags); + if (get_env("EP_JIT_DEBUG", 0) or get_env("EP_JIT_PRINT_COMPILER_COMMAND", 0)) + printf("Running NVCC command: %s\n", command.c_str()); + const auto [return_code, output] = call_external_command(command); + if (return_code != 0) { + printf("NVCC compilation failed: %s\n", output.c_str()); + EP_HOST_ASSERT(false and "NVCC compilation failed"); + } + + // Compile to PTX if needed + if (ptx_path.has_value()) { + const auto ptx_command = fmt::format("cd {} && {} {} -ptx -o {} {}", + compile_dir.c_str(), nvcc_path.c_str(), code_path.c_str(), ptx_path->c_str(), flags); + if (get_env("EP_JIT_DEBUG", 0) or get_env("EP_JIT_PRINT_COMPILER_COMMAND", 0)) + printf("Running NVCC PTX command: %s\n", ptx_command.c_str()); + const auto [ptx_return_code, ptx_output] = call_external_command(ptx_command); + if (ptx_return_code != 0) { + printf("NVCC PTX compilation failed: %s\n", ptx_output.c_str()); + EP_HOST_ASSERT(false and "NVCC PTX compilation failed"); + } + } + + // Check local memory usage + if (get_env("EP_JIT_PTXAS_CHECK", 0)) + EP_HOST_ASSERT(not std::regex_search(output, std::regex(R"(Local memory used)"))); + + // Print PTXAS log + if (get_env("EP_JIT_DEBUG", 0) or get_env("EP_JIT_PTXAS_VERBOSE", 0)) + printf("%s", output.c_str()); + } +}; + +static auto compiler = LazyInit([]() -> std::shared_ptr { + return std::make_shared(); +}); + +} // namespace deep_ep::jit diff --git a/csrc/jit/device_runtime.hpp b/csrc/jit/device_runtime.hpp new file mode 100644 index 0000000..3ed5328 --- /dev/null +++ b/csrc/jit/device_runtime.hpp @@ -0,0 +1,67 @@ +#pragma once + +#include + +#include "../utils/lazy_init.hpp" + +namespace deep_ep::jit { + +class DeviceRuntime { + int64_t cached_clock_rate = 0; + std::shared_ptr cached_prop; + + std::shared_ptr get_prop() { + if (cached_prop == nullptr) { + int device_idx; + cudaDeviceProp prop; + CUDA_RUNTIME_CHECK(cudaGetDevice(&device_idx)); + CUDA_RUNTIME_CHECK(cudaGetDeviceProperties(&prop, device_idx)); + cached_prop = std::make_shared(prop); + } + return cached_prop; + } + +public: + int64_t get_clock_rate() { + if (cached_clock_rate == 0) { + // NOTES: we should convert kHz into Hz + int device_idx, rate; + CUDA_RUNTIME_CHECK(cudaGetDevice(&device_idx)); + CUDA_RUNTIME_CHECK(cudaDeviceGetAttribute(&rate, cudaDevAttrClockRate, device_idx)); + cached_clock_rate = static_cast(rate) * 1000ll; + } + return cached_clock_rate; + } + + int get_num_smem_bytes() { + return static_cast(get_prop()->sharedMemPerBlockOptin); + } + + int get_num_sms() { + return get_prop()->multiProcessorCount; + } + + std::pair get_arch_pair() { + const auto prop = get_prop(); + return {prop->major, prop->minor}; + } + + std::string get_arch(const bool& number_only = false, + const bool& support_arch_family = false) { + const auto [major, minor] = get_arch_pair(); + if (major == 10 and minor != 1) { + if (number_only) + return "100"; + return support_arch_family ? "100f" : "100a"; + } + return std::to_string(major * 10 + minor) + (number_only ? "" : "a"); + } + + int get_arch_major() { + return get_arch_pair().first; + } +}; + +static auto device_runtime = LazyInit([](){ return std::make_shared(); }); + +} // namespace deep_ep::jit diff --git a/csrc/jit/handle.hpp b/csrc/jit/handle.hpp new file mode 100644 index 0000000..ba135e4 --- /dev/null +++ b/csrc/jit/handle.hpp @@ -0,0 +1,155 @@ +#pragma once + +#include +#include +#include + +#include "../utils/lazy_driver.hpp" + +namespace deep_ep::jit { + +#if CUDART_VERSION >= 12080 and defined(EP_JIT_USE_RUNTIME_API) + +// Use CUDA runtime API +using LibraryHandle = cudaLibrary_t; +using KernelHandle = cudaKernel_t; +using LaunchConfigHandle = cudaLaunchConfig_t; +using LaunchAttrHandle = cudaLaunchAttribute; + +#define EP_CUDA_UNIFIED_CHECK CUDA_RUNTIME_CHECK + +static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const std::string& func_name, + LibraryHandle *library_opt = nullptr) { + LibraryHandle library; + KernelHandle kernel{}; + CUDA_RUNTIME_CHECK(cudaLibraryLoadFromFile(&library, cubin_path.c_str(), nullptr, nullptr, 0, nullptr, nullptr, 0)); + nvshmemx_culibrary_init(library); + CUDA_RUNTIME_CHECK(cudaLibraryGetKernel(&kernel, library, func_name.c_str())); + + if (library_opt != nullptr) + *library_opt = library; + return kernel; +} + +static void unload_library(const LibraryHandle& library) { + nvshmemx_culibrary_finalize(library); + const auto error = cudaLibraryUnload(library); + EP_HOST_ASSERT(error == cudaSuccess or error == cudaErrorCudartUnloading); +} + +static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel, + const cudaStream_t& stream, const int& smem_size, + const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim) { + if (smem_size > 0) + CUDA_RUNTIME_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + LaunchConfigHandle config; + config.gridDim = grid_dim; + config.blockDim = block_dim; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + config.numAttrs = 0; + config.attrs = nullptr; + + // TODO: support cooperative and dependent kernel launch + // NOTES: must use `static` or the `attr` will be deconstructed + static LaunchAttrHandle attr; + if (cluster_dim > 1) { + attr.id = cudaLaunchAttributeClusterDimension; + attr.val.clusterDim = {static_cast(cluster_dim), 1, 1}; + config.attrs = &attr; + config.numAttrs = 1; + } + return config; +} + +template +static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) { + void *ptr_args[] = { &args... }; + return cudaLaunchKernelExC(&config, kernel, ptr_args); +} + +#else + +// Use CUDA driver API +using LibraryHandle = CUmodule; +using KernelHandle = CUfunction; +using LaunchConfigHandle = CUlaunchConfig; +using LaunchAttrHandle = CUlaunchAttribute; + +#define EP_CUDA_UNIFIED_CHECK CUDA_DRIVER_CHECK + +static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const std::string& func_name, + LibraryHandle *library_opt = nullptr) { + LibraryHandle library; + KernelHandle kernel; + CUDA_DRIVER_CHECK(lazy_cuModuleLoad(&library, cubin_path.c_str())); + CUDA_DRIVER_CHECK(lazy_cuModuleGetFunction(&kernel, library, func_name.c_str())); + + if (library_opt != nullptr) + *library_opt = library; + return kernel; +} + +static void unload_library(const LibraryHandle& library) { + const auto error = lazy_cuModuleUnload(library); + EP_HOST_ASSERT(error == CUDA_SUCCESS or error == CUDA_ERROR_DEINITIALIZED); +} + +static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel, + const cudaStream_t& stream, const int& smem_size, + const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim, + const bool& cooperative, const bool& enable_pdl) { + if (smem_size > 0) + CUDA_DRIVER_CHECK(lazy_cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size)); + + LaunchConfigHandle config; + config.gridDimX = grid_dim.x; + config.gridDimY = grid_dim.y; + config.gridDimZ = grid_dim.z; + config.blockDimX = block_dim.x; + config.blockDimY = block_dim.y; + config.blockDimZ = block_dim.z; + config.sharedMemBytes = smem_size; + config.hStream = stream; + + // Create attributes + static LaunchAttrHandle attrs[3]; + config.attrs = attrs; + config.numAttrs = 0; + + // Cooperative launch + if (cooperative) { + auto& attr = attrs[config.numAttrs ++]; + attr.id = CU_LAUNCH_ATTRIBUTE_COOPERATIVE; + attr.value.cooperative = 1; + } + + // Cluster size + // NOTES: must use `static` or the `attr` will be deconstructed + if (cluster_dim > 1) { + auto& attr = attrs[config.numAttrs ++]; + attr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attr.value.clusterDim.x = cluster_dim; + attr.value.clusterDim.y = 1; + attr.value.clusterDim.z = 1; + } + + // Dependent kernel launch + if (enable_pdl) { + auto& attr = attrs[config.numAttrs ++]; + attr.id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attr.value.programmaticStreamSerializationAllowed = 1; + } + return config; +} + +template +static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) { + void *ptr_args[] = { &args... }; + return lazy_cuLaunchKernelEx(&config, kernel, ptr_args, nullptr); +} + +#endif + +} // namespace deep_ep::jit diff --git a/csrc/jit/include_parser.hpp b/csrc/jit/include_parser.hpp new file mode 100644 index 0000000..6df87ff --- /dev/null +++ b/csrc/jit/include_parser.hpp @@ -0,0 +1,80 @@ +#pragma once + +#include +#include +#include +#include + +#include "../utils/format.hpp" +#include "../utils/system.hpp" + +namespace deep_ep::jit { + +class IncludeParser { + std::unordered_map> cache; + + static std::vector get_includes(const std::string& code, const std::filesystem::path& file_path = "") { + std::vector includes; + const std::regex pattern(R"(#\s*include\s*[<"][^>"]+[>"])"); + std::sregex_iterator iter(code.begin(), code.end(), pattern); + const std::sregex_iterator end; + + // TODO: parse relative paths as well + for (; iter != end; ++ iter) { + const auto include_str = iter->str(); + const int len = include_str.length(); + if (include_str.substr(0, 10) == "#include <" and include_str[len - 1] == '>' and include_str[10] != ' ' and include_str[len - 2] != ' ') { + std::string filename = include_str.substr(10, len - 11); + if (filename.substr(0, 7) == "deep_ep") // We only parse `` + includes.push_back(filename); + } else { + std::string error_info = fmt::format("Non-standard include: {}", include_str); + if (file_path != "") + error_info += fmt::format(" ({})", file_path.string()); + EP_HOST_UNREACHABLE(error_info); + } + } + return includes; + } + +public: + static std::filesystem::path library_include_path; + + static void prepare_init(const std::string& library_root_path) { + library_include_path = std::filesystem::path(library_root_path) / "include"; + } + + std::string get_hash_value(const std::string& code, const bool& exclude_code = true) { + std::stringstream ss; + for (const auto& i: get_includes(code)) + ss << get_hash_value_by_path(library_include_path / i) << "$"; + if (not exclude_code) + ss << "#" << get_hex_digest(code); + return get_hex_digest(ss.str()); + } + + std::string get_hash_value_by_path(const std::filesystem::path& path) { + // Check whether hit in cache + // ReSharper disable once CppUseAssociativeContains + if (cache.count(path) > 0) { + const auto opt = cache[path]; + if (not opt.has_value()) + EP_HOST_UNREACHABLE(fmt::format("Circular include may occur: {}", path.string())); + return opt.value(); + } + + // Read file and calculate hash recursively + std::ifstream in(path); + if (not in.is_open()) + EP_HOST_UNREACHABLE(fmt::format("Failed to open: {}", path.string())); + std::string code((std::istreambuf_iterator(in)), std::istreambuf_iterator()); + cache[path] = std::nullopt; + return (cache[path] = get_hash_value(code, false)).value(); + } +}; + +EP_DECLARE_STATIC_VAR_IN_CLASS(IncludeParser, library_include_path); + +static auto include_parser = std::make_shared(); + +} // namespace deep_ep::jit diff --git a/csrc/jit/kernel_runtime.hpp b/csrc/jit/kernel_runtime.hpp new file mode 100644 index 0000000..e610959 --- /dev/null +++ b/csrc/jit/kernel_runtime.hpp @@ -0,0 +1,87 @@ +#pragma once + +#include + +#include + +#include "../utils/format.hpp" +#include "../utils/lazy_init.hpp" +#include "handle.hpp" + +namespace deep_ep::jit { + +class KernelRuntime final { +public: + static std::filesystem::path cuda_home; + + LibraryHandle library; + KernelHandle kernel; + + explicit KernelRuntime(const std::filesystem::path& dir_path) { + EP_HOST_ASSERT(not cuda_home.empty()); + + // NOLINT(*-pro-type-member-init) + const auto cuobjdump_path = cuda_home / "bin" / "cuobjdump"; + const auto cubin_path = dir_path / "kernel.cubin"; + if (get_env("EP_JIT_DEBUG")) + printf("Loading CUBIN: %s\n", cubin_path.c_str()); + + // Find the only symbol + // TODO: use kernel enumeration for newer drivers + const std::vector illegal_names = {"vprintf", "__instantiate_kernel", "__internal", "__assertfail"}; + const auto [exit_code, symbols] = call_external_command(fmt::format("{} -symbols {}", cuobjdump_path.c_str(), cubin_path.c_str())); + EP_HOST_ASSERT(exit_code == 0); + std::istringstream iss(symbols); + std::vector symbol_names; + for (std::string line; std::getline(iss, line); ) { + if (line.find("STT_FUNC") == 0 and line.find("STO_ENTRY") != std::string::npos and + std::none_of(illegal_names.begin(), illegal_names.end(), + [&](const auto name) { return line.find(name) != std::string::npos; })) { + const auto last_space = line.rfind(' '); + symbol_names.push_back(line.substr(last_space + 1)); + } + } + + // Print symbols + if (symbol_names.size() != 1) { + printf("Corrupted JIT cache directory (expected 1 kernel symbol, found %zu): %s, " + "please run `rm -rf %s` and restart your task.\n", + symbol_names.size(), dir_path.c_str(), dir_path.c_str()); + printf("Symbol names: "); + for (const auto& symbol: symbol_names) + printf("%s, ", symbol.c_str()); + printf("\n"); + EP_HOST_ASSERT(false and "Corrupted JIT cache directory"); + } + + // Load from the library + kernel = load_kernel(cubin_path, symbol_names[0], &library); + } + + static void prepare_init(const std::string& cuda_home_path_by_python) { + cuda_home = cuda_home_path_by_python; + } + + static bool check_validity(const std::filesystem::path& dir_path) { + if (not std::filesystem::exists(dir_path)) + return false; + // NOTES: if the directory exists, kernel.cu and kernel.cubin must both exist, + // because the directory is created atomically via rename + if (not std::filesystem::exists(dir_path / "kernel.cu") or + not std::filesystem::exists(dir_path / "kernel.cubin")) { + printf("Corrupted JIT cache directory (missing kernel.cu or kernel.cubin): %s, " + "please run `rm -rf %s` and restart your task.\n", + dir_path.c_str(), dir_path.c_str()); + EP_HOST_ASSERT(false and "Corrupted JIT cache directory"); + } + return true; + } + + ~KernelRuntime() noexcept(false) { + unload_library(library); + } +}; + +EP_DECLARE_STATIC_VAR_IN_CLASS(KernelRuntime, cuda_home); + +} // namespace deep_ep::jit diff --git a/csrc/jit/launch_runtime.hpp b/csrc/jit/launch_runtime.hpp new file mode 100644 index 0000000..c01c75a --- /dev/null +++ b/csrc/jit/launch_runtime.hpp @@ -0,0 +1,74 @@ +#pragma once + +#include + +#include + +#include "../utils/format.hpp" +#include "../utils/system.hpp" +#include "compiler.hpp" +#include "include_parser.hpp" + +namespace deep_ep::jit { + +struct LaunchArgs { + std::pair grid_dim; + int num_threads; + int smem_size; + int cluster_dim; + bool cooperative; + bool pdl_enabled; + + LaunchArgs(const int& grid_dim_x, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1, const bool& cooperative = false, const bool& pdl_enabled = false): + grid_dim({grid_dim_x, 1}), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim), cooperative(cooperative), pdl_enabled(pdl_enabled) {} + + LaunchArgs(const std::pair& grid_dim, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1, const bool& cooperative = false, const bool& pdl_enabled = false): + grid_dim(grid_dim), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim), cooperative(cooperative), pdl_enabled(pdl_enabled) {} +}; + +template +class LaunchRuntime { +public: + template + static std::string generate(const Args& args) { + auto code = Derived::generate_impl(args); + + // NOTES: we require that `generate_impl`'s includes never change + static std::string include_hash; + if (include_hash.empty()) + include_hash = include_parser->get_hash_value(code); + + // TODO: optimize string concat performance + code = fmt::format("// Includes' hash value: {}\n{}", include_hash, code); + if (get_env("EP_JIT_DEBUG", 0)) + printf("Generated kernel code:\n%s\n", code.c_str()); + return code; + } + + template + static void launch(const std::shared_ptr& kernel_runtime, const Args& args, + const std::optional& stream_opt = std::nullopt) { + const auto kernel = kernel_runtime->kernel; + const auto stream = stream_opt.value_or(at::cuda::getCurrentCUDAStream()); + const LaunchArgs& launch_args = args.launch_args; + + const dim3& grid_dim = {static_cast(launch_args.grid_dim.first), + static_cast(launch_args.grid_dim.second), + 1}; + const dim3& block_dim = {static_cast(launch_args.num_threads), 1, 1}; + auto config = construct_launch_config(kernel, stream, launch_args.smem_size, + grid_dim, block_dim, launch_args.cluster_dim, + launch_args.cooperative, launch_args.pdl_enabled); + + // Launch in the derived class + if (get_env("EP_JIT_DEBUG")) { + printf("Launch kernel with {%d, %d} x %d (cooperative: %d), shared memory: %d bytes, cluster: %d, stream: %ld\n", + launch_args.grid_dim.first, launch_args.grid_dim.second, launch_args.num_threads, + launch_args.cooperative, + launch_args.smem_size, launch_args.cluster_dim, stream.id()); + } + Derived::launch_impl(kernel, config, args); + } +}; + +} // namespace deep_ep::jit diff --git a/csrc/kernels/CMakeLists.txt b/csrc/kernels/CMakeLists.txt new file mode 100644 index 0000000..c480371 --- /dev/null +++ b/csrc/kernels/CMakeLists.txt @@ -0,0 +1,15 @@ +function(add_deep_ep_library target_name source_file) + add_library(${target_name} STATIC ${source_file}) + set_target_properties(${target_name} PROPERTIES + POSITION_INDEPENDENT_CODE ON + CXX_STANDARD_REQUIRED ON + CUDA_STANDARD_REQUIRED ON + CXX_STANDARD 17 + CUDA_STANDARD 17 + CUDA_SEPARABLE_COMPILATION ON + ) + target_link_libraries(${target_name} PUBLIC nvshmem_host nvshmem_device cudart cudadevrt mlx5) +endfunction() + +add_subdirectory(legacy) +add_subdirectory(backend) diff --git a/csrc/kernels/backend/CMakeLists.txt b/csrc/kernels/backend/CMakeLists.txt new file mode 100644 index 0000000..1cd71b8 --- /dev/null +++ b/csrc/kernels/backend/CMakeLists.txt @@ -0,0 +1,6 @@ +add_deep_ep_library(runtime_nccl_cuda nccl.cu) +add_deep_ep_library(runtime_nvshmem_cuda nvshmem.cu) +add_deep_ep_library(runtime_cuda_driver cuda_driver.cu) + +# Link these libraries later +set(RUNTIME_CUDA_LIBRARIES runtime_nccl_cuda runtime_nvshmem_cuda runtime_cuda_driver CACHE INTERNAL "Runtime kernels") diff --git a/csrc/kernels/backend/api.cuh b/csrc/kernels/backend/api.cuh new file mode 100644 index 0000000..ba04ad7 --- /dev/null +++ b/csrc/kernels/backend/api.cuh @@ -0,0 +1,104 @@ +#pragma once + +#include +#include +#include + +#include +#include + +#include "symmetric.hpp" + +// TODO: make a unified API +namespace deep_ep::nvshmem { + +std::vector get_unique_id(); + +int init(const std::vector& root_unique_id_val, + const int& rank, + const int& num_ranks, + const int& team_split_stride); + +void* alloc(const size_t& size, const size_t& alignment); + +void free(void* ptr); + +void barrier(const bool& with_cpu_sync, const std::optional& stream_opt = std::nullopt); + +void finalize(); + +} // deep_ep::nvshmem + +namespace deep_ep::nccl { + +pybind11::bytearray get_local_unique_id(); + +int64_t create_nccl_comm(const pybind11::bytearray& root_unique_id_bytes, + const int& num_ranks, const int& rank_idx); + +void destroy_nccl_comm(const int64_t& nccl_comm); + +std::tuple get_physical_domain_size(const int64_t& nccl_comm); + +std::tuple get_logical_domain_size(const int64_t& nccl_comm, const bool& allow_hybrid_mode); + +// TODO: make it header only? +struct NCCLSymmetricMemoryContext { +private: + // Can not use this unmapped pointer from outside + void* raw_window_ptr; + std::shared_ptr symmetric_memory; + +public: + // Global + int rank_idx; + int num_ranks; + + // Logical + int num_scaleout_ranks, num_scaleup_ranks; + int scaleout_rank_idx, scaleup_rank_idx; + + // Physical + int num_rdma_ranks, num_nvl_ranks; + int rdma_rank_idx, nvl_rank_idx; + bool is_scaleup_nvlink; + + // NCCL handles + ncclComm_t comm; + ncclDevComm_t dev_comm; + ncclWindow_t window; + void* mapped_window_ptr; + std::vector nvl_window_ptrs; + + // Configs + int num_allocated_qps; + + // Buffer size + int64_t num_gpu_bytes; + int64_t num_cpu_bytes; + + NCCLSymmetricMemoryContext(const int64_t& nccl_comm, const symmetric::cpu_comm_t& cpu_comm, + const int& num_ranks, const int& rank_idx, + const int64_t& num_bytes, const int64_t& num_cpu_bytes, + const bool& allow_hybrid_mode, + const int& sl_idx, const int& num_allocated_qps); + + // TODO: finish this with `explicit_destroy` + // ~NCCLSymmetricMemoryContext(); + + void* get_sym_ptr(void* ptr, const int& dst_rank_idx) const; + + void finalize(); +}; + +} // deep_ep::nccl + +namespace deep_ep::cuda_driver { + +void batched_write(CUstream stream, const std::vector& ptrs, const int& value); + +void batched_wait(CUstream stream, const std::vector& ptrs, const int& value); + +void batched_write_and_wait(CUstream stream, const std::vector& write_ptrs, const std::vector& wait_ptrs, const int& value); + +} // namespace deep_ep::cuda_driver diff --git a/csrc/kernels/backend/cuda_driver.cu b/csrc/kernels/backend/cuda_driver.cu new file mode 100644 index 0000000..9b6fc99 --- /dev/null +++ b/csrc/kernels/backend/cuda_driver.cu @@ -0,0 +1,54 @@ +#include +#include +#include + +#include + +#include "api.cuh" +#include "../../utils/lazy_driver.hpp" + +namespace deep_ep::cuda_driver { + +static CUstreamBatchMemOpParams create_mem_op( + void *ptr, const int& value, + const CUstreamBatchMemOpType& type, + const CUstreamWaitValue_flags& wait_flag = CU_STREAM_WAIT_VALUE_EQ) { + CUstreamBatchMemOpParams params; + if (type == CU_STREAM_MEM_OP_WRITE_VALUE_32) { + params.operation = CU_STREAM_MEM_OP_WRITE_VALUE_32; + params.writeValue.address = reinterpret_cast(ptr); + params.writeValue.value = value; + params.writeValue.flags = 0; + } else { + params.operation = CU_STREAM_MEM_OP_WAIT_VALUE_32; + params.waitValue.address = reinterpret_cast(ptr); + params.waitValue.value = value; + params.waitValue.flags = wait_flag; + } + return params; +} + +void batched_write(CUstream stream, const std::vector& ptrs, const int& value) { + std::vector ops(ptrs.size()); + for (int i = 0; i < ptrs.size(); ++ i) + ops[i] = create_mem_op(ptrs[i], value, CU_STREAM_MEM_OP_WRITE_VALUE_32); + CUDA_DRIVER_CHECK(lazy_cuStreamBatchMemOp(stream, ops.size(), ops.data(), 0)); +} + +void batched_wait(CUstream stream, const std::vector& ptrs, const int& value) { + std::vector ops(ptrs.size()); + for (int i = 0; i < ptrs.size(); ++ i) + ops[i] = create_mem_op(ptrs[i], value, CU_STREAM_MEM_OP_WAIT_VALUE_32, CU_STREAM_WAIT_VALUE_GEQ); + CUDA_DRIVER_CHECK(lazy_cuStreamBatchMemOp(stream, ops.size(), ops.data(), 0)); +} + +void batched_write_and_wait(CUstream stream, const std::vector& write_ptrs, const std::vector& wait_ptrs, const int& value) { + std::vector ops(write_ptrs.size() + wait_ptrs.size()); + for (int i = 0; i < write_ptrs.size(); ++ i) + ops[i] = create_mem_op(write_ptrs[i], value, CU_STREAM_MEM_OP_WRITE_VALUE_32); + for (int i = 0; i < wait_ptrs.size(); ++ i) + ops[write_ptrs.size() + i] = create_mem_op(wait_ptrs[i], value, CU_STREAM_MEM_OP_WAIT_VALUE_32, CU_STREAM_WAIT_VALUE_GEQ); + CUDA_DRIVER_CHECK(lazy_cuStreamBatchMemOp(stream, ops.size(), ops.data(), 0)); +} + +} // namespace deep_ep::cuda_driver diff --git a/csrc/kernels/backend/nccl.cu b/csrc/kernels/backend/nccl.cu new file mode 100644 index 0000000..969b136 --- /dev/null +++ b/csrc/kernels/backend/nccl.cu @@ -0,0 +1,156 @@ +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include + +#include "api.cuh" +#include "../../utils/system.hpp" + + +namespace deep_ep::nccl { + +pybind11::bytearray get_local_unique_id() { + ncclUniqueId unique_id; + NCCL_CHECK(ncclGetUniqueId(&unique_id)); + std::vector result(sizeof(ncclUniqueId)); + std::memcpy(result.data(), &unique_id, sizeof(ncclUniqueId)); + return {result.data(), result.size()}; +} + +int64_t create_nccl_comm(const pybind11::bytearray& root_unique_id_bytes, + const int& num_ranks, const int& rank_idx) { + // Copy unique ID + ncclUniqueId root_unique_id; + const auto root_unique_id_str = root_unique_id_bytes.cast(); + std::memcpy(&root_unique_id, root_unique_id_str.c_str(), sizeof(ncclUniqueId)); + + // Init + ncclComm_t comm; + NCCL_CHECK(ncclCommInitRank(&comm, num_ranks, root_unique_id, rank_idx)); + if (get_env("EP_BUFFER_DEBUG")) + printf("New NCCL host communicator created (%d/%d)\n", rank_idx, num_ranks); + return reinterpret_cast(comm); +} + +void destroy_nccl_comm(const int64_t& nccl_comm) { + NCCL_CHECK(ncclCommAbort(reinterpret_cast(nccl_comm))); + if (get_env("EP_BUFFER_DEBUG")) + printf("NCCL host communicator aborted\n"); +} + +std::tuple get_physical_domain_size(const int64_t& nccl_comm) { + const auto comm = reinterpret_cast(nccl_comm); + const int num_ranks = ncclTeamWorld(comm).nRanks, num_nvl_ranks = ncclTeamLsa(comm).nRanks; + EP_HOST_ASSERT(num_ranks % num_nvl_ranks == 0); + return {num_ranks / num_nvl_ranks, num_nvl_ranks}; +} + +std::tuple get_logical_domain_size(const int64_t& nccl_comm, const bool& allow_hybrid_mode) { + const auto [num_rdma_ranks, num_nvl_ranks] = get_physical_domain_size(nccl_comm); + return {allow_hybrid_mode ? num_rdma_ranks : 1, + allow_hybrid_mode ? num_nvl_ranks : num_rdma_ranks * num_nvl_ranks}; +} + +NCCLSymmetricMemoryContext::NCCLSymmetricMemoryContext(const int64_t& nccl_comm, const symmetric::cpu_comm_t& cpu_comm, + const int& num_ranks, const int& rank_idx, + const int64_t& num_bytes, const int64_t& num_cpu_bytes, + const bool& allow_hybrid_mode, + const int& sl_idx, const int& num_allocated_qps): + rank_idx(rank_idx), num_ranks(num_ranks), num_allocated_qps(num_allocated_qps) { + if (get_env("EP_BUFFER_DEBUG", 0)) { + int nccl_version; + NCCL_CHECK(ncclGetVersion(&nccl_version)); + printf("DeepEP initialized with NCCL version: %d.%d.%d (loaded library)\n", + nccl_version / 10000, (nccl_version % 10000) / 100, nccl_version % 100); + } + + // Reuse the NCCL communicator + comm = reinterpret_cast(nccl_comm); + + // Print number of allocated QPs + if (get_env("EP_BUFFER_DEBUG")) + printf("EP NCCL device communicator has %d allocated QPs\n", num_allocated_qps); + + // Initialize NCCL device communicator + ncclDevCommRequirements_t reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER; + if (num_ranks > 1 and get_env("EP_DISABLE_GIN", 0) == 0) { + // Query NCCL supported Gin Type + ncclCommProperties props = NCCL_COMM_PROPERTIES_INITIALIZER; + NCCL_CHECK(ncclCommQueryProperties(comm, &props)); + EP_HOST_ASSERT( + (allow_hybrid_mode ? props.railedGinType : props.ginType) != NCCL_GIN_TYPE_NONE and + "NCCL GIN is unavailable. This is usually due to a network configuration issue, " + "such as `allow_hybrid_mode=0` (disable direct RDMA kernels) in multi-plane network."); + + reqs.ginContextCount = num_allocated_qps; + reqs.ginExclusiveContexts = true; + reqs.ginQueueDepth = kGinQPDepth; + reqs.ginTrafficClass = sl_idx; + // Customized RDMA barrier needs extra signals + reqs.ginSignalCount = num_ranks + 2 * 2; + reqs.ginConnectionType = allow_hybrid_mode ? NCCL_GIN_CONNECTION_RAIL: NCCL_GIN_CONNECTION_FULL; + } + NCCL_CHECK(ncclDevCommCreate(comm, &reqs, &dev_comm)); + + // Now we know the NVLink domain size + num_nvl_ranks = dev_comm.lsaSize, nvl_rank_idx = dev_comm.lsaRank; + num_rdma_ranks = num_ranks / num_nvl_ranks, rdma_rank_idx = rank_idx / num_nvl_ranks; + EP_HOST_ASSERT(num_ranks % num_nvl_ranks == 0 and nvl_rank_idx == rank_idx % num_nvl_ranks); + EP_HOST_ASSERT(rank_idx == rdma_rank_idx * num_nvl_ranks + nvl_rank_idx); + + // Calculate scaleout/up domain size + if (allow_hybrid_mode) { + num_scaleout_ranks = num_rdma_ranks, num_scaleup_ranks = num_nvl_ranks; + scaleout_rank_idx = rdma_rank_idx, scaleup_rank_idx = nvl_rank_idx; + } else { + num_scaleout_ranks = 1, num_scaleup_ranks = num_ranks; + scaleout_rank_idx = 0, scaleup_rank_idx = rank_idx; + } + is_scaleup_nvlink = num_scaleup_ranks == num_nvl_ranks; + + // Create symmetric memory + // num_bytes = GPU + CPU, derive GPU portion + this->symmetric_memory = symmetric::alloc( + num_bytes - num_cpu_bytes, num_cpu_bytes, + allow_hybrid_mode, num_scaleup_ranks, scaleout_rank_idx, + cpu_comm); + + // Create window + // NOTES: `ncclCommWindowRegister` is collective: it internally calls bootstrapBarrier + // across all ranks, so no explicit barrier is needed after this call. + raw_window_ptr = this->symmetric_memory->ptr; + this->num_gpu_bytes = this->symmetric_memory->num_gpu_bytes; + this->num_cpu_bytes = this->symmetric_memory->num_cpu_bytes; + NCCL_CHECK(ncclCommWindowRegister(comm, raw_window_ptr, this->symmetric_memory->num_bytes, &window, NCCL_WIN_STRICT_ORDERING)); + NCCL_CHECK(ncclGetLsaDevicePointer(window, 0, nvl_rank_idx, &mapped_window_ptr)); + + // Get LSA pointers for all LSA peers + // TODO: check whether this is correct for network with RDMA + nvl_window_ptrs.resize(num_nvl_ranks); + for (int i = 0; i < num_nvl_ranks; ++ i) + NCCL_CHECK(ncclGetLsaDevicePointer(window, 0, i, &nvl_window_ptrs[i])); +} + +void* NCCLSymmetricMemoryContext::get_sym_ptr(void* ptr, const int& dst_rank_idx) const { + const auto offset = static_cast(ptr) - static_cast(mapped_window_ptr); + return static_cast(nvl_window_ptrs[dst_rank_idx]) + offset; +} + +void NCCLSymmetricMemoryContext::finalize() { + // Deregister window + NCCL_CHECK(ncclCommWindowDeregister(comm, window)); + symmetric_memory.reset(); + + // Destroy device communicator + NCCL_CHECK(ncclDevCommDestroy(comm, &dev_comm)); +} + +} // namespace deep_ep::nccl diff --git a/csrc/kernels/backend/nvshmem.cu b/csrc/kernels/backend/nvshmem.cu new file mode 100644 index 0000000..837d8fe --- /dev/null +++ b/csrc/kernels/backend/nvshmem.cu @@ -0,0 +1,87 @@ +#include +#include +#include + +#include +#include + +#include + +namespace deep_ep::nvshmem { + +nvshmem_team_t cpu_rdma_team = NVSHMEM_TEAM_INVALID; +nvshmem_team_config_t cpu_rdma_team_config; + +std::vector get_unique_id() { + nvshmemx_uniqueid_t unique_id; + nvshmemx_get_uniqueid(&unique_id); + std::vector result(sizeof(nvshmemx_uniqueid_t)); + std::memcpy(result.data(), &unique_id, sizeof(nvshmemx_uniqueid_t)); + return result; +} + +void* alloc(const size_t& size, const size_t& alignment) { + return nvshmem_align(alignment, size); +} + +void free(void* ptr) { + nvshmem_free(ptr); +} + +void barrier(const bool& with_cpu_sync, + const std::optional& stream_opt = std::nullopt) { + // Wait all streams to finish on this GPU + if (with_cpu_sync) + CUDA_RUNTIME_CHECK(cudaDeviceSynchronize()); + + // NOTES: this only launches kernels at GPU + if (stream_opt.has_value()) { + nvshmemx_barrier_all_on_stream(stream_opt.value()); + } else { + nvshmem_barrier_all(); + } + + // Let CPU wait + if (with_cpu_sync) + CUDA_RUNTIME_CHECK(cudaDeviceSynchronize()); +} + +int init(const std::vector& root_unique_id_val, + const int& rank, + const int& num_ranks, + const int& team_split_stride) { + nvshmemx_uniqueid_t root_unique_id; + nvshmemx_init_attr_t attr; + std::memcpy(&root_unique_id, root_unique_id_val.data(), sizeof(nvshmemx_uniqueid_t)); + nvshmemx_set_attr_uniqueid_args(rank, num_ranks, &root_unique_id, &attr); + nvshmemx_init_attr(NVSHMEMX_INIT_WITH_UNIQUEID, &attr); + + // Create sub-RDMA teams + if (team_split_stride > 0 and num_ranks > team_split_stride) { + EP_HOST_ASSERT(cpu_rdma_team == NVSHMEM_TEAM_INVALID); + EP_HOST_ASSERT(num_ranks % team_split_stride == 0); + EP_HOST_ASSERT(nvshmem_team_split_strided(NVSHMEM_TEAM_WORLD, + rank % team_split_stride, + team_split_stride, + num_ranks / team_split_stride, + &cpu_rdma_team_config, + 0, + &cpu_rdma_team) == 0); + EP_HOST_ASSERT(cpu_rdma_team != NVSHMEM_TEAM_INVALID); + } + + // Wait all GPUs to get ready + barrier(true); + return nvshmem_my_pe(); +} + +void finalize() { + barrier(true); + if (cpu_rdma_team != NVSHMEM_TEAM_INVALID) { + nvshmem_team_destroy(cpu_rdma_team); + cpu_rdma_team = NVSHMEM_TEAM_INVALID; + } + nvshmem_finalize(); +} + +} // namespace deep_ep::nvshmem diff --git a/csrc/kernels/backend/symmetric.hpp b/csrc/kernels/backend/symmetric.hpp new file mode 100644 index 0000000..c82f851 --- /dev/null +++ b/csrc/kernels/backend/symmetric.hpp @@ -0,0 +1,319 @@ +#pragma once + +#include +#include + +#include +#include +#include +#include +#include + +#include "../../utils/lazy_driver.hpp" + +namespace deep_ep::symmetric { + +static constexpr int64_t kNumAlignmentBytes = 2097152; + +// CPU communicator types: (pid, fd) handle and list of handles from all ranks +using cpu_handle_t = std::pair; +using cpu_comm_t = std::vector; + +struct DeviceContext { + int device_idx; + CUdevice device; + int numa_idx; + + DeviceContext() { + CUDA_RUNTIME_CHECK(cudaGetDevice(&device_idx)); + CUDA_DRIVER_CHECK(lazy_cuDeviceGet(&device, device_idx)); + CUDA_DRIVER_CHECK(lazy_cuDeviceGetAttribute(&numa_idx, CU_DEVICE_ATTRIBUTE_HOST_NUMA_ID, device)); + } + + CUmemAllocationProp gpu_alloc_prop() const { return build_alloc_prop(CU_MEM_LOCATION_TYPE_DEVICE, device_idx); } + CUmemAllocationProp cpu_alloc_prop() const { return build_alloc_prop(CU_MEM_LOCATION_TYPE_HOST_NUMA, numa_idx); } + +private: + CUmemAllocationProp build_alloc_prop(const CUmemLocationType& location_type, const int& location_idx) const { + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = location_type; + prop.location.id = location_idx; + + int requested_handle_types = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; + int flag = 0; +#if CUDART_VERSION >= 13000 + CUDA_DRIVER_CHECK(lazy_cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device)); + if (flag) + requested_handle_types |= CU_MEM_HANDLE_TYPE_FABRIC; +#endif + prop.requestedHandleTypes = static_cast(requested_handle_types); + + if (location_type == CU_MEM_LOCATION_TYPE_DEVICE) { + flag = 0; +#if CUDART_VERSION >= 13000 + CUDA_DRIVER_CHECK(lazy_cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED, device)); +#else + CUDA_DRIVER_CHECK(lazy_cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_SUPPORTED, device)); +#endif + EP_HOST_ASSERT(flag and "GPUDirect RDMA with CUDA VMM is not supported on this device"); + prop.allocFlags.gpuDirectRDMACapable = 1; + } + + // Check granularity + size_t num_granularity_bytes = 0; + CUDA_DRIVER_CHECK(lazy_cuMemGetAllocationGranularity(&num_granularity_bytes, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); + EP_HOST_ASSERT((num_granularity_bytes == 0 or kNumAlignmentBytes % num_granularity_bytes == 0) and + "Alignment must be a multiple of CUDA allocation granularity"); + return prop; + } +}; + +// Try cuMemCreate with FABRIC fallback (mirrors ncclMemAlloc logic) +static void cumem_create_with_fallback(CUmemGenericAllocationHandle* handle, + const int64_t& num_bytes, CUmemAllocationProp* prop) { + if (prop->requestedHandleTypes & CU_MEM_HANDLE_TYPE_FABRIC) { + CUresult err = lazy_cuMemCreate(handle, num_bytes, prop, 0); + if (err == CUDA_ERROR_NOT_PERMITTED or err == CUDA_ERROR_NOT_SUPPORTED) { + prop->requestedHandleTypes = static_cast( + prop->requestedHandleTypes & ~CU_MEM_HANDLE_TYPE_FABRIC); + CUDA_DRIVER_CHECK(lazy_cuMemCreate(handle, num_bytes, prop, 0)); + } else { + CUDA_DRIVER_CHECK(err); + } + } else { + CUDA_DRIVER_CHECK(lazy_cuMemCreate(handle, num_bytes, prop, 0)); + } +} + +// Set read/write access for all peer GPUs, and optionally a NUMA node +static void set_access(const CUdeviceptr& addr, const int64_t& num_bytes, + const int& device_idx, const int& numa_idx = -1) { + const bool with_cpu = (numa_idx >= 0); + CUmemAccessDesc desc[2]; + desc[0].location.type = CU_MEM_LOCATION_TYPE_DEVICE; + desc[0].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + if (with_cpu) { + desc[1].location.type = CU_MEM_LOCATION_TYPE_HOST_NUMA; + desc[1].location.id = numa_idx; + desc[1].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + } + + int num_devices; + CUDA_RUNTIME_CHECK(cudaGetDeviceCount(&num_devices)); + for (int i = 0; i < num_devices; ++ i) { + int can_access_peer = 0; + if (i == device_idx or (cudaDeviceCanAccessPeer(&can_access_peer, i, device_idx) == cudaSuccess and can_access_peer)) { + desc[0].location.id = i; + CUDA_DRIVER_CHECK(lazy_cuMemSetAccess(addr, num_bytes, desc, with_cpu ? 2 : 1)); + } + } +} + +class SymmetricMemory { +public: + void* ptr = nullptr; + int64_t num_bytes = 0; + int64_t num_gpu_bytes = 0; + int64_t num_cpu_bytes = 0; + + virtual ~SymmetricMemory() noexcept(0) = default; +}; + +// Wraps ncclMemAlloc/ncclMemFree (pure GPU, current default) +class GPUSymmetricMemory final : public SymmetricMemory { +public: + explicit GPUSymmetricMemory(const int64_t& num_bytes) { + EP_HOST_ASSERT(num_bytes > 0 and num_bytes % kNumAlignmentBytes == 0); + NCCL_CHECK(ncclMemAlloc(&ptr, num_bytes)); + EP_HOST_ASSERT(reinterpret_cast(ptr) % kNumAlignmentBytes == 0); + this->num_bytes = num_bytes; + this->num_gpu_bytes = num_bytes; + } + + ~GPUSymmetricMemory() override { + if (ptr != nullptr) { + NCCL_CHECK(ncclMemFree(ptr)); + ptr = nullptr; + } + } +}; + +// GPU + CPU mixed allocation via CUDA Driver API +// Memory layout: [GPU VRAM (front)] [CPU RAM / NUMA-local (back)] +// The entire contiguous VA range is compatible with `ncclCommWindowRegister`. +class ElasticSymmetricMemory : public SymmetricMemory { + CUmemGenericAllocationHandle gpu_handle = {}; + CUmemGenericAllocationHandle cpu_handle = {}; + +public: + ElasticSymmetricMemory(const int64_t& num_gpu_bytes, const int64_t& num_cpu_bytes) { + EP_HOST_ASSERT(num_gpu_bytes > 0 and num_gpu_bytes % kNumAlignmentBytes == 0); + EP_HOST_ASSERT(num_cpu_bytes > 0 and num_cpu_bytes % kNumAlignmentBytes == 0); + + DeviceContext ctx; + auto gpu_prop = ctx.gpu_alloc_prop(); + auto cpu_prop = ctx.cpu_alloc_prop(); + + this->num_gpu_bytes = num_gpu_bytes; + this->num_cpu_bytes = num_cpu_bytes; + this->num_bytes = num_gpu_bytes + num_cpu_bytes; + + // Reserve VA and map segments + CUdeviceptr addr; + CUDA_DRIVER_CHECK(lazy_cuMemAddressReserve(&addr, this->num_bytes, kNumAlignmentBytes, 0, 0)); + this->ptr = reinterpret_cast(static_cast(addr)); + + cumem_create_with_fallback(&gpu_handle, num_gpu_bytes, &gpu_prop); + CUDA_DRIVER_CHECK(lazy_cuMemMap(addr, num_gpu_bytes, 0, gpu_handle, 0)); + set_access(addr, num_gpu_bytes, ctx.device_idx); + + cumem_create_with_fallback(&cpu_handle, num_cpu_bytes, &cpu_prop); + CUDA_DRIVER_CHECK(lazy_cuMemMap(addr + num_gpu_bytes, num_cpu_bytes, 0, cpu_handle, 0)); + set_access(addr + num_gpu_bytes, num_cpu_bytes, ctx.device_idx, ctx.numa_idx); + + EP_HOST_ASSERT(reinterpret_cast(ptr) % kNumAlignmentBytes == 0); + } + + ~ElasticSymmetricMemory() override { + auto addr = static_cast(reinterpret_cast(ptr)); + CUDA_DRIVER_CHECK(lazy_cuMemUnmap(addr, num_gpu_bytes)); + CUDA_DRIVER_CHECK(lazy_cuMemRelease(gpu_handle)); + CUDA_DRIVER_CHECK(lazy_cuMemUnmap(addr + num_gpu_bytes, num_cpu_bytes)); + CUDA_DRIVER_CHECK(lazy_cuMemRelease(cpu_handle)); + CUDA_DRIVER_CHECK(lazy_cuMemAddressFree(addr, num_bytes)); + } +}; + +// Each rank creates its own NUMA-local CPU segment; +// CPU segments are imported and mapped contiguously into every rank's VA space. +// Layout: [GPU VRAM (front)] [CPU rank0 | CPU rank1 | ... | CPU rank(N-1) (back)] +class HybridElasticSymmetricMemory final : public SymmetricMemory { + int num_scaleup_ranks; + CUmemGenericAllocationHandle gpu_handle = {}; + std::vector cpu_handles; + int local_export_fd = -1; + +public: + HybridElasticSymmetricMemory(const cpu_comm_t& cpu_comm, + const int64_t& num_gpu_bytes, const int64_t& num_cpu_bytes, + const int& num_scaleup_ranks, const int& scaleout_rank_idx): + num_scaleup_ranks(num_scaleup_ranks), + cpu_handles(num_scaleup_ranks) { + EP_HOST_ASSERT(num_gpu_bytes > 0 and num_gpu_bytes % kNumAlignmentBytes == 0); + EP_HOST_ASSERT(num_cpu_bytes > 0 and num_cpu_bytes % kNumAlignmentBytes == 0); + + DeviceContext ctx; + auto gpu_prop = ctx.gpu_alloc_prop(); + + this->num_gpu_bytes = num_gpu_bytes; + this->num_cpu_bytes = num_cpu_bytes; + this->num_bytes = num_gpu_bytes + num_cpu_bytes * num_scaleup_ranks; + + // Reserve VA + CUdeviceptr addr; + CUDA_DRIVER_CHECK(lazy_cuMemAddressReserve(&addr, this->num_bytes, kNumAlignmentBytes, 0, 0)); + this->ptr = reinterpret_cast(static_cast(addr)); + + // Map GPU segment + cumem_create_with_fallback(&gpu_handle, num_gpu_bytes, &gpu_prop); + CUDA_DRIVER_CHECK(lazy_cuMemMap(addr, num_gpu_bytes, 0, gpu_handle, 0)); + set_access(addr, num_gpu_bytes, ctx.device_idx); + + // Import and map all intra-node CPU segments + const auto local_pid = getpid(); + for (int i = 0; i < num_scaleup_ranks; ++ i) { + auto [pid, fd] = cpu_comm[num_scaleup_ranks * scaleout_rank_idx + i]; + int local_fd = fd; + if (pid != local_pid) { + int pidfd = syscall(SYS_pidfd_open, pid, 0); + EP_HOST_ASSERT(pidfd >= 0 and "`pidfd_open` failed"); + local_fd = syscall(SYS_pidfd_getfd, pidfd, fd, 0); + EP_HOST_ASSERT(local_fd >= 0 and "`pidfd_getfd` failed"); + close(pidfd); + } + + const auto offset = num_gpu_bytes + i * num_cpu_bytes; + CUDA_DRIVER_CHECK(lazy_cuMemImportFromShareableHandle( + &cpu_handles[i], + reinterpret_cast(static_cast(local_fd)), + CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR)); + CUDA_DRIVER_CHECK(lazy_cuMemMap(addr + offset, num_cpu_bytes, 0, cpu_handles[i], 0)); + set_access(addr + offset, num_cpu_bytes, ctx.device_idx, ctx.numa_idx); + + if (pid != local_pid) { + close(local_fd); + } else { + local_export_fd = local_fd; + } + } + + EP_HOST_ASSERT(reinterpret_cast(ptr) % kNumAlignmentBytes == 0); + } + + ~HybridElasticSymmetricMemory() override { + auto addr = static_cast(reinterpret_cast(ptr)); + + CUDA_DRIVER_CHECK(lazy_cuMemUnmap(addr, num_gpu_bytes)); + CUDA_DRIVER_CHECK(lazy_cuMemRelease(gpu_handle)); + + if (local_export_fd >= 0) + close(local_export_fd); + CUDA_DRIVER_CHECK(lazy_cuMemUnmap(addr + num_gpu_bytes, num_cpu_bytes * num_scaleup_ranks)); + for (int i = 0; i < num_scaleup_ranks; ++ i) + CUDA_DRIVER_CHECK(lazy_cuMemRelease(cpu_handles[i])); + + CUDA_DRIVER_CHECK(lazy_cuMemAddressFree(addr, num_bytes)); + } + + // Create a NUMA-local CPU segment and export its POSIX FD handle. + // Returns (pid, fd) handle for cross-process sharing. + static cpu_handle_t create_cpu_handle(const int64_t& num_cpu_bytes) { + EP_HOST_ASSERT(num_cpu_bytes > 0 and num_cpu_bytes % kNumAlignmentBytes == 0); + + DeviceContext ctx; + auto cpu_prop = ctx.cpu_alloc_prop(); + + CUmemGenericAllocationHandle handle; + cumem_create_with_fallback(&handle, num_cpu_bytes, &cpu_prop); + + int fd = -1; + CUDA_DRIVER_CHECK(lazy_cuMemExportToShareableHandle( + &fd, handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0)); + + // Release the allocation handle -- the POSIX FD keeps the physical memory alive + CUDA_DRIVER_CHECK(lazy_cuMemRelease(handle)); + + return {getpid(), fd}; + } +}; + +static std::shared_ptr alloc(const int64_t& num_gpu_bytes, const int64_t& num_cpu_bytes, + const bool& allow_hybrid_mode = false, + const int& num_scaleup_ranks = 0, const int& scaleout_rank_idx = 0, + const cpu_comm_t& cpu_comm = {}) { + EP_HOST_ASSERT(num_gpu_bytes > 0 and num_gpu_bytes % kNumAlignmentBytes == 0); + EP_HOST_ASSERT(num_cpu_bytes >= 0 and num_cpu_bytes % kNumAlignmentBytes == 0); + + std::shared_ptr result; + if (num_cpu_bytes > 0) { + if (allow_hybrid_mode) { + result = std::make_shared( + cpu_comm, num_gpu_bytes, num_cpu_bytes, + num_scaleup_ranks, scaleout_rank_idx); + } else { + result = std::make_shared( + num_gpu_bytes, num_cpu_bytes); + } + } else { + result = std::make_shared(num_gpu_bytes); + } + + // TODO: move all variables into NCCL runtime + // Enable NCCL elastic buffer register if the allocator may produce CPU-backed segments + if (dynamic_cast(result.get()) != nullptr) + setenv("NCCL_ELASTIC_BUFFER_REGISTER", "1", 0); + return result; +} + +} // namespace deep_ep::symmetric diff --git a/csrc/kernels/elastic/api.hpp b/csrc/kernels/elastic/api.hpp new file mode 100644 index 0000000..29023db --- /dev/null +++ b/csrc/kernels/elastic/api.hpp @@ -0,0 +1,9 @@ +#pragma once + +#include + +#include "barrier.hpp" +#include "dispatch.hpp" +#include "combine.hpp" +#include "engram.hpp" +#include "pp_send_recv.hpp" diff --git a/csrc/kernels/elastic/barrier.hpp b/csrc/kernels/elastic/barrier.hpp new file mode 100644 index 0000000..d5840e6 --- /dev/null +++ b/csrc/kernels/elastic/barrier.hpp @@ -0,0 +1,86 @@ +#pragma once + +#include +#include + +#include +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/launch_runtime.hpp" + +namespace deep_ep::elastic { + +class BarrierRuntime final : public jit::LaunchRuntime { +public: + struct Args { + // Templated arguments + bool is_scaleup_nvlink; + int num_scaleout_ranks, num_scaleup_ranks; + int64_t num_timeout_cycles; + bool sequential; + + // Parameters + ncclDevComm_t nccl_dev_comm; + ncclWindow_t nccl_window; + void* workspace; + int scaleout_rank_idx, scaleup_rank_idx; + + jit::LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_ep::elastic; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&barrier_impl<{}, {}, {}, {}, {}, {}, {}>); +}} +)", args.is_scaleup_nvlink, + args.launch_args.grid_dim.first, args.launch_args.num_threads, + args.num_scaleout_ranks, args.num_scaleup_ranks, + args.num_timeout_cycles, args.sequential); + } + + static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) { + EP_CUDA_UNIFIED_CHECK(jit::launch_kernel( + kernel, config, + args.nccl_dev_comm, args.nccl_window, + args.workspace, args.scaleout_rank_idx, args.scaleup_rank_idx + )); + } +}; + +static void launch_barrier(const ncclDevComm_t& nccl_dev_comm, const ncclWindow_t& nccl_window, + void* workspace, + const int& scaleout_rank_idx, const int& scaleup_rank_idx, + const int& num_scaleout_ranks, const int& num_scaleup_ranks, + const int64_t& num_timeout_cycles, + const bool& is_scaleup_nvlink, + const bool& sequential, + const at::cuda::CUDAStream& stream) { + // Number of threads equals to the number of ranks + constexpr auto kNumThreads = 512; + + // Generate, build and launch + // NOTES: only the parallel hybrid kernel needs 2 SMs; the sequential mode does scaleout and + // scaleup one after another, so a single SM is sufficient. + const auto num_sms = (not sequential and num_scaleout_ranks > 1) ? 2 : 1; + const BarrierRuntime::Args args = { + .is_scaleup_nvlink = is_scaleup_nvlink, + .num_scaleout_ranks = num_scaleout_ranks, .num_scaleup_ranks = num_scaleup_ranks, + .num_timeout_cycles = num_timeout_cycles, + .sequential = sequential, + .nccl_dev_comm = nccl_dev_comm, + .nccl_window = nccl_window, + .workspace = workspace, + .scaleout_rank_idx = scaleout_rank_idx, .scaleup_rank_idx = scaleup_rank_idx, + .launch_args = jit::LaunchArgs(num_sms, kNumThreads, 0, 1, true)}; + const auto code = BarrierRuntime::generate(args); + const auto runtime = jit::compiler->build("barrier", code); + BarrierRuntime::launch(runtime, args, stream); +} + +} // namespace deep_ep::elastic diff --git a/csrc/kernels/elastic/combine.hpp b/csrc/kernels/elastic/combine.hpp new file mode 100644 index 0000000..3764da9 --- /dev/null +++ b/csrc/kernels/elastic/combine.hpp @@ -0,0 +1,289 @@ +#pragma once + +#include + +#include +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/launch_runtime.hpp" + +namespace deep_ep::elastic { + +class CombineRuntime final : public jit::LaunchRuntime { +public: + struct Args { + // Templated arguments + bool is_scaleup_nvlink; + bool use_expanded_layout, allow_multiple_reduction; + int num_scaleup_warps, num_forward_warps; + int num_scaleout_ranks, num_scaleup_ranks; + int hidden; + int num_max_tokens_per_rank; + int num_experts; + int num_topk; + int num_qps; + int64_t num_timeout_cycles; + + // Parameters + nv_bfloat16* x; + float* topk_weights; + int* src_metadata; + int* psum_num_recv_tokens_per_scaleup_rank; + int* token_metadata_at_forward; + int* channel_linked_list; + ncclDevComm_t nccl_dev_comm; + ncclWindow_t nccl_window; + void* buffer; + void* workspace; + int scaleout_rank_idx, scaleup_rank_idx; + int num_reduced_tokens; + + jit::LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + std::string header_name, func_name; + if (args.num_scaleout_ranks == 1) { + header_name = "combine"; + func_name = fmt::format("combine_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>", + args.is_scaleup_nvlink, + args.use_expanded_layout, args.allow_multiple_reduction, + args.launch_args.grid_dim.first, + args.launch_args.num_threads / 32, + args.num_scaleup_ranks * args.num_scaleout_ranks, + args.hidden, + args.num_max_tokens_per_rank, + args.num_experts, + args.num_topk, + args.num_qps, args.num_timeout_cycles); + } else { + header_name = "hybrid_combine"; + func_name = fmt::format("hybrid_combine_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>", + args.use_expanded_layout, args.allow_multiple_reduction, + args.launch_args.grid_dim.first, + args.num_scaleup_warps, args.num_forward_warps, + args.num_scaleout_ranks, args.num_scaleup_ranks, + args.hidden, + args.num_max_tokens_per_rank, + args.num_experts, + args.num_topk, + args.num_qps, + args.num_timeout_cycles); + } + return fmt::format(R"( +#include + +using namespace deep_ep::elastic; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&{}); +}} +)", header_name, func_name); + } + + static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) { + if (args.num_scaleout_ranks == 1) { + EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(kernel, config, + args.x, args.topk_weights, + args.src_metadata, args.psum_num_recv_tokens_per_scaleup_rank, + args.nccl_dev_comm, args.nccl_window, + args.buffer, args.workspace, + args.scaleup_rank_idx, + args.num_reduced_tokens)); + } else { + EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(kernel, config, + args.x, args.topk_weights, + args.src_metadata, + args.psum_num_recv_tokens_per_scaleup_rank, + args.token_metadata_at_forward, + args.channel_linked_list, + args.nccl_dev_comm, args.nccl_window, + args.buffer, args.workspace, + args.scaleout_rank_idx, args.scaleup_rank_idx, + args.num_reduced_tokens)); + } + } +}; + +static layout::TokenLayout get_combine_token_layout( + const int& hidden, const int& elem_size, const int& num_topk) { + return layout::TokenLayout(hidden * elem_size, 0, num_topk, false); +} + +static void* launch_combine(void* x, + void* topk_weights, + int* src_metadata, + int* psum_num_recv_tokens_per_scaleup_rank, + int* token_metadata_at_forward, + int* channel_linked_list, + const ncclDevComm_t& nccl_dev_comm, const ncclWindow_t& nccl_window, + void* buffer, void* workspace, + const int& num_reduced_tokens, const int& num_max_tokens_per_rank, + const int& hidden, + const int& num_experts, const int& num_topk, + const int& num_qps, const int64_t& num_timeout_cycles, + const int& num_scaleout_ranks, const int& num_scaleup_ranks, + const int& scaleout_rank_idx, const int& scaleup_rank_idx, + const bool& is_scaleup_nvlink, + const int& num_sms, const int& num_smem_bytes, + const int& num_channels, + const bool& use_expanded_layout, const bool& allow_multiple_reduction, + const at::cuda::CUDAStream& stream) { + // Maximize shared memory utilization + const auto token_layout = get_combine_token_layout(hidden, sizeof(nv_bfloat16), num_topk); + auto num_warps = std::min(num_smem_bytes / token_layout.get_num_bytes(), 32); + + // Decide warps + int num_scaleup_warps = 0, num_forward_warps = 0; + if (num_scaleout_ranks > 1) { + EP_HOST_ASSERT(num_channels % num_sms == 0 and + "Invalid number of channels or SMs, you may use a different SM count than dispatch"); + EP_HOST_ASSERT(num_channels / num_sms <= 16); + + num_scaleup_warps = num_forward_warps = num_channels / num_sms; + num_warps = num_scaleup_warps + num_forward_warps; + EP_HOST_ASSERT(num_warps * token_layout.get_num_bytes() <= num_smem_bytes and + "Invalid combine SM count, please try to match your dispatch config"); + } + + // Generate, build and launch + const auto num_threads = num_warps * 32; + const CombineRuntime::Args args = { + .is_scaleup_nvlink = is_scaleup_nvlink, + .use_expanded_layout = use_expanded_layout, + .allow_multiple_reduction = allow_multiple_reduction, + .num_scaleup_warps = num_scaleup_warps, .num_forward_warps = num_forward_warps, + .num_scaleout_ranks = num_scaleout_ranks, .num_scaleup_ranks = num_scaleup_ranks, + .hidden = hidden, + .num_max_tokens_per_rank = num_max_tokens_per_rank, + .num_experts = num_experts, + .num_topk = num_topk, + .num_qps = num_qps, .num_timeout_cycles = num_timeout_cycles, + .x = static_cast(x), + .topk_weights = static_cast(topk_weights), + .src_metadata = src_metadata, + .psum_num_recv_tokens_per_scaleup_rank = psum_num_recv_tokens_per_scaleup_rank, + .token_metadata_at_forward = token_metadata_at_forward, + .channel_linked_list = channel_linked_list, + .nccl_dev_comm = nccl_dev_comm, .nccl_window = nccl_window, + .buffer = buffer, .workspace = workspace, + .scaleout_rank_idx = scaleout_rank_idx, .scaleup_rank_idx = scaleup_rank_idx, + .num_reduced_tokens = num_reduced_tokens, + // NOTES: make cluster dim 2 to overlap with clustered computation kernels + .launch_args = jit::LaunchArgs(num_sms, num_threads, num_smem_bytes, 2 - (num_sms % 2), true) + }; + const auto code = CombineRuntime::generate(args); + const auto runtime = jit::compiler->build("combine", code); + CombineRuntime::launch(runtime, args, stream); + + // Return the buffer to be reduced + if (num_scaleout_ranks == 1) + return buffer; + + // For hybrid mode, we have to skip the scale-up buffer + const bool is_scaleup_buffer_rank_layout = + allow_multiple_reduction ? (num_scaleup_ranks <= num_topk) : false; + const auto scaleup_buffer = layout::BufferLayout( + token_layout, + is_scaleup_buffer_rank_layout ? num_scaleup_ranks : num_topk, + num_scaleout_ranks * num_max_tokens_per_rank, + buffer); + return scaleup_buffer.get_buffer_end_ptr(); +} + +class CombineReduceEpilogueRuntime final : public jit::LaunchRuntime { +public: + struct Args { + // Templated arguments + bool use_expanded_layout, allow_multiple_reduction; + int num_scaleout_ranks, num_scaleup_ranks; + int hidden; + int num_max_tokens_per_rank; + int num_experts, num_topk; + + // Parameters + nv_bfloat16* combined_x; + float* combined_topk_weights; + topk_idx_t* combined_topk_idx; + void* reduce_buffer; + void* bias_0; + void* bias_1; + int num_combined_tokens; + int scaleout_rank_idx, scaleup_rank_idx; + + jit::LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_ep::elastic; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&combine_reduce_epilogue_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}>); +}} +)", args.use_expanded_layout, args.allow_multiple_reduction, + args.launch_args.grid_dim.first, + args.launch_args.num_threads / 32, + args.num_scaleout_ranks, args.num_scaleup_ranks, + args.hidden, + args.num_max_tokens_per_rank, + args.num_experts, args.num_topk); + } + + static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) { + EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(kernel, config, + args.combined_x, + args.combined_topk_weights, + args.combined_topk_idx, + args.reduce_buffer, + args.bias_0, args.bias_1, + args.num_combined_tokens, + args.scaleout_rank_idx, args.scaleup_rank_idx)); + } +}; + +static void launch_combine_reduce_epilogue(void* combined_x, + float* combined_topk_weights, + topk_idx_t* combined_topk_idx, + const int& num_combined_tokens, const int& num_max_tokens_per_rank, + const int& hidden, + const int& num_experts, const int& num_topk, + void* reduce_buffer, + void* bias_0, void* bias_1, + const int& num_scaleout_ranks, const int& num_scaleup_ranks, + const int& scaleout_rank_idx, const int& scaleup_rank_idx, + const int& num_sms, const int& num_smem_bytes, + const bool& use_expanded_layout, const bool& allow_multiple_reduction, + const at::cuda::CUDAStream& stream) { + // Maximize shared memory utilization + // Too many warps may cause performance degrade, so we limit into 1024 + const auto token_layout = layout::TokenLayout(hidden * sizeof(nv_bfloat16), 0, 0, false); + const auto num_warps = std::min(num_smem_bytes / token_layout.get_num_bytes(), 32); + const auto num_threads = num_warps * 32; + + // Generate, build and launch + const CombineReduceEpilogueRuntime::Args args = { + .use_expanded_layout = use_expanded_layout, + .allow_multiple_reduction = allow_multiple_reduction, + .num_scaleout_ranks = num_scaleout_ranks, .num_scaleup_ranks = num_scaleup_ranks, + .hidden = hidden, + .num_max_tokens_per_rank = num_max_tokens_per_rank, + .num_experts = num_experts, .num_topk = num_topk, + .combined_x = static_cast(combined_x), + .combined_topk_weights = combined_topk_weights, + .combined_topk_idx = combined_topk_idx, + .reduce_buffer = reduce_buffer, + .bias_0 = bias_0, .bias_1 = bias_1, + .num_combined_tokens = num_combined_tokens, + .scaleout_rank_idx = scaleout_rank_idx, .scaleup_rank_idx = scaleup_rank_idx, + .launch_args = jit::LaunchArgs(num_sms, num_threads, num_smem_bytes, 1, false, true) + }; + const auto code = CombineReduceEpilogueRuntime::generate(args); + const auto runtime = jit::compiler->build("combine_reduce_epilogue", code); + CombineReduceEpilogueRuntime::launch(runtime, args, stream); +} + +} // namespace deep_ep::elastic diff --git a/csrc/kernels/elastic/dispatch.hpp b/csrc/kernels/elastic/dispatch.hpp new file mode 100644 index 0000000..37d2847 --- /dev/null +++ b/csrc/kernels/elastic/dispatch.hpp @@ -0,0 +1,342 @@ +#pragma once + +#include +#include + +#include +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/launch_runtime.hpp" + +namespace deep_ep::elastic { + +class DispatchRuntime final : public jit::LaunchRuntime { +public: + struct Args { + // Templated arguments + bool is_scaleup_nvlink; + bool do_cpu_sync; + bool reuse_slot_indices; + int num_notify_warps; + int num_dispatch_warps; // For hybrid dispatch + int num_scaleout_warps, num_forward_warps; // For direct dispatch + int num_scaleout_ranks, num_scaleup_ranks; + int num_hidden_bytes, num_sf_packs; + int num_max_tokens_per_rank; + int num_experts, num_topk, expert_alignment; + int num_qps; + int64_t num_timeout_cycles; + + // Parameters + void* x; sf_pack_t* sf; topk_idx_t* topk_idx; float* topk_weights; + topk_idx_t* copied_topk_idx; + int* cumulative_local_expert_recv_stats; + int* psum_num_recv_tokens_per_scaleup_rank; + int* psum_num_recv_tokens_per_expert; + int* num_unaligned_recv_tokens_per_expert; + int* dst_buffer_slot_idx; + int* token_metadata_at_forward; + int num_tokens; + int sf_token_stride, sf_hidden_stride; + ncclDevComm_t nccl_dev_comm; + ncclWindow_t nccl_window; + void* buffer; + void* workspace; void* mapped_host_workspace; + int scaleout_rank_idx, scaleup_rank_idx; + + jit::LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + std::string header_name, func_name; + if (args.num_scaleout_ranks == 1) { + header_name = "dispatch"; + func_name = fmt::format("dispatch_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>", + args.is_scaleup_nvlink, + args.do_cpu_sync, + args.reuse_slot_indices, + args.launch_args.grid_dim.first, + args.num_notify_warps, args.num_dispatch_warps, + args.num_scaleup_ranks, + args.num_hidden_bytes, args.num_sf_packs, + args.num_max_tokens_per_rank, + args.num_experts, args.num_topk, args.expert_alignment, + args.num_qps, args.num_timeout_cycles); + } else { + header_name = "hybrid_dispatch"; + func_name = fmt::format("hybrid_dispatch_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>", + args.do_cpu_sync, + args.reuse_slot_indices, + args.launch_args.grid_dim.first, + args.num_notify_warps, args.num_scaleout_warps, args.num_forward_warps, + args.num_scaleout_ranks, args.num_scaleup_ranks, + args.num_hidden_bytes, args.num_sf_packs, + args.num_max_tokens_per_rank, + args.num_experts, args.num_topk, args.expert_alignment, + args.num_qps, args.num_timeout_cycles); + } + + return fmt::format(R"( +#include + +using namespace deep_ep::elastic; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&{}); +}} +)", header_name, func_name); + } + + static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) { + if (args.num_scaleout_ranks == 1) { + EP_CUDA_UNIFIED_CHECK(jit::launch_kernel( + kernel, config, + args.x, args.sf, args.topk_idx, args.topk_weights, + args.copied_topk_idx, + args.cumulative_local_expert_recv_stats, + args.psum_num_recv_tokens_per_scaleup_rank, + args.psum_num_recv_tokens_per_expert, + args.num_unaligned_recv_tokens_per_expert, + args.dst_buffer_slot_idx, + args.num_tokens, + args.sf_token_stride, args.sf_hidden_stride, + args.nccl_dev_comm, args.nccl_window, + args.buffer, + args.workspace, args.mapped_host_workspace, + args.scaleup_rank_idx)); + } else { + EP_CUDA_UNIFIED_CHECK(jit::launch_kernel( + kernel, config, + args.x, args.sf, args.topk_idx, args.topk_weights, + args.copied_topk_idx, + args.cumulative_local_expert_recv_stats, + args.psum_num_recv_tokens_per_scaleup_rank, + args.psum_num_recv_tokens_per_expert, + args.num_unaligned_recv_tokens_per_expert, + args.dst_buffer_slot_idx, + args.token_metadata_at_forward, + args.num_tokens, + args.sf_token_stride, args.sf_hidden_stride, + args.nccl_dev_comm, args.nccl_window, + args.buffer, + args.workspace, args.mapped_host_workspace, + args.scaleout_rank_idx, args.scaleup_rank_idx + )); + } + } +}; + +constexpr int kNumNotifyWarps = 4; + +static int get_num_notify_smem_bytes(const int& num_ranks, const int& num_experts) { + return math::align(num_ranks + num_experts, kNumNotifyWarps * 32) * sizeof(int); +} + +static layout::TokenLayout get_dispatch_token_layout( + const int& hidden, const int& elem_size, const int& num_sf_packs, const int& num_topk) { + return layout::TokenLayout(hidden * elem_size, num_sf_packs * sizeof(sf_pack_t), num_topk, true); +} + +static void launch_dispatch(void* x, void* sf, + topk_idx_t* topk_idx, float* topk_weights, + topk_idx_t* copied_topk_idx, + int* cumulative_local_expert_recv_stats, + int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, + int* num_unaligned_recv_tokens_per_expert, + int* dst_buffer_slot_idx, + int* token_metadata_at_forward, + const int& num_tokens, const int& num_max_tokens_per_rank, + const int& hidden, const int& elem_size, + const int& num_sf_packs, const int& sf_token_stride, const int& sf_hidden_stride, + const int& num_experts, const int& num_topk, const int& expert_alignment, + const ncclDevComm_t& nccl_dev_comm, const ncclWindow_t& nccl_window, + void* buffer, + void* workspace, void* mapped_host_workspace, + const int& scaleout_rank_idx, const int& scaleup_rank_idx, + const int& num_scaleout_ranks, const int& num_scaleup_ranks, + const bool& is_scaleup_nvlink, + const int& num_sms, const int& num_channels_per_sm, + const int& num_smem_bytes, + const int& num_qps, const int64_t& num_timeout_cycles, + const bool& cached_mode, + const bool& do_cpu_sync, + const at::cuda::CUDAStream& stream) { + // Cached mode does not support expert token counting + if (cached_mode) + EP_HOST_ASSERT(cumulative_local_expert_recv_stats == nullptr); + + // Utils + const auto num_ranks = num_scaleout_ranks * num_scaleup_ranks; + + // Notify warps + // TODO: why don't we use 4 notify warps? + const int num_notify_warps = cached_mode ? 0 : kNumNotifyWarps; + const bool reuse_slot_indices = cached_mode; + const int num_notify_smem_bytes = cached_mode ? 0 : get_num_notify_smem_bytes(num_ranks, num_experts); + EP_HOST_ASSERT(num_notify_warps % 4 == 0); + + // Other warps + int num_dispatch_warps = 0; + int num_scaleout_warps = 0, num_forward_warps = 0; + int num_threads = 0; + + // Maximize shared memory utilization + if (num_scaleout_ranks == 1) { + const auto token_layout = get_dispatch_token_layout(hidden, elem_size, num_sf_packs, num_topk); + num_dispatch_warps = std::min( + (num_smem_bytes - num_notify_smem_bytes) / token_layout.get_num_bytes(), 32 - num_notify_warps); + num_threads = (num_notify_warps + num_dispatch_warps) * 32; + } else { + // Hybrid kernels + num_scaleout_warps = num_channels_per_sm; + num_forward_warps = num_channels_per_sm; + num_threads = (num_notify_warps + num_scaleout_warps + num_forward_warps) * 32; + } + + // Generate, build and launch + const DispatchRuntime::Args args = { + .is_scaleup_nvlink = is_scaleup_nvlink, + .do_cpu_sync = do_cpu_sync, + .reuse_slot_indices = reuse_slot_indices, + .num_notify_warps = num_notify_warps, + .num_dispatch_warps = num_dispatch_warps, + .num_scaleout_warps = num_scaleout_warps, .num_forward_warps = num_forward_warps, + .num_scaleout_ranks = num_scaleout_ranks, .num_scaleup_ranks = num_scaleup_ranks, + .num_hidden_bytes = hidden * elem_size, .num_sf_packs = num_sf_packs, + .num_max_tokens_per_rank = num_max_tokens_per_rank, + .num_experts = num_experts, .num_topk = num_topk, .expert_alignment = expert_alignment, + .num_qps = num_qps, .num_timeout_cycles = num_timeout_cycles, + .x = x, .sf = static_cast(sf), .topk_idx = topk_idx, .topk_weights = topk_weights, + .copied_topk_idx = copied_topk_idx, + .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats, + .psum_num_recv_tokens_per_scaleup_rank = psum_num_recv_tokens_per_scaleup_rank, + .psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert, + .num_unaligned_recv_tokens_per_expert = num_unaligned_recv_tokens_per_expert, + .dst_buffer_slot_idx = dst_buffer_slot_idx, + .token_metadata_at_forward = token_metadata_at_forward, + .num_tokens = num_tokens, + .sf_token_stride = sf_token_stride, .sf_hidden_stride = sf_hidden_stride, + .nccl_dev_comm = nccl_dev_comm, .nccl_window = nccl_window, + .buffer = buffer, + .workspace = workspace, .mapped_host_workspace = mapped_host_workspace, + .scaleout_rank_idx = scaleout_rank_idx, .scaleup_rank_idx = scaleup_rank_idx, + // NOTES: make cluster dim 2 to overlap with clustered computation kernels + .launch_args = jit::LaunchArgs(num_sms, num_threads, num_smem_bytes, 2 - (num_sms % 2), true)}; + const auto code = DispatchRuntime::generate(args); + const auto runtime = jit::compiler->build("dispatch", code); + DispatchRuntime::launch(runtime, args, stream); +} + +class DispatchCopyEpilogueRuntime final : public jit::LaunchRuntime { +public: + struct Args { + // Templated arguments + bool do_expand, cached_mode, do_zero_padding; + int num_channels; + int num_warps; + int num_scaleout_ranks, num_scaleup_ranks; + int num_hidden_bytes, num_sf_packs; + int num_max_tokens_per_rank; + int num_experts, num_topk, expert_alignment; + + // Parameters + void *buffer, *workspace; + int* psum_num_recv_tokens_per_scaleup_rank; + int* psum_num_recv_tokens_per_expert; + void* recv_x; void* recv_sf; + topk_idx_t* recv_topk_idx; float* recv_topk_weights; + int* recv_src_metadata; + int* channel_linked_list; + int* num_unaligned_recv_tokens_per_expert; + int num_recv_tokens; + int recv_sf_token_stride, recv_sf_hidden_stride; + int scaleout_rank_idx, scaleup_rank_idx; + + jit::LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_ep::elastic; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&dispatch_copy_epilogue_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>); +}} +)", + args.do_expand, args.cached_mode, args.do_zero_padding, + args.launch_args.grid_dim.first, args.num_channels, args.num_warps, + args.num_scaleout_ranks, args.num_scaleup_ranks, + args.num_hidden_bytes, args.num_sf_packs, + args.num_max_tokens_per_rank, + args.num_experts, args.num_topk, args.expert_alignment); + } + + static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) { + EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(kernel, config, + args.buffer, args.workspace, + args.psum_num_recv_tokens_per_scaleup_rank, + args.psum_num_recv_tokens_per_expert, + args.recv_x, args.recv_sf, args.recv_topk_idx, args.recv_topk_weights, + args.recv_src_metadata, + args.channel_linked_list, + args.num_unaligned_recv_tokens_per_expert, + args.num_recv_tokens, + args.recv_sf_token_stride, args.recv_sf_hidden_stride, + args.scaleout_rank_idx, args.scaleup_rank_idx)); + } +}; + +static void launch_dispatch_copy_epilogue(void* buffer, void* workspace, + int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, + void* recv_x, void* recv_sf, + topk_idx_t* recv_topk_idx, float* recv_topk_weights, + int* recv_src_metadata, + int* channel_linked_list, + int* num_unaligned_recv_tokens_per_expert, + const int& num_recv_tokens, const int& num_max_tokens_per_rank, + const int& num_hidden_bytes, + const int& num_sf_packs, const int& recv_sf_token_stride, const int& recv_sf_hidden_stride, + const int& num_experts, const int& num_topk, const int& expert_alignment, + const int& scaleout_rank_idx, const int& scaleup_rank_idx, + const int& num_scaleout_ranks, const int& num_scaleup_ranks, + const int& num_sms, const int& num_smem_bytes, + const int& num_channels, + const bool& do_expand, const bool& cached_mode, + const bool& do_zero_padding, + const at::cuda::CUDAStream& stream) { + // Maximize shared memory utilization + const auto token_layout = layout::TokenLayout(num_hidden_bytes, num_sf_packs * sizeof(sf_pack_t), num_topk, true); + const auto num_warps = std::min(num_smem_bytes / token_layout.get_num_bytes(), 32); + const auto num_threads = num_warps * 32; + + // Generate, build and launch + const DispatchCopyEpilogueRuntime::Args args = { + .do_expand = do_expand, .cached_mode = cached_mode, .do_zero_padding = do_zero_padding, + .num_channels = num_channels, .num_warps = num_warps, + .num_scaleout_ranks = num_scaleout_ranks, .num_scaleup_ranks = num_scaleup_ranks, + .num_hidden_bytes = num_hidden_bytes, .num_sf_packs = num_sf_packs, + .num_max_tokens_per_rank = num_max_tokens_per_rank, + .num_experts = num_experts, .num_topk = num_topk, .expert_alignment = expert_alignment, + .buffer = buffer, .workspace = workspace, + .psum_num_recv_tokens_per_scaleup_rank = psum_num_recv_tokens_per_scaleup_rank, + .psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert, + .recv_x = recv_x, .recv_sf = recv_sf, + .recv_topk_idx = recv_topk_idx, .recv_topk_weights = recv_topk_weights, + .recv_src_metadata = recv_src_metadata, + .channel_linked_list = channel_linked_list, + .num_unaligned_recv_tokens_per_expert = num_unaligned_recv_tokens_per_expert, + .num_recv_tokens = num_recv_tokens, + .recv_sf_token_stride = recv_sf_token_stride, .recv_sf_hidden_stride = recv_sf_hidden_stride, + .scaleout_rank_idx = scaleout_rank_idx, .scaleup_rank_idx = scaleup_rank_idx, + .launch_args = jit::LaunchArgs(num_sms, num_threads, num_smem_bytes, 1, false, true)}; + const auto code = DispatchCopyEpilogueRuntime::generate(args); + const auto runtime = jit::compiler->build("dispatch_copy_epilogue", code); + DispatchCopyEpilogueRuntime::launch(runtime, args, stream); +} + +} // namespace deep_ep::elastic diff --git a/csrc/kernels/elastic/engram.hpp b/csrc/kernels/elastic/engram.hpp new file mode 100644 index 0000000..dcbc7be --- /dev/null +++ b/csrc/kernels/elastic/engram.hpp @@ -0,0 +1,191 @@ +#pragma once + +#include +#include + +#include +#include +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/launch_runtime.hpp" + +namespace deep_ep::elastic { + +class EngramFetchRuntime final : public jit::LaunchRuntime { +public: + struct Args { + // Templated arguments + int num_entries_per_rank; + int num_hidden_bytes, num_sf_packs; + int num_entries_per_token; + int num_scaleout_ranks, num_scaleup_ranks; + int64_t num_cpu_bytes_per_rank; + int num_qps; + bool allow_hybrid_mode; + + // Parameters + ncclDevComm_t nccl_dev_comm; + ncclWindow_t nccl_window; + void* storage; + void* fetched; + int* indices; + ncclGinRequest_t* last_gin_requests; + sf_pack_t* sf_table; sf_pack_t* fetched_sf; + int sf_token_stride; int sf_hidden_stride; + int num_tokens; + + jit::LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + int num_rdma_peers, num_ranks_per_rdma_peer; + std::string team_tag; + if (args.allow_hybrid_mode) { + num_rdma_peers = args.num_scaleout_ranks; + num_ranks_per_rdma_peer = args.num_scaleup_ranks; + team_tag = "ncclTeamTagRail"; + } else { + num_rdma_peers = args.num_scaleout_ranks * args.num_scaleup_ranks; + num_ranks_per_rdma_peer = 1; + team_tag = "ncclTeamTagWorld"; + } + auto func_name = fmt::format("engram_fetch_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}>", + args.num_qps, args.num_entries_per_rank, args.num_hidden_bytes, args.num_sf_packs, + args.num_entries_per_token, num_rdma_peers, num_ranks_per_rdma_peer, + args.num_cpu_bytes_per_rank, args.launch_args.num_threads, team_tag); + + return fmt::format(R"( +#include + +using namespace deep_ep::elastic; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&{}); +}} +)", func_name); + } + + static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) { + EP_CUDA_UNIFIED_CHECK(jit::launch_kernel( + kernel, config, + args.nccl_dev_comm, args.nccl_window, + args.storage, args.fetched, + args.indices, + args.last_gin_requests, + args.sf_table, args.fetched_sf, + args.sf_token_stride, args.sf_hidden_stride, + args.num_tokens + )); + } +}; + +static void launch_engram_fetch(const ncclDevComm_t& nccl_dev_comm, const ncclWindow_t& nccl_window, + void* storage, void* fetched, + int* indices, + ncclGinRequest_t* last_gin_requests, + void* sf_table, void* fetched_sf, + const int& sf_token_stride, const int& sf_hidden_stride, + const int& num_entries_per_rank, + const int& hidden, const int& elem_size, const int& num_sf_packs, + const int& num_entries_per_token, + const int& num_tokens, + const int& num_scaleout_ranks, const int& num_scaleup_ranks, + const int64_t& num_cpu_bytes_per_rank, + const int& num_qps, + const bool& allow_hybrid_mode, + const at::cuda::CUDAStream& stream) { + constexpr int kNumEngramFetchThreads = 1024; + + // Generate, build and launch + const EngramFetchRuntime::Args args = { + .num_entries_per_rank = num_entries_per_rank, + .num_hidden_bytes = hidden * elem_size, + .num_sf_packs = num_sf_packs, + .num_entries_per_token = num_entries_per_token, + .num_scaleout_ranks = num_scaleout_ranks, + .num_scaleup_ranks = num_scaleup_ranks, + .num_cpu_bytes_per_rank = num_cpu_bytes_per_rank, + .num_qps = num_qps, + .allow_hybrid_mode = allow_hybrid_mode, + .nccl_dev_comm = nccl_dev_comm, + .nccl_window = nccl_window, + .storage = storage, + .fetched = fetched, + .indices = indices, + .last_gin_requests = last_gin_requests, + .sf_table = static_cast(sf_table), + .fetched_sf = static_cast(fetched_sf), + .sf_token_stride = sf_token_stride, + .sf_hidden_stride = sf_hidden_stride, + .num_tokens = num_tokens, + .launch_args = jit::LaunchArgs(num_qps, kNumEngramFetchThreads)}; + const auto code = EngramFetchRuntime::generate(args); + const auto runtime = jit::compiler->build("engram_fetch", code); + EngramFetchRuntime::launch(runtime, args, stream); +} + +class EngramFetchWaitRuntime final : public jit::LaunchRuntime { +public: + struct Args { + // Templated arguments + int num_scaleout_ranks, num_scaleup_ranks; + bool allow_hybrid_mode; + + ncclDevComm_t nccl_dev_comm; + ncclWindow_t nccl_window; + ncclGinRequest_t* last_gin_requests; + + jit::LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + const int num_rdma_peers = args.allow_hybrid_mode + ? args.num_scaleout_ranks + : args.num_scaleout_ranks * args.num_scaleup_ranks; + auto func_name = fmt::format("engram_fetch_wait_impl<{}, {}>", + num_rdma_peers, args.launch_args.num_threads); + + return fmt::format(R"( +#include + +using namespace deep_ep::elastic; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&{}); +}} +)", func_name); + } + + static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) { + EP_CUDA_UNIFIED_CHECK(jit::launch_kernel( + kernel, config, + args.nccl_dev_comm, args.nccl_window, + args.last_gin_requests + )); + } +}; + +static void launch_engram_fetch_wait(ncclGinRequest_t* last_gin_requests, + const ncclDevComm_t& nccl_dev_comm, const ncclWindow_t& nccl_window, + const int& num_scaleout_ranks, const int& num_scaleup_ranks, + const int& num_qps, + const bool& allow_hybrid_mode, + const at::cuda::CUDAStream& stream) { + constexpr int kNumEngramFetchWaitThreads = 1024; + + // Generate, build and launch + const EngramFetchWaitRuntime::Args args = { + .num_scaleout_ranks = num_scaleout_ranks, + .num_scaleup_ranks = num_scaleup_ranks, + .allow_hybrid_mode = allow_hybrid_mode, + .nccl_dev_comm = nccl_dev_comm, + .nccl_window = nccl_window, + .last_gin_requests = last_gin_requests, + .launch_args = jit::LaunchArgs(num_qps, kNumEngramFetchWaitThreads)}; + const auto code = EngramFetchWaitRuntime::generate(args); + const auto runtime = jit::compiler->build("engram_fetch_wait", code); + EngramFetchWaitRuntime::launch(runtime, args, stream); +} + +} // namespace deep_ep::elastic diff --git a/csrc/kernels/elastic/pp_send_recv.hpp b/csrc/kernels/elastic/pp_send_recv.hpp new file mode 100644 index 0000000..06cbea4 --- /dev/null +++ b/csrc/kernels/elastic/pp_send_recv.hpp @@ -0,0 +1,182 @@ +#pragma once + +#include +#include + +#include +#include +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/launch_runtime.hpp" + +namespace deep_ep::elastic { + +class PPSendRuntime final : public jit::LaunchRuntime { +public: + struct Args { + // Templated arguments + int num_ranks; + int num_smem_bytes; + int64_t num_timeout_cycles; + + // Parameters + ncclDevComm_t nccl_dev_comm; + ncclWindow_t nccl_window; + void* x; + int64_t num_x_bytes; + void* buffer; + void* workspace; + int rank_idx, dst_rank_idx; + int64_t num_max_tensor_bytes; + int num_max_inflight_tensors; + + jit::LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_ep::elastic; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&pp_send_impl<{}, {}, {}, {}>); +}} +)", args.launch_args.grid_dim.first, + args.num_ranks, + args.num_smem_bytes, + args.num_timeout_cycles); + } + + static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) { + EP_CUDA_UNIFIED_CHECK(jit::launch_kernel( + kernel, config, + args.nccl_dev_comm, args.nccl_window, + args.x, args.num_x_bytes, + args.buffer, args.workspace, + args.rank_idx, args.dst_rank_idx, + args.num_max_tensor_bytes, args.num_max_inflight_tensors + )); + } +}; + +static void launch_pp_send(const ncclDevComm_t& nccl_dev_comm, + const ncclWindow_t& nccl_window, + void* x, const int64_t& num_x_bytes, + void* buffer, void* workspace, + const int& rank_idx, const int& dst_rank_idx, const int& num_ranks, + const int64_t& num_max_tensor_bytes, + const int num_max_inflight_tensors, + const int& num_sms, + const int64_t& num_timeout_cycles, + const int& num_smem_bytes, + const at::cuda::CUDAStream& stream) { + // Generate, build and launch + const PPSendRuntime::Args args = { + .num_ranks = num_ranks, + .num_smem_bytes = num_smem_bytes, + .num_timeout_cycles = num_timeout_cycles, + .nccl_dev_comm = nccl_dev_comm, + .nccl_window = nccl_window, + .x = x, + .num_x_bytes = num_x_bytes, + .buffer = buffer, + .workspace = workspace, + .rank_idx = rank_idx, + .dst_rank_idx = dst_rank_idx, + .num_max_tensor_bytes = num_max_tensor_bytes, + .num_max_inflight_tensors = num_max_inflight_tensors, + .launch_args = jit::LaunchArgs(num_sms, 32, num_smem_bytes, 1, true) + }; + const auto code = PPSendRuntime::generate(args); + const auto runtime = jit::compiler->build("pp_send", code); + PPSendRuntime::launch(runtime, args, stream); +} + +class PPRecvRuntime final : public jit::LaunchRuntime { +public: + struct Args { + // Templated arguments + int num_ranks; + int num_smem_bytes; + int64_t num_timeout_cycles; + + // Parameters + ncclDevComm_t nccl_dev_comm; + ncclWindow_t nccl_window; + void* x; + int64_t num_x_bytes; + void* buffer; + void* workspace; + int rank_idx; + int src_rank_idx; + int64_t num_max_tensor_bytes; + int num_max_inflight_tensors; + + jit::LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_ep::elastic; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&pp_recv_impl<{}, {}, {}, {}>); +}} +)", args.launch_args.grid_dim.first, + args.num_ranks, + args.num_smem_bytes, + args.num_timeout_cycles); + } + + static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) { + EP_CUDA_UNIFIED_CHECK(jit::launch_kernel( + kernel, config, + args.nccl_dev_comm, args.nccl_window, + args.x, args.num_x_bytes, + args.buffer, args.workspace, + args.rank_idx, args.src_rank_idx, + args.num_max_tensor_bytes, + args.num_max_inflight_tensors + )); + } +}; + +static void launch_pp_recv(const ncclDevComm_t& nccl_dev_comm, + const ncclWindow_t& nccl_window, + void* x, + const int64_t& num_x_bytes, + void* buffer, void* workspace, + const int& rank_idx, const int& src_rank_idx, const int& num_ranks, + const int64_t& num_max_tensor_bytes, + const int& num_max_inflight_tensors, + const int& num_sms, + const int64_t& num_timeout_cycles, + const int& num_smem_bytes, + const at::cuda::CUDAStream& stream) { + // Generate, build and launch + const PPRecvRuntime::Args args = { + .num_ranks = num_ranks, + .num_smem_bytes = num_smem_bytes, + .num_timeout_cycles = num_timeout_cycles, + .nccl_dev_comm = nccl_dev_comm, + .nccl_window = nccl_window, + .x = x, + .num_x_bytes = num_x_bytes, + .buffer = buffer, + .workspace = workspace, + .rank_idx = rank_idx, + .src_rank_idx = src_rank_idx, + .num_max_tensor_bytes = num_max_tensor_bytes, + .num_max_inflight_tensors = num_max_inflight_tensors, + .launch_args = jit::LaunchArgs(num_sms, 32, num_smem_bytes, 1, true) + }; + const auto code = PPRecvRuntime::generate(args); + const auto runtime = jit::compiler->build("pp_recv", code); + PPRecvRuntime::launch(runtime, args, stream); +} + +} // namespace deep_ep::elastic diff --git a/csrc/kernels/legacy/CMakeLists.txt b/csrc/kernels/legacy/CMakeLists.txt new file mode 100644 index 0000000..29b20b1 --- /dev/null +++ b/csrc/kernels/legacy/CMakeLists.txt @@ -0,0 +1,7 @@ +add_deep_ep_library(legacy_layout_cuda layout.cu) +add_deep_ep_library(legacy_intranode_cuda intranode.cu) +add_deep_ep_library(legacy_internode_cuda internode.cu) +add_deep_ep_library(legacy_internode_ll_cuda internode_ll.cu) + +# Link these libraries later +set(LEGACY_CUDA_LIBRARIES legacy_layout_cuda legacy_intranode_cuda legacy_internode_cuda legacy_internode_ll_cuda CACHE INTERNAL "Legacy kernels") diff --git a/csrc/kernels/legacy/api.cuh b/csrc/kernels/legacy/api.cuh new file mode 100644 index 0000000..641028d --- /dev/null +++ b/csrc/kernels/legacy/api.cuh @@ -0,0 +1,326 @@ +#pragma once + +#include "compiled.cuh" + +namespace deep_ep::legacy { + +// 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 + +// Intranode kernels +namespace intranode { + +void barrier(int** barrier_signal_ptrs, int rank, int num_ranks, cudaStream_t stream); + +void notify_dispatch(const int* num_tokens_per_rank, + int* moe_recv_counter_mapped, + int num_ranks, + const int* num_tokens_per_expert, + int* moe_recv_expert_counter_mapped, + int num_experts, + int num_tokens, + const bool* is_token_in_rank, + int* channel_prefix_matrix, + int* rank_prefix_matrix_copy, + int num_memset_int, + int expert_alignment, + void** buffer_ptrs, + int** barrier_signal_ptrs, + int rank, + cudaStream_t stream, + int num_sms); + +void cached_notify_dispatch(const int* rank_prefix_matrix, + int num_memset_int, + void** buffer_ptrs, + int** barrier_signal_ptrs, + int rank, + int num_ranks, + cudaStream_t stream); + +void dispatch(void* recv_x, + float* recv_x_scales, + int* recv_src_idx, + topk_idx_t* recv_topk_idx, + float* recv_topk_weights, + int* recv_channel_offset, + int* send_head, + const void* x, + const float* x_scales, + const topk_idx_t* topk_idx, + const float* topk_weights, + const bool* is_token_in_rank, + const int* channel_prefix_matrix, + int num_tokens, + int num_worst_tokens, + int hidden_int4, + int num_topk, + int num_experts, + int num_scales, + int scale_token_stride, + int scale_hidden_stride, + void** buffer_ptrs, + int rank, + int num_ranks, + cudaStream_t stream, + int num_sms, + int num_max_send_tokens, + int num_recv_buffer_tokens); + +void cached_notify_combine(void** buffer_ptrs, + int* send_head, + int num_channels, + int num_recv_tokens, + int num_memset_int, + int** barrier_signal_ptrs, + int rank, + int num_ranks, + cudaStream_t stream); + +void combine(cudaDataType_t type, + void* recv_x, + float* recv_topk_weights, + const void* x, + const float* topk_weights, + const void* bias_0, + const void* bias_1, + const int* src_idx, + const int* rank_prefix_matrix, + const int* channel_prefix_matrix, + int* send_head, + int num_tokens, + int num_recv_tokens, + int hidden, + int num_topk, + void** buffer_ptrs, + int rank, + int num_ranks, + cudaStream_t stream, + int num_sms, + int num_max_send_tokens, + int num_recv_buffer_tokens); + +} // namespace intranode + +// Internode kernels +namespace internode { + +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_worst_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); + +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 num_worst_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); + +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); + +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); + +} // namespace internode + +// Internode low-latency kernels +namespace internode_ll { + +void clean_low_latency_buffer(int* clean_0, + int num_clean_int_0, + int* clean_1, + int num_clean_int_1, + int rank, + int num_ranks, + int* mask_buffer, + int* sync_buffer, + cudaStream_t stream); + +void dispatch(void* packed_recv_x, + void* packed_recv_x_scales, + int* packed_recv_src_info, + int64_t* packed_recv_layout_range, + int* packed_recv_count, + int* mask_buffer, + int* cumulative_local_expert_recv_stats, + int64_t* dispatch_wait_recv_cost_stats, + void* rdma_recv_x, + int* rdma_recv_count, + void* rdma_x, + const void* x, + const topk_idx_t* topk_idx, + int* next_clean, + int num_next_clean_int, + int num_tokens, + int hidden, + int num_max_dispatch_tokens_per_rank, + int num_topk, + int num_experts, + int rank, + int num_ranks, + bool use_fp8, + bool round_scale, + bool use_ue8m0, + void* workspace, + int num_device_sms, + cudaStream_t stream, + int phases); + +void combine(void* combined_x, + void* rdma_recv_x, + int* rdma_recv_flag, + void* rdma_send_x, + const void* x, + const topk_idx_t* topk_idx, + const float* topk_weights, + const int* src_info, + const int64_t* layout_range, + int* mask_buffer, + int64_t* combine_wait_recv_cost_stats, + int* next_clean, + int num_next_clean_int, + int num_combined_tokens, + int hidden, + int num_max_dispatch_tokens_per_rank, + int num_topk, + int num_experts, + int rank, + int num_ranks, + bool use_logfmt, + void* workspace, + int num_device_sms, + cudaStream_t stream, + int phases, + bool zero_copy); + +void query_mask_buffer(int* mask_buffer_ptr, int num_ranks, int* output_mask_tensor, cudaStream_t stream); + +void update_mask_buffer(int* mask_buffer_ptr, int rank_to_mask, bool mask, cudaStream_t stream); + +void clean_mask_buffer(int* mask_buffer_ptr, int num_ranks, cudaStream_t stream); + +} // namespace internode_ll + +} // namespace deep_ep::legacy diff --git a/csrc/kernels/legacy/buffer.cuh b/csrc/kernels/legacy/buffer.cuh new file mode 100644 index 0000000..6b9942d --- /dev/null +++ b/csrc/kernels/legacy/buffer.cuh @@ -0,0 +1,132 @@ +#pragma once + +#include + +#include "compiled.cuh" + +namespace deep_ep::legacy { + +template +struct Buffer { +private: + uint8_t* ptr; + +public: + int64_t 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 = static_cast(gbl_ptr) + offset * sizeof(dtype_t); + gbl_ptr = static_cast(gbl_ptr) + total_bytes; + } + + __device__ __forceinline__ Buffer advance_also(void*& gbl_ptr) { + gbl_ptr = static_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]; + int64_t num_bytes; + +public: + int64_t 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); + + int64_t per_channel_bytes = num_bytes * num_ranks; + total_bytes = per_channel_bytes * num_sms; + ptrs[0] = static_cast(gbl_ptr) + per_channel_bytes * sm_id + num_bytes * offset; + gbl_ptr = static_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); + + int64_t per_channel_bytes = num_bytes * num_ranks; + total_bytes = per_channel_bytes * num_sms; + for (int i = 0; i < kNumRanks; ++i) { + ptrs[i] = static_cast(gbl_ptrs[i]) + per_channel_bytes * sm_id + num_bytes * offset; + gbl_ptrs[i] = static_cast(gbl_ptrs[i]) + total_bytes; + } + } + + __device__ __forceinline__ void advance(int shift) { + #pragma unroll + 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 = static_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] = static_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; + int64_t num_bytes; + +public: + int64_t 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); + + int64_t per_channel_bytes = num_bytes * num_ranks; + total_bytes = per_channel_bytes * num_sms * (static_cast(kDecoupled) + 1); + send_ptr = static_cast(gbl_ptr) + per_channel_bytes * sm_id; + recv_ptr = static_cast(gbl_ptr) + per_channel_bytes * (sm_id + num_sms); + gbl_ptr = static_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 deep_ep::legacy diff --git a/csrc/kernels/legacy/compiled.cuh b/csrc/kernels/legacy/compiled.cuh new file mode 100644 index 0000000..1f59d6d --- /dev/null +++ b/csrc/kernels/legacy/compiled.cuh @@ -0,0 +1,18 @@ +#pragma once + +#include + +#define LEGACY_NUM_MAX_NVL_PEERS 8 +#define LEGACY_NUM_MAX_RDMA_PEERS 20 +#define LEGACY_NUM_WORKSPACE_BYTES (32 * 1024 * 1024) +#define LEGACY_NUM_MAX_LOCAL_EXPERTS 1024 +#define LEGACY_NUM_BUFFER_ALIGNMENT_BYTES 128 + +#define LEGACY_LOW_LATENCY_SEND_PHASE 1 +#define LEGACY_LOW_LATENCY_RECV_PHASE 2 + +#define LEGACY_FINISHED_SUM_TAG 1024 +#define LEGACY_NUM_WAIT_NANOSECONDS 500 + +#define LEGACY_NUM_CPU_TIMEOUT_SECS 100 +#define LEGACY_NUM_TIMEOUT_CYCLES 200000000000ull // 200G cycles ~= 100s diff --git a/csrc/kernels/legacy/ibgda_device.cuh b/csrc/kernels/legacy/ibgda_device.cuh new file mode 100644 index 0000000..819bdd1 --- /dev/null +++ b/csrc/kernels/legacy/ibgda_device.cuh @@ -0,0 +1,496 @@ +// Portions derived from NVSHMEM (https://developer.nvidia.com/nvshmem) +// Copyright (c) NVIDIA Corporation. +// Licensed under the NVSHMEM Software License Agreement (version: September 3, 2019). +// See full license at: https://docs.nvidia.com/nvshmem/api/sla.html +// +// Modified from original source: +// - nvshmem/src/include/non_abi/device/pt-to-pt/ibgda_device.cuh +#pragma once + + +#include +#include +#include + +#include + +#include "utils.cuh" + +namespace deep_ep::legacy { + +EP_STATIC_ASSERT(NVSHMEMI_IBGDA_MIN_QP_DEPTH >= 64, "Invalid QP minimum depth"); + +__device__ static __forceinline__ uint64_t HtoBE64(uint64_t x) { + uint64_t ret; + asm("{\n\t" + ".reg .b32 ign;\n\t" + ".reg .b32 lo;\n\t" + ".reg .b32 hi;\n\t" + ".reg .b32 new_lo;\n\t" + ".reg .b32 new_hi;\n\t" + "mov.b64 {lo,hi}, %1;\n\t" + "prmt.b32 new_hi, lo, ign, 0x0123;\n\t" + "prmt.b32 new_lo, hi, ign, 0x0123;\n\t" + "mov.b64 %0, {new_lo,new_hi};\n\t" + "}" + : "=l"(ret) + : "l"(x)); + return ret; +} + +__device__ static __forceinline__ uint32_t HtoBE32(uint32_t x) { + uint32_t ret; + asm("{\n\t" + ".reg .b32 ign;\n\t" + "prmt.b32 %0, %1, ign, 0x0123;\n\t" + "}" + : "=r"(ret) + : "r"(x)); + return ret; +} + +__device__ static __forceinline__ uint16_t HtoBE16(uint16_t x) { + // TODO: simplify PTX using 16-bit instructions + auto a = static_cast(x); + uint32_t d; + asm volatile( + "{\n\t" + ".reg .b32 mask;\n\t" + ".reg .b32 ign;\n\t" + "mov.b32 mask, 0x4401;\n\t" + "mov.b32 ign, 0x0;\n\t" + "prmt.b32 %0, %1, ign, mask;\n\t" + "}" + : "=r"(d) + : "r"(a)); + return static_cast(d); +} + +typedef struct mlx5_wqe_ctrl_seg __attribute__((__aligned__(8))) ibgda_ctrl_seg_t; + +typedef struct { + uint32_t add_data; + uint32_t field_boundary; + uint64_t reserved; +} __attribute__((__packed__)) ibgda_atomic_32_masked_fa_seg_t; + +__device__ static __forceinline__ nvshmemi_ibgda_device_state_t* ibgda_get_state() { + return &nvshmemi_ibgda_device_state_d; +} + +__device__ static __forceinline__ nvshmemi_ibgda_device_qp_t* ibgda_get_rc(int pe, int id) { + auto state = ibgda_get_state(); + const auto num_rc_per_pe = ibgda_get_state()->num_rc_per_pe; + return &state->globalmem + .rcs[pe * num_rc_per_pe * state->num_devices_initialized + id % (num_rc_per_pe * state->num_devices_initialized)]; +} + +__device__ static __forceinline__ void ibgda_lock_acquire(int* lock) { + while (atomicCAS(lock, 0, 1) == 1) + ; + + // Prevent reordering before the lock is acquired + memory_fence_cta(); +} + +__device__ static __forceinline__ void ibgda_lock_release(int* lock) { + memory_fence_cta(); + + // Prevent reordering before lock is released + st_na_relaxed(lock, 0); +} + +__device__ static __forceinline__ void ibgda_update_dbr(nvshmemi_ibgda_device_qp_t* qp, uint32_t dbrec_head) { + // `DBREC` contains the index of the next empty `WQEBB` + __be32 dbrec_val; + __be32* dbrec_ptr = qp->tx_wq.dbrec; + + // This is equivalent to `WRITE_ONCE(dbrec_ptr, HtoBE32(dbrec_head & 0xffff))` + asm("{\n\t" + ".reg .b32 dbrec_head_16b;\n\t" + ".reg .b32 ign;\n\t" + "and.b32 dbrec_head_16b, %1, 0xffff;\n\t" + "prmt.b32 %0, dbrec_head_16b, ign, 0x123;\n\t" + "}" + : "=r"(dbrec_val) + : "r"(dbrec_head)); + st_na_release(dbrec_ptr, dbrec_val); +} + +__device__ static __forceinline__ void ibgda_ring_db(nvshmemi_ibgda_device_qp_t* qp, uint16_t prod_idx) { + auto bf_ptr = reinterpret_cast(qp->tx_wq.bf); + ibgda_ctrl_seg_t ctrl_seg = {.opmod_idx_opcode = HtoBE32(prod_idx << 8), .qpn_ds = HtoBE32(qp->qpn << 8)}; + + EP_STATIC_ASSERT(sizeof(decltype(&ctrl_seg)) == sizeof(uint64_t), ""); + st_na_release(bf_ptr, *(reinterpret_cast(&ctrl_seg))); +} + +__device__ static __forceinline__ void ibgda_post_send(nvshmemi_ibgda_device_qp_t* qp, uint64_t new_prod_idx) { + nvshmemi_ibgda_device_qp_management_t* mvars = &qp->mvars; + uint64_t old_prod_idx; + + // Update `prod_idx` before ringing the doorbell, so that we know which index is needed in quiet/fence + ibgda_lock_acquire(&mvars->post_send_lock); + + old_prod_idx = atomicMax(reinterpret_cast(&mvars->tx_wq.prod_idx), new_prod_idx); + if (new_prod_idx > old_prod_idx) { + ibgda_update_dbr(qp, new_prod_idx); + ibgda_ring_db(qp, new_prod_idx); + } + ibgda_lock_release(&mvars->post_send_lock); +} + +template +__device__ static __forceinline__ void ibgda_submit_requests(nvshmemi_ibgda_device_qp_t* qp, + uint64_t base_wqe_idx, + uint32_t num_wqes, + int message_idx = 0) { + auto state = ibgda_get_state(); + nvshmemi_ibgda_device_qp_management_t* mvars = &qp->mvars; + uint64_t new_wqe_idx = base_wqe_idx + num_wqes; + + // WQE writes must be finished first + __threadfence(); + + unsigned long long int* ready_idx = + (unsigned long long int*)(state->use_async_postsend ? qp->tx_wq.prod_idx : &mvars->tx_wq.ready_head); + + // Wait for prior WQE slots to be filled first + while (atomicCAS(ready_idx, base_wqe_idx, new_wqe_idx) != base_wqe_idx) + ; + + // Always post, not in batch + if (!state->use_async_postsend) { + constexpr int kNumRequestInBatch = 4; + if (kAlwaysDoPostSend or (message_idx + 1) % kNumRequestInBatch == 0) + ibgda_post_send(qp, new_wqe_idx); + } +} + +__device__ static __forceinline__ void ibgda_write_rdma_write_inl_wqe( + nvshmemi_ibgda_device_qp_t* qp, const uint32_t* val, uint64_t raddr, __be32 rkey, uint16_t wqe_idx, void** out_wqes, uint32_t imm) { + ibgda_ctrl_seg_t ctrl_seg; + struct mlx5_wqe_raddr_seg raddr_seg; + struct mlx5_wqe_inl_data_seg inl_seg; + + auto* ctrl_seg_ptr = reinterpret_cast(out_wqes[0]); + auto* raddr_seg_ptr = reinterpret_cast(reinterpret_cast(ctrl_seg_ptr) + sizeof(*ctrl_seg_ptr)); + auto* inl_seg_ptr = reinterpret_cast(reinterpret_cast(raddr_seg_ptr) + sizeof(*raddr_seg_ptr)); + auto* wqe_data_ptr = reinterpret_cast(reinterpret_cast(inl_seg_ptr) + sizeof(*inl_seg_ptr)); + + raddr_seg.raddr = HtoBE64(raddr); + raddr_seg.rkey = rkey; + raddr_seg.reserved = 0; + + inl_seg.byte_count = HtoBE32(4 | MLX5_INLINE_SEG); + + // `imm == std::numeric_limits::max()` means no imm writes + ctrl_seg = {0}; + ctrl_seg.qpn_ds = HtoBE32((qp->qpn << 8) | 3); + ctrl_seg.fm_ce_se = MLX5_WQE_CTRL_CQ_UPDATE; + ctrl_seg.opmod_idx_opcode = + HtoBE32((wqe_idx << 8) | (imm != std::numeric_limits::max() ? MLX5_OPCODE_RDMA_WRITE_IMM : MLX5_OPCODE_RDMA_WRITE)); + if (imm != std::numeric_limits::max()) + ctrl_seg.imm = HtoBE32(imm); + + EP_STATIC_ASSERT(sizeof(*ctrl_seg_ptr) == 16, "sizeof(*ctrl_seg_ptr) == 16"); + EP_STATIC_ASSERT(sizeof(*raddr_seg_ptr) == 16, "sizeof(*raddr_seg_ptr) == 16"); + EP_STATIC_ASSERT(sizeof(*inl_seg_ptr) == 4, "sizeof(*inl_seg_ptr) == 4"); + st_na_relaxed(reinterpret_cast(ctrl_seg_ptr), *reinterpret_cast(&ctrl_seg)); + st_na_relaxed(reinterpret_cast(raddr_seg_ptr), *reinterpret_cast(&raddr_seg)); + st_na_relaxed(reinterpret_cast(inl_seg_ptr), *reinterpret_cast(&inl_seg)); + st_na_relaxed(reinterpret_cast(wqe_data_ptr), *reinterpret_cast(val)); +} + +__device__ static __forceinline__ uint64_t +ibgda_get_lkey_and_rkey(uint64_t laddr, __be32* lkey, uint64_t raddr, int dst_pe, uint64_t* out_raddr, __be32* out_rkey, uint32_t dev_idx) { + auto state = ibgda_get_state(); + auto heap_start = reinterpret_cast(nvshmemi_device_state_d.heap_base); + auto log2_cumem_granularity = state->log2_cumem_granularity; + + // Local key + uint64_t idx = ((laddr - heap_start) >> log2_cumem_granularity) * state->num_devices_initialized + dev_idx; + auto device_key = state->constmem.lkeys[idx]; + auto lchunk_size = device_key.next_addr - laddr; + *lkey = device_key.key; + + // Remote key + uint64_t roffset = raddr - heap_start; + + idx = ((roffset >> log2_cumem_granularity) * nvshmemi_device_state_d.npes) * state->num_devices_initialized + + dst_pe * state->num_devices_initialized + dev_idx; + if (idx < NVSHMEMI_IBGDA_MAX_CONST_RKEYS) { + device_key = state->constmem.rkeys[idx]; + } else { + device_key = state->globalmem.rkeys[idx - NVSHMEMI_IBGDA_MAX_CONST_RKEYS]; + } + *out_raddr = reinterpret_cast(nvshmemi_device_state_d.peer_heap_base_remote[dst_pe]) + roffset; + *out_rkey = device_key.key; + + // Return the minimum of local and remote chunk sizes + auto rchunk_size = device_key.next_addr - roffset; + return min(lchunk_size, rchunk_size); +} + +__device__ static __forceinline__ void ibgda_get_rkey(uint64_t addr, int dst_pe, uint64_t* out_raddr, __be32* out_rkey, uint32_t dev_idx) { + auto state = ibgda_get_state(); + auto heap_start = reinterpret_cast(nvshmemi_device_state_d.heap_base); + + uint64_t roffset = addr - heap_start; + uint64_t idx = ((roffset >> state->log2_cumem_granularity) * nvshmemi_device_state_d.npes * state->num_devices_initialized) + + dst_pe * state->num_devices_initialized + dev_idx; + nvshmemi_ibgda_device_key_t device_key; + if (idx < NVSHMEMI_IBGDA_MAX_CONST_RKEYS) + device_key = state->constmem.rkeys[idx]; + else + device_key = state->globalmem.rkeys[idx - NVSHMEMI_IBGDA_MAX_CONST_RKEYS]; + *out_raddr = reinterpret_cast(nvshmemi_device_state_d.peer_heap_base_remote[dst_pe]) + roffset; + *out_rkey = device_key.key; +} + +__device__ static __forceinline__ uint64_t ibgda_reserve_wqe_slots(nvshmemi_ibgda_device_qp_t* qp, uint32_t num_wqes) { + auto mvars = &qp->mvars; + return atomicAdd(reinterpret_cast(&mvars->tx_wq.resv_head), static_cast(num_wqes)); +} + +__device__ static __forceinline__ void* ibgda_get_wqe_ptr(nvshmemi_ibgda_device_qp_t* qp, uint16_t wqe_idx) { + uint16_t cnt = qp->tx_wq.nwqes; + uint16_t idx = wqe_idx & (cnt - 1); + return reinterpret_cast(reinterpret_cast(qp->tx_wq.wqe) + (idx << MLX5_SEND_WQE_SHIFT)); +} + +__device__ static __forceinline__ void nvshmemi_ibgda_rma_p( + int* rptr, const int value, int dst_pe, int qp_id, uint32_t imm = std::numeric_limits::max()) { + // Get rkey + // NOTES: the `p` operation will not cross multiple remote chunks + __be32 rkey; + uint64_t raddr; + auto qp = ibgda_get_rc(dst_pe, qp_id); + ibgda_get_rkey(reinterpret_cast(rptr), dst_pe, &raddr, &rkey, qp->dev_idx); + + // Write WQEs + uint64_t base_wqe_idx = ibgda_reserve_wqe_slots(qp, 1); + void* wqe_ptrs; + wqe_ptrs = ibgda_get_wqe_ptr(qp, base_wqe_idx); + ibgda_write_rdma_write_inl_wqe(qp, reinterpret_cast(&value), raddr, rkey, base_wqe_idx, &wqe_ptrs, imm); + + // Submit requests + ibgda_submit_requests(qp, base_wqe_idx, 1); +} + +__device__ static __forceinline__ void ibgda_write_rdma_write_wqe(nvshmemi_ibgda_device_qp_t* qp, + uint64_t laddr, + __be32 lkey, + uint64_t raddr, + __be32 rkey, + uint32_t bytes, + uint16_t wqe_idx, + void** out_wqes) { + ibgda_ctrl_seg_t ctrl_seg; + struct mlx5_wqe_raddr_seg raddr_seg; + struct mlx5_wqe_data_seg data_seg; + + auto* ctrl_seg_ptr = reinterpret_cast(out_wqes[0]); + void* av_seg_ptr = reinterpret_cast(reinterpret_cast(ctrl_seg_ptr) + sizeof(*ctrl_seg_ptr)); + struct mlx5_wqe_raddr_seg* raddr_seg_ptr; + struct mlx5_wqe_data_seg* data_seg_ptr; + + raddr_seg_ptr = reinterpret_cast(reinterpret_cast(av_seg_ptr)); + data_seg_ptr = reinterpret_cast(reinterpret_cast(raddr_seg_ptr) + sizeof(*raddr_seg_ptr)); + + raddr_seg.raddr = HtoBE64(raddr); + raddr_seg.rkey = rkey; + raddr_seg.reserved = 0; + + data_seg.byte_count = HtoBE32(bytes); + data_seg.lkey = lkey; + data_seg.addr = HtoBE64(laddr); + + ctrl_seg = {0}; + ctrl_seg.qpn_ds = HtoBE32((qp->qpn << 8) | 3); + ctrl_seg.fm_ce_se = MLX5_WQE_CTRL_CQ_UPDATE; + ctrl_seg.opmod_idx_opcode = HtoBE32((wqe_idx << 8) | MLX5_OPCODE_RDMA_WRITE); + + EP_STATIC_ASSERT(sizeof(*ctrl_seg_ptr) == 16, "sizeof(*ctrl_seg_ptr) == 16"); + EP_STATIC_ASSERT(sizeof(*raddr_seg_ptr) == 16, "sizeof(*raddr_seg_ptr) == 16"); + EP_STATIC_ASSERT(sizeof(*data_seg_ptr) == 16, "sizeof(*data_seg_ptr) == 16"); + st_na_relaxed(reinterpret_cast(ctrl_seg_ptr), *reinterpret_cast(&ctrl_seg)); + st_na_relaxed(reinterpret_cast(raddr_seg_ptr), *reinterpret_cast(&raddr_seg)); + st_na_relaxed(reinterpret_cast(data_seg_ptr), *reinterpret_cast(&data_seg)); +} + +__device__ static __forceinline__ void ibgda_write_empty_recv_wqe(void* out_wqe) { + auto* data_seg_ptr = reinterpret_cast(out_wqe); + struct mlx5_wqe_data_seg data_seg; + + // Make the first segment in the WQE invalid, then the entire list will be invalid + data_seg.byte_count = 0; + data_seg.lkey = HtoBE64(MLX5_INVALID_LKEY); + data_seg.addr = 0; + + EP_STATIC_ASSERT(sizeof(mlx5_wqe_data_seg) == sizeof(int4), "Invalid data type length"); + st_na_relaxed(reinterpret_cast(data_seg_ptr), *reinterpret_cast(&data_seg)); +} + +template +__device__ static __forceinline__ void nvshmemi_ibgda_put_nbi_warp( + uint64_t req_rptr, uint64_t req_lptr, size_t bytes, int dst_pe, int qp_id, int lane_id, int message_idx) { + // Get lkey and rkey, store them into lanes + uint32_t num_wqes = 0; + __be32 my_lkey = 0; + uint64_t my_laddr = 0; + __be32 my_rkey = 0; + uint64_t my_raddr = 0; + uint64_t my_chunk_size = 0; + + auto qp = ibgda_get_rc(dst_pe, qp_id); + + // Decide how many messages (theoretically 3 for maximum) + auto remaining_bytes = bytes; + while (remaining_bytes > 0) { + if (lane_id == num_wqes) { + my_chunk_size = min(remaining_bytes, + ibgda_get_lkey_and_rkey(my_laddr = req_lptr, &my_lkey, req_rptr, dst_pe, &my_raddr, &my_rkey, qp->dev_idx)); + } + + // Move one more message + auto chunk_size = __shfl_sync(0xffffffff, my_chunk_size, static_cast(num_wqes)); + remaining_bytes -= chunk_size; + req_lptr += chunk_size; + req_rptr += chunk_size; + ++num_wqes; + } + EP_DEVICE_ASSERT(num_wqes <= 32); + + // Process WQE + uint64_t base_wqe_idx = 0; + if (lane_id == 0) + base_wqe_idx = ibgda_reserve_wqe_slots(qp, num_wqes); + base_wqe_idx = __shfl_sync(0xffffffff, base_wqe_idx, 0); + if (lane_id < num_wqes) { + auto wqe_idx = base_wqe_idx + lane_id; + auto wqe_ptr = ibgda_get_wqe_ptr(qp, wqe_idx); + ibgda_write_rdma_write_wqe(qp, my_laddr, my_lkey, my_raddr, my_rkey, my_chunk_size, wqe_idx, &wqe_ptr); + } + __syncwarp(); + + // Submit + if (lane_id == 0) + ibgda_submit_requests(qp, base_wqe_idx, num_wqes, message_idx); + __syncwarp(); +} + +__device__ static __forceinline__ void ibgda_write_amo_add_wqe(nvshmemi_ibgda_device_qp_t* qp, + const int& value, + uint64_t laddr, + __be32 lkey, + uint64_t raddr, + __be32 rkey, + uint16_t wqe_idx, + void** out_wqes) { + ibgda_ctrl_seg_t ctrl_seg = {0}; + struct mlx5_wqe_raddr_seg raddr_seg; + struct mlx5_wqe_atomic_seg atomic_seg_1; + struct mlx5_wqe_data_seg data_seg; + + auto ctrl_seg_ptr = reinterpret_cast(out_wqes[0]); + auto raddr_seg_ptr = reinterpret_cast(reinterpret_cast(ctrl_seg_ptr) + sizeof(*ctrl_seg_ptr)); + auto atomic_seg_ptr = reinterpret_cast(reinterpret_cast(raddr_seg_ptr) + sizeof(*raddr_seg_ptr)); + auto data_seg_ptr = reinterpret_cast(reinterpret_cast(atomic_seg_ptr) + sizeof(*atomic_seg_ptr)); + + raddr_seg.raddr = HtoBE64(raddr); + raddr_seg.rkey = rkey; + raddr_seg.reserved = 0; + + // NOTES: `0x08000000` means `IBGDA_4_BYTE_EXT_AMO_OPMOD` + ctrl_seg.opmod_idx_opcode = HtoBE32(MLX5_OPCODE_ATOMIC_MASKED_FA | (wqe_idx << 8) | 0x08000000); + auto atomic_32_masked_fa_seg = reinterpret_cast(&atomic_seg_1); + atomic_32_masked_fa_seg->add_data = HtoBE32(value); + atomic_32_masked_fa_seg->field_boundary = 0; + + ctrl_seg.qpn_ds = HtoBE32((qp->qpn << 8) | 4); + ctrl_seg.fm_ce_se = MLX5_WQE_CTRL_CQ_UPDATE; + + data_seg.byte_count = HtoBE32(sizeof(int)); + data_seg.lkey = lkey; + data_seg.addr = HtoBE64(laddr); + + EP_STATIC_ASSERT(sizeof(*ctrl_seg_ptr) == sizeof(int4), "Invalid vectorization"); + EP_STATIC_ASSERT(sizeof(*raddr_seg_ptr) == sizeof(int4), "Invalid vectorization"); + EP_STATIC_ASSERT(sizeof(*atomic_seg_ptr) == sizeof(int4), "Invalid vectorization"); + EP_STATIC_ASSERT(sizeof(*data_seg_ptr) == sizeof(int4), "Invalid vectorization"); + st_na_relaxed(reinterpret_cast(ctrl_seg_ptr), *reinterpret_cast(&ctrl_seg)); + st_na_relaxed(reinterpret_cast(raddr_seg_ptr), *reinterpret_cast(&raddr_seg)); + st_na_relaxed(reinterpret_cast(atomic_seg_ptr), *reinterpret_cast(&atomic_seg_1)); + st_na_relaxed(reinterpret_cast(data_seg_ptr), *reinterpret_cast(&data_seg)); +} + +__device__ __forceinline__ void nvshmemi_ibgda_amo_nonfetch_add( + void* rptr, const int& value, int pe, int qp_id, bool is_local_copy = false) { + if (is_local_copy) { + atomicAdd(static_cast(rptr), value); + } else { + nvshmemi_ibgda_device_qp_t* qp = ibgda_get_rc(pe, qp_id); + + __be32 rkey; + uint64_t raddr; + ibgda_get_rkey(reinterpret_cast(rptr), pe, &raddr, &rkey, qp->dev_idx); + + uint64_t my_wqe_idx = ibgda_reserve_wqe_slots(qp, 1); + void* wqe_ptrs = ibgda_get_wqe_ptr(qp, my_wqe_idx); + + ibgda_write_amo_add_wqe(qp, value, reinterpret_cast(qp->ibuf.buf), qp->ibuf.lkey, raddr, rkey, my_wqe_idx, &wqe_ptrs); + + ibgda_submit_requests(qp, my_wqe_idx, 1); + } +} + +__device__ __forceinline__ uint64_t nvshmemi_get_p2p_ptr(const uint64_t& ptr, const int& rank, const int& dst_rank) { + // Local rank, no need for mapping + if (rank == dst_rank) + return ptr; + auto peer_base = __ldg(reinterpret_cast(nvshmemi_device_state_d.peer_heap_base_p2p) + dst_rank); + + // RDMA connected + if (peer_base == 0) + return 0; + + // NVLink P2P is enabled + return peer_base + (ptr - reinterpret_cast(nvshmemi_device_state_d.heap_base)); +} + +// This is a simplified version of NVSHMEM's `ibgda_poll_cq`. +// Note that this implementation does not guarantee thread safety, +// so we must ensure that no other threads are concurrently using the same QP. +__device__ static __forceinline__ void ibgda_poll_cq(nvshmemi_ibgda_device_cq_t* cq, uint64_t idx) { + const auto cqe64 = static_cast(cq->cqe); + const uint32_t ncqes = cq->ncqes; + memory_fence_cta(); + if (*cq->cons_idx >= idx) + return; + // NOTES: this while loop is part of do-while below. + // `wqe_counter` is the HW consumer index. However, we always maintain `index + 1`. + // To be able to compare with the index, we need to use `wqe_counter + 1`. + // Because `wqe_counter` is `uint16_t`, it may be overflow. Still, we know for + // sure that if `idx - wqe_counter - 1 < ncqes`, `wqe_counter + 1 is less than + // idx, and thus we need to wait. We don't need to wait when `idx == wqe_counter + 1` + // That's why we use `- 2` here to make this case overflow. + uint16_t wqe_counter; + do { + wqe_counter = HtoBE16(ld_na_relaxed(&cqe64->wqe_counter)); + } while ((static_cast(static_cast(idx) - wqe_counter - static_cast(2)) < ncqes)); + *cq->cons_idx = idx; + + // Prevent reordering of this function and later instructions + memory_fence_cta(); +} + +// Wait until wqe `idx - 1` is completed. +__device__ static __forceinline__ void nvshmemi_ibgda_quiet(int dst_pe, int qp_id) { + auto qp = ibgda_get_rc(dst_pe, qp_id); + auto state = ibgda_get_state(); + uint64_t prod_idx = state->use_async_postsend ? ld_na_relaxed(qp->tx_wq.prod_idx) : ld_na_relaxed(&qp->mvars.tx_wq.ready_head); + ibgda_poll_cq(qp->tx_wq.cq, prod_idx); +} + +} // namespace deep_ep::legacy diff --git a/csrc/kernels/legacy/internode.cu b/csrc/kernels/legacy/internode.cu new file mode 100644 index 0000000..adc2629 --- /dev/null +++ b/csrc/kernels/legacy/internode.cu @@ -0,0 +1,2384 @@ +#include +#include + +#include "buffer.cuh" +#include "compiled.cuh" +#include "ibgda_device.cuh" +#include "launch.cuh" +#include "utils.cuh" + +namespace deep_ep { + +namespace nvshmem { + +extern nvshmem_team_t cpu_rdma_team; + +} // namespace nvshmem + +namespace legacy { + +namespace internode { + +struct SourceMeta { + int src_rdma_rank, is_token_in_nvl_rank_bits; + + EP_STATIC_ASSERT(LEGACY_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 < LEGACY_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), + (LEGACY_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 kLowLatencyMode ? (dst_rdma_rank * LEGACY_NUM_MAX_NVL_PEERS + nvl_rank) : dst_rdma_rank; +} + +template +__forceinline__ __device__ void nvshmem_sync_with_same_gpu_idx(const nvshmem_team_t& rdma_team) { + kLowLatencyMode ? void(nvshmem_sync(rdma_team)) : nvshmem_sync_all(); +} + +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_worst_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, + const nvshmem_team_t rdma_team) { + 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 / LEGACY_NUM_MAX_NVL_PEERS, nvl_rank = rank % LEGACY_NUM_MAX_NVL_PEERS; + auto num_rdma_experts = num_experts / kNumRDMARanks, num_nvl_experts = num_rdma_experts / LEGACY_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); + + // waiting for all previous inflight wrs to complete, + // in case of rewriting cleared rdma_buffer + auto qps_per_rdma_rank = ibgda_get_state()->num_rc_per_pe * ibgda_get_state()->num_devices_initialized; + for (int i = thread_id; i < qps_per_rdma_rank * (kNumRDMARanks - 1); i += num_threads) { + auto dst_rdma_rank = (i / qps_per_rdma_rank + rdma_rank + 1) % kNumRDMARanks; + auto qp_id = i % qps_per_rdma_rank; + nvshmemi_ibgda_quiet(translate_dst_rdma_rank(dst_rdma_rank, nvl_rank), qp_id); + } + __syncthreads(); + + if (thread_id == 32) + nvshmem_sync_with_same_gpu_idx(rdma_team); + 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, LEGACY_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 / LEGACY_NUM_MAX_NVL_PEERS)[i % LEGACY_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)[LEGACY_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)[LEGACY_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) { + nvshmemi_ibgda_put_nbi_warp(reinterpret_cast(rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank)), + reinterpret_cast(rdma_recv_num_tokens_mixed.send_buffer(i)), + (LEGACY_NUM_MAX_NVL_PEERS + num_rdma_experts + 1) * sizeof(int), + translate_dst_rdma_rank(i, nvl_rank), + 0, + lane_id, + 0); + } else { + UNROLLED_WARP_COPY(1, + lane_id, + LEGACY_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(); + + // Wait previous operations to be finished + if (thread_id < kNumRDMARanks and thread_id != rdma_rank) + nvshmemi_ibgda_quiet(translate_dst_rdma_rank(thread_id, nvl_rank), 0); + __syncthreads(); + + // Barrier + if (thread_id == 0) + nvshmem_sync_with_same_gpu_idx(rdma_team); + __syncthreads(); + + // NVL buffers + auto nvl_send_buffer = thread_id < LEGACY_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, LEGACY_NUM_MAX_NVL_PEERS); + auto nvl_send_num_tokens_per_expert = AsymBuffer(nvl_send_buffer, num_nvl_experts, LEGACY_NUM_MAX_NVL_PEERS); + auto nvl_recv_num_tokens_per_rank = AsymBuffer(nvl_recv_buffer, kNumRDMARanks, LEGACY_NUM_MAX_NVL_PEERS); + auto nvl_recv_num_tokens_per_expert = AsymBuffer(nvl_recv_buffer, num_nvl_experts, LEGACY_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)[LEGACY_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)[LEGACY_NUM_MAX_NVL_PEERS + num_rdma_experts]; + recv_rdma_rank_prefix_sum[i] = sum; + } + if (num_worst_tokens == 0) { + 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(LEGACY_NUM_MAX_NVL_PEERS <= num_threads); + if (thread_id < LEGACY_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 / LEGACY_NUM_MAX_NVL_PEERS, src_nvl_rank = i % LEGACY_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; + } + if (num_worst_tokens == 0) { + 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 < LEGACY_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; + if (num_worst_tokens == 0) { + while (ld_volatile_global(moe_recv_expert_counter_mapped + thread_id) != -1) + ; + moe_recv_expert_counter_mapped[thread_id] = sum; + } + } + + // Finally barrier + if (thread_id == 32) + nvshmem_sync_with_same_gpu_idx(rdma_team); + 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[LEGACY_NUM_MAX_NVL_PEERS] = {0}; + for (int64_t i = token_start_idx + lane_id; i < token_end_idx; i += 32) { + EP_STATIC_ASSERT(LEGACY_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 * LEGACY_NUM_MAX_NVL_PEERS); + auto is_token_in_rank_values = reinterpret_cast(&is_token_in_rank_uint64); + #pragma unroll + for (int j = 0; j < LEGACY_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 < LEGACY_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 < LEGACY_NUM_MAX_NVL_PEERS; ++i) + gbl_channel_prefix_matrix[(dst_rdma_rank * LEGACY_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(LEGACY_NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); + if (thread_id < LEGACY_NUM_MAX_NVL_PEERS) { + auto prefix_row = gbl_channel_prefix_matrix + (dst_rdma_rank * LEGACY_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_worst_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) { +#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_worst_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, \ + nvshmem::cpu_rdma_team); \ + } \ + break + + constexpr int kNumThreads = 512; + const auto num_rdma_ranks = num_ranks / LEGACY_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, + LEGACY_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 + LEGACY_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 num_worst_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) { + 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 / LEGACY_NUM_MAX_NVL_PEERS, nvl_rank = rank % LEGACY_NUM_MAX_NVL_PEERS; + + EP_DEVICE_ASSERT(ibgda_get_state()->num_rc_per_pe == num_channels or ibgda_get_state()->num_rc_per_pe >= num_sms); + + const auto role_meta = [=]() -> std::pair { + if (is_forwarder) { + if (warp_id < LEGACY_NUM_MAX_NVL_PEERS) { + return {WarpRole::kRDMAAndNVLForwarder, (warp_id + channel_id) % LEGACY_NUM_MAX_NVL_PEERS}; + } else { + return {WarpRole::kForwarderCoordinator, warp_id - LEGACY_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) % LEGACY_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 + LEGACY_NUM_MAX_NVL_PEERS); + + // Data checks + EP_DEVICE_ASSERT(num_topk <= 32); + + // RDMA symmetric layout + EP_STATIC_ASSERT(LEGACY_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, LEGACY_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, + LEGACY_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, LEGACY_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, LEGACY_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, LEGACY_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, LEGACY_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]; + 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[LEGACY_NUM_MAX_NVL_PEERS][kNumRDMARanks]; + __shared__ volatile bool forward_channel_retired[LEGACY_NUM_MAX_NVL_PEERS]; + auto sync_forwarder_smem = []() { asm volatile("barrier.sync 1, %0;" ::"r"((LEGACY_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(LEGACY_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 < LEGACY_NUM_MAX_NVL_PEERS) { + dst_ptr[lane_id] = + -(channel_id == 0 + ? 0 + : gbl_channel_prefix_matrix[(dst_rdma_rank * LEGACY_NUM_MAX_NVL_PEERS + lane_id) * num_channels + channel_id - 1]) - + 1; + } else if (lane_id < LEGACY_NUM_MAX_NVL_PEERS * 2) { + dst_ptr[lane_id] = + -gbl_channel_prefix_matrix[(dst_rdma_rank * LEGACY_NUM_MAX_NVL_PEERS + lane_id - LEGACY_NUM_MAX_NVL_PEERS) * num_channels + + channel_id] - + 1; + } else if (lane_id == LEGACY_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 == LEGACY_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) { + nvshmemi_ibgda_put_nbi_warp(reinterpret_cast(rdma_channel_meta.recv_buffer(rdma_rank)), + reinterpret_cast(rdma_channel_meta.send_buffer(dst_rdma_rank)), + sizeof(int) * (LEGACY_NUM_MAX_NVL_PEERS * 2 + 2), + translate_dst_rdma_rank(dst_rdma_rank, nvl_rank), + channel_id, + lane_id, + 0); + } + } + 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 * LEGACY_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 >= LEGACY_NUM_TIMEOUT_CYCLES) { + printf("DeepEP 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 > LEGACY_NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) { + printf("DeepEP 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); + nvshmemi_ibgda_put_nbi_warp(dst_ptr, + src_ptr, + num_bytes_per_msg, + translate_dst_rdma_rank(dst_rdma_rank, nvl_rank), + channel_id, + lane_id, + 0); + } 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; + nvshmemi_ibgda_amo_nonfetch_add(rdma_channel_tail.buffer(rdma_rank), + num_tokens_to_issue, + translate_dst_rdma_rank(dst_rdma_rank, nvl_rank), + channel_id, + dst_rdma_rank == rdma_rank); + } + __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) + LEGACY_NUM_MAX_NVL_PEERS + dst_nvl_rank); + auto meta_2 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + LEGACY_NUM_MAX_NVL_PEERS * 2); + auto meta_3 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + LEGACY_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 > LEGACY_NUM_TIMEOUT_CYCLES) { + printf( + "DeepEP 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 * LEGACY_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 > LEGACY_NUM_TIMEOUT_CYCLES) { + printf( + "DeepEP 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 > LEGACY_NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) { + printf( + "DeepEP 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 * LEGACY_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(LEGACY_NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); + #pragma unroll + for (int i = lane_id; i < kNumRDMARanks * LEGACY_NUM_MAX_NVL_PEERS; i += 32) + forward_channel_head[i % LEGACY_NUM_MAX_NVL_PEERS][i / LEGACY_NUM_MAX_NVL_PEERS] = 0; + if (lane_id < LEGACY_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 < LEGACY_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) { + nvshmemi_ibgda_amo_nonfetch_add(rdma_channel_head.buffer(rdma_rank), + min_head - last_head, + translate_dst_rdma_rank(lane_id, nvl_rank), + channel_id + num_channels, + lane_id == rdma_rank); + last_head = min_head; + } + + // Nanosleep and let other warps work + __nanosleep(LEGACY_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 * LEGACY_NUM_MAX_NVL_PEERS + src_nvl_rank > 0) + total_offset = recv_gbl_rank_prefix_sum[lane_id * LEGACY_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 > LEGACY_NUM_TIMEOUT_CYCLES) { + printf( + "DeepEP 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 * LEGACY_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 > LEGACY_NUM_TIMEOUT_CYCLES) { + printf("DeepEP 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); + } + } + + // Clean unused `recv_topk_idx` as -1 + if (num_worst_tokens > 0) { + if (is_forwarder) + return; + // get the actual number of num_recv_tokens on the current rank + int num_recv_tokens = recv_gbl_rank_prefix_sum[num_ranks - 1]; + // some ForwarderCoordinator threads exit early, so we only use non-forwarder in clean-up + // channel_id * num_threads is the offset of the current non-forwarder sms + const auto clean_start = num_recv_tokens * num_topk + channel_id * num_threads; + const auto clean_end = num_worst_tokens * num_topk; + const auto clean_stride = num_channels * num_threads; + #pragma unroll + for (int i = clean_start + thread_id; i < clean_end; i += clean_stride) + recv_topk_idx[i] = -1; + } +} + +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 num_worst_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) { + constexpr int kNumDispatchRDMASenderWarps = 7; + constexpr int kNumTMABytesPerWarp = 16384; + constexpr int smem_size = kNumTMABytesPerWarp * LEGACY_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, \ + num_worst_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); \ + } \ + 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 + LEGACY_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, + const nvshmem_team_t rdma_team) { + 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 % LEGACY_NUM_MAX_NVL_PEERS; + auto num_rdma_ranks = num_ranks / LEGACY_NUM_MAX_NVL_PEERS; + auto rdma_rank = rank / LEGACY_NUM_MAX_NVL_PEERS; + + // Using two SMs, which clean the RDMA/NVL buffer respectively + if (sm_id == 0) { + auto qps_per_rdma_rank = ibgda_get_state()->num_rc_per_pe * ibgda_get_state()->num_devices_initialized; + for (int i = thread_id; i < qps_per_rdma_rank * (num_rdma_ranks - 1); i += num_threads) { + auto dst_rdma_rank = (i / qps_per_rdma_rank + rdma_rank + 1) % num_rdma_ranks; + auto qp_id = i % qps_per_rdma_rank; + nvshmemi_ibgda_quiet(translate_dst_rdma_rank(dst_rdma_rank, nvl_rank), qp_id); + } + __syncthreads(); + + // Barrier for RDMA + if (thread_id == 32) + nvshmem_sync_with_same_gpu_idx(rdma_team); + + // 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 (thread_id == 32) + nvshmem_sync_with_same_gpu_idx(rdma_team); + 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(LEGACY_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) * LEGACY_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 * LEGACY_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 < LEGACY_NUM_MAX_NVL_PEERS) { + auto current_head = + reinterpret_cast(tma_buffer)[(token_idx - batch_start_idx) * LEGACY_NUM_MAX_NVL_PEERS + lane_id]; + if (current_head < 0) { + reinterpret_cast(tma_buffer)[(token_idx - batch_start_idx) * LEGACY_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 * LEGACY_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) { + const int num_threads = std::max(128, 32 * num_channels); + const int num_warps = num_threads / 32; + const auto num_rdma_ranks = num_ranks / LEGACY_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, + LEGACY_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, + nvshmem::cpu_rdma_team); +} + +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 * (LEGACY_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 + LEGACY_NUM_MAX_NVL_PEERS * kNumTMALoadBytes); + }; + auto tma_mbarrier = [=](const int& i) -> uint64_t* { + return reinterpret_cast(smem_ptr + i * kNumTMABufferBytesPerStage + (LEGACY_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 - LEGACY_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) { + 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 / LEGACY_NUM_MAX_NVL_PEERS, nvl_rank = rank % LEGACY_NUM_MAX_NVL_PEERS; + auto role_meta = [=]() -> std::pair { + auto warp_id = thread_id / 32; + if (not is_forwarder_sm) { + if (warp_id < LEGACY_NUM_MAX_NVL_PEERS) { + auto shuffled_warp_id = warp_id; + shuffled_warp_id = (shuffled_warp_id + channel_id) % LEGACY_NUM_MAX_NVL_PEERS; + return {WarpRole::kNVLSender, shuffled_warp_id}; + } else if (warp_id < kNumForwarders) { + return {WarpRole::kRDMAReceiver, warp_id - LEGACY_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, + LEGACY_NUM_MAX_NVL_PEERS, + channel_id, + num_channels, + nvl_rank) + .advance_also(local_buffer_ptr); + auto nvl_channel_head = AsymBuffer(local_buffer_ptr, kNumRDMARanks, LEGACY_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, LEGACY_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 * LEGACY_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 > LEGACY_NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) { + printf( + "DeepEP 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[LEGACY_NUM_MAX_NVL_PEERS]; + #pragma unroll + for (int i = 0; i < LEGACY_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, LEGACY_NUM_MAX_NVL_PEERS, channel_id, num_channels) + .advance_also(nvl_buffers); + auto nvl_channel_head = + AsymBuffer(nvl_buffers, kNumRDMARanks, LEGACY_NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank) + .advance_also(local_nvl_buffer); + auto nvl_channel_tail = AsymBuffer(local_nvl_buffer, kNumRDMARanks, LEGACY_NUM_MAX_NVL_PEERS, channel_id, num_channels) + .advance_also(nvl_buffers); + + // Combiner warp synchronization + __shared__ volatile int forwarder_nvl_head[kNumForwarders][LEGACY_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 * (LEGACY_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 * (LEGACY_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(LEGACY_NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); + lane_id < LEGACY_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 * LEGACY_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 > LEGACY_NUM_TIMEOUT_CYCLES) { + printf( + "DeepEP 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 < LEGACY_NUM_MAX_NVL_PEERS) { + expected_head = ld_nc_global(combined_nvl_head + token_idx * LEGACY_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 > LEGACY_NUM_TIMEOUT_CYCLES and lane_id < LEGACY_NUM_MAX_NVL_PEERS) { + printf( + "DeepEP 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 < LEGACY_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); + nvshmemi_ibgda_put_nbi_warp(dst_ptr, + src_ptr, + num_bytes_per_msg, + translate_dst_rdma_rank(dst_rdma_rank, nvl_rank), + channel_id, + lane_id, + 0); + } else { + memory_fence(); + } + + // Write new RDMA tail + __syncwarp(); + if (elect_one_sync()) { + nvshmemi_ibgda_amo_nonfetch_add(rdma_channel_tail.buffer(rdma_rank), + num_chunked_tokens, + translate_dst_rdma_rank(dst_rdma_rank, nvl_rank), + channel_id, + dst_rdma_rank == rdma_rank); + } + } + } + + // 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 > LEGACY_NUM_TIMEOUT_CYCLES) { + printf( + "DeepEP 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 < LEGACY_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) { + nvshmemi_ibgda_amo_nonfetch_add(rdma_channel_head.buffer(rdma_rank), + min_head - last_rdma_head, + translate_dst_rdma_rank(dst_rdma_rank, nvl_rank), + channel_id + num_channels, + dst_rdma_rank == rdma_rank); + 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 < LEGACY_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(LEGACY_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) { + constexpr int kNumCombineForwarderWarps = 24; + constexpr int kNumTMABytesPerSenderWarp = 16384; + constexpr int kNumTMABytesPerForwarderWarp = 9248; + constexpr int smem_size = + std::max(kNumTMABytesPerSenderWarp * LEGACY_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); \ + } \ + break + + int num_rdma_ranks = num_ranks / LEGACY_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 > LEGACY_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 internode + +} // namespace legacy + +} // namespace deep_ep diff --git a/csrc/kernels/legacy/internode_ll.cu b/csrc/kernels/legacy/internode_ll.cu new file mode 100644 index 0000000..38ff722 --- /dev/null +++ b/csrc/kernels/legacy/internode_ll.cu @@ -0,0 +1,1289 @@ +#include "compiled.cuh" +#include "ibgda_device.cuh" +#include "launch.cuh" + +namespace deep_ep::legacy { + +namespace internode_ll { + +template +__forceinline__ __device__ bool is_rank_masked(int* mask_buffer_ptr, int rank) { + if (mask_buffer_ptr == nullptr) { + return false; + } + if constexpr (use_warp_sync) { + return __shfl_sync(0xffffffff, ld_acquire_global(mask_buffer_ptr + rank), 0) != 0; + } else { + return ld_acquire_global(mask_buffer_ptr + rank) != 0; + } +} + +template +__forceinline__ __device__ void barrier(int thread_id, int rank, int num_ranks, int* mask_buffer_ptr, int* sync_buffer_ptr) { + EP_DEVICE_ASSERT(kNumThreads >= num_ranks); + + // Quiet all QPs + auto qps_per_rank = ibgda_get_state()->num_rc_per_pe * ibgda_get_state()->num_devices_initialized; + + for (int i = thread_id; i < qps_per_rank * (num_ranks - 1); i += kNumThreads) { + auto dst_rank = (rank + 1 + i / qps_per_rank) % num_ranks; + auto qp_id = i % qps_per_rank; + nvshmemi_ibgda_quiet(dst_rank, qp_id); + } + + // Update local counter + if (thread_id == 0) + atomicAdd(sync_buffer_ptr + rank, -1); + __syncthreads(); + + int cnt = sync_buffer_ptr[rank]; + // Update remote counter and wait for local counter to be updated + if (thread_id < num_ranks && thread_id != rank) { + const auto dst_rank = thread_id; + const auto dst_ptr = reinterpret_cast(sync_buffer_ptr + rank); + const auto dst_p2p_ptr = nvshmemi_get_p2p_ptr(dst_ptr, rank, dst_rank); + + if (not is_rank_masked(mask_buffer_ptr, dst_rank)) { + if (dst_p2p_ptr == 0) { + nvshmemi_ibgda_rma_p(reinterpret_cast(dst_ptr), cnt, dst_rank, 0); + } else { + st_release_sys_global(reinterpret_cast(dst_p2p_ptr), cnt); + } + + auto start_time = clock64(); + uint64_t wait_recv_cost = 0; + while (ld_acquire_sys_global(sync_buffer_ptr + dst_rank) != cnt // remote is not ready + && (wait_recv_cost = clock64() - start_time) <= LEGACY_NUM_TIMEOUT_CYCLES // not timeout + ) + ; + // Mask rank if timeout + if (wait_recv_cost > LEGACY_NUM_TIMEOUT_CYCLES) { + printf("Warning: DeepEP timeout for barrier, rank %d, dst_rank %d\n", rank, dst_rank); + if (mask_buffer_ptr == nullptr) + trap(); + atomicExch(mask_buffer_ptr + dst_rank, 1); + } + } + } + __syncthreads(); +} + +template +__launch_bounds__(kNumThreads, 1) __global__ void clean_low_latency_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_ptr, + int* sync_buffer_ptr) { + auto thread_id = static_cast(threadIdx.x); + + // Barrier before cleaning (in case of unfinished chunked EP) + if (sync_buffer_ptr == nullptr) + nvshmemx_barrier_all_block(); + else + barrier(thread_id, rank, num_ranks, mask_buffer_ptr, sync_buffer_ptr); + + // Clean + #pragma unroll + for (int i = thread_id; i < num_clean_int_0; i += kNumThreads) + clean_0[i] = 0; + #pragma unroll + for (int i = thread_id; i < num_clean_int_1; i += kNumThreads) + clean_1[i] = 0; + + // Barrier after cleaning (make sure the low-latency mode works fine) + if (sync_buffer_ptr == nullptr) + nvshmemx_barrier_all_block(); + else + barrier(thread_id, rank, num_ranks, mask_buffer_ptr, sync_buffer_ptr); +} + +void clean_low_latency_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_ptr, + int* sync_buffer_ptr, + cudaStream_t stream) { + constexpr int kNumThreads = 256; + + SETUP_LAUNCH_CONFIG(1, kNumThreads, stream); + + LAUNCH_KERNEL(&cfg, + clean_low_latency_buffer, + clean_0, + num_clean_int_0, + clean_1, + num_clean_int_1, + rank, + num_ranks, + mask_buffer_ptr, + sync_buffer_ptr); +} + +template +__global__ __launch_bounds__(1024, 1) void dispatch(void* packed_recv_x, + void* packed_recv_x_scales, + int* packed_recv_src_info, + int64_t* packed_recv_layout_range, + int* packed_recv_count, + int* mask_buffer_ptr, + int* cumulative_local_expert_recv_stats, + int64_t* dispatch_wait_recv_cost_stats, + void* rdma_recv_x, + int* rdma_recv_count, + void* rdma_x, + const void* x, + const topk_idx_t* topk_idx, + int* atomic_counter_per_expert, + int* atomic_finish_counter_per_expert, + int* next_clean, + int num_next_clean_int, + int num_tokens, + int num_max_dispatch_tokens_per_rank, + int num_topk, + int num_experts, + int rank, + int num_ranks, + int num_warp_groups, + int num_warps_per_group, + bool round_scale, + int phases) { + 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(); + const auto num_sms = static_cast(gridDim.x); + const auto num_warps = num_warp_groups * num_warps_per_group; + const auto num_local_experts = num_experts / num_ranks; + const auto warp_group_id = warp_id / num_warps_per_group; + const auto sub_warp_id = warp_id % num_warps_per_group; + const auto responsible_expert_idx = sm_id * num_warp_groups + warp_group_id; + + // May extract UE8M0 from the scales + using scale_t = std::conditional_t; + using packed_t = std::conditional_t; + EP_STATIC_ASSERT(sizeof(packed_t) % sizeof(scale_t) == 0, "Invalid vector length"); + + // FP8 staffs + constexpr int kNumPerChannels = 128; + const int num_scales = kHidden / kNumPerChannels; + const size_t hidden_bytes = kHidden * (kUseFP8 ? sizeof(__nv_fp8_storage_t) : sizeof(nv_bfloat16)); + const size_t hidden_int4 = hidden_bytes / sizeof(int4); + + // Message package: index at source (int), 3 reserved int fields, hidden data, FP8 scales + // NOTES: currently we have 3 reserved int fields for future use + using vec_t = std::conditional_t; + const size_t num_bytes_per_msg = sizeof(int4) + (kUseFP8 ? (kHidden + num_scales * sizeof(float)) : (kHidden * sizeof(nv_bfloat16))); + const size_t num_int4_per_msg = num_bytes_per_msg / sizeof(int4); + EP_DEVICE_ASSERT(num_bytes_per_msg % sizeof(int4) == 0); + + // Expert counts + constexpr int kNumMaxWarpGroups = 32; + __shared__ int shared_num_tokens_sent_per_expert[kNumMaxWarpGroups]; + + // Sending phase + if ((phases & LEGACY_LOW_LATENCY_SEND_PHASE) == 0) + goto LOW_LATENCY_DISPATCH_RECV; + + // There are 2 kinds of warps in this part: + // 1. The first-kind warps for FP8 cast and sending top-k tokens + // 2. The last warp for reading `topk_idx` and count for per-expert information + if (warp_id < num_warps - 1) { + constexpr int kNumElemsPerRead = sizeof(int4) / sizeof(nv_bfloat16); + EP_STATIC_ASSERT(kHidden % (32 * kNumElemsPerRead) == 0, "Invalid hidden"); + EP_STATIC_ASSERT(kNumElemsPerRead * 32 % kNumPerChannels == 0, "Invalid vectorization"); + const auto num_threads = (num_warps - 1) * 32; + const size_t hidden_bf16_int4 = kHidden / kNumElemsPerRead; + + for (int token_idx = sm_id; token_idx < num_tokens; token_idx += num_sms) { + const auto x_int4 = static_cast(x) + token_idx * hidden_bf16_int4; + const auto rdma_x_src_idx = reinterpret_cast(static_cast(rdma_x) + token_idx * num_bytes_per_msg); + const auto rdma_x_vec = reinterpret_cast(reinterpret_cast(rdma_x_src_idx) + sizeof(int4)); + const auto rdma_x_scales = reinterpret_cast(reinterpret_cast(rdma_x_vec) + hidden_bytes); + + // Overlap top-k index read and source token index writes + auto dst_expert_idx = warp_id < num_topk ? static_cast(__ldg(topk_idx + token_idx * num_topk + warp_id)) : -1; + thread_id == 0 ? (*rdma_x_src_idx = token_idx) : 0; + + // FP8 cast + EP_STATIC_ASSERT(hidden_bf16_int4 % 32 == 0, "Must use the full warp to reduce"); + #pragma unroll + for (int i = thread_id; i < hidden_bf16_int4; i += num_threads) { + // Read + auto int4_value = __ldg(x_int4 + i); + + if constexpr (kUseFP8) { + // Calculate local amax + auto bf16_values = reinterpret_cast(&int4_value); + float fp32_values[kNumElemsPerRead]; + float amax = kFP8Margin, scale, scale_inv; + #pragma unroll + for (int j = 0; j < kNumElemsPerRead; ++j) { + fp32_values[j] = static_cast(bf16_values[j]); + amax = fmaxf(amax, fabsf(fp32_values[j])); + } + + // Reduce amax and scale + EP_STATIC_ASSERT(kNumElemsPerRead * 32 / kNumPerChannels == 2, "Invalid vectorization"); + amax = warp_reduce_max<16>(amax); + calculate_fp8_scales(amax, scale, scale_inv, round_scale); + if (lane_id == 0 or lane_id == 16) + rdma_x_scales[i * kNumElemsPerRead / 128] = scale_inv; + + // Cast into send buffer + vec_t int2_value; + auto fp8x2_values = reinterpret_cast<__nv_fp8x2_storage_t*>(&int2_value); + #pragma unroll + for (int j = 0; j < kNumElemsPerRead; j += 2) { + float2 fp32x2 = {fp32_values[j] * scale, fp32_values[j + 1] * scale}; + fp8x2_values[j / 2] = __nv_cvt_float2_to_fp8x2(fp32x2, __NV_SATFINITE, __NV_E4M3); + } + rdma_x_vec[i] = int2_value; + } else { + // Reinterpret-cast is for C++14 compatibility + rdma_x_vec[i] = *reinterpret_cast(&int4_value); + } + } + asm volatile("bar.sync 1, %0;" ::"r"(num_threads)); + + // Issue IBGDA sends + if (dst_expert_idx >= 0) { + int slot_idx = lane_id == 0 ? atomicAdd(atomic_counter_per_expert + dst_expert_idx, 1) : 0; + slot_idx = __shfl_sync(0xffffffff, slot_idx, 0); + const auto dst_rank = dst_expert_idx / num_local_experts; + const auto dst_expert_local_idx = dst_expert_idx % num_local_experts; + const auto src_ptr = reinterpret_cast(rdma_x_src_idx); + const auto dst_ptr = reinterpret_cast(rdma_recv_x) + + dst_expert_local_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_msg + + rank * num_max_dispatch_tokens_per_rank * num_bytes_per_msg + slot_idx * num_bytes_per_msg; + const auto dst_p2p_ptr = nvshmemi_get_p2p_ptr(dst_ptr, rank, dst_rank); + if (not is_rank_masked(mask_buffer_ptr, dst_rank)) { + if (dst_p2p_ptr == 0) { + nvshmemi_ibgda_put_nbi_warp(dst_ptr, src_ptr, num_bytes_per_msg, dst_rank, dst_expert_local_idx, lane_id, slot_idx); + } else { + // NOTES: only 2 load iterations for 7K hidden with 8 unrolls + const auto* src_int4_ptr = reinterpret_cast(src_ptr); + const auto* dst_int4_ptr = reinterpret_cast(dst_p2p_ptr); + UNROLLED_WARP_COPY(8, lane_id, num_int4_per_msg, dst_int4_ptr, src_int4_ptr, ld_nc_global, st_na_global); + } + } + + // Increase counter after finishing + __syncwarp(); + lane_id == 0 ? atomic_add_release_global(atomic_finish_counter_per_expert + dst_expert_idx, 1) : 0; + } + } + } else if (warp_id == num_warps - 1) { + EP_DEVICE_ASSERT(num_sms > 1); + if (sm_id == 0) { + // The first SM is also responsible for checking QPs + EP_DEVICE_ASSERT(ibgda_get_state()->num_rc_per_pe >= num_local_experts); + + // The first SM is also responsible for cleaning the next buffer + #pragma unroll + for (int i = lane_id; i < num_next_clean_int; i += 32) + next_clean[i] = 0; + + // Notify before executing `int_p` + __syncwarp(); + #pragma unroll + for (int i = lane_id; i < num_experts; i += 32) + atomic_add_release_global(atomic_finish_counter_per_expert + i, LEGACY_FINISHED_SUM_TAG); + } + + // This SM should be responsible for some destination experts, read `topk_idx` for them + int expert_count[kNumMaxWarpGroups] = {0}; + const auto expert_begin_idx = sm_id * num_warp_groups; + const auto expert_end_idx = min(expert_begin_idx + num_warp_groups, num_experts); + + // Per lane count + #pragma unroll 8 + for (int i = lane_id; i < num_tokens * num_topk; i += 32) { + auto idx = static_cast(__ldg(topk_idx + i)); + if (idx >= expert_begin_idx and idx < expert_end_idx) + expert_count[idx - expert_begin_idx]++; + } + + // Warp reduce + #pragma unroll + for (int i = expert_begin_idx; i < expert_end_idx; ++i) { + auto sum = warp_reduce_sum(expert_count[i - expert_begin_idx]); + if (lane_id == 0) { + shared_num_tokens_sent_per_expert[i - expert_begin_idx] = sum; + atomic_add_release_global(atomic_finish_counter_per_expert + i, LEGACY_FINISHED_SUM_TAG - sum); + } + } + } + __syncthreads(); + + // Issue count sends + if (responsible_expert_idx < num_experts and sub_warp_id == 0 and lane_id == 0) { + const auto dst_rank = responsible_expert_idx / num_local_experts; + const auto dst_expert_local_idx = responsible_expert_idx % num_local_experts; + const auto num_tokens_sent = shared_num_tokens_sent_per_expert[responsible_expert_idx - sm_id * num_warp_groups]; + + // Wait local sends issued and send expert counts + while (ld_acquire_global(atomic_finish_counter_per_expert + responsible_expert_idx) != LEGACY_FINISHED_SUM_TAG * 2) + ; + auto dst_ptr = reinterpret_cast(rdma_recv_count + dst_expert_local_idx * num_ranks + rank); + auto dst_p2p_ptr = nvshmemi_get_p2p_ptr(dst_ptr, rank, dst_rank); + if (not is_rank_masked(mask_buffer_ptr, dst_rank)) { + if (dst_p2p_ptr == 0) { + nvshmemi_ibgda_amo_nonfetch_add(reinterpret_cast(dst_ptr), -num_tokens_sent - 1, dst_rank, dst_expert_local_idx); + } else { + st_release_sys_global(reinterpret_cast(dst_p2p_ptr), -num_tokens_sent - 1); + } + } + + // Clean workspace for next use + atomic_counter_per_expert[responsible_expert_idx] = 0; + atomic_finish_counter_per_expert[responsible_expert_idx] = 0; + + // Clean `packed_recv_count` + if (dst_rank == 0) + packed_recv_count[dst_expert_local_idx] = 0; + } + __syncwarp(); + +// Receiving phase +LOW_LATENCY_DISPATCH_RECV: + if ((phases & LEGACY_LOW_LATENCY_RECV_PHASE) == 0) + return; + + // For send-and-recv kernels, we need a grid sync for making `packed_recv_count` visible + if (phases & LEGACY_LOW_LATENCY_SEND_PHASE) + cg::this_grid().sync(); + + // Receiving and packing + if (responsible_expert_idx < num_experts) { + const auto src_rank = responsible_expert_idx / num_local_experts; + const auto local_expert_idx = responsible_expert_idx % num_local_experts; + const auto rdma_recv_x_uint8 = static_cast(rdma_recv_x) + + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_msg + + src_rank * num_max_dispatch_tokens_per_rank * num_bytes_per_msg; + const auto recv_x_int4 = + static_cast(packed_recv_x) + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * hidden_int4; + const auto recv_src_info = packed_recv_src_info + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank; + const auto recv_range = packed_recv_layout_range + local_expert_idx * num_ranks; + const auto num_aligned_scales = align_up(num_scales, sizeof(float) / sizeof(scale_t)); + const auto recv_x_scales = static_cast(packed_recv_x_scales) + + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_aligned_scales; + + // Shared between sub-warps in warp groups + __shared__ int shared_num_recv_tokens[kNumMaxWarpGroups], shared_recv_token_begin_idx[kNumMaxWarpGroups]; + + // Wait tokens to arrive + // NOTES: using sub-warp 1 to overlap with sub-warp 0 + int num_recv_tokens = 0, recv_token_begin_idx; + EP_DEVICE_ASSERT(num_warps_per_group > 1 and num_warp_groups < 15); + if (sub_warp_id == 1 and lane_id == 0) { + auto start_time = clock64(); + uint64_t wait_recv_cost = 0; + if (not is_rank_masked(mask_buffer_ptr, src_rank)) { + while ((num_recv_tokens = ld_acquire_sys_global(rdma_recv_count + local_expert_idx * num_ranks + src_rank)) == + 0 // data not arrived + && (wait_recv_cost = clock64() - start_time) <= LEGACY_NUM_TIMEOUT_CYCLES // not timeout + ) + ; + } + // Do not receive tokens if rank timeout or masked + if (num_recv_tokens == 0) + num_recv_tokens = -1; + // Mask rank if timeout + if (wait_recv_cost > LEGACY_NUM_TIMEOUT_CYCLES) { + printf("Warning: DeepEP timeout for dispatch receive, rank %d, local_expert_idx %d, src_rank %d\n", + rank, + local_expert_idx, + src_rank); + if (mask_buffer_ptr == nullptr) + trap(); + atomicExch(mask_buffer_ptr + src_rank, 1); + } + + num_recv_tokens = -num_recv_tokens - 1; + recv_token_begin_idx = atomicAdd(packed_recv_count + local_expert_idx, num_recv_tokens); + shared_num_recv_tokens[warp_group_id] = num_recv_tokens; + shared_recv_token_begin_idx[warp_group_id] = recv_token_begin_idx; + recv_range[src_rank] = pack2(num_recv_tokens, recv_token_begin_idx); + + // Add stats for diagnosis + if (cumulative_local_expert_recv_stats != nullptr) + atomicAdd(cumulative_local_expert_recv_stats + local_expert_idx, num_recv_tokens); + if (dispatch_wait_recv_cost_stats != nullptr) + atomicAdd(reinterpret_cast(dispatch_wait_recv_cost_stats + src_rank), wait_recv_cost); + } + asm volatile("bar.sync %0, %1;" ::"r"(warp_group_id + 2), "r"(num_warps_per_group * 32)); + num_recv_tokens = shared_num_recv_tokens[warp_group_id]; + recv_token_begin_idx = shared_recv_token_begin_idx[warp_group_id]; + + // Copy tokens + EP_DEVICE_ASSERT(num_scales <= 64); + for (int i = sub_warp_id; i < num_recv_tokens; i += num_warps_per_group) { + // Copy source info + const auto src_src_idx = reinterpret_cast(rdma_recv_x_uint8 + i * num_bytes_per_msg); + if (lane_id == 0) + recv_src_info[recv_token_begin_idx + i] = ld_nc_global(src_src_idx); + __syncwarp(); + + // Copy data + // NOTES: only 2 load iterations for 7K hidden with 7 unrolls + const auto src_data = reinterpret_cast(reinterpret_cast(src_src_idx) + sizeof(int4)); + const auto dst_data = recv_x_int4 + (recv_token_begin_idx + i) * hidden_int4; + UNROLLED_WARP_COPY(7, lane_id, hidden_int4, dst_data, src_data, ld_nc_global, st_na_global); + + // Copy scales + if constexpr (kUseFP8) { + // Equivalent CuTe layout: + // (num_tokens, (num_packed, num_elems_per_pack)):(num_elems_per_pack, (num_tokens * num_elems_per_pack, 1)) + const auto src_scales = reinterpret_cast(reinterpret_cast(src_data) + hidden_bytes); + const auto num_elems_per_pack = static_cast(sizeof(packed_t) / sizeof(scale_t)); + const auto token_idx = recv_token_begin_idx + i; + const auto token_stride = num_elems_per_pack; + const auto pack_stride = num_ranks * num_max_dispatch_tokens_per_rank * num_elems_per_pack; + if (lane_id < num_scales) { + const auto pack_idx = lane_id / num_elems_per_pack; + const auto elem_idx = lane_id % num_elems_per_pack; + auto scale = extract_required_scale_format(ld_nc_global(src_scales + lane_id)); + recv_x_scales[token_idx * token_stride + pack_idx * pack_stride + elem_idx] = scale; + } + if (lane_id + 32 < num_scales) { + const auto pack_idx = (lane_id + 32) / num_elems_per_pack; + const auto elem_idx = (lane_id + 32) % num_elems_per_pack; + auto scale = extract_required_scale_format(ld_nc_global(src_scales + lane_id + 32)); + recv_x_scales[token_idx * token_stride + pack_idx * pack_stride + elem_idx] = scale; + } + } + } + } +} + +void dispatch(void* packed_recv_x, + void* packed_recv_x_scales, + int* packed_recv_src_info, + int64_t* packed_recv_layout_range, + int* packed_recv_count, + int* mask_buffer_ptr, + int* cumulative_local_expert_recv_stats, + int64_t* dispatch_wait_recv_cost_stats, + void* rdma_recv_x, + int* rdma_recv_count, + void* rdma_x, + const void* x, + const topk_idx_t* topk_idx, + int* next_clean, + int num_next_clean_int, + int num_tokens, + int hidden, + int num_max_dispatch_tokens_per_rank, + int num_topk, + int num_experts, + int rank, + int num_ranks, + bool use_fp8, + bool round_scale, + bool use_ue8m0, + void* workspace, + int num_device_sms, + cudaStream_t stream, + int phases) { + 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; + EP_HOST_ASSERT(num_warp_groups > 0 and num_warps_per_group > 0); + EP_HOST_ASSERT(kNumMaxTopK + 1 <= num_warp_groups * num_warps_per_group); + + const auto num_warps = num_warp_groups * num_warps_per_group; + const auto num_sms = ceil_div(num_experts, num_warp_groups); + EP_HOST_ASSERT(num_topk <= kNumMaxTopK); + + // Workspace checks + auto atomic_counter_per_expert = static_cast(workspace); + auto atomic_finish_counter_per_expert = atomic_counter_per_expert + num_experts; + EP_HOST_ASSERT(num_experts * sizeof(int) * 2 <= LEGACY_NUM_WORKSPACE_BYTES); + + // FP8 checks + if (use_ue8m0) + EP_HOST_ASSERT(round_scale and "UE8M0 SF requires `round_scale=True`"); + +#define DISPATCH_LAUNCH_CASE(hidden) \ + { \ + auto dispatch_func = dispatch