diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 000000000..7b75c6869 --- /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/.gitignore b/.gitignore index fd2a10383..9f2fd3195 100644 --- a/.gitignore +++ b/.gitignore @@ -1,8 +1,28 @@ +# VS Code compile_commands.json -.idea -.DS_Store -*.pyc -build/ +.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 000000000..4c15fffab --- /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 000000000..02e1ed349 --- /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/README.md b/README.md index 2db47e4ec..507e2ba2d 100644 --- a/README.md +++ b/README.md @@ -1,360 +1,440 @@ # DeepEP -DeepEP is a communication library tailored for Mixture-of-Experts (MoE) and expert parallelism (EP). It provides high-throughput and low-latency all-to-all GPU kernels, which are also known as MoE dispatch and combine. The library also supports low-precision operations, including FP8. +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. -To align with the group-limited gating algorithm proposed in the [DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3) paper, DeepEP offers a set of kernels optimized for asymmetric-domain bandwidth forwarding, such as forwarding data from NVLink domain to RDMA domain. These kernels deliver high throughput, making them suitable for both training and inference prefilling tasks. Additionally, they support SM (Streaming Multiprocessors) number control. +Despite its lightweight design, DeepEP's performance matches or exceeds hardware bandwidth limits across various configurations. -For latency-sensitive inference decoding, DeepEP includes a set of low-latency kernels with pure RDMA to minimize delays. The library also introduces a hook-based communication-computation overlapping method that does not occupy any SM resource. +## News -Notice: the implementation in this library may have some slight differences from the [DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3) paper. +- **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**. -## Performance +### 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 -### Normal kernels with NVLink and RDMA forwarding +- **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 -We test normal kernels on H800 (~160 GB/s NVLink maximum bandwidth), with each connected to a CX7 InfiniBand 400 Gb/s RDMA network card (~50 GB/s maximum bandwidth). And we follow the DeepSeek-V3/R1 pretraining setting (4096 tokens per batch, 7168 hidden, top-4 groups, top-8 experts, FP8 dispatching and BF16 combining). +For the legacy V1 documentation (NVSHMEM-based), see [docs/legacy.md](docs/legacy.md). -| Type | Dispatch #EP | Bottleneck bandwidth | Combine #EP | Bottleneck bandwidth | -|:---------:|:------------:|:--------------------:|:-----------:|:--------------------:| -| Intranode | 8 | 153 GB/s (NVLink) | 8 | 158 GB/s (NVLink) | -| Internode | 16 | 43 GB/s (RDMA) | 16 | 43 GB/s (RDMA) | -| Internode | 32 | 58 GB/s (RDMA) | 32 | 57 GB/s (RDMA) | -| Internode | 64 | 51 GB/s (RDMA) | 64 | 50 GB/s (RDMA) | +## 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: -**News (2025.04.22)**: with optimizations from Tencent Network Platform Department, performance was enhanced by up to 30%, see [#130](https://github.com/deepseek-ai/DeepEP/pull/130) for more details. Thanks for the contribution! +| 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) | -### Low-latency kernels with pure RDMA +Notes: the results are logical bandwidth. For example, under the `EP 8 x 2` case, 90 GB/s actually contains local rank traffic. -We test low-latency kernels on H800 with each connected to a CX7 InfiniBand 400 Gb/s RDMA network card (~50 GB/s maximum bandwidth). And we follow a typical DeepSeek-V3/R1 production setting (128 tokens per batch, 7168 hidden, top-8 experts, FP8 dispatching and BF16 combining). +Comparing with V1, **V2 achieves up to 1.3x peak performance, while saving up to 4x SM count**. -| Dispatch #EP | Latency | RDMA bandwidth | Combine #EP | Latency | RDMA bandwidth | -|:------------:|:-------:|:--------------:|:-----------:|:-------:|:--------------:| -| 8 | 77 us | 98 GB/s | 8 | 114 us | 127 GB/s | -| 16 | 118 us | 63 GB/s | 16 | 195 us | 74 GB/s | -| 32 | 155 us | 48 GB/s | 32 | 273 us | 53 GB/s | -| 64 | 173 us | 43 GB/s | 64 | 314 us | 46 GB/s | -| 128 | 192 us | 39 GB/s | 128 | 369 us | 39 GB/s | -| 256 | 194 us | 39 GB/s | 256 | 360 us | 40 GB/s | +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. -**News (2025.06.05)**: low-latency kernels now leverage NVLink as much as possible, see [#173](https://github.com/deepseek-ai/DeepEP/pull/173) for more details. Thanks for the contribution! +For V1 performance data, see [docs/legacy.md](docs/legacy.md#performance). ## Quick start ### Requirements -- Ampere (SM80), Hopper (SM90) GPUs, or other architectures with SM90 PTX ISA support +- Hopper (SM90) GPUs, or other architectures with SM90 PTX ISA support - Python 3.8 and above - CUDA version - - CUDA 11.0 and above for SM80 GPUs - CUDA 12.3 and above for SM90 GPUs -- PyTorch 2.1 and above +- PyTorch 2.10 and above +- NCCL 2.30.4 and above - NVLink for intranode communication - RDMA network for internode communication -### Download and install NVSHMEM dependency +### 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. Please refer to our [NVSHMEM Installation Guide](third-party/README.md) for instructions. +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 -NVSHMEM_DIR=/path/to/installed/nvshmem python setup.py build +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.py` +# 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/test_intranode.py -python tests/test_internode.py -python tests/test_low_latency.py +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 -NVSHMEM_DIR=/path/to/installed/nvshmem python setup.py install +python setup.py install ``` -#### Installation environment variables - -- `NVSHMEM_DIR`: the path to the NVSHMEM directory, disable all internode and low-latency features if not specified -- `DISABLE_SM90_FEATURES`: 0 or 1, whether to disable SM90 features, it is required for SM90 devices or CUDA 11 -- `TORCH_CUDA_ARCH_LIST`: the list of target architectures, e.g. `TORCH_CUDA_ARCH_LIST="9.0"` -- `DISABLE_AGGRESSIVE_PTX_INSTRS`: 0 or 1, whether to disable aggressive load/store instructions, see [Undefined-behavior PTX usage](#undefined-behavior-ptx-usage) for more details - Then, import `deep_ep` in your Python project, and enjoy! -## 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: - -- workloads using normal kernels -- workloads using low-latency kernels -- other workloads - -For DeepEP, you can control the virtual lane assignment by setting the `NVSHMEM_IB_SL` environment variable. - -### Adaptive routing - -Adaptive routing is an advanced routing feature provided by InfiniBand switches that can evenly distribute traffic across multiple paths. Enabling adaptive routing can completely eliminate network congestion caused by routing conflicts, but it also introduces additional latency. We recommend the following configuration for optimal performance: - -- enable adaptive routing in environments with heavy network loads -- use static routing in environments with light network loads - -### Congestion control - -Congestion control is disabled as we have not observed significant congestion in our production environment. - ## Interfaces and examples -### Example use in model training or inference prefilling +### Buffer initialization -The normal kernels can be used in model training or the inference prefilling phase (without the backward part) as the below example code shows. +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 List, Tuple, Optional, Union +from typing import Optional -from deep_ep import Buffer, EventOverlap +from deep_ep import ElasticBuffer # Communication buffer (will allocate at runtime) -_buffer: Optional[Buffer] = None +_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) -# Set the number of SMs to use -# NOTES: this is a static variable -Buffer.set_num_sms(24) - - -# You may call this function at the framework initialization -def get_buffer(group: dist.ProcessGroup, hidden_bytes: int) -> Buffer: - global _buffer + return _buffer +``` - # NOTES: you may also replace `get_*_config` with your auto-tuned results via all the tests - num_nvl_bytes, num_rdma_bytes = 0, 0 - for config in (Buffer.get_dispatch_config(group.size()), Buffer.get_combine_config(group.size())): - num_nvl_bytes = max(config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes) - num_rdma_bytes = max(config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes) +### Example use in model training or inference prefilling - # Allocate a buffer if not existed or not enough buffer size - if _buffer is None or _buffer.group != group or _buffer.num_nvl_bytes < num_nvl_bytes or _buffer.num_rdma_bytes < num_rdma_bytes: - _buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes) - return _buffer +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 -def get_hidden_bytes(x: torch.Tensor) -> int: - t = x[0] if isinstance(x, tuple) else x - return t.size(1) * max(t.element_size(), 2) +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, previous_event: Optional[EventOverlap] = None) -> \ - Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], torch.Tensor, torch.Tensor, List, Tuple, EventOverlap]: - # NOTES: an optional `previous_event` means a CUDA event captured that you want to make it as a dependency - # of the dispatch kernel, it may be useful with communication-computation overlap. For more information, please - # refer to the docs of `Buffer.dispatch` - global _buffer - - # Calculate layout before actual dispatch - num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, previous_event = \ - _buffer.get_dispatch_layout(topk_idx, num_experts, - previous_event=previous_event, async_finish=True, - allocate_on_comm_stream=previous_event is not None) - # Do MoE dispatch - # NOTES: the CPU will wait for GPU's signal to arrive, so this is not compatible with CUDA graph - # Unless you specify `num_worst_tokens`, but this flag is for intranode only - # For more advanced usages, please refer to the docs of the `dispatch` function - recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event = \ - _buffer.dispatch(x, topk_idx=topk_idx, topk_weights=topk_weights, - num_tokens_per_rank=num_tokens_per_rank, num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, - is_token_in_rank=is_token_in_rank, num_tokens_per_expert=num_tokens_per_expert, - previous_event=previous_event, async_finish=True, - allocate_on_comm_stream=True) - # For event management, please refer to the docs of the `EventOverlap` class - return recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event - - -def dispatch_backward(grad_recv_x: torch.Tensor, grad_recv_topk_weights: torch.Tensor, handle: Tuple) -> \ - Tuple[torch.Tensor, torch.Tensor, EventOverlap]: - global _buffer - - # The backward process of MoE dispatch is actually a combine - # For more advanced usages, please refer to the docs of the `combine` function - combined_grad_x, combined_grad_recv_topk_weights, event = \ - _buffer.combine(grad_recv_x, handle, topk_weights=grad_recv_topk_weights, async_finish=True) - - # For event management, please refer to the docs of the `EventOverlap` class - return combined_grad_x, combined_grad_recv_topk_weights, event - - -def combine_forward(x: torch.Tensor, handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \ - Tuple[torch.Tensor, EventOverlap]: - global _buffer - - # Do MoE combine - # For more advanced usages, please refer to the docs of the `combine` function - combined_x, _, event = _buffer.combine(x, handle, async_finish=True, previous_event=previous_event, - allocate_on_comm_stream=previous_event is not None) - - # For event management, please refer to the docs of the `EventOverlap` class + 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: Tuple, previous_event: Optional[EventOverlap] = None) -> \ + handle: EPHandle) -> \ Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], EventOverlap]: - global _buffer + """The backward pass of MoE combine is actually a dispatch.""" + global _buffer, _num_comm_sms - # The backward process of MoE combine is actually a dispatch - # For more advanced usages, please refer to the docs of the `dispatch` function - grad_x, _, _, _, _, event = _buffer.dispatch(grad_combined_x, handle=handle, async_finish=True, - previous_event=previous_event, - allocate_on_comm_stream=previous_event is not None) + grad_x, _, _, _, event = _buffer.dispatch( + grad_combined_x, + handle=handle, + num_sms=_num_comm_sms, + async_with_compute_stream=True, + ) - # For event management, please refer to the docs of the `EventOverlap` class return grad_x, event ``` -Moreover, inside the dispatch function, we may not know how many tokens to receive for the current rank. So an implicit CPU wait for GPU received count signal will be involved, as the following figure shows. - -![normal](figures/normal.png) - -### Example use in inference decoding - -The low latency kernels can be used in the inference decoding phase as the below example code shows. +For communication-computation overlap, use the `EventOverlap` interface to manage dependencies between the communication stream and the compute stream: ```python -import torch -import torch.distributed as dist -from typing import Tuple, Optional - -from deep_ep import Buffer +# After dispatch, overlap computation while communication is in-flight +recv_x, recv_topk_idx, recv_topk_weights, handle, event = dispatch_forward(...) -# Communication buffer (will allocate at runtime) -# NOTES: there is no SM control API for the low-latency kernels -_buffer: Optional[Buffer] = None - - -# You may call this function at the framework initialization -def get_buffer(group: dist.ProcessGroup, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int) -> Buffer: - # NOTES: the low-latency mode will consume much more space than the normal mode - # So we recommend that `num_max_dispatch_tokens_per_rank` (the actual batch size in the decoding engine) should be less than 256 - global _buffer - num_rdma_bytes = Buffer.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, group.size(), num_experts) - - # Allocate a buffer if not existed or not enough buffer size - if _buffer is None or _buffer.group != group or not _buffer.low_latency_mode or _buffer.num_rdma_bytes < num_rdma_bytes: - # NOTES: for the best performance, the QP number **must** be equal to the number of the local experts - assert num_experts % group.size() == 0 - _buffer = Buffer(group, 0, num_rdma_bytes, low_latency_mode=True, num_qps_per_rank=num_experts // group.size()) - return _buffer +# ... do some independent computation here ... +# Wait for communication to finish before using results +event.current_stream_wait() -def low_latency_dispatch(hidden_states: torch.Tensor, topk_idx: torch.Tensor, num_max_dispatch_tokens_per_rank: int, num_experts: int): - global _buffer +# Now safe to use recv_x, recv_topk_idx, recv_topk_weights +``` - # Do MoE dispatch, compatible with CUDA graph (but you may restore some buffer status once you replay) - recv_hidden_states, recv_expert_count, handle, event, hook = \ - _buffer.low_latency_dispatch(hidden_states, topk_idx, num_max_dispatch_tokens_per_rank, num_experts, - async_finish=False, return_recv_hook=True) +### Example use in inference decoding - # NOTES: the actual tensor will not be received only if you call `hook()`, - # it is useful for double-batch overlapping, but **without any SM occupation** - # If you don't want to overlap, please set `return_recv_hook=False` - # Later, you can use our GEMM library to do the computation with this specific format - return recv_hidden_states, recv_expert_count, handle, event, hook +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, + ) -def low_latency_combine(hidden_states: torch.Tensor, - topk_idx: torch.Tensor, topk_weights: torch.Tensor, handle: Tuple): - global _buffer + return combined_x, event +``` - # Do MoE combine, compatible with CUDA graph (but you may restore some buffer status once you replay) - combined_hidden_states, event_overlap, hook = \ - _buffer.low_latency_combine(hidden_states, topk_idx, topk_weights, handle, - async_finish=False, return_recv_hook=True) +### 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. - # NOTES: the same behavior as described in the dispatch kernel - return combined_hidden_states, event_overlap, hook -``` +## Network configurations -For two-micro-batch overlapping, you can refer to the following figure. With our receiving hook interface, the RDMA network traffic is happening in the background, without costing any GPU SMs from the computation part. But notice, the overlapped parts can be adjusted, i.e., the 4 parts of attention/dispatch/MoE/combine may not have the exact same execution time. You may adjust the stage settings according to your workload. +DeepEP is fully tested with InfiniBand networks. However, it is theoretically compatible with RDMA over Converged Ethernet (RoCE) as well. -![low-latency](figures/low-latency.png) +### Traffic isolation -## Roadmap +Traffic isolation is supported by InfiniBand through Virtual Lanes (VL). -- [x] AR support -- [x] Refactor low-latency mode AR code -- [x] A100 support (intranode only) -- [x] Support BF16 for the low-latency dispatch kernel -- [x] Support NVLink protocol for intranode low-latency kernels -- [ ] TMA copy instead of LD/ST - - [x] Intranode kernels - - [ ] Internode kernels - - [ ] Low-latency kernels -- [ ] SM-free kernels and refactors -- [ ] Fully remove undefined-behavior PTX instructions +To prevent interference between different types of traffic, we recommend segregating workloads across different virtual lanes as follows: -## Notices +- expert-parallel workloads +- other workloads -#### Easier potential overall design +For DeepEP V2, you can control the virtual lane assignment by setting the `sl_idx` argument or the `EP_OVERRIDE_RDMA_SL` environment variable. -The current DeepEP implementation uses queues for communication buffers which save memory but introduce complexity and potential deadlocks. If you're implementing your own version based on DeepEP, consider using fixed-size buffers allocated to maximum capacity for simplicity and better performance. For a detailed discussion of this alternative approach, see https://github.com/deepseek-ai/DeepEP/issues/39. +### Adaptive routing -#### Undefined-behavior PTX usage +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. -- For extreme performance, we discover and use an undefined-behavior PTX usage: using read-only PTX `ld.global.nc.L1::no_allocate.L2::256B` to **read volatile data**. The PTX modifier `.nc` indicates that a non-coherent cache is used. But the correctness is tested to be guaranteed with `.L1::no_allocate` on Hopper architectures, and performance will be much better. The reason we guess may be: the non-coherent cache is unified with L1, and the L1 modifier is not just a hint but a strong option, so that the correctness can be guaranteed by no dirty data in L1. -- Initially, because NVCC could not automatically unroll volatile read PTX, we tried using `__ldg` (i.e., `ld.nc`). Even compared to manually unrolled volatile reads, it was significantly faster (likely due to additional compiler optimizations). However, the results could be incorrect or dirty. After consulting the PTX documentation, we discovered that L1 and non-coherent cache are unified on Hopper architectures. We speculated that `.L1::no_allocate` might resolve the issue, leading to this discovery. -- If you find kernels not working on some other platforms, you may add `DISABLE_AGGRESSIVE_PTX_INSTRS=1` to `setup.py` and disable this, or file an issue. +### Congestion control -#### Auto-tuning on your cluster +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. -For better performance on your cluster, we recommend to run all the tests and use the best auto-tuned configuration. The default configurations are optimized on the DeepSeek's internal cluster. +### PCI atomic mode -## License +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: -This code repository is released under [the MIT License](LICENSE), except for codes that reference NVSHMEM (including `csrc/kernels/ibgda_device.cuh` and `third-party/nvshmem.patch`), which are subject to [NVSHMEM SLA](https://docs.nvidia.com/nvshmem/api/sla.html). +```bash +sudo mlxconfig -y -d mlx5_$i set PCI_ATOMIC_MODE=4 +``` -## Experimental Branches +## 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** + - 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 + - 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 + - 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 + - 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) + - ROCm/AMD GPU support powered by [MORI](https://github.com/ROCm/mori) backend (low-latency mode) -## Community Forks +## 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) -## Citation +## 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! -If you use this codebase or otherwise find our work valuable, please cite: +## License + +This code repository is released under [the MIT License](LICENSE). + +## Citation ```bibtex @misc{deepep2025, diff --git a/build.sh b/build.sh new file mode 100755 index 000000000..abdfc4067 --- /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 index 3f51c2713..94e1eba1a 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -1,36 +1 @@ -# 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") - -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) -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) - -set(CMAKE_CXX_STANDARD 17) -set(CMAKE_CUDA_STANDARD 17) - -include_directories(${CUDA_TOOLKIT_ROOT_DIR}/include ${TORCH_INCLUDE_DIRS} ${PYTHON_INCLUDE_DIRS} ${NVSHMEM_INCLUDE_DIR}) -link_directories(${TORCH_INSTALL_PREFIX}/lib ${CUDA_TOOLKIT_ROOT_DIR}/lib ${NVSHMEM_LIB_DIR}) - add_subdirectory(kernels) - -# Link CPP and CUDA together -pybind11_add_module(deep_ep_cpp deep_ep.cpp) -target_link_libraries(deep_ep_cpp PRIVATE ${EP_CUDA_LIBRARIES} ${TORCH_LIBRARIES} torch_python) diff --git a/csrc/deep_ep.cpp b/csrc/deep_ep.cpp deleted file mode 100644 index 954c9ffb6..000000000 --- a/csrc/deep_ep.cpp +++ /dev/null @@ -1,1893 +0,0 @@ -#include "deep_ep.hpp" - -#include -#include -#include -#include -#include - -#include -#include - -#include "kernels/api.cuh" -#include "kernels/configs.cuh" - -namespace shared_memory { -void cu_mem_set_access_all(void* ptr, size_t size) { - int device_count; - CUDA_CHECK(cudaGetDeviceCount(&device_count)); - - CUmemAccessDesc access_desc[device_count]; - for (int idx = 0; idx < device_count; ++idx) { - access_desc[idx].location.type = CU_MEM_LOCATION_TYPE_DEVICE; - access_desc[idx].location.id = idx; - access_desc[idx].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; - } - - CU_CHECK(cuMemSetAccess((CUdeviceptr)ptr, size, access_desc, device_count)); -} - -void cu_mem_free(void* ptr) { - CUmemGenericAllocationHandle handle; - CU_CHECK(cuMemRetainAllocationHandle(&handle, ptr)); - - size_t size = 0; - CU_CHECK(cuMemGetAddressRange(NULL, &size, (CUdeviceptr)ptr)); - - CU_CHECK(cuMemUnmap((CUdeviceptr)ptr, size)); - CU_CHECK(cuMemAddressFree((CUdeviceptr)ptr, size)); - CU_CHECK(cuMemRelease(handle)); -} - -size_t get_size_align_to_granularity(size_t size_raw, size_t granularity) { - size_t size = (size_raw + granularity - 1) & ~(granularity - 1); - if (size == 0) - size = granularity; - return size; -} - -SharedMemoryAllocator::SharedMemoryAllocator(bool use_fabric) : use_fabric(use_fabric) {} - -void SharedMemoryAllocator::malloc(void** ptr, size_t size_raw) { - if (use_fabric) { - CUdevice device; - CU_CHECK(cuCtxGetDevice(&device)); - - CUmemAllocationProp prop = {}; - prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; - prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; - prop.location.id = device; - - size_t granularity = 0; - CU_CHECK(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM)); - - size_t size = get_size_align_to_granularity(size_raw, granularity); - - CUmemGenericAllocationHandle handle; - CU_CHECK(cuMemCreate(&handle, size, &prop, 0)); - - CU_CHECK(cuMemAddressReserve((CUdeviceptr*)ptr, size, granularity, 0, 0)); - CU_CHECK(cuMemMap((CUdeviceptr)*ptr, size, 0, handle, 0)); - cu_mem_set_access_all(*ptr, size); - } else { - CUDA_CHECK(cudaMalloc(ptr, size_raw)); - } -} - -void SharedMemoryAllocator::free(void* ptr) { - if (use_fabric) { - cu_mem_free(ptr); - } else { - CUDA_CHECK(cudaFree(ptr)); - } -} - -void SharedMemoryAllocator::get_mem_handle(MemHandle* mem_handle, void* ptr) { - size_t size = 0; - CU_CHECK(cuMemGetAddressRange(NULL, &size, (CUdeviceptr)ptr)); - - mem_handle->size = size; - - if (use_fabric) { - CUmemGenericAllocationHandle handle; - CU_CHECK(cuMemRetainAllocationHandle(&handle, ptr)); - - CU_CHECK(cuMemExportToShareableHandle(&mem_handle->inner.cu_mem_fabric_handle, handle, CU_MEM_HANDLE_TYPE_FABRIC, 0)); - } else { - CUDA_CHECK(cudaIpcGetMemHandle(&mem_handle->inner.cuda_ipc_mem_handle, ptr)); - } -} - -void SharedMemoryAllocator::open_mem_handle(void** ptr, MemHandle* mem_handle) { - if (use_fabric) { - size_t size = mem_handle->size; - - CUmemGenericAllocationHandle handle; - CU_CHECK(cuMemImportFromShareableHandle(&handle, &mem_handle->inner.cu_mem_fabric_handle, CU_MEM_HANDLE_TYPE_FABRIC)); - - CU_CHECK(cuMemAddressReserve((CUdeviceptr*)ptr, size, 0, 0, 0)); - CU_CHECK(cuMemMap((CUdeviceptr)*ptr, size, 0, handle, 0)); - cu_mem_set_access_all(*ptr, size); - } else { - CUDA_CHECK(cudaIpcOpenMemHandle(ptr, mem_handle->inner.cuda_ipc_mem_handle, cudaIpcMemLazyEnablePeerAccess)); - } -} - -void SharedMemoryAllocator::close_mem_handle(void* ptr) { - if (use_fabric) { - cu_mem_free(ptr); - } else { - CUDA_CHECK(cudaIpcCloseMemHandle(ptr)); - } -} -} // namespace shared_memory - -namespace deep_ep { - -Buffer::Buffer(int rank, - int num_ranks, - int64_t num_nvl_bytes, - int64_t num_rdma_bytes, - bool low_latency_mode, - bool explicitly_destroy, - bool enable_shrink, - bool use_fabric) - : rank(rank), - num_ranks(num_ranks), - num_nvl_bytes(num_nvl_bytes), - num_rdma_bytes(num_rdma_bytes), - enable_shrink(enable_shrink), - low_latency_mode(low_latency_mode), - explicitly_destroy(explicitly_destroy), - comm_stream(at::cuda::getStreamFromPool(true)), - shared_memory_allocator(use_fabric) { - // Metadata memory - int64_t barrier_signal_bytes = NUM_MAX_NVL_PEERS * sizeof(int); - int64_t buffer_ptr_bytes = NUM_MAX_NVL_PEERS * sizeof(void*); - int64_t barrier_signal_ptr_bytes = NUM_MAX_NVL_PEERS * sizeof(int*); - - // Common checks - EP_STATIC_ASSERT(NUM_BUFFER_ALIGNMENT_BYTES % sizeof(int4) == 0, "Invalid alignment"); - EP_HOST_ASSERT(num_nvl_bytes % NUM_BUFFER_ALIGNMENT_BYTES == 0 and - (num_nvl_bytes <= std::numeric_limits::max() or num_rdma_bytes == 0)); - EP_HOST_ASSERT(num_rdma_bytes % NUM_BUFFER_ALIGNMENT_BYTES == 0 and - (low_latency_mode or num_rdma_bytes <= std::numeric_limits::max())); - EP_HOST_ASSERT(num_nvl_bytes / sizeof(int4) < std::numeric_limits::max()); - EP_HOST_ASSERT(num_rdma_bytes / sizeof(int4) < std::numeric_limits::max()); - EP_HOST_ASSERT(0 <= rank and rank < num_ranks and (num_ranks <= NUM_MAX_NVL_PEERS * NUM_MAX_RDMA_PEERS or low_latency_mode)); - EP_HOST_ASSERT(num_ranks < NUM_MAX_NVL_PEERS or num_ranks % NUM_MAX_NVL_PEERS == 0); - if (num_rdma_bytes > 0) - EP_HOST_ASSERT(num_ranks > NUM_MAX_NVL_PEERS or low_latency_mode); - - // Get ranks - CUDA_CHECK(cudaGetDevice(&device_id)); - rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS; - num_rdma_ranks = std::max(1, num_ranks / NUM_MAX_NVL_PEERS), num_nvl_ranks = std::min(num_ranks, NUM_MAX_NVL_PEERS); -#ifdef DISABLE_NVSHMEM - EP_HOST_ASSERT(num_rdma_ranks == 1 and not low_latency_mode and "NVSHMEM is disabled during compilation"); -#endif - - // Get device info - cudaDeviceProp device_prop = {}; - CUDA_CHECK(cudaGetDeviceProperties(&device_prop, device_id)); - num_device_sms = device_prop.multiProcessorCount; - - // Number of per-channel bytes cannot be large - EP_HOST_ASSERT(ceil_div(num_nvl_bytes, num_device_sms / 2) < std::numeric_limits::max()); - EP_HOST_ASSERT(ceil_div(num_rdma_bytes, num_device_sms / 2) < std::numeric_limits::max()); - - if (num_nvl_bytes > 0) { - // Local IPC: alloc local memory and set local IPC handles - shared_memory_allocator.malloc(&buffer_ptrs[nvl_rank], - num_nvl_bytes + barrier_signal_bytes + buffer_ptr_bytes + barrier_signal_ptr_bytes); - shared_memory_allocator.get_mem_handle(&ipc_handles[nvl_rank], buffer_ptrs[nvl_rank]); - buffer_ptrs_gpu = reinterpret_cast(static_cast(buffer_ptrs[nvl_rank]) + num_nvl_bytes + barrier_signal_bytes); - - // Set barrier signals - barrier_signal_ptrs[nvl_rank] = reinterpret_cast(static_cast(buffer_ptrs[nvl_rank]) + num_nvl_bytes); - barrier_signal_ptrs_gpu = - reinterpret_cast(static_cast(buffer_ptrs[nvl_rank]) + num_nvl_bytes + barrier_signal_bytes + buffer_ptr_bytes); - - // No need to synchronize, will do a full device sync during `sync` - CUDA_CHECK(cudaMemsetAsync(barrier_signal_ptrs[nvl_rank], 0, barrier_signal_bytes, comm_stream)); - } - - // Create 32 MiB workspace - CUDA_CHECK(cudaMalloc(&workspace, NUM_WORKSPACE_BYTES)); - CUDA_CHECK(cudaMemsetAsync(workspace, 0, NUM_WORKSPACE_BYTES, comm_stream)); - - // MoE counter - CUDA_CHECK(cudaMallocHost(&moe_recv_counter, sizeof(int64_t), cudaHostAllocMapped)); - CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_counter_mapped, const_cast(moe_recv_counter), 0)); - *moe_recv_counter = -1; - - // MoE expert-level counter - CUDA_CHECK(cudaMallocHost(&moe_recv_expert_counter, sizeof(int) * NUM_MAX_LOCAL_EXPERTS, cudaHostAllocMapped)); - CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_expert_counter_mapped, const_cast(moe_recv_expert_counter), 0)); - for (int i = 0; i < NUM_MAX_LOCAL_EXPERTS; ++i) - moe_recv_expert_counter[i] = -1; - - // MoE RDMA-level counter - if (num_rdma_ranks > 0) { - CUDA_CHECK(cudaMallocHost(&moe_recv_rdma_counter, sizeof(int), cudaHostAllocMapped)); - CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_rdma_counter_mapped, const_cast(moe_recv_rdma_counter), 0)); - *moe_recv_rdma_counter = -1; - } -} - -Buffer::~Buffer() noexcept(false) { - if (not explicitly_destroy) { - destroy(); - } else if (not destroyed) { - printf("WARNING: destroy() was not called before DeepEP buffer destruction, which can leak resources.\n"); - fflush(stdout); - } -} - -bool Buffer::is_available() const { - return available; -} - -bool Buffer::is_internode_available() const { - return is_available() and num_ranks > NUM_MAX_NVL_PEERS; -} - -int Buffer::get_num_rdma_ranks() const { - return num_rdma_ranks; -} - -int Buffer::get_rdma_rank() const { - return rdma_rank; -} - -int Buffer::get_root_rdma_rank(bool global) const { - return global ? nvl_rank : 0; -} - -int Buffer::get_local_device_id() const { - return device_id; -} - -pybind11::bytearray Buffer::get_local_ipc_handle() const { - const shared_memory::MemHandle& handle = ipc_handles[nvl_rank]; - return {reinterpret_cast(&handle), sizeof(handle)}; -} - -pybind11::bytearray Buffer::get_local_nvshmem_unique_id() const { -#ifndef DISABLE_NVSHMEM - EP_HOST_ASSERT(rdma_rank == 0 and "Only RDMA rank 0 can get NVSHMEM unique ID"); - auto unique_id = internode::get_unique_id(); - return {reinterpret_cast(unique_id.data()), unique_id.size()}; -#else - EP_HOST_ASSERT(false and "NVSHMEM is disabled during compilation"); -#endif -} - -torch::Tensor Buffer::get_local_buffer_tensor(const pybind11::object& dtype, int64_t offset, bool use_rdma_buffer) const { - torch::ScalarType casted_dtype = torch::python::detail::py_object_to_dtype(dtype); - auto element_bytes = static_cast(elementSize(casted_dtype)); - auto base_ptr = static_cast(use_rdma_buffer ? rdma_buffer_ptr : buffer_ptrs[nvl_rank]) + offset; - auto num_bytes = use_rdma_buffer ? num_rdma_bytes : num_nvl_bytes; - return torch::from_blob(base_ptr, num_bytes / element_bytes, torch::TensorOptions().dtype(casted_dtype).device(at::kCUDA)); -} - -torch::Stream Buffer::get_comm_stream() const { - return comm_stream; -} - -void Buffer::destroy() { - EP_HOST_ASSERT(not destroyed); - - // Synchronize - CUDA_CHECK(cudaDeviceSynchronize()); - - if (num_nvl_bytes > 0) { - // Barrier - intranode::barrier(barrier_signal_ptrs_gpu, nvl_rank, num_nvl_ranks, comm_stream); - CUDA_CHECK(cudaDeviceSynchronize()); - - // Close remote IPC - if (is_available()) { - for (int i = 0; i < num_nvl_ranks; ++i) - if (i != nvl_rank) - shared_memory_allocator.close_mem_handle(buffer_ptrs[i]); - } - - // Free local buffer and error flag - shared_memory_allocator.free(buffer_ptrs[nvl_rank]); - } - - // Free NVSHMEM -#ifndef DISABLE_NVSHMEM - if (is_available() and num_rdma_bytes > 0) { - CUDA_CHECK(cudaDeviceSynchronize()); - internode::barrier(); - internode::free(rdma_buffer_ptr); - if (enable_shrink) { - internode::free(mask_buffer_ptr); - internode::free(sync_buffer_ptr); - } - internode::finalize(); - } -#endif - - // Free workspace and MoE counter - CUDA_CHECK(cudaFree(workspace)); - CUDA_CHECK(cudaFreeHost(const_cast(moe_recv_counter))); - - // Free chunked mode staffs - CUDA_CHECK(cudaFreeHost(const_cast(moe_recv_expert_counter))); - - destroyed = true; - available = false; -} - -void Buffer::sync(const std::vector& device_ids, - const std::vector>& all_gathered_handles, - const std::optional& root_unique_id_opt) { - EP_HOST_ASSERT(not is_available()); - - // Sync IPC handles - if (num_nvl_bytes > 0) { - EP_HOST_ASSERT(num_ranks == device_ids.size()); - EP_HOST_ASSERT(device_ids.size() == all_gathered_handles.size()); - for (int i = 0, offset = rdma_rank * num_nvl_ranks; i < num_nvl_ranks; ++i) { - EP_HOST_ASSERT(all_gathered_handles[offset + i].has_value()); - auto handle_str = std::string(all_gathered_handles[offset + i].value()); - EP_HOST_ASSERT(handle_str.size() == shared_memory::HANDLE_SIZE); - if (offset + i != rank) { - std::memcpy(&ipc_handles[i], handle_str.c_str(), shared_memory::HANDLE_SIZE); - shared_memory_allocator.open_mem_handle(&buffer_ptrs[i], &ipc_handles[i]); - barrier_signal_ptrs[i] = reinterpret_cast(static_cast(buffer_ptrs[i]) + num_nvl_bytes); - } else { - EP_HOST_ASSERT(std::memcmp(&ipc_handles[i], handle_str.c_str(), shared_memory::HANDLE_SIZE) == 0); - } - } - - // Copy all buffer and barrier signal pointers to GPU - CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs, sizeof(void*) * NUM_MAX_NVL_PEERS, cudaMemcpyHostToDevice)); - CUDA_CHECK(cudaMemcpy(barrier_signal_ptrs_gpu, barrier_signal_ptrs, sizeof(int*) * NUM_MAX_NVL_PEERS, cudaMemcpyHostToDevice)); - CUDA_CHECK(cudaDeviceSynchronize()); - } - - // Sync NVSHMEM handles and allocate memory -#ifndef DISABLE_NVSHMEM - if (num_rdma_bytes > 0) { - // Initialize NVSHMEM - EP_HOST_ASSERT(root_unique_id_opt.has_value()); - std::vector root_unique_id(root_unique_id_opt->size()); - auto root_unique_id_str = root_unique_id_opt->cast(); - std::memcpy(root_unique_id.data(), root_unique_id_str.c_str(), root_unique_id_opt->size()); - auto nvshmem_rank = low_latency_mode ? rank : rdma_rank; - auto num_nvshmem_ranks = low_latency_mode ? num_ranks : num_rdma_ranks; - EP_HOST_ASSERT(nvshmem_rank == internode::init(root_unique_id, nvshmem_rank, num_nvshmem_ranks, low_latency_mode)); - internode::barrier(); - - // Allocate - rdma_buffer_ptr = internode::alloc(num_rdma_bytes, NUM_BUFFER_ALIGNMENT_BYTES); - - // Clean buffer (mainly for low-latency mode) - CUDA_CHECK(cudaMemset(rdma_buffer_ptr, 0, num_rdma_bytes)); - - // Allocate and clean shrink buffer - if (enable_shrink) { - int num_mask_buffer_bytes = num_ranks * sizeof(int); - int num_sync_buffer_bytes = num_ranks * sizeof(int); - mask_buffer_ptr = reinterpret_cast(internode::alloc(num_mask_buffer_bytes, NUM_BUFFER_ALIGNMENT_BYTES)); - sync_buffer_ptr = reinterpret_cast(internode::alloc(num_sync_buffer_bytes, NUM_BUFFER_ALIGNMENT_BYTES)); - CUDA_CHECK(cudaMemset(mask_buffer_ptr, 0, num_mask_buffer_bytes)); - CUDA_CHECK(cudaMemset(sync_buffer_ptr, 0, num_sync_buffer_bytes)); - } - - // Barrier - internode::barrier(); - CUDA_CHECK(cudaDeviceSynchronize()); - } -#endif - - // Ready to use - available = true; -} - -std::tuple, torch::Tensor, torch::Tensor, std::optional> -Buffer::get_dispatch_layout( - const torch::Tensor& topk_idx, int num_experts, std::optional& previous_event, bool async, bool allocate_on_comm_stream) { - EP_HOST_ASSERT(topk_idx.dim() == 2); - EP_HOST_ASSERT(topk_idx.is_contiguous()); - EP_HOST_ASSERT(num_experts > 0); - - // Allocate all tensors on comm stream if set - // NOTES: do not allocate tensors upfront! - auto compute_stream = at::cuda::getCurrentCUDAStream(); - if (allocate_on_comm_stream) { - EP_HOST_ASSERT(previous_event.has_value() and async); - at::cuda::setCurrentCUDAStream(comm_stream); - } - - // Wait previous tasks to be finished - if (previous_event.has_value()) { - stream_wait(comm_stream, previous_event.value()); - } else { - stream_wait(comm_stream, compute_stream); - } - - auto num_tokens = static_cast(topk_idx.size(0)), num_topk = static_cast(topk_idx.size(1)); - auto num_tokens_per_rank = torch::empty({num_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); - auto num_tokens_per_rdma_rank = std::optional(); - auto num_tokens_per_expert = torch::empty({num_experts}, dtype(torch::kInt32).device(torch::kCUDA)); - auto is_token_in_rank = torch::empty({num_tokens, num_ranks}, dtype(torch::kBool).device(torch::kCUDA)); - if (is_internode_available()) - num_tokens_per_rdma_rank = torch::empty({num_rdma_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); - - layout::get_dispatch_layout(topk_idx.data_ptr(), - num_tokens_per_rank.data_ptr(), - num_tokens_per_rdma_rank.has_value() ? num_tokens_per_rdma_rank.value().data_ptr() : nullptr, - num_tokens_per_expert.data_ptr(), - is_token_in_rank.data_ptr(), - num_tokens, - num_topk, - num_ranks, - num_experts, - comm_stream); - - // Wait streams - std::optional event; - if (async) { - event = EventHandle(comm_stream); - for (auto& t : {topk_idx, num_tokens_per_rank, num_tokens_per_expert, is_token_in_rank}) { - t.record_stream(comm_stream); - if (allocate_on_comm_stream) - t.record_stream(compute_stream); - } - for (auto& to : {num_tokens_per_rdma_rank}) { - to.has_value() ? to->record_stream(comm_stream) : void(); - if (allocate_on_comm_stream) - to.has_value() ? to->record_stream(compute_stream) : void(); - } - } else { - stream_wait(compute_stream, comm_stream); - } - - // Switch back compute stream - if (allocate_on_comm_stream) - at::cuda::setCurrentCUDAStream(compute_stream); - - return {num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, event}; -} - -std::tuple, - std::optional, - std::optional, - std::vector, - torch::Tensor, - torch::Tensor, - torch::Tensor, - torch::Tensor, - torch::Tensor, - std::optional> -Buffer::intranode_dispatch(const torch::Tensor& x, - const std::optional& x_scales, - const std::optional& topk_idx, - const std::optional& topk_weights, - const std::optional& num_tokens_per_rank, - const torch::Tensor& is_token_in_rank, - const std::optional& num_tokens_per_expert, - int cached_num_recv_tokens, - const std::optional& cached_rank_prefix_matrix, - const std::optional& cached_channel_prefix_matrix, - int expert_alignment, - int num_worst_tokens, - const Config& config, - std::optional& previous_event, - bool async, - bool allocate_on_comm_stream) { - bool cached_mode = cached_rank_prefix_matrix.has_value(); - - // One channel use two blocks, even-numbered blocks for sending, odd-numbered blocks for receiving. - EP_HOST_ASSERT(config.num_sms % 2 == 0); - int num_channels = config.num_sms / 2; - if (cached_mode) { - EP_HOST_ASSERT(cached_rank_prefix_matrix.has_value()); - EP_HOST_ASSERT(cached_channel_prefix_matrix.has_value()); - } else { - EP_HOST_ASSERT(num_tokens_per_rank.has_value()); - EP_HOST_ASSERT(num_tokens_per_expert.has_value()); - } - - // Type checks - EP_HOST_ASSERT(is_token_in_rank.scalar_type() == torch::kBool); - if (cached_mode) { - EP_HOST_ASSERT(cached_rank_prefix_matrix->scalar_type() == torch::kInt32); - EP_HOST_ASSERT(cached_channel_prefix_matrix->scalar_type() == torch::kInt32); - } else { - EP_HOST_ASSERT(num_tokens_per_expert->scalar_type() == torch::kInt32); - EP_HOST_ASSERT(num_tokens_per_rank->scalar_type() == torch::kInt32); - } - - // Shape and contiguous checks - EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous()); - EP_HOST_ASSERT((x.size(1) * x.element_size()) % sizeof(int4) == 0); - EP_HOST_ASSERT(is_token_in_rank.dim() == 2 and is_token_in_rank.is_contiguous()); - EP_HOST_ASSERT(is_token_in_rank.size(0) == x.size(0) and is_token_in_rank.size(1) == num_ranks); - if (cached_mode) { - EP_HOST_ASSERT(cached_rank_prefix_matrix->dim() == 2 and cached_rank_prefix_matrix->is_contiguous()); - EP_HOST_ASSERT(cached_rank_prefix_matrix->size(0) == num_ranks and cached_rank_prefix_matrix->size(1) == num_ranks); - EP_HOST_ASSERT(cached_channel_prefix_matrix->dim() == 2 and cached_channel_prefix_matrix->is_contiguous()); - EP_HOST_ASSERT(cached_channel_prefix_matrix->size(0) == num_ranks and cached_channel_prefix_matrix->size(1) == num_channels); - } else { - EP_HOST_ASSERT(num_tokens_per_expert->dim() == 1 and num_tokens_per_expert->is_contiguous()); - EP_HOST_ASSERT(num_tokens_per_expert->size(0) % num_ranks == 0); - EP_HOST_ASSERT(num_tokens_per_expert->size(0) / num_ranks <= NUM_MAX_LOCAL_EXPERTS); - EP_HOST_ASSERT(num_tokens_per_rank->dim() == 1 and num_tokens_per_rank->is_contiguous()); - EP_HOST_ASSERT(num_tokens_per_rank->size(0) == num_ranks); - } - - auto num_tokens = static_cast(x.size(0)), hidden = static_cast(x.size(1)); - auto num_experts = cached_mode ? 0 : static_cast(num_tokens_per_expert->size(0)), num_local_experts = num_experts / num_ranks; - - // Top-k checks - int num_topk = 0; - topk_idx_t* topk_idx_ptr = nullptr; - float* topk_weights_ptr = nullptr; - EP_HOST_ASSERT(topk_idx.has_value() == topk_weights.has_value()); - if (topk_idx.has_value()) { - num_topk = static_cast(topk_idx->size(1)); - EP_HOST_ASSERT(num_experts > 0); - EP_HOST_ASSERT(topk_idx->dim() == 2 and topk_idx->is_contiguous()); - EP_HOST_ASSERT(topk_weights->dim() == 2 and topk_weights->is_contiguous()); - EP_HOST_ASSERT(num_tokens == topk_idx->size(0) and num_tokens == topk_weights->size(0)); - EP_HOST_ASSERT(num_topk == topk_weights->size(1)); - EP_HOST_ASSERT(topk_weights->scalar_type() == torch::kFloat32); - topk_idx_ptr = topk_idx->data_ptr(); - topk_weights_ptr = topk_weights->data_ptr(); - } - - // FP8 scales checks - float* x_scales_ptr = nullptr; - int num_scales = 0, scale_token_stride = 0, scale_hidden_stride = 0; - if (x_scales.has_value()) { - EP_HOST_ASSERT(x.element_size() == 1); - EP_HOST_ASSERT(x_scales->scalar_type() == torch::kFloat32 or x_scales->scalar_type() == torch::kInt); - EP_HOST_ASSERT(x_scales->dim() == 2); - EP_HOST_ASSERT(x_scales->size(0) == num_tokens); - num_scales = x_scales->dim() == 1 ? 1 : static_cast(x_scales->size(1)); - x_scales_ptr = static_cast(x_scales->data_ptr()); - scale_token_stride = static_cast(x_scales->stride(0)); - scale_hidden_stride = static_cast(x_scales->stride(1)); - } - - // Allocate all tensors on comm stream if set - // NOTES: do not allocate tensors upfront! - auto compute_stream = at::cuda::getCurrentCUDAStream(); - if (allocate_on_comm_stream) { - EP_HOST_ASSERT(previous_event.has_value() and async); - at::cuda::setCurrentCUDAStream(comm_stream); - } - - // Wait previous tasks to be finished - if (previous_event.has_value()) { - stream_wait(comm_stream, previous_event.value()); - } else { - stream_wait(comm_stream, compute_stream); - } - - // Create handles (only return for non-cached mode) - int num_recv_tokens = -1; - auto rank_prefix_matrix = torch::Tensor(); - auto channel_prefix_matrix = torch::Tensor(); - std::vector num_recv_tokens_per_expert_list; - - // Barrier or send sizes - // To clean: channel start/end offset, head and tail - int num_memset_int = num_channels * num_ranks * 4; - if (cached_mode) { - num_recv_tokens = cached_num_recv_tokens; - rank_prefix_matrix = cached_rank_prefix_matrix.value(); - channel_prefix_matrix = cached_channel_prefix_matrix.value(); - - // Copy rank prefix matrix and clean flags - intranode::cached_notify_dispatch( - rank_prefix_matrix.data_ptr(), num_memset_int, buffer_ptrs_gpu, barrier_signal_ptrs_gpu, rank, num_ranks, comm_stream); - } else { - rank_prefix_matrix = torch::empty({num_ranks, num_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); - channel_prefix_matrix = torch::empty({num_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); - - // Send sizes - // Meta information: - // - Size prefix by ranks, shaped as `[num_ranks, num_ranks]` - // - Size prefix by experts (not used later), shaped as `[num_ranks, num_local_experts]` - // NOTES: no more token dropping in this version - *moe_recv_counter = -1; - for (int i = 0; i < num_local_experts; ++i) - moe_recv_expert_counter[i] = -1; - EP_HOST_ASSERT(num_ranks * (num_ranks + num_local_experts) * sizeof(int) <= num_nvl_bytes); - intranode::notify_dispatch(num_tokens_per_rank->data_ptr(), - moe_recv_counter_mapped, - num_ranks, - num_tokens_per_expert->data_ptr(), - moe_recv_expert_counter_mapped, - num_experts, - num_tokens, - is_token_in_rank.data_ptr(), - channel_prefix_matrix.data_ptr(), - rank_prefix_matrix.data_ptr(), - num_memset_int, - expert_alignment, - buffer_ptrs_gpu, - barrier_signal_ptrs_gpu, - rank, - comm_stream, - num_channels); - - if (num_worst_tokens > 0) { - // No CPU sync, just allocate the worst case - num_recv_tokens = num_worst_tokens; - - // Must be forward with top-k stuffs - EP_HOST_ASSERT(topk_idx.has_value()); - EP_HOST_ASSERT(topk_weights.has_value()); - } else { - // Synchronize total received tokens and tokens per expert - auto start_time = std::chrono::high_resolution_clock::now(); - while (true) { - // Read total count - num_recv_tokens = static_cast(*moe_recv_counter); - - // Read per-expert count - bool ready = (num_recv_tokens >= 0); - for (int i = 0; i < num_local_experts and ready; ++i) - ready &= moe_recv_expert_counter[i] >= 0; - - if (ready) - break; - - // Timeout check - if (std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - start_time).count() > - NUM_CPU_TIMEOUT_SECS) - throw std::runtime_error("DeepEP error: CPU recv timeout"); - } - num_recv_tokens_per_expert_list = std::vector(moe_recv_expert_counter, moe_recv_expert_counter + num_local_experts); - } - } - - // Allocate new tensors - auto recv_x = torch::empty({num_recv_tokens, hidden}, x.options()); - auto recv_src_idx = torch::empty({num_recv_tokens}, dtype(torch::kInt32).device(torch::kCUDA)); - auto recv_topk_idx = std::optional(), recv_topk_weights = std::optional(), - recv_x_scales = std::optional(); - auto recv_channel_prefix_matrix = torch::empty({num_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); - auto send_head = torch::empty({num_tokens, num_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); - - // Assign pointers - topk_idx_t* recv_topk_idx_ptr = nullptr; - float* recv_topk_weights_ptr = nullptr; - float* recv_x_scales_ptr = nullptr; - if (topk_idx.has_value()) { - recv_topk_idx = torch::empty({num_recv_tokens, num_topk}, topk_idx->options()); - recv_topk_weights = torch::empty({num_recv_tokens, num_topk}, topk_weights->options()); - recv_topk_idx_ptr = recv_topk_idx->data_ptr(); - recv_topk_weights_ptr = recv_topk_weights->data_ptr(); - } - if (x_scales.has_value()) { - recv_x_scales = x_scales->dim() == 1 ? torch::empty({num_recv_tokens}, x_scales->options()) - : torch::empty({num_recv_tokens, num_scales}, x_scales->options()); - recv_x_scales_ptr = static_cast(recv_x_scales->data_ptr()); - } - - // Dispatch - EP_HOST_ASSERT( - num_ranks * num_ranks * sizeof(int) + // Size prefix matrix - num_channels * num_ranks * sizeof(int) + // Channel start offset - num_channels * num_ranks * sizeof(int) + // Channel end offset - num_channels * num_ranks * sizeof(int) * 2 + // Queue head and tail - num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * hidden * recv_x.element_size() + // Data buffer - num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * sizeof(int) + // Source index buffer - num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * num_topk * sizeof(topk_idx_t) + // Top-k index buffer - num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * num_topk * sizeof(float) + // Top-k weight buffer - num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * sizeof(float) * num_scales // FP8 scale buffer - <= num_nvl_bytes); - intranode::dispatch(recv_x.data_ptr(), - recv_x_scales_ptr, - recv_src_idx.data_ptr(), - recv_topk_idx_ptr, - recv_topk_weights_ptr, - recv_channel_prefix_matrix.data_ptr(), - send_head.data_ptr(), - x.data_ptr(), - x_scales_ptr, - topk_idx_ptr, - topk_weights_ptr, - is_token_in_rank.data_ptr(), - channel_prefix_matrix.data_ptr(), - num_tokens, - num_worst_tokens, - static_cast(hidden * recv_x.element_size() / sizeof(int4)), - num_topk, - num_experts, - num_scales, - scale_token_stride, - scale_hidden_stride, - buffer_ptrs_gpu, - rank, - num_ranks, - comm_stream, - config.num_sms, - config.num_max_nvl_chunked_send_tokens, - config.num_max_nvl_chunked_recv_tokens); - - // Wait streams - std::optional event; - if (async) { - event = EventHandle(comm_stream); - for (auto& t : {x, - is_token_in_rank, - rank_prefix_matrix, - channel_prefix_matrix, - recv_x, - recv_src_idx, - recv_channel_prefix_matrix, - send_head}) { - t.record_stream(comm_stream); - if (allocate_on_comm_stream) - t.record_stream(compute_stream); - } - for (auto& to : {x_scales, - topk_idx, - topk_weights, - num_tokens_per_rank, - num_tokens_per_expert, - cached_channel_prefix_matrix, - cached_rank_prefix_matrix, - recv_topk_idx, - recv_topk_weights, - recv_x_scales}) { - to.has_value() ? to->record_stream(comm_stream) : void(); - if (allocate_on_comm_stream) - to.has_value() ? to->record_stream(compute_stream) : void(); - } - } else { - stream_wait(compute_stream, comm_stream); - } - - // Switch back compute stream - if (allocate_on_comm_stream) - at::cuda::setCurrentCUDAStream(compute_stream); - - // Return values - return {recv_x, - recv_x_scales, - recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, - rank_prefix_matrix, - channel_prefix_matrix, - recv_channel_prefix_matrix, - recv_src_idx, - send_head, - event}; -} - -std::tuple, std::optional> Buffer::intranode_combine( - const torch::Tensor& x, - const std::optional& topk_weights, - const std::optional& bias_0, - const std::optional& bias_1, - const torch::Tensor& src_idx, - const torch::Tensor& rank_prefix_matrix, - const torch::Tensor& channel_prefix_matrix, - const torch::Tensor& send_head, - const Config& config, - std::optional& previous_event, - bool async, - bool allocate_on_comm_stream) { - EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous()); - EP_HOST_ASSERT(src_idx.dim() == 1 and src_idx.is_contiguous() and src_idx.scalar_type() == torch::kInt32); - EP_HOST_ASSERT(send_head.dim() == 2 and send_head.is_contiguous() and send_head.scalar_type() == torch::kInt32); - EP_HOST_ASSERT(rank_prefix_matrix.dim() == 2 and rank_prefix_matrix.is_contiguous() and - rank_prefix_matrix.scalar_type() == torch::kInt32); - EP_HOST_ASSERT(channel_prefix_matrix.dim() == 2 and channel_prefix_matrix.is_contiguous() and - channel_prefix_matrix.scalar_type() == torch::kInt32); - - // One channel use two blocks, even-numbered blocks for sending, odd-numbered blocks for receiving. - EP_HOST_ASSERT(config.num_sms % 2 == 0); - int num_channels = config.num_sms / 2; - - auto num_tokens = static_cast(x.size(0)), hidden = static_cast(x.size(1)); - auto num_recv_tokens = static_cast(send_head.size(0)); - EP_HOST_ASSERT(src_idx.size(0) == num_tokens); - EP_HOST_ASSERT(send_head.size(1) == num_ranks); - EP_HOST_ASSERT(rank_prefix_matrix.size(0) == num_ranks and rank_prefix_matrix.size(1) == num_ranks); - EP_HOST_ASSERT(channel_prefix_matrix.size(0) == num_ranks and channel_prefix_matrix.size(1) == num_channels); - EP_HOST_ASSERT((hidden * x.element_size()) % sizeof(int4) == 0); - - // Allocate all tensors on comm stream if set - // NOTES: do not allocate tensors upfront! - auto compute_stream = at::cuda::getCurrentCUDAStream(); - if (allocate_on_comm_stream) { - EP_HOST_ASSERT(previous_event.has_value() and async); - at::cuda::setCurrentCUDAStream(comm_stream); - } - - // Wait previous tasks to be finished - if (previous_event.has_value()) { - stream_wait(comm_stream, previous_event.value()); - } else { - stream_wait(comm_stream, compute_stream); - } - - int num_topk = 0; - auto recv_topk_weights = std::optional(); - float* topk_weights_ptr = nullptr; - float* recv_topk_weights_ptr = nullptr; - if (topk_weights.has_value()) { - EP_HOST_ASSERT(topk_weights->dim() == 2 and topk_weights->is_contiguous()); - EP_HOST_ASSERT(topk_weights->size(0) == num_tokens); - EP_HOST_ASSERT(topk_weights->scalar_type() == torch::kFloat32); - num_topk = static_cast(topk_weights->size(1)); - topk_weights_ptr = topk_weights->data_ptr(); - recv_topk_weights = torch::empty({num_recv_tokens, num_topk}, topk_weights->options()); - recv_topk_weights_ptr = recv_topk_weights->data_ptr(); - } - - // Launch barrier and reset queue head and tail - EP_HOST_ASSERT(num_channels * num_ranks * sizeof(int) * 2 <= num_nvl_bytes); - intranode::cached_notify_combine(buffer_ptrs_gpu, - send_head.data_ptr(), - num_channels, - num_recv_tokens, - num_channels * num_ranks * 2, - barrier_signal_ptrs_gpu, - rank, - num_ranks, - comm_stream); - - // Assign bias pointers - auto bias_opts = std::vector>({bias_0, bias_1}); - void* bias_ptrs[2] = {nullptr, nullptr}; - for (int i = 0; i < 2; ++i) - if (bias_opts[i].has_value()) { - auto bias = bias_opts[i].value(); - EP_HOST_ASSERT(bias.dim() == 2 and bias.is_contiguous()); - EP_HOST_ASSERT(bias.scalar_type() == x.scalar_type()); - EP_HOST_ASSERT(bias.size(0) == num_recv_tokens and bias.size(1) == hidden); - bias_ptrs[i] = bias.data_ptr(); - } - - // Combine data - auto recv_x = torch::empty({num_recv_tokens, hidden}, x.options()); - EP_HOST_ASSERT(num_channels * num_ranks * sizeof(int) * 2 + // Queue head and tail - num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * hidden * x.element_size() + // Data buffer - num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * sizeof(int) + // Source index buffer - num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * num_topk * sizeof(float) // Top-k weight buffer - <= num_nvl_bytes); - intranode::combine(at::cuda::ScalarTypeToCudaDataType(x.scalar_type()), - recv_x.data_ptr(), - recv_topk_weights_ptr, - x.data_ptr(), - topk_weights_ptr, - bias_ptrs[0], - bias_ptrs[1], - src_idx.data_ptr(), - rank_prefix_matrix.data_ptr(), - channel_prefix_matrix.data_ptr(), - send_head.data_ptr(), - num_tokens, - num_recv_tokens, - hidden, - num_topk, - buffer_ptrs_gpu, - rank, - num_ranks, - comm_stream, - config.num_sms, - config.num_max_nvl_chunked_send_tokens, - config.num_max_nvl_chunked_recv_tokens); - - // Wait streams - std::optional event; - if (async) { - event = EventHandle(comm_stream); - for (auto& t : {x, src_idx, send_head, rank_prefix_matrix, channel_prefix_matrix, recv_x}) { - t.record_stream(comm_stream); - if (allocate_on_comm_stream) - t.record_stream(compute_stream); - } - for (auto& to : {topk_weights, recv_topk_weights, bias_0, bias_1}) { - to.has_value() ? to->record_stream(comm_stream) : void(); - if (allocate_on_comm_stream) - to.has_value() ? to->record_stream(compute_stream) : void(); - } - } else { - stream_wait(compute_stream, comm_stream); - } - - // Switch back compute stream - if (allocate_on_comm_stream) - at::cuda::setCurrentCUDAStream(compute_stream); - - return {recv_x, recv_topk_weights, event}; -} - -std::tuple, - std::optional, - std::optional, - std::vector, - torch::Tensor, - torch::Tensor, - std::optional, - torch::Tensor, - std::optional, - torch::Tensor, - std::optional, - std::optional, - std::optional, - std::optional> -Buffer::internode_dispatch(const torch::Tensor& x, - const std::optional& x_scales, - const std::optional& topk_idx, - const std::optional& topk_weights, - const std::optional& num_tokens_per_rank, - const std::optional& num_tokens_per_rdma_rank, - const torch::Tensor& is_token_in_rank, - const std::optional& num_tokens_per_expert, - int cached_num_recv_tokens, - int cached_num_rdma_recv_tokens, - const std::optional& cached_rdma_channel_prefix_matrix, - const std::optional& cached_recv_rdma_rank_prefix_sum, - const std::optional& cached_gbl_channel_prefix_matrix, - const std::optional& cached_recv_gbl_rank_prefix_sum, - int expert_alignment, - int num_worst_tokens, - const Config& config, - std::optional& previous_event, - bool async, - bool allocate_on_comm_stream) { -#ifndef DISABLE_NVSHMEM - // In dispatch, CPU will busy-wait until GPU receive tensor size metadata from other ranks, which can be quite long. - // If users of DeepEP need to execute other Python code on other threads, such as KV transfer, their code will get stuck due to GIL - // unless we release GIL here. - pybind11::gil_scoped_release release; - - const int num_channels = config.num_sms / 2; - EP_HOST_ASSERT(config.num_sms % 2 == 0); - EP_HOST_ASSERT(0 < get_num_rdma_ranks() and get_num_rdma_ranks() <= NUM_MAX_RDMA_PEERS); - - bool cached_mode = cached_rdma_channel_prefix_matrix.has_value(); - if (cached_mode) { - EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix.has_value()); - EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum.has_value()); - EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix.has_value()); - EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum.has_value()); - } else { - EP_HOST_ASSERT(num_tokens_per_rank.has_value()); - EP_HOST_ASSERT(num_tokens_per_rdma_rank.has_value()); - EP_HOST_ASSERT(num_tokens_per_expert.has_value()); - } - - // Type checks - if (cached_mode) { - EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix->scalar_type() == torch::kInt32); - EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum->scalar_type() == torch::kInt32); - EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix->scalar_type() == torch::kInt32); - EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum->scalar_type() == torch::kInt32); - } else { - EP_HOST_ASSERT(num_tokens_per_rank->scalar_type() == torch::kInt32); - EP_HOST_ASSERT(num_tokens_per_rdma_rank->scalar_type() == torch::kInt32); - EP_HOST_ASSERT(num_tokens_per_expert->scalar_type() == torch::kInt32); - } - - // Shape and contiguous checks - EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous()); - EP_HOST_ASSERT((x.size(1) * x.element_size()) % sizeof(int4) == 0); - if (cached_mode) { - EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix->dim() == 2 and cached_rdma_channel_prefix_matrix->is_contiguous()); - EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix->size(0) == num_rdma_ranks and - cached_rdma_channel_prefix_matrix->size(1) == num_channels); - EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum->dim() == 1 and cached_recv_rdma_rank_prefix_sum->is_contiguous()); - EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum->size(0) == num_rdma_ranks); - EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix->dim() == 2 and cached_gbl_channel_prefix_matrix->is_contiguous()); - EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix->size(0) == num_ranks and - cached_gbl_channel_prefix_matrix->size(1) == num_channels); - EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum->dim() == 1 and cached_recv_gbl_rank_prefix_sum->is_contiguous()); - EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum->size(0) == num_ranks); - } else { - EP_HOST_ASSERT(num_tokens_per_rank->dim() == 1 and num_tokens_per_rank->is_contiguous()); - EP_HOST_ASSERT(num_tokens_per_rdma_rank->dim() == 1 and num_tokens_per_rdma_rank->is_contiguous()); - EP_HOST_ASSERT(num_tokens_per_expert->dim() == 1 and num_tokens_per_expert->is_contiguous()); - EP_HOST_ASSERT(num_tokens_per_rank->size(0) == num_ranks); - EP_HOST_ASSERT(num_tokens_per_rdma_rank->size(0) == num_rdma_ranks); - EP_HOST_ASSERT(num_tokens_per_expert->size(0) % num_ranks == 0); - EP_HOST_ASSERT(num_tokens_per_expert->size(0) / num_ranks <= NUM_MAX_LOCAL_EXPERTS); - } - - auto num_tokens = static_cast(x.size(0)), hidden = static_cast(x.size(1)), - hidden_int4 = static_cast(x.size(1) * x.element_size() / sizeof(int4)); - auto num_experts = cached_mode ? 0 : static_cast(num_tokens_per_expert->size(0)), num_local_experts = num_experts / num_ranks; - - // Top-k checks - int num_topk = 0; - topk_idx_t* topk_idx_ptr = nullptr; - float* topk_weights_ptr = nullptr; - EP_HOST_ASSERT(topk_idx.has_value() == topk_weights.has_value()); - if (topk_idx.has_value()) { - num_topk = static_cast(topk_idx->size(1)); - EP_HOST_ASSERT(num_experts > 0); - EP_HOST_ASSERT(topk_idx->dim() == 2 and topk_idx->is_contiguous()); - EP_HOST_ASSERT(topk_weights->dim() == 2 and topk_weights->is_contiguous()); - EP_HOST_ASSERT(num_tokens == topk_idx->size(0) and num_tokens == topk_weights->size(0)); - EP_HOST_ASSERT(num_topk == topk_weights->size(1)); - EP_HOST_ASSERT(topk_weights->scalar_type() == torch::kFloat32); - topk_idx_ptr = topk_idx->data_ptr(); - topk_weights_ptr = topk_weights->data_ptr(); - } - - // FP8 scales checks - float* x_scales_ptr = nullptr; - int num_scales = 0, scale_token_stride = 0, scale_hidden_stride = 0; - if (x_scales.has_value()) { - EP_HOST_ASSERT(x.element_size() == 1); - EP_HOST_ASSERT(x_scales->scalar_type() == torch::kFloat32 or x_scales->scalar_type() == torch::kInt); - EP_HOST_ASSERT(x_scales->dim() == 2); - EP_HOST_ASSERT(x_scales->size(0) == num_tokens); - num_scales = x_scales->dim() == 1 ? 1 : static_cast(x_scales->size(1)); - x_scales_ptr = static_cast(x_scales->data_ptr()); - scale_token_stride = static_cast(x_scales->stride(0)); - scale_hidden_stride = static_cast(x_scales->stride(1)); - } - - // Allocate all tensors on comm stream if set - // NOTES: do not allocate tensors upfront! - auto compute_stream = at::cuda::getCurrentCUDAStream(); - if (allocate_on_comm_stream) { - EP_HOST_ASSERT(previous_event.has_value() and async); - at::cuda::setCurrentCUDAStream(comm_stream); - } - - // Wait previous tasks to be finished - if (previous_event.has_value()) { - stream_wait(comm_stream, previous_event.value()); - } else { - stream_wait(comm_stream, compute_stream); - } - - // Create handles (only return for non-cached mode) - int num_recv_tokens = -1, num_rdma_recv_tokens = -1; - auto rdma_channel_prefix_matrix = torch::Tensor(); - auto recv_rdma_rank_prefix_sum = torch::Tensor(); - auto gbl_channel_prefix_matrix = torch::Tensor(); - auto recv_gbl_rank_prefix_sum = torch::Tensor(); - std::vector num_recv_tokens_per_expert_list; - - // Barrier or send sizes - if (cached_mode) { - num_recv_tokens = cached_num_recv_tokens; - num_rdma_recv_tokens = cached_num_rdma_recv_tokens; - rdma_channel_prefix_matrix = cached_rdma_channel_prefix_matrix.value(); - recv_rdma_rank_prefix_sum = cached_recv_rdma_rank_prefix_sum.value(); - gbl_channel_prefix_matrix = cached_gbl_channel_prefix_matrix.value(); - recv_gbl_rank_prefix_sum = cached_recv_gbl_rank_prefix_sum.value(); - - // Just a barrier and clean flags - internode::cached_notify(hidden_int4, - num_scales, - num_topk, - num_topk, - num_ranks, - num_channels, - 0, - nullptr, - nullptr, - nullptr, - nullptr, - rdma_buffer_ptr, - config.num_max_rdma_chunked_recv_tokens, - buffer_ptrs_gpu, - config.num_max_nvl_chunked_recv_tokens, - barrier_signal_ptrs_gpu, - rank, - comm_stream, - config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks), - num_nvl_bytes, - true, - low_latency_mode); - } else { - rdma_channel_prefix_matrix = torch::empty({num_rdma_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); - recv_rdma_rank_prefix_sum = torch::empty({num_rdma_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); - gbl_channel_prefix_matrix = torch::empty({num_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); - recv_gbl_rank_prefix_sum = torch::empty({num_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); - - // Send sizes - *moe_recv_counter = -1, *moe_recv_rdma_counter = -1; - for (int i = 0; i < num_local_experts; ++i) - moe_recv_expert_counter[i] = -1; - internode::notify_dispatch(num_tokens_per_rank->data_ptr(), - moe_recv_counter_mapped, - num_ranks, - num_tokens_per_rdma_rank->data_ptr(), - moe_recv_rdma_counter_mapped, - num_tokens_per_expert->data_ptr(), - moe_recv_expert_counter_mapped, - num_experts, - is_token_in_rank.data_ptr(), - num_tokens, - num_worst_tokens, - num_channels, - hidden_int4, - num_scales, - num_topk, - expert_alignment, - rdma_channel_prefix_matrix.data_ptr(), - recv_rdma_rank_prefix_sum.data_ptr(), - gbl_channel_prefix_matrix.data_ptr(), - recv_gbl_rank_prefix_sum.data_ptr(), - rdma_buffer_ptr, - config.num_max_rdma_chunked_recv_tokens, - buffer_ptrs_gpu, - config.num_max_nvl_chunked_recv_tokens, - barrier_signal_ptrs_gpu, - rank, - comm_stream, - config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks), - num_nvl_bytes, - low_latency_mode); - - // Synchronize total received tokens and tokens per expert - if (num_worst_tokens > 0) { - num_recv_tokens = num_worst_tokens; - num_rdma_recv_tokens = num_worst_tokens; - } else { - auto start_time = std::chrono::high_resolution_clock::now(); - while (true) { - // Read total count - num_recv_tokens = static_cast(*moe_recv_counter); - num_rdma_recv_tokens = static_cast(*moe_recv_rdma_counter); - - // Read per-expert count - bool ready = (num_recv_tokens >= 0) and (num_rdma_recv_tokens >= 0); - for (int i = 0; i < num_local_experts and ready; ++i) - ready &= moe_recv_expert_counter[i] >= 0; - - if (ready) - break; - - // Timeout check - if (std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - start_time).count() > - NUM_CPU_TIMEOUT_SECS) { - printf("Global rank: %d, num_recv_tokens: %d, num_rdma_recv_tokens: %d\n", rank, num_recv_tokens, num_rdma_recv_tokens); - for (int i = 0; i < num_local_experts; ++i) - printf("moe_recv_expert_counter[%d]: %d\n", i, moe_recv_expert_counter[i]); - throw std::runtime_error("DeepEP error: timeout (dispatch CPU)"); - } - } - num_recv_tokens_per_expert_list = std::vector(moe_recv_expert_counter, moe_recv_expert_counter + num_local_experts); - } - } - - // Allocate new tensors - auto recv_x = torch::empty({num_recv_tokens, hidden}, x.options()); - auto recv_topk_idx = std::optional(), recv_topk_weights = std::optional(), - recv_x_scales = std::optional(); - auto recv_src_meta = std::optional(); - auto recv_rdma_channel_prefix_matrix = std::optional(); - auto recv_gbl_channel_prefix_matrix = std::optional(); - auto send_rdma_head = std::optional(); - auto send_nvl_head = std::optional(); - if (not cached_mode) { - recv_src_meta = torch::empty({num_recv_tokens, internode::get_source_meta_bytes()}, dtype(torch::kByte).device(torch::kCUDA)); - recv_rdma_channel_prefix_matrix = torch::empty({num_rdma_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); - recv_gbl_channel_prefix_matrix = torch::empty({num_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); - send_rdma_head = torch::empty({num_tokens, num_rdma_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); - send_nvl_head = torch::empty({num_rdma_recv_tokens, NUM_MAX_NVL_PEERS}, dtype(torch::kInt32).device(torch::kCUDA)); - } - - // Assign pointers - topk_idx_t* recv_topk_idx_ptr = nullptr; - float* recv_topk_weights_ptr = nullptr; - float* recv_x_scales_ptr = nullptr; - if (topk_idx.has_value()) { - recv_topk_idx = torch::empty({num_recv_tokens, num_topk}, topk_idx->options()); - recv_topk_weights = torch::empty({num_recv_tokens, num_topk}, topk_weights->options()); - recv_topk_idx_ptr = recv_topk_idx->data_ptr(); - recv_topk_weights_ptr = recv_topk_weights->data_ptr(); - } - if (x_scales.has_value()) { - recv_x_scales = x_scales->dim() == 1 ? torch::empty({num_recv_tokens}, x_scales->options()) - : torch::empty({num_recv_tokens, num_scales}, x_scales->options()); - recv_x_scales_ptr = static_cast(recv_x_scales->data_ptr()); - } - - // Launch data dispatch - // NOTES: the buffer size checks are moved into the `.cu` file - internode::dispatch(recv_x.data_ptr(), - recv_x_scales_ptr, - recv_topk_idx_ptr, - recv_topk_weights_ptr, - cached_mode ? nullptr : recv_src_meta->data_ptr(), - x.data_ptr(), - x_scales_ptr, - topk_idx_ptr, - topk_weights_ptr, - cached_mode ? nullptr : send_rdma_head->data_ptr(), - cached_mode ? nullptr : send_nvl_head->data_ptr(), - cached_mode ? nullptr : recv_rdma_channel_prefix_matrix->data_ptr(), - cached_mode ? nullptr : recv_gbl_channel_prefix_matrix->data_ptr(), - rdma_channel_prefix_matrix.data_ptr(), - recv_rdma_rank_prefix_sum.data_ptr(), - gbl_channel_prefix_matrix.data_ptr(), - recv_gbl_rank_prefix_sum.data_ptr(), - is_token_in_rank.data_ptr(), - num_tokens, - num_worst_tokens, - hidden_int4, - num_scales, - num_topk, - num_experts, - scale_token_stride, - scale_hidden_stride, - rdma_buffer_ptr, - config.num_max_rdma_chunked_send_tokens, - config.num_max_rdma_chunked_recv_tokens, - buffer_ptrs_gpu, - config.num_max_nvl_chunked_send_tokens, - config.num_max_nvl_chunked_recv_tokens, - rank, - num_ranks, - cached_mode, - comm_stream, - num_channels, - low_latency_mode); - - // Wait streams - std::optional event; - if (async) { - event = EventHandle(comm_stream); - for (auto& t : {x, - is_token_in_rank, - recv_x, - rdma_channel_prefix_matrix, - recv_rdma_rank_prefix_sum, - gbl_channel_prefix_matrix, - recv_gbl_rank_prefix_sum}) { - t.record_stream(comm_stream); - if (allocate_on_comm_stream) - t.record_stream(compute_stream); - } - for (auto& to : {x_scales, - topk_idx, - topk_weights, - num_tokens_per_rank, - num_tokens_per_rdma_rank, - num_tokens_per_expert, - cached_rdma_channel_prefix_matrix, - cached_recv_rdma_rank_prefix_sum, - cached_gbl_channel_prefix_matrix, - cached_recv_gbl_rank_prefix_sum, - recv_topk_idx, - recv_topk_weights, - recv_x_scales, - recv_rdma_channel_prefix_matrix, - recv_gbl_channel_prefix_matrix, - send_rdma_head, - send_nvl_head, - recv_src_meta}) { - to.has_value() ? to->record_stream(comm_stream) : void(); - if (allocate_on_comm_stream) - to.has_value() ? to->record_stream(compute_stream) : void(); - } - } else { - stream_wait(compute_stream, comm_stream); - } - - // Switch back compute stream - if (allocate_on_comm_stream) - at::cuda::setCurrentCUDAStream(compute_stream); - - // Return values - return {recv_x, - recv_x_scales, - recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, - rdma_channel_prefix_matrix, - gbl_channel_prefix_matrix, - recv_rdma_channel_prefix_matrix, - recv_rdma_rank_prefix_sum, - recv_gbl_channel_prefix_matrix, - recv_gbl_rank_prefix_sum, - recv_src_meta, - send_rdma_head, - send_nvl_head, - event}; -#else - EP_HOST_ASSERT(false and "NVSHMEM is disabled during compilation"); - return {}; -#endif -} - -std::tuple, std::optional> Buffer::internode_combine( - const torch::Tensor& x, - const std::optional& topk_weights, - const std::optional& bias_0, - const std::optional& bias_1, - const torch::Tensor& src_meta, - const torch::Tensor& is_combined_token_in_rank, - const torch::Tensor& rdma_channel_prefix_matrix, - const torch::Tensor& rdma_rank_prefix_sum, - const torch::Tensor& gbl_channel_prefix_matrix, - const torch::Tensor& combined_rdma_head, - const torch::Tensor& combined_nvl_head, - const Config& config, - std::optional& previous_event, - bool async, - bool allocate_on_comm_stream) { -#ifndef DISABLE_NVSHMEM - const int num_channels = config.num_sms / 2; - EP_HOST_ASSERT(config.num_sms % 2 == 0); - - // Shape and contiguous checks - EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous()); - EP_HOST_ASSERT(src_meta.dim() == 2 and src_meta.is_contiguous() and src_meta.scalar_type() == torch::kByte); - EP_HOST_ASSERT(is_combined_token_in_rank.dim() == 2 and is_combined_token_in_rank.is_contiguous() and - is_combined_token_in_rank.scalar_type() == torch::kBool); - EP_HOST_ASSERT(rdma_channel_prefix_matrix.dim() == 2 and rdma_channel_prefix_matrix.is_contiguous() and - rdma_channel_prefix_matrix.scalar_type() == torch::kInt32); - EP_HOST_ASSERT(rdma_rank_prefix_sum.dim() == 1 and rdma_rank_prefix_sum.is_contiguous() and - rdma_rank_prefix_sum.scalar_type() == torch::kInt32); - EP_HOST_ASSERT(gbl_channel_prefix_matrix.dim() == 2 and gbl_channel_prefix_matrix.is_contiguous() and - gbl_channel_prefix_matrix.scalar_type() == torch::kInt32); - EP_HOST_ASSERT(combined_rdma_head.dim() == 2 and combined_rdma_head.is_contiguous() and - combined_rdma_head.scalar_type() == torch::kInt32); - EP_HOST_ASSERT(combined_nvl_head.dim() == 2 and combined_nvl_head.is_contiguous() and combined_nvl_head.scalar_type() == torch::kInt32); - - auto num_tokens = static_cast(x.size(0)), hidden = static_cast(x.size(1)), - hidden_int4 = static_cast(x.size(1) * x.element_size() / sizeof(int4)); - auto num_combined_tokens = static_cast(is_combined_token_in_rank.size(0)); - EP_HOST_ASSERT((hidden * x.element_size()) % sizeof(int4) == 0); - EP_HOST_ASSERT(src_meta.size(1) == internode::get_source_meta_bytes()); - EP_HOST_ASSERT(is_combined_token_in_rank.size(1) == num_ranks); - EP_HOST_ASSERT(rdma_channel_prefix_matrix.size(0) == num_rdma_ranks and rdma_channel_prefix_matrix.size(1) == num_channels); - EP_HOST_ASSERT(rdma_rank_prefix_sum.size(0) == num_rdma_ranks); - EP_HOST_ASSERT(gbl_channel_prefix_matrix.size(0) == num_ranks and gbl_channel_prefix_matrix.size(1) == num_channels); - EP_HOST_ASSERT(combined_rdma_head.dim() == 2 and combined_rdma_head.size(0) == num_combined_tokens and - combined_rdma_head.size(1) == num_rdma_ranks); - EP_HOST_ASSERT(combined_nvl_head.dim() == 2 and combined_nvl_head.size(1) == NUM_MAX_NVL_PEERS); - - // Allocate all tensors on comm stream if set - // NOTES: do not allocate tensors upfront! - auto compute_stream = at::cuda::getCurrentCUDAStream(); - if (allocate_on_comm_stream) { - EP_HOST_ASSERT(previous_event.has_value() and async); - at::cuda::setCurrentCUDAStream(comm_stream); - } - - // Wait previous tasks to be finished - if (previous_event.has_value()) { - stream_wait(comm_stream, previous_event.value()); - } else { - stream_wait(comm_stream, compute_stream); - } - - // Top-k checks - int num_topk = 0; - auto combined_topk_weights = std::optional(); - float* topk_weights_ptr = nullptr; - float* combined_topk_weights_ptr = nullptr; - if (topk_weights.has_value()) { - EP_HOST_ASSERT(topk_weights->dim() == 2 and topk_weights->is_contiguous()); - EP_HOST_ASSERT(topk_weights->size(0) == num_tokens); - EP_HOST_ASSERT(topk_weights->scalar_type() == torch::kFloat32); - num_topk = static_cast(topk_weights->size(1)); - topk_weights_ptr = topk_weights->data_ptr(); - combined_topk_weights = torch::empty({num_combined_tokens, num_topk}, topk_weights->options()); - combined_topk_weights_ptr = combined_topk_weights->data_ptr(); - } - - // Extra check for avoid-dead-lock design - EP_HOST_ASSERT(config.num_max_nvl_chunked_recv_tokens % num_rdma_ranks == 0); - EP_HOST_ASSERT(config.num_max_nvl_chunked_send_tokens <= config.num_max_nvl_chunked_recv_tokens / num_rdma_ranks); - - // Launch barrier and reset queue head and tail - internode::cached_notify(hidden_int4, - 0, - 0, - num_topk, - num_ranks, - num_channels, - num_combined_tokens, - combined_rdma_head.data_ptr(), - rdma_channel_prefix_matrix.data_ptr(), - rdma_rank_prefix_sum.data_ptr(), - combined_nvl_head.data_ptr(), - rdma_buffer_ptr, - config.num_max_rdma_chunked_recv_tokens, - buffer_ptrs_gpu, - config.num_max_nvl_chunked_recv_tokens, - barrier_signal_ptrs_gpu, - rank, - comm_stream, - config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks), - num_nvl_bytes, - false, - low_latency_mode); - - // Assign bias pointers - auto bias_opts = std::vector>({bias_0, bias_1}); - void* bias_ptrs[2] = {nullptr, nullptr}; - for (int i = 0; i < 2; ++i) - if (bias_opts[i].has_value()) { - auto bias = bias_opts[i].value(); - EP_HOST_ASSERT(bias.dim() == 2 and bias.is_contiguous()); - EP_HOST_ASSERT(bias.scalar_type() == x.scalar_type()); - EP_HOST_ASSERT(bias.size(0) == num_combined_tokens and bias.size(1) == hidden); - bias_ptrs[i] = bias.data_ptr(); - } - - // Launch data combine - auto combined_x = torch::empty({num_combined_tokens, hidden}, x.options()); - internode::combine(at::cuda::ScalarTypeToCudaDataType(x.scalar_type()), - combined_x.data_ptr(), - combined_topk_weights_ptr, - is_combined_token_in_rank.data_ptr(), - x.data_ptr(), - topk_weights_ptr, - bias_ptrs[0], - bias_ptrs[1], - combined_rdma_head.data_ptr(), - combined_nvl_head.data_ptr(), - src_meta.data_ptr(), - rdma_channel_prefix_matrix.data_ptr(), - rdma_rank_prefix_sum.data_ptr(), - gbl_channel_prefix_matrix.data_ptr(), - num_tokens, - num_combined_tokens, - hidden, - num_topk, - rdma_buffer_ptr, - config.num_max_rdma_chunked_send_tokens, - config.num_max_rdma_chunked_recv_tokens, - buffer_ptrs_gpu, - config.num_max_nvl_chunked_send_tokens, - config.num_max_nvl_chunked_recv_tokens, - rank, - num_ranks, - comm_stream, - num_channels, - low_latency_mode); - - // Wait streams - std::optional event; - if (async) { - event = EventHandle(comm_stream); - for (auto& t : {x, - src_meta, - is_combined_token_in_rank, - rdma_channel_prefix_matrix, - rdma_rank_prefix_sum, - gbl_channel_prefix_matrix, - combined_x, - combined_rdma_head, - combined_nvl_head}) { - t.record_stream(comm_stream); - if (allocate_on_comm_stream) - t.record_stream(compute_stream); - } - for (auto& to : {topk_weights, combined_topk_weights, bias_0, bias_1}) { - to.has_value() ? to->record_stream(comm_stream) : void(); - if (allocate_on_comm_stream) - to.has_value() ? to->record_stream(compute_stream) : void(); - } - } else { - stream_wait(compute_stream, comm_stream); - } - - // Switch back compute stream - if (allocate_on_comm_stream) - at::cuda::setCurrentCUDAStream(compute_stream); - - // Return values - return {combined_x, combined_topk_weights, event}; -#else - EP_HOST_ASSERT(false and "NVSHMEM is disabled during compilation"); - return {}; -#endif -} - -void Buffer::clean_low_latency_buffer(int num_max_dispatch_tokens_per_rank, int hidden, int num_experts) { -#ifndef DISABLE_NVSHMEM - EP_HOST_ASSERT(low_latency_mode); - - auto layout = LowLatencyLayout(rdma_buffer_ptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts); - auto clean_meta_0 = layout.buffers[0].clean_meta(); - auto clean_meta_1 = layout.buffers[1].clean_meta(); - - auto check_boundary = [=](void* ptr, size_t num_bytes) { - auto offset = reinterpret_cast(ptr) - reinterpret_cast(rdma_buffer_ptr); - EP_HOST_ASSERT(0 <= offset and offset + num_bytes <= num_rdma_bytes); - }; - check_boundary(clean_meta_0.first, clean_meta_0.second * sizeof(int)); - check_boundary(clean_meta_1.first, clean_meta_1.second * sizeof(int)); - - internode_ll::clean_low_latency_buffer(clean_meta_0.first, - clean_meta_0.second, - clean_meta_1.first, - clean_meta_1.second, - rank, - num_ranks, - mask_buffer_ptr, - sync_buffer_ptr, - at::cuda::getCurrentCUDAStream()); -#else - EP_HOST_ASSERT(false and "NVSHMEM is disabled during compilation"); -#endif -} - -std::tuple, - torch::Tensor, - torch::Tensor, - torch::Tensor, - std::optional, - std::optional>> -Buffer::low_latency_dispatch(const torch::Tensor& x, - const torch::Tensor& topk_idx, - const std::optional& cumulative_local_expert_recv_stats, - const std::optional& dispatch_wait_recv_cost_stats, - int num_max_dispatch_tokens_per_rank, - int num_experts, - bool use_fp8, - bool round_scale, - bool use_ue8m0, - bool async, - bool return_recv_hook) { -#ifndef DISABLE_NVSHMEM - EP_HOST_ASSERT(low_latency_mode); - - // Tensor checks - // By default using `ptp128c` FP8 cast - EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous() and x.scalar_type() == torch::kBFloat16); - EP_HOST_ASSERT(x.size(1) % sizeof(int4) == 0 and x.size(1) % 128 == 0); - EP_HOST_ASSERT(topk_idx.dim() == 2 and topk_idx.is_contiguous()); - EP_HOST_ASSERT(x.size(0) == topk_idx.size(0) and x.size(0) <= num_max_dispatch_tokens_per_rank); - EP_HOST_ASSERT(topk_idx.scalar_type() == c10::CppTypeToScalarType::value); - EP_HOST_ASSERT(num_experts % num_ranks == 0); - - // Diagnosis tensors - if (cumulative_local_expert_recv_stats.has_value()) { - EP_HOST_ASSERT(cumulative_local_expert_recv_stats->scalar_type() == torch::kInt); - EP_HOST_ASSERT(cumulative_local_expert_recv_stats->dim() == 1 and cumulative_local_expert_recv_stats->is_contiguous()); - EP_HOST_ASSERT(cumulative_local_expert_recv_stats->size(0) == num_experts / num_ranks); - } - if (dispatch_wait_recv_cost_stats.has_value()) { - EP_HOST_ASSERT(dispatch_wait_recv_cost_stats->scalar_type() == torch::kInt64); - EP_HOST_ASSERT(dispatch_wait_recv_cost_stats->dim() == 1 and dispatch_wait_recv_cost_stats->is_contiguous()); - EP_HOST_ASSERT(dispatch_wait_recv_cost_stats->size(0) == num_ranks); - } - - auto num_tokens = static_cast(x.size(0)), hidden = static_cast(x.size(1)); - auto num_topk = static_cast(topk_idx.size(1)); - auto num_local_experts = num_experts / num_ranks; - - // Buffer control - LowLatencyLayout layout(rdma_buffer_ptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts); - EP_HOST_ASSERT(layout.total_bytes <= num_rdma_bytes); - auto buffer = layout.buffers[low_latency_buffer_idx]; - auto next_buffer = layout.buffers[low_latency_buffer_idx ^= 1]; - - // Wait previous tasks to be finished - // NOTES: the hook mode will always use the default stream - auto compute_stream = at::cuda::getCurrentCUDAStream(); - auto launch_stream = return_recv_hook ? compute_stream : comm_stream; - EP_HOST_ASSERT(not(async and return_recv_hook)); - if (not return_recv_hook) - stream_wait(launch_stream, compute_stream); - - // Allocate packed tensors - auto packed_recv_x = torch::empty({num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank, hidden}, - x.options().dtype(use_fp8 ? torch::kFloat8_e4m3fn : torch::kBFloat16)); - auto packed_recv_src_info = - torch::empty({num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank}, torch::dtype(torch::kInt32).device(torch::kCUDA)); - auto packed_recv_layout_range = torch::empty({num_local_experts, num_ranks}, torch::dtype(torch::kInt64).device(torch::kCUDA)); - auto packed_recv_count = torch::empty({num_local_experts}, torch::dtype(torch::kInt32).device(torch::kCUDA)); - - // Allocate column-majored scales - auto packed_recv_x_scales = std::optional(); - void* packed_recv_x_scales_ptr = nullptr; - EP_HOST_ASSERT((num_ranks * num_max_dispatch_tokens_per_rank) % 4 == 0 and "TMA requires the number of tokens to be multiple of 4"); - - if (use_fp8) { - // TODO: support unaligned cases - EP_HOST_ASSERT(hidden % 512 == 0); - if (not use_ue8m0) { - packed_recv_x_scales = torch::empty({num_local_experts, hidden / 128, num_ranks * num_max_dispatch_tokens_per_rank}, - torch::dtype(torch::kFloat32).device(torch::kCUDA)); - } else { - EP_HOST_ASSERT(round_scale); - packed_recv_x_scales = torch::empty({num_local_experts, hidden / 512, num_ranks * num_max_dispatch_tokens_per_rank}, - torch::dtype(torch::kInt).device(torch::kCUDA)); - } - packed_recv_x_scales = torch::transpose(packed_recv_x_scales.value(), 1, 2); - packed_recv_x_scales_ptr = packed_recv_x_scales->data_ptr(); - } - - // Kernel launch - auto next_clean_meta = next_buffer.clean_meta(); - auto launcher = [=](int phases) { - internode_ll::dispatch( - packed_recv_x.data_ptr(), - packed_recv_x_scales_ptr, - packed_recv_src_info.data_ptr(), - packed_recv_layout_range.data_ptr(), - packed_recv_count.data_ptr(), - mask_buffer_ptr, - cumulative_local_expert_recv_stats.has_value() ? cumulative_local_expert_recv_stats->data_ptr() : nullptr, - dispatch_wait_recv_cost_stats.has_value() ? dispatch_wait_recv_cost_stats->data_ptr() : nullptr, - buffer.dispatch_rdma_recv_data_buffer, - buffer.dispatch_rdma_recv_count_buffer, - buffer.dispatch_rdma_send_buffer, - x.data_ptr(), - topk_idx.data_ptr(), - next_clean_meta.first, - next_clean_meta.second, - num_tokens, - hidden, - num_max_dispatch_tokens_per_rank, - num_topk, - num_experts, - rank, - num_ranks, - use_fp8, - round_scale, - use_ue8m0, - workspace, - num_device_sms, - launch_stream, - phases); - }; - launcher(return_recv_hook ? LOW_LATENCY_SEND_PHASE : (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE)); - - // Wait streams - std::optional event; - if (async) { - // NOTES: we must ensure the all tensors will not be deallocated before the stream-wait happens, - // so in Python API, we must wrap all tensors into the event handle. - event = EventHandle(launch_stream); - } else if (not return_recv_hook) { - stream_wait(compute_stream, launch_stream); - } - - // Receiver callback - std::optional> recv_hook = std::nullopt; - if (return_recv_hook) - recv_hook = [=]() { launcher(LOW_LATENCY_RECV_PHASE); }; - - // Return values - return {packed_recv_x, packed_recv_x_scales, packed_recv_count, packed_recv_src_info, packed_recv_layout_range, event, recv_hook}; -#else - EP_HOST_ASSERT(false and "NVSHMEM is disabled during compilation"); - return {}; -#endif -} - -std::tuple, std::optional>> Buffer::low_latency_combine( - const torch::Tensor& x, - const torch::Tensor& topk_idx, - const torch::Tensor& topk_weights, - const torch::Tensor& src_info, - const torch::Tensor& layout_range, - const std::optional& combine_wait_recv_cost_stats, - int num_max_dispatch_tokens_per_rank, - int num_experts, - bool use_logfmt, - bool zero_copy, - bool async, - bool return_recv_hook, - const std::optional& out) { -#ifndef DISABLE_NVSHMEM - EP_HOST_ASSERT(low_latency_mode); - - // Tensor checks - EP_HOST_ASSERT(x.dim() == 3 and x.is_contiguous() and x.scalar_type() == torch::kBFloat16); - EP_HOST_ASSERT(x.size(0) == num_experts / num_ranks); - EP_HOST_ASSERT(x.size(1) == num_ranks * num_max_dispatch_tokens_per_rank); - EP_HOST_ASSERT(x.size(2) % sizeof(int4) == 0 and x.size(2) % 128 == 0); - EP_HOST_ASSERT(topk_idx.dim() == 2 and topk_idx.is_contiguous()); - EP_HOST_ASSERT(topk_idx.size(0) == topk_weights.size(0) and topk_idx.size(1) == topk_weights.size(1)); - EP_HOST_ASSERT(topk_idx.scalar_type() == c10::CppTypeToScalarType::value); - EP_HOST_ASSERT(topk_weights.dim() == 2 and topk_weights.is_contiguous()); - EP_HOST_ASSERT(topk_weights.size(0) <= num_max_dispatch_tokens_per_rank); - EP_HOST_ASSERT(topk_weights.scalar_type() == torch::kFloat32); - EP_HOST_ASSERT(src_info.dim() == 2 and src_info.is_contiguous()); - EP_HOST_ASSERT(src_info.scalar_type() == torch::kInt32 and x.size(0) == src_info.size(0)); - EP_HOST_ASSERT(layout_range.dim() == 2 and layout_range.is_contiguous()); - EP_HOST_ASSERT(layout_range.scalar_type() == torch::kInt64); - EP_HOST_ASSERT(layout_range.size(0) == num_experts / num_ranks and layout_range.size(1) == num_ranks); - - if (combine_wait_recv_cost_stats.has_value()) { - EP_HOST_ASSERT(combine_wait_recv_cost_stats->scalar_type() == torch::kInt64); - EP_HOST_ASSERT(combine_wait_recv_cost_stats->dim() == 1 and combine_wait_recv_cost_stats->is_contiguous()); - EP_HOST_ASSERT(combine_wait_recv_cost_stats->size(0) == num_ranks); - } - - auto hidden = static_cast(x.size(2)); - auto num_topk = static_cast(topk_weights.size(1)); - auto num_combined_tokens = static_cast(topk_weights.size(0)); - - // Buffer control - LowLatencyLayout layout(rdma_buffer_ptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts); - EP_HOST_ASSERT(layout.total_bytes <= num_rdma_bytes); - auto buffer = layout.buffers[low_latency_buffer_idx]; - auto next_buffer = layout.buffers[low_latency_buffer_idx ^= 1]; - - // Wait previous tasks to be finished - // NOTES: the hook mode will always use the default stream - auto compute_stream = at::cuda::getCurrentCUDAStream(); - auto launch_stream = return_recv_hook ? compute_stream : comm_stream; - EP_HOST_ASSERT(not(async and return_recv_hook)); - if (not return_recv_hook) - stream_wait(launch_stream, compute_stream); - - // Allocate output tensor - torch::Tensor combined_x; - if (out.has_value()) { - EP_HOST_ASSERT(out->dim() == 2 and out->is_contiguous()); - EP_HOST_ASSERT(out->size(0) == num_combined_tokens and out->size(1) == hidden); - EP_HOST_ASSERT(out->scalar_type() == x.scalar_type()); - combined_x = out.value(); - } else { - combined_x = torch::empty({num_combined_tokens, hidden}, x.options()); - } - - // Kernel launch - auto next_clean_meta = next_buffer.clean_meta(); - auto launcher = [=](int phases) { - internode_ll::combine(combined_x.data_ptr(), - buffer.combine_rdma_recv_data_buffer, - buffer.combine_rdma_recv_flag_buffer, - buffer.combine_rdma_send_buffer, - x.data_ptr(), - topk_idx.data_ptr(), - topk_weights.data_ptr(), - src_info.data_ptr(), - layout_range.data_ptr(), - mask_buffer_ptr, - combine_wait_recv_cost_stats.has_value() ? combine_wait_recv_cost_stats->data_ptr() : nullptr, - next_clean_meta.first, - next_clean_meta.second, - num_combined_tokens, - hidden, - num_max_dispatch_tokens_per_rank, - num_topk, - num_experts, - rank, - num_ranks, - use_logfmt, - workspace, - num_device_sms, - launch_stream, - phases, - zero_copy); - }; - launcher(return_recv_hook ? LOW_LATENCY_SEND_PHASE : (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE)); - - // Wait streams - std::optional event; - if (async) { - // NOTES: we must ensure the all tensors will not be deallocated before the stream-wait happens, - // so in Python API, we must wrap all tensors into the event handle. - event = EventHandle(launch_stream); - } else if (not return_recv_hook) { - stream_wait(compute_stream, launch_stream); - } - - // Receiver callback - std::optional> recv_hook = std::nullopt; - if (return_recv_hook) - recv_hook = [=]() { launcher(LOW_LATENCY_RECV_PHASE); }; - - // Return values - return {combined_x, event, recv_hook}; -#else - EP_HOST_ASSERT(false and "NVSHMEM is disabled during compilation"); - return {}; -#endif -} - -torch::Tensor Buffer::get_next_low_latency_combine_buffer(int num_max_dispatch_tokens_per_rank, int hidden, int num_experts) const { -#ifndef DISABLE_NVSHMEM - LowLatencyLayout layout(rdma_buffer_ptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts); - - auto buffer = layout.buffers[low_latency_buffer_idx]; - auto dtype = torch::kBFloat16; - auto num_msg_elems = static_cast(buffer.num_bytes_per_combine_msg / elementSize(torch::kBFloat16)); - - EP_HOST_ASSERT(buffer.num_bytes_per_combine_msg % elementSize(torch::kBFloat16) == 0); - return torch::from_blob(buffer.combine_rdma_send_buffer_data_start, - {num_experts / num_ranks, num_ranks * num_max_dispatch_tokens_per_rank, hidden}, - {num_ranks * num_max_dispatch_tokens_per_rank * num_msg_elems, num_msg_elems, 1}, - torch::TensorOptions().dtype(dtype).device(torch::kCUDA)); -#else - EP_HOST_ASSERT(false and "NVSHMEM is disabled during compilation"); - return {}; -#endif -} - -bool is_sm90_compiled() { -#ifndef DISABLE_SM90_FEATURES - return true; -#else - return false; -#endif -} - -void Buffer::low_latency_update_mask_buffer(int rank_to_mask, bool mask) { - EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled"); - EP_HOST_ASSERT(rank_to_mask >= 0 and rank_to_mask < num_ranks); - internode_ll::update_mask_buffer(mask_buffer_ptr, rank_to_mask, mask, at::cuda::getCurrentCUDAStream()); -} - -void Buffer::low_latency_query_mask_buffer(const torch::Tensor& mask_status) { - EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled"); - EP_HOST_ASSERT(mask_status.numel() == num_ranks && mask_status.scalar_type() == torch::kInt32); - - internode_ll::query_mask_buffer( - mask_buffer_ptr, num_ranks, reinterpret_cast(mask_status.data_ptr()), at::cuda::getCurrentCUDAStream()); -} - -void Buffer::low_latency_clean_mask_buffer() { - EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled"); - internode_ll::clean_mask_buffer(mask_buffer_ptr, num_ranks, at::cuda::getCurrentCUDAStream()); -} - -} // namespace deep_ep - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.doc() = "DeepEP: an efficient expert-parallel communication library"; - - pybind11::class_(m, "Config") - .def(pybind11::init(), - py::arg("num_sms") = 20, - py::arg("num_max_nvl_chunked_send_tokens") = 6, - py::arg("num_max_nvl_chunked_recv_tokens") = 256, - py::arg("num_max_rdma_chunked_send_tokens") = 6, - py::arg("num_max_rdma_chunked_recv_tokens") = 256) - .def("get_nvl_buffer_size_hint", &deep_ep::Config::get_nvl_buffer_size_hint) - .def("get_rdma_buffer_size_hint", &deep_ep::Config::get_rdma_buffer_size_hint); - m.def("get_low_latency_rdma_size_hint", &deep_ep::get_low_latency_rdma_size_hint); - - pybind11::class_(m, "EventHandle") - .def(pybind11::init<>()) - .def("current_stream_wait", &deep_ep::EventHandle::current_stream_wait); - - pybind11::class_(m, "Buffer") - .def(pybind11::init()) - .def("is_available", &deep_ep::Buffer::is_available) - .def("get_num_rdma_ranks", &deep_ep::Buffer::get_num_rdma_ranks) - .def("get_rdma_rank", &deep_ep::Buffer::get_rdma_rank) - .def("get_root_rdma_rank", &deep_ep::Buffer::get_root_rdma_rank) - .def("get_local_device_id", &deep_ep::Buffer::get_local_device_id) - .def("get_local_ipc_handle", &deep_ep::Buffer::get_local_ipc_handle) - .def("get_local_nvshmem_unique_id", &deep_ep::Buffer::get_local_nvshmem_unique_id) - .def("get_local_buffer_tensor", &deep_ep::Buffer::get_local_buffer_tensor) - .def("get_comm_stream", &deep_ep::Buffer::get_comm_stream) - .def("sync", &deep_ep::Buffer::sync) - .def("destroy", &deep_ep::Buffer::destroy) - .def("get_dispatch_layout", &deep_ep::Buffer::get_dispatch_layout) - .def("intranode_dispatch", &deep_ep::Buffer::intranode_dispatch) - .def("intranode_combine", &deep_ep::Buffer::intranode_combine) - .def("internode_dispatch", &deep_ep::Buffer::internode_dispatch) - .def("internode_combine", &deep_ep::Buffer::internode_combine) - .def("clean_low_latency_buffer", &deep_ep::Buffer::clean_low_latency_buffer) - .def("low_latency_dispatch", &deep_ep::Buffer::low_latency_dispatch) - .def("low_latency_combine", &deep_ep::Buffer::low_latency_combine) - .def("low_latency_update_mask_buffer", &deep_ep::Buffer::low_latency_update_mask_buffer) - .def("low_latency_query_mask_buffer", &deep_ep::Buffer::low_latency_query_mask_buffer) - .def("low_latency_clean_mask_buffer", &deep_ep::Buffer::low_latency_clean_mask_buffer) - .def("get_next_low_latency_combine_buffer", &deep_ep::Buffer::get_next_low_latency_combine_buffer); - - m.def("is_sm90_compiled", deep_ep::is_sm90_compiled); - m.attr("topk_idx_t") = - py::reinterpret_borrow((PyObject*)torch::getTHPDtype(c10::CppTypeToScalarType::value)); -} diff --git a/csrc/deep_ep.hpp b/csrc/deep_ep.hpp deleted file mode 100644 index 5fb90bfff..000000000 --- a/csrc/deep_ep.hpp +++ /dev/null @@ -1,300 +0,0 @@ -#pragma once - -// Forcibly disable NDEBUG -#ifdef NDEBUG -#undef NDEBUG -#endif - -#include -#include -#include - -#include -#include - -#include "config.hpp" -#include "event.hpp" -#include "kernels/configs.cuh" -#include "kernels/exception.cuh" - -#ifndef TORCH_EXTENSION_NAME -#define TORCH_EXTENSION_NAME deep_ep_cpp -#endif - -namespace shared_memory { - -union MemHandleInner { - cudaIpcMemHandle_t cuda_ipc_mem_handle; - CUmemFabricHandle cu_mem_fabric_handle; -}; - -struct MemHandle { - MemHandleInner inner; - size_t size; -}; - -constexpr size_t HANDLE_SIZE = sizeof(MemHandle); - -class SharedMemoryAllocator { -public: - SharedMemoryAllocator(bool use_fabric); - void malloc(void** ptr, size_t size); - void free(void* ptr); - void get_mem_handle(MemHandle* mem_handle, void* ptr); - void open_mem_handle(void** ptr, MemHandle* mem_handle); - void close_mem_handle(void* ptr); - -private: - bool use_fabric; -}; -} // namespace shared_memory - -namespace deep_ep { - -struct Buffer { - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS == 8, "The number of maximum NVLink peers must be 8"); - -private: - // Low-latency mode buffer - int low_latency_buffer_idx = 0; - bool low_latency_mode = false; - - // NVLink Buffer - int64_t num_nvl_bytes; - void* buffer_ptrs[NUM_MAX_NVL_PEERS] = {nullptr}; - void** buffer_ptrs_gpu = nullptr; - - // NVSHMEM Buffer - int64_t num_rdma_bytes; - void* rdma_buffer_ptr = nullptr; - - // Shrink mode buffer - bool enable_shrink = false; - int* mask_buffer_ptr = nullptr; - int* sync_buffer_ptr = nullptr; - - // Device info and communication - int device_id; - int num_device_sms; - int rank, rdma_rank, nvl_rank; - int num_ranks, num_rdma_ranks, num_nvl_ranks; - shared_memory::MemHandle ipc_handles[NUM_MAX_NVL_PEERS]; - - // Stream for communication - at::cuda::CUDAStream comm_stream; - - // After IPC/NVSHMEM synchronization, this flag will be true - bool available = false; - - // Whether explicit `destroy()` is required. - bool explicitly_destroy; - // After `destroy()` be called, this flag will be true - bool destroyed = false; - - // Barrier signals - int* barrier_signal_ptrs[NUM_MAX_NVL_PEERS] = {nullptr}; - int** barrier_signal_ptrs_gpu = nullptr; - - // Workspace - void* workspace = nullptr; - - // Host-side MoE info - volatile int* moe_recv_counter = nullptr; - int* moe_recv_counter_mapped = nullptr; - - // Host-side expert-level MoE info - volatile int* moe_recv_expert_counter = nullptr; - int* moe_recv_expert_counter_mapped = nullptr; - - // Host-side RDMA-level MoE info - volatile int* moe_recv_rdma_counter = nullptr; - int* moe_recv_rdma_counter_mapped = nullptr; - - shared_memory::SharedMemoryAllocator shared_memory_allocator; - -public: - Buffer(int rank, - int num_ranks, - int64_t num_nvl_bytes, - int64_t num_rdma_bytes, - bool low_latency_mode, - bool explicitly_destroy, - bool enable_shrink, - bool use_fabric); - - ~Buffer() noexcept(false); - - bool is_available() const; - - bool is_internode_available() const; - - int get_num_rdma_ranks() const; - - int get_rdma_rank() const; - - int get_root_rdma_rank(bool global) const; - - int get_local_device_id() const; - - pybind11::bytearray get_local_ipc_handle() const; - - pybind11::bytearray get_local_nvshmem_unique_id() const; - - torch::Tensor get_local_buffer_tensor(const pybind11::object& dtype, int64_t offset, bool use_rdma_buffer) const; - - torch::Stream get_comm_stream() const; - - void sync(const std::vector& device_ids, - const std::vector>& all_gathered_handles, - const std::optional& root_unique_id_opt); - - void destroy(); - - std::tuple, torch::Tensor, torch::Tensor, std::optional> get_dispatch_layout( - const torch::Tensor& topk_idx, - int num_experts, - std::optional& previous_event, - bool async, - bool allocate_on_comm_stream); - - std::tuple, - std::optional, - std::optional, - std::vector, - torch::Tensor, - torch::Tensor, - torch::Tensor, - torch::Tensor, - torch::Tensor, - std::optional> - intranode_dispatch(const torch::Tensor& x, - const std::optional& x_scales, - const std::optional& topk_idx, - const std::optional& topk_weights, - const std::optional& num_tokens_per_rank, - const torch::Tensor& is_token_in_rank, - const std::optional& num_tokens_per_expert, - int cached_num_recv_tokens, - const std::optional& cached_rank_prefix_matrix, - const std::optional& cached_channel_prefix_matrix, - int expert_alignment, - int num_worst_tokens, - const Config& config, - std::optional& previous_event, - bool async, - bool allocate_on_comm_stream); - - std::tuple, std::optional> intranode_combine( - const torch::Tensor& x, - const std::optional& topk_weights, - const std::optional& bias_0, - const std::optional& bias_1, - const torch::Tensor& src_idx, - const torch::Tensor& rank_prefix_matrix, - const torch::Tensor& channel_prefix_matrix, - const torch::Tensor& send_head, - const Config& config, - std::optional& previous_event, - bool async, - bool allocate_on_comm_stream); - - std::tuple, - std::optional, - std::optional, - std::vector, - torch::Tensor, - torch::Tensor, - std::optional, - torch::Tensor, - std::optional, - torch::Tensor, - std::optional, - std::optional, - std::optional, - std::optional> - internode_dispatch(const torch::Tensor& x, - const std::optional& x_scales, - const std::optional& topk_idx, - const std::optional& topk_weights, - const std::optional& num_tokens_per_rank, - const std::optional& num_tokens_per_rdma_rank, - const torch::Tensor& is_token_in_rank, - const std::optional& num_tokens_per_expert, - int cached_num_recv_tokens, - int cached_num_rdma_recv_tokens, - const std::optional& cached_rdma_channel_prefix_matrix, - const std::optional& cached_recv_rdma_rank_prefix_sum, - const std::optional& cached_gbl_channel_prefix_matrix, - const std::optional& cached_recv_gbl_rank_prefix_sum, - int expert_alignment, - int num_worst_tokens, - const Config& config, - std::optional& previous_event, - bool async, - bool allocate_on_comm_stream); - - std::tuple, std::optional> internode_combine( - const torch::Tensor& x, - const std::optional& topk_weights, - const std::optional& bias_0, - const std::optional& bias_1, - const torch::Tensor& src_meta, - const torch::Tensor& is_combined_token_in_rank, - const torch::Tensor& rdma_channel_prefix_matrix, - const torch::Tensor& rdma_rank_prefix_sum, - const torch::Tensor& gbl_channel_prefix_matrix, - const torch::Tensor& combined_rdma_head, - const torch::Tensor& combined_nvl_head, - const Config& config, - std::optional& previous_event, - bool async, - bool allocate_on_comm_stream); - - void clean_low_latency_buffer(int num_max_dispatch_tokens_per_rank, int hidden, int num_experts); - - std::tuple, - torch::Tensor, - torch::Tensor, - torch::Tensor, - std::optional, - std::optional>> - low_latency_dispatch(const torch::Tensor& x, - const torch::Tensor& topk_idx, - const std::optional& cumulative_local_expert_recv_stats, - const std::optional& dispatch_wait_recv_cost_stats, - int num_max_dispatch_tokens_per_rank, - int num_experts, - bool use_fp8, - bool round_scale, - bool use_ue8m0, - bool async, - bool return_recv_hook); - - std::tuple, std::optional>> low_latency_combine( - const torch::Tensor& x, - const torch::Tensor& topk_idx, - const torch::Tensor& topk_weights, - const torch::Tensor& src_info, - const torch::Tensor& layout_range, - const std::optional& combine_wait_recv_cost_stats, - int num_max_dispatch_tokens_per_rank, - int num_experts, - bool use_logfmt, - bool zero_copy, - bool async, - bool return_recv_hook, - const std::optional& out = std::nullopt); - - torch::Tensor get_next_low_latency_combine_buffer(int num_max_dispatch_tokens_per_rank, int hidden, int num_experts) const; - - void low_latency_update_mask_buffer(int rank_to_mask, bool mask); - - void low_latency_query_mask_buffer(const torch::Tensor& mask_status); - - void low_latency_clean_mask_buffer(); -}; - -} // namespace deep_ep diff --git a/csrc/elastic/buffer.hpp b/csrc/elastic/buffer.hpp new file mode 100644 index 000000000..f52296079 --- /dev/null +++ b/csrc/elastic/buffer.hpp @@ -0,0 +1,1307 @@ +#pragma once + +#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 registered for both scaleout and scaleup + int64_t num_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 deterministic algorithms + bool deterministic; + + // 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; + int64_t num_engram_storage_bytes = 0; + int64_t num_engram_recv_bytes = 0; + + // 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 int64_t& num_buffer_bytes, + const bool& deterministic, + 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), + explicitly_destroy(explicitly_destroy), + comm_stream(get_global_comm_stream()), + deterministic(deterministic), + allow_hybrid_mode(allow_hybrid_mode), + allow_multiple_reduction(allow_multiple_reduction), + prefer_overlap_with_compute(prefer_overlap_with_compute) { + // Init NCCL runtime + static constexpr int kBufferAlignment = 16; + this->nccl_context = std::make_shared( + nccl_comm, num_ranks, rank_idx, + layout::WorkspaceLayout::get_num_bytes() + num_buffer_bytes, kBufferAlignment, + allow_hybrid_mode, sl_idx, num_allocated_qps); + + // 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) + layout::WorkspaceLayout::get_num_bytes(); + CUDA_RUNTIME_CHECK(cudaMemset(workspace, 0, layout::WorkspaceLayout::get_num_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 { + 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, + 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) { + // 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); + EP_HOST_ASSERT(storage.is_cuda() and storage.is_contiguous()); + num_engram_entries = num_entries, engram_hidden = hidden; + + // Write storage and ensure the received buffer is aligned + EP_HOST_ASSERT(storage.nbytes() <= num_buffer_bytes); + CUDA_RUNTIME_CHECK(cudaMemcpyAsync( + buffer, storage.data_ptr(), storage.nbytes(), + cudaMemcpyDeviceToDevice, compute_stream)); + num_engram_storage_bytes = math::align(storage.nbytes(), 32); + num_engram_recv_bytes = num_buffer_bytes - num_engram_storage_bytes; + + // Ensure data is visible for all ranks + barrier(false, true); + } + + std::function engram_fetch(const torch::Tensor& indices, int num_qps) const { + const auto [num_tokens] = get_shape<1>(indices); + EP_HOST_ASSERT(indices.scalar_type() == torch::kInt); + EP_HOST_ASSERT(indices.is_cuda() and indices.is_contiguous()); + EP_HOST_ASSERT(num_tokens * engram_hidden * sizeof(nv_bfloat16) <= num_engram_recv_bytes); + + // Calculate a QP count + if (num_qps == 0) + num_qps = nccl_context->num_allocated_qps; + + // Return tensor from the raw buffer + EP_HOST_ASSERT(num_engram_entries > 0); + const auto fetched = torch::from_blob( + math::advance_ptr(buffer, num_engram_storage_bytes), + {num_tokens, engram_hidden}, + torch::TensorOptions().dtype(torch::kBFloat16).device(torch::kCUDA) + ); + + // Last issued Gin requests + const auto last_gin_requests = torch::empty( + {nccl_context->num_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, + buffer, + fetched.data_ptr(), + indices.data_ptr(), + static_cast(last_gin_requests.data_ptr()), + num_engram_entries, engram_hidden, + num_tokens, + nccl_context->num_ranks, num_qps, + at::cuda::getCurrentCUDAStream() + ); + + return [=, this]() { + // 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_ranks, num_qps, + at::cuda::getCurrentCUDAStream() + ); + return fetched; + }; + } + + 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(num_bytes % 32 == 0); + 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 += num_bytes * nccl_context->num_ranks; + } + 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()); + EP_HOST_ASSERT(x.is_cuda() and x.nbytes() % 32 == 0); + + 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 += x.nbytes() * nccl_context->num_ranks; + 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 + return std::max(num_dispatch_bytes, num_combine_bytes); + } + + std::tuple, + std::optional, std::optional, + std::optional, + std::vector, + 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_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_dst_buffer_slot_idx, + const std::optional& cached_token_metadata_at_forward, + 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& use_tma_aligned_col_major_sf) const { + // Check SM count + EP_HOST_ASSERT(num_sms > 0); + + // 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_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 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 + std::optional deterministic_rank_count_buffer = std::nullopt; + 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 if (deterministic) { + const auto prologue_num_sms = jit::device_runtime->get_num_sms(); + + // Allocate new tensors + deterministic_rank_count_buffer = torch::empty( + {prologue_num_sms, nccl_context->num_scaleup_ranks}, + torch::TensorOptions(torch::kCUDA).dtype(torch::kInt)); + dst_buffer_slot_idx = torch::empty( + {num_tokens, num_topk}, torch::TensorOptions(torch::kCUDA).dtype(torch::kInt)); + + // Launch a kernel to preprocess the destination slot indices + launch_dispatch_deterministic_prologue(topk_idx.data_ptr(), + deterministic_rank_count_buffer->data_ptr(), + dst_buffer_slot_idx.data_ptr(), + num_tokens, num_max_tokens_per_rank, + num_experts, num_topk, + nccl_context->scaleup_rank_idx, nccl_context->num_scaleup_ranks, + prologue_num_sms, + jit::device_runtime->get_num_smem_bytes(), + comm_stream); + } 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) { + EP_HOST_ASSERT(not deterministic); + + // 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(), + 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, deterministic, 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 + // TODO: support to expand for MoE training backward with cached handles from non-expanding forward, + // which requires maintaining the same expanding order between forward and backward + EP_HOST_ASSERT(not do_expand and "Cannot do expand with 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(); + } 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 = 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 (do_expand) { + // Slice and exclusive part and do atomic additions into inclusive + EP_HOST_ASSERT(not cached_mode); + psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert.slice(0, 0, num_local_experts); + } else if (not cached_mode) { + // 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_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, + 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, + 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, + recv_src_metadata, + deterministic_rank_count_buffer, + 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_per_expert_list, + psum_num_recv_tokens_per_scaleup_rank, + psum_num_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 (use_expanded_layout) { + // Reduction should be done with SwiGLU + EP_HOST_ASSERT(not topk_weights.has_value()); + } else if (topk_weights.has_value()) { + 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("calculate_elastic_buffer_size", &ElasticBuffer::calculate_buffer_size); + + // 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 000000000..f919b5729 --- /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 000000000..1bf6fc05b --- /dev/null +++ b/csrc/indexing/main.cu @@ -0,0 +1,21 @@ +// Utils +#include + +// EP +#include +#include +#include +#include +#include +#include +#include + +// Engram +#include + +// PP +#include + +int main() { + return 0; +} diff --git a/csrc/jit/api.hpp b/csrc/jit/api.hpp new file mode 100644 index 000000000..8332d9141 --- /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 000000000..96d89e501 --- /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 000000000..ad01b3cac --- /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 000000000..3ed532871 --- /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 000000000..c7387a017 --- /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 000000000..6df87ff37 --- /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 000000000..e6109590a --- /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 000000000..c01c75a7a --- /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 index 22e34a38c..c48037170 100644 --- a/csrc/kernels/CMakeLists.txt +++ b/csrc/kernels/CMakeLists.txt @@ -8,14 +8,8 @@ function(add_deep_ep_library target_name source_file) CUDA_STANDARD 17 CUDA_SEPARABLE_COMPILATION ON ) - target_link_libraries(${target_name} PUBLIC nvshmem cudart cudadevrt mlx5) + target_link_libraries(${target_name} PUBLIC nvshmem_host nvshmem_device cudart cudadevrt mlx5) endfunction() -add_deep_ep_library(runtime_cuda runtime.cu) -add_deep_ep_library(layout_cuda layout.cu) -add_deep_ep_library(intranode_cuda intranode.cu) -add_deep_ep_library(internode_cuda internode.cu) -add_deep_ep_library(internode_ll_cuda internode_ll.cu) - -# Later, we should link all libraries in `EP_CUDA_LIBRARIES` -set(EP_CUDA_LIBRARIES runtime_cuda layout_cuda intranode_cuda internode_cuda internode_ll_cuda PARENT_SCOPE) +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 000000000..1cd71b80f --- /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 000000000..de00668fc --- /dev/null +++ b/csrc/kernels/backend/api.cuh @@ -0,0 +1,97 @@ +#pragma once + +#include +#include +#include + +#include +#include + +// 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; + +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; + + NCCLSymmetricMemoryContext(const int64_t& nccl_comm, + const int& num_ranks, const int& rank_idx, + const size_t& size, const size_t& alignment, + 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() const; +}; + +} // 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 000000000..9b6fc99c1 --- /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 000000000..baf38c7d9 --- /dev/null +++ b/csrc/kernels/backend/nccl.cu @@ -0,0 +1,154 @@ +#include +#include +#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 int& num_ranks, const int& rank_idx, + const size_t& size, const size_t& alignment, + 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); + + // 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."); + + // Initialize NCCL device communicator + ncclDevCommRequirements_t reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER; + if (num_ranks > 1 and get_env("EP_DISABLE_GIN", 0) == 0) { + reqs.ginContextCount = num_allocated_qps; + reqs.ginExclusiveContexts = true; + reqs.ginQueueDepth = 1024; + 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 window + // NOTES: `ncclCommWindowRegister` is collective: it internally calls bootstrapBarrier + // across all ranks, so no explicit barrier is needed after this call. + NCCL_CHECK(ncclMemAlloc(&raw_window_ptr, size)); + NCCL_CHECK(ncclCommWindowRegister(comm, raw_window_ptr, size, &window, NCCL_WIN_DEFAULT)); + 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])); + + // TODO: push NCCL team to support aligned allocation + EP_HOST_ASSERT(size % alignment == 0); + EP_HOST_ASSERT(reinterpret_cast(raw_window_ptr) % alignment == 0); + EP_HOST_ASSERT(reinterpret_cast(mapped_window_ptr) % alignment == 0); +} + +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() const { + // Deregister window and free buffer + NCCL_CHECK(ncclCommWindowDeregister(comm, window)); + NCCL_CHECK(ncclMemFree(raw_window_ptr)); + + // 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 000000000..837d8fe22 --- /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/configs.cuh b/csrc/kernels/configs.cuh deleted file mode 100644 index 6485967cf..000000000 --- a/csrc/kernels/configs.cuh +++ /dev/null @@ -1,90 +0,0 @@ -#pragma once - -#define NUM_MAX_NVL_PEERS 8 -#define NUM_MAX_RDMA_PEERS 20 -#define NUM_WORKSPACE_BYTES (32 * 1024 * 1024) -#define NUM_MAX_LOCAL_EXPERTS 1024 -#define NUM_BUFFER_ALIGNMENT_BYTES 128 - -#define FINISHED_SUM_TAG 1024 -#define NUM_WAIT_NANOSECONDS 500 - -#ifndef ENABLE_FAST_DEBUG -#define NUM_CPU_TIMEOUT_SECS 100 -#define NUM_TIMEOUT_CYCLES 200000000000ull // 200G cycles ~= 100s -#else -#define NUM_CPU_TIMEOUT_SECS 10 -#define NUM_TIMEOUT_CYCLES 20000000000ull // 20G cycles ~= 10s -#endif - -#define LOW_LATENCY_SEND_PHASE 1 -#define LOW_LATENCY_RECV_PHASE 2 - -// Make CLion CUDA indexing work -#ifdef __CLION_IDE__ -#define __CUDA_ARCH__ 900 // NOLINT(*-reserved-identifier) -#define __CUDACC_RDC__ // NOLINT(*-reserved-identifier) -#endif - -// Define __CUDACC_RDC__ to ensure proper extern declarations for NVSHMEM device symbols -#ifndef DISABLE_NVSHMEM -#ifndef __CUDACC_RDC__ -#define __CUDACC_RDC__ // NOLINT(*-reserved-identifier) -#endif -#endif - -// Remove Torch restrictions -#ifdef __CUDA_NO_HALF_CONVERSIONS__ -#undef __CUDA_NO_HALF_CONVERSIONS__ -#endif -#ifdef __CUDA_NO_HALF_OPERATORS__ -#undef __CUDA_NO_HALF_OPERATORS__ -#endif -#ifdef __CUDA_NO_HALF2_OPERATORS__ -#undef __CUDA_NO_HALF2_OPERATORS__ -#endif -#ifdef __CUDA_NO_BFLOAT16_CONVERSIONS__ -#undef __CUDA_NO_BFLOAT16_CONVERSIONS__ -#endif -#ifdef __CUDA_NO_BFLOAT162_OPERATORS__ -#undef __CUDA_NO_BFLOAT162_OPERATORS__ -#endif - -#include -#include - -#include - -#ifndef DISABLE_SM90_FEATURES -#include -#else -// Ampere does not support FP8 features -#define __NV_E4M3 0 -#define __NV_E5M2 1 -typedef int __nv_fp8_interpretation_t; -typedef int __nv_fp8x4_e4m3; -typedef uint8_t __nv_fp8_storage_t; -#endif - -namespace deep_ep { - -#ifndef TOPK_IDX_BITS -#define TOPK_IDX_BITS 64 -#endif - -#define INT_BITS_T2(bits) int##bits##_t -#define INT_BITS_T(bits) INT_BITS_T2(bits) -typedef INT_BITS_T(TOPK_IDX_BITS) topk_idx_t; // int32_t or int64_t -#undef INT_BITS_T -#undef INT_BITS_T2 - -} // namespace deep_ep - -#ifndef DISABLE_NVSHMEM -#include -#include -#include -#include - -#include -#endif diff --git a/csrc/kernels/elastic/api.hpp b/csrc/kernels/elastic/api.hpp new file mode 100644 index 000000000..29023dbd1 --- /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 000000000..1a495e1fa --- /dev/null +++ b/csrc/kernels/elastic/barrier.hpp @@ -0,0 +1,82 @@ +#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; + + // 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); + } + + 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 at::cuda::CUDAStream& stream) { + // Number of threads equals to the number of ranks + constexpr auto kNumThreads = 512; + + // Generate, build and launch + // NOTES: only the hybrid kernel needs 2 SMs + const auto num_sms = 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, + .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 000000000..3764da926 --- /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 000000000..ddac5db8d --- /dev/null +++ b/csrc/kernels/elastic/dispatch.hpp @@ -0,0 +1,414 @@ +#pragma once + +#include +#include + +#include +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/launch_runtime.hpp" + +namespace deep_ep::elastic { + +class DispatchPrologueRuntime final : public jit::LaunchRuntime { +public: + struct Args { + // Templated arguments + int num_warps; + int num_ranks; + int num_max_tokens_per_rank; + int num_experts, num_topk; + + // Parameters + topk_idx_t* topk_idx; + int* rank_count_buffer; + int* dst_buffer_slot_idx; + int num_tokens; + int 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_deterministic_prologue_impl<{}, {}, {}, {}, {}, {}>); +}} +)", + args.launch_args.grid_dim.first, + args.num_warps, + args.num_ranks, + 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.topk_idx, + args.rank_count_buffer, + args.dst_buffer_slot_idx, + args.num_tokens, + args.rank_idx)); + } +}; + +static void launch_dispatch_deterministic_prologue(topk_idx_t* topk_idx, int* rank_count_buffer, + int* dst_buffer_slot_idx, + const int& num_tokens, const int& num_max_tokens_per_rank, + const int& num_experts, const int& num_topk, + const int& rank_idx, const int& num_ranks, + const int& num_sms, const int& num_smem_bytes, + const at::cuda::CUDAStream& stream) { + constexpr auto num_warps = 8; + constexpr auto num_threads = num_warps * 32; + EP_HOST_ASSERT((2 * num_warps + 1) * num_ranks * sizeof(int) <= num_smem_bytes and + "Insufficient shared memory"); + + // Generate, build and launch + const DispatchPrologueRuntime::Args args = { + .num_warps = num_warps, + .num_ranks = num_ranks, + .num_max_tokens_per_rank = num_max_tokens_per_rank, + .num_experts = num_experts, .num_topk = num_topk, + .topk_idx = topk_idx, + .rank_count_buffer = rank_count_buffer, + .dst_buffer_slot_idx = dst_buffer_slot_idx, + .num_tokens = num_tokens, + .rank_idx = rank_idx, + .launch_args = jit::LaunchArgs(num_sms, num_threads, num_smem_bytes, 1, true)}; + const auto code = DispatchPrologueRuntime::generate(args); + const auto runtime = jit::compiler->build("dispatch_deterministic_prologue", code); + DispatchPrologueRuntime::launch(runtime, args, stream); +} + +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* 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.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.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* 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& deterministic, + 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 or deterministic; + 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) { + // Too many warps may cause performance degrade, so we limit the total warps into 512 + const auto token_layout = get_dispatch_token_layout(hidden, elem_size, num_sf_packs, num_topk); + num_dispatch_warps = std::min(std::min( + (num_smem_bytes - num_notify_smem_bytes) / token_layout.get_num_bytes(), 32 - num_notify_warps), + math::ceil_div(512, num_sms)); + num_threads = (num_notify_warps + num_dispatch_warps) * 32; + } else { + // Hybrid kernels + // Some unimplemented assertions + EP_HOST_ASSERT(not deterministic); + + 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, + .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; + 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; + + // 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_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.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); + } + + 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_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, + 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& 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 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, + .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, + .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_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 000000000..d588248a4 --- /dev/null +++ b/csrc/kernels/elastic/engram.hpp @@ -0,0 +1,143 @@ +#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 hidden; + int num_ranks; + int num_qps; + + // Parameters + ncclDevComm_t nccl_dev_comm; + ncclWindow_t nccl_window; + void* storage; + void* fetched; + int* indices; + ncclGinRequest_t* last_gin_requests; + int num_tokens; + + 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(&engram_fetch_impl<{}, {}, {}, {}, {}>); +}} +)", args.num_qps, + args.num_entries_per_rank, args.hidden, + args.num_ranks, args.launch_args.num_threads); + } + + 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.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, + const int& num_entries_per_rank, const int& hidden, + const int& num_tokens, + const int& num_ranks, const int& num_qps, + 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, + .hidden = hidden, + .num_ranks = num_ranks, + .num_qps = num_qps, + .nccl_dev_comm = nccl_dev_comm, + .nccl_window = nccl_window, + .storage = storage, + .fetched = fetched, + .indices = indices, + .last_gin_requests = last_gin_requests, + .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_ranks; + + 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) { + return fmt::format(R"( +#include + +using namespace deep_ep::elastic; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&engram_fetch_wait_impl<{}, {}>); +}} +)", args.num_ranks, args.launch_args.num_threads); + } + + 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_ranks, const int& num_qps, + const at::cuda::CUDAStream& stream) { + constexpr int kNumEngramFetchWaitThreads = 1024; + + // Generate, build and launch + const EngramFetchWaitRuntime::Args args = { + .num_ranks = num_ranks, + .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 000000000..06cbea439 --- /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/exception.cuh b/csrc/kernels/exception.cuh deleted file mode 100644 index 507efa285..000000000 --- a/csrc/kernels/exception.cuh +++ /dev/null @@ -1,63 +0,0 @@ -#pragma once - -#include -#include - -#include "configs.cuh" - -#ifndef EP_STATIC_ASSERT -#define EP_STATIC_ASSERT(cond, reason) static_assert(cond, reason) -#endif - -class EPException : public std::exception { -private: - std::string message = {}; - -public: - explicit EPException(const char* name, const char* file, const int line, const std::string& error) { - message = std::string("Failed: ") + name + " error " + file + ":" + std::to_string(line) + " '" + error + "'"; - } - - const char* what() const noexcept override { return message.c_str(); } -}; - -#ifndef CUDA_CHECK -#define CUDA_CHECK(cmd) \ - do { \ - cudaError_t e = (cmd); \ - if (e != cudaSuccess) { \ - throw EPException("CUDA", __FILE__, __LINE__, cudaGetErrorString(e)); \ - } \ - } while (0) -#endif - -#ifndef CU_CHECK -#define CU_CHECK(cmd) \ - do { \ - CUresult e = (cmd); \ - if (e != CUDA_SUCCESS) { \ - const char* error_str = NULL; \ - cuGetErrorString(e, &error_str); \ - throw EPException("CU", __FILE__, __LINE__, std::string(error_str)); \ - } \ - } while (0) -#endif - -#ifndef EP_HOST_ASSERT -#define EP_HOST_ASSERT(cond) \ - do { \ - if (not(cond)) { \ - throw EPException("Assertion", __FILE__, __LINE__, #cond); \ - } \ - } while (0) -#endif - -#ifndef EP_DEVICE_ASSERT -#define EP_DEVICE_ASSERT(cond) \ - do { \ - if (not(cond)) { \ - printf("Assertion failed: %s:%d, condition: %s\n", __FILE__, __LINE__, #cond); \ - asm("trap;"); \ - } \ - } while (0) -#endif diff --git a/csrc/kernels/legacy/CMakeLists.txt b/csrc/kernels/legacy/CMakeLists.txt new file mode 100644 index 000000000..29b20b1e9 --- /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/api.cuh b/csrc/kernels/legacy/api.cuh similarity index 96% rename from csrc/kernels/api.cuh rename to csrc/kernels/legacy/api.cuh index 9bbe096a8..641028d99 100644 --- a/csrc/kernels/api.cuh +++ b/csrc/kernels/legacy/api.cuh @@ -1,34 +1,8 @@ #pragma once -#include +#include "compiled.cuh" -#include "configs.cuh" - -namespace deep_ep { - -// Intranode runtime -namespace intranode { - -void barrier(int** barrier_signal_ptrs, int rank, int num_ranks, cudaStream_t stream); - -} // namespace intranode - -// Internode runtime -namespace internode { - -std::vector get_unique_id(); - -int init(const std::vector& root_unique_id_val, int rank, int num_ranks, bool low_latency_mode); - -void* alloc(size_t size, size_t alignment); - -void free(void* ptr); - -void barrier(); - -void finalize(); - -} // namespace internode +namespace deep_ep::legacy { // Layout kernels namespace layout { @@ -49,6 +23,8 @@ void get_dispatch_layout(const topk_idx_t* topk_idx, // 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, @@ -347,4 +323,4 @@ void clean_mask_buffer(int* mask_buffer_ptr, int num_ranks, cudaStream_t stream) } // namespace internode_ll -} // namespace deep_ep +} // namespace deep_ep::legacy diff --git a/csrc/kernels/buffer.cuh b/csrc/kernels/legacy/buffer.cuh similarity index 97% rename from csrc/kernels/buffer.cuh rename to csrc/kernels/legacy/buffer.cuh index 222f42ac4..6b9942d1b 100644 --- a/csrc/kernels/buffer.cuh +++ b/csrc/kernels/legacy/buffer.cuh @@ -1,9 +1,10 @@ #pragma once -#include "configs.cuh" -#include "exception.cuh" +#include -namespace deep_ep { +#include "compiled.cuh" + +namespace deep_ep::legacy { template struct Buffer { @@ -128,4 +129,4 @@ public: } }; -} // namespace deep_ep +} // namespace deep_ep::legacy diff --git a/csrc/kernels/legacy/compiled.cuh b/csrc/kernels/legacy/compiled.cuh new file mode 100644 index 000000000..1f59d6d53 --- /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/ibgda_device.cuh b/csrc/kernels/legacy/ibgda_device.cuh similarity index 96% rename from csrc/kernels/ibgda_device.cuh rename to csrc/kernels/legacy/ibgda_device.cuh index 421ec2a85..819bdd1c8 100644 --- a/csrc/kernels/ibgda_device.cuh +++ b/csrc/kernels/legacy/ibgda_device.cuh @@ -7,13 +7,16 @@ // - nvshmem/src/include/non_abi/device/pt-to-pt/ibgda_device.cuh #pragma once -#include -#include "configs.cuh" -#include "exception.cuh" +#include +#include +#include + +#include + #include "utils.cuh" -namespace deep_ep { +namespace deep_ep::legacy { EP_STATIC_ASSERT(NVSHMEMI_IBGDA_MIN_QP_DEPTH >= 64, "Invalid QP minimum depth"); @@ -75,24 +78,11 @@ __device__ static __forceinline__ nvshmemi_ibgda_device_state_t* ibgda_get_state return &nvshmemi_ibgda_device_state_d; } -// Template helper to get RC - uses compile-time type checking with if constexpr (C++17) -template -__device__ static __forceinline__ nvshmemi_ibgda_device_qp_t* ibgda_get_rc_impl(StateType* state, int pe, int id) { - const auto num_rc_per_pe = state->num_rc_per_pe; - - if constexpr (std::is_same_v) { - // v1 implementation - return &state->globalmem - .rcs[pe * num_rc_per_pe * state->num_devices_initialized + id % (num_rc_per_pe * state->num_devices_initialized)]; - } else { - // v2 implementation (or any other type) - return &state->globalmem.rcs[pe + nvshmemi_device_state_d.npes * id]; - } -} - __device__ static __forceinline__ nvshmemi_ibgda_device_qp_t* ibgda_get_rc(int pe, int id) { auto state = ibgda_get_state(); - return ibgda_get_rc_impl(state, pe, id); + 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) { @@ -503,4 +493,4 @@ __device__ static __forceinline__ void nvshmemi_ibgda_quiet(int dst_pe, int qp_i ibgda_poll_cq(qp->tx_wq.cq, prod_idx); } -} // namespace deep_ep +} // namespace deep_ep::legacy diff --git a/csrc/kernels/internode.cu b/csrc/kernels/legacy/internode.cu similarity index 91% rename from csrc/kernels/internode.cu rename to csrc/kernels/legacy/internode.cu index 48c6c0018..adc26296b 100644 --- a/csrc/kernels/internode.cu +++ b/csrc/kernels/legacy/internode.cu @@ -2,22 +2,27 @@ #include #include "buffer.cuh" -#include "configs.cuh" -#include "exception.cuh" +#include "compiled.cuh" #include "ibgda_device.cuh" #include "launch.cuh" #include "utils.cuh" namespace deep_ep { -namespace internode { +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(NUM_MAX_NVL_PEERS == 8, "Invalid number of maximum NVL peers"); + EP_STATIC_ASSERT(LEGACY_NUM_MAX_NVL_PEERS == 8, "Invalid number of maximum NVL peers"); __forceinline__ SourceMeta() = default; @@ -26,7 +31,7 @@ struct SourceMeta { src_rdma_rank = rdma_rank; is_token_in_nvl_rank_bits = is_token_in_nvl_ranks[0]; #pragma unroll - for (int i = 1; i < NUM_MAX_NVL_PEERS; ++i) + for (int i = 1; i < LEGACY_NUM_MAX_NVL_PEERS; ++i) is_token_in_nvl_rank_bits |= is_token_in_nvl_ranks[i] << i; } @@ -56,7 +61,7 @@ __host__ __device__ __forceinline__ std::pair get_rdma_clean_meta(int return {(get_num_bytes_per_token(hidden_int4, num_scales, num_topk_idx, num_topk_weights) * num_rdma_recv_buffer_tokens * num_rdma_ranks * 2 * num_channels) / sizeof(int), - (NUM_MAX_NVL_PEERS * 2 + 4) * num_rdma_ranks * 2 * num_channels}; + (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, @@ -81,7 +86,7 @@ __host__ __device__ __forceinline__ std::pair get_nvl_clean_meta(int h template __forceinline__ __device__ int translate_dst_rdma_rank(const int dst_rdma_rank, const int nvl_rank) { - return kLowLatencyMode ? (dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank) : dst_rdma_rank; + return kLowLatencyMode ? (dst_rdma_rank * LEGACY_NUM_MAX_NVL_PEERS + nvl_rank) : dst_rdma_rank; } template @@ -120,8 +125,8 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, auto thread_id = static_cast(threadIdx.x), warp_id = thread_id / 32, lane_id = get_lane_id(); auto num_threads = static_cast(blockDim.x), num_warps = num_threads / 32; - auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS; - auto num_rdma_experts = num_experts / kNumRDMARanks, num_nvl_experts = num_rdma_experts / NUM_MAX_NVL_PEERS; + 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 @@ -141,11 +146,11 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, if (thread_id == 32) nvshmem_sync_with_same_gpu_idx(rdma_team); - barrier_block(barrier_signal_ptrs, nvl_rank); + barrier_block(barrier_signal_ptrs, nvl_rank); // Send numbers of tokens per rank/expert to RDMA ranks auto rdma_buffer_ptr_int = static_cast(rdma_buffer_ptr); - auto rdma_recv_num_tokens_mixed = SymBuffer(rdma_buffer_ptr, NUM_MAX_NVL_PEERS + num_rdma_experts + 1, kNumRDMARanks); + 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)); @@ -156,13 +161,13 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, // Copy to send buffer #pragma unroll for (int i = thread_id; i < num_ranks; i += num_threads) - rdma_recv_num_tokens_mixed.send_buffer(i / NUM_MAX_NVL_PEERS)[i % NUM_MAX_NVL_PEERS] = num_tokens_per_rank[i]; + 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)[NUM_MAX_NVL_PEERS + i % num_rdma_experts] = + 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)[NUM_MAX_NVL_PEERS + num_rdma_experts] = num_tokens_per_rdma_rank[thread_id]; + 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 @@ -172,7 +177,7 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, 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)), - (NUM_MAX_NVL_PEERS + num_rdma_experts + 1) * sizeof(int), + (LEGACY_NUM_MAX_NVL_PEERS + num_rdma_experts + 1) * sizeof(int), translate_dst_rdma_rank(i, nvl_rank), 0, lane_id, @@ -180,7 +185,7 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, } else { UNROLLED_WARP_COPY(1, lane_id, - NUM_MAX_NVL_PEERS + num_rdma_experts + 1, + 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, @@ -200,13 +205,13 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, __syncthreads(); // NVL buffers - auto nvl_send_buffer = thread_id < NUM_MAX_NVL_PEERS ? buffer_ptrs[thread_id] : nullptr; + 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, NUM_MAX_NVL_PEERS); - auto nvl_send_num_tokens_per_expert = AsymBuffer(nvl_send_buffer, num_nvl_experts, NUM_MAX_NVL_PEERS); - auto nvl_recv_num_tokens_per_rank = AsymBuffer(nvl_recv_buffer, kNumRDMARanks, NUM_MAX_NVL_PEERS); - auto nvl_recv_num_tokens_per_expert = AsymBuffer(nvl_recv_buffer, num_nvl_experts, NUM_MAX_NVL_PEERS); + 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]); @@ -224,7 +229,7 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, int sum = 0; #pragma unroll for (int i = 0; i < kNumRDMARanks; ++i) - sum += rdma_recv_num_tokens_mixed.recv_buffer(i)[NUM_MAX_NVL_PEERS + thread_id]; + 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(); @@ -234,7 +239,7 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, int sum = 0; #pragma unroll for (int i = 0; i < kNumRDMARanks; ++i) { - sum += rdma_recv_num_tokens_mixed.recv_buffer(i)[NUM_MAX_NVL_PEERS + num_rdma_experts]; + 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) { @@ -245,8 +250,8 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, } // Send numbers of tokens per rank/expert to NVL ranks - EP_DEVICE_ASSERT(NUM_MAX_NVL_PEERS <= num_threads); - if (thread_id < NUM_MAX_NVL_PEERS) { + 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]; @@ -254,7 +259,7 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, 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); + barrier_block(barrier_signal_ptrs, nvl_rank); // Reduce the number of tokens per rank/expert EP_DEVICE_ASSERT(num_nvl_experts <= num_threads); @@ -262,7 +267,7 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, int sum = 0; #pragma unroll for (int i = 0; i < num_ranks; ++i) { - int src_rdma_rank = i / NUM_MAX_NVL_PEERS, src_nvl_rank = i % NUM_MAX_NVL_PEERS; + 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; } @@ -275,7 +280,7 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, if (thread_id < num_nvl_experts) { int sum = 0; #pragma unroll - for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) + 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) { @@ -288,7 +293,7 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, // Finally barrier if (thread_id == 32) nvshmem_sync_with_same_gpu_idx(rdma_team); - barrier_block(barrier_signal_ptrs, nvl_rank); + barrier_block(barrier_signal_ptrs, nvl_rank); } else { // Calculate meta data int dst_rdma_rank = sm_id - 1; @@ -297,14 +302,14 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, get_channel_task_range(num_tokens, num_channels, channel_id, token_start_idx, token_end_idx); // Iterate over tokens - int total_count = 0, per_nvl_rank_count[NUM_MAX_NVL_PEERS] = {0}; + 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(NUM_MAX_NVL_PEERS * sizeof(bool) == sizeof(uint64_t), "Invalid number of NVL peers"); + 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 * NUM_MAX_NVL_PEERS); + *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 < NUM_MAX_NVL_PEERS; ++j) + 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); } @@ -312,14 +317,14 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, // Warp reduce total_count = warp_reduce_sum(total_count); #pragma unroll - for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) + 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 < NUM_MAX_NVL_PEERS; ++i) - gbl_channel_prefix_matrix[(dst_rdma_rank * NUM_MAX_NVL_PEERS + i) * num_channels + channel_id] = per_nvl_rank_count[i]; + 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; } } @@ -333,9 +338,9 @@ __global__ void notify_dispatch(const int* num_tokens_per_rank, prefix_row[i] += prefix_row[i - 1]; } - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); - if (thread_id < NUM_MAX_NVL_PEERS) { - auto prefix_row = gbl_channel_prefix_matrix + (dst_rdma_rank * NUM_MAX_NVL_PEERS + thread_id) * num_channels; + 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]; @@ -403,12 +408,12 @@ void notify_dispatch(const int* num_tokens_per_rank, buffer_ptrs, \ barrier_signal_ptrs, \ rank, \ - cpu_rdma_team); \ + nvshmem::cpu_rdma_team); \ } \ break constexpr int kNumThreads = 512; - const auto num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; + const auto num_rdma_ranks = num_ranks / LEGACY_NUM_MAX_NVL_PEERS; // Get clean meta auto rdma_clean_meta = @@ -418,7 +423,7 @@ void notify_dispatch(const int* num_tokens_per_rank, num_topk, num_topk, num_rdma_ranks, - NUM_MAX_NVL_PEERS, + LEGACY_NUM_MAX_NVL_PEERS, num_max_nvl_chunked_recv_tokens, num_channels, true); @@ -444,7 +449,7 @@ template -__global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NVL_PEERS) * 32), 1) +__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, @@ -487,40 +492,40 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV const auto thread_id = static_cast(threadIdx.x), warp_id = thread_id / 32, lane_id = get_lane_id(); const auto num_channels = num_sms / 2, channel_id = sm_id / 2; const bool is_forwarder = sm_id % 2 == 0; - const auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS; + const auto 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 < NUM_MAX_NVL_PEERS) { - return {WarpRole::kRDMAAndNVLForwarder, (warp_id + channel_id) % NUM_MAX_NVL_PEERS}; + 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 - NUM_MAX_NVL_PEERS}; + 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) % NUM_MAX_NVL_PEERS}; + 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 + NUM_MAX_NVL_PEERS); + 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(NUM_MAX_NVL_PEERS * sizeof(bool) == sizeof(uint64_t), "Invalid number of NVL peers"); + 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, NUM_MAX_NVL_PEERS * 2 + 2, 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); @@ -539,20 +544,20 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV // Allocate buffers auto nvl_channel_x = AsymBuffer(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens * num_bytes_per_token, - NUM_MAX_NVL_PEERS, + 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, NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) + 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, NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) + 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, NUM_MAX_NVL_PEERS, channel_id, num_channels, ws_rr_rank).advance_also(ws_rr_buffer_ptr); + 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, NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank).advance_also(rs_wr_buffer_ptr); + 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 @@ -575,9 +580,9 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV __syncwarp(); // Forward warp synchronization - __shared__ volatile int forward_channel_head[NUM_MAX_NVL_PEERS][kNumRDMARanks]; - __shared__ volatile bool forward_channel_retired[NUM_MAX_NVL_PEERS]; - auto sync_forwarder_smem = []() { asm volatile("barrier.sync 1, %0;" ::"r"((NUM_MAX_NVL_PEERS + 1) * 32)); }; + __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 @@ -585,24 +590,24 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV get_channel_task_range(num_tokens, num_channels, channel_id, token_start_idx, token_end_idx); // Send number of tokens in this channel by `-value - 1` - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS * 2 + 2 <= 32, "Invalid number of NVL peers"); + 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 < NUM_MAX_NVL_PEERS) { + if (lane_id < LEGACY_NUM_MAX_NVL_PEERS) { dst_ptr[lane_id] = -(channel_id == 0 ? 0 - : gbl_channel_prefix_matrix[(dst_rdma_rank * NUM_MAX_NVL_PEERS + lane_id) * num_channels + channel_id - 1]) - + : gbl_channel_prefix_matrix[(dst_rdma_rank * LEGACY_NUM_MAX_NVL_PEERS + lane_id) * num_channels + channel_id - 1]) - 1; - } else if (lane_id < NUM_MAX_NVL_PEERS * 2) { + } else if (lane_id < LEGACY_NUM_MAX_NVL_PEERS * 2) { dst_ptr[lane_id] = - -gbl_channel_prefix_matrix[(dst_rdma_rank * NUM_MAX_NVL_PEERS + lane_id - NUM_MAX_NVL_PEERS) * num_channels + + -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 == NUM_MAX_NVL_PEERS * 2) { + } 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 == NUM_MAX_NVL_PEERS * 2 + 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(); @@ -611,7 +616,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV 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) * (NUM_MAX_NVL_PEERS * 2 + 2), + sizeof(int) * (LEGACY_NUM_MAX_NVL_PEERS * 2 + 2), translate_dst_rdma_rank(dst_rdma_rank, nvl_rank), channel_id, lane_id, @@ -629,7 +634,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV uint64_t is_token_in_rank_uint64 = 0; if (lane_id < kNumRDMARanks) { is_token_in_rank_uint64 = - __ldg(reinterpret_cast(is_token_in_rank + token_idx * num_ranks + lane_id * NUM_MAX_NVL_PEERS)); + __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(); @@ -645,7 +650,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV cached_rdma_channel_head = static_cast(ld_volatile_global(rdma_channel_head.buffer(lane_id))); // Timeout check - if (clock64() - start_time >= NUM_TIMEOUT_CYCLES) { + 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, @@ -776,7 +781,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV auto start_time = clock64(); while (__any_sync(0xffffffff, num_tokens_to_send > 0)) { // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) { + 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, @@ -852,9 +857,9 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV if (lane_id < kNumRDMARanks) { while (true) { auto meta_0 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + dst_nvl_rank); - auto meta_1 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + NUM_MAX_NVL_PEERS + dst_nvl_rank); - auto meta_2 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + NUM_MAX_NVL_PEERS * 2); - auto meta_3 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + NUM_MAX_NVL_PEERS * 2 + 1); + 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; @@ -874,7 +879,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV } // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { + 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", @@ -894,7 +899,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV __syncwarp(); // Shift cached head - send_nvl_head += src_rdma_channel_prefix * NUM_MAX_NVL_PEERS + dst_nvl_rank; + send_nvl_head += src_rdma_channel_prefix * LEGACY_NUM_MAX_NVL_PEERS + dst_nvl_rank; // Wait shared memory to be cleaned sync_forwarder_smem(); @@ -914,7 +919,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV cached_nvl_channel_head = __shfl_sync(0xffffffffu, ld_volatile_global(nvl_channel_head.buffer()), 0); // Timeout check - if (elect_one_sync() and clock64() - start_time > NUM_TIMEOUT_CYCLES) { + 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, @@ -939,7 +944,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV } // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) { + 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", @@ -968,7 +973,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV auto cached_head = is_in_dst_nvl_rank ? rdma_nvl_token_idx : -1; rdma_nvl_token_idx += is_in_dst_nvl_rank; if (not kCachedMode) - send_nvl_head[i * NUM_MAX_NVL_PEERS] = cached_head; + send_nvl_head[i * LEGACY_NUM_MAX_NVL_PEERS] = cached_head; } if (not is_in_dst_nvl_rank) continue; @@ -1020,11 +1025,11 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); // Clean shared memory - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); + EP_STATIC_ASSERT(LEGACY_NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); #pragma unroll - for (int i = lane_id; i < kNumRDMARanks * NUM_MAX_NVL_PEERS; i += 32) - forward_channel_head[i % NUM_MAX_NVL_PEERS][i / NUM_MAX_NVL_PEERS] = 0; - if (lane_id < NUM_MAX_NVL_PEERS) + 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(); @@ -1033,7 +1038,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV // Find minimum head int min_head = std::numeric_limits::max(); #pragma unroll - for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) + 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())) @@ -1051,7 +1056,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV } // Nanosleep and let other warps work - __nanosleep(NUM_WAIT_NANOSECONDS); + __nanosleep(LEGACY_NUM_WAIT_NANOSECONDS); } } else { // NVL consumers @@ -1061,8 +1066,8 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV const int local_expert_end = local_expert_begin + (num_experts / num_ranks); EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); - if (lane_id < kNumRDMARanks and lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank > 0) - total_offset = recv_gbl_rank_prefix_sum[lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank - 1]; + 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; @@ -1077,7 +1082,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV } // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { + 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, @@ -1094,7 +1099,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV // Save for combine usage if (lane_id < kNumRDMARanks and not kCachedMode) - recv_gbl_channel_prefix_matrix[(lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank) * num_channels + channel_id] = total_offset; + 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; @@ -1108,7 +1113,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV cached_channel_tail_idx = __shfl_sync(0xffffffff, ld_acquire_sys_global(nvl_channel_tail.buffer()), 0); // Timeout check - if (elect_one_sync() and clock64() - start_time > NUM_TIMEOUT_CYCLES) { + 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, @@ -1247,7 +1252,7 @@ void dispatch(void* recv_x, bool low_latency_mode) { constexpr int kNumDispatchRDMASenderWarps = 7; constexpr int kNumTMABytesPerWarp = 16384; - constexpr int smem_size = kNumTMABytesPerWarp * NUM_MAX_NVL_PEERS; + 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()); @@ -1302,7 +1307,7 @@ void dispatch(void* recv_x, EP_HOST_ASSERT((topk_idx == nullptr) == (topk_weights == nullptr)); EP_HOST_ASSERT((recv_topk_idx == nullptr) == (recv_topk_weights == nullptr)); - SETUP_LAUNCH_CONFIG(num_channels * 2, (kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NVL_PEERS) * 32, stream); + 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 } @@ -1332,9 +1337,9 @@ __global__ void cached_notify(const int rdma_clean_offset, auto warp_id = thread_id / 32; auto lane_id = get_lane_id(); - auto nvl_rank = rank % NUM_MAX_NVL_PEERS; - auto num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; - auto rdma_rank = rank / NUM_MAX_NVL_PEERS; + 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) { @@ -1351,7 +1356,7 @@ __global__ void cached_notify(const int rdma_clean_offset, nvshmem_sync_with_same_gpu_idx(rdma_team); // Barrier for NVL - barrier_block(barrier_signal_ptrs, nvl_rank); + barrier_block(barrier_signal_ptrs, nvl_rank); // Clean RDMA buffer auto rdma_buffer_ptr_int = static_cast(rdma_buffer_ptr); @@ -1369,7 +1374,7 @@ __global__ void cached_notify(const int rdma_clean_offset, // Barrier again if (thread_id == 32) nvshmem_sync_with_same_gpu_idx(rdma_team); - barrier_block(barrier_signal_ptrs, nvl_rank); + barrier_block(barrier_signal_ptrs, nvl_rank); } else if (sm_id == 1) { if (is_cached_dispatch) return; @@ -1399,11 +1404,11 @@ __global__ void cached_notify(const int rdma_clean_offset, EP_DEVICE_ASSERT(num_warps >= num_channels); EP_DEVICE_ASSERT(rdma_channel_prefix_matrix != nullptr and rdma_rank_prefix_sum != nullptr); - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 32, "Too many NVL peers"); + 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) * NUM_MAX_NVL_PEERS; + 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"); @@ -1432,7 +1437,7 @@ __global__ void cached_notify(const int rdma_clean_offset, if (elect_one_sync()) { tma_load_1d(tma_buffer, - combined_nvl_head + batch_start_idx * NUM_MAX_NVL_PEERS, + 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); @@ -1441,11 +1446,11 @@ __global__ void cached_notify(const int rdma_clean_offset, __syncwarp(); for (int token_idx = batch_end_idx - 1; token_idx >= batch_start_idx; --token_idx) { - if (lane_id < NUM_MAX_NVL_PEERS) { + if (lane_id < LEGACY_NUM_MAX_NVL_PEERS) { auto current_head = - reinterpret_cast(tma_buffer)[(token_idx - batch_start_idx) * NUM_MAX_NVL_PEERS + lane_id]; + 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) * NUM_MAX_NVL_PEERS + lane_id] = + reinterpret_cast(tma_buffer)[(token_idx - batch_start_idx) * LEGACY_NUM_MAX_NVL_PEERS + lane_id] = -last_head - 1; } else { last_head = current_head; @@ -1457,7 +1462,7 @@ __global__ void cached_notify(const int rdma_clean_offset, if (elect_one_sync()) tma_store_1d(tma_buffer, - combined_nvl_head + batch_start_idx * NUM_MAX_NVL_PEERS, + 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(); @@ -1491,7 +1496,7 @@ void cached_notify(int hidden_int4, 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 / NUM_MAX_NVL_PEERS; + const auto num_rdma_ranks = num_ranks / LEGACY_NUM_MAX_NVL_PEERS; const int kNumTMABytesPerWarp = 8192; const int smem_size = kNumTMABytesPerWarp * num_warps; @@ -1503,7 +1508,7 @@ void cached_notify(int hidden_int4, num_topk_idx, num_topk_weights, num_rdma_ranks, - NUM_MAX_NVL_PEERS, + LEGACY_NUM_MAX_NVL_PEERS, num_max_nvl_chunked_recv_tokens, num_channels, is_cached_dispatch); @@ -1535,7 +1540,7 @@ void cached_notify(int hidden_int4, rank, num_ranks, is_cached_dispatch, - cpu_rdma_team); + nvshmem::cpu_rdma_team); } template int4* { return reinterpret_cast(smem_ptr + i * kNumTMABufferBytesPerStage + j * kNumTMALoadBytes); }; auto tma_store_buffer = [=](const int& i) -> int4* { - return reinterpret_cast(smem_ptr + i * kNumTMABufferBytesPerStage + NUM_MAX_NVL_PEERS * kNumTMALoadBytes); + 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 + (NUM_MAX_NVL_PEERS + 1) * kNumTMALoadBytes); + return reinterpret_cast(smem_ptr + i * kNumTMABufferBytesPerStage + (LEGACY_NUM_MAX_NVL_PEERS + 1) * kNumTMALoadBytes); }; // Prefetch @@ -1712,7 +1717,7 @@ template 0) ? kNumCombineForwarderWarps / kNumRDMARanks : 1, int kNumForwarders = kNumRDMARanks* kNumWarpsPerForwarder, - int kNumRDMAReceivers = kNumForwarders - NUM_MAX_NVL_PEERS> + 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, @@ -1753,16 +1758,16 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co const auto num_bytes_per_token = get_num_bytes_per_token(hidden_int4, 0, 0, num_topk); // NOTES: we decouple a channel into 2 SMs - const auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS; + 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 < NUM_MAX_NVL_PEERS) { + if (warp_id < LEGACY_NUM_MAX_NVL_PEERS) { auto shuffled_warp_id = warp_id; - shuffled_warp_id = (shuffled_warp_id + channel_id) % NUM_MAX_NVL_PEERS; + 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 - NUM_MAX_NVL_PEERS}; + return {WarpRole::kRDMAReceiver, warp_id - LEGACY_NUM_MAX_NVL_PEERS}; } else { return {WarpRole::kCoordinator, 0}; } @@ -1790,14 +1795,14 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co auto dst_buffer_ptr = buffer_ptrs[dst_nvl_rank], local_buffer_ptr = buffer_ptrs[nvl_rank]; auto nvl_channel_x = AsymBuffer(dst_buffer_ptr, num_max_nvl_chunked_recv_tokens * num_bytes_per_token, - NUM_MAX_NVL_PEERS, + 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, NUM_MAX_NVL_PEERS, channel_id, num_channels, dst_nvl_rank) + 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, NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank) + 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 @@ -1815,7 +1820,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co // Get tasks for each RDMA lane int token_start_idx = 0, token_end_idx = 0; if (lane_id < kNumRDMARanks) { - int prefix_idx = (lane_id * NUM_MAX_NVL_PEERS + dst_nvl_rank) * num_channels + channel_id; + 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]; } @@ -1847,7 +1852,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co cached_channel_head_idx = ld_volatile_global(nvl_channel_head.buffer() + lane_id); // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) { + 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", @@ -1930,22 +1935,22 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co // NVL layouts void* local_nvl_buffer = buffer_ptrs[nvl_rank]; - void* nvl_buffers[NUM_MAX_NVL_PEERS]; + void* nvl_buffers[LEGACY_NUM_MAX_NVL_PEERS]; #pragma unroll - for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) + 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, NUM_MAX_NVL_PEERS, channel_id, num_channels) - .advance_also(nvl_buffers); + 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, NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank) + 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, NUM_MAX_NVL_PEERS, channel_id, num_channels) - .advance_also(nvl_buffers); + 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][NUM_MAX_NVL_PEERS]; + __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]; @@ -1971,13 +1976,13 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co // TMA stuffs constexpr int kNumStages = 2; constexpr int kNumTMALoadBytes = sizeof(int4) * 32; - constexpr int kNumTMABufferBytesPerStage = kNumTMALoadBytes * (NUM_MAX_NVL_PEERS + 1) + 16; + 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 * (NUM_MAX_NVL_PEERS + 1)); + return reinterpret_cast(smem_ptr + i * kNumTMABufferBytesPerStage + kNumTMALoadBytes * (LEGACY_NUM_MAX_NVL_PEERS + 1)); }; uint32_t tma_phase[kNumStages] = {0}; if (lane_id < kNumStages) { @@ -1992,8 +1997,8 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co nvl_channel_tail.advance(dst_rdma_rank); // Clean shared memory and sync - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); - lane_id < NUM_MAX_NVL_PEERS ? (forwarder_nvl_head[warp_id][lane_id] = 0) : 0; + 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(); @@ -2003,7 +2008,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co int num_tokens_prefix = channel_id == 0 ? 0 : rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id - 1]; num_tokens_to_combine -= num_tokens_prefix; num_tokens_prefix += dst_rdma_rank == 0 ? 0 : rdma_rank_prefix_sum[dst_rdma_rank - 1]; - combined_nvl_head += num_tokens_prefix * NUM_MAX_NVL_PEERS; + 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) { @@ -2019,7 +2024,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co break; // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { + 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", @@ -2040,8 +2045,8 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co // Read expected head EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); int expected_head = -1; - if (lane_id < NUM_MAX_NVL_PEERS) { - expected_head = ld_nc_global(combined_nvl_head + token_idx * NUM_MAX_NVL_PEERS + lane_id); + 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); } @@ -2052,7 +2057,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co cached_nvl_channel_tail_idx = ld_acquire_sys_global(nvl_channel_tail.buffer(lane_id)); // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES and lane_id < NUM_MAX_NVL_PEERS) { + 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", @@ -2083,7 +2088,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co hidden_bytes + sizeof(SourceMeta)) + topk_idx); }; - combine_token( + combine_token( expected_head >= 0, expected_head, lane_id, @@ -2100,7 +2105,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co tma_phase); // Update head - if (lane_id < NUM_MAX_NVL_PEERS) + 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); } @@ -2172,7 +2177,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co cached_channel_tail_idx = static_cast(ld_acquire_sys_global(rdma_channel_tail.buffer(lane_id))); // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { + 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", @@ -2229,7 +2234,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co int last_rdma_head = 0; int last_nvl_head[kNumRDMARanks] = {0}; int dst_rdma_rank = lane_id < kNumRDMARanks ? lane_id : 0; - int dst_nvl_rank = lane_id < NUM_MAX_NVL_PEERS ? lane_id : 0; + 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 @@ -2263,13 +2268,13 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * 32, 1) combine(int4* co for (int j = 0; j < num_warps_per_rdma_rank; ++j) if (not forwarder_retired[i * num_warps_per_rdma_rank + j]) min_head = min(min_head, forwarder_nvl_head[i * num_warps_per_rdma_rank + j][dst_nvl_rank]); - if (min_head != std::numeric_limits::max() and min_head > last_nvl_head[i] and lane_id < NUM_MAX_NVL_PEERS) + 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(NUM_WAIT_NANOSECONDS); + __nanosleep(LEGACY_NUM_WAIT_NANOSECONDS); } } } @@ -2308,7 +2313,7 @@ void combine(cudaDataType_t type, constexpr int kNumTMABytesPerSenderWarp = 16384; constexpr int kNumTMABytesPerForwarderWarp = 9248; constexpr int smem_size = - std::max(kNumTMABytesPerSenderWarp * NUM_MAX_NVL_PEERS, kNumTMABytesPerForwarderWarp * kNumCombineForwarderWarps); + std::max(kNumTMABytesPerSenderWarp * LEGACY_NUM_MAX_NVL_PEERS, kNumTMABytesPerForwarderWarp * kNumCombineForwarderWarps); #define COMBINE_LAUNCH_CASE(num_rdma_ranks) \ { \ @@ -2355,11 +2360,11 @@ void combine(cudaDataType_t type, } \ break - int num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; + 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 > NUM_MAX_NVL_PEERS and num_forwarder_warps % num_rdma_ranks == 0); + 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)); @@ -2374,4 +2379,6 @@ void combine(cudaDataType_t type, } // namespace internode +} // namespace legacy + } // namespace deep_ep diff --git a/csrc/kernels/internode_ll.cu b/csrc/kernels/legacy/internode_ll.cu similarity index 96% rename from csrc/kernels/internode_ll.cu rename to csrc/kernels/legacy/internode_ll.cu index e9fd473bc..9ede8e40f 100644 --- a/csrc/kernels/internode_ll.cu +++ b/csrc/kernels/legacy/internode_ll.cu @@ -1,9 +1,8 @@ -#include "configs.cuh" -#include "exception.cuh" +#include "compiled.cuh" #include "ibgda_device.cuh" #include "launch.cuh" -namespace deep_ep { +namespace deep_ep::legacy { namespace internode_ll { @@ -54,11 +53,11 @@ __forceinline__ __device__ void barrier(int thread_id, int rank, int num_ranks, 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) <= NUM_TIMEOUT_CYCLES // not timeout + && (wait_recv_cost = clock64() - start_time) <= LEGACY_NUM_TIMEOUT_CYCLES // not timeout ) ; // Mask rank if timeout - if (wait_recv_cost > NUM_TIMEOUT_CYCLES) { + 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(); @@ -187,7 +186,7 @@ __global__ __launch_bounds__(1024, 1) void dispatch(void* packed_recv_x, __shared__ int shared_num_tokens_sent_per_expert[kNumMaxWarpGroups]; // Sending phase - if ((phases & LOW_LATENCY_SEND_PHASE) == 0) + if ((phases & LEGACY_LOW_LATENCY_SEND_PHASE) == 0) goto LOW_LATENCY_DISPATCH_RECV; // There are 2 kinds of warps in this part: @@ -293,7 +292,7 @@ __global__ __launch_bounds__(1024, 1) void dispatch(void* packed_recv_x, __syncwarp(); #pragma unroll for (int i = lane_id; i < num_experts; i += 32) - atomic_add_release_global(atomic_finish_counter_per_expert + i, FINISHED_SUM_TAG); + 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 @@ -315,7 +314,7 @@ __global__ __launch_bounds__(1024, 1) void dispatch(void* packed_recv_x, auto sum = warp_reduce_sum(expert_count[i - expert_begin_idx]); if (lane_id == 0) { shared_num_tokens_sent_per_expert[i - expert_begin_idx] = sum; - atomic_add_release_global(atomic_finish_counter_per_expert + i, FINISHED_SUM_TAG - sum); + atomic_add_release_global(atomic_finish_counter_per_expert + i, LEGACY_FINISHED_SUM_TAG - sum); } } } @@ -328,7 +327,7 @@ __global__ __launch_bounds__(1024, 1) void dispatch(void* packed_recv_x, const auto num_tokens_sent = shared_num_tokens_sent_per_expert[responsible_expert_idx - sm_id * num_warp_groups]; // Wait local sends issued and send expert counts - while (ld_acquire_global(atomic_finish_counter_per_expert + responsible_expert_idx) != FINISHED_SUM_TAG * 2) + 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); @@ -352,11 +351,11 @@ __global__ __launch_bounds__(1024, 1) void dispatch(void* packed_recv_x, // Receiving phase LOW_LATENCY_DISPATCH_RECV: - if ((phases & LOW_LATENCY_RECV_PHASE) == 0) + 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 & LOW_LATENCY_SEND_PHASE) + if (phases & LEGACY_LOW_LATENCY_SEND_PHASE) cg::this_grid().sync(); // Receiving and packing @@ -387,7 +386,7 @@ LOW_LATENCY_DISPATCH_RECV: 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) <= NUM_TIMEOUT_CYCLES // not timeout + && (wait_recv_cost = clock64() - start_time) <= LEGACY_NUM_TIMEOUT_CYCLES // not timeout ) ; } @@ -395,7 +394,7 @@ LOW_LATENCY_DISPATCH_RECV: if (num_recv_tokens == 0) num_recv_tokens = -1; // Mask rank if timeout - if (wait_recv_cost > NUM_TIMEOUT_CYCLES) { + 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, @@ -504,7 +503,7 @@ void dispatch(void* packed_recv_x, // 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 <= NUM_WORKSPACE_BYTES); + EP_HOST_ASSERT(num_experts * sizeof(int) * 2 <= LEGACY_NUM_WORKSPACE_BYTES); // FP8 checks if (use_ue8m0) @@ -563,12 +562,12 @@ __forceinline__ __device__ int logfmt_encode(void* buffer, nv_bfloat162* shared_ constexpr int kNumValues = 1 << (kNumBits - 1); int4 int4_values[kNumSendUnrolls]; - const auto& uint32_values = reinterpret_cast(int4_values); - const auto& bf162_values = reinterpret_cast(int4_values); + const auto uint32_values = reinterpret_cast(int4_values); + const auto bf162_values = reinterpret_cast(int4_values); // Calculate lane offset - const auto& ld_buffer = reinterpret_cast(static_cast(buffer) + lane_id * (kNumSendUnrolls * sizeof(int4))); - const auto& st_buffer = + const auto ld_buffer = reinterpret_cast(static_cast(buffer) + lane_id * (kNumSendUnrolls * sizeof(int4))); + const auto st_buffer = reinterpret_cast(static_cast(buffer) + lane_id * (kNumSendUnrolls * sizeof(int4) * 10 / 16)); // Local log amax @@ -601,9 +600,9 @@ __forceinline__ __device__ int logfmt_encode(void* buffer, nv_bfloat162* shared_ __syncwarp(); // Calculate log amin/amax float - const auto& log_amax = log2f_approx(amax); - const auto& log_amin = fmaxf(log2f_approx(amin), log_amax - kMinClip); - const bool& enable_cast = warp_reduce_and(log_amax < kLogThreshold and log_amin < log_amax); + const auto log_amax = log2f_approx(amax); + const auto log_amin = fmaxf(log2f_approx(amin), log_amax - kMinClip); + const bool enable_cast = warp_reduce_and(log_amax < kLogThreshold and log_amin < log_amax); // Case into LogFMT-10 if satisfied if (enable_cast) { @@ -619,7 +618,7 @@ __forceinline__ __device__ int logfmt_encode(void* buffer, nv_bfloat162* shared_ for (int i = 0; i < kNumSendUnrolls / 2; ++i) { #pragma unroll for (int k = 0; k < kNumElemsPerInt4; ++k) { - const auto& [x, y] = __bfloat1622float2(bf162_values[i * kNumElemsPerInt4 + k]); + const auto [x, y] = __bfloat1622float2(bf162_values[i * kNumElemsPerInt4 + k]); encoded[k * 2 + 0] = __float2uint_rd(fmaxf(log2f_approx(x) * step_inv + fused_rounding, 0)); encoded[k * 2 + 1] = __float2uint_rd(fmaxf(log2f_approx(y) * step_inv + fused_rounding, 0)); } @@ -648,7 +647,7 @@ __forceinline__ __device__ void logfmt_check_amaxmin( if (lane_id < kNumLanes) { // Calculate log amin/amax float auto amaxmin2 = reinterpret_cast(meta_buffer)[lane_id]; - const auto& bf162_amaxmin = reinterpret_cast<__nv_bfloat162*>(&amaxmin2); + const auto bf162_amaxmin = reinterpret_cast<__nv_bfloat162*>(&amaxmin2); float log_amax[2], log_amin[2]; #pragma unroll for (int i = 0; i < 2; ++i) { @@ -662,8 +661,8 @@ __forceinline__ __device__ void logfmt_check_amaxmin( shared_log_amin[lane_id] = make_float2(log_amin[0], log_amin[1]); } - const auto& casted = warp_reduce_and(enable_cast) ? 1u << (lane_id / kNumRecvUnrolls) : 0u; - const auto& num_casted_prefix = __popc(warp_reduce_or(casted) & ((1u << (lane_id / kNumRecvUnrolls)) - 1)); + const auto casted = warp_reduce_and(enable_cast) ? 1u << (lane_id / kNumRecvUnrolls) : 0u; + const auto num_casted_prefix = __popc(warp_reduce_or(casted) & ((1u << (lane_id / kNumRecvUnrolls)) - 1)); if (lane_id < kNumLanes and lane_id % kNumRecvUnrolls == 0) shared_cast_info[lane_id / kNumRecvUnrolls] = (num_casted_prefix << 1) | (casted ? 1u : 0u); @@ -677,7 +676,7 @@ __forceinline__ __device__ void decode_and_accumulate( constexpr int kNumBits = 10; constexpr int kNumValues = 1 << (kNumBits - 1); - const auto& step = (log_amax - log_amin) / static_cast(kNumValues - 2); + const auto step = (log_amax - log_amin) / static_cast(kNumValues - 2); auto decode = [=](const uint32_t& encoded, const uint32_t& sign) { const auto decoded = encoded == 0 ? .0f : exp2f_approx((encoded - 1) * step + log_amin); return sign ? -decoded : decoded; @@ -693,7 +692,7 @@ __forceinline__ __device__ void decode_and_accumulate( concat[k] = (ld_buffer[i * 5 + k - 1] >> (32 - k * 5)) | (ld_buffer[i * 5 + k] << (k * 5)); concat[5] = ld_buffer[i * 5 + 4] >> 7; - const uint32_t& local_signs = ld_buffer[i * 5 + 4] >> 16; + const uint32_t local_signs = ld_buffer[i * 5 + 4] >> 16; #pragma unroll for (int k = 0; k < 5; ++k) { accum[i * 16 + k * 3 + 0] += decode((concat[k] >> 0) & 0x1ff, (local_signs >> (k * 3 + 0)) & 1) * weight; @@ -771,7 +770,7 @@ __global__ __launch_bounds__(1024, 1) void combine(void* combined_x, EP_STATIC_ASSERT(num_bytes_per_slot % sizeof(int4) == 0, "Invalid vectorization"); // Sending phase - if ((phases & LOW_LATENCY_SEND_PHASE) == 0) + if ((phases & LEGACY_LOW_LATENCY_SEND_PHASE) == 0) goto LOW_LATENCY_COMBINE_RECV; // Clean up next buffer @@ -862,11 +861,11 @@ __global__ __launch_bounds__(1024, 1) void combine(void* combined_x, #pragma unroll for (int i = lane_id * kNumSendUnrolls, iter_idx = 0; i < hidden_bf16_int4_pad; i += 32 * kNumSendUnrolls, ++iter_idx) { // Load the next iteration - const int& stage_idx = iter_idx % kNumStages; - const int& next_stage_idx = (iter_idx + 1) % kNumStages; + const int stage_idx = iter_idx % kNumStages; + const int next_stage_idx = (iter_idx + 1) % kNumStages; if (iter_idx + 1 < kNumIters and elect_one_sync()) { tma_store_wait(); - const auto& offset_int4 = i + 32 * kNumSendUnrolls; + const auto offset_int4 = i + 32 * kNumSendUnrolls; tma_load_and_arrive(next_stage_idx, cpy_src_int4_ptr + offset_int4, get_num_tma_bytes(offset_int4)); } __syncwarp(); @@ -942,7 +941,7 @@ __global__ __launch_bounds__(1024, 1) void combine(void* combined_x, // Receiving phase LOW_LATENCY_COMBINE_RECV: - if ((phases & LOW_LATENCY_RECV_PHASE) == 0) + if ((phases & LEGACY_LOW_LATENCY_RECV_PHASE) == 0) return; // Wait all ranks to arrive @@ -954,12 +953,12 @@ LOW_LATENCY_COMBINE_RECV: uint64_t wait_recv_cost = 0; if (not is_rank_masked(mask_buffer_ptr, src_rank)) { while (ld_acquire_sys_global(rdma_recv_flag + responsible_expert_idx) == 0 // recv not ready - && (wait_recv_cost = clock64() - start_time) <= NUM_TIMEOUT_CYCLES // not timeout + && (wait_recv_cost = clock64() - start_time) <= LEGACY_NUM_TIMEOUT_CYCLES // not timeout ) ; } // Mask rank if timeout - if (wait_recv_cost > NUM_TIMEOUT_CYCLES) { + if (wait_recv_cost > LEGACY_NUM_TIMEOUT_CYCLES) { printf("Warning: DeepEP timeout for combine receive, rank %d, local_expert_idx %d, src_rank %d\n", rank, responsible_expert_idx % num_local_experts, @@ -1057,7 +1056,7 @@ LOW_LATENCY_COMBINE_RECV: if (elect_one_sync()) { int num_casted = 0; if constexpr (kUseLogFMT) { - const auto& info = cast_info_buffers[stage_idx][num_decode_warps - 1]; + const auto info = cast_info_buffers[stage_idx][num_decode_warps - 1]; num_casted = (info >> 1) + (info & 1); } int num_tma_bytes = num_casted * kNumLogFMTPerWarpBytes + (num_decode_warps - num_casted) * kNumBF16PerWarpBytes; @@ -1086,11 +1085,11 @@ LOW_LATENCY_COMBINE_RECV: continue; if (is_rank_masked(mask_buffer_ptr, topk_idx_reg / num_local_experts)) continue; - const auto& topk_weight = __shfl_sync(0xffffffff, topk_weights_by_lane, i); + const auto topk_weight = __shfl_sync(0xffffffff, topk_weights_by_lane, i); mbarrier_wait(full_barriers[stage_idx], tma_phase, stage_idx); if constexpr (kUseLogFMT) { - const auto& info = cast_info_buffers[stage_idx][decode_warp_idx]; + const auto info = cast_info_buffers[stage_idx][decode_warp_idx]; bool enable_cast = info & 1; int num_casted_prefix = info >> 1; int tma_offset = @@ -1176,7 +1175,7 @@ void combine(void* combined_x, // Check workspace auto atomic_clean_flag = static_cast(workspace); - EP_HOST_ASSERT(sizeof(int) <= NUM_WORKSPACE_BYTES); + EP_HOST_ASSERT(sizeof(int) <= LEGACY_NUM_WORKSPACE_BYTES); EP_HOST_ASSERT(num_topk <= kNumMaxTopk); // Online cast cannot use zero-copy @@ -1244,9 +1243,8 @@ __launch_bounds__(kNumThreads, 1) __global__ void query_mask_buffer(int* mask_bu const auto sm_id = static_cast(blockIdx.x); const auto num_threads = num_sms * kNumThreads; const auto thread_id = sm_id * kNumThreads + static_cast(threadIdx.x); - for (int rank_id = thread_id; rank_id < num_ranks; rank_id += num_threads) { + for (int rank_id = thread_id; rank_id < num_ranks; rank_id += num_threads) mask_tensor[rank_id] = mask_buffer_ptr[rank_id]; - } } void query_mask_buffer(int* mask_buffer_ptr, int num_ranks, int* mask_tensor, cudaStream_t stream) { @@ -1260,9 +1258,8 @@ template __launch_bounds__(kNumThreads, 1) __global__ void update_mask_buffer(int* mask_buffer_ptr, int rank_to_mask, bool mask) { const auto sm_id = static_cast(blockIdx.x); const auto thread_id = static_cast(threadIdx.x); - if (sm_id == 0 && thread_id == 0) { + if (sm_id == 0 and thread_id == 0) atomicExch(mask_buffer_ptr + rank_to_mask, mask ? 1 : 0); - } } void update_mask_buffer(int* mask_buffer_ptr, int rank, bool mask, cudaStream_t stream) { @@ -1289,4 +1286,4 @@ void clean_mask_buffer(int* mask_buffer_ptr, int num_ranks, cudaStream_t stream) } // namespace internode_ll -} // namespace deep_ep +} // namespace deep_ep::legacy diff --git a/csrc/kernels/intranode.cu b/csrc/kernels/legacy/intranode.cu similarity index 98% rename from csrc/kernels/intranode.cu rename to csrc/kernels/legacy/intranode.cu index 44b0cef01..682f17bb5 100644 --- a/csrc/kernels/intranode.cu +++ b/csrc/kernels/legacy/intranode.cu @@ -1,13 +1,27 @@ #include "buffer.cuh" -#include "configs.cuh" -#include "exception.cuh" +#include "compiled.cuh" #include "launch.cuh" #include "utils.cuh" -namespace deep_ep { +namespace deep_ep::legacy { namespace intranode { +template +__global__ void barrier(int** barrier_signal_ptrs, int rank) { + barrier_block(barrier_signal_ptrs, rank); +} + +void barrier(int** barrier_signal_ptrs, int rank, int num_ranks, cudaStream_t stream) { +#define BARRIER_LAUNCH_CASE(ranks) \ +LAUNCH_KERNEL(&cfg, barrier, barrier_signal_ptrs, rank); \ +break + + SETUP_LAUNCH_CONFIG(1, 32, stream); + SWITCH_RANKS(BARRIER_LAUNCH_CASE); +#undef BARRIER_LAUNCH_CASE +} + template __global__ void notify_dispatch(const int* num_tokens_per_rank, int* moe_recv_counter_mapped, @@ -329,7 +343,7 @@ __global__ void __launch_bounds__(kNumThreads, 1) dispatch(int4* recv_x, break; // Rare cases to loop again - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { + if (clock64() - start_time > LEGACY_NUM_TIMEOUT_CYCLES) { printf("DeepEP timeout for dispatch senders, rank %d, responsible_channel = %d\n", rank, responsible_channel); trap(); } @@ -443,7 +457,7 @@ __global__ void __launch_bounds__(kNumThreads, 1) dispatch(int4* recv_x, } // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { + if (clock64() - start_time > LEGACY_NUM_TIMEOUT_CYCLES) { printf("DeepEP timeout for dispatch receivers, rank %d, responsible_channel = %d, tokens remained: %d\n", rank, responsible_channel, @@ -786,7 +800,7 @@ __global__ void __launch_bounds__(kNumThreads, 1) combine(dtype_t* recv_x, break; // Rare cases to loop again - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { + if (clock64() - start_time > LEGACY_NUM_TIMEOUT_CYCLES) { printf("DeepEP timeout for combine senders, rank %d, responsible_channel = %d\n", rank, responsible_channel); trap(); } @@ -911,7 +925,7 @@ __global__ void __launch_bounds__(kNumThreads, 1) combine(dtype_t* recv_x, auto start_time = clock64(); while (__any_sync(0xffffffff, channel_tail_idx[lane_id] <= expected_head and expected_head >= 0)) { // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { + if (clock64() - start_time > LEGACY_NUM_TIMEOUT_CYCLES) { printf("DeepEP timeout for combine receivers, rank %d, responsible_channel = %d, expect = %d\n", rank, responsible_channel, @@ -1100,4 +1114,4 @@ void combine(cudaDataType_t type, } // namespace intranode -} // namespace deep_ep +} // namespace deep_ep::legacy diff --git a/csrc/kernels/launch.cuh b/csrc/kernels/legacy/launch.cuh similarity index 97% rename from csrc/kernels/launch.cuh rename to csrc/kernels/legacy/launch.cuh index 856370087..60e308039 100644 --- a/csrc/kernels/launch.cuh +++ b/csrc/kernels/legacy/launch.cuh @@ -1,7 +1,6 @@ #pragma once -#include "configs.cuh" -#include "exception.cuh" +#include "compiled.cuh" #ifndef SETUP_LAUNCH_CONFIG #ifndef DISABLE_SM90_FEATURES @@ -26,7 +25,7 @@ #ifndef LAUNCH_KERNEL #ifndef DISABLE_SM90_FEATURES -#define LAUNCH_KERNEL(config, kernel, ...) CUDA_CHECK(cudaLaunchKernelEx(config, kernel, ##__VA_ARGS__)) +#define LAUNCH_KERNEL(config, kernel, ...) CUDA_RUNTIME_CHECK(cudaLaunchKernelEx(config, kernel, ##__VA_ARGS__)) #else #define LAUNCH_KERNEL(config, kernel, ...) \ do { \ @@ -65,7 +64,7 @@ while (false) #define SWITCH_RDMA_RANKS(case_macro) \ - switch (num_ranks / NUM_MAX_NVL_PEERS) { \ + switch (num_ranks / LEGACY_NUM_MAX_NVL_PEERS) { \ case 2: \ case_macro(2); \ case 3: \ diff --git a/csrc/kernels/layout.cu b/csrc/kernels/legacy/layout.cu similarity index 91% rename from csrc/kernels/layout.cu rename to csrc/kernels/legacy/layout.cu index c3a16aed8..91e39ddff 100644 --- a/csrc/kernels/layout.cu +++ b/csrc/kernels/legacy/layout.cu @@ -1,8 +1,9 @@ -#include "configs.cuh" -#include "exception.cuh" +#include + +#include "compiled.cuh" #include "launch.cuh" -namespace deep_ep { +namespace deep_ep::legacy { namespace layout { @@ -52,15 +53,15 @@ __global__ void get_dispatch_layout(const topk_idx_t* topk_idx, } if (num_tokens_per_rdma_rank != nullptr) - EP_DEVICE_ASSERT(num_ranks % NUM_MAX_NVL_PEERS == 0 and num_ranks > NUM_MAX_NVL_PEERS); + EP_DEVICE_ASSERT(num_ranks % LEGACY_NUM_MAX_NVL_PEERS == 0 and num_ranks > LEGACY_NUM_MAX_NVL_PEERS); // Count rank statistics - constexpr int kNumRDMARanksPerSM = kNumRanksPerSM / NUM_MAX_NVL_PEERS; + constexpr int kNumRDMARanksPerSM = kNumRanksPerSM / LEGACY_NUM_MAX_NVL_PEERS; __shared__ int num_tokens_per_rank_per_thread[kNumThreads][kNumRanksPerSM]; __shared__ int num_tokens_per_rdma_rank_per_thread[kNumThreads][kNumRDMARanksPerSM]; auto sm_begin = (num_experts + kNumExpertsPerSM - 1) / kNumExpertsPerSM; int rank_begin_idx = (sm_id - sm_begin) * kNumRanksPerSM, rank_end_idx = min(rank_begin_idx + kNumRanksPerSM, num_ranks); - int rdma_rank_begin_idx = rank_begin_idx / NUM_MAX_NVL_PEERS, rdma_rank_end_idx = rank_end_idx / NUM_MAX_NVL_PEERS; + int rdma_rank_begin_idx = rank_begin_idx / LEGACY_NUM_MAX_NVL_PEERS, rdma_rank_end_idx = rank_end_idx / LEGACY_NUM_MAX_NVL_PEERS; if (rank_begin_idx < rank_end_idx) { const auto num_expert_per_rank = num_experts / num_ranks; auto expert_begin = rank_begin_idx * num_expert_per_rank; @@ -83,7 +84,7 @@ __global__ void get_dispatch_layout(const topk_idx_t* topk_idx, if (expert_begin <= expert_idx and expert_idx < expert_end) { // Count single rank rank_idx = expert_idx / num_expert_per_rank - rank_begin_idx; - is_in_rank[rank_idx]++, is_in_rdma_rank[rank_idx / NUM_MAX_NVL_PEERS]++; + is_in_rank[rank_idx]++, is_in_rdma_rank[rank_idx / LEGACY_NUM_MAX_NVL_PEERS]++; } } @@ -132,7 +133,7 @@ void get_dispatch_layout(const topk_idx_t* topk_idx, cudaStream_t stream) { constexpr int kNumThreads = 256, kNumExpertsPerSM = 4, kNumRanksPerSM = 8; int num_sms = ((num_experts + kNumExpertsPerSM - 1) / kNumExpertsPerSM) + (num_ranks + kNumRanksPerSM - 1) / kNumRanksPerSM; - EP_STATIC_ASSERT(kNumRanksPerSM % NUM_MAX_NVL_PEERS == 0, "Invalid number of ranks per SM"); + EP_STATIC_ASSERT(kNumRanksPerSM % LEGACY_NUM_MAX_NVL_PEERS == 0, "Invalid number of ranks per SM"); SETUP_LAUNCH_CONFIG(num_sms, kNumThreads, stream); LAUNCH_KERNEL(&cfg, @@ -150,4 +151,4 @@ void get_dispatch_layout(const topk_idx_t* topk_idx, } // namespace layout -} // namespace deep_ep +} // namespace deep_ep::legacy diff --git a/csrc/kernels/utils.cuh b/csrc/kernels/legacy/utils.cuh similarity index 98% rename from csrc/kernels/utils.cuh rename to csrc/kernels/legacy/utils.cuh index 0c2eec020..d6ee6a418 100644 --- a/csrc/kernels/utils.cuh +++ b/csrc/kernels/legacy/utils.cuh @@ -1,6 +1,7 @@ #pragma once -#include "exception.cuh" +#include +#include #define UNROLLED_WARP_COPY(UNROLL_FACTOR, LANE_ID, N, DST, SRC, LD_FUNC, ST_FUNC) \ { \ @@ -27,7 +28,7 @@ } \ } -namespace deep_ep { +namespace deep_ep::legacy { template struct VecInt {}; @@ -352,7 +353,7 @@ __device__ __forceinline__ void mbarrier_inval(uint64_t* mbar_ptr) { template __device__ __forceinline__ void mbarrier_wait(uint64_t* mbar_ptr, uint32_t& phase, int stage_idx = 0) { auto mbar_int_ptr = static_cast(__cvta_generic_to_shared(mbar_ptr)); - const auto& wait = kWithMultiStages ? (phase >> stage_idx) & 1 : phase; + const auto wait = kWithMultiStages ? (phase >> stage_idx) & 1 : phase; asm volatile( "{\n\t" ".reg .pred P1; \n\t" @@ -411,7 +412,7 @@ __device__ __forceinline__ void tma_store_1d(const void* smem_ptr, const void* g template __device__ __forceinline__ void tma_store_wait() { - asm volatile("cp.async.bulk.wait_group.read %0;" ::"n"(N) : "memory"); + asm volatile("cp.async.bulk.wait_group %0;" ::"n"(N) : "memory"); } #endif @@ -513,8 +514,8 @@ __forceinline__ __device__ void barrier_block(int** barrier_signal_ptrs, int ran // Add self-ranks, sub other ranks if (thread_id < kNumRanks) { - atomicAdd_system(barrier_signal_ptrs[rank] + thread_id, FINISHED_SUM_TAG); - atomicSub_system(barrier_signal_ptrs[thread_id] + rank, FINISHED_SUM_TAG); + atomicAdd_system(barrier_signal_ptrs[rank] + thread_id, LEGACY_FINISHED_SUM_TAG); + atomicSub_system(barrier_signal_ptrs[thread_id] + rank, LEGACY_FINISHED_SUM_TAG); } EP_DEVICE_ASSERT(kNumRanks <= blockDim.x); @@ -525,7 +526,7 @@ __forceinline__ __device__ void barrier_block(int** barrier_signal_ptrs, int ran if (__all_sync(0xffffffff, value <= 0)) break; - if (clock64() - start_time > NUM_TIMEOUT_CYCLES and thread_id < kNumRanks) { + if (clock64() - start_time > LEGACY_NUM_TIMEOUT_CYCLES and thread_id < kNumRanks) { printf("DeepEP timeout check failed: rank = %d, thread = %d, value = %d)\n", rank, thread_id, value); trap(); } @@ -637,4 +638,4 @@ __forceinline__ __device__ T warp_reduce_or(T value) { return warp_reduce(value, ReduceOr{}); } -} // namespace deep_ep +} // namespace deep_ep::legacy diff --git a/csrc/kernels/runtime.cu b/csrc/kernels/runtime.cu deleted file mode 100644 index c4fbb8ed2..000000000 --- a/csrc/kernels/runtime.cu +++ /dev/null @@ -1,98 +0,0 @@ -#include -#include - -#include "configs.cuh" -#include "exception.cuh" -#include "launch.cuh" -#include "utils.cuh" - -#ifndef DISABLE_NVSHMEM -#include "ibgda_device.cuh" -#include "nvshmem.h" -#endif - -namespace deep_ep { - -namespace intranode { - -template -__global__ void barrier(int** barrier_signal_ptrs, int rank) { - barrier_block(barrier_signal_ptrs, rank); -} - -void barrier(int** barrier_signal_ptrs, int rank, int num_ranks, cudaStream_t stream) { -#define BARRIER_LAUNCH_CASE(ranks) \ - LAUNCH_KERNEL(&cfg, barrier, barrier_signal_ptrs, rank); \ - break - - SETUP_LAUNCH_CONFIG(1, 32, stream); - SWITCH_RANKS(BARRIER_LAUNCH_CASE); -#undef BARRIER_LAUNCH_CASE -} - -} // namespace intranode - -namespace internode { - -#ifndef DISABLE_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; -} - -int init(const std::vector& root_unique_id_val, int rank, int num_ranks, bool low_latency_mode) { - 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 - // NOTES: if `num_ranks <= NUM_MAX_NVL_PEERS` then only low-latency kernels are used - if (low_latency_mode and num_ranks > NUM_MAX_NVL_PEERS) { - EP_HOST_ASSERT(cpu_rdma_team == NVSHMEM_TEAM_INVALID); - EP_HOST_ASSERT(num_ranks % NUM_MAX_NVL_PEERS == 0); - EP_HOST_ASSERT(nvshmem_team_split_strided(NVSHMEM_TEAM_WORLD, - rank % NUM_MAX_NVL_PEERS, - NUM_MAX_NVL_PEERS, - num_ranks / NUM_MAX_NVL_PEERS, - &cpu_rdma_team_config, - 0, - &cpu_rdma_team) == 0); - EP_HOST_ASSERT(cpu_rdma_team != NVSHMEM_TEAM_INVALID); - } - - nvshmem_barrier_all(); - return nvshmem_my_pe(); -} - -void* alloc(size_t size, size_t alignment) { - return nvshmem_align(alignment, size); -} - -void free(void* ptr) { - nvshmem_free(ptr); -} - -void barrier() { - nvshmem_barrier_all(); -} - -void finalize() { - if (cpu_rdma_team != NVSHMEM_TEAM_INVALID) { - nvshmem_team_destroy(cpu_rdma_team); - cpu_rdma_team = NVSHMEM_TEAM_INVALID; - } - nvshmem_finalize(); -} -#endif - -} // namespace internode - -} // namespace deep_ep diff --git a/csrc/legacy/buffer.hpp b/csrc/legacy/buffer.hpp new file mode 100644 index 000000000..39ddcee3f --- /dev/null +++ b/csrc/legacy/buffer.hpp @@ -0,0 +1,1794 @@ +#pragma once + +#include +#include +#include +#include + +#include +#include + +#include +#include + +#include "../utils/event.hpp" +#include "../utils/shared_memory.hpp" +#include "../kernels/legacy/api.cuh" +#include "../kernels/backend/api.cuh" +#include "config.hpp" + +namespace deep_ep::legacy { + +struct Buffer { + EP_STATIC_ASSERT(LEGACY_NUM_MAX_NVL_PEERS == 8, "The number of maximum NVLink peers must be 8"); + +private: + // Low-latency mode buffer + int low_latency_buffer_idx = 0; + bool low_latency_mode = false; + + // NVLink Buffer + int64_t num_nvl_bytes; + void* buffer_ptrs[LEGACY_NUM_MAX_NVL_PEERS] = {nullptr}; + void** buffer_ptrs_gpu = nullptr; + + // NVSHMEM Buffer + int64_t num_rdma_bytes; + void* rdma_buffer_ptr = nullptr; + + // Shrink mode buffer + bool enable_shrink = false; + int* mask_buffer_ptr = nullptr; + int* sync_buffer_ptr = nullptr; + + // Device info and communication + int device_id; + int num_device_sms; + int rank, rdma_rank, nvl_rank; + int num_ranks, num_rdma_ranks, num_nvl_ranks; + shared_memory::MemHandle ipc_handles[LEGACY_NUM_MAX_NVL_PEERS]; + + // Stream for communication + at::cuda::CUDAStream comm_stream; + + // After IPC/NVSHMEM synchronization, this flag will be true + bool available = false; + + // Whether explicit `destroy()` is required. + bool explicitly_destroy; + // After `destroy()` be called, this flag will be true + bool destroyed = false; + + // Barrier signals + int* barrier_signal_ptrs[LEGACY_NUM_MAX_NVL_PEERS] = {nullptr}; + int** barrier_signal_ptrs_gpu = nullptr; + + // Workspace + void* workspace = nullptr; + + // Host-side MoE info + volatile int* moe_recv_counter = nullptr; + int* moe_recv_counter_mapped = nullptr; + + // Host-side expert-level MoE info + volatile int* moe_recv_expert_counter = nullptr; + int* moe_recv_expert_counter_mapped = nullptr; + + // Host-side RDMA-level MoE info + volatile int* moe_recv_rdma_counter = nullptr; + int* moe_recv_rdma_counter_mapped = nullptr; + + shared_memory::SharedMemoryAllocator shared_memory_allocator; + +public: + Buffer(int rank, + int num_ranks, + int64_t num_nvl_bytes, + int64_t num_rdma_bytes, + bool low_latency_mode, + bool explicitly_destroy, + bool enable_shrink, + bool use_fabric) : rank(rank), + num_ranks(num_ranks), + num_nvl_bytes(num_nvl_bytes), + num_rdma_bytes(num_rdma_bytes), + enable_shrink(enable_shrink), + low_latency_mode(low_latency_mode), + explicitly_destroy(explicitly_destroy), + comm_stream(at::cuda::getStreamFromPool(true)), + shared_memory_allocator(use_fabric) { + // Metadata memory + int64_t barrier_signal_bytes = LEGACY_NUM_MAX_NVL_PEERS * sizeof(int); + int64_t buffer_ptr_bytes = LEGACY_NUM_MAX_NVL_PEERS * sizeof(void*); + int64_t barrier_signal_ptr_bytes = LEGACY_NUM_MAX_NVL_PEERS * sizeof(int*); + + // Common checks + EP_STATIC_ASSERT(LEGACY_NUM_BUFFER_ALIGNMENT_BYTES % sizeof(int4) == 0, "Invalid alignment"); + EP_HOST_ASSERT(num_nvl_bytes % LEGACY_NUM_BUFFER_ALIGNMENT_BYTES == 0 and + (num_nvl_bytes <= std::numeric_limits::max() or num_rdma_bytes == 0)); + EP_HOST_ASSERT(num_rdma_bytes % LEGACY_NUM_BUFFER_ALIGNMENT_BYTES == 0 and + (low_latency_mode or num_rdma_bytes <= std::numeric_limits::max())); + EP_HOST_ASSERT(num_nvl_bytes / sizeof(int4) < std::numeric_limits::max()); + EP_HOST_ASSERT(num_rdma_bytes / sizeof(int4) < std::numeric_limits::max()); + EP_HOST_ASSERT(0 <= rank and rank < num_ranks and (num_ranks <= LEGACY_NUM_MAX_NVL_PEERS * LEGACY_NUM_MAX_RDMA_PEERS or low_latency_mode)); + EP_HOST_ASSERT(num_ranks < LEGACY_NUM_MAX_NVL_PEERS or num_ranks % LEGACY_NUM_MAX_NVL_PEERS == 0); + if (num_rdma_bytes > 0) + EP_HOST_ASSERT(num_ranks > LEGACY_NUM_MAX_NVL_PEERS or low_latency_mode); + + // Get ranks + CUDA_RUNTIME_CHECK(cudaGetDevice(&device_id)); + rdma_rank = rank / LEGACY_NUM_MAX_NVL_PEERS, nvl_rank = rank % LEGACY_NUM_MAX_NVL_PEERS; + num_rdma_ranks = std::max(1, num_ranks / LEGACY_NUM_MAX_NVL_PEERS), num_nvl_ranks = std::min(num_ranks, LEGACY_NUM_MAX_NVL_PEERS); + + // Get device info + cudaDeviceProp device_prop = {}; + CUDA_RUNTIME_CHECK(cudaGetDeviceProperties(&device_prop, device_id)); + num_device_sms = device_prop.multiProcessorCount; + + // Number of per-channel bytes cannot be large + EP_HOST_ASSERT(ceil_div(num_nvl_bytes, num_device_sms / 2) < std::numeric_limits::max()); + EP_HOST_ASSERT(ceil_div(num_rdma_bytes, num_device_sms / 2) < std::numeric_limits::max()); + + if (num_nvl_bytes > 0) { + // Local IPC: alloc local memory and set local IPC handles + shared_memory_allocator.malloc(&buffer_ptrs[nvl_rank], + num_nvl_bytes + barrier_signal_bytes + buffer_ptr_bytes + barrier_signal_ptr_bytes); + shared_memory_allocator.get_mem_handle(&ipc_handles[nvl_rank], buffer_ptrs[nvl_rank]); + buffer_ptrs_gpu = reinterpret_cast(static_cast(buffer_ptrs[nvl_rank]) + num_nvl_bytes + barrier_signal_bytes); + + // Set barrier signals + barrier_signal_ptrs[nvl_rank] = reinterpret_cast(static_cast(buffer_ptrs[nvl_rank]) + num_nvl_bytes); + barrier_signal_ptrs_gpu = + reinterpret_cast(static_cast(buffer_ptrs[nvl_rank]) + num_nvl_bytes + barrier_signal_bytes + buffer_ptr_bytes); + + // No need to synchronize, will do a full device sync during `sync` + CUDA_RUNTIME_CHECK(cudaMemsetAsync(barrier_signal_ptrs[nvl_rank], 0, barrier_signal_bytes, comm_stream)); + } + + // Create 32 MiB workspace + CUDA_RUNTIME_CHECK(cudaMalloc(&workspace, LEGACY_NUM_WORKSPACE_BYTES)); + CUDA_RUNTIME_CHECK(cudaMemsetAsync(workspace, 0, LEGACY_NUM_WORKSPACE_BYTES, comm_stream)); + + // MoE counter + CUDA_RUNTIME_CHECK(cudaMallocHost(&moe_recv_counter, sizeof(int64_t), cudaHostAllocMapped)); + CUDA_RUNTIME_CHECK(cudaHostGetDevicePointer(&moe_recv_counter_mapped, const_cast(moe_recv_counter), 0)); + *moe_recv_counter = -1; + + // MoE expert-level counter + CUDA_RUNTIME_CHECK(cudaMallocHost(&moe_recv_expert_counter, sizeof(int) * LEGACY_NUM_MAX_LOCAL_EXPERTS, cudaHostAllocMapped)); + CUDA_RUNTIME_CHECK(cudaHostGetDevicePointer(&moe_recv_expert_counter_mapped, const_cast(moe_recv_expert_counter), 0)); + for (int i = 0; i < LEGACY_NUM_MAX_LOCAL_EXPERTS; ++i) + moe_recv_expert_counter[i] = -1; + + // MoE RDMA-level counter + if (num_rdma_ranks > 0) { + CUDA_RUNTIME_CHECK(cudaMallocHost(&moe_recv_rdma_counter, sizeof(int), cudaHostAllocMapped)); + CUDA_RUNTIME_CHECK(cudaHostGetDevicePointer(&moe_recv_rdma_counter_mapped, const_cast(moe_recv_rdma_counter), 0)); + *moe_recv_rdma_counter = -1; + } + } + + ~Buffer() noexcept(false) { + if (not explicitly_destroy) { + destroy(); + } else if (not destroyed) { + printf("WARNING: destroy() was not called before DeepEP buffer destruction, which can leak resources.\n"); + fflush(stdout); + } + } + + bool is_available() const { + return available; + } + + bool is_internode_available() const { + return is_available() and num_ranks > LEGACY_NUM_MAX_NVL_PEERS; + } + + int get_num_rdma_ranks() const { + return num_rdma_ranks; + } + + int get_rdma_rank() const { + return rdma_rank; + } + + int get_root_rdma_rank(bool global) const { + return global ? nvl_rank : 0; + } + + int get_local_device_id() const { + return device_id; + } + + pybind11::bytearray get_local_ipc_handle() const { + const shared_memory::MemHandle& handle = ipc_handles[nvl_rank]; + return {reinterpret_cast(&handle), sizeof(handle)}; + } + + pybind11::bytearray get_local_nvshmem_unique_id() const { + EP_HOST_ASSERT(rdma_rank == 0 and "Only RDMA rank 0 can get NVSHMEM unique ID"); + const auto unique_id = nvshmem::get_unique_id(); + return {reinterpret_cast(unique_id.data()), unique_id.size()}; + } + + torch::Tensor get_local_buffer_tensor(const pybind11::object& dtype, int64_t offset, bool use_rdma_buffer) const { + torch::ScalarType casted_dtype = torch::python::detail::py_object_to_dtype(dtype); + auto element_bytes = static_cast(elementSize(casted_dtype)); + auto base_ptr = static_cast(use_rdma_buffer ? rdma_buffer_ptr : buffer_ptrs[nvl_rank]) + offset; + auto num_bytes = use_rdma_buffer ? num_rdma_bytes : num_nvl_bytes; + return torch::from_blob(base_ptr, num_bytes / element_bytes, torch::TensorOptions().dtype(casted_dtype).device(at::kCUDA)); + } + + torch::Stream get_comm_stream() const { + return comm_stream; + } + + void sync(const std::vector& device_ids, + const std::vector>& all_gathered_handles, + const std::optional& root_unique_id_opt) { + EP_HOST_ASSERT(not is_available()); + + // Sync IPC handles + if (num_nvl_bytes > 0) { + EP_HOST_ASSERT(num_ranks == device_ids.size()); + EP_HOST_ASSERT(device_ids.size() == all_gathered_handles.size()); + for (int i = 0, offset = rdma_rank * num_nvl_ranks; i < num_nvl_ranks; ++i) { + EP_HOST_ASSERT(all_gathered_handles[offset + i].has_value()); + auto handle_str = std::string(all_gathered_handles[offset + i].value()); + EP_HOST_ASSERT(handle_str.size() == sizeof(shared_memory::MemHandle)); + if (offset + i != rank) { + std::memcpy(&ipc_handles[i], handle_str.c_str(), sizeof(shared_memory::MemHandle)); + shared_memory_allocator.open_mem_handle(&buffer_ptrs[i], &ipc_handles[i]); + barrier_signal_ptrs[i] = reinterpret_cast(static_cast(buffer_ptrs[i]) + num_nvl_bytes); + } else { + EP_HOST_ASSERT(std::memcmp(&ipc_handles[i], handle_str.c_str(), sizeof(shared_memory::MemHandle)) == 0); + } + } + + // Copy all buffer and barrier signal pointers to GPU + CUDA_RUNTIME_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs, sizeof(void*) * LEGACY_NUM_MAX_NVL_PEERS, cudaMemcpyHostToDevice)); + CUDA_RUNTIME_CHECK(cudaMemcpy(barrier_signal_ptrs_gpu, barrier_signal_ptrs, sizeof(int*) * LEGACY_NUM_MAX_NVL_PEERS, cudaMemcpyHostToDevice)); + CUDA_RUNTIME_CHECK(cudaDeviceSynchronize()); + } + + // Sync NVSHMEM handles and allocate memory + if (num_rdma_bytes > 0) { + // Initialize NVSHMEM + EP_HOST_ASSERT(root_unique_id_opt.has_value()); + std::vector root_unique_id(root_unique_id_opt->size()); + auto root_unique_id_str = root_unique_id_opt->cast(); + std::memcpy(root_unique_id.data(), root_unique_id_str.c_str(), root_unique_id_opt->size()); + auto nvshmem_rank = low_latency_mode ? rank : rdma_rank; + auto num_nvshmem_ranks = low_latency_mode ? num_ranks : num_rdma_ranks; + EP_HOST_ASSERT(nvshmem_rank == nvshmem::init(root_unique_id, nvshmem_rank, num_nvshmem_ranks, + low_latency_mode ? LEGACY_NUM_MAX_NVL_PEERS : 0)); + + // Allocate + rdma_buffer_ptr = nvshmem::alloc(num_rdma_bytes, LEGACY_NUM_BUFFER_ALIGNMENT_BYTES); + + // Clean buffer (mainly for low-latency mode) + CUDA_RUNTIME_CHECK(cudaMemset(rdma_buffer_ptr, 0, num_rdma_bytes)); + + // Allocate and clean shrink buffer + if (enable_shrink) { + int num_mask_buffer_bytes = num_ranks * sizeof(int); + int num_sync_buffer_bytes = num_ranks * sizeof(int); + mask_buffer_ptr = static_cast(nvshmem::alloc(num_mask_buffer_bytes, LEGACY_NUM_BUFFER_ALIGNMENT_BYTES)); + sync_buffer_ptr = static_cast(nvshmem::alloc(num_sync_buffer_bytes, LEGACY_NUM_BUFFER_ALIGNMENT_BYTES)); + CUDA_RUNTIME_CHECK(cudaMemset(mask_buffer_ptr, 0, num_mask_buffer_bytes)); + CUDA_RUNTIME_CHECK(cudaMemset(sync_buffer_ptr, 0, num_sync_buffer_bytes)); + } + + // Barrier + nvshmem::barrier(true); + CUDA_RUNTIME_CHECK(cudaDeviceSynchronize()); + } + + // Ready to use + available = true; + } + + void destroy() { + EP_HOST_ASSERT(not destroyed); + + // Synchronize + CUDA_RUNTIME_CHECK(cudaDeviceSynchronize()); + + if (num_nvl_bytes > 0) { + // Barrier + intranode::barrier(barrier_signal_ptrs_gpu, nvl_rank, num_nvl_ranks, comm_stream); + CUDA_RUNTIME_CHECK(cudaDeviceSynchronize()); + + // Close remote IPC + if (is_available()) { + for (int i = 0; i < num_nvl_ranks; ++i) + if (i != nvl_rank) + shared_memory_allocator.close_mem_handle(buffer_ptrs[i]); + } + + // Free local buffer and error flag + shared_memory_allocator.free(buffer_ptrs[nvl_rank]); + } + + // Free NVSHMEM + if (is_available() and num_rdma_bytes > 0) { + CUDA_RUNTIME_CHECK(cudaDeviceSynchronize()); + nvshmem::barrier(true); + nvshmem::free(rdma_buffer_ptr); + if (enable_shrink) { + nvshmem::free(mask_buffer_ptr); + nvshmem::free(sync_buffer_ptr); + } + nvshmem::finalize(); + } + + // Free workspace and MoE counter + CUDA_RUNTIME_CHECK(cudaFree(workspace)); + CUDA_RUNTIME_CHECK(cudaFreeHost(const_cast(moe_recv_counter))); + + // Free chunked mode staffs + CUDA_RUNTIME_CHECK(cudaFreeHost(const_cast(moe_recv_expert_counter))); + + destroyed = true; + available = false; + } + + std::tuple, torch::Tensor, torch::Tensor, std::optional> get_dispatch_layout( + const torch::Tensor& topk_idx, + int num_experts, + const std::optional& previous_event, + bool async, + bool allocate_on_comm_stream) { + EP_HOST_ASSERT(topk_idx.dim() == 2); + EP_HOST_ASSERT(topk_idx.is_contiguous()); + EP_HOST_ASSERT(num_experts > 0); + + // Allocate all tensors on comm stream if set + // NOTES: do not allocate tensors upfront! + auto compute_stream = at::cuda::getCurrentCUDAStream(); + if (allocate_on_comm_stream) { + EP_HOST_ASSERT(previous_event.has_value() and async); + at::cuda::setCurrentCUDAStream(comm_stream); + } + + // Wait previous tasks to be finished + if (previous_event.has_value()) { + stream_wait(comm_stream, previous_event.value()); + } else { + stream_wait(comm_stream, compute_stream); + } + + auto num_tokens = static_cast(topk_idx.size(0)), num_topk = static_cast(topk_idx.size(1)); + auto num_tokens_per_rank = torch::empty({num_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); + auto num_tokens_per_rdma_rank = std::optional(); + auto num_tokens_per_expert = torch::empty({num_experts}, dtype(torch::kInt32).device(torch::kCUDA)); + auto is_token_in_rank = torch::empty({num_tokens, num_ranks}, dtype(torch::kBool).device(torch::kCUDA)); + if (is_internode_available()) + num_tokens_per_rdma_rank = torch::empty({num_rdma_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); + + layout::get_dispatch_layout(topk_idx.data_ptr(), + num_tokens_per_rank.data_ptr(), + num_tokens_per_rdma_rank.has_value() ? num_tokens_per_rdma_rank.value().data_ptr() : nullptr, + num_tokens_per_expert.data_ptr(), + is_token_in_rank.data_ptr(), + num_tokens, + num_topk, + num_ranks, + num_experts, + comm_stream); + + // Wait streams + std::optional event; + if (async) { + event = EventHandle(comm_stream); + for (auto& t : {topk_idx, num_tokens_per_rank, num_tokens_per_expert, is_token_in_rank}) { + t.record_stream(comm_stream); + if (allocate_on_comm_stream) + t.record_stream(compute_stream); + } + for (auto& to : {num_tokens_per_rdma_rank}) { + to.has_value() ? to->record_stream(comm_stream) : void(); + if (allocate_on_comm_stream) + to.has_value() ? to->record_stream(compute_stream) : void(); + } + } else { + stream_wait(compute_stream, comm_stream); + } + + // Switch back compute stream + if (allocate_on_comm_stream) + at::cuda::setCurrentCUDAStream(compute_stream); + + return {num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, event}; + } + + std::tuple, + std::optional, + std::optional, + std::vector, + torch::Tensor, + torch::Tensor, + torch::Tensor, + torch::Tensor, + torch::Tensor, + std::optional> + intranode_dispatch(const torch::Tensor& x, + const std::optional& x_scales, + const std::optional& topk_idx, + const std::optional& topk_weights, + const std::optional& num_tokens_per_rank, + const torch::Tensor& is_token_in_rank, + const std::optional& num_tokens_per_expert, + int cached_num_recv_tokens, + const std::optional& cached_rank_prefix_matrix, + const std::optional& cached_channel_prefix_matrix, + int expert_alignment, + int num_worst_tokens, + const Config& config, + std::optional& previous_event, + bool async, + bool allocate_on_comm_stream) { + bool cached_mode = cached_rank_prefix_matrix.has_value(); + + // One channel use two blocks, even-numbered blocks for sending, odd-numbered blocks for receiving. + EP_HOST_ASSERT(config.num_sms % 2 == 0); + int num_channels = config.num_sms / 2; + if (cached_mode) { + EP_HOST_ASSERT(cached_rank_prefix_matrix.has_value()); + EP_HOST_ASSERT(cached_channel_prefix_matrix.has_value()); + } else { + EP_HOST_ASSERT(num_tokens_per_rank.has_value()); + EP_HOST_ASSERT(num_tokens_per_expert.has_value()); + } + + // Type checks + EP_HOST_ASSERT(is_token_in_rank.scalar_type() == torch::kBool); + if (cached_mode) { + EP_HOST_ASSERT(cached_rank_prefix_matrix->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(cached_channel_prefix_matrix->scalar_type() == torch::kInt32); + } else { + EP_HOST_ASSERT(num_tokens_per_expert->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(num_tokens_per_rank->scalar_type() == torch::kInt32); + } + + // Shape and contiguous checks + EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous()); + EP_HOST_ASSERT((x.size(1) * x.element_size()) % sizeof(int4) == 0); + EP_HOST_ASSERT(is_token_in_rank.dim() == 2 and is_token_in_rank.is_contiguous()); + EP_HOST_ASSERT(is_token_in_rank.size(0) == x.size(0) and is_token_in_rank.size(1) == num_ranks); + if (cached_mode) { + EP_HOST_ASSERT(cached_rank_prefix_matrix->dim() == 2 and cached_rank_prefix_matrix->is_contiguous()); + EP_HOST_ASSERT(cached_rank_prefix_matrix->size(0) == num_ranks and cached_rank_prefix_matrix->size(1) == num_ranks); + EP_HOST_ASSERT(cached_channel_prefix_matrix->dim() == 2 and cached_channel_prefix_matrix->is_contiguous()); + EP_HOST_ASSERT(cached_channel_prefix_matrix->size(0) == num_ranks and cached_channel_prefix_matrix->size(1) == num_channels); + } else { + EP_HOST_ASSERT(num_tokens_per_expert->dim() == 1 and num_tokens_per_expert->is_contiguous()); + EP_HOST_ASSERT(num_tokens_per_expert->size(0) % num_ranks == 0); + EP_HOST_ASSERT(num_tokens_per_expert->size(0) / num_ranks <= LEGACY_NUM_MAX_LOCAL_EXPERTS); + EP_HOST_ASSERT(num_tokens_per_rank->dim() == 1 and num_tokens_per_rank->is_contiguous()); + EP_HOST_ASSERT(num_tokens_per_rank->size(0) == num_ranks); + } + + auto num_tokens = static_cast(x.size(0)), hidden = static_cast(x.size(1)); + auto num_experts = cached_mode ? 0 : static_cast(num_tokens_per_expert->size(0)), num_local_experts = num_experts / num_ranks; + + // Top-k checks + int num_topk = 0; + topk_idx_t* topk_idx_ptr = nullptr; + float* topk_weights_ptr = nullptr; + EP_HOST_ASSERT(topk_idx.has_value() == topk_weights.has_value()); + if (topk_idx.has_value()) { + num_topk = static_cast(topk_idx->size(1)); + EP_HOST_ASSERT(num_experts > 0); + EP_HOST_ASSERT(topk_idx->dim() == 2 and topk_idx->is_contiguous()); + EP_HOST_ASSERT(topk_weights->dim() == 2 and topk_weights->is_contiguous()); + EP_HOST_ASSERT(num_tokens == topk_idx->size(0) and num_tokens == topk_weights->size(0)); + EP_HOST_ASSERT(num_topk == topk_weights->size(1)); + EP_HOST_ASSERT(topk_weights->scalar_type() == torch::kFloat32); + topk_idx_ptr = topk_idx->data_ptr(); + topk_weights_ptr = topk_weights->data_ptr(); + } + + // FP8 scales checks + float* x_scales_ptr = nullptr; + int num_scales = 0, scale_token_stride = 0, scale_hidden_stride = 0; + if (x_scales.has_value()) { + EP_HOST_ASSERT(x.element_size() == 1); + EP_HOST_ASSERT(x_scales->scalar_type() == torch::kFloat32 or x_scales->scalar_type() == torch::kInt); + EP_HOST_ASSERT(x_scales->dim() == 2); + EP_HOST_ASSERT(x_scales->size(0) == num_tokens); + num_scales = x_scales->dim() == 1 ? 1 : static_cast(x_scales->size(1)); + x_scales_ptr = static_cast(x_scales->data_ptr()); + scale_token_stride = static_cast(x_scales->stride(0)); + scale_hidden_stride = static_cast(x_scales->stride(1)); + } + + // Allocate all tensors on comm stream if set + // NOTES: do not allocate tensors upfront! + auto compute_stream = at::cuda::getCurrentCUDAStream(); + if (allocate_on_comm_stream) { + EP_HOST_ASSERT(previous_event.has_value() and async); + at::cuda::setCurrentCUDAStream(comm_stream); + } + + // Wait previous tasks to be finished + if (previous_event.has_value()) { + stream_wait(comm_stream, previous_event.value()); + } else { + stream_wait(comm_stream, compute_stream); + } + + // Create handles (only return for non-cached mode) + int num_recv_tokens = -1; + auto rank_prefix_matrix = torch::Tensor(); + auto channel_prefix_matrix = torch::Tensor(); + std::vector num_recv_tokens_per_expert_list; + + // Barrier or send sizes + // To clean: channel start/end offset, head and tail + int num_memset_int = num_channels * num_ranks * 4; + if (cached_mode) { + num_recv_tokens = cached_num_recv_tokens; + rank_prefix_matrix = cached_rank_prefix_matrix.value(); + channel_prefix_matrix = cached_channel_prefix_matrix.value(); + + // Copy rank prefix matrix and clean flags + intranode::cached_notify_dispatch( + rank_prefix_matrix.data_ptr(), num_memset_int, buffer_ptrs_gpu, barrier_signal_ptrs_gpu, rank, num_ranks, comm_stream); + } else { + rank_prefix_matrix = torch::empty({num_ranks, num_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); + channel_prefix_matrix = torch::empty({num_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); + + // Send sizes + // Meta information: + // - Size prefix by ranks, shaped as `[num_ranks, num_ranks]` + // - Size prefix by experts (not used later), shaped as `[num_ranks, num_local_experts]` + // NOTES: no more token dropping in this version + *moe_recv_counter = -1; + for (int i = 0; i < num_local_experts; ++i) + moe_recv_expert_counter[i] = -1; + EP_HOST_ASSERT(num_ranks * (num_ranks + num_local_experts) * sizeof(int) <= num_nvl_bytes); + intranode::notify_dispatch(num_tokens_per_rank->data_ptr(), + moe_recv_counter_mapped, + num_ranks, + num_tokens_per_expert->data_ptr(), + moe_recv_expert_counter_mapped, + num_experts, + num_tokens, + is_token_in_rank.data_ptr(), + channel_prefix_matrix.data_ptr(), + rank_prefix_matrix.data_ptr(), + num_memset_int, + expert_alignment, + buffer_ptrs_gpu, + barrier_signal_ptrs_gpu, + rank, + comm_stream, + num_channels); + + if (num_worst_tokens > 0) { + // No CPU sync, just allocate the worst case + num_recv_tokens = num_worst_tokens; + + // Must be forward with top-k stuffs + EP_HOST_ASSERT(topk_idx.has_value()); + EP_HOST_ASSERT(topk_weights.has_value()); + } else { + // Synchronize total received tokens and tokens per expert + auto start_time = std::chrono::high_resolution_clock::now(); + while (true) { + // Read total count + num_recv_tokens = static_cast(*moe_recv_counter); + + // Read per-expert count + bool ready = (num_recv_tokens >= 0); + for (int i = 0; i < num_local_experts and ready; ++i) + ready &= moe_recv_expert_counter[i] >= 0; + + if (ready) + break; + + // Timeout check + if (std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - start_time).count() > + LEGACY_NUM_CPU_TIMEOUT_SECS) + throw std::runtime_error("DeepEP error: CPU recv timeout"); + } + num_recv_tokens_per_expert_list = std::vector(moe_recv_expert_counter, moe_recv_expert_counter + num_local_experts); + } + } + + // Allocate new tensors + auto recv_x = torch::empty({num_recv_tokens, hidden}, x.options()); + auto recv_src_idx = torch::empty({num_recv_tokens}, dtype(torch::kInt32).device(torch::kCUDA)); + auto recv_topk_idx = std::optional(), recv_topk_weights = std::optional(), + recv_x_scales = std::optional(); + auto recv_channel_prefix_matrix = torch::empty({num_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); + auto send_head = torch::empty({num_tokens, num_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); + + // Assign pointers + topk_idx_t* recv_topk_idx_ptr = nullptr; + float* recv_topk_weights_ptr = nullptr; + float* recv_x_scales_ptr = nullptr; + if (topk_idx.has_value()) { + recv_topk_idx = torch::empty({num_recv_tokens, num_topk}, topk_idx->options()); + recv_topk_weights = torch::empty({num_recv_tokens, num_topk}, topk_weights->options()); + recv_topk_idx_ptr = recv_topk_idx->data_ptr(); + recv_topk_weights_ptr = recv_topk_weights->data_ptr(); + } + if (x_scales.has_value()) { + recv_x_scales = x_scales->dim() == 1 ? torch::empty({num_recv_tokens}, x_scales->options()) + : torch::empty({num_recv_tokens, num_scales}, x_scales->options()); + recv_x_scales_ptr = static_cast(recv_x_scales->data_ptr()); + } + + // Dispatch + EP_HOST_ASSERT( + num_ranks * num_ranks * sizeof(int) + // Size prefix matrix + num_channels * num_ranks * sizeof(int) + // Channel start offset + num_channels * num_ranks * sizeof(int) + // Channel end offset + num_channels * num_ranks * sizeof(int) * 2 + // Queue head and tail + num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * hidden * recv_x.element_size() + // Data buffer + num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * sizeof(int) + // Source index buffer + num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * num_topk * sizeof(topk_idx_t) + // Top-k index buffer + num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * num_topk * sizeof(float) + // Top-k weight buffer + num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * sizeof(float) * num_scales // FP8 scale buffer + <= num_nvl_bytes); + intranode::dispatch(recv_x.data_ptr(), + recv_x_scales_ptr, + recv_src_idx.data_ptr(), + recv_topk_idx_ptr, + recv_topk_weights_ptr, + recv_channel_prefix_matrix.data_ptr(), + send_head.data_ptr(), + x.data_ptr(), + x_scales_ptr, + topk_idx_ptr, + topk_weights_ptr, + is_token_in_rank.data_ptr(), + channel_prefix_matrix.data_ptr(), + num_tokens, + num_worst_tokens, + static_cast(hidden * recv_x.element_size() / sizeof(int4)), + num_topk, + num_experts, + num_scales, + scale_token_stride, + scale_hidden_stride, + buffer_ptrs_gpu, + rank, + num_ranks, + comm_stream, + config.num_sms, + config.num_max_nvl_chunked_send_tokens, + config.num_max_nvl_chunked_recv_tokens); + + // Wait streams + std::optional event; + if (async) { + event = EventHandle(comm_stream); + for (auto& t : {x, + is_token_in_rank, + rank_prefix_matrix, + channel_prefix_matrix, + recv_x, + recv_src_idx, + recv_channel_prefix_matrix, + send_head}) { + t.record_stream(comm_stream); + if (allocate_on_comm_stream) + t.record_stream(compute_stream); + } + for (auto& to : {x_scales, + topk_idx, + topk_weights, + num_tokens_per_rank, + num_tokens_per_expert, + cached_channel_prefix_matrix, + cached_rank_prefix_matrix, + recv_topk_idx, + recv_topk_weights, + recv_x_scales}) { + to.has_value() ? to->record_stream(comm_stream) : void(); + if (allocate_on_comm_stream) + to.has_value() ? to->record_stream(compute_stream) : void(); + } + } else { + stream_wait(compute_stream, comm_stream); + } + + // Switch back compute stream + if (allocate_on_comm_stream) + at::cuda::setCurrentCUDAStream(compute_stream); + + // Return values + return {recv_x, + recv_x_scales, + recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert_list, + rank_prefix_matrix, + channel_prefix_matrix, + recv_channel_prefix_matrix, + recv_src_idx, + send_head, + event}; + } + + std::tuple, std::optional> intranode_combine( + const torch::Tensor& x, + const std::optional& topk_weights, + const std::optional& bias_0, + const std::optional& bias_1, + const torch::Tensor& src_idx, + const torch::Tensor& rank_prefix_matrix, + const torch::Tensor& channel_prefix_matrix, + const torch::Tensor& send_head, + const Config& config, + const std::optional& previous_event, + bool async, + bool allocate_on_comm_stream) { + EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous()); + EP_HOST_ASSERT(src_idx.dim() == 1 and src_idx.is_contiguous() and src_idx.scalar_type() == torch::kInt32); + EP_HOST_ASSERT(send_head.dim() == 2 and send_head.is_contiguous() and send_head.scalar_type() == torch::kInt32); + EP_HOST_ASSERT(rank_prefix_matrix.dim() == 2 and rank_prefix_matrix.is_contiguous() and + rank_prefix_matrix.scalar_type() == torch::kInt32); + EP_HOST_ASSERT(channel_prefix_matrix.dim() == 2 and channel_prefix_matrix.is_contiguous() and + channel_prefix_matrix.scalar_type() == torch::kInt32); + + // One channel use two blocks, even-numbered blocks for sending, odd-numbered blocks for receiving. + EP_HOST_ASSERT(config.num_sms % 2 == 0); + int num_channels = config.num_sms / 2; + + auto num_tokens = static_cast(x.size(0)), hidden = static_cast(x.size(1)); + auto num_recv_tokens = static_cast(send_head.size(0)); + EP_HOST_ASSERT(src_idx.size(0) == num_tokens); + EP_HOST_ASSERT(send_head.size(1) == num_ranks); + EP_HOST_ASSERT(rank_prefix_matrix.size(0) == num_ranks and rank_prefix_matrix.size(1) == num_ranks); + EP_HOST_ASSERT(channel_prefix_matrix.size(0) == num_ranks and channel_prefix_matrix.size(1) == num_channels); + EP_HOST_ASSERT((hidden * x.element_size()) % sizeof(int4) == 0); + + // Allocate all tensors on comm stream if set + // NOTES: do not allocate tensors upfront! + auto compute_stream = at::cuda::getCurrentCUDAStream(); + if (allocate_on_comm_stream) { + EP_HOST_ASSERT(previous_event.has_value() and async); + at::cuda::setCurrentCUDAStream(comm_stream); + } + + // Wait previous tasks to be finished + if (previous_event.has_value()) { + stream_wait(comm_stream, previous_event.value()); + } else { + stream_wait(comm_stream, compute_stream); + } + + int num_topk = 0; + auto recv_topk_weights = std::optional(); + float* topk_weights_ptr = nullptr; + float* recv_topk_weights_ptr = nullptr; + if (topk_weights.has_value()) { + EP_HOST_ASSERT(topk_weights->dim() == 2 and topk_weights->is_contiguous()); + EP_HOST_ASSERT(topk_weights->size(0) == num_tokens); + EP_HOST_ASSERT(topk_weights->scalar_type() == torch::kFloat32); + num_topk = static_cast(topk_weights->size(1)); + topk_weights_ptr = topk_weights->data_ptr(); + recv_topk_weights = torch::empty({num_recv_tokens, num_topk}, topk_weights->options()); + recv_topk_weights_ptr = recv_topk_weights->data_ptr(); + } + + // Launch barrier and reset queue head and tail + EP_HOST_ASSERT(num_channels * num_ranks * sizeof(int) * 2 <= num_nvl_bytes); + intranode::cached_notify_combine(buffer_ptrs_gpu, + send_head.data_ptr(), + num_channels, + num_recv_tokens, + num_channels * num_ranks * 2, + barrier_signal_ptrs_gpu, + rank, + num_ranks, + comm_stream); + + // Assign bias pointers + auto bias_opts = std::vector>({bias_0, bias_1}); + void* bias_ptrs[2] = {nullptr, nullptr}; + for (int i = 0; i < 2; ++i) + if (bias_opts[i].has_value()) { + auto bias = bias_opts[i].value(); + EP_HOST_ASSERT(bias.dim() == 2 and bias.is_contiguous()); + EP_HOST_ASSERT(bias.scalar_type() == x.scalar_type()); + EP_HOST_ASSERT(bias.size(0) == num_recv_tokens and bias.size(1) == hidden); + bias_ptrs[i] = bias.data_ptr(); + } + + // Combine data + auto recv_x = torch::empty({num_recv_tokens, hidden}, x.options()); + EP_HOST_ASSERT(num_channels * num_ranks * sizeof(int) * 2 + // Queue head and tail + num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * hidden * x.element_size() + // Data buffer + num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * sizeof(int) + // Source index buffer + num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * num_topk * sizeof(float) // Top-k weight buffer + <= num_nvl_bytes); + intranode::combine(at::cuda::ScalarTypeToCudaDataType(x.scalar_type()), + recv_x.data_ptr(), + recv_topk_weights_ptr, + x.data_ptr(), + topk_weights_ptr, + bias_ptrs[0], + bias_ptrs[1], + src_idx.data_ptr(), + rank_prefix_matrix.data_ptr(), + channel_prefix_matrix.data_ptr(), + send_head.data_ptr(), + num_tokens, + num_recv_tokens, + hidden, + num_topk, + buffer_ptrs_gpu, + rank, + num_ranks, + comm_stream, + config.num_sms, + config.num_max_nvl_chunked_send_tokens, + config.num_max_nvl_chunked_recv_tokens); + + // Wait streams + std::optional event; + if (async) { + event = EventHandle(comm_stream); + for (auto& t : {x, src_idx, send_head, rank_prefix_matrix, channel_prefix_matrix, recv_x}) { + t.record_stream(comm_stream); + if (allocate_on_comm_stream) + t.record_stream(compute_stream); + } + for (auto& to : {topk_weights, recv_topk_weights, bias_0, bias_1}) { + to.has_value() ? to->record_stream(comm_stream) : void(); + if (allocate_on_comm_stream) + to.has_value() ? to->record_stream(compute_stream) : void(); + } + } else { + stream_wait(compute_stream, comm_stream); + } + + // Switch back compute stream + if (allocate_on_comm_stream) + at::cuda::setCurrentCUDAStream(compute_stream); + + return {recv_x, recv_topk_weights, event}; + } + + std::tuple, + std::optional, + std::optional, + std::vector, + torch::Tensor, + torch::Tensor, + std::optional, + torch::Tensor, + std::optional, + torch::Tensor, + std::optional, + std::optional, + std::optional, + std::optional> + internode_dispatch(const torch::Tensor& x, + const std::optional& x_scales, + const std::optional& topk_idx, + const std::optional& topk_weights, + const std::optional& num_tokens_per_rank, + const std::optional& num_tokens_per_rdma_rank, + const torch::Tensor& is_token_in_rank, + const std::optional& num_tokens_per_expert, + int cached_num_recv_tokens, + int cached_num_rdma_recv_tokens, + const std::optional& cached_rdma_channel_prefix_matrix, + const std::optional& cached_recv_rdma_rank_prefix_sum, + const std::optional& cached_gbl_channel_prefix_matrix, + const std::optional& cached_recv_gbl_rank_prefix_sum, + int expert_alignment, + int num_worst_tokens, + const Config& config, + std::optional& previous_event, + bool async, + bool allocate_on_comm_stream) { + // In dispatch, CPU will busy-wait until GPU receive tensor size metadata from other ranks, which can be quite long. + // If users of DeepEP need to execute other Python code on other threads, such as KV transfer, their code will get stuck due to GIL + // unless we release GIL here. + pybind11::gil_scoped_release release; + + const int num_channels = config.num_sms / 2; + EP_HOST_ASSERT(config.num_sms % 2 == 0); + EP_HOST_ASSERT(0 < get_num_rdma_ranks() and get_num_rdma_ranks() <= LEGACY_NUM_MAX_RDMA_PEERS); + + bool cached_mode = cached_rdma_channel_prefix_matrix.has_value(); + if (cached_mode) { + EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix.has_value()); + EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum.has_value()); + EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix.has_value()); + EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum.has_value()); + } else { + EP_HOST_ASSERT(num_tokens_per_rank.has_value()); + EP_HOST_ASSERT(num_tokens_per_rdma_rank.has_value()); + EP_HOST_ASSERT(num_tokens_per_expert.has_value()); + } + + // Type checks + if (cached_mode) { + EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum->scalar_type() == torch::kInt32); + } else { + EP_HOST_ASSERT(num_tokens_per_rank->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(num_tokens_per_rdma_rank->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(num_tokens_per_expert->scalar_type() == torch::kInt32); + } + + // Shape and contiguous checks + EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous()); + EP_HOST_ASSERT((x.size(1) * x.element_size()) % sizeof(int4) == 0); + if (cached_mode) { + EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix->dim() == 2 and cached_rdma_channel_prefix_matrix->is_contiguous()); + EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix->size(0) == num_rdma_ranks and + cached_rdma_channel_prefix_matrix->size(1) == num_channels); + EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum->dim() == 1 and cached_recv_rdma_rank_prefix_sum->is_contiguous()); + EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum->size(0) == num_rdma_ranks); + EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix->dim() == 2 and cached_gbl_channel_prefix_matrix->is_contiguous()); + EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix->size(0) == num_ranks and + cached_gbl_channel_prefix_matrix->size(1) == num_channels); + EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum->dim() == 1 and cached_recv_gbl_rank_prefix_sum->is_contiguous()); + EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum->size(0) == num_ranks); + } else { + EP_HOST_ASSERT(num_tokens_per_rank->dim() == 1 and num_tokens_per_rank->is_contiguous()); + EP_HOST_ASSERT(num_tokens_per_rdma_rank->dim() == 1 and num_tokens_per_rdma_rank->is_contiguous()); + EP_HOST_ASSERT(num_tokens_per_expert->dim() == 1 and num_tokens_per_expert->is_contiguous()); + EP_HOST_ASSERT(num_tokens_per_rank->size(0) == num_ranks); + EP_HOST_ASSERT(num_tokens_per_rdma_rank->size(0) == num_rdma_ranks); + EP_HOST_ASSERT(num_tokens_per_expert->size(0) % num_ranks == 0); + EP_HOST_ASSERT(num_tokens_per_expert->size(0) / num_ranks <= LEGACY_NUM_MAX_LOCAL_EXPERTS); + } + + auto num_tokens = static_cast(x.size(0)), hidden = static_cast(x.size(1)), + hidden_int4 = static_cast(x.size(1) * x.element_size() / sizeof(int4)); + auto num_experts = cached_mode ? 0 : static_cast(num_tokens_per_expert->size(0)), num_local_experts = num_experts / num_ranks; + + // Top-k checks + int num_topk = 0; + topk_idx_t* topk_idx_ptr = nullptr; + float* topk_weights_ptr = nullptr; + EP_HOST_ASSERT(topk_idx.has_value() == topk_weights.has_value()); + if (topk_idx.has_value()) { + num_topk = static_cast(topk_idx->size(1)); + EP_HOST_ASSERT(num_experts > 0); + EP_HOST_ASSERT(topk_idx->dim() == 2 and topk_idx->is_contiguous()); + EP_HOST_ASSERT(topk_weights->dim() == 2 and topk_weights->is_contiguous()); + EP_HOST_ASSERT(num_tokens == topk_idx->size(0) and num_tokens == topk_weights->size(0)); + EP_HOST_ASSERT(num_topk == topk_weights->size(1)); + EP_HOST_ASSERT(topk_weights->scalar_type() == torch::kFloat32); + topk_idx_ptr = topk_idx->data_ptr(); + topk_weights_ptr = topk_weights->data_ptr(); + } + + // FP8 scales checks + float* x_scales_ptr = nullptr; + int num_scales = 0, scale_token_stride = 0, scale_hidden_stride = 0; + if (x_scales.has_value()) { + EP_HOST_ASSERT(x.element_size() == 1); + EP_HOST_ASSERT(x_scales->scalar_type() == torch::kFloat32 or x_scales->scalar_type() == torch::kInt); + EP_HOST_ASSERT(x_scales->dim() == 2); + EP_HOST_ASSERT(x_scales->size(0) == num_tokens); + num_scales = x_scales->dim() == 1 ? 1 : static_cast(x_scales->size(1)); + x_scales_ptr = static_cast(x_scales->data_ptr()); + scale_token_stride = static_cast(x_scales->stride(0)); + scale_hidden_stride = static_cast(x_scales->stride(1)); + } + + // Allocate all tensors on comm stream if set + // NOTES: do not allocate tensors upfront! + auto compute_stream = at::cuda::getCurrentCUDAStream(); + if (allocate_on_comm_stream) { + EP_HOST_ASSERT(previous_event.has_value() and async); + at::cuda::setCurrentCUDAStream(comm_stream); + } + + // Wait previous tasks to be finished + if (previous_event.has_value()) { + stream_wait(comm_stream, previous_event.value()); + } else { + stream_wait(comm_stream, compute_stream); + } + + // Create handles (only return for non-cached mode) + int num_recv_tokens = -1, num_rdma_recv_tokens = -1; + auto rdma_channel_prefix_matrix = torch::Tensor(); + auto recv_rdma_rank_prefix_sum = torch::Tensor(); + auto gbl_channel_prefix_matrix = torch::Tensor(); + auto recv_gbl_rank_prefix_sum = torch::Tensor(); + std::vector num_recv_tokens_per_expert_list; + + // Barrier or send sizes + if (cached_mode) { + num_recv_tokens = cached_num_recv_tokens; + num_rdma_recv_tokens = cached_num_rdma_recv_tokens; + rdma_channel_prefix_matrix = cached_rdma_channel_prefix_matrix.value(); + recv_rdma_rank_prefix_sum = cached_recv_rdma_rank_prefix_sum.value(); + gbl_channel_prefix_matrix = cached_gbl_channel_prefix_matrix.value(); + recv_gbl_rank_prefix_sum = cached_recv_gbl_rank_prefix_sum.value(); + + // Just a barrier and clean flags + internode::cached_notify(hidden_int4, + num_scales, + num_topk, + num_topk, + num_ranks, + num_channels, + 0, + nullptr, + nullptr, + nullptr, + nullptr, + rdma_buffer_ptr, + config.num_max_rdma_chunked_recv_tokens, + buffer_ptrs_gpu, + config.num_max_nvl_chunked_recv_tokens, + barrier_signal_ptrs_gpu, + rank, + comm_stream, + config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks), + num_nvl_bytes, + true, + low_latency_mode); + } else { + rdma_channel_prefix_matrix = torch::empty({num_rdma_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); + recv_rdma_rank_prefix_sum = torch::empty({num_rdma_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); + gbl_channel_prefix_matrix = torch::empty({num_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); + recv_gbl_rank_prefix_sum = torch::empty({num_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); + + // Send sizes + *moe_recv_counter = -1; + *moe_recv_rdma_counter = -1; + for (int i = 0; i < num_local_experts; ++i) + moe_recv_expert_counter[i] = -1; + internode::notify_dispatch(num_tokens_per_rank->data_ptr(), + moe_recv_counter_mapped, + num_ranks, + num_tokens_per_rdma_rank->data_ptr(), + moe_recv_rdma_counter_mapped, + num_tokens_per_expert->data_ptr(), + moe_recv_expert_counter_mapped, + num_experts, + is_token_in_rank.data_ptr(), + num_tokens, + num_worst_tokens, + num_channels, + hidden_int4, + num_scales, + num_topk, + expert_alignment, + rdma_channel_prefix_matrix.data_ptr(), + recv_rdma_rank_prefix_sum.data_ptr(), + gbl_channel_prefix_matrix.data_ptr(), + recv_gbl_rank_prefix_sum.data_ptr(), + rdma_buffer_ptr, + config.num_max_rdma_chunked_recv_tokens, + buffer_ptrs_gpu, + config.num_max_nvl_chunked_recv_tokens, + barrier_signal_ptrs_gpu, + rank, + comm_stream, + config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks), + num_nvl_bytes, + low_latency_mode); + + // Synchronize total received tokens and tokens per expert + auto start_time = std::chrono::high_resolution_clock::now(); + while (true) { + // Read total count + num_recv_tokens = static_cast(*moe_recv_counter); + num_rdma_recv_tokens = static_cast(*moe_recv_rdma_counter); + + // Read per-expert count + bool ready = (num_recv_tokens >= 0) and (num_rdma_recv_tokens >= 0); + for (int i = 0; i < num_local_experts and ready; ++i) + ready &= moe_recv_expert_counter[i] >= 0; + + if (ready) + break; + + // Timeout check + if (std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - start_time).count() > + LEGACY_NUM_CPU_TIMEOUT_SECS) { + printf("Global rank: %d, num_recv_tokens: %d, num_rdma_recv_tokens: %d\n", rank, num_recv_tokens, num_rdma_recv_tokens); + for (int i = 0; i < num_local_experts; ++i) + printf("moe_recv_expert_counter[%d]: %d\n", i, moe_recv_expert_counter[i]); + throw std::runtime_error("DeepEP error: timeout (dispatch CPU)"); + } + } + num_recv_tokens_per_expert_list = std::vector(moe_recv_expert_counter, moe_recv_expert_counter + num_local_experts); + } + + // Allocate new tensors + auto recv_x = torch::empty({num_recv_tokens, hidden}, x.options()); + auto recv_topk_idx = std::optional(), recv_topk_weights = std::optional(), + recv_x_scales = std::optional(); + auto recv_src_meta = std::optional(); + auto recv_rdma_channel_prefix_matrix = std::optional(); + auto recv_gbl_channel_prefix_matrix = std::optional(); + auto send_rdma_head = std::optional(); + auto send_nvl_head = std::optional(); + if (not cached_mode) { + recv_src_meta = torch::empty({num_recv_tokens, internode::get_source_meta_bytes()}, dtype(torch::kByte).device(torch::kCUDA)); + recv_rdma_channel_prefix_matrix = torch::empty({num_rdma_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); + recv_gbl_channel_prefix_matrix = torch::empty({num_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA)); + send_rdma_head = torch::empty({num_tokens, num_rdma_ranks}, dtype(torch::kInt32).device(torch::kCUDA)); + send_nvl_head = torch::empty({num_rdma_recv_tokens, LEGACY_NUM_MAX_NVL_PEERS}, dtype(torch::kInt32).device(torch::kCUDA)); + } + + // Assign pointers + topk_idx_t* recv_topk_idx_ptr = nullptr; + float* recv_topk_weights_ptr = nullptr; + float* recv_x_scales_ptr = nullptr; + if (topk_idx.has_value()) { + recv_topk_idx = torch::empty({num_recv_tokens, num_topk}, topk_idx->options()); + recv_topk_weights = torch::empty({num_recv_tokens, num_topk}, topk_weights->options()); + recv_topk_idx_ptr = recv_topk_idx->data_ptr(); + recv_topk_weights_ptr = recv_topk_weights->data_ptr(); + } + if (x_scales.has_value()) { + recv_x_scales = x_scales->dim() == 1 ? torch::empty({num_recv_tokens}, x_scales->options()) + : torch::empty({num_recv_tokens, num_scales}, x_scales->options()); + recv_x_scales_ptr = static_cast(recv_x_scales->data_ptr()); + } + + // Launch data dispatch + // NOTES: the buffer size checks are moved into the `.cu` file + internode::dispatch(recv_x.data_ptr(), + recv_x_scales_ptr, + recv_topk_idx_ptr, + recv_topk_weights_ptr, + cached_mode ? nullptr : recv_src_meta->data_ptr(), + x.data_ptr(), + x_scales_ptr, + topk_idx_ptr, + topk_weights_ptr, + cached_mode ? nullptr : send_rdma_head->data_ptr(), + cached_mode ? nullptr : send_nvl_head->data_ptr(), + cached_mode ? nullptr : recv_rdma_channel_prefix_matrix->data_ptr(), + cached_mode ? nullptr : recv_gbl_channel_prefix_matrix->data_ptr(), + rdma_channel_prefix_matrix.data_ptr(), + recv_rdma_rank_prefix_sum.data_ptr(), + gbl_channel_prefix_matrix.data_ptr(), + recv_gbl_rank_prefix_sum.data_ptr(), + is_token_in_rank.data_ptr(), + num_tokens, + num_worst_tokens, + hidden_int4, + num_scales, + num_topk, + num_experts, + scale_token_stride, + scale_hidden_stride, + rdma_buffer_ptr, + config.num_max_rdma_chunked_send_tokens, + config.num_max_rdma_chunked_recv_tokens, + buffer_ptrs_gpu, + config.num_max_nvl_chunked_send_tokens, + config.num_max_nvl_chunked_recv_tokens, + rank, + num_ranks, + cached_mode, + comm_stream, + num_channels, + low_latency_mode); + + // Wait streams + std::optional event; + if (async) { + event = EventHandle(comm_stream); + for (auto& t : {x, + is_token_in_rank, + recv_x, + rdma_channel_prefix_matrix, + recv_rdma_rank_prefix_sum, + gbl_channel_prefix_matrix, + recv_gbl_rank_prefix_sum}) { + t.record_stream(comm_stream); + if (allocate_on_comm_stream) + t.record_stream(compute_stream); + } + for (auto& to : {x_scales, + topk_idx, + topk_weights, + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + cached_rdma_channel_prefix_matrix, + cached_recv_rdma_rank_prefix_sum, + cached_gbl_channel_prefix_matrix, + cached_recv_gbl_rank_prefix_sum, + recv_topk_idx, + recv_topk_weights, + recv_x_scales, + recv_rdma_channel_prefix_matrix, + recv_gbl_channel_prefix_matrix, + send_rdma_head, + send_nvl_head, + recv_src_meta}) { + to.has_value() ? to->record_stream(comm_stream) : void(); + if (allocate_on_comm_stream) + to.has_value() ? to->record_stream(compute_stream) : void(); + } + } else { + stream_wait(compute_stream, comm_stream); + } + + // Switch back compute stream + if (allocate_on_comm_stream) + at::cuda::setCurrentCUDAStream(compute_stream); + + // Return values + return {recv_x, + recv_x_scales, + recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert_list, + rdma_channel_prefix_matrix, + gbl_channel_prefix_matrix, + recv_rdma_channel_prefix_matrix, + recv_rdma_rank_prefix_sum, + recv_gbl_channel_prefix_matrix, + recv_gbl_rank_prefix_sum, + recv_src_meta, + send_rdma_head, + send_nvl_head, + event}; + } + + std::tuple, std::optional> internode_combine( + const torch::Tensor& x, + const std::optional& topk_weights, + const std::optional& bias_0, + const std::optional& bias_1, + const torch::Tensor& src_meta, + const torch::Tensor& is_combined_token_in_rank, + const torch::Tensor& rdma_channel_prefix_matrix, + const torch::Tensor& rdma_rank_prefix_sum, + const torch::Tensor& gbl_channel_prefix_matrix, + const torch::Tensor& combined_rdma_head, + const torch::Tensor& combined_nvl_head, + const Config& config, + const std::optional& previous_event, + bool async, + bool allocate_on_comm_stream) { + const int num_channels = config.num_sms / 2; + EP_HOST_ASSERT(config.num_sms % 2 == 0); + + // Shape and contiguous checks + EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous()); + EP_HOST_ASSERT(src_meta.dim() == 2 and src_meta.is_contiguous() and src_meta.scalar_type() == torch::kByte); + EP_HOST_ASSERT(is_combined_token_in_rank.dim() == 2 and is_combined_token_in_rank.is_contiguous() and + is_combined_token_in_rank.scalar_type() == torch::kBool); + EP_HOST_ASSERT(rdma_channel_prefix_matrix.dim() == 2 and rdma_channel_prefix_matrix.is_contiguous() and + rdma_channel_prefix_matrix.scalar_type() == torch::kInt32); + EP_HOST_ASSERT(rdma_rank_prefix_sum.dim() == 1 and rdma_rank_prefix_sum.is_contiguous() and + rdma_rank_prefix_sum.scalar_type() == torch::kInt32); + EP_HOST_ASSERT(gbl_channel_prefix_matrix.dim() == 2 and gbl_channel_prefix_matrix.is_contiguous() and + gbl_channel_prefix_matrix.scalar_type() == torch::kInt32); + EP_HOST_ASSERT(combined_rdma_head.dim() == 2 and combined_rdma_head.is_contiguous() and + combined_rdma_head.scalar_type() == torch::kInt32); + EP_HOST_ASSERT(combined_nvl_head.dim() == 2 and combined_nvl_head.is_contiguous() and combined_nvl_head.scalar_type() == torch::kInt32); + + auto num_tokens = static_cast(x.size(0)), hidden = static_cast(x.size(1)), + hidden_int4 = static_cast(x.size(1) * x.element_size() / sizeof(int4)); + auto num_combined_tokens = static_cast(is_combined_token_in_rank.size(0)); + EP_HOST_ASSERT((hidden * x.element_size()) % sizeof(int4) == 0); + EP_HOST_ASSERT(src_meta.size(1) == internode::get_source_meta_bytes()); + EP_HOST_ASSERT(is_combined_token_in_rank.size(1) == num_ranks); + EP_HOST_ASSERT(rdma_channel_prefix_matrix.size(0) == num_rdma_ranks and rdma_channel_prefix_matrix.size(1) == num_channels); + EP_HOST_ASSERT(rdma_rank_prefix_sum.size(0) == num_rdma_ranks); + EP_HOST_ASSERT(gbl_channel_prefix_matrix.size(0) == num_ranks and gbl_channel_prefix_matrix.size(1) == num_channels); + EP_HOST_ASSERT(combined_rdma_head.dim() == 2 and combined_rdma_head.size(0) == num_combined_tokens and + combined_rdma_head.size(1) == num_rdma_ranks); + EP_HOST_ASSERT(combined_nvl_head.dim() == 2 and combined_nvl_head.size(1) == LEGACY_NUM_MAX_NVL_PEERS); + + // Allocate all tensors on comm stream if set + // NOTES: do not allocate tensors upfront! + auto compute_stream = at::cuda::getCurrentCUDAStream(); + if (allocate_on_comm_stream) { + EP_HOST_ASSERT(previous_event.has_value() and async); + at::cuda::setCurrentCUDAStream(comm_stream); + } + + // Wait previous tasks to be finished + if (previous_event.has_value()) { + stream_wait(comm_stream, previous_event.value()); + } else { + stream_wait(comm_stream, compute_stream); + } + + // Top-k checks + int num_topk = 0; + auto combined_topk_weights = std::optional(); + float* topk_weights_ptr = nullptr; + float* combined_topk_weights_ptr = nullptr; + if (topk_weights.has_value()) { + EP_HOST_ASSERT(topk_weights->dim() == 2 and topk_weights->is_contiguous()); + EP_HOST_ASSERT(topk_weights->size(0) == num_tokens); + EP_HOST_ASSERT(topk_weights->scalar_type() == torch::kFloat32); + num_topk = static_cast(topk_weights->size(1)); + topk_weights_ptr = topk_weights->data_ptr(); + combined_topk_weights = torch::empty({num_combined_tokens, num_topk}, topk_weights->options()); + combined_topk_weights_ptr = combined_topk_weights->data_ptr(); + } + + // Extra check for avoid-dead-lock design + EP_HOST_ASSERT(config.num_max_nvl_chunked_recv_tokens % num_rdma_ranks == 0); + EP_HOST_ASSERT(config.num_max_nvl_chunked_send_tokens <= config.num_max_nvl_chunked_recv_tokens / num_rdma_ranks); + + // Launch barrier and reset queue head and tail + internode::cached_notify(hidden_int4, + 0, + 0, + num_topk, + num_ranks, + num_channels, + num_combined_tokens, + combined_rdma_head.data_ptr(), + rdma_channel_prefix_matrix.data_ptr(), + rdma_rank_prefix_sum.data_ptr(), + combined_nvl_head.data_ptr(), + rdma_buffer_ptr, + config.num_max_rdma_chunked_recv_tokens, + buffer_ptrs_gpu, + config.num_max_nvl_chunked_recv_tokens, + barrier_signal_ptrs_gpu, + rank, + comm_stream, + config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks), + num_nvl_bytes, + false, + low_latency_mode); + + // Assign bias pointers + auto bias_opts = std::vector>({bias_0, bias_1}); + void* bias_ptrs[2] = {nullptr, nullptr}; + for (int i = 0; i < 2; ++i) + if (bias_opts[i].has_value()) { + auto bias = bias_opts[i].value(); + EP_HOST_ASSERT(bias.dim() == 2 and bias.is_contiguous()); + EP_HOST_ASSERT(bias.scalar_type() == x.scalar_type()); + EP_HOST_ASSERT(bias.size(0) == num_combined_tokens and bias.size(1) == hidden); + bias_ptrs[i] = bias.data_ptr(); + } + + // Launch data combine + auto combined_x = torch::empty({num_combined_tokens, hidden}, x.options()); + internode::combine(at::cuda::ScalarTypeToCudaDataType(x.scalar_type()), + combined_x.data_ptr(), + combined_topk_weights_ptr, + is_combined_token_in_rank.data_ptr(), + x.data_ptr(), + topk_weights_ptr, + bias_ptrs[0], + bias_ptrs[1], + combined_rdma_head.data_ptr(), + combined_nvl_head.data_ptr(), + src_meta.data_ptr(), + rdma_channel_prefix_matrix.data_ptr(), + rdma_rank_prefix_sum.data_ptr(), + gbl_channel_prefix_matrix.data_ptr(), + num_tokens, + num_combined_tokens, + hidden, + num_topk, + rdma_buffer_ptr, + config.num_max_rdma_chunked_send_tokens, + config.num_max_rdma_chunked_recv_tokens, + buffer_ptrs_gpu, + config.num_max_nvl_chunked_send_tokens, + config.num_max_nvl_chunked_recv_tokens, + rank, + num_ranks, + comm_stream, + num_channels, + low_latency_mode); + + // Wait streams + std::optional event; + if (async) { + event = EventHandle(comm_stream); + for (auto& t : {x, + src_meta, + is_combined_token_in_rank, + rdma_channel_prefix_matrix, + rdma_rank_prefix_sum, + gbl_channel_prefix_matrix, + combined_x, + combined_rdma_head, + combined_nvl_head}) { + t.record_stream(comm_stream); + if (allocate_on_comm_stream) + t.record_stream(compute_stream); + } + for (auto& to : {topk_weights, combined_topk_weights, bias_0, bias_1}) { + to.has_value() ? to->record_stream(comm_stream) : void(); + if (allocate_on_comm_stream) + to.has_value() ? to->record_stream(compute_stream) : void(); + } + } else { + stream_wait(compute_stream, comm_stream); + } + + // Switch back compute stream + if (allocate_on_comm_stream) + at::cuda::setCurrentCUDAStream(compute_stream); + + // Return values + return {combined_x, combined_topk_weights, event}; + } + + void clean_low_latency_buffer(int num_max_dispatch_tokens_per_rank, int hidden, int num_experts) { + EP_HOST_ASSERT(low_latency_mode); + + auto layout = LowLatencyLayout(rdma_buffer_ptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts); + auto clean_meta_0 = layout.buffers[0].clean_meta(); + auto clean_meta_1 = layout.buffers[1].clean_meta(); + + auto check_boundary = [=, this](void* ptr, size_t num_bytes) { + auto offset = reinterpret_cast(ptr) - reinterpret_cast(rdma_buffer_ptr); + EP_HOST_ASSERT(0 <= offset and offset + num_bytes <= num_rdma_bytes); + }; + check_boundary(clean_meta_0.first, clean_meta_0.second * sizeof(int)); + check_boundary(clean_meta_1.first, clean_meta_1.second * sizeof(int)); + + internode_ll::clean_low_latency_buffer(clean_meta_0.first, + clean_meta_0.second, + clean_meta_1.first, + clean_meta_1.second, + rank, + num_ranks, + mask_buffer_ptr, + sync_buffer_ptr, + at::cuda::getCurrentCUDAStream()); + } + + std::tuple, + torch::Tensor, + torch::Tensor, + torch::Tensor, + std::optional, + std::optional>> + low_latency_dispatch(const torch::Tensor& x, + const torch::Tensor& topk_idx, + const std::optional& cumulative_local_expert_recv_stats, + const std::optional& dispatch_wait_recv_cost_stats, + int num_max_dispatch_tokens_per_rank, + int num_experts, + bool use_fp8, + bool round_scale, + bool use_ue8m0, + bool async, + bool return_recv_hook) { + EP_HOST_ASSERT(low_latency_mode); + + // Tensor checks + // By default using `ptp128c` FP8 cast + EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous() and x.scalar_type() == torch::kBFloat16); + EP_HOST_ASSERT(x.size(1) % sizeof(int4) == 0 and x.size(1) % 128 == 0); + EP_HOST_ASSERT(topk_idx.dim() == 2 and topk_idx.is_contiguous()); + EP_HOST_ASSERT(x.size(0) == topk_idx.size(0) and x.size(0) <= num_max_dispatch_tokens_per_rank); + EP_HOST_ASSERT(topk_idx.scalar_type() == c10::CppTypeToScalarType::value); + EP_HOST_ASSERT(num_experts % num_ranks == 0); + + // Diagnosis tensors + if (cumulative_local_expert_recv_stats.has_value()) { + EP_HOST_ASSERT(cumulative_local_expert_recv_stats->scalar_type() == torch::kInt); + EP_HOST_ASSERT(cumulative_local_expert_recv_stats->dim() == 1 and cumulative_local_expert_recv_stats->is_contiguous()); + EP_HOST_ASSERT(cumulative_local_expert_recv_stats->size(0) == num_experts / num_ranks); + } + if (dispatch_wait_recv_cost_stats.has_value()) { + EP_HOST_ASSERT(dispatch_wait_recv_cost_stats->scalar_type() == torch::kInt64); + EP_HOST_ASSERT(dispatch_wait_recv_cost_stats->dim() == 1 and dispatch_wait_recv_cost_stats->is_contiguous()); + EP_HOST_ASSERT(dispatch_wait_recv_cost_stats->size(0) == num_ranks); + } + + auto num_tokens = static_cast(x.size(0)), hidden = static_cast(x.size(1)); + auto num_topk = static_cast(topk_idx.size(1)); + auto num_local_experts = num_experts / num_ranks; + + // Buffer control + LowLatencyLayout layout(rdma_buffer_ptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts); + EP_HOST_ASSERT(layout.total_bytes <= num_rdma_bytes); + auto buffer = layout.buffers[low_latency_buffer_idx]; + auto next_buffer = layout.buffers[low_latency_buffer_idx ^= 1]; + + // Wait previous tasks to be finished + // NOTES: the hook mode will always use the default stream + auto compute_stream = at::cuda::getCurrentCUDAStream(); + auto launch_stream = return_recv_hook ? compute_stream : comm_stream; + EP_HOST_ASSERT(not(async and return_recv_hook)); + if (not return_recv_hook) + stream_wait(launch_stream, compute_stream); + + // Allocate packed tensors + auto packed_recv_x = torch::empty({num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank, hidden}, + x.options().dtype(use_fp8 ? torch::kFloat8_e4m3fn : torch::kBFloat16)); + auto packed_recv_src_info = + torch::empty({num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank}, torch::dtype(torch::kInt32).device(torch::kCUDA)); + auto packed_recv_layout_range = torch::empty({num_local_experts, num_ranks}, torch::dtype(torch::kInt64).device(torch::kCUDA)); + auto packed_recv_count = torch::empty({num_local_experts}, torch::dtype(torch::kInt32).device(torch::kCUDA)); + + // Allocate column-majored scales + auto packed_recv_x_scales = std::optional(); + void* packed_recv_x_scales_ptr = nullptr; + EP_HOST_ASSERT((num_ranks * num_max_dispatch_tokens_per_rank) % 4 == 0 and "TMA requires the number of tokens to be multiple of 4"); + + if (use_fp8) { + // TODO: support unaligned cases + EP_HOST_ASSERT(hidden % 512 == 0); + if (not use_ue8m0) { + packed_recv_x_scales = torch::empty({num_local_experts, hidden / 128, num_ranks * num_max_dispatch_tokens_per_rank}, + torch::dtype(torch::kFloat32).device(torch::kCUDA)); + } else { + EP_HOST_ASSERT(round_scale); + packed_recv_x_scales = torch::empty({num_local_experts, hidden / 512, num_ranks * num_max_dispatch_tokens_per_rank}, + torch::dtype(torch::kInt).device(torch::kCUDA)); + } + packed_recv_x_scales = torch::transpose(packed_recv_x_scales.value(), 1, 2); + packed_recv_x_scales_ptr = packed_recv_x_scales->data_ptr(); + } + + // Kernel launch + auto next_clean_meta = next_buffer.clean_meta(); + auto launcher = [=, this](int phases) { + internode_ll::dispatch( + packed_recv_x.data_ptr(), + packed_recv_x_scales_ptr, + packed_recv_src_info.data_ptr(), + packed_recv_layout_range.data_ptr(), + packed_recv_count.data_ptr(), + mask_buffer_ptr, + cumulative_local_expert_recv_stats.has_value() ? cumulative_local_expert_recv_stats->data_ptr() : nullptr, + dispatch_wait_recv_cost_stats.has_value() ? dispatch_wait_recv_cost_stats->data_ptr() : nullptr, + buffer.dispatch_rdma_recv_data_buffer, + buffer.dispatch_rdma_recv_count_buffer, + buffer.dispatch_rdma_send_buffer, + x.data_ptr(), + topk_idx.data_ptr(), + next_clean_meta.first, + next_clean_meta.second, + num_tokens, + hidden, + num_max_dispatch_tokens_per_rank, + num_topk, + num_experts, + rank, + num_ranks, + use_fp8, + round_scale, + use_ue8m0, + workspace, + num_device_sms, + launch_stream, + phases); + }; + launcher(return_recv_hook ? LEGACY_LOW_LATENCY_SEND_PHASE : (LEGACY_LOW_LATENCY_SEND_PHASE | LEGACY_LOW_LATENCY_RECV_PHASE)); + + // Wait streams + std::optional event; + if (async) { + // NOTES: we must ensure the all tensors will not be deallocated before the stream-wait happens, + // so in Python API, we must wrap all tensors into the event handle. + event = EventHandle(launch_stream); + } else if (not return_recv_hook) { + stream_wait(compute_stream, launch_stream); + } + + // Receiver callback + std::optional> recv_hook = std::nullopt; + if (return_recv_hook) + recv_hook = [=]() { launcher(LEGACY_LOW_LATENCY_RECV_PHASE); }; + + // Return values + return {packed_recv_x, packed_recv_x_scales, packed_recv_count, packed_recv_src_info, packed_recv_layout_range, event, recv_hook}; + } + + std::tuple, std::optional>> low_latency_combine( + const torch::Tensor& x, + const torch::Tensor& topk_idx, + const torch::Tensor& topk_weights, + const torch::Tensor& src_info, + const torch::Tensor& layout_range, + const std::optional& combine_wait_recv_cost_stats, + int num_max_dispatch_tokens_per_rank, + int num_experts, + bool use_logfmt, + bool zero_copy, + bool async, + bool return_recv_hook, + const std::optional& out = std::nullopt) { + EP_HOST_ASSERT(low_latency_mode); + + // Tensor checks + EP_HOST_ASSERT(x.dim() == 3 and x.is_contiguous() and x.scalar_type() == torch::kBFloat16); + EP_HOST_ASSERT(x.size(0) == num_experts / num_ranks); + EP_HOST_ASSERT(x.size(1) == num_ranks * num_max_dispatch_tokens_per_rank); + EP_HOST_ASSERT(x.size(2) % sizeof(int4) == 0 and x.size(2) % 128 == 0); + EP_HOST_ASSERT(topk_idx.dim() == 2 and topk_idx.is_contiguous()); + EP_HOST_ASSERT(topk_idx.size(0) == topk_weights.size(0) and topk_idx.size(1) == topk_weights.size(1)); + EP_HOST_ASSERT(topk_idx.scalar_type() == c10::CppTypeToScalarType::value); + EP_HOST_ASSERT(topk_weights.dim() == 2 and topk_weights.is_contiguous()); + EP_HOST_ASSERT(topk_weights.size(0) <= num_max_dispatch_tokens_per_rank); + EP_HOST_ASSERT(topk_weights.scalar_type() == torch::kFloat32); + EP_HOST_ASSERT(src_info.dim() == 2 and src_info.is_contiguous()); + EP_HOST_ASSERT(src_info.scalar_type() == torch::kInt32 and x.size(0) == src_info.size(0)); + EP_HOST_ASSERT(layout_range.dim() == 2 and layout_range.is_contiguous()); + EP_HOST_ASSERT(layout_range.scalar_type() == torch::kInt64); + EP_HOST_ASSERT(layout_range.size(0) == num_experts / num_ranks and layout_range.size(1) == num_ranks); + + if (combine_wait_recv_cost_stats.has_value()) { + EP_HOST_ASSERT(combine_wait_recv_cost_stats->scalar_type() == torch::kInt64); + EP_HOST_ASSERT(combine_wait_recv_cost_stats->dim() == 1 and combine_wait_recv_cost_stats->is_contiguous()); + EP_HOST_ASSERT(combine_wait_recv_cost_stats->size(0) == num_ranks); + } + + auto hidden = static_cast(x.size(2)); + auto num_topk = static_cast(topk_weights.size(1)); + auto num_combined_tokens = static_cast(topk_weights.size(0)); + + // Buffer control + LowLatencyLayout layout(rdma_buffer_ptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts); + EP_HOST_ASSERT(layout.total_bytes <= num_rdma_bytes); + auto buffer = layout.buffers[low_latency_buffer_idx]; + auto next_buffer = layout.buffers[low_latency_buffer_idx ^= 1]; + + // Wait previous tasks to be finished + // NOTES: the hook mode will always use the default stream + auto compute_stream = at::cuda::getCurrentCUDAStream(); + auto launch_stream = return_recv_hook ? compute_stream : comm_stream; + EP_HOST_ASSERT(not(async and return_recv_hook)); + if (not return_recv_hook) + stream_wait(launch_stream, compute_stream); + + // Allocate output tensor + torch::Tensor combined_x; + if (out.has_value()) { + EP_HOST_ASSERT(out->dim() == 2 and out->is_contiguous()); + EP_HOST_ASSERT(out->size(0) == num_combined_tokens and out->size(1) == hidden); + EP_HOST_ASSERT(out->scalar_type() == x.scalar_type()); + combined_x = out.value(); + } else { + combined_x = torch::empty({num_combined_tokens, hidden}, x.options()); + } + + // Kernel launch + auto next_clean_meta = next_buffer.clean_meta(); + auto launcher = [=, this](int phases) { + internode_ll::combine(combined_x.data_ptr(), + buffer.combine_rdma_recv_data_buffer, + buffer.combine_rdma_recv_flag_buffer, + buffer.combine_rdma_send_buffer, + x.data_ptr(), + topk_idx.data_ptr(), + topk_weights.data_ptr(), + src_info.data_ptr(), + layout_range.data_ptr(), + mask_buffer_ptr, + combine_wait_recv_cost_stats.has_value() ? combine_wait_recv_cost_stats->data_ptr() : nullptr, + next_clean_meta.first, + next_clean_meta.second, + num_combined_tokens, + hidden, + num_max_dispatch_tokens_per_rank, + num_topk, + num_experts, + rank, + num_ranks, + use_logfmt, + workspace, + num_device_sms, + launch_stream, + phases, + zero_copy); + }; + launcher(return_recv_hook ? LEGACY_LOW_LATENCY_SEND_PHASE : (LEGACY_LOW_LATENCY_SEND_PHASE | LEGACY_LOW_LATENCY_RECV_PHASE)); + + // Wait streams + std::optional event; + if (async) { + // NOTES: we must ensure the all tensors will not be deallocated before the stream-wait happens, + // so in Python API, we must wrap all tensors into the event handle. + event = EventHandle(launch_stream); + } else if (not return_recv_hook) { + stream_wait(compute_stream, launch_stream); + } + + // Receiver callback + std::optional> recv_hook = std::nullopt; + if (return_recv_hook) + recv_hook = [=]() { launcher(LEGACY_LOW_LATENCY_RECV_PHASE); }; + + // Return values + return {combined_x, event, recv_hook}; + } + + torch::Tensor get_next_low_latency_combine_buffer(int num_max_dispatch_tokens_per_rank, int hidden, int num_experts) const { + LowLatencyLayout layout(rdma_buffer_ptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts); + + auto buffer = layout.buffers[low_latency_buffer_idx]; + auto dtype = torch::kBFloat16; + auto num_msg_elems = static_cast(buffer.num_bytes_per_combine_msg / elementSize(torch::kBFloat16)); + + EP_HOST_ASSERT(buffer.num_bytes_per_combine_msg % elementSize(torch::kBFloat16) == 0); + return torch::from_blob(buffer.combine_rdma_send_buffer_data_start, + {num_experts / num_ranks, num_ranks * num_max_dispatch_tokens_per_rank, hidden}, + {num_ranks * num_max_dispatch_tokens_per_rank * num_msg_elems, num_msg_elems, 1}, + torch::TensorOptions().dtype(dtype).device(torch::kCUDA)); + } + + void low_latency_update_mask_buffer(int rank_to_mask, bool mask) const { + EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled"); + EP_HOST_ASSERT(rank_to_mask >= 0 and rank_to_mask < num_ranks); + internode_ll::update_mask_buffer(mask_buffer_ptr, rank_to_mask, mask, at::cuda::getCurrentCUDAStream()); + } + + void low_latency_query_mask_buffer(const torch::Tensor& mask_status) const { + EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled"); + EP_HOST_ASSERT(mask_status.numel() == num_ranks && mask_status.scalar_type() == torch::kInt32); + + internode_ll::query_mask_buffer( + mask_buffer_ptr, num_ranks, static_cast(mask_status.data_ptr()), at::cuda::getCurrentCUDAStream()); + } + + void low_latency_clean_mask_buffer() const { + EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled"); + internode_ll::clean_mask_buffer(mask_buffer_ptr, num_ranks, at::cuda::getCurrentCUDAStream()); + } +}; + +static void register_apis(pybind11::module_& m) { + pybind11::class_(m, "Config") + .def(pybind11::init(), + py::arg("num_sms") = 20, + py::arg("num_max_nvl_chunked_send_tokens") = 6, + py::arg("num_max_nvl_chunked_recv_tokens") = 256, + py::arg("num_max_rdma_chunked_send_tokens") = 6, + py::arg("num_max_rdma_chunked_recv_tokens") = 256) + .def("get_nvl_buffer_size_hint", &Config::get_nvl_buffer_size_hint) + .def("get_rdma_buffer_size_hint", &Config::get_rdma_buffer_size_hint); + m.def("get_low_latency_rdma_size_hint", &get_low_latency_rdma_size_hint); + + pybind11::class_(m, "EventHandle") + .def(pybind11::init<>()) + .def("current_stream_wait", &EventHandle::current_stream_wait); + + pybind11::class_(m, "Buffer") + .def(pybind11::init()) + .def("is_available", &Buffer::is_available) + .def("get_num_rdma_ranks", &Buffer::get_num_rdma_ranks) + .def("get_rdma_rank", &Buffer::get_rdma_rank) + .def("get_root_rdma_rank", &Buffer::get_root_rdma_rank) + .def("get_local_device_id", &Buffer::get_local_device_id) + .def("get_local_ipc_handle", &Buffer::get_local_ipc_handle) + .def("get_local_nvshmem_unique_id", &Buffer::get_local_nvshmem_unique_id) + .def("get_local_buffer_tensor", &Buffer::get_local_buffer_tensor) + .def("get_comm_stream", &Buffer::get_comm_stream) + .def("sync", &Buffer::sync) + .def("destroy", &Buffer::destroy) + .def("get_dispatch_layout", &Buffer::get_dispatch_layout) + .def("intranode_dispatch", &Buffer::intranode_dispatch) + .def("intranode_combine", &Buffer::intranode_combine) + .def("internode_dispatch", &Buffer::internode_dispatch) + .def("internode_combine", &Buffer::internode_combine) + .def("clean_low_latency_buffer", &Buffer::clean_low_latency_buffer) + .def("low_latency_dispatch", &Buffer::low_latency_dispatch) + .def("low_latency_combine", &Buffer::low_latency_combine) + .def("low_latency_update_mask_buffer", &Buffer::low_latency_update_mask_buffer) + .def("low_latency_query_mask_buffer", &Buffer::low_latency_query_mask_buffer) + .def("low_latency_clean_mask_buffer", &Buffer::low_latency_clean_mask_buffer) + .def("get_next_low_latency_combine_buffer", &Buffer::get_next_low_latency_combine_buffer); +} + +} // namespace deep_ep::legacy diff --git a/csrc/config.hpp b/csrc/legacy/config.hpp similarity index 90% rename from csrc/config.hpp rename to csrc/legacy/config.hpp index 0e4f5b065..a0e0b0687 100644 --- a/csrc/config.hpp +++ b/csrc/legacy/config.hpp @@ -1,9 +1,10 @@ #pragma once -#include "kernels/api.cuh" -#include "kernels/exception.cuh" +#include -namespace deep_ep { +#include "../kernels/legacy/api.cuh" + +namespace deep_ep::legacy { template dtype_t ceil_div(dtype_t a, dtype_t b) { @@ -54,18 +55,16 @@ struct Config { // TODO: add assertions constexpr int kNumMaxTopK = 128; constexpr int kNumMaxScales = 128; - EP_HOST_ASSERT(num_ranks < NUM_MAX_NVL_PEERS or num_ranks % NUM_MAX_NVL_PEERS == 0); - EP_HOST_ASSERT(num_ranks <= NUM_MAX_NVL_PEERS or num_sms % 2 == 0); - const auto num_rdma_ranks = std::max(num_ranks / NUM_MAX_NVL_PEERS, 1); - const auto num_nvl_ranks = std::min(num_ranks, NUM_MAX_NVL_PEERS); + EP_HOST_ASSERT(num_ranks < LEGACY_NUM_MAX_NVL_PEERS or num_ranks % LEGACY_NUM_MAX_NVL_PEERS == 0); + EP_HOST_ASSERT(num_ranks <= LEGACY_NUM_MAX_NVL_PEERS or num_sms % 2 == 0); + const auto num_rdma_ranks = std::max(num_ranks / LEGACY_NUM_MAX_NVL_PEERS, 1); + const auto num_nvl_ranks = std::min(num_ranks, LEGACY_NUM_MAX_NVL_PEERS); const int num_channels = num_sms / 2; size_t num_bytes = 0; num_bytes += num_channels * num_nvl_ranks * (2 * num_rdma_ranks + 3) * sizeof(int); num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * hidden_bytes; -#ifndef DISABLE_NVSHMEM num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * internode::get_source_meta_bytes(); -#endif num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxTopK * sizeof(topk_idx_t); num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxTopK * sizeof(float); num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxScales * sizeof(float); @@ -74,22 +73,21 @@ struct Config { } size_t get_rdma_buffer_size_hint(int64_t hidden_bytes, int num_ranks) const { -#ifndef DISABLE_NVSHMEM // Legacy mode - if (num_ranks <= NUM_MAX_NVL_PEERS) + if (num_ranks <= LEGACY_NUM_MAX_NVL_PEERS) return 0; // Below are some assumptions // TODO: add assertions constexpr int kNumMaxTopK = 128; constexpr int kNumMaxScales = 128; - EP_HOST_ASSERT(num_ranks % NUM_MAX_NVL_PEERS == 0); + EP_HOST_ASSERT(num_ranks % LEGACY_NUM_MAX_NVL_PEERS == 0); EP_HOST_ASSERT(num_sms % 2 == 0); - const int num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; + const int num_rdma_ranks = num_ranks / LEGACY_NUM_MAX_NVL_PEERS; const int num_channels = num_sms / 2; size_t num_bytes = 0; - num_bytes += num_channels * num_rdma_ranks * (NUM_MAX_NVL_PEERS * 2 + 2) * 2 * sizeof(int); + num_bytes += num_channels * num_rdma_ranks * (LEGACY_NUM_MAX_NVL_PEERS * 2 + 2) * 2 * sizeof(int); num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * hidden_bytes * 2; num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * internode::get_source_meta_bytes() * 2; num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * kNumMaxTopK * sizeof(topk_idx_t) * 2; @@ -98,9 +96,6 @@ struct Config { num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * sizeof(int4) * 2; num_bytes = ((num_bytes + 127) / 128) * 128; return num_bytes; -#else - EP_HOST_ASSERT(false and "NVSHMEM is disable during compilation"); -#endif } }; @@ -189,7 +184,7 @@ struct LowLatencyLayout { size_t get_low_latency_rdma_size_hint(int num_max_dispatch_tokens_per_rank, int hidden, int num_ranks, int num_experts) { auto num_bytes = LowLatencyLayout(nullptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts).total_bytes; - return ((num_bytes + NUM_BUFFER_ALIGNMENT_BYTES) / NUM_BUFFER_ALIGNMENT_BYTES) * NUM_BUFFER_ALIGNMENT_BYTES; + return ((num_bytes + LEGACY_NUM_BUFFER_ALIGNMENT_BYTES) / LEGACY_NUM_BUFFER_ALIGNMENT_BYTES) * LEGACY_NUM_BUFFER_ALIGNMENT_BYTES; } } // namespace deep_ep diff --git a/csrc/python_api.cpp b/csrc/python_api.cpp new file mode 100644 index 000000000..c8c2378ef --- /dev/null +++ b/csrc/python_api.cpp @@ -0,0 +1,39 @@ +#include +#include + +#include + +#include "jit/api.hpp" +#include "elastic/buffer.hpp" +#include "legacy/buffer.hpp" + +#ifndef TORCH_EXTENSION_NAME +#define TORCH_EXTENSION_NAME _C +#endif + +bool is_sm90_compiled() { +#ifndef DISABLE_SM90_FEATURES + return true; +#else + return false; +#endif +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.doc() = "DeepEP: an efficient expert-parallel communication library"; + + // Whether support FP8 and TMA features + m.def("is_sm90_compiled", []() { return deep_ep::kEnableSM90Features; }); + + // The integer type of top-k indices + m.attr("topk_idx_t") = py::cast(c10::CppTypeToScalarType::value); + + // JIT API + deep_ep::jit::register_apis(m); + + // Register legacy buffer APIs + deep_ep::legacy::register_apis(m); + + // Register elastic buffer (DeepEP V2) APIs + deep_ep::elastic::register_apis(m); +} diff --git a/csrc/event.hpp b/csrc/utils/event.hpp similarity index 70% rename from csrc/event.hpp rename to csrc/utils/event.hpp index b0b4383bf..f0199484a 100644 --- a/csrc/event.hpp +++ b/csrc/utils/event.hpp @@ -1,13 +1,15 @@ -#include +#pragma once +#include #include -#include "kernels/exception.cuh" +#include namespace deep_ep { struct EventHandle { std::shared_ptr event; + std::vector> tensors_to_record; EventHandle() { event = std::make_shared(torch::kCUDA); @@ -24,18 +26,18 @@ struct EventHandle { void current_stream_wait() const { at::cuda::getCurrentCUDAStream().unwrap().wait(*event); } }; -torch::Event create_event(const at::cuda::CUDAStream& s) { +static torch::Event create_event(const at::cuda::CUDAStream& s) { auto event = torch::Event(torch::kCUDA); event.record(s); return event; } -void stream_wait(const at::cuda::CUDAStream& s_0, const at::cuda::CUDAStream& s_1) { +static void stream_wait(const at::cuda::CUDAStream& s_0, const at::cuda::CUDAStream& s_1) { EP_HOST_ASSERT(s_0.id() != s_1.id()); s_0.unwrap().wait(create_event(s_1)); } -void stream_wait(const at::cuda::CUDAStream& s, const EventHandle& event) { +static void stream_wait(const at::cuda::CUDAStream& s, const EventHandle& event) { s.unwrap().wait(*event.event); } diff --git a/csrc/utils/format.hpp b/csrc/utils/format.hpp new file mode 100644 index 000000000..bf617372b --- /dev/null +++ b/csrc/utils/format.hpp @@ -0,0 +1,6 @@ +#pragma once + +// Just a wrapper for the `fmt` headers +#define FMT_HEADER_ONLY +#include +#include diff --git a/csrc/utils/hash.hpp b/csrc/utils/hash.hpp new file mode 100644 index 000000000..61458ba30 --- /dev/null +++ b/csrc/utils/hash.hpp @@ -0,0 +1,40 @@ +#pragma once + +#include +#include + +namespace deep_ep { + +static uint64_t fnv1a(const std::vector& data, const uint64_t& seed) { + uint64_t h = seed; + constexpr uint64_t prime = 0x100000001b3ull; + for (const char& c: data) { + h ^= static_cast(c); + h *= prime; + } + return h; +} + +static std::string get_hex_digest(const std::vector& data) { + const auto state_0 = fnv1a(data, 0xc6a4a7935bd1e995ull); + const auto state_1 = fnv1a(data, 0x9e3779b97f4a7c15ull); + + // Split-mix 64 + const auto split_mix = [](uint64_t z) { + z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ull; + z = (z ^ (z >> 27)) * 0x94d049bb133111ebull; + return z ^ (z >> 31); + }; + + std::ostringstream oss; + oss << std::hex << std::setfill('0') + << std::setw(16) << split_mix(state_0) + << std::setw(16) << split_mix(state_1); + return oss.str(); +} + +static std::string get_hex_digest(const std::string& data) { + return get_hex_digest(std::vector{data.begin(), data.end()}); +} + +} // namespace deep_ep diff --git a/csrc/utils/lazy_driver.hpp b/csrc/utils/lazy_driver.hpp new file mode 100644 index 000000000..d4c0fa582 --- /dev/null +++ b/csrc/utils/lazy_driver.hpp @@ -0,0 +1,55 @@ +#pragma once + +#include +#include + +#include + +namespace deep_ep { + +// Lazy loading all driver symbols +static void* get_driver_handle() { + static void* handle = nullptr; + if (handle == nullptr) { + handle = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL); + EP_HOST_ASSERT(handle != nullptr and "Failed to load CUDA driver `libcuda.so.1`"); + } + return handle; +} + +// Macro to define wrapper functions named `lazy_cu{API name}` +#define DECL_LAZY_CUDA_DRIVER_FUNCTION(name) \ +template \ +static auto lazy_##name(Args&&... args) -> decltype(name(args...)) { \ + using FuncType = decltype(&name); \ + static FuncType func = nullptr; \ + if (func == nullptr) { \ + func = reinterpret_cast(dlsym(get_driver_handle(), #name)); \ + EP_HOST_ASSERT(func != nullptr and "Failed to load CUDA driver API"); \ + } \ + return func(std::forward(args)...); \ +} + +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorName); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorString); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuFuncSetAttribute); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleLoad); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleUnload); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleGetFunction); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLaunchKernelEx); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemSetAccess); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemRetainAllocationHandle); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemGetAddressRange_v2); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemAddressReserve); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemAddressFree); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemMap); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemUnmap); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemCreate); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemRelease); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuCtxGetDevice); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemImportFromShareableHandle); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemExportToShareableHandle); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemGetAllocationGranularity); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuStreamBatchMemOp); + +} // namespace deep_ep diff --git a/csrc/utils/lazy_init.hpp b/csrc/utils/lazy_init.hpp new file mode 100644 index 000000000..f9bd722c0 --- /dev/null +++ b/csrc/utils/lazy_init.hpp @@ -0,0 +1,27 @@ +#pragma once + +#include +#include + +#define EP_DECLARE_STATIC_VAR_IN_CLASS(cls, name) decltype(cls::name) cls::name + +namespace deep_ep { + +template +class LazyInit { +public: + explicit LazyInit(std::function()> factory) + : factory(std::move(factory)) {} + + T* operator -> () { + if (ptr == nullptr) + ptr = factory(); + return ptr.get(); + } + +private: + std::shared_ptr ptr; + std::function()> factory; +}; + +} // namespace deep_ep diff --git a/csrc/utils/shared_memory.hpp b/csrc/utils/shared_memory.hpp new file mode 100644 index 000000000..00fd256c4 --- /dev/null +++ b/csrc/utils/shared_memory.hpp @@ -0,0 +1,127 @@ +#pragma once + +#include +#include + +#include "lazy_driver.hpp" + +namespace deep_ep::shared_memory { + +union MemHandleInner { + cudaIpcMemHandle_t cuda_ipc_mem_handle; + CUmemFabricHandle cu_mem_fabric_handle; +}; + +struct MemHandle { + MemHandleInner inner; + size_t size; +}; + +static void cu_mem_set_access_all(void* ptr, size_t size) { + int device_count; + CUDA_RUNTIME_CHECK(cudaGetDeviceCount(&device_count)); + + constexpr int kMaxDeviceCount = 8; + EP_HOST_ASSERT(0 < device_count and device_count <= kMaxDeviceCount); + + CUmemAccessDesc access_desc[kMaxDeviceCount]; + for (int i = 0; i < device_count; ++ i) { + access_desc[i].location.type = CU_MEM_LOCATION_TYPE_DEVICE; + access_desc[i].location.id = i; + access_desc[i].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + } + CUDA_DRIVER_CHECK(lazy_cuMemSetAccess(reinterpret_cast(ptr), size, access_desc, device_count)); +} + +static void cu_mem_free(void* ptr) { + CUmemGenericAllocationHandle handle; + CUDA_DRIVER_CHECK(lazy_cuMemRetainAllocationHandle(&handle, ptr)); + + size_t size = 0; + CUDA_DRIVER_CHECK(lazy_cuMemGetAddressRange_v2(nullptr, &size, reinterpret_cast(ptr))); + + CUDA_DRIVER_CHECK(lazy_cuMemUnmap(reinterpret_cast(ptr), size)); + CUDA_DRIVER_CHECK(lazy_cuMemAddressFree(reinterpret_cast(ptr), size)); + CUDA_DRIVER_CHECK(lazy_cuMemRelease(handle)); +} + +class SharedMemoryAllocator { +public: + explicit SharedMemoryAllocator(const bool& use_fabric) : use_fabric(use_fabric) {} + + void malloc(void** ptr, size_t size) const { + if (use_fabric) { + CUdevice device; + CUDA_DRIVER_CHECK(lazy_cuCtxGetDevice(&device)); + + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; + prop.location.id = device; + + size_t alignment = 0; + EP_HOST_ASSERT(size > 0); + CUDA_DRIVER_CHECK(lazy_cuMemGetAllocationGranularity(&alignment, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM)); + size = ((size + alignment - 1) / alignment) * alignment; + + CUmemGenericAllocationHandle handle; + CUDA_DRIVER_CHECK(lazy_cuMemCreate(&handle, size, &prop, 0)); + CUDA_DRIVER_CHECK(lazy_cuMemAddressReserve(reinterpret_cast(ptr), size, alignment, 0, 0)); + CUDA_DRIVER_CHECK(lazy_cuMemMap(reinterpret_cast(*ptr), size, 0, handle, 0)); + cu_mem_set_access_all(*ptr, size); + } else { + CUDA_RUNTIME_CHECK(cudaMalloc(ptr, size)); + } + } + + void free(void* ptr) const { + if (use_fabric) { + cu_mem_free(ptr); + } else { + CUDA_RUNTIME_CHECK(cudaFree(ptr)); + } + } + + void get_mem_handle(MemHandle* mem_handle, void* ptr) const { + size_t size = 0; + CUDA_DRIVER_CHECK(lazy_cuMemGetAddressRange_v2(nullptr, &size, reinterpret_cast(ptr))); + mem_handle->size = size; + + if (use_fabric) { + CUmemGenericAllocationHandle handle; + CUDA_DRIVER_CHECK(lazy_cuMemRetainAllocationHandle(&handle, ptr)); + CUDA_DRIVER_CHECK(lazy_cuMemExportToShareableHandle(&mem_handle->inner.cu_mem_fabric_handle, handle, CU_MEM_HANDLE_TYPE_FABRIC, 0)); + } else { + CUDA_RUNTIME_CHECK(cudaIpcGetMemHandle(&mem_handle->inner.cuda_ipc_mem_handle, ptr)); + } + } + + void open_mem_handle(void** ptr, MemHandle* mem_handle) const { + if (use_fabric) { + size_t size = mem_handle->size; + + CUmemGenericAllocationHandle handle; + CUDA_DRIVER_CHECK(lazy_cuMemImportFromShareableHandle(&handle, &mem_handle->inner.cu_mem_fabric_handle, CU_MEM_HANDLE_TYPE_FABRIC)); + + CUDA_DRIVER_CHECK(lazy_cuMemAddressReserve(reinterpret_cast(ptr), size, 0, 0, 0)); + CUDA_DRIVER_CHECK(lazy_cuMemMap(reinterpret_cast(*ptr), size, 0, handle, 0)); + cu_mem_set_access_all(*ptr, size); + } else { + CUDA_RUNTIME_CHECK(cudaIpcOpenMemHandle(ptr, mem_handle->inner.cuda_ipc_mem_handle, cudaIpcMemLazyEnablePeerAccess)); + } + } + + void close_mem_handle(void* ptr) const { + if (use_fabric) { + cu_mem_free(ptr); + } else { + CUDA_RUNTIME_CHECK(cudaIpcCloseMemHandle(ptr)); + } + } + +private: + bool use_fabric; +}; + +} // namespace deep_ep::shared_memory diff --git a/csrc/utils/system.hpp b/csrc/utils/system.hpp new file mode 100644 index 000000000..7b410d467 --- /dev/null +++ b/csrc/utils/system.hpp @@ -0,0 +1,111 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "format.hpp" + +namespace deep_ep { + +// ReSharper disable once CppNotAllPathsReturnValue +template +static dtype_t get_env(const std::string& name, const dtype_t& default_value = dtype_t()) { + const auto c_str = std::getenv(name.c_str()); + if (c_str == nullptr) + return default_value; + + // Read the env and convert to the desired type + if constexpr (std::is_same_v) { + return std::string(c_str); + } else if constexpr (std::is_same_v) { + int value; + std::sscanf(c_str, "%d", &value); + return value; + } else { + EP_HOST_ASSERT(false and "Unexpected type"); + } +} + +static std::tuple call_external_command(std::string command) { + command = command + " 2>&1"; + const auto deleter = [](FILE* f) { if (f) pclose(f); }; + std::unique_ptr pipe(popen(command.c_str(), "r"), deleter); + EP_HOST_ASSERT(pipe != nullptr); + + std::array buffer; + std::string output; + while (fgets(buffer.data(), buffer.size(), pipe.get())) + output += buffer.data(); + const auto status = pclose(pipe.release()); + // NOTES: if the child was killed by a signal (e.g., SIGINT from Ctrl+C), + // WEXITSTATUS would incorrectly return 0. Treat signal death as failure. + const auto exit_code = WIFEXITED(status) ? WEXITSTATUS(status) : 128 + WTERMSIG(status); + return {exit_code, output}; +} + +static std::filesystem::path make_dirs(const std::filesystem::path& path) { + // OK if existed + std::error_code capture; + const bool created = std::filesystem::create_directories(path, capture); + if (not (created or capture.value() == 0)) { + EP_HOST_UNREACHABLE(fmt::format("Failed to make directory: {}, created: {}, value: {}", + path.c_str(), created, capture.value())); + } + if (created and get_env("EP_JIT_DEBUG")) + printf("Create directory: %s\n", path.c_str()); + return path; +} + +static std::string get_uuid() { + static std::random_device rd; + static std::mt19937 gen([]() { + return rd() ^ std::chrono::steady_clock::now().time_since_epoch().count(); + }()); + static std::uniform_int_distribution dist; + + std::stringstream ss; + ss << getpid() << "-" + << std::hex << std::setfill('0') + << std::setw(8) << dist(gen) << "-" + << std::setw(8) << dist(gen) << "-" + << std::setw(8) << dist(gen); + return ss.str(); +} + +static void safe_remove_all(const std::filesystem::path& path) { + std::error_code ec; + if (not std::filesystem::exists(path, ec) or ec) + return; + + // A single file + if (not std::filesystem::is_directory(path, ec) or ec) { + std::filesystem::remove(path, ec); + return; + } + + // Remove directory + auto it = std::filesystem::directory_iterator(path, + std::filesystem::directory_options::skip_permission_denied, ec); + for (auto end = std::filesystem::directory_iterator(); it != end and not ec;) { + const auto entry_path = it->path(); + + // Increase firstly to avoid failures + it.increment(ec); + if (ec) + break; + + // Recursively clean + safe_remove_all(entry_path); + } + std::filesystem::remove(path, ec); +} + +} // deep_ep diff --git a/deep_ep/__init__.py b/deep_ep/__init__.py index 2b20d832f..e0c439f98 100644 --- a/deep_ep/__init__.py +++ b/deep_ep/__init__.py @@ -1,7 +1,97 @@ +import filecmp +import functools +import glob +import subprocess import torch +import os -from .utils import EventOverlap -from .buffer import Buffer +from .utils.find_pkgs import find_nccl_root +# Set some default environment provided at setup +try: + # noinspection PyUnresolvedReferences + from .envs import persistent_envs + for key, value in persistent_envs.items(): + if key not in os.environ: + os.environ[key] = value +except ImportError: + pass + +# Initialize +@functools.lru_cache() +def find_cuda_home() -> str: + """ + Find the CUDA installation directory, cached. + + Returns: + cuda_home: the CUDA installation path. + """ + # TODO: reuse PyTorch API later + # For some PyTorch versions, the original `_find_cuda_home` will initialize CUDA, which is incompatible with process forks + cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH') + if cuda_home is None: + # noinspection PyBroadException + try: + with open(os.devnull, 'w') as devnull: + nvcc = subprocess.check_output(['which', 'nvcc'], stderr=devnull).decode().rstrip('\r\n') + cuda_home = os.path.dirname(os.path.dirname(nvcc)) + except Exception: + cuda_home = '/usr/local/cuda' + if not os.path.exists(cuda_home): + cuda_home = None + assert cuda_home is not None + return cuda_home + + +def check_nccl_so(): + """ + Verify that the NCCL library loaded at runtime matches the linked version. + Aborts if duplicate NCCL libraries are found or if versions mismatch. + """ + if int(os.environ.get('EP_SUPPRESS_NCCL_CHECK', 0)): + return + + # PyTorch may load another NCCL library, which is different to the linked one + with open('/proc/self/maps', 'r') as f: + loaded_nccl_so = None + for so in [line.strip().split(' ')[-1] for line in f if 'nccl' in line]: + loaded_nccl_so = so if loaded_nccl_so is None else loaded_nccl_so + assert so == loaded_nccl_so, f'Duplicate NCCL runtime found in the current system: {so} and {loaded_nccl_so}' + linked_nccl_so_candidates = sorted(glob.glob(f'{find_nccl_root()}/lib/libnccl.so*')) + assert linked_nccl_so_candidates, f'No libnccl.so found in {find_nccl_root()}/lib/' + linked_nccl_so = linked_nccl_so_candidates[0] + + # So checking binary-level equalness is necessary + # noinspection PyTypeChecker + assert filecmp.cmp(loaded_nccl_so, linked_nccl_so, shallow=False), \ + (f'Invalid NCCL versions: {loaded_nccl_so} (loaded) v.s. {linked_nccl_so} (expected), ' + f'please contact Chenggang or Shangyan to upgrade PyTorch NCCL version') + + +def init_jit(): + """ + Initialize the JIT compilation runtime. Sets up CUDA and NCCL root paths for the JIT compiler. + """ + # noinspection PyUnresolvedReferences + import deep_ep._C as _C + library_root_path = os.path.dirname(os.path.abspath(__file__)) + _C.init_jit(library_root_path, # Library root directory path + find_cuda_home(), # CUDA home + find_nccl_root()) # NCCL root + +# Run initialization +check_nccl_so() +init_jit() + + +# Import APIs after initialization +from .buffers.legacy import Buffer +from .buffers.elastic import ElasticBuffer, EPHandle # noinspection PyUnresolvedReferences -from deep_ep_cpp import Config, topk_idx_t +from .utils.event import EventOverlap, EventHandle +from .utils.envs import get_physical_domain_size, get_logical_domain_size + +# noinspection PyUnresolvedReferences +from deep_ep._C import Config, topk_idx_t + +__version__ = '2.0.0' diff --git a/deep_ep/buffers/__init__.py b/deep_ep/buffers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/deep_ep/buffers/elastic.py b/deep_ep/buffers/elastic.py new file mode 100644 index 000000000..cebd8d52b --- /dev/null +++ b/deep_ep/buffers/elastic.py @@ -0,0 +1,883 @@ +import os +import math +import torch +import torch.distributed as dist +from typing import Callable, Optional, Tuple, Union, List, Sequence +from contextlib import contextmanager + +# noinspection PyUnresolvedReferences +import deep_ep._C as _C +# noinspection PyUnresolvedReferences +from deep_ep._C import EventHandle + +from ..utils.event import EventOverlap +from ..utils.math import ceil_div, align +from ..utils.semantic import value_or, weak_lru +from ..utils.envs import ( + check_fast_rdma_atomic_support, + check_nvlink_connections, check_torch_deterministic, + get_nvlink_gbs, get_rdma_gbs +) +from ..utils.comm import get_nccl_comm_handle + + +class EPHandle: + """ + Communication handle returned by `ElasticBuffer.dispatch`. + Can be reused as a cached handle in subsequent `ElasticBuffer.dispatch` calls to skip layout recomputation, + and is consumed by `ElasticBuffer.combine` to reverse the token routing. + + Attributes: + do_expand: whether the expanding (one-token-per-expert-slot) layout is used. + num_experts: the number of all experts. + expert_alignment: align the number of tokens received by each local expert to this variable. + num_max_tokens_per_rank: the maximum number of tokens per rank, all the ranks must hold the same value. + num_sms: the SM count used during dispatch (reused in combine). + topk_idx: cloned top-k expert indices from dispatch, `[num_tokens, num_topk]`. + psum_num_recv_tokens_per_scaleup_rank: inclusive prefix sum of deduplicated received token counts + per scaleup rank, shape `[num_scaleup_ranks]`. A token is counted once per rank even if + multiple of its top-k experts land on the same rank. The last element equals the total number + of received tokens. + psum_num_recv_tokens_per_expert: prefix sum of alignment-padded received token counts per local + expert, shape `[num_local_experts]`. Each expert's count is padded to `expert_alignment`. + In non-expand mode, this is the inclusive prefix sum. In expand mode, `psum[i]` equals + the aligned cumulative count of experts before `i` plus the actual (unaligned) token count + of expert `i` — so `psum[i] - align(psum[i-1], expert_alignment)` recovers the real + count for expert `i`, and `align(psum[i], expert_alignment)` gives expert `i+1`'s + starting offset. + num_recv_tokens_per_expert_list: Python list of per-expert received token counts (CPU-side). + recv_src_metadata: source token indices and buffer slot indices. + dst_buffer_slot_idx: destination buffer slot indices from dispatch. + token_metadata_at_forward: per-channel forwarded token metadata (hybrid mode only). + channel_linked_list: per-channel per-scaleup-peer linked list (hybrid mode only). + num_recv_tokens: the total number of received tokens. + """ + + def __init__(self, + do_expand: bool, + num_experts: int, expert_alignment: int, + num_max_tokens_per_rank: int, + num_sms: int, + topk_idx: torch.Tensor, + num_recv_tokens_per_expert_list: list, + psum_num_recv_tokens_per_scaleup_rank: torch.Tensor, + psum_num_recv_tokens_per_expert: torch.Tensor, + recv_src_metadata: torch.Tensor, + dst_buffer_slot_idx: torch.Tensor, + token_metadata_at_forward: Optional[torch.Tensor], + channel_linked_list: Optional[torch.Tensor]): + # NOTES: remember to copy the original users' input to prevent uncasual modifications on them + assert topk_idx is not None + + self.do_expand = do_expand + self.num_experts = num_experts + self.expert_alignment = expert_alignment + self.num_max_tokens_per_rank = num_max_tokens_per_rank + self.num_sms = num_sms + self.topk_idx = topk_idx + self.psum_num_recv_tokens_per_scaleup_rank = psum_num_recv_tokens_per_scaleup_rank + self.psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert + self.num_recv_tokens_per_expert_list = num_recv_tokens_per_expert_list + self.recv_src_metadata = recv_src_metadata + self.dst_buffer_slot_idx = dst_buffer_slot_idx + self.token_metadata_at_forward = token_metadata_at_forward + self.channel_linked_list = channel_linked_list + + # Inferred value, may not accurate without CPU sync + self.num_recv_tokens = recv_src_metadata.shape[0] + + +class ElasticBuffer: + """ + The elastic communication buffer, which supports: + - high-throughput expert-parallel all-to-all (dispatch and combine, using NVLink and/or RDMA) + - Engram (remote KV cache fetch, using RDMA) + - pipeline-parallel send/recv (PP, using NVLink) + - all-gather reduce-scatter (AGRS, using NVLink) + "Elastic" refers to the flexibility of underlying memory: currently GPU-only, with CPU and mixed + (GPU+CPU) backends on the roadmap + + Attributes: + group: the communication group. + rank_idx: the rank index. + num_ranks: the number of ranks in the group. + allow_hybrid_mode: whether to enable hybrid mode for multi-node communication. Hybrid mode uses + hierarchical RDMA + NVLink communication to achieve higher bandwidth, and is more friendly + to multi-plane/multi-rail networks. + allow_multiple_reduction: whether to allow multiple reductions in combine. If disabled, + only one reduction will be done in the combine epilogue for best precision, + but it may increase data transfer size. + prefer_overlap_with_compute: whether to prefer overlapping communication with compute. + If enabled, we tend to use fewer SMs. + num_bytes: the total buffer size in bytes. + num_max_tokens_per_rank: the default maximum tokens per rank. + num_scaleout_ranks: the number of scaleout ranks. + num_scaleup_ranks: the number of scaleup ranks. + scaleout_rank_idx: the scaleout rank index of this rank. + scaleup_rank_idx: the scaleup rank index of this rank. + num_rdma_ranks: the number of physical RDMA ranks. + num_nvlink_ranks: the number of physical NVLink ranks. + runtime: the C++ runtime. + """ + + def __init__(self, + group: dist.ProcessGroup, + # Provide `num_bytes` + num_bytes: Optional[int] = None, + # Or provide MoE settings (BF16 by default) + num_max_tokens_per_rank: int = 0, + hidden: int = 0, + num_topk: int = 0, + use_fp8_dispatch: bool = False, + # Configs + deterministic: bool = False, + allow_hybrid_mode: bool = True, + allow_multiple_reduction: bool = True, + prefer_overlap_with_compute: bool = True, + sl_idx: int = 3, + num_allocated_qps: int = 0, + num_cpu_timeout_secs: int = 300, num_gpu_timeout_secs: int = 100, + explicitly_destroy: bool = False): + """ + Initialize the elastic communication buffer. + + Arguments: + group: the communication group. + num_bytes: the total buffer size in bytes, if set, overrides MoE-based calculation. + num_max_tokens_per_rank: the maximum number of tokens per rank, used for buffer size calculation. + hidden: the hidden dimension of each token. + num_topk: the number of top-k experts per token. + use_fp8_dispatch: whether to enable FP8 casting, with this, the received data will be a tuple of FP8 tensor and scaling factors. + deterministic: whether to use deterministic routing algorithms. + allow_hybrid_mode: whether to enable hybrid mode. + allow_multiple_reduction: whether to allow multiple reductions in combine. + prefer_overlap_with_compute: whether to prefer overlapping communication with compute. + sl_idx: the RDMA service level index, can be overridden by `EP_OVERRIDE_RDMA_SL` env var. + num_allocated_qps: the number of QPs to allocate for RDMA (0 for automatic). + num_cpu_timeout_secs: CPU-side timeout in seconds for CPU sync. + num_gpu_timeout_secs: GPU-side timeout in seconds for GPU operations. + explicitly_destroy: If this flag is set to True, you need to explicitly call `destroy()` to release resources; + otherwise, the resources will be released by the destructor. + """ + # Some useful utilities + self.group = group + self.rank_idx = group.rank() + self.num_ranks = group.size() + self.allow_hybrid_mode = allow_hybrid_mode + self.allow_multiple_reduction = allow_multiple_reduction + self.prefer_overlap_with_compute = prefer_overlap_with_compute + self.nccl_comm_handle = get_nccl_comm_handle(group) + + # Calculate buffer size + if num_bytes is None: + # NOTES: we allow `num_topk == 0`, as the buffer size can also be calculated by number of ranks (maybe bigger though) + num_bytes = _C.calculate_elastic_buffer_size( + self.nccl_comm_handle.get(), + num_max_tokens_per_rank, hidden, num_topk, use_fp8_dispatch, + allow_hybrid_mode, allow_multiple_reduction) + if os.environ.get('EP_BUFFER_DEBUG', 0): + print(f'Initializing EP elastic buffer with {num_bytes} bytes at rank EP {group.rank()}/{group.size()}') + self.num_bytes = num_bytes + + # Store default values + self.num_max_tokens_per_rank = num_max_tokens_per_rank + + # Check PCIe GPUs + check_nvlink_connections(group) + + # RDMA SL + if 'EP_OVERRIDE_RDMA_SL' in os.environ: + sl_idx = int(os.environ['EP_OVERRIDE_RDMA_SL']) + + # Automatic maximum QP count allowed + # TODO(tianr22): revise the QP count in consideration of Engram + if num_allocated_qps == 0: + # Hybrid mode will consume more QPs + # The extra QP is for notify warps + if self.allow_hybrid_mode: + num_allocated_qps = 65 if check_fast_rdma_atomic_support() else 129 + else: + num_allocated_qps = 17 + self.num_allocated_qps = num_allocated_qps + + # Create CPP handle + self.explicitly_destroy = explicitly_destroy + self.runtime = _C.ElasticBuffer(group.rank(), group.size(), + self.nccl_comm_handle.get(), + num_bytes, + deterministic, + allow_hybrid_mode, + allow_multiple_reduction, + prefer_overlap_with_compute, + sl_idx, num_allocated_qps, + num_cpu_timeout_secs, num_gpu_timeout_secs, + self.explicitly_destroy) + + # Logical rank indices + self.num_scaleout_ranks, self.num_scaleup_ranks = self.get_logical_domain_size() + self.scaleout_rank_idx = self.rank_idx // self.num_scaleup_ranks + self.scaleup_rank_idx = self.rank_idx % self.num_scaleup_ranks + + # Physical rank indices + self.num_rdma_ranks, self.num_nvlink_ranks = self.get_physical_domain_size() + + # Call a barrier to ensure initialization visibility for all peers + torch.cuda.synchronize() + group.barrier() + torch.cuda.synchronize() + + def destroy(self) -> None: + """ + Destroy the C++ runtime and release resources. Requires `explicitly_destroy=True` at construction. + """ + assert self.explicitly_destroy + + if self.runtime is not None: + self.runtime.destroy() + self.runtime = None # Cannot use anymore + self.nccl_comm_handle = None + + @staticmethod + def get_buffer_size_hint(group: dist.ProcessGroup, + num_max_tokens_per_rank: int, hidden: int, + num_topk: int = 0, use_fp8_dispatch: bool = False, + allow_hybrid_mode: bool = True, + allow_multiple_reduction: bool = True) -> int: + """ + Get a recommended buffer size (in bytes) for the given MoE settings, without constructing the buffer. + + Arguments: + group: the communication group. + num_max_tokens_per_rank: the maximum number of tokens per rank. + hidden: the hidden dimension of each token. + num_topk: the number of top-k experts per token. + use_fp8_dispatch: whether to use FP8 for dispatch. + allow_hybrid_mode: whether to enable hybrid mode. + allow_multiple_reduction: whether to allow multiple reductions in combine. + + Returns: + size: the recommended buffer size in bytes. + """ + return _C.calculate_elastic_buffer_size( + get_nccl_comm_handle(group).get(), + num_max_tokens_per_rank, hidden, num_topk, use_fp8_dispatch, + allow_hybrid_mode, allow_multiple_reduction) + + @staticmethod + def get_engram_storage_size_hint(num_entries: int, hidden: int, + num_max_tokens_per_rank: int, + dtype: torch.dtype = torch.bfloat16) -> int: + """ + (Experimental) Get a minimum buffer size requirement for Engram storage. + + Arguments: + num_entries: the number of entries in the Engram storage. + hidden: the hidden dimension of each entry. + num_max_tokens_per_rank: the maximum number of tokens per rank (reserved for receive space). + dtype: the data type, defaults to `torch.bfloat16`. + + Returns: + size: the recommended Engram storage size in bytes. + """ + # TODO: refactor all APIs to allow more parallelism + # TODO: consider FP4 + num_sf_packs = ceil_div(hidden, 32) if dtype.itemsize <= 1 else 0 + # NOTES: we align per-entry size with 32 bytes (LDG.256) + num_bytes_per_entry = align(hidden * dtype.itemsize + num_sf_packs * 4, 32) + return num_bytes_per_entry * (num_entries + num_max_tokens_per_rank) + + @staticmethod + def get_pp_buffer_size_hint(num_max_tensor_bytes: int, + num_max_inflight_tensors: int) -> int: + """ + (Experimental) Get a minimum buffer size requirement for pipeline-parallel (PP) send/recv. + + Arguments: + num_max_tensor_bytes: the maximum tensor size in bytes per send/recv operation. + num_max_inflight_tensors: the maximum number of in-flight tensors at once. + + Returns: + size: the recommended PP buffer size in bytes. + """ + # Align with `LDG.256` + num_max_tensor_bytes = align(num_max_tensor_bytes, 32) + + # Each buffer (send and recv, * 2) contains prev and next rank (* 2) in the ring + return num_max_tensor_bytes * num_max_inflight_tensors * 2 * 2 + + @staticmethod + def get_agrs_buffer_size_hint(group: dist.ProcessGroup, + num_max_session_bytes: int) -> int: + """ + (Experimental) Get a minimum buffer size requirement for all-gather reduce-scatter (AGRS) sessions. + + Arguments: + group: the communication group. + num_max_session_bytes: the maximum total bytes of all gathered tensors in a single session. + + Returns: + size: the recommended AGRS buffer size in bytes. + """ + return num_max_session_bytes + + def barrier(self, use_comm_stream: bool = True, with_cpu_sync: bool = False) -> None: + """ + Perform a GPU-level barrier across all ranks, optionally with CPU synchronization. + + Arguments: + use_comm_stream: whether to use the communication stream (otherwise uses the current compute stream). + with_cpu_sync: whether to also call `cudaDeviceSynchronize` before and after the barrier. + """ + self.runtime.barrier(use_comm_stream, with_cpu_sync) + + @staticmethod + def _unpack_handle(handle: Optional[EPHandle] = None) \ + -> Tuple[Optional[int], Optional[list], + Optional[torch.Tensor], Optional[torch.Tensor], + Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + if handle is None: + return None, None, None, None, None, None, None + return (handle.num_recv_tokens, + handle.num_recv_tokens_per_expert_list, + handle.psum_num_recv_tokens_per_scaleup_rank, + handle.psum_num_recv_tokens_per_expert, + handle.dst_buffer_slot_idx, + handle.token_metadata_at_forward, + handle.channel_linked_list) + + @staticmethod + def capture() -> EventHandle: + """ + Capture a CUDA event on the current stream, i.e. `torch.cuda.current_stream()`. + + Returns: + event_handle: the captured event handle. + """ + return EventHandle() + + def get_comm_stream(self) -> torch.Stream: + """ + Get the communication stream. + + Returns: + stream: the communication stream. + """ + ts: torch.Stream = self.runtime.get_comm_stream() + return torch.cuda.Stream(stream_id=ts.stream_id, device_index=ts.device_index, device_type=ts.device_type) + + def get_physical_domain_size(self) -> Tuple[int, int]: + """ + Get the physical domain sizes (RDMA ranks and NVLink ranks). + + Returns: + num_rdma_ranks: the number of physical RDMA ranks. + num_nvlink_ranks: the number of physical NVLink ranks. + """ + return self.runtime.get_physical_domain_size() + + def get_logical_domain_size(self) -> Tuple[int, int]: + """ + Get the logical domain sizes (scaleout ranks and scaleup ranks). + + Returns: + num_scaleout_ranks: the number of logical scaleout ranks. + num_scaleup_ranks: the number of logical scaleup ranks. + """ + return self.runtime.get_logical_domain_size() + + def engram_write(self, storage: torch.Tensor) -> None: + """ + (Experimental) Write Engram storage data into the buffer. + This call includes a barrier before and after the write to ensure visibility. + + Arguments: + storage: `[num_entries, hidden]` with `torch.bfloat16`, the Engram storage tensor. + """ + # TODO: support FP8 + self.runtime.engram_write(storage) + + def engram_fetch(self, indices: torch.Tensor, num_qps: int = 0) -> Callable: + """ + (Experimental) Fetch Engram entries from remote ranks via RDMA. + Returns a callable that, when invoked, waits for the RDMA gets to complete and returns the fetched tensor. + + Arguments: + indices: `[num_tokens]` with `torch.int`, the entry indices to fetch. + num_qps: the number of QPs to use (0 for all allocated QPs). + + Returns: + hook: a callable that blocks until data arrives and returns the fetched tensor + with shape `[num_tokens, hidden]` and type `torch.bfloat16`. + """ + return self.runtime.engram_fetch(indices, num_qps) + + def pp_set_config(self, num_max_tensor_bytes: int, num_max_inflight_tensors: int): + """ + (Experimental) Configure pipeline-parallel (PP) send/recv parameters. Includes a barrier to flush previous operations. + + Arguments: + num_max_tensor_bytes: the maximum tensor size in bytes per send/recv operation. + num_max_inflight_tensors: the maximum number of in-flight tensors at once. + """ + self.runtime.pp_set_config(num_max_tensor_bytes, num_max_inflight_tensors) + + def pp_send(self, t: torch.Tensor, dst_rank_idx: int, num_sms: int = 0) -> None: + """ + (Experimental) Send a tensor to an adjacent rank in the PP ring (prev or next rank only). + + Arguments: + t: the tensor to send, must be contiguous and fit within `num_max_tensor_bytes`. + dst_rank_idx: the destination rank index (must be prev or next rank in the ring). + num_sms: the number of SMs to use (0 for all SMs). + """ + self.runtime.pp_send(t, dst_rank_idx, num_sms) + + def pp_recv(self, t: torch.Tensor, src_rank_idx: int, num_sms: int = 0) -> None: + """ + (Experimental) Receive a tensor from an adjacent rank in the PP ring (prev or next rank only). + + Arguments: + t: the output tensor to receive into, must be contiguous and fit within `num_max_tensor_bytes`. + src_rank_idx: the source rank index (must be prev or next rank in the ring). + num_sms: the number of SMs to use (0 for all SMs). + """ + self.runtime.pp_recv(t, src_rank_idx, num_sms) + + def create_agrs_session(self) -> None: + """ + (Experimental) Begin a new all-gather reduce-scatter (AGRS) session. Must be paired with `destroy_agrs_session`. + + """ + self.runtime.create_agrs_session() + + def destroy_agrs_session(self) -> None: + """ + (Experimental) End the current AGRS session. Waits for the compute stream, signals session completion to all peers. + + """ + self.runtime.destroy_agrs_session() + + @contextmanager + def agrs_new_session(self, enabled: bool = True): + """ + (Experimental) Context manager that wraps `create_agrs_session` and `destroy_agrs_session`. + + Arguments: + enabled: if `False`, the context manager is a no-op. + """ + if not enabled: + yield + return + + self.runtime.create_agrs_session() + try: + yield + finally: + self.runtime.destroy_agrs_session() + + def agrs_set_config(self, num_max_session_bytes: int, + num_max_all_gathers_per_session: int) -> None: + """ + (Experimental) Configure AGRS session parameters. Includes a barrier to flush previous operations. + + Arguments: + num_max_session_bytes: the maximum total bytes of gathered tensors per session. + num_max_all_gathers_per_session: the maximum number of all-gather operations per session. + """ + self.runtime.agrs_set_config(num_max_session_bytes, num_max_all_gathers_per_session) + + # noinspection PyTypeChecker + def agrs_get_inplace_tensor(self, + shapes: Union[Tuple[int, ...], torch.Size, Sequence[Union[Tuple[int, ...], torch.Size]]], + dtype: torch.dtype) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + """ + (Experimental) Get in-place tensor(s) from the AGRS buffer for this rank's slot, without copying. + Must be called within an active AGRS session. + + Arguments: + shapes: the shape(s) of tensor(s) to allocate. Pass a single shape tuple, or a sequence of shape tuples + for batched mode. + dtype: the data type for the tensor(s). + + Returns: + tensor: a single tensor if a single shape is given, or a tuple of tensors for batched mode. + """ + is_batched_mode = isinstance(shapes[0], tuple) + if not is_batched_mode: + shapes = (shapes, ) + tensors = self.runtime.agrs_get_inplace_tensor( + (math.prod(shape) * dtype.itemsize for shape in shapes) + ) + out = tuple(tensor.view(dtype).view(shape) for tensor, shape in zip(tensors, shapes, strict=True)) + return out if is_batched_mode else out[0] + + def all_gather(self, t: Union[torch.Tensor, Sequence[torch.Tensor]]): + """ + (Experimental) Perform an all-gather operation within an active AGRS session. + Each rank's data is gathered to all ranks via NVLink symmetric memory. + + Arguments: + t: a single tensor or a sequence of tensors to all-gather. Each tensor must be contiguous and + CUDA-allocated, with `nbytes` aligned to 32 bytes. + + Returns: + For a single tensor: `(gathered, handle)` where `gathered` has an extra leading dimension of + `num_ranks`, and `handle` is a callable to wait for data arrival. + For a sequence: `(*gathered_tensors, handle)` with one gathered tensor per input. + """ + if isinstance(t, torch.Tensor): + tensors, handle = self.runtime.all_gather((t,)) + return tensors[0], handle + + # Batched + tensors, handle = self.runtime.all_gather(t) + return *tensors, handle + + @weak_lru(maxsize=None) + def get_theoretical_num_sms(self, num_experts: int, num_topk: int, + num_scaleout_topk: int = 0, + rdma_gbs: float = 0, nvlink_gbs: float = 0, + # TODO: use different values for other architectures + sm_read_gbs: float = 200, sm_write_gbs: float = 50) -> int: + """ + Estimate the optimal number of SMs for dispatch/combine kernels based on bandwidth modeling. + The result is cached. This assumes a balanced gate distribution. + + Arguments: + num_experts: the number of all experts. + num_topk: the number of top-k experts per token. + num_scaleout_topk: reserved for balanced gate (must be 0 currently). + rdma_gbs: the RDMA bandwidth in GB/s (0 for auto-detect). + nvlink_gbs: the NVLink bandwidth in GB/s (0 for auto-detect). + sm_read_gbs: the per-SM HBM read bandwidth in GB/s. + sm_write_gbs: the per-SM HBM write bandwidth in GB/s. + + Returns: + num_sms: the recommended SM count (even, at least 4). + """ + # TODO: support `do_expand` and `allow_multiple_reduction` + + # The `1` in this function means scale-up traffic + # i.e. the HBM read volume of the dispatch copy epilogue, equals to "the number of tokens" * "num_expected_topk" * "data size per token" + # NOTES: this is for balanced gate + # For V3.0's group-limited gate, please do not use this function + # TODO: support this + assert num_scaleout_topk == 0 + + # Get bandwidth + if rdma_gbs == 0 and self.num_rdma_ranks > 1: + rdma_gbs = get_rdma_gbs() + if nvlink_gbs == 0: + nvlink_gbs = get_nvlink_gbs() + + # Initial count + # NOTES: we don't count HBM traffic + sm_read, sm_write = 0, 0 + rdma_traffic, nvlink_traffic = 0, 0 + + def get_expected_topk(num_groups: int) -> float: + assert num_experts % num_groups == 0 + return num_groups * (1 - math.comb(num_experts - num_experts // num_groups, num_topk) / math.comb(num_experts, num_topk)) + + # Expected top-k scale-out ranks + num_expected_scaleout_topk = get_expected_topk(self.num_scaleout_ranks) if self.num_scaleout_ranks > 1 else 0 + + # Expected top-k scale-up ranks + num_expected_topk = get_expected_topk(self.num_ranks) + + # Read tokens + sm_read += 1 / num_expected_topk + + # NOTES: we don't consider the skip-send-buffer cases (all selections fall in the local) + if self.num_scaleout_ranks > 1: + # Scaleup warps: write send buffer + sm_write += 1 / num_expected_topk + + # Scaleout traffic + sm_write += (1 / num_expected_topk) * (num_expected_scaleout_topk / self.num_scaleout_ranks) # Local bypass + rdma_traffic += (1 / num_expected_topk) * (num_expected_scaleout_topk * (1 - 1 / self.num_scaleout_ranks)) + + # Forward warps + sm_read += num_expected_scaleout_topk / num_expected_topk + sm_write += 1 # Issue scaleup + nvlink_traffic += 1 - (1 / self.num_scaleup_ranks) + else: + # Write send buffer + if self.num_rdma_ranks > 1: + sm_write += 1 / num_expected_topk + + # Issue NVLink + sm_write += self.num_nvlink_ranks / self.num_ranks + + # NVLink and RDMA traffic + nvlink_traffic += self.num_nvlink_ranks / self.num_ranks * (1 - 1 / self.num_nvlink_ranks) # Except local bypass + rdma_traffic += (self.num_ranks - self.num_nvlink_ranks) / self.num_ranks + + # Found the bounded one + if self.num_scaleout_ranks > 1 and (rdma_traffic / rdma_gbs) > (nvlink_traffic / nvlink_gbs): + bounded_traffic, bounded_gbs = rdma_traffic, rdma_gbs + else: + bounded_traffic, bounded_gbs = nvlink_traffic, nvlink_gbs + + # Calculate SM count + # NOTES: will try to use more SMs if not overlap with compute + num_device_sms = torch.cuda.get_device_properties('cuda').multi_processor_count + num_sms = num_device_sms # No traffic, e.g., EP=1 + if bounded_traffic > 0: + num_sms = max( + bounded_gbs / bounded_traffic * sm_read / sm_read_gbs, + bounded_gbs / bounded_traffic * sm_write / sm_write_gbs, + ) + num_sms = align(max(4, math.ceil(num_sms * 1.25)), 2) + num_sms = num_sms if self.prefer_overlap_with_compute else max(num_sms, 64) + num_sms = min(num_sms, num_device_sms) + + # Summary + if os.environ.get('EP_BUFFER_DEBUG', 0): + print(f'EP SM approximation: ' + f'{sm_read=}, {sm_write=}, {rdma_traffic=}, {nvlink_traffic=}, ' + f'{rdma_gbs=}, {nvlink_gbs=}, ' + f'{num_expected_scaleout_topk=}, {num_expected_topk=}, ' + f'{bounded_traffic=}, {bounded_gbs=}, {num_sms=}') + return num_sms + + def get_theoretical_num_qps(self, num_sms: int) -> int: + """ + Estimate the optimal number of RDMA QPs based on SM count and mode. + + Arguments: + num_sms: the number of SMs used for the dispatch/combine kernel. + + Returns: + num_qps: the recommended QP count, capped by `num_allocated_qps`. + """ + # For direct mode, we encourage less QPs to reduce DB ringing overhead + num_qps = min(num_sms, 8 + 1) + + # For hybrid mode, we encourage every channel (and notify) to have an independent QP + if self.allow_hybrid_mode: + num_qps = num_sms * 16 + 1 + + return min(num_qps, self.num_allocated_qps) + + def dispatch(self, + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + topk_idx: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, + num_max_tokens_per_rank: Optional[int] = None, + expert_alignment: Optional[int] = None, + num_sms: int = 0, num_qps: int = 0, + previous_event: Optional[EventHandle] = None, + previous_event_before_epilogue: Optional[EventHandle] = None, + async_with_compute_stream: bool = False, + allocate_on_comm_stream: bool = False, + handle: Optional[EPHandle] = None, + do_handle_copy: bool = True, + do_cpu_sync: Optional[bool] = None, + do_expand: bool = False, + use_tma_aligned_col_major_sf: bool = False) \ + -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + Optional[torch.Tensor], Optional[torch.Tensor], + EPHandle, EventOverlap]: + """ + Dispatch tokens to different ranks. Supports both single-node and multi-node settings. + SM and QP counts are automatically determined if not specified. + + Arguments: + x: `torch.Tensor` or tuple of `torch.Tensor`, for the first type, the shape must be + `[num_tokens, hidden]`, and type must be `torch.bfloat16`; for the second type (FP8 mode), + the first element of the tuple must be `[num_tokens, hidden]` with type `torch.float8_e4m3fn`, + the second is the scale factors. + topk_idx: `[num_tokens, num_topk]` with `deep_ep.topk_idx_t` (typically `torch.int64`), the expert + indices selected by each token, `-1` means no selections. + Must be `None` if `handle` is provided. + topk_weights: `[num_tokens, num_topk]` with `torch.float`, the expert weights of each token to dispatch. + Must be `None` if `handle` is provided. + cumulative_local_expert_recv_stats: `[num_local_experts]` with `torch.int`, a cumulative expert count + tensor for statistics, useful for online EP load balance monitoring. + num_experts: the number of all experts. Inferred from `handle` if provided. + num_max_tokens_per_rank: the maximum number of tokens per rank. Inferred from constructor default + or `handle` if provided. + expert_alignment: align the number of tokens received by each local expert to this variable. + num_sms: the number of SMs to use (0 for automatic via `get_theoretical_num_sms`). + num_qps: the number of RDMA QPs to use (0 for automatic via `get_theoretical_num_qps`). + previous_event: the event to wait before actually executing the kernel. + If set, `allocate_on_comm_stream` must also be `True`. + previous_event_before_epilogue: the event to wait before actually executing the copy epilogue. + async_with_compute_stream: the current stream will not wait for the communication kernels to be + finished if set. + allocate_on_comm_stream: control whether all the allocated tensors' ownership to be on the + communication stream. + handle: an optional cached `EPHandle` from a previous dispatch, if set, the CPU will reuse the layout + information to save some time. `topk_idx` and `topk_weights` must be `None`. + do_handle_copy: whether to clone `topk_idx` in the returned handle (to prevent user modification). + do_cpu_sync: whether to synchronize with CPU to get exact received token counts. + `None` defaults to `True` unless `handle` is provided. + do_expand: whether to use the expanding layout (one slot per expert per token). + use_tma_aligned_col_major_sf: whether to use TMA-aligned column-major layout for scale factors. + + Returns: + recv_x: received tokens, the same type and tuple as the input `x` + recv_topk_idx: received expert indices + recv_topk_weights: received expert weights (`None` if `topk_weights` was not provided). + handle: the returned communication handle. + event: the event after executing the kernel (valid only if `async_with_compute_stream` is set). + """ + check_torch_deterministic() + + # Automatic decide SM and QP count + num_topk = (handle.topk_idx if topk_idx is None else topk_idx).shape[1] + num_sms = self.get_theoretical_num_sms(num_experts, num_topk) if num_sms == 0 else num_sms + num_qps = self.get_theoretical_num_qps(num_sms) if num_qps == 0 else num_qps + assert num_qps <= self.num_allocated_qps, f'Allocated QPs are not enough' + + # Unpack SF + x, sf = x if isinstance(x, tuple) else (x, None) + + # Unpack handles + # Reuse some values if possible + if handle is not None: + assert topk_idx is None and topk_weights is None + assert do_cpu_sync is None or not do_cpu_sync, 'Cannot do CPU sync with cached handle' + topk_idx = handle.topk_idx + num_max_tokens_per_rank = value_or(num_max_tokens_per_rank, handle.num_max_tokens_per_rank) + num_experts = value_or(num_experts, handle.num_experts) + expert_alignment = value_or(expert_alignment, handle.expert_alignment) + do_cpu_sync = False + + # Should be aligned with the handle context + assert (num_experts, expert_alignment, num_max_tokens_per_rank) == \ + (handle.num_experts, handle.expert_alignment, handle.num_max_tokens_per_rank) + (cached_num_recv_tokens, cached_num_recv_tokens_per_expert_list, + cached_psum_num_recv_tokens_per_scaleup_rank, cached_psum_num_recv_tokens_per_expert, + cached_dst_buffer_slot_idx, + cached_token_metadata_at_forward, + cached_channel_linked_list) = self._unpack_handle(handle) + + # Some default values + num_max_tokens_per_rank = value_or(num_max_tokens_per_rank, self.num_max_tokens_per_rank) + expert_alignment = value_or(expert_alignment, 1) + do_cpu_sync = value_or(do_cpu_sync, True) + + # Do dispatch + (recv_x, recv_sf, + recv_topk_idx, recv_topk_weights, + cloned_topk_idx, + num_recv_tokens_per_expert_list, + psum_num_recv_tokens_per_scaleup_rank, + psum_num_recv_tokens_per_expert, + recv_src_metadata, + dst_buffer_slot_idx, + token_metadata_at_forward, + channel_linked_list, + event) = self.runtime.dispatch(x, sf, topk_idx, topk_weights, + cumulative_local_expert_recv_stats, + cached_num_recv_tokens, + cached_num_recv_tokens_per_expert_list, + cached_psum_num_recv_tokens_per_scaleup_rank, + cached_psum_num_recv_tokens_per_expert, + cached_dst_buffer_slot_idx, + cached_token_metadata_at_forward, + cached_channel_linked_list, + num_max_tokens_per_rank, + num_experts, expert_alignment, + num_sms, num_qps, + previous_event, + previous_event_before_epilogue, + async_with_compute_stream, allocate_on_comm_stream, + do_handle_copy, do_cpu_sync, do_expand, + use_tma_aligned_col_major_sf) + if handle is None: + handle = EPHandle(do_expand, + num_experts, expert_alignment, + num_max_tokens_per_rank, + num_sms, + cloned_topk_idx if do_handle_copy else topk_idx, + num_recv_tokens_per_expert_list, + psum_num_recv_tokens_per_scaleup_rank, + psum_num_recv_tokens_per_expert, + recv_src_metadata, + dst_buffer_slot_idx, + token_metadata_at_forward, + channel_linked_list) + + # Repack SF + recv_x = (recv_x, recv_sf) if recv_sf is not None else recv_x + + # Return + return recv_x, recv_topk_idx, recv_topk_weights, handle, EventOverlap(event) + + @staticmethod + def _unpack_bias(bias: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]) \ + -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + bias_0, bias_1 = None, None + if isinstance(bias, torch.Tensor): + bias_0 = bias + elif isinstance(bias, tuple): + assert len(bias) == 2 + bias_0, bias_1 = bias + return bias_0, bias_1 + + def combine(self, + x: torch.Tensor, + handle: EPHandle, + topk_weights: Optional[torch.Tensor] = None, + bias: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] = None, + num_sms: int = 0, num_qps: int = 0, + previous_event: EventHandle = None, + previous_event_before_epilogue: Optional[EventHandle] = None, + async_with_compute_stream: bool = False, + allocate_on_comm_stream: bool = False) \ + -> Tuple[torch.Tensor, Optional[torch.Tensor], EventOverlap]: + """ + Combine (reduce) tokens from different ranks back to their original ranks. + Supports both single-node and multi-node settings. + + Arguments: + x: `[num_tokens, hidden]` with `torch.bfloat16`, the tokens to send for reducing to its original ranks. + handle: a must-set communication handle, you can obtain this from the `dispatch` function. + topk_weights: `[num_tokens, num_topk]` with `torch.float`, the tokens' top-k weights for reducing to + its original ranks. Not used in expand mode. + bias: 0, 1 or 2 `[num_combined_tokens, hidden]` with `torch.bfloat16` final bias to the output. + num_sms: the number of SMs to use (0 to reuse the SM count from the dispatch handle). + num_qps: the number of RDMA QPs to use (0 for automatic via `get_theoretical_num_qps`). + previous_event: the event to wait before actually executing the kernel. + If set, `allocate_on_comm_stream` must also be `True`. + previous_event_before_epilogue: the event to wait before actually executing the reduce epilogue. + async_with_compute_stream: the current stream will not wait for the communication kernels to be + finished if set. + allocate_on_comm_stream: control whether all the allocated tensors' ownership to be on the + communication stream. + + Returns: + combined_x: the reduced token tensor, with shape `[num_combined_tokens, hidden]` and type `torch.bfloat16`. + combined_topk_weights: the reduced top-k weights, with shape `[num_combined_tokens, num_topk]` and type `torch.float`. + event: the event after executing the kernel (valid only if `async_with_compute_stream` is set). + """ + check_torch_deterministic() + + # Automatic decide SM and QP count + num_sms = handle.num_sms if num_sms == 0 else num_sms + num_qps = self.get_theoretical_num_qps(num_sms) if num_qps == 0 else num_qps + assert num_qps <= self.num_allocated_qps, f'Allocated QPs are not enough' + + bias_0, bias_1 = ElasticBuffer._unpack_bias(bias) + combined_x, combined_topk_weights, event = \ + self.runtime.combine(x, topk_weights, + bias_0, bias_1, + handle.recv_src_metadata, + handle.topk_idx, + handle.psum_num_recv_tokens_per_scaleup_rank, + handle.token_metadata_at_forward, + handle.channel_linked_list, + handle.num_experts, + handle.num_max_tokens_per_rank, + num_sms, num_qps, + previous_event, + previous_event_before_epilogue, + async_with_compute_stream, + allocate_on_comm_stream, + handle.do_expand) + return combined_x, combined_topk_weights, EventOverlap(event) diff --git a/deep_ep/buffer.py b/deep_ep/buffers/legacy.py similarity index 96% rename from deep_ep/buffer.py rename to deep_ep/buffers/legacy.py index 37512ee91..af8821598 100644 --- a/deep_ep/buffer.py +++ b/deep_ep/buffers/legacy.py @@ -4,10 +4,11 @@ from typing import Callable, List, Tuple, Optional, Union # noinspection PyUnresolvedReferences -import deep_ep_cpp +import deep_ep._C as _C # noinspection PyUnresolvedReferences -from deep_ep_cpp import Config, EventHandle -from .utils import EventOverlap, check_nvlink_connections +from deep_ep._C import Config, EventHandle +from ..utils.event import EventOverlap +from ..utils.envs import check_nvlink_connections, check_torch_deterministic class Buffer: @@ -37,7 +38,6 @@ def __init__(self, num_qps_per_rank: int = 24, allow_nvlink_for_low_latency_mode: bool = True, allow_mnnvl: bool = False, - use_fabric: bool = False, explicitly_destroy: bool = False, enable_shrink: bool = False, comm: Optional["mpi4py.MPI.Comm"] = None) -> None: # noqa: F821 @@ -89,8 +89,8 @@ def all_gather_object(obj): self.low_latency_mode = low_latency_mode self.explicitly_destroy = explicitly_destroy self.enable_shrink = enable_shrink - self.runtime = deep_ep_cpp.Buffer(self.rank, self.group_size, num_nvl_bytes, num_rdma_bytes, low_latency_mode, explicitly_destroy, - enable_shrink, use_fabric) + self.runtime = _C.Buffer(self.rank, self.group_size, num_nvl_bytes, num_rdma_bytes, low_latency_mode, + explicitly_destroy, enable_shrink, allow_mnnvl) # Synchronize device IDs local_device_id = self.runtime.get_local_device_id() @@ -116,7 +116,7 @@ def all_gather_object(obj): # Reduce gpu memory usage # 6 default teams + 1 extra team os.environ['NVSHMEM_MAX_TEAMS'] = '7' - # Disable NVLink SHArP + # Disable NVLink SHARP os.environ['NVSHMEM_DISABLE_NVLS'] = '1' # NOTES: NVSHMEM initialization requires at least 256 MiB os.environ['NVSHMEM_CUMEM_GRANULARITY'] = f'{2 ** 29}' @@ -148,7 +148,7 @@ def destroy(self): @staticmethod def is_sm90_compiled(): - return deep_ep_cpp.is_sm90_compiled() + return _C.is_sm90_compiled() @staticmethod def set_num_sms(new_num_sms: int) -> None: @@ -186,7 +186,7 @@ def get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank: int, hidden Returns: size: the RDMA buffer size recommended. """ - return deep_ep_cpp.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts) + return _C.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts) def get_comm_stream(self) -> torch.Stream: """ @@ -329,7 +329,7 @@ def dispatch(self, x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], previous_event: Optional[EventOverlap] = None, async_finish: bool = False, allocate_on_comm_stream: bool = False) -> \ Tuple[Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor], Optional[torch.Tensor], - Optional[torch.Tensor], List[int], Tuple, EventOverlap]: + Optional[torch.Tensor], List[int], Tuple, EventOverlap]: """ Dispatch tokens to different ranks, both intranode and internode settings are supported. Intranode kernels require all the ranks should be visible via NVLink. @@ -369,14 +369,17 @@ def dispatch(self, x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], handle: the returned communication handle. event: the event after executing the kernel (valid only if `async_finish` is set). """ + check_torch_deterministic() + # Default config config = self.get_dispatch_config(self.group_size) if config is None else config # Internode if self.runtime.get_num_rdma_ranks() > 1: + assert num_worst_tokens == 0, 'Internode dispatch does not support `num_worst_tokens > 0`' return self.internode_dispatch(x, handle, num_tokens_per_rank, num_tokens_per_rdma_rank, is_token_in_rank, - num_tokens_per_expert, topk_idx, topk_weights, expert_alignment, num_worst_tokens, config, - previous_event, async_finish, allocate_on_comm_stream) + num_tokens_per_expert, topk_idx, topk_weights, expert_alignment, config, previous_event, + async_finish, allocate_on_comm_stream) # Launch the kernel with cached or non-cached mode x, x_scales = x if isinstance(x, tuple) else (x, None) @@ -431,6 +434,8 @@ def combine(self, x: torch.Tensor, handle: Tuple, recv_topk_weights: the reduced top-k weights from its dispatch ranks. event: the event after executing the kernel (valid only if `async_finish` is set). """ + check_torch_deterministic() + # Default config config = self.get_combine_config(self.group_size) if config is None else config @@ -455,7 +460,7 @@ def internode_dispatch(self, x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Te num_tokens_per_rank: Optional[torch.Tensor] = None, num_tokens_per_rdma_rank: Optional[torch.Tensor] = None, is_token_in_rank: Optional[torch.Tensor] = None, num_tokens_per_expert: Optional[torch.Tensor] = None, topk_idx: Optional[torch.Tensor] = None, topk_weights: Optional[torch.Tensor] = None, expert_alignment: int = 1, - num_worst_tokens: int = 0, config: Optional[Config] = None, + config: Optional[Config] = None, previous_event: Optional[EventOverlap] = None, async_finish: bool = False, allocate_on_comm_stream: bool = False) -> \ Tuple[Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor], Optional[torch.Tensor], @@ -479,7 +484,7 @@ def internode_dispatch(self, x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Te recv_x, recv_x_scales, _, _, _, _, _, _, _, _, _, _, _, _, event = self.runtime.internode_dispatch( x, x_scales, topk_idx, topk_weights, None, None, is_token_in_rank, None, num_recv_tokens, num_rdma_recv_tokens, rdma_channel_prefix_matrix, recv_rdma_rank_prefix_sum, gbl_channel_prefix_matrix, recv_gbl_rank_prefix_sum, - expert_alignment, num_worst_tokens, config, getattr(previous_event, 'event', None), async_finish, allocate_on_comm_stream) + expert_alignment, config, getattr(previous_event, 'event', None), async_finish, allocate_on_comm_stream) return (recv_x, recv_x_scales) if x_scales is not None else recv_x, None, None, None, None, EventOverlap(event) else: assert num_tokens_per_rank is not None and is_token_in_rank is not None and num_tokens_per_expert is not None @@ -491,7 +496,7 @@ def internode_dispatch(self, x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Te x, x_scales, topk_idx, topk_weights, num_tokens_per_rank, num_tokens_per_rdma_rank, is_token_in_rank, num_tokens_per_expert, 0, 0, None, None, None, None, - expert_alignment, num_worst_tokens, config, getattr(previous_event, 'event', None), async_finish, allocate_on_comm_stream) + expert_alignment, config, getattr(previous_event, 'event', None), async_finish, allocate_on_comm_stream) handle = (is_token_in_rank, rdma_channel_prefix_matrix, gbl_channel_prefix_matrix, recv_rdma_channel_prefix_matrix, recv_rdma_rank_prefix_sum, recv_gbl_channel_prefix_matrix, recv_gbl_rank_prefix_sum, recv_src_meta, send_rdma_head, send_nvl_head) @@ -599,6 +604,8 @@ def low_latency_dispatch(self, x: torch.Tensor, topk_idx: torch.Tensor, event: the event after executing the kernel (valid only if `async_finish` is set). hook: the receiving hook function (valid only if `return_recv_hook` is set). """ + check_torch_deterministic() + assert self.nvshmem_qp_depth >= (num_max_dispatch_tokens_per_rank + 1) * 2 packed_recv_x, packed_recv_x_scales, packed_recv_count, packed_recv_src_info, packed_recv_layout_range, event, hook = \ self.runtime.low_latency_dispatch(x, topk_idx, @@ -645,13 +652,15 @@ def low_latency_combine(self, x: torch.Tensor, topk_idx: torch.Tensor, topk_weig out: the in-place output tensor, if set, the kernel will write the result to this tensor and return it directly. combine_wait_recv_cost_stats: a cumulative time spent waiting to receive each token tensor for statistics, which should have shape `[num_ranks, num_ranks]` and be typed as `torch.int64`. - This is useful for detecting and pre-cisely localizing slow anomalies. + This is useful for detecting and precisely localizing slow anomalies. Returns: combined_x: the reduced token tensor, with shape `[num_combined_tokens, hidden]` and type `torch.bfloat16`. event: the event after executing the kernel (valid only if `async_finish` is set). hook: the receiving hook function (valid only if `return_recv_hook` is set). """ + check_torch_deterministic() + src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle assert self.nvshmem_qp_depth >= (num_max_dispatch_tokens_per_rank + 1) * 2 combined_x, event, hook = self.runtime.low_latency_combine(x, topk_idx, topk_weights, src_info, layout_range, @@ -665,8 +674,8 @@ def low_latency_update_mask_buffer(self, rank_to_mask: int, mask: bool = False): Mask (unmask) a rank during communication (dispatch, combine, and clean) Arguments: - rank: the rank to mask (unmask). - mask: if True, will mask the rank (do not recvfrom/sendto the rank), otherwise will unmask the rank. + rank_to_mask: the rank to mask (unmask). + mask: if True, will mask the rank (do not recv from/send to the rank), otherwise will unmask the rank. """ self.runtime.low_latency_update_mask_buffer(rank_to_mask, mask) diff --git a/deep_ep/include/deep_ep/common/comm.cuh b/deep_ep/include/deep_ep/common/comm.cuh new file mode 100644 index 000000000..04a77d418 --- /dev/null +++ b/deep_ep/include/deep_ep/common/comm.cuh @@ -0,0 +1,266 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include + +namespace deep_ep::elastic::comm { + +static constexpr int64_t kNumOneSecCycles = 2000000000; // An approximation of the GPU clock at 2000 MHz + +// Some reserved tags +static constexpr int kDeviceBarrierTag = 0; +static constexpr int kKernelBarrierTag = 1; +static constexpr int kDispatchTag0 = 2; +static constexpr int kDispatchTag1 = 3; +static constexpr int kCombineTag0 = 4; +static constexpr int kCombineTag1 = 5; +static constexpr int kHybridDispatchTag0 = 6; +static constexpr int kHybridDispatchTag1 = 7; +static constexpr int kHybridCombineTag0 = 8; +static constexpr int kHybridCombineTag1 = 9; + +// Some reserved count +static constexpr int kFlushAllAllocatedQPs = -1; + +template +__device__ __forceinline__ void timeout_while(const bool& condition, const func_t& func, + int64_t start_clock = 0) { + // User may share a start clock for multiple waits + if (start_clock == 0) + start_clock = clock64(); + + while (condition) { + const bool timeout = clock64() - start_clock >= kNumTimeoutCycles; + if (func(timeout)) + break; + + if (timeout) { + // Wait another 1 second to let all threads print information and trap + start_clock = clock64(); + while (clock64() - start_clock < kNumOneSecCycles) {} + ptx::trap(); + } + } +} + +template +__device__ __forceinline__ void timeout_while(const func_t& func, const int64_t& start_clock = 0) { + timeout_while(true, func, start_clock); +} + +template +__device__ __forceinline__ std::pair get_qp_mode( + const int& sm_idx, const int& channel_in_sm_idx, const bool& is_notify_warp = false) { + constexpr auto kSharingCTA = NCCL_GIN_RESOURCE_SHARING_CTA; + constexpr auto kSharingGrid = kNumSMs == 1 ? NCCL_GIN_RESOURCE_SHARING_CTA : NCCL_GIN_RESOURCE_SHARING_GPU; + + // Only one QP + if constexpr (kNumQPs == 1) + return {0, kSharingGrid}; + + // The notify warp always use 1 SM and 1 QP + if (is_notify_warp) + return {0, kSharingCTA}; + + // Data channels + constexpr int kQPStartIdx = static_cast(kWithNotifyWarps); + constexpr int kNumAvailableQPs = kNumQPs - kQPStartIdx; + if constexpr (kNumSMs <= kNumAvailableQPs) { + // A single SM uses an entire QP + // e.g., 3 SMs and 10 QPs + // SM 0: 0 3 6 9 + // SM 1: 1 4 7 + // SM 2: 2 5 8 + const int num_qps_in_sm = (kNumAvailableQPs / kNumSMs) + (sm_idx < (kNumAvailableQPs % kNumSMs)); + return {kQPStartIdx + sm_idx + (channel_in_sm_idx % num_qps_in_sm) * kNumSMs, kSharingCTA}; + } else { + // All SMs share all QPs + const auto global_channel_idx = sm_idx * kNumChannelsPerSM + channel_in_sm_idx; + return {kQPStartIdx + (global_channel_idx % kNumAvailableQPs), kSharingGrid}; + } +} + +template +__forceinline__ __device__ void nvlink_barrier_wo_local_sync( + const handle::NCCLGin& gin, + const layout::WorkspaceLayout& workspace, + const int& rank_idx, const int& sm_idx, const int& thread_idx) { + // This barrier only uses 1 SM + if (kNumSMs > 1 and sm_idx > 0) + return; + + // Read the current barrier phase first + const int status = static_cast((*workspace.get_nvl_barrier_counter_ptr()) & 3); + const int phase = status & 1, sign = status >> 1; + + EP_STATIC_ASSERT(kNumRanks <= kNumThreads, "Insufficient threads"); + if (thread_idx < kNumRanks) { + const auto dst_ptr = + gin.get_sym_ptr(workspace.get_nvl_barrier_signal_ptr(phase), thread_idx); + ptx::red_add_rel_sys(dst_ptr, sign ? -1 : 1); + } + __syncthreads(); + + // NOTES: we need `2^64 / 1e6 / 3600 / 24 / 365 = 571000` years to make the counter overflow (1 barrier per us) + // Add the phase counter + if (thread_idx == 0) + atomicAdd(workspace.get_nvl_barrier_counter_ptr(), 1); + + // Check timeout + const auto target = sign ? 0 : kNumRanks; + timeout_while(thread_idx == 0, [=](const bool& is_last_check) { + const auto signal = ptx::ld_acquire_sys(workspace.get_nvl_barrier_signal_ptr(phase)); + if (signal == target) + return true; + + if (is_last_check) { + printf("DeepEP NVLink barrier timeout, tag: %d, nvl: %d, thread: %d, " + "status: %d, signal: %d, phase: %d, target: %d, counter: %llu\n", + kTag, rank_idx, thread_idx, status, signal, phase, target, + *workspace.get_nvl_barrier_counter_ptr()); + } + return false; + }); +} + +template +__forceinline__ __device__ void gin_barrier_wo_local_sync( + const ncclDevComm_t& nccl_dev_comm, + const int& scaleout_rank_idx, const int& scaleup_rank_idx, + const int& sm_idx, const int& thread_idx) { + const auto global_warp_idx = sm_idx * kNumWarps + (thread_idx / 32); + const int& rank_idx = (std::is_same_v) ? scaleup_rank_idx : scaleout_rank_idx; + const int num_qps = kNumQPs == kFlushAllAllocatedQPs ? nccl_dev_comm.ginContextCount : kNumQPs; + + // Flush all QPs by all SMs (only needed for release semantics) + if constexpr (kFlushStores) { + for (int i = global_warp_idx; i < num_qps; i += kNumSMs * kNumWarps) { + ncclGin(nccl_dev_comm, i, NCCL_GIN_RESOURCE_SHARING_CTA).flush(ncclCoopWarp()); + } + // NOTES: we can not use `kNumSMs` to judge, as maybe only part of the SMs will call this function + (gridDim.x > 1) ? cooperative_groups::this_grid().sync() : __syncthreads(); + } + + if (sm_idx == 0) { + // Use QP 0 to do barrier + const auto team = (std::is_same_v) ? + ncclTeamWorld(nccl_dev_comm) : ncclTeamRail(nccl_dev_comm); + const ncclGin gin(nccl_dev_comm, 0, NCCL_GIN_RESOURCE_SHARING_CTA); + for (int i = thread_idx; i < kNumRanks; i += kNumThreads) + gin.signal(team, i, ncclGin_SignalInc{static_cast(rank_idx)}); + + // TODO(NCCL): Using the official NCCL wait signal API, after they added timeout check. + for (int i = thread_idx; i < kNumRanks; i += kNumThreads) { + const auto signal_idx = static_cast(i); + const auto shadow_ptr = gin.getSignalShadowPtr(signal_idx); + const auto target = ++(*shadow_ptr); + + const auto gdaki = static_cast(gin._ginHandle) + gin.contextId; + const auto signal_ptr = reinterpret_cast(__ldg(reinterpret_cast(&gdaki->signals_table.buffer))) + signal_idx; + timeout_while([=](const bool& is_last_check) { + const auto signal = ptx::ld_acquire_sys(signal_ptr); + if (signal >= target) + return true; + + if (is_last_check) { + printf("DeepEP Gin barrier timeout, tag: %d, scaleout: %d, scaleup: %d, thread: %d, " + "signal: %lu, target: %lu\n", kTag, scaleout_rank_idx, scaleup_rank_idx, thread_idx, signal, target); + } + return false; + }); + } + } +} + +template +__forceinline__ __device__ void scaleup_barrier_wo_local_sync( + const handle::NCCLGin& gin, + const layout::WorkspaceLayout& workspace, + const int& rank_idx, const int& sm_idx, const int& thread_idx) { + if constexpr (kIsScaleupNVLink) { + nvlink_barrier_wo_local_sync( + gin, workspace, rank_idx, sm_idx, thread_idx); + } else { + gin_barrier_wo_local_sync( + gin.nccl_dev_comm, 1, rank_idx, sm_idx, thread_idx); + } +} + +template +__forceinline__ __device__ void scaleout_barrier_wo_local_sync( + const handle::NCCLGin& gin, + const int& scaleout_rank_idx, const int& scaleup_rank_idx, + const int& sm_idx, const int& thread_idx) { + gin_barrier_wo_local_sync( + gin.nccl_dev_comm, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx); +} + +template +__forceinline__ __device__ void gpu_barrier(const handle::NCCLGin& gin, + const layout::WorkspaceLayout& workspace, + const int& scaleout_rank_idx, const int& scaleup_rank_idx, + const int& sm_idx, const int& thread_idx, + bool do_scaleout = true, bool do_scaleup = true) { + // A general TMA store wait to prevent proxy memory issues + if constexpr (kFlushStores) { + ptx::tma_store_commit(); + ptx::tma_store_wait(); + __syncwarp(); + } + + // All the SMs should wait + if constexpr (kSyncAtStart) { + cooperative_groups::this_grid().sync(); + } else { + EP_STATIC_ASSERT(not kFlushStores, "No data to be flushed"); + } + + do_scaleout &= kNumScaleoutRanks > 1; + do_scaleup &= kNumScaleupRanks > 1; + if (do_scaleup and do_scaleout) { + // Do scaleup and scaleout barrier in parallel + EP_DEVICE_ASSERT(kNumSMs >= 2 and "At least 2 SMs for a hybrid barrier"); + if (sm_idx == 0) { + // First SM do the scaleup barrier + scaleup_barrier_wo_local_sync( + gin, workspace, scaleup_rank_idx, sm_idx, thread_idx); + + // We need an extra grid sync, as the scaleout barrier will do a sync after flush, before the barrier + // NOTES: this is kind of hacky + if constexpr (kFlushStores) + cooperative_groups::this_grid().sync(); + } else { + // The remaining SMs do the scaleout barrier + scaleout_barrier_wo_local_sync( + gin, scaleout_rank_idx, scaleup_rank_idx, sm_idx - 1, thread_idx); + } + } else if (do_scaleup) { + // Scaleup only + scaleup_barrier_wo_local_sync( + gin, workspace, scaleup_rank_idx, sm_idx, thread_idx); + } else if (do_scaleout) { + // Scaleout only + scaleout_barrier_wo_local_sync( + gin, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx); + } + + // All the SMs should wait + if constexpr (kSyncAtEnd) + cooperative_groups::this_grid().sync(); +} + +} // namespace deep_ep::elastic::comm diff --git a/deep_ep/include/deep_ep/common/compiled.cuh b/deep_ep/include/deep_ep/common/compiled.cuh new file mode 100644 index 000000000..8ea94ccf6 --- /dev/null +++ b/deep_ep/include/deep_ep/common/compiled.cuh @@ -0,0 +1,85 @@ +#pragma once + +// Make CLion CUDA indexing work +#ifdef __CLION_IDE__ +#define __CUDA_ARCH__ 900 +#define __CUDACC_RDC__ +#define __CUDACC__ +#endif + +// Remove Torch restrictions +#ifdef __CUDA_NO_HALF_CONVERSIONS__ +#undef __CUDA_NO_HALF_CONVERSIONS__ +#endif +#ifdef __CUDA_NO_HALF_OPERATORS__ +#undef __CUDA_NO_HALF_OPERATORS__ +#endif +#ifdef __CUDA_NO_HALF2_OPERATORS__ +#undef __CUDA_NO_HALF2_OPERATORS__ +#endif +#ifdef __CUDA_NO_BFLOAT16_CONVERSIONS__ +#undef __CUDA_NO_BFLOAT16_CONVERSIONS__ +#endif +#ifdef __CUDA_NO_BFLOAT162_OPERATORS__ +#undef __CUDA_NO_BFLOAT162_OPERATORS__ +#endif + +#include +#include +#include + +#ifndef DISABLE_SM90_FEATURES +#include +#else +// Ampere does not support FP8 features +#define __NV_E4M3 0 +#define __NV_E5M2 1 +typedef int __nv_fp8_interpretation_t; +typedef int __nv_fp8x4_e4m3; +typedef uint8_t __nv_fp8_storage_t; +#endif + +// Compatibility: 256 bits LD/ST instructions +#if defined(CUDART_VERSION) and CUDART_VERSION >= 13000 +using longlong4_t = longlong4_32a; +#define make_longlong4_t make_longlong4_32a +#else +struct alignas(32) longlong4_t { long long x, y, z, w; }; +__device__ __forceinline__ longlong4_t make_longlong4_t( + const long long& x, const long long& y, const long long& z, const long long& w) { + return {x, y, z, w}; +} +#endif + +#ifndef EP_NUM_TOPK_IDX_BITS +#define EP_NUM_TOPK_IDX_BITS 64 +#endif + +namespace deep_ep { + +#ifndef DISABLE_SM90_FEATURES +constexpr bool kEnableSM90Features = true; +#else +constexpr bool kEnableSM90Features = false; +#endif + +template struct int_with_bits; +template <> struct int_with_bits<8> { using type = int8_t; }; +template <> struct int_with_bits<16> { using type = int16_t; }; +template <> struct int_with_bits<32> { using type = int32_t; }; +template <> struct int_with_bits<64> { using type = int64_t; }; + +using topk_idx_t = int_with_bits::type; + +union sf_pack_t { + float fp32; + int ue8m0x4; +}; + +constexpr int kNumTMAAlignedBytes = 16; +constexpr int kNumAlignedSFPacks = 16 / sizeof(sf_pack_t); + +// Some communication channel settings +constexpr int kNumMaxChannels = 1024; + +} // namespace deep_ep diff --git a/deep_ep/include/deep_ep/common/exception.cuh b/deep_ep/include/deep_ep/common/exception.cuh new file mode 100644 index 000000000..405e573bd --- /dev/null +++ b/deep_ep/include/deep_ep/common/exception.cuh @@ -0,0 +1,94 @@ +#pragma once + +#include +#include +#include + +#ifndef EP_STATIC_ASSERT +#define EP_STATIC_ASSERT(cond, reason) static_assert(cond, reason) +#endif + +class EPException : public std::exception { +private: + std::string message = {}; + +public: + explicit EPException(const char* name, const char* file, const int line, const std::string& error) { + std::stringstream ss; + ss << name << " exception (" << file << ":" << line << "): " << error; + message = ss.str(); + } + + const char* what() const noexcept override { return message.c_str(); } +}; + +#define EPExceptionWithLineInfo(name, message) EPException(name, __FILE__, __LINE__, message) + +#ifndef EP_HOST_ASSERT +#define EP_HOST_ASSERT(cond) \ + do { \ + if (not(cond)) { \ + throw EPException("Assertion", __FILE__, __LINE__, #cond); \ + } \ + } while (0) +#endif + +#ifndef EP_HOST_UNREACHABLE +#define EP_HOST_UNREACHABLE(reason) (throw EPException("Assertion", __FILE__, __LINE__, reason)) +#endif + +#ifndef EP_DEVICE_ASSERT +#define EP_DEVICE_ASSERT(cond) \ + do { \ + if (not(cond)) { \ + printf("Assertion failed: %s:%d, condition: %s\n", __FILE__, __LINE__, #cond); \ + asm("trap;"); \ + } \ + } while (0) +#endif + +#ifndef EP_UNIFIED_ASSERT +#ifdef __CUDA_ARCH__ +#define EP_UNIFIED_ASSERT(cond) EP_DEVICE_ASSERT(cond) +#else +#define EP_UNIFIED_ASSERT(cond) EP_HOST_ASSERT(cond) +#endif +#endif + +#ifndef CUDA_RUNTIME_CHECK +#define CUDA_RUNTIME_CHECK(cmd) \ +do { \ + const auto e = (cmd); \ + if (e != cudaSuccess) { \ + std::stringstream ss; \ + ss << static_cast(e) << " (" << cudaGetErrorName(e) << ", " << cudaGetErrorString(e) << ")"; \ + throw EPException("CUDA runtime", __FILE__, __LINE__, ss.str()); \ + } \ +} while (0) +#endif + +#ifndef CUDA_DRIVER_CHECK +#define CUDA_DRIVER_CHECK(cmd) \ +do { \ + const auto e = (cmd); \ + if (e != CUDA_SUCCESS) { \ + std::stringstream ss; \ + const char *name, *info; \ + lazy_cuGetErrorName(e, &name), lazy_cuGetErrorString(e, &info); \ + ss << static_cast(e) << " (" << name << ", " << info << ")"; \ + throw EPException("CUDA driver", __FILE__, __LINE__, ss.str()); \ + } \ +} while (0) +#endif + +#ifndef NCCL_CHECK +#define NCCL_CHECK(cmd) \ +do { \ + const auto e = (cmd); \ + if (e != ncclSuccess) { \ + std::stringstream ss; \ + ss << static_cast(e) << " (" << ncclGetLastError(nullptr) << ")"; \ + throw EPException("NCCL", __FILE__, __LINE__, ss.str()); \ + } \ +} while (0) +#endif diff --git a/deep_ep/include/deep_ep/common/handle.cuh b/deep_ep/include/deep_ep/common/handle.cuh new file mode 100644 index 000000000..fbadc88cd --- /dev/null +++ b/deep_ep/include/deep_ep/common/handle.cuh @@ -0,0 +1,221 @@ +#pragma once + +#include +#include + +#include +#include + +namespace deep_ep::elastic::handle { + +struct NCCLGin { +#define IS_TEAM_WORLD(code) if constexpr (std::is_same_v) { code } +#define IS_TEAM_LSA(code) if constexpr (std::is_same_v) { code } +#define IS_TEAM_RAIL(code) if constexpr (std::is_same_v) { code } +#define IS_TEAM_WORLD_RAIL(code) if constexpr (std::is_same_v or std::is_same_v) { code } +#define IS_TEAM_WORLD_LSA(code) if constexpr (std::is_same_v or std::is_same_v) { code } +#define TEAM_WORLD_RAIL() ((std::is_same_v) ? team_world : team_rail) + + const ncclDevComm_t& nccl_dev_comm; + const ncclWindow_t& nccl_window; + ncclGin gin; + ncclTeam team_world, team_lsa, team_rail; + uint64_t lsa_base_ptr; + + // TODO(NCCL): QP index should just be a hint or the users maintain the mapping? + __device__ __forceinline__ + NCCLGin(const ncclDevComm_t& nccl_dev_comm, const ncclWindow_t& nccl_window, + const int& qp_idx = 0, + const ncclGinResourceSharingMode& resource_sharing_mode = NCCL_GIN_RESOURCE_SHARING_GPU): + nccl_dev_comm(nccl_dev_comm), nccl_window(nccl_window), + gin(ncclGin(nccl_dev_comm, qp_idx, resource_sharing_mode)), + team_world(ncclTeamWorld(nccl_dev_comm)), team_lsa(ncclTeamLsa(nccl_dev_comm)), team_rail(ncclTeamRail(nccl_dev_comm)) { + // TODO: what if we only have 1 NVLink rank + lsa_base_ptr = reinterpret_cast(ncclGetLsaPointer(nccl_window, 0, team_lsa.rank)); + } + + template + __device__ __forceinline__ bool is_nvlink_accessible(const int& dst_rank_idx) const { + IS_TEAM_LSA({ + return true; + }) + + IS_TEAM_WORLD({ + // TODO(NCCL): optimize this function's cycles + // return ncclTeamRankIsMember(team_lsa, team_world, dst_rank_idx); + return team_rail.rank * team_lsa.nRanks <= dst_rank_idx and + dst_rank_idx < (team_rail.rank + 1) * team_lsa.nRanks; + }) + + IS_TEAM_RAIL({ + // TODO(NCCL): some ranks may be connected by NVLink, e.g., "2 + 2 + 4" + return team_rail.rank == dst_rank_idx; + }) + } + + // ReSharper disable once CppNotAllPathsReturnValue + template + __device__ __forceinline__ + uint64_t get_sym_offset(dtype_t* ptr) const { + return reinterpret_cast(ptr) - lsa_base_ptr; + } + + // ReSharper disable once CppNotAllPathsReturnValue + template + __device__ __forceinline__ + dtype_t* get_sym_ptr(dtype_t* ptr, const int& dst_rank_idx) const { + IS_TEAM_RAIL({ + return team_rail.rank == dst_rank_idx ? ptr : nullptr; + }) + + IS_TEAM_WORLD_LSA({ + constexpr bool kIsTeamLSA = (std::is_same_v); + + // Team world and not accessible by symmetric pointers + if (not is_nvlink_accessible(dst_rank_idx)) + return nullptr; + + // Translate into NVLink rank index + const auto dst_nvl_rank_idx = kIsTeamLSA ? + dst_rank_idx : (dst_rank_idx - team_rail.rank * team_lsa.nRanks); + + // Local rank bypass + // TODO(NCCL): support this + if (dst_nvl_rank_idx == team_lsa.rank) + return ptr; + + // Get base ptr + const auto dst_ptr = ncclGetLsaPointer( + nccl_window, get_sym_offset(ptr), dst_nvl_rank_idx); + return static_cast(dst_ptr); + }); + } + + // NOTES: take care of this function when `team_t` is not LSA + // Do not mix atomic add with gin signal into a single position + template + __device__ __forceinline__ + void red_add_rel(dtype_t* sym_ptr, const dtype_t& value, const int& dst_rank_idx, + const int& extra_options = 0) const { + const auto dst_ptr = get_sym_ptr(sym_ptr, dst_rank_idx); + // Use symmetric pointers as much as possible, RDMA otherwise + if (dst_ptr != nullptr) { + // NOTES: local rank (or even NVLink-connected) for tag rail can also bypass + ptx::red_add_rel_sys(dst_ptr, value); + } else { + EP_DEVICE_ASSERT((not std::is_same_v)); + EP_DEVICE_ASSERT((std::is_same_v) or (std::is_same_v)); + // TODO(NCCL): support all dtypes + gin.signal(TEAM_WORLD_RAIL(), dst_rank_idx, + ncclGin_VASignalAdd(nccl_window, reinterpret_cast(sym_ptr) - lsa_base_ptr, static_cast(value)), + ncclCoopThread(), + ncclGin_None(), + cuda::thread_scope_thread, + cuda::thread_scope_device, + ncclGinOptFlagsDefault | extra_options); + } + } + + __device__ __forceinline__ + void wait(ncclGinRequest_t& request) const { + gin.wait(request); + } + + template + __device__ __forceinline__ + void get(void* src_ptr, void* dst_ptr, const int& num_bytes, const int& src_rank_idx, + const int& extra_options = 0) const { + IS_TEAM_WORLD_RAIL({ + gin.get( + TEAM_WORLD_RAIL(), + src_rank_idx, + nccl_window, reinterpret_cast(src_ptr) - lsa_base_ptr, + nccl_window, reinterpret_cast(dst_ptr) - lsa_base_ptr, + num_bytes, + coop_t(), + ncclGin_None(), + ncclGinOptFlagsDefault | extra_options, + segment_t() + ); + }); + } + + template + __device__ __forceinline__ + void flush_async(const int& src_rank_idx, ncclGinRequest_t* request, + const int& extra_options = 0) const { + IS_TEAM_WORLD_RAIL({ + gin.flushAsync( + TEAM_WORLD_RAIL(), + src_rank_idx, + request, + coop_t(), + ncclGinOptFlagsDefault | extra_options + ); + }); + } + + template + __device__ __forceinline__ + void signal(const int& dst_rank_idx, const remote_action_t& remote_action) const { + IS_TEAM_WORLD_RAIL({ + gin.signal(TEAM_WORLD_RAIL(), dst_rank_idx, remote_action); + }); + } + + template + __device__ __forceinline__ + void put(void* recv_sym_ptr, void* send_sym_ptr, const int& num_bytes, const int& dst_rank_idx, + const int& extra_options = 0, + const remote_action_t& remote_action = remote_action_t()) const { + // NOTES: local or NVLink put will also go through NIC via this API + IS_TEAM_WORLD_RAIL({ + gin.put(TEAM_WORLD_RAIL(), + dst_rank_idx, + // TODO: can we don't repeat the window? + // TODO: can we pass raw pointers? + nccl_window, reinterpret_cast(recv_sym_ptr) - lsa_base_ptr, + nccl_window, reinterpret_cast(send_sym_ptr) - lsa_base_ptr, + num_bytes, + remote_action, + ncclGin_None(), + ncclCoopThread(), + ncclGin_None(), + cuda::thread_scope_thread, + cuda::thread_scope_device, + ncclGinOptFlagsDefault | extra_options); + }); + } + + template + __device__ __forceinline__ + void put_value(dtype_t* sym_ptr, const dtype_t& value, const int& dst_rank_idx, + const int& extra_options = 0) const { + const auto dst_ptr = get_sym_ptr(sym_ptr, dst_rank_idx); + if (dst_ptr != nullptr) { + ptx::st_relaxed_sys(dst_ptr, value); + } else { + EP_DEVICE_ASSERT((not std::is_same_v)); + gin.putValue(TEAM_WORLD_RAIL(), + dst_rank_idx, + nccl_window, reinterpret_cast(sym_ptr) - lsa_base_ptr, + value, + ncclGin_None(), + ncclCoopThread(), + ncclGin_None(), + cuda::thread_scope_thread, + cuda::thread_scope_device, + ncclGinOptFlagsDefault | extra_options); + } + } + +#undef IS_TEAM_WORLD +#undef IS_TEAM_LSA +#undef IS_TEAM_RAIL +#undef IS_TEAM_WORLD_RAIL +#undef IS_TEAM_WORLD_LSA +#undef TEAM_WORLD_RAIL +}; + +} // namespace deep_ep::elastic::handle diff --git a/deep_ep/include/deep_ep/common/layout.cuh b/deep_ep/include/deep_ep/common/layout.cuh new file mode 100644 index 000000000..7c8fa50ed --- /dev/null +++ b/deep_ep/include/deep_ep/common/layout.cuh @@ -0,0 +1,314 @@ +#pragma once + +#include +#include +#include +#include + +namespace deep_ep::elastic::layout { + +struct WorkspaceLayout { + void* workspace; + + int num_ranks; + int num_scaleout_ranks, num_scaleup_ranks; + int num_experts, num_experts_per_rank; + + // We want to fix the layout position for all settings, + // so that one buffer can be reused for all cases + static constexpr int kNumMaxRanks = 1024; + static constexpr int kNumMaxExperts = 2048; + static constexpr int kNumMaxExpertsPerRank = 256; + static constexpr int kNumMaxInflightAGRS = 32; + + static constexpr int64_t kNumBarrierSignalBytes = 16; + + __forceinline__ __device__ __host__ + WorkspaceLayout(void* workspace, + const int& num_scaleout_ranks, + const int& num_scaleup_ranks, + const int& num_experts): + workspace(workspace), + num_ranks(num_scaleout_ranks * num_scaleup_ranks), + num_scaleout_ranks(num_scaleout_ranks), + num_scaleup_ranks(num_scaleup_ranks), + num_experts(num_experts) { + num_experts_per_rank = num_experts / num_ranks; + EP_UNIFIED_ASSERT(num_experts % num_ranks == 0); + EP_UNIFIED_ASSERT(num_ranks <= kNumMaxRanks); + EP_UNIFIED_ASSERT(num_experts <= kNumMaxExperts); + EP_UNIFIED_ASSERT(num_experts_per_rank <= kNumMaxExpertsPerRank); + } + + static int64_t get_num_bytes() { + // Pure NVLink scaleup barrier signals + int64_t num_bytes = 0; + num_bytes += kNumBarrierSignalBytes; + + // Notify reduction workspace + num_bytes += (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t); + + // Scaleup notify threads + // Rank send/recv count + num_bytes += kNumMaxRanks * sizeof(int64_t) * 2; + // Expert send/recv count + num_bytes += kNumMaxExperts * sizeof(int64_t) * 2; + + // Scaleup atomic sender count + num_bytes += kNumMaxRanks * sizeof(int); + + // Scaleout notify threads + // Rank send/recv count + num_bytes += kNumMaxRanks * sizeof(int) * 2; + // Expert send/recv count + num_bytes += kNumMaxExperts * sizeof(int) * 2; + + // Scaleout channel metadata (finish flag and tails) + num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int64_t); + + // Channel aggregated into the scaleup domains + // Also reused for channel scaleup tail + num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int); + + // Rank send/recv count, for PP prev/next ranks + num_bytes += 2 * 2 * sizeof(int64_t); + + // AGRS signals + num_bytes += (kNumMaxInflightAGRS + 1) * kNumMaxRanks * sizeof(int); + + // Ensure LDG.256 work + return math::align(num_bytes, 32); + } + + __forceinline__ __device__ __host__ unsigned long long* get_nvl_barrier_counter_ptr() const { + return static_cast(workspace); + } + + __forceinline__ __device__ __host__ int* get_nvl_barrier_signal_ptr(const int& phase) const { + return math::advance_ptr(workspace, (2 + phase) * sizeof(int)); + } + + __forceinline__ __device__ __host__ int64_t* get_notify_reduction_workspace_ptr() const { + return math::advance_ptr(workspace, kNumBarrierSignalBytes); + } + + template + __forceinline__ __device__ __host__ int64_t* get_scaleup_rank_expert_count_ptr() const { + const auto base_ptr = + math::advance_ptr(get_notify_reduction_workspace_ptr(), (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t)); + return base_ptr + (kIsSendBuffer ? 0 : kNumMaxRanks + kNumMaxExperts); + } + + template + __forceinline__ __device__ __host__ int64_t* get_scaleup_rank_count_ptr() const { + return get_scaleup_rank_expert_count_ptr(); + } + + template + __forceinline__ __device__ __host__ int64_t* get_scaleup_expert_count_ptr() const { + return get_scaleup_rank_expert_count_ptr() + num_scaleup_ranks; + } + + __forceinline__ __device__ __host__ int* get_scaleup_atomic_sender_counter() const { + return math::advance_ptr( + get_scaleup_rank_expert_count_ptr(), 2 * (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t)); + } + + template + __forceinline__ __device__ __host__ int* get_scaleout_rank_expert_count_ptr() const { + const auto base_ptr = + math::advance_ptr(get_scaleup_atomic_sender_counter(), kNumMaxRanks * sizeof(int)); + return base_ptr + (kIsSendBuffer ? 0 : kNumMaxRanks + kNumMaxExperts); + } + + template + __forceinline__ __device__ __host__ int* get_scaleout_rank_count_ptr( + const int& scaleout_rank_idx = 0, const int& scaleup_rank_idx = 0) const { + const auto base_ptr = get_scaleout_rank_expert_count_ptr(); + return base_ptr + scaleout_rank_idx * num_scaleup_ranks + scaleup_rank_idx; + } + + template + __forceinline__ __device__ __host__ int* get_scaleout_expert_count_ptr( + const int& scaleout_rank_idx = 0, const int& expert_idx = 0) const { + const auto base_ptr = get_scaleout_rank_expert_count_ptr() + num_ranks; + return base_ptr + scaleout_rank_idx * (num_scaleup_ranks * num_experts_per_rank) + expert_idx; + } + + __forceinline__ __device__ __host__ int64_t* get_scaleout_channel_signaled_tail_ptr( + const int& channel_idx, const int& scaleout_rank_idx) const { + const auto base_ptr = math::advance_ptr( + get_scaleout_rank_expert_count_ptr(), + (kNumMaxRanks + kNumMaxExperts) * sizeof(int) * 2); + return base_ptr + (channel_idx * num_scaleout_ranks + scaleout_rank_idx); + } + + __forceinline__ __device__ __host__ int* get_channel_scaleup_tail_ptr( + const int& channel_idx, const int& scaleup_rank_idx) const { + const auto base_ptr = math::advance_ptr( + get_scaleout_channel_signaled_tail_ptr(0, 0), + kNumMaxRanks * kNumMaxChannels * sizeof(int64_t)); + return base_ptr + (channel_idx * num_scaleup_ranks + scaleup_rank_idx); + } + + __forceinline__ __device__ __host__ int64_t* get_pp_send_count_ptr(const int& offset) const { + const auto base_ptr = math::advance_ptr( + get_channel_scaleup_tail_ptr(0, 0), + kNumMaxRanks * kNumMaxChannels * sizeof(int)); + return base_ptr + offset; + } + + __forceinline__ __device__ __host__ int64_t* get_pp_recv_count_ptr(const int& offset) const { + const auto base_ptr = math::advance_ptr( + get_pp_send_count_ptr(0), 2 * sizeof(int64_t)); + return base_ptr + offset; + } + + __forceinline__ __device__ __host__ int* get_agrs_recv_signal_ptr(const int& slot, const int& rank_idx) const { + const auto base_ptr = math::advance_ptr( + get_pp_recv_count_ptr(0), 2 * sizeof(int64_t)); + return base_ptr + slot * kNumMaxRanks + rank_idx; + } + + __forceinline__ __device__ __host__ int* get_agrs_session_signal_ptr(const int& rank_idx) const { + const auto base_ptr = math::advance_ptr( + get_agrs_recv_signal_ptr(0, 0), kNumMaxInflightAGRS * kNumMaxRanks * sizeof(int)); + return base_ptr + rank_idx; + } +}; + +struct TokenLayout { + int num_hidden_bytes, num_sf_bytes; + // NOTES: the top-k index is always 32-bit + bool with_metadata; + int num_topk, num_metadata_bytes; + void* base; + + __forceinline__ __device__ __host__ + TokenLayout(const int& num_hidden_bytes, const int& num_sf_bytes, + const int& num_topk, const bool& with_metadata, void* base = nullptr) : + num_hidden_bytes(num_hidden_bytes), + num_sf_bytes(num_sf_bytes), + // Metadata includes: top-k indices, weight and source rank/token index + with_metadata(with_metadata), + num_topk(num_topk), + num_metadata_bytes(num_topk * (sizeof(int) + sizeof(float)) + + (with_metadata ? (1 + num_topk) * sizeof(int) : 0)), + base(base) { + EP_STATIC_ASSERT(sizeof(int) == sizeof(float), "Invalid size assumption"); + EP_UNIFIED_ASSERT(num_hidden_bytes % ptx::kNumTMAAlignBytes == 0); + } + + template + __forceinline__ __device__ __host__ dtype_t get_num_bytes() const { + const auto num_bytes = math::align(num_hidden_bytes, ptx::kNumTMAAlignBytes) + + math::align(num_sf_bytes, ptx::kNumTMAAlignBytes) + + math::align(num_metadata_bytes, ptx::kNumTMAAlignBytes) + + math::align(kWithMBarrier ? sizeof(ptx::mbarrier) : 0, ptx::kNumTMAAlignBytes); + return static_cast(num_bytes); + } + + __forceinline__ __device__ __host__ void* get_base_ptr() const { + return base; + } + + __forceinline__ __device__ __host__ void set_base_ptr(void* ptr) { + base = ptr; + } + + __forceinline__ __device__ __host__ void* get_hidden_ptr() const { + return get_base_ptr(); + } + + __forceinline__ __device__ __host__ sf_pack_t* get_sf_ptr() const { + return math::advance_ptr(base, math::align(num_hidden_bytes, ptx::kNumTMAAlignBytes)); + } + + __forceinline__ __device__ __host__ int* get_metadata_ptr() const { + return math::advance_ptr(get_sf_ptr(), math::align(num_sf_bytes, ptx::kNumTMAAlignBytes)); + } + + __forceinline__ __device__ __host__ int* get_topk_idx_ptr() const { + return get_metadata_ptr(); + } + + __forceinline__ __device__ __host__ float* get_topk_weights_ptr() const { + return math::advance_ptr(get_metadata_ptr(), num_topk * sizeof(int)); + } + + __forceinline__ __device__ __host__ int* get_src_token_global_idx_ptr() const { + return math::advance_ptr(get_topk_weights_ptr(), num_topk * sizeof(float)); + } + + __forceinline__ __device__ __host__ int* get_linked_list_idx_ptr() const { + return get_src_token_global_idx_ptr() + 1; + } + + __forceinline__ __device__ ptx::mbarrier* get_mbarrier_ptr() const { + return math::advance_ptr(get_metadata_ptr(), math::align(num_metadata_bytes, ptx::kNumTMAAlignBytes)); + } +}; + +template +struct BufferLayout { + TokenLayout token_layout; + int num_ranks; + int num_max_tokens_per_rank; + + void* base; + + __forceinline__ __device__ __host__ + BufferLayout(const TokenLayout& token_layout, + const int& num_ranks, + const int& max_num_tokens_per_rank, + void* base = nullptr) : + token_layout(token_layout), + num_ranks(num_ranks), num_max_tokens_per_rank(max_num_tokens_per_rank), + base(base) {} + + __forceinline__ __device__ __host__ + int64_t get_num_bytes_per_token() const { + return token_layout.get_num_bytes(); + } + + __forceinline__ __device__ __host__ + int64_t get_num_bytes_per_rank() const { + return num_max_tokens_per_rank * get_num_bytes_per_token(); + } + + __forceinline__ __device__ __host__ + int64_t get_num_bytes() const { + return get_num_bytes_per_rank() * num_ranks; + } + + __forceinline__ __device__ __host__ + void* get_buffer_end_ptr() const { + return math::advance_ptr(base, get_num_bytes()); + } + + __forceinline__ __device__ __host__ + BufferLayout get_rank_buffer(const int& rank_idx) const { + return BufferLayout(token_layout, + 1, num_max_tokens_per_rank, + static_cast(base) + get_num_bytes_per_rank() * rank_idx); + } + + template + __forceinline__ __device__ __host__ + BufferLayout get_channel_buffer(const int& channel_idx) const { + EP_UNIFIED_ASSERT(num_max_tokens_per_rank % kNumTokensPerChannel == 0); + return BufferLayout(token_layout, + // Do not use `num_max_tokens_per_rank / kNumTokensPerChannel` as the false stride + num_ranks, num_max_tokens_per_rank, + static_cast(base) + get_num_bytes_per_token() * kNumTokensPerChannel * channel_idx); + } + + __forceinline__ __device__ __host__ + TokenLayout get_token_buffer(const int& token_idx, const bool& global = false) const { + EP_UNIFIED_ASSERT(num_ranks == 1 or global); + return TokenLayout(token_layout.num_hidden_bytes, token_layout.num_sf_bytes, token_layout.num_topk, token_layout.with_metadata, + static_cast(base) + token_layout.get_num_bytes() * token_idx); + } +}; + +} // namespace deep_ep::elastic diff --git a/deep_ep/include/deep_ep/common/math.cuh b/deep_ep/include/deep_ep/common/math.cuh new file mode 100644 index 000000000..a042c650c --- /dev/null +++ b/deep_ep/include/deep_ep/common/math.cuh @@ -0,0 +1,68 @@ +#pragma once + +namespace deep_ep::elastic::math { + +template +__forceinline__ __device__ __host__ T ceil_div(T a, T b) { + return (a + b - 1) / b; +} + +template +__forceinline__ __device__ __host__ constexpr T constexpr_ceil_div(T a, T b) { + return (a + b - 1) / b; +} + +template +__forceinline__ __device__ __host__ T align(T a, T b) { + return (kDoCeilAlignment ? ceil_div(a, b) : (a / b)) * b; +} + +template +__forceinline__ __device__ __host__ constexpr T constexpr_align(T a, T b) { + return (kDoCeilAlignment ? constexpr_ceil_div(a, b) : (a / b)) * b; +} + +template +__forceinline__ __device__ __host__ bool is_decoded_positive_ready(const dtype_t& value) { + return value >= 0; +} + +template +__forceinline__ __device__ __host__ dtype_t encode_decode_positive(const dtype_t& value) { + return -value - static_cast(1); +} + +template +__forceinline__ __device__ __host__ dtype_t* advance_ptr(void* ptr, const int64_t num_bytes) { + return reinterpret_cast(static_cast(ptr) + num_bytes); +} + +__forceinline__ __device__ __host__ ptrdiff_t ptr_diff(const void* ptr, const void* base) { + return static_cast(ptr) - static_cast(base); +} + +template +__device__ __forceinline__ dtype_b_t pack2(const dtype_a_t& x, const dtype_a_t& y) { + EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), "Invalid dtypes"); + dtype_b_t packed; + auto unpacked_ptr = reinterpret_cast(&packed); + unpacked_ptr[0] = x, unpacked_ptr[1] = y; + return packed; +} + +template +__device__ __forceinline__ std::tuple unpack2(const dtype_b_t& packed) { + EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), "Invalid dtypes"); + auto unpacked_ptr = reinterpret_cast(&packed); + dtype_a_t x = unpacked_ptr[0], y = unpacked_ptr[1]; + return {x, y}; +} + +template +__device__ __forceinline__ void unpack2(const dtype_b_t& packed, dtype_a_t& x, dtype_a_t& y) { + EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), "Invalid dtypes"); + auto unpacked_ptr = reinterpret_cast(&packed); + x = unpacked_ptr[0], y = unpacked_ptr[1]; +} + +} // namespace deep_ep::elastic::math diff --git a/deep_ep/include/deep_ep/common/ptx.cuh b/deep_ep/include/deep_ep/common/ptx.cuh new file mode 100644 index 000000000..7d2a4b4c7 --- /dev/null +++ b/deep_ep/include/deep_ep/common/ptx.cuh @@ -0,0 +1,428 @@ +#pragma once + +#include + +#include +#include + +namespace deep_ep::elastic::ptx { + +// Host-side placeholder with the same size/alignment as cuda::barrier +// (a single uint64_t atomic), so that sizeof(mbarrier) is consistent across host and device. +struct alignas(8) mbarrier { uint64_t __placeholder; }; +using arrival_phase = uint32_t; + +// More than TMA, `longlong4` requires 32 bytes aligned +static constexpr int kNumTMAAlignBytes = 32; + +#ifdef __CUDACC__ + +/// Exceptions +__forceinline__ __device__ void trap() { + asm volatile("trap;"); +} + +/// Thread layout +__forceinline__ __device__ int get_warp_idx() { + return __shfl_sync(0xffffffff, threadIdx.x / 32, 0); +} + +__forceinline__ __device__ int get_lane_idx() { + int lane_idx; + asm volatile("mov.s32 %0, %laneid;" : "=r"(lane_idx)); + return lane_idx; +} + +/// Election +__forceinline__ __device__ int elect_one_sync() { +#ifndef DISABLE_SM90_FEATURES + int pred = 0; + asm volatile( + "{\n" + ".reg .b32 %%rx;\n" + ".reg .pred %%px;\n" + " elect.sync %%rx|%%px, %1;\n" + "@%%px mov.s32 %0, 1;\n" + "}\n" + : "+r"(pred) + : "r"(0xffffffff)); + return pred; +#else + return get_lane_idx() == 0; +#endif +} + +/// TMA and `cp.async` +__forceinline__ __device__ void mbarrier_init_with_fence(mbarrier* ptr, const int& arrive_count = 1) { + asm volatile("mbarrier.init.shared::cta.b64 [%1], %0;" :: + "r"(arrive_count), "r"(static_cast(__cvta_generic_to_shared(ptr)))); + asm volatile("fence.mbarrier_init.release.cluster;" ::); +} + +__forceinline__ __device__ void mbarrier_invalidate(mbarrier* ptr) { + asm volatile("mbarrier.inval.shared::cta.b64 [%0];" :: + "r"(static_cast(__cvta_generic_to_shared(ptr)))); +} + +__forceinline__ __device__ void mbarrier_arrive(mbarrier* ptr) { + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0]; \n\t" :: + "r"(static_cast(__cvta_generic_to_shared(ptr)))); +} + +__forceinline__ __device__ void mbarrier_arrive_and_set_tx(mbarrier* ptr, const int& num_bytes) { + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%1], %0; \n\t" :: + "r"(num_bytes), "r"(static_cast(__cvta_generic_to_shared(ptr)))); +} + +__forceinline__ __device__ void mbarrier_wait_and_flip_phase(mbarrier* ptr, arrival_phase& phase) { + asm volatile( + "{\n\t" + ".reg .pred P1; \n\t" + "LAB_WAIT: \n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2; \n\t" + "@P1 bra DONE; \n\t" + "bra LAB_WAIT; \n\t" + "DONE: \n\t" + "}" :: + "r"(static_cast(__cvta_generic_to_shared(ptr))), + "r"(phase), "r"(0x989680)); + phase ^= 1; +} + +__forceinline__ __device__ void tma_store_fence() { + asm volatile("fence.proxy.async.shared::cta;"); +} + +template +__forceinline__ __device__ void tma_store_wait() { + asm volatile("cp.async.bulk.wait_group %0;" ::"n"(kNumRemainingWaits) : "memory"); +} + +enum TMACacheHint: int64_t { + kEvictFirst = 0x12f0000000000000ll, + kEvictNormal = 0x1000000000000000ll +}; + +__forceinline__ __device__ void tma_load_1d( + const void* dst_ptr, const void* src_ptr, mbarrier* ptr, const int& num_bytes, + const TMACacheHint& hint = TMACacheHint::kEvictFirst) { + // NOTES: normally, the loaded part will be evicted soon + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;\n" :: + "r"(static_cast(__cvta_generic_to_shared(dst_ptr))), + "l"(src_ptr), + "r"(num_bytes), + "r"(static_cast(__cvta_generic_to_shared(ptr))), + "l"(hint) + : "memory"); +} + +__forceinline__ __device__ void tma_store_1d( + const void* dst_ptr, const void* src_ptr, const int& num_bytes, + const TMACacheHint& hint = TMACacheHint::kEvictNormal) { + // NOTES: normally, the stored part will be used soon + asm volatile("cp.async.bulk.global.shared::cta.bulk_group.L2::cache_hint [%0], [%1], %2, %3;\n" :: + "l"(dst_ptr), + "r"(static_cast(__cvta_generic_to_shared(src_ptr))), + "r"(num_bytes), + "l"(hint) + : "memory"); +} + +__forceinline__ __device__ void tma_store_commit() { + asm volatile("cp.async.bulk.commit_group;"); +} + +template +__forceinline__ __device__ void cp_async_ca(const dtype_t* gmem_src, const dtype_t* smem_dst) { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8 or sizeof(dtype_t) == 16, "Invalid dtype bytes"); + asm volatile("cp.async.ca.shared::cta.global.L2::128B [%0], [%1], %2;\n" :: + "r"(static_cast(__cvta_generic_to_shared(smem_dst))), + "l"(gmem_src), + "n"(sizeof(dtype_t))); +} + +__forceinline__ __device__ void cp_async_mbarrier_arrive(mbarrier* ptr) { + asm volatile("cp.async.mbarrier.arrive.shared::cta.b64 [%0];\n" :: + "r"(static_cast(__cvta_generic_to_shared(ptr)))); +} + +/// Barriers +template +__forceinline__ __device__ void named_barrier(const int& idx) { + // Equivalent to `barrier.sync.aligned`, which requires all threads run the same location of code + asm volatile("bar.sync %0, %1;" ::"r"(idx), "r"(kNumThreads)); +} + +/// LD/ST instructions +__forceinline__ __device__ int4 ldg_with_gez_pred(const int4* ptr, const int& value, const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) { + int4 ret = make_int4(0, 0, 0, 0); + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.ge.s32 p, %5, 0;\n\t" + " @p ld.L1::no_allocate.L2::cache_hint.global.nc.v4.s32 {%0, %1, %2, %3}, [%4], %6;\n\t" + "}" + : "+r"(ret.x), "+r"(ret.y), "+r"(ret.z), "+r"(ret.w) + : "l"(ptr), "r"(value), "l"(cache_hint) + : "memory"); + return ret; +} + +__forceinline__ __device__ int4 ldg_with_gtz_pred(const int4* ptr, const int& value, const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) { + int4 ret = make_int4(0, 0, 0, 0); + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.gt.s32 p, %5, 0;\n\t" + " @p ld.L1::no_allocate.L2::cache_hint.global.nc.v4.s32 {%0, %1, %2, %3}, [%4], %6;\n\t" + "}" + : "+r"(ret.x), "+r"(ret.y), "+r"(ret.z), "+r"(ret.w) + : "l"(ptr), "r"(value), "l"(cache_hint) + : "memory"); + return ret; +} + +__forceinline__ __device__ int4 ld_with_gez_pred(const int4* ptr, const int& value, const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) { + int4 ret = make_int4(0, 0, 0, 0); + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.ge.s32 p, %5, 0;\n\t" + " @p ld.L1::no_allocate.L2::cache_hint.global.v4.s32 {%0, %1, %2, %3}, [%4], %6;\n\t" + "}" + : "+r"(ret.x), "+r"(ret.y), "+r"(ret.z), "+r"(ret.w) + : "l"(ptr), "r"(value), "l"(cache_hint) + : "memory"); + return ret; +} + +#if defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) +__forceinline__ __device__ longlong4_t ldg_with_gez_pred(const longlong4_t* ptr, const int& value, const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) { + longlong4_t ret = make_longlong4_t(0, 0, 0, 0); + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.ge.s32 p, %5, 0;\n\t" + " @p ld.L1::no_allocate.L2::cache_hint.global.nc.v4.s64 {%0, %1, %2, %3}, [%4], %6;\n\t" + "}" + : "+l"(ret.x), "+l"(ret.y), "+l"(ret.z), "+l"(ret.w) + : "l"(ptr), "r"(value), "l"(cache_hint) + : "memory"); + return ret; +} + +__forceinline__ __device__ longlong4_t ldg(const longlong4_t* ptr) { + longlong4_t ret; + asm volatile( + "ld.L1::no_allocate.global.nc.v4.s64 {%0, %1, %2, %3}, [%4];\n\t" + : "=l"(ret.x), "=l"(ret.y), "=l"(ret.z), "=l"(ret.w) + : "l"(ptr) + : "memory"); + return ret; +} +#endif + +__forceinline__ __device__ int4 ldg(const int4* ptr) { + return __ldg(ptr); +} + +template +__forceinline__ __device__ void st_with_gez_pred(dtype_t* ptr, dtype_t value, const int& condition) { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4, "Invalid data type"); + auto view = *reinterpret_cast(&value); + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.ge.s32 p, %2, 0;\n\t" + " @p st.global.s32 [%0], %1;\n\t" + "}" + :: "l"(ptr), "r"(view), "r"(condition) + : "memory"); +} + +template +__forceinline__ __device__ dtype_t ld_volatile(const void* ptr) { + if constexpr (sizeof(dtype_t) == 4) { + uint32_t value; + asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(value) : "l"(ptr)); + return reinterpret_cast(value); + } else if constexpr (sizeof(dtype_t) == 8) { + uint64_t value; + asm volatile("ld.volatile.global.u64 %0, [%1];" : "=l"(value) : "l"(ptr)); + return reinterpret_cast(value); + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, "Invalid data type length"); + } +} + +__forceinline__ __device__ void red_add(const int64_t* ptr, const int64_t& value) { + // TODO(NVCC): why don't NVCC support `s64`? + asm volatile("red.gpu.global.add.u64 [%0], %1;" :: "l"(ptr), "l"(value)); +} + +__forceinline__ __device__ void red_add_rel_sys(const int* ptr, const int& value) { + asm volatile("red.release.sys.global.add.s32 [%0], %1;" :: "l"(ptr), "r"(value)); +} + +__forceinline__ __device__ void red_add_rel_sys(const int64_t* ptr, const int64_t& value) { + asm volatile("red.release.sys.global.add.u64 [%0], %1;" :: "l"(ptr), "l"(value)); +} + +template +__forceinline__ __device__ dtype_t ld_acquire_sys(const dtype_t* ptr) { + if constexpr (sizeof(dtype_t) == 4) { + uint32_t value; + asm volatile("ld.acquire.sys.L1::no_allocate.global.u32 %0, [%1];" : "=r"(value) : "l"(ptr)); + return reinterpret_cast(value); + } else if constexpr (sizeof(dtype_t) == 8) { + uint64_t value; + asm volatile("ld.acquire.sys.L1::no_allocate.global.u64 %0, [%1];" : "=l"(value) : "l"(ptr)); + return reinterpret_cast(value); + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, "Invalid data type length"); + } +} + +template +__forceinline__ __device__ void st_relaxed_sys(void* ptr, dtype_t value) { + if constexpr (sizeof(dtype_t) == 4) { + uint32_t int_value = reinterpret_cast(value); + asm volatile("st.relaxed.sys.global.u32 [%0], %1;" :: "l"(ptr), "r"(int_value)); + } else if constexpr (sizeof(dtype_t) == 8) { + uint64_t int_value = reinterpret_cast(value); + asm volatile("st.relaxed.sys.global.u64 [%0], %1;" :: "l"(ptr), "l"(int_value)); + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, "Invalid data type length"); + } +} + +template +__forceinline__ __device__ void st_release_sys(void* ptr, dtype_t value) { + if constexpr (sizeof(dtype_t) == 4) { + uint32_t int_value = reinterpret_cast(value); + asm volatile("st.release.sys.global.u32 [%0], %1;" :: "l"(ptr), "r"(int_value)); + } else if constexpr (sizeof(dtype_t) == 8) { + uint64_t int_value = reinterpret_cast(value); + asm volatile("st.release.sys.global.u64 [%0], %1;" :: "l"(ptr), "l"(int_value)); + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, "Invalid data type length"); + } +} + +// Adjust registers +template +__device__ __forceinline__ void warpgroup_reg_alloc(){ + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(kNumRegs)); +} + +template +__device__ __forceinline__ void warpgroup_reg_dealloc(){ + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(kNumRegs)); +} + +/// General fences +__device__ __forceinline__ void fence_acq_rel_sys() { + asm volatile("fence.acq_rel.sys;" ::: "memory"); +} + +/// Intrinsics +template +__device__ __forceinline__ dtype_t exchange(dtype_t ptr, const int& src_lane_idx) { + EP_STATIC_ASSERT(sizeof(dtype_t) % sizeof(int) == 0, ""); + const auto send_int_values = reinterpret_cast(&ptr); + dtype_t recv_dtype; + auto recv_int_values = reinterpret_cast(&recv_dtype); + #pragma unroll + for (int i = 0; i < sizeof(dtype_t) / sizeof(int); ++i) + recv_int_values[i] = __shfl_sync(0xffffffff, send_int_values[i], src_lane_idx); + return recv_dtype; +} + +__device__ __forceinline__ unsigned gather(const bool& value) { + return __ballot_sync(0xffffffff, value); +} + +__device__ __forceinline__ bool all(const bool& value) { + return __all_sync(0xffffffff, value); +} + +__device__ __forceinline__ bool any(const bool& value) { + return __any_sync(0xffffffff, value); +} + +__device__ __forceinline__ unsigned reduce_or(const unsigned& value) { + return __reduce_or_sync(0xffffffff, value); +} + +__device__ __forceinline__ unsigned long long reduce_or(const unsigned long long& value) { + const auto low = __reduce_or_sync(0xffffffff, static_cast(value)); + const auto high = __reduce_or_sync(0xffffffff, static_cast(value >> 32)); + return (static_cast(high) << 32) | low; +} + +__device__ __forceinline__ int reduce_add(const int& value) { + return __reduce_add_sync(0xffffffff, value); +} + +__device__ __forceinline__ unsigned match(const int& value) { + return __match_any_sync(0xffffffff, value); +} + +__device__ __forceinline__ int fns(const unsigned& value, const int& offset) { + return __fns(value, 0, offset); +} + +template +__device__ __forceinline__ auto ffs(const dtype_t& value) { + if constexpr (sizeof(dtype_t) == 4) { + return __ffs(static_cast(value)) - 1; + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 8, "Invalid data type"); + return __ffsll(static_cast(value)) - 1; + } +} + +__device__ __forceinline__ int get_master_lane_idx(const unsigned& mask) { + // Equivalent to `31 - __clz(mask)` + int highest_idx; + asm volatile("bfind.u32 %0, %1;" : "=r"(highest_idx) : "r"(mask)); + return highest_idx; +} + +__device__ __forceinline__ bool deduplicate(const int& value, const int& lane_idx) { + return get_master_lane_idx(match(value)) == lane_idx; +} + +__device__ __forceinline__ int warp_inclusive_sum(int value, const int& lane_idx) { + #pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { + const auto synced = __shfl_up_sync(0xffffffff, value, offset); + if (lane_idx >= offset) + value += synced; + } + return value; +} + +__device__ __forceinline__ float2 fadd2(const float2& a, const float2& b) { +#if defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) + return __fadd2_rn(a, b); +#else + return {a.x + b.x, a.y + b.y}; +#endif +} + +__device__ __forceinline__ void accumulate(float2& a, nv_bfloat162 b) { +#if defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) + // Use `add.rn.f32.bf16` instruction to perform fused (cast + add) operation on SM100 + asm("add.rn.f32.bf16 %0, %1, %0;\n" : "+f"(a.x) : "h"(*reinterpret_cast(&b.x))); + asm("add.rn.f32.bf16 %0, %1, %0;\n" : "+f"(a.y) : "h"(*reinterpret_cast(&b.y))); +#else + const auto [x, y] = __bfloat1622float2(b); + a.x += x, a.y += y; +#endif +} + +#endif + +} // namespace deep_ep::elastic::ptx diff --git a/deep_ep/include/deep_ep/impls/barrier.cuh b/deep_ep/include/deep_ep/impls/barrier.cuh new file mode 100644 index 000000000..b0ef5ec08 --- /dev/null +++ b/deep_ep/include/deep_ep/impls/barrier.cuh @@ -0,0 +1,28 @@ +#pragma once + +#include +#include +#include +#include + + +namespace deep_ep::elastic { + +template +__global__ void __launch_bounds__(kNumThreads, 1) +barrier_impl(const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window, void* workspace, + const int scaleout_rank_idx, const int scaleup_rank_idx) { + const auto sm_idx = static_cast(blockIdx.x), thread_idx = static_cast(threadIdx.x); + + // Barrier only uses the first part of workspace, so making `num_experts` as 0 is fine + const auto workspace_layout = layout::WorkspaceLayout(workspace, kNumScaleoutRanks, kNumScaleupRanks, 0); + const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, 0); + comm::gpu_barrier( + gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx); +} + +} // namespace deep_ep::elastic diff --git a/deep_ep/include/deep_ep/impls/combine.cuh b/deep_ep/include/deep_ep/impls/combine.cuh new file mode 100644 index 000000000..79f3931c7 --- /dev/null +++ b/deep_ep/include/deep_ep/impls/combine.cuh @@ -0,0 +1,239 @@ +#pragma once + +#include + +#include +#include +#include +#include + +#include + + +namespace deep_ep::elastic { + +template (), + int kNumTokensInLayout = get_num_tokens_in_layout(), + typename team_t = std::conditional_t> +__global__ void __launch_bounds__(kNumThreads, 1) +combine_impl(nv_bfloat16* x, + float* topk_weights, + int* src_metadata, int* psum_num_recv_tokens_per_scaleup_rank, + const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window, + void* buffer, void* workspace, + const int rank_idx, + int num_reduced_tokens) { + // Utils + const auto sm_idx = static_cast(blockIdx.x); + const auto thread_idx = static_cast(threadIdx.x); + const auto warp_idx = (ptx::get_warp_idx() + rank_idx) % kNumWarps; + const auto lane_idx = ptx::get_lane_idx(); + const auto global_warp_idx = warp_idx * kNumSMs + sm_idx; + constexpr bool kDoExpandedSend = not kAllowMultipleReduction and kUseExpandedLayout; + + // We should assign the real number of received tokens if without CPU sync + if (num_reduced_tokens == kNumMaxTokensPerRank * kNumRanks) + num_reduced_tokens = __ldg(psum_num_recv_tokens_per_scaleup_rank + kNumRanks - 1); + + // Buffer layouts + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + const auto token_layout = layout::TokenLayout(kNumHiddenBytes, 0, kNumTopk, false); + const auto tma_buffer = layout::BufferLayout(token_layout, kNumWarps, 1, smem) + .get_rank_buffer(warp_idx).get_token_buffer(0); + const auto recv_buffer = layout::BufferLayout( + token_layout, kNumTokensInLayout, kNumMaxTokensPerRank, buffer); + const auto send_buffer = layout::BufferLayout( + token_layout, kNumRanks, + kNumMaxTokensPerRank * (kDoExpandedSend ? kNumTopk : 1), + recv_buffer.get_buffer_end_ptr()); + + // Init TMA + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (ptx::elect_one_sync()) + ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Expanding mode must not be backward + if constexpr (kUseExpandedLayout) + EP_DEVICE_ASSERT(topk_weights == nullptr); + + // Gin handle + // We treat each warp as a "channel" + const auto [qp_idx, sharing_mode] = comm::get_qp_mode(sm_idx, warp_idx); + const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, sharing_mode); + + // Full barrier to ensure the remote buffer is available + const auto workspace_layout = layout::WorkspaceLayout(workspace, 1, kNumRanks, kNumExperts); + comm::gpu_barrier( + gin, workspace_layout, 0, rank_idx, sm_idx, thread_idx); + + // Do TMA writes into the remote buffers + int num_tokens_per_warp = math::ceil_div(num_reduced_tokens, kNumSMs * kNumWarps); + const int token_start_idx = num_tokens_per_warp * global_warp_idx; + const int token_end_idx = min(token_start_idx + num_tokens_per_warp, num_reduced_tokens); + for (int i = token_start_idx; i < token_end_idx; ++ i) { + // The master slot index during dispatch + constexpr int kMetadataStride = 2 + kNumTopk; + const int src_token_idx = __ldg(src_metadata + i * kMetadataStride) % kNumMaxTokensPerRank; + const int src_rank_topk_idx = __ldg(src_metadata + i * kMetadataStride + 1); + const int src_rank_idx = src_rank_topk_idx / kNumTopk; + const int src_topk_idx = src_rank_topk_idx % kNumTopk; + + // Directly to the remote or via RDMA + const bool nvlink_bypass = gin.is_nvlink_accessible(src_rank_idx); + layout::TokenLayout master_token_buffer = [=]() { + // NVLink bypass + if (nvlink_bypass) { + auto token_buffer = recv_buffer.get_rank_buffer(kUseRankLayout ? rank_idx : src_topk_idx).get_token_buffer(src_token_idx); + token_buffer.set_base_ptr(gin.get_sym_ptr(token_buffer.get_base_ptr(), src_rank_idx)); + return token_buffer; + } + + // Use RDMA + return send_buffer.get_rank_buffer(src_rank_idx).get_token_buffer(src_token_idx); + }(); + + // Hidden requirements + EP_STATIC_ASSERT(kHidden % (32 * sizeof(int4) / sizeof(nv_bfloat16)) == 0, "Invalid hidden"); + using combine_vec_t = typename CombineVecTraits::vec_t; + constexpr int kHiddenVec = kHidden * sizeof(nv_bfloat16) / sizeof(combine_vec_t); + + // Read source indices for expand mode + int stored_topk_slot_idx = -1; + if constexpr (kUseExpandedLayout) { + if (lane_idx < kNumTopk) + stored_topk_slot_idx = __ldg(src_metadata + i * kMetadataStride + (2 + lane_idx)); + __syncwarp(); + } + + // 3 cases: + // - no expand + no reduce, or expand + no reduce + // - expand + reduce + // - expand + send all + auto reduce_valid_mask = ptx::gather(stored_topk_slot_idx >= 0); + auto no_local_reduce = not kUseExpandedLayout or (kAllowMultipleReduction and __popc(reduce_valid_mask) == 1); + if (no_local_reduce) { + int token_idx_in_tensor = i; + if constexpr (kUseExpandedLayout) + token_idx_in_tensor = ptx::exchange(stored_topk_slot_idx, ptx::get_master_lane_idx(reduce_valid_mask)); + + // No reduce + if (ptx::elect_one_sync()) { + const auto load_ptr = + math::advance_ptr(x, static_cast(token_idx_in_tensor) * kNumHiddenBytes); + ptx::tma_store_wait(); + ptx::tma_load_1d(tma_buffer.get_base_ptr(), load_ptr, mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + ptx::tma_store_1d(master_token_buffer.get_base_ptr(), tma_buffer.get_base_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } + __syncwarp(); + } else if constexpr (kAllowMultipleReduction) { + // Do local reduction + // Sort valid top-k indices to front + int topk_slot_idx[kNumTopk]; + compute_topk_slots( + topk_slot_idx, reduce_valid_mask, + [=](const int& idx) { + return ptx::exchange(stored_topk_slot_idx, idx); + } + ); + + // Reduce into shared memory + constexpr int kUnrollFactor = get_max_unroll_factor(); + combine_reduce( + lane_idx, topk_slot_idx, static_cast(tma_buffer.get_base_ptr()), + /* Get source base */ [=](const int& slot_idx) { + return math::advance_ptr( + x, slot_idx * static_cast(kNumHiddenBytes)); + }, + /* Wait buffer release */ [=]() { + ptx::tma_store_wait(); + __syncwarp(); + } + ); + ptx::tma_store_fence(); + __syncwarp(); + + // Issue TMA stores + if (ptx::elect_one_sync()) { + ptx::tma_store_1d(master_token_buffer.get_base_ptr(), tma_buffer.get_base_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } + __syncwarp(); + } else { + // No local reduction, send all data (expanded send) + #pragma unroll + for (int k = 0; k < kNumTopk; ++ k) { + const auto slot_idx = ptx::exchange(stored_topk_slot_idx, k); + if (slot_idx >= 0) { + const auto src_token_ptr = math::advance_ptr(x, slot_idx * static_cast(kNumHiddenBytes)); + const auto token_buffer = recv_buffer.get_rank_buffer(k).get_token_buffer(src_token_idx); + if (ptx::elect_one_sync()) { + // Load + ptx::tma_store_wait(); + ptx::tma_load_1d(tma_buffer.get_base_ptr(), src_token_ptr, mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + + if (nvlink_bypass) { + // Write into the same position + ptx::tma_store_1d(gin.get_sym_ptr(token_buffer.get_base_ptr(), src_rank_idx), + tma_buffer.get_base_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } else { + // Write to the RDMA send buffer + const auto send_token_buffer = + send_buffer.get_rank_buffer(src_rank_idx).get_token_buffer(src_token_idx * kNumTopk + k); + ptx::tma_store_1d(send_token_buffer.get_base_ptr(), tma_buffer.get_base_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + ptx::tma_store_wait(); + + // Issue RDMA + gin.put(token_buffer.get_base_ptr(), send_token_buffer.get_base_ptr(), + kNumHiddenBytes, src_rank_idx); + } + } + __syncwarp(); + } + } + } + + // Write topk weights + if (not kUseExpandedLayout and topk_weights != nullptr and lane_idx < kNumTopk) { + const float value = __ldg(topk_weights + (i * kNumTopk + lane_idx)); + master_token_buffer.get_topk_weights_ptr()[lane_idx] = value; + } + __syncwarp(); + + // Wait send buffer's TMA store and issue RDMA send + // NOTES: `kDoExpandedSend` mode has already issued + if (not kDoExpandedSend and not nvlink_bypass and ptx::elect_one_sync()) { + ptx::tma_store_wait(); + const auto dst_ptr = recv_buffer.get_rank_buffer(kUseRankLayout ? rank_idx : src_topk_idx) + .get_token_buffer(src_token_idx).get_base_ptr(); + gin.put(dst_ptr, master_token_buffer.get_base_ptr(), + master_token_buffer.get_num_bytes(), src_rank_idx); + } + } + + // Final barrier to ensure data arrival + comm::gpu_barrier( + gin, workspace_layout, 0, rank_idx, sm_idx, thread_idx); +} + +} // deep_ep::elastic diff --git a/deep_ep/include/deep_ep/impls/combine_reduce_epilogue.cuh b/deep_ep/include/deep_ep/impls/combine_reduce_epilogue.cuh new file mode 100644 index 000000000..b57415011 --- /dev/null +++ b/deep_ep/include/deep_ep/impls/combine_reduce_epilogue.cuh @@ -0,0 +1,145 @@ +#pragma once + +#include +#include +#include + +#include + + +namespace deep_ep::elastic { + +template (), + int kNumTokensInLayout = get_num_tokens_in_layout()> +__global__ void __launch_bounds__(kNumThreads, 1) +combine_reduce_epilogue_impl(nv_bfloat16* combined_x, + float* combined_topk_weights, + topk_idx_t* combined_topk_idx, + void* recv_buffer, + void* bias_0, void* bias_1, + const int num_combined_tokens, + const int scaleout_rank_idx, const int scaleup_rank_idx) { + constexpr int kNumExpertsPerScaleout = kNumExperts / kNumScaleoutRanks; + constexpr int kNumExpertsPerRank = kNumExperts / (kNumScaleupRanks * kNumScaleoutRanks); + EP_STATIC_ASSERT(kNumExperts % (kNumScaleupRanks * kNumScaleoutRanks) == 0, "Invalid number of experts or ranks"); + + // Utils + const auto sm_idx = static_cast(blockIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + const auto global_warp_idx = warp_idx * kNumSMs + sm_idx; // NOTES: Here we prioritize distributing tasks to different SMs to ensure that the last wave is evenly concentrated on each SM. + + // Load buffers from scale-out or scale-up ranks + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + const auto comm_token_layout = layout::TokenLayout(kNumHiddenBytes, 0, kNumTopk, false); + const auto comm_buffer = layout::BufferLayout( + comm_token_layout, kNumTokensInLayout, kNumMaxTokensPerRank, recv_buffer); + + // Store buffers + const auto output_token_layout = layout::TokenLayout(kNumHiddenBytes, 0, 0, false); + const auto output_buffer = layout::BufferLayout(output_token_layout, 1, num_combined_tokens, combined_x); + const auto tma_buffer = layout::BufferLayout(output_token_layout, kNumWarps, 1, smem) + .get_rank_buffer(warp_idx).get_token_buffer(0); + + // Bias layout + const auto bias_0_buffer = layout::BufferLayout(output_token_layout, 1, num_combined_tokens, bias_0); + const auto bias_1_buffer = layout::BufferLayout(output_token_layout, 1, num_combined_tokens, bias_1); + + // Will block until the main combine kernel has finished and all data are visible + // NOTES: PDL is used, please do not use `__ldg` + cudaGridDependencySynchronize(); + + // Read from buffers and do reduction + for (int token_idx = global_warp_idx; token_idx < num_combined_tokens; token_idx += kNumWarps * kNumSMs) { + // Preprocess all indices + int stored_dst_rank_idx = -1, stored_dst_expert_idx = -1; + EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k selections"); + if (lane_idx < kNumTopk) { + stored_dst_expert_idx = static_cast(combined_topk_idx[token_idx * kNumTopk + lane_idx]); + stored_dst_rank_idx = stored_dst_expert_idx >= 0 ? + stored_dst_expert_idx / (kNumScaleoutRanks == 1 ? kNumExpertsPerRank : kNumExpertsPerScaleout) : -1; + } + __syncwarp(); + + // Sort valid top-k indices to front + const auto [should_deduplicate, deduplicate_key] = [&]() -> std::pair { + if constexpr (kUseExpandedLayout and not kAllowMultipleReduction) { + // Activations are never reduced before + return {false, 0}; + } else if constexpr (kNumScaleoutRanks != 1 and not kUseExpandedLayout and not kAllowMultipleReduction) { + // Hybrid mode without expanded layout and multiple reduction. Should deduplicate on a per-rank basis + return {true, stored_dst_expert_idx >= 0 ? stored_dst_expert_idx / kNumExpertsPerRank : -1}; + } else { + // Should deduplicate on a per-rank (for non-hybrid mode) or a per-scale-rank (for hybrid mode) basis + return {true, stored_dst_rank_idx}; + } + }(); + auto reduce_valid_mask = should_deduplicate ? + ptx::gather(ptx::deduplicate(deduplicate_key, lane_idx) and stored_dst_rank_idx >= 0) : + ptx::gather(stored_dst_rank_idx >= 0); + int topk_slot_idx[kNumTokensInLayout]; + compute_topk_slots( + topk_slot_idx, reduce_valid_mask, + [=](const int& idx) { + return kUseRankLayout ? ptx::exchange(stored_dst_rank_idx, idx) : idx; + } + ); + + // Iterate over per-hidden-chunk stage + using combine_vec_t = typename CombineVecTraits::vec_t; + constexpr int kHiddenVec = kHidden * sizeof(nv_bfloat16) / sizeof(combine_vec_t); + constexpr int kUnrollFactor = get_max_unroll_factor(); + combine_reduce( + lane_idx, topk_slot_idx, static_cast(tma_buffer.get_base_ptr()), + /* Get source base */ [=](const int& slot_idx) { + return static_cast( + comm_buffer.get_rank_buffer(slot_idx).get_token_buffer(token_idx).get_base_ptr()); + }, + /* Wait buffer release */ [=]() { + ptx::tma_store_wait(); + __syncwarp(); + }, + /* Bias 0 */ bias_0 == nullptr ? + nullptr : static_cast(bias_0_buffer.get_token_buffer(token_idx).get_base_ptr()), + /* Bias 1 */ bias_1 == nullptr ? + nullptr : static_cast(bias_1_buffer.get_token_buffer(token_idx).get_base_ptr()) + ); + ptx::tma_store_fence(); + __syncwarp(); + + // Issue TMA copy + if (ptx::elect_one_sync()) { + ptx::tma_store_1d(output_buffer.get_token_buffer(token_idx).get_base_ptr(), + tma_buffer.get_base_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } + __syncwarp(); + + // Write top-k weights + if (combined_topk_weights != nullptr) { + const auto master_lane_idx = ptx::get_master_lane_idx(ptx::match(stored_dst_rank_idx)); + if (lane_idx < kNumTopk) { + float value = 0; + if (stored_dst_rank_idx >= 0) { + const auto dst_ptr = comm_buffer + .get_rank_buffer(kUseRankLayout ? stored_dst_rank_idx : master_lane_idx) + .get_token_buffer(token_idx).get_topk_weights_ptr() + lane_idx; + value = *dst_ptr; + } + combined_topk_weights[token_idx * kNumTopk + lane_idx] = value; + } + __syncwarp(); + } + } +} + +} // deep_ep::elastic diff --git a/deep_ep/include/deep_ep/impls/combine_utils.cuh b/deep_ep/include/deep_ep/impls/combine_utils.cuh new file mode 100644 index 000000000..116068718 --- /dev/null +++ b/deep_ep/include/deep_ep/impls/combine_utils.cuh @@ -0,0 +1,172 @@ +#pragma once + +#include + + +namespace deep_ep::elastic { + +template +constexpr bool use_rank_layout() { + if constexpr (not kAllowMultipleReduction) + return false; + return kNumRanks <= kNumTopk; +} + +template +constexpr int get_num_tokens_in_layout() { + return use_rank_layout() ? kNumRanks : kNumTopk; +} + +template +constexpr int get_max_unroll_factor() { + for (int i = kMaxUnrollFactor; i >= 1; -- i) + if (kLength % (kWarpSize * i) == 0) + return i; + throw std::logic_error("Invalid length, cannot find unrolling factor"); +} + +// Determine the vector type for combine loads/stores based on arch and hidden size alignment +template +struct CombineVecTraits { +#if defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) + // On SM100+, use longlong4_t (32 bytes) if hidden is aligned, otherwise fall back to int4 (16 bytes) + static constexpr bool kUseLonglong4 = (kHiddenBytes % sizeof(longlong4_t) == 0) and + ((kHiddenBytes / sizeof(longlong4_t)) % 32 == 0); + using vec_t = std::conditional_t; +#else + using vec_t = int4; +#endif +}; + +template +__device__ __forceinline__ +void compute_topk_slots(int (&topk_slot_idx)[kNumValidTopk], uint32_t mask, + const fetch_func_t& fetch_func) { + #pragma unroll + for (int k = 0; k < kNumValidTopk; ++ k) { + const int lowest_idx = __ffs(mask) - 1; + // Here we perform the exchange unconditionally to avoid `BRA.DIV` + const auto fetched = fetch_func(lowest_idx); + mask &= mask - 1; + topk_slot_idx[k] = lowest_idx >= 0 ? fetched : -1; + } +} + +template +__device__ __forceinline__ +void combine_reduce(const int& lane_idx, int (&topk_slot_idx)[kNumValidTopk], + vec_t* dst_buffer_ptr, + const get_src_buffer_ptr_func_t& get_src_buffer_ptr_func, + const wait_buffer_func_t& wait_buffer_func, + vec_t* bias_0 = nullptr, vec_t* bias_1 = nullptr) { + constexpr int kNumElemsPerVec = sizeof(vec_t) / sizeof(nv_bfloat16); + EP_STATIC_ASSERT(kNumElemsPerVec % 2 == 0, "Invalid number of elements"); + EP_STATIC_ASSERT(kHiddenVec % (kUnrollFactor * 32) == 0, "Invalid unrolling"); + + // We use BF16 add as much as possible, as casting is slow + const bool enable_hadd_bypass = + (bias_0 == nullptr and bias_1 == nullptr) and + (kNumValidTopk <= 2 or topk_slot_idx[2] < 0); + EP_STATIC_ASSERT(kNumValidTopk > 0, "Invalid top-k"); + + if (enable_hadd_bypass) { + #pragma unroll 1 + for (int i = 0; i < kHiddenVec / (kUnrollFactor * 32); ++ i) { + // Read values 0 + const auto slot_0 = topk_slot_idx[0]; + const auto src_base_ptr_0 = get_src_buffer_ptr_func(slot_0); + vec_t values_0[kUnrollFactor] = {}; + #pragma unroll + for (int j = 0; j < kUnrollFactor; ++ j) { + values_0[j] = ptx::ldg_with_gez_pred( + src_base_ptr_0 + (i * (kUnrollFactor * 32) + j * 32 + lane_idx), slot_0); + } + + // Read values 1 + vec_t values_1[kUnrollFactor] = {}; + const auto slot_1 = kNumValidTopk == 1 ? -1 : topk_slot_idx[1]; + const auto src_base_ptr_1 = get_src_buffer_ptr_func(slot_1); + #pragma unroll + for (int j = 0; j < kUnrollFactor; ++ j) { + values_1[j] = ptx::ldg_with_gez_pred( + src_base_ptr_1 + (i * (kUnrollFactor * 32) + j * 32 + lane_idx), slot_1); + } + + // Wait buffer releases for the first write + if (i == 0) + wait_buffer_func(); + + // Reduce into shared memory + const auto bf162_view_0 = reinterpret_cast(values_0); + const auto bf162_view_1 = reinterpret_cast(values_1); + #pragma unroll + for (int j = 0; j < kUnrollFactor; ++ j) { + #pragma unroll + for (int l = 0; l < kNumElemsPerVec / 2; ++ l) + bf162_view_0[j * (kNumElemsPerVec / 2) + l] += bf162_view_1[j * (kNumElemsPerVec / 2) + l]; + dst_buffer_ptr[i * (kUnrollFactor * 32) + j * 32 + lane_idx] = values_0[j]; + } + } + } else { + #pragma unroll 1 + for (int i = 0; i < kHiddenVec / (kUnrollFactor * 32); ++ i) { + // Add bias + float2 reduced[kUnrollFactor * kNumElemsPerVec / 2] = {}; + const auto add_bias = [&](const vec_t* base_ptr) { + // Read + vec_t values[kUnrollFactor]; + #pragma unroll + for (int j = 0; j < kUnrollFactor; ++ j) + values[j] = ptx::ldg(base_ptr + i * (kUnrollFactor * 32) + j * 32 + lane_idx); + + // Reduce + const auto bf162_view = reinterpret_cast(values); + #pragma unroll + for (int j = 0; j < kUnrollFactor * kNumElemsPerVec / 2; ++ j) + ptx::accumulate(reduced[j], bf162_view[j]); + }; + bias_0 != nullptr ? add_bias(bias_0) : void(); + bias_1 != nullptr ? add_bias(bias_1) : void(); + + #pragma unroll + for (int k = 0; k < kNumValidTopk; ++ k) { + // We have a limitation on `k` to reduce the branch instruction count + if (k >= kNumExpectedTopk and topk_slot_idx[k] < 0) + break; + + // Read values + const auto src_base_ptr = get_src_buffer_ptr_func(topk_slot_idx[k]); + vec_t values[kUnrollFactor] = {}; + #pragma unroll + for (int j = 0; j < kUnrollFactor; ++ j) { + values[j] = ptx::ldg_with_gez_pred( + src_base_ptr + (i * (kUnrollFactor * 32) + j * 32 + lane_idx), topk_slot_idx[k]); + } + + // Reduce + const auto bf162_view = reinterpret_cast(values); + #pragma unroll + for (int j = 0; j < kUnrollFactor * kNumElemsPerVec / 2; ++ j) + ptx::accumulate(reduced[j], bf162_view[j]); + } + + // Wait buffer releases for the first write + if (i == 0) + wait_buffer_func(); + + // Cast into shared memory + #pragma unroll + for (int j = 0; j < kUnrollFactor; ++ j) { + vec_t casted_value; + auto bf162_view = reinterpret_cast(&casted_value); + #pragma unroll + for (int l = 0; l < kNumElemsPerVec / 2; ++ l) + bf162_view[l] = __float22bfloat162_rn(reduced[j * (kNumElemsPerVec / 2) + l]); + dst_buffer_ptr[i * (kUnrollFactor * 32) + j * 32 + lane_idx] = casted_value; + } + } + } +} + +} // namespace deep_ep::elastic diff --git a/deep_ep/include/deep_ep/impls/dispatch.cuh b/deep_ep/include/deep_ep/impls/dispatch.cuh new file mode 100644 index 000000000..cf21a18b0 --- /dev/null +++ b/deep_ep/include/deep_ep/impls/dispatch.cuh @@ -0,0 +1,405 @@ +#pragma once + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + + +namespace deep_ep::elastic { + +template > +__global__ void __launch_bounds__(kNumThreads, 1) +dispatch_impl( + 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* dst_buffer_slot_idx, + const int num_tokens, + const int sf_token_stride, const int sf_hidden_stride, + const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window, void* buffer, + void* workspace, void* mapped_host_workspace, + const int rank_idx +) { + constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks; + EP_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks"); + EP_STATIC_ASSERT(kNumNotifyWarps % 4 == 0, "Invalid warpgroup size"); + + // Utils + const auto sm_idx = static_cast(blockIdx.x), thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + + // Workspaces + const auto workspace_layout = layout::WorkspaceLayout(workspace, 1, kNumRanks, kNumExperts); + const auto host_workspace_layout = layout::WorkspaceLayout(mapped_host_workspace, 1, kNumRanks, kNumExperts); + + // The kernel uses a fixed space of dynamic shared memory (no static shared memory) + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + constexpr int kNumSmemBytesForNotify = kNumNotifyThreads > 0 ? + math::constexpr_align(kNumRanks + kNumExperts, kNumNotifyThreads) * sizeof(int) : 0; + EP_STATIC_ASSERT(kNumSmemBytesForNotify % ptx::kNumTMAAlignBytes == 0, "Invalid TMA alignment"); + + // Named barrier indices + constexpr int kNotifyBarrierIndex = 1; + + // Gin handle + // We treat each warp as a "channel" + const auto [qp_idx, sharing_mode] = comm::get_qp_mode 0)>( + sm_idx, warp_idx - kNumNotifyWarps, warp_idx < kNumNotifyWarps); + const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, sharing_mode); + + // Barrier without TMA store flush, without prologue grid sync + comm::gpu_barrier( + gin, workspace_layout, 0, rank_idx, sm_idx, thread_idx); + + // Different warp roles + if (warp_idx < kNumNotifyWarps) { + // Assign shared memory + constexpr int kNumAlignedElems = kNumSmemBytesForNotify / sizeof(int); + const auto rank_expert_count = math::advance_ptr(smem, 0); + + // Clean initial counts + // NOTES: if you want to change the order of different warp roles, please take care of the `thread_idx` + int *rank_count = rank_expert_count, *expert_count = rank_expert_count + kNumRanks; + #pragma unroll + for (int i = 0; i < kNumAlignedElems / kNumNotifyThreads; ++ i) + rank_expert_count[i * kNumNotifyThreads + thread_idx] = 0; + ptx::named_barrier(kNotifyBarrierIndex); + + // Atomic add on shared memory + EP_STATIC_ASSERT(kNumTopk <= 32, "Insufficient lanes"); + const auto global_warp_idx = warp_idx * kNumSMs + sm_idx; + for (int i = global_warp_idx; i < num_tokens; i += kNumNotifyWarps * kNumSMs) { + // Expert choice can not be redundant + // NOTES: no assertions here as they are expensive + const auto dst_expert_idx = lane_idx < kNumTopk ? + static_cast(__ldg(topk_idx + i * kNumTopk + lane_idx)) : -1; + if (dst_expert_idx >= 0) + atomicAdd_block(expert_count + dst_expert_idx, 1); + + // Rank choice should do deduplication here + const auto dst_rank_idx = dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerRank : -1; + if (ptx::deduplicate(dst_rank_idx, lane_idx) and dst_rank_idx >= 0) + atomicAdd_block(rank_count + dst_rank_idx, 1); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // Do full-grid reduction + #pragma unroll + for (int i = thread_idx; i < kNumRanks + kNumExperts; i += kNumNotifyThreads) { + const int64_t counter = (1ll << 32ll) | rank_expert_count[i]; + ptx::red_add(workspace_layout.get_notify_reduction_workspace_ptr() + i, counter); + } + + // Do the remaining work by SM 0 + if (sm_idx == 0) { + // Reduce all SM's count + // Wait all SMs' arrival + #pragma unroll + for (int i = thread_idx; i < kNumRanks + kNumExperts; i += kNumNotifyThreads) { + comm::timeout_while(true, [=](const bool& is_last_check) { + const auto status = ptx::ld_volatile(workspace_layout.get_notify_reduction_workspace_ptr() + i); + if ((status >> 32) == kNumSMs) { + // Write into shared memory + // Write into send buffer if with RDMA + const auto encoded = + math::encode_decode_positive(static_cast(status & 0xffffffffll)); + rank_expert_count[i] = encoded; + if constexpr (not kIsScaleupNVLink) + workspace_layout.get_scaleup_rank_expert_count_ptr()[i] = encoded; + + // Clean for the next usage + workspace_layout.get_notify_reduction_workspace_ptr()[i] = 0; + return true; + } + + if (is_last_check) { + printf("DeepEP notify (GPU reduction) timeout, rank: %d/%d, " + "thread: %d, status: %d | %d, expected: %d\n", + rank_idx, kNumRanks, thread_idx, + static_cast(status >> 32), static_cast(status & 0xffffffff), kNumSMs); + } + return false; + }); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // TODO: for further optimization, we can fuse rank and expert counters + // Issue scaleup rank count writes to peers + for (int i = thread_idx; i < kNumRanks; i += kNumNotifyThreads) { + // Rank counters + const auto dst_rank_counter = + workspace_layout.get_scaleup_rank_count_ptr() + rank_idx; + gin.put_value(dst_rank_counter, static_cast(rank_count[i]), i, + ncclGinOptFlagsAggregateRequests); + } + __syncwarp(); + + // Issue scaleup expert count writes to peers + if constexpr (kIsScaleupNVLink) { + // NVLink per-element copy + // We don't use TMA as the dtype of shared memory and global is different + for (int i = thread_idx; i < kNumExperts; i += kNumNotifyThreads) { + const auto idx = kNumExpertsPerRank * rank_idx + (i % kNumExpertsPerRank); + gin.put_value( + workspace_layout.get_scaleup_expert_count_ptr() + idx, + static_cast(expert_count[i]), i / kNumExpertsPerRank); + } + } else { + // RDMA bulk copy + for (int i = thread_idx; i < kNumRanks; i += kNumNotifyThreads) { + const auto src_ptr = workspace_layout.get_scaleup_expert_count_ptr() + kNumExpertsPerRank * i; + const auto dst_ptr = workspace_layout.get_scaleup_expert_count_ptr() + kNumExpertsPerRank * rank_idx; + gin.put(dst_ptr, src_ptr, kNumExpertsPerRank * sizeof(int64_t), i); + } + } + + // This is necessary, as the waited results will rewrite the shared memory + ptx::named_barrier(kNotifyBarrierIndex); + + // Wait for rank and expert count + const auto start_clock = clock64(); + for (int i = thread_idx; i < kNumRanks + kNumExperts; i += kNumNotifyThreads) { + comm::timeout_while([=](const bool& is_last_check) { + // NOTES: the global memory type has 64 bits + const auto count = static_cast( + ptx::ld_volatile(workspace_layout.get_scaleup_rank_expert_count_ptr() + i)); + const auto decoded = math::encode_decode_positive(count); + if (math::is_decoded_positive_ready(decoded)) { + workspace_layout.get_scaleup_rank_expert_count_ptr()[i] = 0; + rank_expert_count[i] = decoded; + return true; + } + + if (is_last_check) + printf("DeepEP notify timeout, rank: %d, thread: %d, count: %d\n", rank_idx, i, decoded); + return false; + }, start_clock); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // Reduce expert count and add stats + for (int i = thread_idx; i < kNumExpertsPerRank; i += kNumNotifyThreads) { + int sum = 0; + #pragma unroll + for (int j = 0; j < kNumRanks; ++ j) + sum += expert_count[j * kNumExpertsPerRank + i]; + expert_count[i] = math::align(sum, kExpertAlignment); + + // Update statistics counters + if (cumulative_local_expert_recv_stats != nullptr) + atomicAdd(cumulative_local_expert_recv_stats + i, sum); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // Write host workspace + if constexpr (kDoCPUSync) { + for (int i = thread_idx; i < kNumRanks + kNumExpertsPerRank; i += kNumNotifyThreads) { + host_workspace_layout.get_scaleup_rank_expert_count_ptr()[i] = + math::encode_decode_positive(rank_expert_count[i]); + } + __syncwarp(); + } + + // Do prefix sum by the warps + // NOTES: we may have fast implementation with `cub::BlockScan`, but it is too heavy to use + const auto do_psum = [=](const int* count, int* out, const int n, const int is_exclusive) { + int psum = 0; + #pragma unroll + for (int i = 0; i < math::ceil_div(n + is_exclusive, 32); ++ i) { + const auto idx = i * 32 + lane_idx; + const auto mem_idx = idx - is_exclusive; + const auto value = (0 <= mem_idx and mem_idx < n) ? count[mem_idx] : 0; + const auto sum = psum + ptx::warp_inclusive_sum(value, lane_idx); + + // Store into global memory + if (idx < n + is_exclusive) + out[idx] = sum; + + // Update `psum` by using the last lane's value + psum = ptx::exchange(sum, 31); + } + }; + if (warp_idx == 0) { + // Inclusive prefix sum + do_psum(rank_count, psum_num_recv_tokens_per_scaleup_rank, kNumRanks, 0); + } else if (warp_idx == 1) { + // Exclusive prefix sum for later expanding + do_psum(expert_count, psum_num_recv_tokens_per_expert, kNumExpertsPerRank, 1); + } + } + } else { + const int dispatch_warp_idx = warp_idx - kNumNotifyWarps; + + // Buffer layouts + const auto token_layout = layout::TokenLayout(kNumHiddenBytes, kNumSFPacks * sizeof(sf_pack_t), kNumTopk, true); + const auto tma_buffer = layout::BufferLayout(token_layout, kNumDispatchWarps, 1, + math::advance_ptr(smem, kNumSmemBytesForNotify)).get_rank_buffer(dispatch_warp_idx).get_token_buffer(0); + auto recv_buffer = layout::BufferLayout(token_layout, kNumRanks, kNumMaxTokensPerRank, buffer); + auto send_buffer = layout::BufferLayout(token_layout, 1, kNumMaxTokensPerRank, recv_buffer.get_buffer_end_ptr()); + recv_buffer = recv_buffer.get_rank_buffer(rank_idx); + + // Init TMA + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (ptx::elect_one_sync()) + ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Iterate all tokens + const auto token_start = dispatch_warp_idx * kNumSMs + sm_idx; + const auto token_stride = kNumDispatchWarps * kNumSMs; + for (int token_idx = token_start; token_idx < num_tokens; token_idx += token_stride) { + const auto token_i64_idx = static_cast(token_idx); + + // Wait TMA store arrivals + ptx::tma_store_wait(); + __syncwarp(); + + // Issue data TMA + if (ptx::elect_one_sync()) { + ptx::tma_load_1d(tma_buffer.get_hidden_ptr(), math::advance_ptr(x, token_i64_idx * kNumHiddenBytes), + mbarrier_ptr, kNumHiddenBytes); + } + __syncwarp(); + + // Issue SF TMA or cp.async + if constexpr (kNumSFPacks > 0) { + EP_STATIC_ASSERT(sizeof(sf_pack_t) % 4 == 0, "Unaligned SF element type"); + const auto gmem_src_ptr = math::advance_ptr(sf, token_i64_idx * sf_token_stride * sizeof(sf_pack_t)); + const auto smem_dst_ptr = tma_buffer.get_sf_ptr(); + + constexpr auto kNumFullIters = kNumSFPacks / 32; + #pragma unroll + for (int k = 0; k < kNumFullIters; ++ k) { + ptx::cp_async_ca(gmem_src_ptr + (k * 32 + lane_idx) * sf_hidden_stride, + smem_dst_ptr + k * 32 + lane_idx); + } + if (kNumFullIters * 32 + lane_idx < kNumSFPacks) { + ptx::cp_async_ca(gmem_src_ptr + (kNumFullIters * 32 + lane_idx) * sf_hidden_stride, + smem_dst_ptr + kNumFullIters * 32 + lane_idx); + } + ptx::cp_async_mbarrier_arrive(mbarrier_ptr); + __syncwarp(); + } + + // Load top-k indices and weights + EP_STATIC_ASSERT(kNumTopk <= 32, "Insufficient lanes for loading top-k indices"); + int stored_dst_rank_idx = -1; + if (lane_idx < kNumTopk) { + const auto uncasted_dst_expert_idx = __ldg(topk_idx + token_idx * kNumTopk + lane_idx); + const auto dst_expert_idx = static_cast(uncasted_dst_expert_idx); + stored_dst_rank_idx = dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerRank : -1; + tma_buffer.get_topk_idx_ptr()[lane_idx] = dst_expert_idx; + if (topk_weights != nullptr) + tma_buffer.get_topk_weights_ptr()[lane_idx] = __ldg(topk_weights + token_idx * kNumTopk + lane_idx); + if (copied_topk_idx != nullptr) + copied_topk_idx[token_idx * kNumTopk + lane_idx] = uncasted_dst_expert_idx; + } + __syncwarp(); + + // Add source metadata (rank index and token index) + // Please ensure no TMA buffer shared memory writes after this part + if (ptx::elect_one_sync()) + *tma_buffer.get_src_token_global_idx_ptr() = rank_idx * kNumMaxTokensPerRank + token_idx; + ptx::tma_store_fence(); + __syncwarp(); + + // Deduplicate ranks and assign slots + int stored_dst_slot_idx = -1; + if constexpr (kReuseSlotIndices) { + if (lane_idx < kNumTopk) + stored_dst_slot_idx = __ldg(dst_buffer_slot_idx + token_idx * kNumTopk + lane_idx); + stored_dst_slot_idx = stored_dst_slot_idx >= 0 ? + (stored_dst_slot_idx - rank_idx * kNumMaxTokensPerRank) : -1; + } else { + if (ptx::deduplicate(stored_dst_rank_idx, lane_idx) and stored_dst_rank_idx >= 0) + stored_dst_slot_idx = atomicAdd(workspace_layout.get_scaleup_atomic_sender_counter() + stored_dst_rank_idx, 1); + if (lane_idx < kNumTopk) { + const auto value = stored_dst_slot_idx >= 0 ? + rank_idx * kNumMaxTokensPerRank + stored_dst_slot_idx : -1; + dst_buffer_slot_idx[token_idx * kNumTopk + lane_idx] = value; + } + } + __syncwarp(); + + // Wait TMA load arrival + // NOTES: this arrive must be after the `ptx::cp_async_mbarrier_arrive` + if (ptx::elect_one_sync()) { + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + } + __syncwarp(); + + // TMA store to send buffer + auto send_buffer_ptr = send_buffer.get_token_buffer(token_idx).get_base_ptr(); + if constexpr (not kIsScaleupNVLink) { + if (ptx::elect_one_sync()) + ptx::tma_store_1d(send_buffer_ptr, tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes()); + ptx::tma_store_commit(); + __syncwarp(); + } + + // Issue TMA NVLink stores + EP_STATIC_ASSERT(kNumTopk <= 32, "Invalid top-k selection"); + const auto dst_ptr = stored_dst_slot_idx >= 0 ? + gin.get_sym_ptr(recv_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(), stored_dst_rank_idx) : + nullptr; + if (dst_ptr != nullptr) + ptx::tma_store_1d(dst_ptr, tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes()); + ptx::tma_store_commit(); + __syncwarp(); + + // Issue RDMA put + if constexpr (not kIsScaleupNVLink) { + // Wait the send buffer store to arrive + ptx::tma_store_wait<1>(); + __syncwarp(); + + // NOTES: we should skip the NVLink accessible ranks + if (stored_dst_slot_idx >= 0 and dst_ptr == nullptr) { + gin.put(recv_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(), + send_buffer_ptr, tma_buffer.get_num_bytes(), stored_dst_rank_idx); + } + __syncwarp(); + } + } + } + + // Barrier to ensure data arrival + comm::gpu_barrier( + gin, workspace_layout, 0, rank_idx, sm_idx, thread_idx); + + // Trigger the copy epilogue kernel + cudaTriggerProgrammaticLaunchCompletion(); + + // Clean atomic counters + EP_STATIC_ASSERT(kNumRanks <= kNumThreads, "Insufficient threads"); + if (not kReuseSlotIndices and sm_idx == 0 and thread_idx < kNumRanks) + workspace_layout.get_scaleup_atomic_sender_counter()[thread_idx] = 0; +} + +} // namespace deep_ep::elastic diff --git a/deep_ep/include/deep_ep/impls/dispatch_copy_epilogue.cuh b/deep_ep/include/deep_ep/impls/dispatch_copy_epilogue.cuh new file mode 100644 index 000000000..e6b5bc45c --- /dev/null +++ b/deep_ep/include/deep_ep/impls/dispatch_copy_epilogue.cuh @@ -0,0 +1,226 @@ +#pragma once + +#include +#include +#include +#include + + +namespace deep_ep::elastic { + +template 1 and not kCachedMode)> +__global__ void __launch_bounds__(kNumThreads, 1) +dispatch_copy_epilogue_impl(void* buffer, void* workspace, + int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, + void* recv_x, sf_pack_t* recv_sf, + topk_idx_t* recv_topk_idx, float* recv_topk_weights, + int* recv_src_metadata, + int* channel_linked_list, + int num_recv_tokens, + const int recv_sf_token_stride, const int recv_sf_hidden_stride, + const int scaleout_rank_idx, const int scaleup_rank_idx) { + // Utils + const auto sm_idx = static_cast(blockIdx.x), thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + const auto global_warp_idx = warp_idx * kNumSMs + sm_idx; + + // For top-k index transformations + constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks; + const auto rank_idx = scaleout_rank_idx * kNumScaleupRanks + scaleup_rank_idx; + const auto expert_start_idx = kNumExpertsPerRank * rank_idx, expert_end_idx = kNumExpertsPerRank * (rank_idx + 1); + + // Buffer layouts + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + const auto token_layout = layout::TokenLayout(kNumHiddenBytes, kNumSFPacks * sizeof(sf_pack_t), kNumTopk, true); + const auto tma_buffer = layout::BufferLayout(token_layout, kNumWarps, 1, smem) + .get_rank_buffer(warp_idx).get_token_buffer(0); + const auto scaleup_buffer = layout::BufferLayout(token_layout, kNumScaleupRanks, kNumScaleoutRanks * kNumMaxTokensPerRank, buffer); + + // Init TMA + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (ptx::elect_one_sync()) + ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Will block until the main dispatch kernel has finished and all data are visible + // NOTES: PDL is used, please do not use `__ldg` + cudaGridDependencySynchronize(); + + // For no CPU sync case, the number of received tokens should be read from the GPU tensor + if (num_recv_tokens == kNumMaxTokensPerRank * kNumRanks) + num_recv_tokens = psum_num_recv_tokens_per_scaleup_rank[kNumScaleupRanks - 1]; + + // Current rank indices should be maintained + int current_rank_idx = -1, stored_psum_num_recv_tokens; + int current_rank_start = 0, current_rank_end = 0; + #pragma unroll + for (int i = global_warp_idx; i < num_recv_tokens; i += kNumWarps * kNumSMs) { + // Calculate token index in the buffer + while (i >= current_rank_end) { + current_rank_idx += 1; + EP_DEVICE_ASSERT(current_rank_idx < kNumScaleupRanks); + const auto stored_lane_idx = current_rank_idx % 32; + if (stored_lane_idx == 0 and current_rank_idx + lane_idx < kNumScaleupRanks) + stored_psum_num_recv_tokens = psum_num_recv_tokens_per_scaleup_rank[current_rank_idx + lane_idx]; + current_rank_start = current_rank_end; + current_rank_end = ptx::exchange(stored_psum_num_recv_tokens, stored_lane_idx); + } + const auto buffer_token = scaleup_buffer.get_rank_buffer(current_rank_idx).get_token_buffer(i - current_rank_start); + + // Wait buffer releases + ptx::tma_store_wait(); + __syncwarp(); + + // Issue TMA loads + // Including all stuffs: data, SF, top-k metadata + if (ptx::elect_one_sync()) { + ptx::tma_load_1d(tma_buffer.get_base_ptr(), buffer_token.get_base_ptr(), + mbarrier_ptr, tma_buffer.get_num_bytes()); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, tma_buffer.get_num_bytes()); + } + __syncwarp(); + + // Load target expert indices separately to tolerate TMA load latency + EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k selections"); + int dst_expert_idx = -1; + if (lane_idx < kNumTopk) + dst_expert_idx = buffer_token.get_topk_idx_ptr()[lane_idx]; + __syncwarp(); + + // Validate target expert indices and store for non-expand mode + const auto in_range = expert_start_idx <= dst_expert_idx and dst_expert_idx < expert_end_idx; + const auto master_src_topk_idx = ptx::get_master_lane_idx(ptx::gather(in_range)); + dst_expert_idx = in_range ? dst_expert_idx - expert_start_idx : -1; + EP_DEVICE_ASSERT(ptx::deduplicate(dst_expert_idx, lane_idx) or dst_expert_idx == -1); + if (not kDoExpand and lane_idx < kNumTopk) + recv_topk_idx[i * kNumTopk + lane_idx] = static_cast(dst_expert_idx); + __syncwarp(); + + // Calculate target indices in the tensor + int dst_tensor_idx = -1; + if (not kDoExpand and ptx::elect_one_sync()) { + dst_tensor_idx = i; + } else if (kDoExpand and dst_expert_idx >= 0) { + dst_tensor_idx = atomicAdd(psum_num_recv_tokens_per_expert + dst_expert_idx, 1); + } + __syncwarp(); + + // Wait for TMA arrival + if (ptx::elect_one_sync()) + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + __syncwarp(); + + // Maintain linked list + if constexpr (kDoCreateLinkedList) { + if (ptx::elect_one_sync()) + channel_linked_list[tma_buffer.get_linked_list_idx_ptr()[master_src_topk_idx]] = i; + __syncwarp(); + } + + // Issue TMA stores for data + if (kDoExpand ? (dst_tensor_idx >= 0) : ptx::elect_one_sync()) { + ptx::tma_store_1d(math::advance_ptr(recv_x, static_cast(dst_tensor_idx) * kNumHiddenBytes), + tma_buffer.get_hidden_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } + __syncwarp(); + + // Store SF + if constexpr (kNumSFPacks > 0) { + constexpr auto kNumFullIters = kNumSFPacks / 32; + const bool do_last_iter = (kNumSFPacks % 32 != 0) and (kNumFullIters * 32 + lane_idx < kNumSFPacks); + EP_STATIC_ASSERT(sizeof(sf_pack_t) % 4 == 0, "Unaligned SF element type"); + + // Load into registers + const auto smem_src_ptr = tma_buffer.get_sf_ptr(); + sf_pack_t reg_src[kNumFullIters + 1]; + #pragma unroll + for (int k = 0; k < kNumFullIters; ++ k) + reg_src[k] = smem_src_ptr[k * 32 + lane_idx]; + if (do_last_iter) + reg_src[kNumFullIters] = smem_src_ptr[kNumFullIters * 32 + lane_idx]; + + // Prepare strides + const auto recv_sf_token_stride_i64 = static_cast(recv_sf_token_stride); + const auto recv_sf_hidden_stride_i64 = static_cast(recv_sf_hidden_stride); + + // Iterate through all valid indices and store into output buffer + auto mask = kDoExpand ? ptx::gather(dst_tensor_idx >= 0) : 1; + while (mask) { + const int valid_lane_idx = __ffs(mask) - 1; + const auto gmem_dst = math::advance_ptr(recv_sf, + ptx::exchange(dst_tensor_idx, valid_lane_idx) * (recv_sf_token_stride_i64 * sizeof(sf_pack_t))); + #pragma unroll + for (int k = 0; k < kNumFullIters; ++ k) + gmem_dst[(k * 32 + lane_idx) * recv_sf_hidden_stride_i64] = reg_src[k]; + if (do_last_iter) + gmem_dst[(kNumFullIters * 32 + lane_idx) * recv_sf_hidden_stride_i64] = reg_src[kNumFullIters]; + mask ^= 1 << valid_lane_idx; + } + } + + // Store the top-k weights + if (kDoExpand and recv_topk_weights != nullptr and dst_tensor_idx >= 0) { + recv_topk_weights[dst_tensor_idx] = tma_buffer.get_topk_weights_ptr()[lane_idx]; + } else if (not kDoExpand and recv_topk_weights != nullptr and lane_idx < kNumTopk) { + // For backward, weights are optional + recv_topk_weights[i * kNumTopk + lane_idx] = tma_buffer.get_topk_weights_ptr()[lane_idx]; + } + __syncwarp(); + + // Write source token index + // And: + // - Non-hybrid mode: the source scaleup peer rank index and master top-k lane index + // - Hybrid mode: the slot index and master top-k lane index + constexpr int kMetadataStride = 2 + kNumTopk; + if (ptx::elect_one_sync()) { + recv_src_metadata[i * kMetadataStride + 0] = *tma_buffer.get_src_token_global_idx_ptr(); + if constexpr (kNumScaleoutRanks == 1) { + recv_src_metadata[i * kMetadataStride + 1] = current_rank_idx * kNumTopk + master_src_topk_idx; + } else { + recv_src_metadata[i * kMetadataStride + 1] = (i - current_rank_start) * kNumTopk + master_src_topk_idx; + } + } + __syncwarp(); + + // Write reduction source indices + if (kDoExpand and lane_idx < kNumTopk) + recv_src_metadata[i * kMetadataStride + 2 + lane_idx] = dst_tensor_idx; + __syncwarp(); + } + + // Maintain linked list's ending + // Or you can understand it as writing the tail at once + if constexpr (kDoCreateLinkedList) { + constexpr int kNumScaleupRanksPerLane = math::constexpr_ceil_div(kNumScaleupRanks, 32); + const auto workspace_layout = layout::WorkspaceLayout(workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts); + for (int i = global_warp_idx; i < kNumChannels; i += kNumSMs * kNumWarps) { + #pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) { + if (const auto k = j * 32 + lane_idx; j < (kNumScaleupRanksPerLane - 1) or k < kNumScaleupRanks) { + channel_linked_list[ + *workspace_layout.get_channel_scaleup_tail_ptr(i, k) + ] = -1; + + // Clean for combine usages + *workspace_layout.get_channel_scaleup_tail_ptr(i, k) = 0; + } + } + __syncwarp(); + } + } +} + +} // namespace deep_ep::elastic diff --git a/deep_ep/include/deep_ep/impls/dispatch_deterministic_prologue.cuh b/deep_ep/include/deep_ep/impls/dispatch_deterministic_prologue.cuh new file mode 100644 index 000000000..8e9dd79b3 --- /dev/null +++ b/deep_ep/include/deep_ep/impls/dispatch_deterministic_prologue.cuh @@ -0,0 +1,143 @@ +#pragma once + +#include + +#include +#include +#include + + +namespace deep_ep::elastic { + +// TODO: support scale-out +template +__global__ void __launch_bounds__(kNumThreads, 1) +dispatch_deterministic_prologue_impl( + topk_idx_t* topk_idx, + int* rank_count_buffer, + int* dst_buffer_slot_idx, + const int num_tokens, + const int scaleup_rank_idx +) { + constexpr int kNumExpertsPerRank = kNumExperts / kNumScaleupRanks; + EP_STATIC_ASSERT(kNumExperts % kNumScaleupRanks == 0, "Invalid number of experts or ranks"); + + // Utils + const auto sm_idx = static_cast(blockIdx.x), thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + const auto global_warp_idx = sm_idx * kNumWarps + warp_idx; + + // Token region the current warp is responsible for + const auto num_tokens_per_warp = math::ceil_div(num_tokens, kNumSMs * kNumWarps); + const auto start_token_idx = global_warp_idx * num_tokens_per_warp; + const auto end_token_idx = min(start_token_idx + num_tokens_per_warp, num_tokens); + + // Group configs + // NOTES: Group refers to the tokens that each warp handles concurrently + constexpr int kNumTokensPerGroup = 32 / kNumTopk; + const auto token_idx_offset = lane_idx / kNumTopk; + const unsigned token_mask = ((1u << kNumTopk) - 1) << (token_idx_offset * kNumTopk); + EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k"); + + // Shared memory for reduction + // NOTES: Each warp owns separate shared memory region for separate sum. + extern __shared__ int8_t smem[]; + const auto rank_count_global_psum = math::advance_ptr(smem, 0); + const auto rank_count_warp_sum = math::advance_ptr(rank_count_global_psum, (kNumScaleupRanks + warp_idx * kNumScaleupRanks) * sizeof(int)); + const auto rank_count_warp_psum = math::advance_ptr(rank_count_warp_sum, kNumWarps * kNumScaleupRanks * sizeof(int)); + + // Initialize to zero before reduce + for (int i = thread_idx; i < kNumScaleupRanks * (1 + 2 * kNumWarps); i += kNumThreads) + reinterpret_cast(smem)[i] = 0; + __syncthreads(); + + // Util functions + const auto map_expert_to_rank_idx = [&](const int& expert_idx) { + return expert_idx >= 0 ? expert_idx / kNumExpertsPerRank : -1; + }; + const auto is_unique = [&](const int& rank_idx) { + return ((ptx::match(rank_idx) & token_mask) >> lane_idx) == 1; + }; + const auto count_ones_before = [&](const unsigned& mask, const int& bit_idx) { + return __popc(mask & ((1u << bit_idx) - 1)); + }; + const auto get_other_rank_count_warp_sum = [&](const int& other_warp_idx) { + // NOTES: pass negative num_bytes to advance pointer + return math::advance_ptr(rank_count_warp_sum, (other_warp_idx - warp_idx) * kNumScaleupRanks * sizeof(int)); + }; + + // Each warp scan the tokens separately + for (int i = start_token_idx; i < end_token_idx; i += kNumTokensPerGroup) { + const auto token_idx = i + token_idx_offset; + const auto is_active_thread = lane_idx < kNumTopk * kNumTokensPerGroup and token_idx < end_token_idx; + const int expert_idx = is_active_thread ? static_cast(__ldg(topk_idx + i * kNumTopk + lane_idx)) : -1; + const auto rank_idx = map_expert_to_rank_idx(expert_idx); + + // Avoid duplicate messages to a single rank + const auto deduped_rank_idx = is_unique(rank_idx) ? rank_idx : -1; + const auto rank_idx_mask = ptx::match(deduped_rank_idx); + + // Let the one with the largest lane index send the count + if ((rank_idx_mask >> lane_idx) == 1 and deduped_rank_idx >= 0) + rank_count_warp_sum[deduped_rank_idx] += __popc(rank_idx_mask); + } + __syncthreads(); + + // Get block sum and store to global + for (int rank_idx = thread_idx; rank_idx < kNumScaleupRanks; rank_idx += kNumThreads) { + int rank_count_block_sum = 0; + for (int i = 0; i < kNumWarps; i++) + rank_count_block_sum += get_other_rank_count_warp_sum(i)[rank_idx]; + rank_count_buffer[sm_idx * kNumScaleupRanks + rank_idx] = rank_count_block_sum; + } + cooperative_groups::this_grid().sync(); + + // Get the prefix sum before the current SM + for (int rank_idx = lane_idx; rank_idx < kNumScaleupRanks; rank_idx += 32) { + int rank_count = 0; + for (int i = warp_idx; i < sm_idx; i += kNumWarps) + rank_count += rank_count_buffer[i * kNumScaleupRanks + rank_idx]; + atomicAdd_block(rank_count_global_psum + rank_idx, rank_count); + } + __syncthreads(); + + // Get each warp's prefix sum + for (int rank_idx = lane_idx; rank_idx < kNumScaleupRanks; rank_idx += 32) { + int rank_count = rank_count_global_psum[rank_idx]; + for (int i = 0; i < warp_idx; i++) + rank_count += get_other_rank_count_warp_sum(i)[rank_idx]; + rank_count_warp_psum[rank_idx] = rank_count; + } + __syncwarp(); + + // Each warp scan the tokens separately + for (int i = start_token_idx; i < end_token_idx; i += kNumTokensPerGroup) { + const auto token_idx = i + token_idx_offset; + const auto is_active_thread = lane_idx < kNumTopk * kNumTokensPerGroup and token_idx < end_token_idx; + const auto expert_idx = is_active_thread ? static_cast(__ldg(topk_idx + i * kNumTopk + lane_idx)) : -1; + const auto rank_idx = map_expert_to_rank_idx(expert_idx); + + // Avoid duplicate messages to a single rank + const auto deduped_rank_idx = is_unique(rank_idx) ? rank_idx : -1; + const auto rank_idx_mask = ptx::match(deduped_rank_idx); + + // Store to target buffer + const auto stored_dst_slot_idx = deduped_rank_idx >= 0 ? + rank_count_warp_psum[deduped_rank_idx] + count_ones_before(rank_idx_mask, lane_idx) : -1; + const auto value = stored_dst_slot_idx >= 0 ? + scaleup_rank_idx * kNumMaxTokensPerRank + stored_dst_slot_idx : -1; + if (is_active_thread) + dst_buffer_slot_idx[i * kNumTopk + lane_idx] = value; + + // Let the one with the largest lane index send the count + if ((rank_idx_mask >> lane_idx) == 1 and deduped_rank_idx >= 0) + rank_count_warp_psum[deduped_rank_idx] += __popc(rank_idx_mask); + __syncwarp(); + } +} + +} // namespace deep_ep::elastic diff --git a/deep_ep/include/deep_ep/impls/engram_fetch.cuh b/deep_ep/include/deep_ep/impls/engram_fetch.cuh new file mode 100644 index 000000000..9339f2159 --- /dev/null +++ b/deep_ep/include/deep_ep/impls/engram_fetch.cuh @@ -0,0 +1,99 @@ +#pragma once + +#include +#include +#include +#include + + +namespace deep_ep::elastic { + +template +__global__ void __launch_bounds__(kNumThreads, 1) +engram_fetch_impl(const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window, + void* storage, void* fetched, int* indices, + ncclGinRequest_t* last_gin_requests, + const int num_tokens) { + const auto qp_idx = static_cast(blockIdx.x); + const auto warp_idx = ptx::get_warp_idx(); + const auto global_warp_idx = qp_idx * kNumWarps + warp_idx; + const auto thread_idx = static_cast(threadIdx.x); + + // Gin handle + const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, NCCL_GIN_RESOURCE_SHARING_CTA); + + __shared__ bool sent_to_rank[kNumRanks]; + EP_STATIC_ASSERT(kNumRanks <= kNumThreads, "Too many ranks"); + if (thread_idx < kNumRanks) + sent_to_rank[thread_idx] = false; + __syncthreads(); + + // Issue RDMA + const auto issue_rdma_get = [=](const int& token_idx, const int& src_rank_idx, const int& src_entry_idx, + const int& extra_options = 0) { + gin.get(math::advance_ptr(storage, static_cast(src_entry_idx) * kNumHiddenBytes), + math::advance_ptr(fetched, static_cast(token_idx) * kNumHiddenBytes), + kNumHiddenBytes, src_rank_idx, extra_options); + }; + + // Each warp fetches one token cooperatively via RDMA gin.get + // TODO: deal with padded tokens + if (ptx::elect_one_sync()) { + #pragma unroll 4 + for (int i = global_warp_idx; i < num_tokens; i += kNumQPs * kNumWarps) { + const auto global_idx = __ldg(indices + i); + const auto src_rank_idx = global_idx / kNumEntriesPerRank; + const auto src_entry_idx = global_idx % kNumEntriesPerRank; + + // Delay ring DB + issue_rdma_get(i, src_rank_idx, src_entry_idx, ncclGinOptFlagsAggregateRequests); + sent_to_rank[src_rank_idx] = true; + } + } + __syncthreads(); + + // Issue flush per peer we sent to; its unconditional DB ring flushes all + // prior aggregated gets on the same QP. + if (ptx::elect_one_sync()) { + for (int i = warp_idx; i < kNumRanks; i += kNumWarps) { + const auto request_ptr = last_gin_requests + qp_idx * kNumRanks + i; + if (sent_to_rank[i]) { + gin.flush_async(i, request_ptr); + } else { + EP_STATIC_ASSERT(sizeof(ncclGinRequest_t) == sizeof(int4), "Invalid request size"); + *reinterpret_cast(request_ptr) = make_int4(0, 0, 0, 0); + } + } + } +} + +template +__global__ void __launch_bounds__(kNumThreads, 1) +engram_fetch_wait_impl(const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window, + ncclGinRequest_t* last_gin_requests) { + const auto qp_idx = static_cast(blockIdx.x); + const auto thread_idx = static_cast(threadIdx.x); + + // Gin handle + const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, NCCL_GIN_RESOURCE_SHARING_CTA); + + // Wait for all RDMA gets to complete + for (int i = thread_idx; i < kNumRanks; i += kNumThreads) { + EP_STATIC_ASSERT(sizeof(ncclGinRequest_t) == sizeof(int4), "Invalid request size"); + auto last_gin_req_int4 = __ldg(reinterpret_cast(last_gin_requests + qp_idx * kNumRanks + i)); + if (last_gin_req_int4.x != 0 or last_gin_req_int4.y != 0 or + last_gin_req_int4.z != 0 or last_gin_req_int4.w != 0) { + auto last_gin_req = *reinterpret_cast(&last_gin_req_int4); + gin.wait(last_gin_req); + } + } +} + +} // namespace deep_ep::elastic diff --git a/deep_ep/include/deep_ep/impls/hybrid_combine.cuh b/deep_ep/include/deep_ep/impls/hybrid_combine.cuh new file mode 100644 index 000000000..018f22b77 --- /dev/null +++ b/deep_ep/include/deep_ep/impls/hybrid_combine.cuh @@ -0,0 +1,616 @@ +#pragma once + +#include +#include +#include +#include +#include + +namespace deep_ep::elastic { + +template (), + bool kUseScaleupRankLayout = use_rank_layout(), + int kNumTokensInScaleoutLayout = get_num_tokens_in_layout(), + int kNumTokensInScaleupLayout = get_num_tokens_in_layout()> +__global__ void __launch_bounds__(kNumThreads, 1) +hybrid_combine_impl(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, + const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window, + void* buffer, void* workspace, + const int scaleout_rank_idx, const int scaleup_rank_idx, + int num_reduced_tokens) { + // Utils + const auto sm_idx = static_cast(blockIdx.x); + const auto thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(); + const auto lane_idx = ptx::get_lane_idx(); + constexpr bool kDoExpandedSend = not kAllowMultipleReduction and kUseExpandedLayout; + + // Combine vector type selection + using combine_vec_t = typename CombineVecTraits::vec_t; + constexpr int kHiddenVec = kNumHiddenBytes / sizeof(combine_vec_t); + + // Workspaces + const auto workspace_layout = layout::WorkspaceLayout(workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts); + + // We should assign the real number of received tokens if without CPU sync + if (num_reduced_tokens == kNumMaxTokensPerRank * kNumRanks) + num_reduced_tokens = __ldg(psum_num_recv_tokens_per_scaleup_rank + kNumScaleupRanks - 1); + + // Token layouts + const auto token_layout = layout::TokenLayout(kNumHiddenBytes, 0, kNumTopk, false); + + // TMA buffers + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + const auto tma_buffer = layout::BufferLayout( + token_layout, kNumWarps, 1, smem).get_rank_buffer(warp_idx).get_token_buffer(0); + + // All the buffer layouts + auto scaleup_buffer = layout::BufferLayout( + token_layout, kNumTokensInScaleupLayout, kNumScaleoutRanks * kNumMaxTokensPerRank, + buffer); + auto scaleout_recv_buffer = layout::BufferLayout( + token_layout, kNumTokensInScaleoutLayout, kNumMaxTokensPerRank, + scaleup_buffer.get_buffer_end_ptr()); + auto scaleout_send_buffer = layout::BufferLayout( + token_layout, kAllowMultipleReduction ? 1 : kNumTopk, kNumChannels * (kNumScaleoutRanks * kNumMaxTokensPerChannel), + scaleout_recv_buffer.get_buffer_end_ptr()); + + // Init TMA for scale-up and forward warps + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (ptx::elect_one_sync()) + ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // NCCL Gin handle + // Each warp is a channel + const auto [qp_idx, sharing_mode] = + comm::get_qp_mode(sm_idx, warp_idx % kNumChannelsPerSM); + const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, sharing_mode); + + // Global parallel barriers for scale-out subteam and scale-up subteam + // NOTES: this barrier needs a grid sync, as there are channel scale-up tail cleaning before + comm::gpu_barrier( + gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx); + + // Adjust register count at certain cases + // TODO: support more cases, or try to make channel count more aligned + const bool kAdjustRegisters = (kNumChannelsPerSM == 4 or kNumChannelsPerSM == 8) and not kUseExpandedLayout; + constexpr int kNumRegistersForScaleupWarps = 40; + constexpr int kNumRegistersForForwardWarps = 256 - kNumRegistersForScaleupWarps; + + // Different warp roles + if (warp_idx < kNumScaleupWarps) { + const auto channel_idx = sm_idx * kNumChannelsPerSM + warp_idx; + + // Adjust registers + if constexpr (kAdjustRegisters) + ptx::warpgroup_reg_dealloc(); + + // Shift into the right buffer if using rank layout + if constexpr (kUseScaleupRankLayout) + scaleup_buffer = scaleup_buffer.get_rank_buffer(scaleup_rank_idx); + + // Expanding mode must not be backward + if constexpr (kUseExpandedLayout) + EP_DEVICE_ASSERT(topk_weights == nullptr); + + // Tail issuer + // `st.release.sys` is pretty slow, so do it by an interval + int update_counter = 0; + int stored_num_tokens_sent[kNumScaleupRanksPerLane] = {}; + int stored_old_num_tokens_sent[kNumScaleupRanksPerLane] = {}; + const auto tail_ptr = workspace_layout.get_channel_scaleup_tail_ptr(channel_idx, scaleup_rank_idx); + const auto update_tails = [&](const bool& finish = false) { + ++ update_counter; + if (finish or update_counter == kNumScaleupUpdateInterval) { + // Wait all TMA stores to finish + ptx::tma_store_wait(); + __syncwarp(); + + // Issue + #pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++ i) { + if (const auto j = i * 32 + lane_idx; i < (kNumScaleupRanksPerLane - 1) or j < kNumScaleupRanks) { + // NOTES: save some traffic with `stored_old_num_tokens_sent` + // Also, we cannot rewrite a finished slot, if the peer is going to clean it + if (stored_num_tokens_sent[i] != stored_old_num_tokens_sent[i]) + ptx::st_release_sys(gin.get_sym_ptr(tail_ptr, j), stored_num_tokens_sent[i]); + stored_old_num_tokens_sent[i] = stored_num_tokens_sent[i]; + } + } + update_counter = 0; + } + __syncwarp(); + }; + + // Shape of `channel_linked_list`: `[kNumChannels, kNumMaxTokensPerChannel + 1, kNumScaleupRanks]` + // Iterate until all scale-up peers finish + int dst_scaleup_rank_idx = channel_idx; + int stored_ll_idx[kNumScaleupRanksPerLane] = {}, stored_token_idx[kNumScaleupRanksPerLane] = {}; + #pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++ i) + stored_token_idx[i] = -1; + while (true) { + // Load token indices in the list + #pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++ i) { + const auto j = i * 32 + lane_idx; + stored_token_idx[i] = i < (kNumScaleupRanksPerLane - 1) or j < kNumScaleupRanks ? + __ldg(channel_linked_list + + channel_idx * (kNumScaleoutRanks * kNumMaxTokensPerChannel + 1) * kNumScaleupRanks + + stored_ll_idx[i] * kNumScaleupRanks + j) : -1; + } + __syncwarp(); + + // Check whether all ranks are finished + bool exited = true; + #pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++ i) + exited &= ptx::all(stored_token_idx[i] < 0); + if (exited) + break; + + // Process tokens for all ranks together using bitmask to skip inactive ranks + EP_STATIC_ASSERT(kNumScaleupRanks <= 64, "Too many scale-up ranks for 64-bit mask"); + using mask_t = std::conditional_t<(kNumScaleupRanks <= 32), uint32_t, uint64_t>; + mask_t wip_mask = 0; + #pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) + wip_mask |= static_cast(ptx::gather(stored_token_idx[j] >= 0)) << (j * 32); + while (wip_mask) { + // Find next active rank after `dst_scaleup_rank_idx` (round-robin) + const auto start = (dst_scaleup_rank_idx + 1) % kNumScaleupRanks; + const auto hi_mask = (wip_mask >> start) << start; + dst_scaleup_rank_idx = hi_mask ? ptx::ffs(hi_mask) : ptx::ffs(wip_mask); + wip_mask ^= static_cast(1) << dst_scaleup_rank_idx; + + // Exchange token index from the owning lane using static partition iteration + int token_idx = -1; + #pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) { + const auto src_lane_idx = dst_scaleup_rank_idx - j * 32; + token_idx = src_lane_idx == lane_idx ? stored_token_idx[j] : token_idx; + } + token_idx = ptx::exchange(token_idx, dst_scaleup_rank_idx % 32); + + // Get source metadata and decide the destination buffer + constexpr int kMetadataStride = 2 + kNumTopk; + const auto src_global_token_idx = __ldg(src_metadata + token_idx * kMetadataStride + 0); + const auto src_token_idx = src_global_token_idx % kNumMaxTokensPerRank; + const auto src_scaleout_rank_idx = src_global_token_idx / (kNumMaxTokensPerRank * kNumScaleupRanks); + auto token_buffer = [&]() { + if constexpr (kUseScaleupRankLayout) { + const auto src_slot_idx = __ldg(src_metadata + token_idx * kMetadataStride + 1) / kNumTopk; + return scaleup_buffer.get_token_buffer(src_slot_idx); + } else { + const auto master_topk_idx = __ldg(src_metadata + token_idx * kMetadataStride + 1) % kNumTopk; + return scaleup_buffer + .get_rank_buffer(master_topk_idx) + .get_token_buffer(src_scaleout_rank_idx * kNumMaxTokensPerRank + src_token_idx); + } + }(); + token_buffer.set_base_ptr(gin.get_sym_ptr(token_buffer.get_base_ptr(), dst_scaleup_rank_idx)); + + // Some checks + EP_STATIC_ASSERT(kHidden % (32 * sizeof(int4) / sizeof(nv_bfloat16)) == 0, "Invalid hidden"); + + // Read source indices for expand mode + int stored_topk_slot_idx = -1; + if constexpr (kUseExpandedLayout) { + if (lane_idx < kNumTopk) + stored_topk_slot_idx = __ldg(src_metadata + token_idx * kMetadataStride + (2 + lane_idx)); + __syncwarp(); + } + + // 3 cases: + // - no-expand, expand + no-reduce + // - expand + reduce + // - expand + send all + auto reduce_valid_mask = ptx::gather(stored_topk_slot_idx >= 0); + auto no_local_reduce = not kUseExpandedLayout or (kAllowMultipleReduction and __popc(reduce_valid_mask) == 1); + if (no_local_reduce) { + int token_idx_in_tensor = token_idx; + if constexpr (kUseExpandedLayout) + token_idx_in_tensor = ptx::exchange(stored_topk_slot_idx, ptx::get_master_lane_idx(reduce_valid_mask)); + + // Directly load + if (ptx::elect_one_sync()) { + const auto load_ptr = + math::advance_ptr(x, static_cast(token_idx_in_tensor) * kNumHiddenBytes); + ptx::tma_store_wait(); + ptx::tma_load_1d(tma_buffer.get_base_ptr(), load_ptr, mbarrier_ptr, kNumHiddenBytes); + } + __syncwarp(); + } else if constexpr (kAllowMultipleReduction) { + // Do local reduction + // Sort valid top-k indices to front + int topk_slot_idx[kNumTopk]; + compute_topk_slots( + topk_slot_idx, reduce_valid_mask, + [=](const int& idx) { + return ptx::exchange(stored_topk_slot_idx, idx); + } + ); + + // Reduce into shared memory + constexpr int kUnrollFactor = get_max_unroll_factor(); + combine_reduce( + lane_idx, topk_slot_idx, static_cast(tma_buffer.get_base_ptr()), + /* Get source base */ [=](const int& slot_idx) { + return math::advance_ptr( + x, slot_idx * static_cast(kNumHiddenBytes)); + }, + /* Wait buffer release */ [=]() { + ptx::tma_store_wait(); + __syncwarp(); + } + ); + ptx::tma_store_fence(); + __syncwarp(); + } else { + // No local reduction, send all data (expanded send) + #pragma unroll + for (int k = 0; k < kNumTopk; ++ k) { + int topk_slot_idx = ptx::exchange(stored_topk_slot_idx, k); + if (topk_slot_idx < 0) + continue; + + if (ptx::elect_one_sync()) { + // Load + const auto load_ptr = math::advance_ptr(x, static_cast(kDoExpandedSend ? topk_slot_idx : token_idx) * kNumHiddenBytes); + ptx::tma_store_wait(); + ptx::tma_load_1d(tma_buffer.get_base_ptr(), load_ptr, mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + // NOTES: We don't need to care about `topk_weights` since we are in expand mode + + // Store + const auto dst_token_buffer = scaleup_buffer + .get_rank_buffer(k) + .get_token_buffer(src_scaleout_rank_idx * kNumMaxTokensPerRank + src_token_idx); + ptx::tma_store_1d( + gin.get_sym_ptr(dst_token_buffer.get_base_ptr(), dst_scaleup_rank_idx), + tma_buffer.get_base_ptr(), token_layout.get_num_bytes()); + ptx::tma_store_commit(); + } + __syncwarp(); + } + } + + // Write top-k weights + if (not kUseExpandedLayout and topk_weights != nullptr and lane_idx < kNumTopk) { + const float value = __ldg(topk_weights + (token_idx * kNumTopk + lane_idx)); + tma_buffer.get_topk_weights_ptr()[lane_idx] = value; + ptx::tma_store_fence(); + } + __syncwarp(); + + // Issue TMA stores into remote scale-up buffer + // NOTES: `kDoExpandedSend` mode has already issued + if (not kDoExpandedSend and ptx::elect_one_sync()) { + // Wait TMA arrival (only for non-reduced cases) + if (no_local_reduce) { + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + } + + // Issue stores + ptx::tma_store_1d( + token_buffer.get_base_ptr(), tma_buffer.get_base_ptr(), + token_layout.get_num_bytes()); + ptx::tma_store_commit(); + } + #pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) + stored_num_tokens_sent[j] += (j * 32 + lane_idx) == dst_scaleup_rank_idx; + __syncwarp(); + } + + // Update the tails together + // NOTES: TMA wait is inside + update_tails(); + + // Move linked list + #pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++ i) + stored_ll_idx[i] += (stored_token_idx[i] >= 0); + } + + // Update for the unissued ones + update_tails(true); + } else { + const auto forward_warp_idx = warp_idx - kNumScaleupWarps; + const auto channel_idx = sm_idx * kNumChannelsPerSM + forward_warp_idx; + + // Adjust registers + if constexpr (kAdjustRegisters) + ptx::warpgroup_reg_alloc(); + + // Shift into the right buffer + scaleout_send_buffer = scaleout_send_buffer.get_channel_buffer(channel_idx); + + // Shape of `token_metadata_at_forward`: `[kNumChannels, kNumScaleoutRanks * kNumMaxTokensPerChannel + 1, kNumForwardMetadataDims]` + constexpr int kNumForwardMetadataDims = 2 + kNumTopk * 2; + token_metadata_at_forward += channel_idx * ((kNumScaleoutRanks * kNumMaxTokensPerChannel + 1) * kNumForwardMetadataDims); + + // Overlap TMA stores and reduction + int last_src_scaleout_rank_idx = -1; + int last_is_token_last_in_chunk = 0; + void* last_recv_token_buffer_ptr = nullptr; + void* last_send_token_buffer_ptr = nullptr; + const auto flush_last_tma_and_issue_rdma = [&]() { + if (last_src_scaleout_rank_idx >= 0 and ptx::elect_one_sync()) { + ptx::tma_store_wait(); + + // Issue only if not local rank + if (last_src_scaleout_rank_idx != scaleout_rank_idx) { + gin.put( + last_recv_token_buffer_ptr, + last_send_token_buffer_ptr, + token_layout.get_num_bytes(), + last_src_scaleout_rank_idx, + last_is_token_last_in_chunk ? 0 : ncclGinOptFlagsAggregateRequests + ); + } + } + __syncwarp(); + }; + + // Replay the dispatch + int stored_num_tokens_recv[kNumScaleupRanksPerLane] = {}, stored_cached_scaleup_tail[kNumScaleupRanksPerLane] = {}; + for (int i = 0; ; ++ i) { + const auto src_token_global_idx = __ldg(token_metadata_at_forward + i * kNumForwardMetadataDims); + const auto is_token_last_in_chunk = __ldg(token_metadata_at_forward + i * kNumForwardMetadataDims + 1); + const auto src_rank_idx = src_token_global_idx / kNumMaxTokensPerRank; + const auto src_scaleout_rank_idx = src_rank_idx / kNumScaleupRanks; + const auto src_token_idx = src_token_global_idx % kNumMaxTokensPerRank; + auto stored_src_scaleup_rank_idx = lane_idx < kNumTopk ? + __ldg(token_metadata_at_forward + i * kNumForwardMetadataDims + 2 + lane_idx) : -1; + auto stored_src_slot_idx = lane_idx < kNumTopk ? + __ldg(token_metadata_at_forward + i * kNumForwardMetadataDims + 2 + kNumTopk + lane_idx) : -1; + if (src_token_global_idx < 0) + break; + + // Scaleup rank mask + EP_STATIC_ASSERT(kNumScaleupRanks <= 64, "Too many scale-up peers"); + using mask_t = std::conditional_t; + const auto scaleup_mask = ptx::reduce_or( + stored_src_scaleup_rank_idx >= 0 ? + (mask_t(1) << stored_src_scaleup_rank_idx) : mask_t(0)); + bool stored_is_scaleup_rank_needed[kNumScaleupRanksPerLane]; + #pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) + stored_is_scaleup_rank_needed[j] = (scaleup_mask >> (j * 32 + lane_idx)) & 1; + + // Wait all tails to arrive + comm::timeout_while([&](const bool& is_last_check) { + bool arrived = true; + #pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) + arrived &= not stored_is_scaleup_rank_needed[j] or stored_num_tokens_recv[j] < stored_cached_scaleup_tail[j]; + if (ptx::all(arrived)) + return true; + + // Reload cached + #pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) { + const auto k = j * 32 + lane_idx; + stored_cached_scaleup_tail[j] = j < (kNumScaleupRanksPerLane - 1) or k < kNumScaleupRanks ? + ptx::ld_acquire_sys(workspace_layout.get_channel_scaleup_tail_ptr(channel_idx, k)) : -1; + } + + // Timeout + if (is_last_check) { + #pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) { + printf("DeepEP combine (scale-up wait) timeout, scale-out: %d/%d, scale-up: %d/%d, " + "channel: %d, lane: %d, recv: %d, tail: %d (wait=%d)\n", + scaleout_rank_idx, kNumScaleoutRanks, scaleup_rank_idx, kNumScaleupRanks, + channel_idx, j * 32 + lane_idx, + stored_num_tokens_recv[j], + stored_cached_scaleup_tail[j], + stored_is_scaleup_rank_needed[j]); + } + } + return false; + }); + + // Increase received count + #pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) + stored_num_tokens_recv[j] += static_cast(stored_is_scaleup_rank_needed[j]); + + if constexpr (not kAllowMultipleReduction) { + // Cases where multiple reduction is disabled. We need to forward all data from scaleup peers to scaleout peers + // TODO: Let scale-up warps directly put data into `send_buffer`? + const auto src_slot_idx = src_scaleout_rank_idx * kNumMaxTokensPerRank + src_token_idx; + auto topk_valid_mask = kUseExpandedLayout ? + ptx::gather(stored_src_scaleup_rank_idx >= 0) : + ptx::gather(ptx::deduplicate(stored_src_scaleup_rank_idx, lane_idx) and stored_src_scaleup_rank_idx >= 0); // Deduplicate w.r.t. scaleup rank index if expanded mode is disabled + if (ptx::elect_one_sync()) { + #pragma unroll + for (int k = 0; k < kNumTopk; ++ k) { + if ((topk_valid_mask & (1u << k)) == 0u) + continue; + + // Issue TMA load, and wait + ptx::tma_load_1d( + tma_buffer.get_base_ptr(), scaleup_buffer.get_rank_buffer(k).get_token_buffer(src_slot_idx).get_base_ptr(), + mbarrier_ptr, token_layout.get_num_bytes()); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, token_layout.get_num_bytes()); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + + // Issue TMA store, and wait + const auto recv_buffer_ptr = scaleout_recv_buffer.get_rank_buffer(k).get_token_buffer(src_token_idx).get_base_ptr(); + const auto send_buffer_ptr = src_scaleout_rank_idx == scaleout_rank_idx ? + recv_buffer_ptr : scaleout_send_buffer.get_rank_buffer(k).get_token_buffer(i).get_base_ptr(); + ptx::tma_store_1d(send_buffer_ptr, tma_buffer.get_base_ptr(), token_layout.get_num_bytes()); + ptx::tma_store_commit(); + ptx::tma_store_wait(); + + // Issue IBGDA + topk_valid_mask ^= 1u << k; + if (src_scaleout_rank_idx != scaleout_rank_idx) { + gin.put( + recv_buffer_ptr, + send_buffer_ptr, + token_layout.get_num_bytes(), + src_scaleout_rank_idx, + topk_valid_mask == 0 and is_token_last_in_chunk ? 0 : ncclGinOptFlagsAggregateRequests + ); + } + } + } + __syncwarp(); + } else { + // NOTES: we must do deduplicate and only add once from one rank + auto reduce_valid_mask = ptx::gather( + ptx::deduplicate(stored_src_scaleup_rank_idx, lane_idx) and stored_src_scaleup_rank_idx >= 0); + + // Calculate the source buffer index + int stored_src_buffer_idx = 0; + if constexpr (kUseScaleupRankLayout) { + stored_src_buffer_idx = + stored_src_scaleup_rank_idx * scaleup_buffer.num_max_tokens_per_rank + stored_src_slot_idx; + } else { + const auto src_slot_idx = src_scaleout_rank_idx * kNumMaxTokensPerRank + src_token_idx; + stored_src_buffer_idx = stored_src_slot_idx == -1 ? -1 : + lane_idx * scaleup_buffer.num_max_tokens_per_rank + src_slot_idx; + } + + // Preprocess top-k indices + int topk_slot_idx[kNumTokensInScaleupLayout]; + compute_topk_slots( + topk_slot_idx, reduce_valid_mask, + [=](const int& idx) { + return ptx::exchange(stored_src_buffer_idx, idx); + } + ); + + // Do reduce + constexpr int kUnrollFactor = get_max_unroll_factor(); + combine_reduce( + lane_idx, topk_slot_idx, static_cast(tma_buffer.get_base_ptr()), + /* Get source base */ [=](const int& slot_idx) { + return static_cast(scaleup_buffer.get_token_buffer(slot_idx, true).get_base_ptr()); + }, + /* Wait buffer release */ [=]() { + flush_last_tma_and_issue_rdma(); + } + ); + + // Merge topk weights + // NOTES: the slot indices must follow the master lane + stored_src_buffer_idx = ptx::exchange( + stored_src_buffer_idx, ptx::get_master_lane_idx(ptx::match(stored_src_scaleup_rank_idx))); + if (not kUseExpandedLayout and stored_src_scaleup_rank_idx >= 0) { + tma_buffer.get_topk_weights_ptr()[lane_idx] = + scaleup_buffer.get_token_buffer(stored_src_buffer_idx, true) + .get_topk_weights_ptr()[lane_idx]; + } + ptx::tma_store_fence(); + __syncwarp(); // Necessary to let the leader lane see the writes + + // Assign send and receive buffers + // NOTES: as we only have 1 destination, we will use "send" as "recv" for local transfer + int scaleout_recv_buffer_rank_idx; + if constexpr (kUseScaleoutRankLayout) { + scaleout_recv_buffer_rank_idx = scaleout_rank_idx; + } else { + const int src_topk_idx = ptx::get_master_lane_idx(ptx::gather(stored_src_scaleup_rank_idx >= 0)); + scaleout_recv_buffer_rank_idx = src_topk_idx; + } + const auto recv_token_buffer = scaleout_recv_buffer.get_rank_buffer(scaleout_recv_buffer_rank_idx).get_token_buffer(src_token_idx); + const auto send_token_buffer = src_scaleout_rank_idx == scaleout_rank_idx ? + recv_token_buffer : + scaleout_send_buffer.get_token_buffer(i); + + // Write into scale-out send buffer or local rank recv buffer bypass + if (ptx::elect_one_sync()) { + ptx::tma_store_1d(send_token_buffer.get_base_ptr(), tma_buffer.get_base_ptr(), + token_layout.get_num_bytes()); + ptx::tma_store_commit(); + } + __syncwarp(); + + // Record RDMA info to issue later + last_src_scaleout_rank_idx = src_scaleout_rank_idx; + last_is_token_last_in_chunk = is_token_last_in_chunk; + last_recv_token_buffer_ptr = recv_token_buffer.get_base_ptr(); + last_send_token_buffer_ptr = send_token_buffer.get_base_ptr(); + } + } + + // Issue the last RDMA + if constexpr (kAllowMultipleReduction) + flush_last_tma_and_issue_rdma(); + + // Clean scaleup tails + #pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) { + const auto k = j * 32 + lane_idx; + if (j < (kNumScaleupRanksPerLane - 1) or k < kNumScaleupRanks) + *workspace_layout.get_channel_scaleup_tail_ptr(channel_idx, k) = 0; + } + __syncwarp(); + + // Update, wait and clean + EP_STATIC_ASSERT(kNumScaleoutRanks <= 32, "Invalid ranks"); + if (lane_idx < kNumScaleoutRanks) { + // Update remote tails + const auto expected_signal = math::pack2(1, 0); + gin.red_add_rel( + workspace_layout.get_scaleout_channel_signaled_tail_ptr(channel_idx, scaleout_rank_idx), + expected_signal, lane_idx); + + // Wait tail arrival + const auto wait_ptr = workspace_layout.get_scaleout_channel_signaled_tail_ptr(channel_idx, lane_idx); + comm::timeout_while([=](const bool& is_last_check) { + const auto signal = ptx::ld_acquire_sys(wait_ptr); + if (signal == expected_signal) { + // Clean for next usages + *wait_ptr = 0; + return true; + } + + if (is_last_check) { + printf("DeepEP combine (scale-out wait all) timeout, scale-out: %d/%d, scale-up: %d/%d, " + "channel: %d, lane: %d, signal: %lld, expected: %lld\n", + scaleout_rank_idx, kNumScaleoutRanks, scaleup_rank_idx, kNumScaleupRanks, + channel_idx, lane_idx, + signal, expected_signal); + } + return false; + }); + } + __syncwarp(); + } + + // No barrier at epilogue +} + +} // namespace deep_ep::elastic diff --git a/deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh b/deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh new file mode 100644 index 000000000..7d75b0238 --- /dev/null +++ b/deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh @@ -0,0 +1,672 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + + +namespace deep_ep::elastic { + +template +__global__ void __launch_bounds__(kNumThreads, 1) +hybrid_dispatch_impl( + 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* dst_buffer_slot_idx, + int* token_metadata_at_forward, + const int num_tokens, + const int sf_token_stride, const int sf_hidden_stride, + // TODO(NCCL): so many params, plans to optimize? + 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) { + constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks; + constexpr int kNumExpertsPerScaleout = kNumExperts / kNumScaleoutRanks; + EP_STATIC_ASSERT(kNumExperts % kNumScaleupRanks == 0, "Invalid number of experts or ranks"); + EP_STATIC_ASSERT(kNumNotifyWarps % 4 == 0, "Invalid warpgroup size"); + EP_STATIC_ASSERT(kNumScaleoutWarps == kNumForwardWarps, "Invalid warp size"); + + // Utils + // NOTES: a warp is a channel (different channels may share QPs) + const auto sm_idx = static_cast(blockIdx.x), thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + const auto rank_idx = scaleout_rank_idx * kNumScaleupRanks + scaleup_rank_idx; + + // Workspaces + const auto workspace_layout = layout::WorkspaceLayout(workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts); + const auto host_workspace_layout = layout::WorkspaceLayout(mapped_host_workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts); + + // The kernel uses a fixed space of dynamic shared memory (no static shared memory) + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + constexpr int kNumSmemBytesForNotify = kNumNotifyThreads > 0 ? + math::constexpr_align(kNumRanks + kNumExperts, kNumNotifyThreads) * sizeof(int) : 0; + EP_STATIC_ASSERT(kNumSmemBytesForNotify % ptx::kNumTMAAlignBytes == 0, "Invalid TMA alignment"); + + // Named barrier indices + constexpr int kNotifyBarrierIndex = 1; + + // NCCL Gin handle + // Each warp is a channel + const auto [qp_idx, sharing_mode] = comm::get_qp_mode 0)>( + sm_idx, (warp_idx - kNumNotifyWarps) % kNumChannelsPerSM, warp_idx < kNumNotifyWarps); + const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, sharing_mode); + + // Global parallel barriers for scale-out subteam and scale-up subteam + comm::gpu_barrier( + gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx); + + // The golden layout during the whole process for both scale-out and forward warps + const auto token_layout = layout::TokenLayout(kNumHiddenBytes, kNumSFPacks * sizeof(sf_pack_t), kNumTopk, true); + const auto tma_buffer = layout::BufferLayout(token_layout, kNumScaleoutWarps + kNumForwardWarps, 1, + math::advance_ptr(smem, kNumSmemBytesForNotify)).get_rank_buffer(warp_idx - kNumNotifyWarps).get_token_buffer(0); + + // All the buffers + auto scaleup_buffer = layout::BufferLayout( + token_layout, kNumScaleupRanks, kNumScaleoutRanks * kNumMaxTokensPerRank, buffer); + auto scaleout_send_buffer = layout::BufferLayout( + token_layout, 1, kNumMaxTokensPerRank, scaleup_buffer.get_buffer_end_ptr()); + auto scaleout_recv_buffer = layout::BufferLayout( + token_layout, kNumScaleoutRanks, kNumChannels * kNumMaxTokensPerChannel, scaleout_send_buffer.get_buffer_end_ptr()); + + // Init TMA for scale-out and forward warps + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (warp_idx >= kNumNotifyWarps and ptx::elect_one_sync()) + ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Different warp roles + if (warp_idx < kNumNotifyWarps) { + // Assign shared memory + constexpr int kNumAlignedElems = kNumSmemBytesForNotify / sizeof(int); + const auto rank_expert_count = math::advance_ptr(smem, 0); + + // Clean initial counts + // NOTES: if you want to change the order of different warp roles, please take care of the `thread_idx` + int *rank_count = rank_expert_count, *expert_count = rank_expert_count + kNumRanks; + #pragma unroll + for (int i = 0; i < kNumAlignedElems / kNumNotifyThreads; ++ i) + rank_expert_count[i * kNumNotifyThreads + thread_idx] = 0; + ptx::named_barrier(kNotifyBarrierIndex); + + // Atomic add on shared memory + EP_STATIC_ASSERT(kNumTopk <= 32, "Insufficient lanes"); + const auto global_warp_idx = sm_idx * kNumNotifyWarps + warp_idx; + for (int i = global_warp_idx; i < num_tokens; i += kNumNotifyWarps * kNumSMs) { + // Expert choice can not be redundant + // NOTES: no assertions here as they are expensive + const auto dst_expert_idx = lane_idx < kNumTopk ? + static_cast(__ldg(topk_idx + i * kNumTopk + lane_idx)) : -1; + if (dst_expert_idx >= 0) + atomicAdd_block(expert_count + dst_expert_idx, 1); + + // Rank choice should do deduplication here + const auto dst_rank_idx = dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerRank : -1; + if (ptx::deduplicate(dst_rank_idx, lane_idx) and dst_rank_idx >= 0) + atomicAdd_block(rank_count + dst_rank_idx, 1); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // Do full-grid reduction + #pragma unroll + for (int i = thread_idx; i < kNumRanks + kNumExperts; i += kNumNotifyThreads) { + const int64_t counter = (1ll << 32ll) | rank_expert_count[i]; + ptx::red_add(workspace_layout.get_notify_reduction_workspace_ptr() + i, counter); + } + + // Do the remaining work by SM 0 + if (sm_idx == 0) { + // Reduce all SM's count + // Wait all SMs' arrival + #pragma unroll + for (int i = thread_idx; i < kNumRanks + kNumExperts; i += kNumNotifyThreads) { + comm::timeout_while([=](const bool& is_last_check) { + const auto status = ptx::ld_volatile(workspace_layout.get_notify_reduction_workspace_ptr() + i); + if ((status >> 32) == kNumSMs) { + // Encode and write into the send buffer + workspace_layout.get_scaleout_rank_expert_count_ptr()[i] = + math::encode_decode_positive(status & 0xffffffffll); + + // Clean for the next usage + workspace_layout.get_notify_reduction_workspace_ptr()[i] = 0; + return true; + } + + if (is_last_check) { + printf("DeepEP hybrid notify (GPU reduction) timeout, scale-out: %d/%d, scale-up: %d/%d, " + "thread: %d, status: %d | %d, expected: %d\n", + scaleout_rank_idx, kNumScaleoutRanks, scaleup_rank_idx, kNumScaleupRanks, thread_idx, + static_cast(status >> 32), static_cast(status & 0xffffffff), kNumSMs); + } + return false; + }); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // Issue scaleout writes to peers + EP_STATIC_ASSERT(kReuseSlotIndices or kNumScaleoutRanks <= kNumNotifyThreads, + "kNumScaleoutRanks must be less than kNumNotifyThreads"); + if (thread_idx < kNumScaleoutRanks) { + const auto dst_scaleout_rank_idx = thread_idx; + gin.put( + workspace_layout.get_scaleout_rank_count_ptr(scaleout_rank_idx), + workspace_layout.get_scaleout_rank_count_ptr(dst_scaleout_rank_idx), + kNumScaleupRanks * sizeof(int), dst_scaleout_rank_idx, + ncclGinOptFlagsAggregateRequests); + gin.put( + workspace_layout.get_scaleout_expert_count_ptr(scaleout_rank_idx), + workspace_layout.get_scaleout_expert_count_ptr(dst_scaleout_rank_idx), + kNumExpertsPerScaleout * sizeof(int), dst_scaleout_rank_idx); + } + __syncwarp(); + + // Util functions to get metadata from scale-out peers + // NOTES: this is correct as RDMA operations has a minimum write granularity of 1024 bytes (a whole integer write is atomic) + const auto recv_and_reduce = [=](const auto& get_ptr_func, const bool& is_expert_reduction = false) -> int { + int count = 0; + #pragma unroll + for (int j = 0; j < kNumScaleoutRanks; ++ j) { + const auto ptr = get_ptr_func(j); + int decoded; + comm::timeout_while([&](const bool& is_last_check){ + decoded = math::encode_decode_positive(ptx::ld_acquire_sys(ptr)); + if (math::is_decoded_positive_ready(decoded)) + return true; + + if (is_last_check) { + printf("DeepEP hybrid notify (scale-out %s reduction) timeout, " + "scale-out: %d, scale-up: %d, " + "thread: %d, wait scale-out: %d, decoded: %d\n", + is_expert_reduction ? "expert" : "rank", + scaleout_rank_idx, scaleup_rank_idx, thread_idx, j, + decoded); + } + return false; + }); + + // Add and clean for next usages + count += decoded, *ptr = 0; + } + return count; + }; + + // Write into all scale-up peers' rank-level counters + #pragma unroll + for (int i = thread_idx; i < kNumScaleupRanks; i += kNumNotifyThreads) { + // Wait scale-out arrival and reduce + const auto count = recv_and_reduce([=](const int& scaleout_peer_idx) { + return workspace_layout.get_scaleout_rank_count_ptr(scaleout_peer_idx, i); + }); + + // Write into the remote scale-up peer + const int64_t counter = (static_cast(kNumScaleupRanks) << 32ll) | count; + gin.put_value( + workspace_layout.get_scaleup_rank_count_ptr() + scaleup_rank_idx, + counter, i); + } + __syncwarp(); + + // Atomic add into all scale-up peers' expert-level counters + #pragma unroll + for (int i = thread_idx; i < kNumExpertsPerScaleout; i += kNumNotifyThreads) { + // Wait scale-out arrival and reduce + const auto count = recv_and_reduce([=](const int& scaleout_peer_idx) { + return workspace_layout.get_scaleout_expert_count_ptr(scaleout_peer_idx, i); + }, true); + + // Write into the remote scale-up peer + const int64_t counter = (1ll << 32ll) | count; + const auto dst_scaleup_rank_idx = i / kNumExpertsPerRank; + const auto expert_idx_in_dst_rank = i % kNumExpertsPerRank; + gin.red_add_rel( + workspace_layout.get_scaleup_expert_count_ptr() + expert_idx_in_dst_rank, + counter, dst_scaleup_rank_idx); + } + // There are shared memory reads above, a barrier is necessary + ptx::named_barrier(kNotifyBarrierIndex); + + // NOTES: from now on, the `rank` and `expert`s size change into the local size + expert_count = rank_expert_count + kNumScaleupRanks; + + // Wait local counters to be ready + // NOTES: here we only care the prefix sum by scale-up peers (used for later epilogue), not all ranks + EP_STATIC_ASSERT(kNumNotifyWarps == 0 or kNumScaleupRanks + kNumExpertsPerRank <= kNumNotifyWarps * 32, + "Insufficient notify threads"); + comm::timeout_while(thread_idx < kNumScaleupRanks + kNumExpertsPerRank, + [&](const bool& is_last_check) { + const auto status = ptx::ld_volatile(workspace_layout.get_scaleup_rank_expert_count_ptr() + thread_idx); + if ((status >> 32ull) == kNumScaleupRanks) { + // Clean GPU workspace and write into host workspace + const auto count = static_cast(status & 0xffffffffll); + const auto aligned_count = math::align( + count, thread_idx < kNumScaleupRanks ? 1 : kExpertAlignment); + + workspace_layout.get_scaleup_rank_expert_count_ptr()[thread_idx] = 0; + if constexpr (kDoCPUSync) { + host_workspace_layout.get_scaleup_rank_expert_count_ptr()[thread_idx] = + math::encode_decode_positive(aligned_count); + } + + // Update statistics counters + if (cumulative_local_expert_recv_stats != nullptr and thread_idx >= kNumScaleupRanks) + atomicAdd(cumulative_local_expert_recv_stats + (thread_idx - kNumScaleupRanks), count); + + // Save for later prefix sum calculation + rank_expert_count[thread_idx] = aligned_count; + return true; + } + + if (is_last_check) { + printf("DeepEP hybrid notify (scale-up reduction) timeout," + "scale-out: %d/%d, scale-up: %d/%d, " + "thread: %d, status: %d | %d, expected: %d\n", + scaleout_rank_idx, kNumScaleoutRanks, scaleup_rank_idx, kNumScaleupRanks, thread_idx, + static_cast(status >> 32), static_cast(status & 0xffffffff), kNumScaleupRanks); + } + return false; + }); + ptx::named_barrier(kNotifyBarrierIndex); + + // Do prefix sum by the warps of the first SM + // NOTES: we may have fast implementation with `cub::BlockScan`, but it is too heavy to use + const auto do_psum = [=](const int* count, int* out, const int n, const int is_exclusive) { + int psum = 0; + #pragma unroll + for (int i = 0; i < math::ceil_div(n + is_exclusive, 32); ++ i) { + const auto idx = i * 32 + lane_idx; + const auto mem_idx = idx - is_exclusive; + const auto value = (0 <= mem_idx and mem_idx < n) ? count[mem_idx] : 0; + const auto sum = psum + ptx::warp_inclusive_sum(value, lane_idx); + + // Store into global memory + if (idx < n + is_exclusive) + out[idx] = sum; + + // Update `psum` by using the last lane's value + psum = ptx::exchange(sum, 31); + } + }; + if (warp_idx == 0) { + // Inclusive prefix sum + do_psum(rank_count, psum_num_recv_tokens_per_scaleup_rank, kNumScaleupRanks, 0); + } else if (warp_idx == 1) { + // Exclusive prefix sum for later expanding + do_psum(expert_count, psum_num_recv_tokens_per_expert, kNumExpertsPerRank, 1); + } + } + } else if (warp_idx < kNumNotifyWarps + kNumScaleoutWarps) { + const int scaleout_warp_idx = warp_idx - kNumNotifyWarps; + const int channel_idx = sm_idx * kNumChannelsPerSM + scaleout_warp_idx; + scaleout_recv_buffer = scaleout_recv_buffer.get_rank_buffer(scaleout_rank_idx); + scaleout_recv_buffer = scaleout_recv_buffer.get_channel_buffer(channel_idx); + + // Channel metadata maintenance + EP_STATIC_ASSERT(kNumScaleoutRanks <= 32, "Invalid number of scale-out ranks"); + int stored_scaleout_tail = 0, stored_old_scaleout_tail = 0; + const auto update_scaleout_tail = [&](const bool& finish_flag = false) { + if (lane_idx < kNumScaleoutRanks and + (stored_scaleout_tail >= stored_old_scaleout_tail + kScaleoutUpdateInterval or finish_flag)) { + const auto signaled_tail = math::pack2(finish_flag, stored_scaleout_tail); + const auto ptr = workspace_layout.get_scaleout_channel_signaled_tail_ptr(channel_idx, scaleout_rank_idx); + const auto old_signaled_tail = math::pack2(0, stored_old_scaleout_tail); + + // NOTES: the "release" scope will be `sys` for the local rank (we may involve NVLink so not `gpu`) + // For RDMA requests, "release" is ensured by "atomic" + gin.red_add_rel(ptr, signaled_tail - old_signaled_tail, lane_idx); + stored_old_scaleout_tail = stored_scaleout_tail; + } + __syncwarp(); + }; + + // Preload next token + const auto preload_next_token = [&](const int& token_idx) { + if (token_idx >= num_tokens) + return; + + // Issue TMA load + const auto token_i64_idx = static_cast(token_idx); + if (ptx::elect_one_sync()) { + ptx::tma_load_1d(tma_buffer.get_hidden_ptr(), math::advance_ptr(x, token_i64_idx * kNumHiddenBytes), + mbarrier_ptr, kNumHiddenBytes); + } + __syncwarp(); + + // Issue SF `cp.async` + if constexpr (kNumSFPacks > 0) { + EP_STATIC_ASSERT(sizeof(sf_pack_t) % 4 == 0, "Unaligned SF element type"); + const auto gmem_src_ptr = math::advance_ptr(sf, token_i64_idx * sf_token_stride * sizeof(sf_pack_t)); + const auto smem_dst_ptr = tma_buffer.get_sf_ptr(); + + constexpr auto kNumFullIters = kNumSFPacks / 32; + #pragma unroll + for (int k = 0; k < kNumFullIters; ++ k) { + ptx::cp_async_ca(gmem_src_ptr + (k * 32 + lane_idx) * sf_hidden_stride, + smem_dst_ptr + k * 32 + lane_idx); + } + if (kNumFullIters * 32 + lane_idx < kNumSFPacks) { + ptx::cp_async_ca(gmem_src_ptr + (kNumFullIters * 32 + lane_idx) * sf_hidden_stride, + smem_dst_ptr + kNumFullIters * 32 + lane_idx); + } + ptx::cp_async_mbarrier_arrive(mbarrier_ptr); + __syncwarp(); + } + }; + + // Iterate all tokens + preload_next_token(channel_idx); + for (int token_idx = channel_idx; token_idx < num_tokens; token_idx += kNumChannels) { + // Load top-k indices and weights + EP_STATIC_ASSERT(kNumTopk <= 32, "Insufficient lanes for loading top-k indices"); + int stored_dst_scaleout_rank_idx = -1; + if (lane_idx < kNumTopk) { + const auto uncasted_dst_expert_idx = __ldg(topk_idx + token_idx * kNumTopk + lane_idx); + const auto dst_expert_idx = static_cast(uncasted_dst_expert_idx); + stored_dst_scaleout_rank_idx = dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerScaleout : -1; + tma_buffer.get_topk_idx_ptr()[lane_idx] = dst_expert_idx; + if (topk_weights != nullptr) + tma_buffer.get_topk_weights_ptr()[lane_idx] = __ldg(topk_weights + token_idx * kNumTopk + lane_idx); + if (copied_topk_idx != nullptr) + copied_topk_idx[token_idx * kNumTopk + lane_idx] = uncasted_dst_expert_idx; + } + __syncwarp(); + + // Add source metadata (rank index and token index) + if (ptx::elect_one_sync()) + *tma_buffer.get_src_token_global_idx_ptr() = rank_idx * kNumMaxTokensPerRank + token_idx; + ptx::tma_store_fence(); + __syncwarp(); + + // Deduplicate ranks and assign slots + int stored_dst_slot_idx = -1; + const auto stored_old_slot_idx = ptx::exchange( + stored_scaleout_tail, stored_dst_scaleout_rank_idx >= 0 ? stored_dst_scaleout_rank_idx : 0); + if (ptx::deduplicate(stored_dst_scaleout_rank_idx, lane_idx) and stored_dst_scaleout_rank_idx >= 0) + stored_dst_slot_idx = stored_old_slot_idx; + + // Update scale-out tail + const auto scaleout_rank_mask = ptx::reduce_or(stored_dst_scaleout_rank_idx >= 0 ? (1u << stored_dst_scaleout_rank_idx) : 0u); + stored_scaleout_tail += (scaleout_rank_mask >> lane_idx) & 1; + + // Wait TMA arrival and issue the TMA store into send buffer + if (ptx::elect_one_sync()) { + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + + // So if no ranks will go by RDMA, we skip the send buffer stores + if (scaleout_rank_mask ^ (1 << scaleout_rank_idx)) { + ptx::tma_store_1d(scaleout_send_buffer.get_token_buffer(token_idx).get_base_ptr(), + tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes()); + } + } + __syncwarp(); + + // Local rank can be bypassed + if (stored_dst_slot_idx >= 0 and stored_dst_scaleout_rank_idx == scaleout_rank_idx) { + ptx::tma_store_1d(scaleout_recv_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(), + tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes()); + } + ptx::tma_store_commit(); + ptx::tma_store_wait(); + __syncwarp(); + + // Preload the next token (overlapping with the IBGDA issues) + preload_next_token(token_idx + kNumChannels); + + // Issue IBGDA requests + if (stored_dst_slot_idx >= 0 and stored_dst_scaleout_rank_idx != scaleout_rank_idx) { + gin.put( + scaleout_recv_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(), + scaleout_send_buffer.get_token_buffer(token_idx).get_base_ptr(), + tma_buffer.get_num_bytes(), + stored_dst_scaleout_rank_idx, + ncclGinOptFlagsAggregateRequests); + } + __syncwarp(); + + // Issue scale-out tail update + update_scaleout_tail(); + } + + // Flush unflushed tails + update_scaleout_tail(true); + } else { + const int forward_warp_idx = warp_idx - (kNumNotifyWarps + kNumScaleoutWarps); + const int channel_idx = sm_idx * kNumChannelsPerSM + forward_warp_idx; + scaleout_recv_buffer = scaleout_recv_buffer.get_channel_buffer(channel_idx); + scaleup_buffer = scaleup_buffer.get_rank_buffer(scaleup_rank_idx); + + // Shape of `token_metadata_at_forward`: `[kNumChannels, kNumScaleoutRanks * kNumMaxTokensPerChannel + 1, kNumForwardMetadataDims]` + constexpr int kNumForwardMetadataDims = 2 + kNumTopk * 2; + token_metadata_at_forward += channel_idx * ((kNumScaleoutRanks * kNumMaxTokensPerChannel + 1) * kNumForwardMetadataDims); + + // Shape of `dst_buffer_slot_idx`: `[kNumChannels, kNumScaleoutRanks, kNumMaxTokensPerChannel, kNumTopk]` + dst_buffer_slot_idx += channel_idx * (kNumScaleoutRanks * kNumMaxTokensPerChannel * kNumTopk); + + // Transform linked list index + const auto transform_linked_list_idx = [=](const int& idx) { + constexpr int kNumTokensInLinkedList = kNumMaxTokensPerChannel * kNumScaleoutRanks + 1; + return channel_idx * (kNumTokensInLinkedList * kNumScaleupRanks) + + idx * kNumScaleupRanks + scaleup_rank_idx; + }; + + // Forward tokens from scale-out ranks + EP_STATIC_ASSERT(kNumScaleoutRanks <= 32, "Too many scale-out ranks"); + int num_tokens_processed = 0; + int stored_scaleout_old_tail_idx = 0; + int stored_scaleup_send_counters[kNumScaleupRanksPerLane] = {}; + int stored_finish_flag = lane_idx >= kNumScaleoutRanks; + int stored_scaleout_tail_idx = 0; + int recv_scaleout_rank_idx = channel_idx % kNumScaleoutRanks; + uint32_t wip_mask; + while ((wip_mask = ptx::gather(stored_scaleout_tail_idx > stored_scaleout_old_tail_idx or stored_finish_flag == 0))) { + // Pick next rank in round-robin + const auto offset = (recv_scaleout_rank_idx + 1) % kNumScaleoutRanks; + const auto hi_mask = (wip_mask >> offset) << offset; + recv_scaleout_rank_idx = hi_mask ? ptx::ffs(hi_mask) : ptx::ffs(wip_mask); + + // Wait for this rank to have data (or finish) + comm::timeout_while([&](const bool& is_last_check) { + const uint32_t arrived_or_finished = + stored_scaleout_tail_idx > stored_scaleout_old_tail_idx or stored_finish_flag > 0; + if (ptx::exchange(arrived_or_finished, recv_scaleout_rank_idx)) + return true; + + // Timeout + if (is_last_check) { + if (lane_idx < kNumScaleoutRanks) { + printf("DeepEP hybrid dispatch (forwarding) timeout, scale-out: %d, scale-up: %d, " + "channel: %d, lane: %d, old scale-out tail: %d, scale-out tail: (%d, %d)\n", + scaleout_rank_idx, scaleup_rank_idx, + channel_idx, lane_idx, stored_scaleout_old_tail_idx, + stored_finish_flag, stored_scaleout_tail_idx); + } + return false; + } + + // Read new signaled tails + if (lane_idx < kNumScaleoutRanks) { + const auto signaled_tail = ptx::ld_acquire_sys( + workspace_layout.get_scaleout_channel_signaled_tail_ptr(channel_idx, lane_idx)); + math::unpack2(signaled_tail, stored_finish_flag, stored_scaleout_tail_idx); + } + __syncwarp(); + return false; + }); + + // Process one chunk from the current rank + const auto start_slot_idx = ptx::exchange(stored_scaleout_old_tail_idx, recv_scaleout_rank_idx); + const auto end_slot_idx = std::min( + ptx::exchange(stored_scaleout_tail_idx, recv_scaleout_rank_idx), + start_slot_idx + kNumSlotsPerForwardChunk + ); + if (lane_idx == recv_scaleout_rank_idx) + stored_scaleout_old_tail_idx = end_slot_idx; + + const auto recv_buffer = scaleout_recv_buffer.get_rank_buffer(recv_scaleout_rank_idx); + for (int slot_idx = start_slot_idx; slot_idx < end_slot_idx; ++ slot_idx) { + const auto token_buffer = recv_buffer.get_token_buffer(slot_idx); + + // Wait TMA arrival + ptx::tma_store_wait(); + __syncwarp(); + + // TMA load into shared memory + if (ptx::elect_one_sync()) { + ptx::tma_load_1d(tma_buffer.get_base_ptr(), token_buffer.get_base_ptr(), + mbarrier_ptr, token_layout.get_num_bytes()); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, token_layout.get_num_bytes()); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + } + __syncwarp(); + + // Read top-k indices + EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k selections"); + int stored_dst_scaleup_rank_idx = -1; + auto dst_expert_idx = lane_idx < kNumTopk ? tma_buffer.get_topk_idx_ptr()[lane_idx] : -1; + dst_expert_idx -= scaleout_rank_idx * kNumExpertsPerScaleout; + stored_dst_scaleup_rank_idx = 0 <= dst_expert_idx and dst_expert_idx < kNumExpertsPerScaleout ? + dst_expert_idx / kNumExpertsPerRank : -1; + + // Write the per-scaleup channel index for this token + int linked_list_idx = -1; + #pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) { + const auto src_lane_idx = stored_dst_scaleup_rank_idx - j * 32; + const bool valid = 0 <= src_lane_idx and src_lane_idx < 32; + const auto exchanged = ptx::exchange( + stored_scaleup_send_counters[j], valid ? src_lane_idx : 0); + linked_list_idx = valid ? exchanged : linked_list_idx; + } + if (not kReuseSlotIndices and lane_idx < kNumTopk) { + tma_buffer.get_linked_list_idx_ptr()[lane_idx] = transform_linked_list_idx(linked_list_idx); + ptx::tma_store_fence(); + } + __syncwarp(); + + // Deduplicate for scale-up ranks + int stored_dst_slot_idx = -1; + const auto dst_slot_idx_ptr = dst_buffer_slot_idx + + recv_scaleout_rank_idx * (kNumMaxTokensPerChannel * kNumTopk) + slot_idx * kNumTopk; + if constexpr (kReuseSlotIndices) { + if (lane_idx < kNumTopk) + stored_dst_slot_idx = __ldg(dst_slot_idx_ptr + lane_idx); + } else { + // Deduplicate for NVLink ranks + if (ptx::deduplicate(stored_dst_scaleup_rank_idx, lane_idx) and stored_dst_scaleup_rank_idx >= 0) + stored_dst_slot_idx = atomicAdd(workspace_layout.get_scaleup_atomic_sender_counter() + stored_dst_scaleup_rank_idx, 1); + } + __syncwarp(); + + // Issue TMAs + if (stored_dst_slot_idx >= 0) { + const auto dst_ptr = gin.get_sym_ptr( + scaleup_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(), + stored_dst_scaleup_rank_idx); + ptx::tma_store_1d(dst_ptr, tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes()); + ptx::tma_store_commit(); + } + __syncwarp(); + + // Add per-scale-up counter + EP_STATIC_ASSERT(kNumScaleupRanks <= 64, "Invalid number of scale-up peers"); + using mask_t = std::conditional_t; + const auto scaleup_send_mask = ptx::reduce_or( + stored_dst_scaleup_rank_idx >= 0 ? + (mask_t(1) << stored_dst_scaleup_rank_idx) : mask_t(0)); + #pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) + stored_scaleup_send_counters[j] += (scaleup_send_mask >> (j * 32 + lane_idx)) & 1; + + // Record metadata at forward + if constexpr (not kReuseSlotIndices) { + EP_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of selections"); + const auto metadata_ptr = token_metadata_at_forward + + num_tokens_processed * kNumForwardMetadataDims; + + // Source token index and last token index flag + if (ptx::elect_one_sync()) { + metadata_ptr[0] = tma_buffer.get_src_token_global_idx_ptr()[0]; + metadata_ptr[1] = slot_idx == (end_slot_idx - 1); + } + + // Second, original top-k indices and destination slots + if (lane_idx < kNumTopk) { + metadata_ptr[2 + lane_idx] = stored_dst_scaleup_rank_idx; + metadata_ptr[2 + kNumTopk + lane_idx] = stored_dst_slot_idx; + dst_slot_idx_ptr[lane_idx] = stored_dst_slot_idx; + } + } + num_tokens_processed += 1; + __syncwarp(); + } + } + + // Assign the source token index part of the metadata into `-1` as an ending mark + if (not kReuseSlotIndices and ptx::elect_one_sync()) + token_metadata_at_forward[num_tokens_processed * kNumForwardMetadataDims] = -1; + __syncwarp(); + + // Update linked list's ending position + if constexpr (not kReuseSlotIndices) { + const auto tail_ptr = workspace_layout.get_channel_scaleup_tail_ptr(channel_idx, scaleup_rank_idx); + #pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++ i) { + if (const auto j = i * 32 + lane_idx; i < (kNumScaleupRanksPerLane - 1) or j < kNumScaleupRanks) { + ptx::st_relaxed_sys( + gin.get_sym_ptr(tail_ptr, j), + transform_linked_list_idx(stored_scaleup_send_counters[i])); + } + } + } + __syncwarp(); + + // Clean tails for next usages + if (lane_idx < kNumScaleoutRanks) + *workspace_layout.get_scaleout_channel_signaled_tail_ptr(channel_idx, lane_idx) = 0; + __syncwarp(); + } + + // Scale-up barrier to ensure data arrival + // As scale-out tokens have already been consumed by forwarders, no need to do scale-out barrier again + comm::gpu_barrier( + gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx, /* do not scale-out */ false, true); + + // Trigger the copy epilogue kernel + cudaTriggerProgrammaticLaunchCompletion(); + + // Clean scale-up counters + // All scale-out counters should be cleaned before + EP_STATIC_ASSERT(kNumScaleupRanks <= kNumThreads, "Insufficient threads"); + if (not kReuseSlotIndices and sm_idx == 0 and thread_idx < kNumScaleupRanks) + workspace_layout.get_scaleup_atomic_sender_counter()[thread_idx] = 0; +} + +} // namespace deep_ep::elastic diff --git a/deep_ep/include/deep_ep/impls/pp_send_recv.cuh b/deep_ep/include/deep_ep/impls/pp_send_recv.cuh new file mode 100644 index 000000000..4d42260be --- /dev/null +++ b/deep_ep/include/deep_ep/impls/pp_send_recv.cuh @@ -0,0 +1,213 @@ +#pragma once + +#include +#include +#include +#include + + +namespace deep_ep::elastic { + +template +__device__ __forceinline__ std::pair get_buffer_offset( + const int& src_rank_idx, const int& dst_rank_idx) { + const auto next_rank_idx = (src_rank_idx + 1) % kNumRanks; + return dst_rank_idx == next_rank_idx ? std::make_pair(0, 1) : std::make_pair(1, 0); +} + +template +__device__ __forceinline__ void check_signal( + const handle::NCCLGin& gin, + const ncclGinSignal_t& signal_idx, + const int64_t& target, + const timeout_print_t& timeout_print) { + const auto gdaki = static_cast(gin.gin._ginHandle) + gin.gin.contextId; + const auto signal_ptr = reinterpret_cast( + __ldg(reinterpret_cast(&gdaki->signals_table.buffer))) + signal_idx; + comm::timeout_while([=](const bool& is_last_check) { + const auto signal = ptx::ld_acquire_sys(signal_ptr); + if (signal >= target) + return true; + + if (is_last_check) + timeout_print(); + return false; + }); +} + +template ( + (kNumSmemBytes - kNumStages * sizeof(ptx::mbarrier)) / kNumStages, ptx::kNumTMAAlignBytes), + int kNumTMABlocksPerStage = kNumTMABytesPerStage / ptx::kNumTMAAlignBytes> +__device__ __forceinline__ void tma_copy( + void* src_ptr, void* dst_ptr, + const int64_t& num_bytes, const int& sm_idx) { + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + const auto tma_buffers = smem; + const auto mbarriers = reinterpret_cast(smem + kNumStages * kNumTMABytesPerStage); + EP_STATIC_ASSERT(kNumTMABytesPerStage > 0, "Invalid shared memory bytes"); + EP_STATIC_ASSERT(kNumStages >= 2, "Need at least 2 stages for pipelining"); + + // Init mbarriers + ptx::arrival_phase phases[kNumStages]; + #pragma unroll + for (int s = 0; s < kNumStages; ++ s) + phases[s] = 0, ptx::mbarrier_init_with_fence(mbarriers + s, 1); + + // Work partitioning across SMs + EP_DEVICE_ASSERT(num_bytes % ptx::kNumTMAAlignBytes == 0); + const auto num_tma_blocks = num_bytes / ptx::kNumTMAAlignBytes; + const auto num_tma_blocks_per_sm = math::ceil_div(num_tma_blocks, kNumSMs); + const auto start_block_idx = sm_idx * num_tma_blocks_per_sm; + const auto end_block_idx = std::min(start_block_idx + num_tma_blocks_per_sm, num_tma_blocks); + const auto num_iterations = math::ceil_div(end_block_idx - start_block_idx, kNumTMABlocksPerStage); + + auto get_iter_info = [&](const int64_t& iter_idx) { + const auto i = start_block_idx + iter_idx * kNumTMABlocksPerStage; + const auto offset = i * ptx::kNumTMAAlignBytes; + const auto num_transaction_bytes = + std::min(kNumTMABlocksPerStage, end_block_idx - i) * ptx::kNumTMAAlignBytes; + return std::make_pair(offset, num_transaction_bytes); + }; + + for (int64_t iter_idx = 0; iter_idx < num_iterations; ++ iter_idx) { + const auto stage_idx = static_cast(iter_idx % kNumStages); + const auto [store_offset, num_store_bytes] = get_iter_info(iter_idx); + + // Fill pipeline: issue loads for the first kNumStages iterations + if (iter_idx < kNumStages) { + ptx::tma_load_1d( + tma_buffers + stage_idx * kNumTMABytesPerStage, + math::advance_ptr(src_ptr, store_offset), + mbarriers + stage_idx, num_store_bytes); + ptx::mbarrier_arrive_and_set_tx(mbarriers + stage_idx, num_store_bytes); + } + + // Wait for this stage's load, then store + ptx::mbarrier_wait_and_flip_phase(mbarriers + stage_idx, phases[stage_idx]); + ptx::tma_store_1d( + math::advance_ptr(dst_ptr, store_offset), + tma_buffers + stage_idx * kNumTMABytesPerStage, + num_store_bytes); + ptx::tma_store_commit(); + + // Prefetch: wait until this stage's buffer is safe to reuse, then issue next load + const auto next_iter_idx = iter_idx + kNumStages; + if (next_iter_idx < num_iterations) { + ptx::tma_store_wait(); + const auto [load_offset, num_load_bytes] = get_iter_info(next_iter_idx); + ptx::tma_load_1d( + tma_buffers + stage_idx * kNumTMABytesPerStage, + math::advance_ptr(src_ptr, load_offset), + mbarriers + stage_idx, num_load_bytes); + ptx::mbarrier_arrive_and_set_tx(mbarriers + stage_idx, num_load_bytes); + } + } + + // Drain all outstanding stores + ptx::tma_store_wait(); +} + +template +__global__ void __launch_bounds__(32, 1) +pp_send_impl(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 int64_t num_max_tensor_bytes, + const int num_max_inflight_tensors) { + const auto sm_idx = static_cast(blockIdx.x); + const auto workspace_layout = layout::WorkspaceLayout(workspace, 1, kNumRanks, 0); + const auto [local_idx_in_dst, dst_idx_in_local] = get_buffer_offset(rank_idx, dst_rank_idx); + + // Gin handle + const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, 0, NCCL_GIN_RESOURCE_SHARING_CTA); + + // Buffer offsets + const auto send_count_ptr = workspace_layout.get_pp_send_count_ptr(dst_idx_in_local); + const auto send_count = __ldg(send_count_ptr); + const auto slot_idx = send_count % num_max_inflight_tensors; + auto send_buffer_ptr = math::advance_ptr( + buffer, ((dst_idx_in_local + 2) * num_max_inflight_tensors + slot_idx) * num_max_tensor_bytes); + auto recv_buffer_ptr = math::advance_ptr( + buffer, ((local_idx_in_dst + 0) * num_max_inflight_tensors + slot_idx) * num_max_tensor_bytes); + + // Wait buffer slot release and do TMA + if (ptx::elect_one_sync()) { + check_signal( + gin, + static_cast(kNumRanks + dst_idx_in_local + 2), + send_count - num_max_inflight_tensors + 1, + // TODO: print more info, and control the SM who prints it + []() { printf("DeepEP PP send timeout, recv buffer is full"); } + ); + tma_copy(x, send_buffer_ptr, num_x_bytes, sm_idx); + } + cooperative_groups::this_grid().sync(); + + // Issue RDMA put + if (sm_idx == 0 and ptx::elect_one_sync()) { + gin.put( + recv_buffer_ptr, + send_buffer_ptr, + num_x_bytes, dst_rank_idx, + 0, + // TODO: is this signal highly optimized? + ncclGin_SignalInc(static_cast(local_idx_in_dst + kNumRanks))); + *send_count_ptr += 1; + } +} + +template +__global__ void __launch_bounds__(32, 1) +pp_recv_impl(const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window, + void* x, int64_t num_x_bytes, + void* buffer, void* workspace, + const int rank_idx, const int src_rank_idx, + const int64_t num_max_tensor_bytes, + const int num_max_inflight_tensors) { + const auto sm_idx = static_cast(blockIdx.x); + const auto workspace_layout = layout::WorkspaceLayout(workspace, 1, kNumRanks, 0); + const auto [src_idx_in_local, local_idx_in_src] = get_buffer_offset(src_rank_idx, rank_idx); + + // Gin handle + const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, 0, NCCL_GIN_RESOURCE_SHARING_CTA); + + // Buffer offsets + const auto recv_count_ptr = workspace_layout.get_pp_recv_count_ptr(src_idx_in_local); + const auto recv_count = __ldg(recv_count_ptr); + const auto slot_idx = recv_count % num_max_inflight_tensors; + const auto recv_buffer_ptr = math::advance_ptr( + buffer, ((src_idx_in_local + 0) * num_max_inflight_tensors + slot_idx) * num_max_tensor_bytes); + + // Copy from the buffer into a new tensor + if (ptx::elect_one_sync()) { + check_signal( + gin, + static_cast(src_idx_in_local + kNumRanks), + recv_count + 1, + // TODO: print more info, and control the SM who prints it + []() { printf("DeepEP PP recv timeout, recv buffer is empty\n"); } + ); + tma_copy(recv_buffer_ptr, x, num_x_bytes, sm_idx); + } + cooperative_groups::this_grid().sync(); + + // TODO: add a comment + if (sm_idx == 0 and ptx::elect_one_sync()) { + gin.signal( + src_rank_idx, ncclGin_SignalInc(static_cast(kNumRanks + local_idx_in_src + 2)) + ); + *recv_count_ptr += 1; + } +} + +} // namespace deep_ep::elastic diff --git a/deep_ep/utils.py b/deep_ep/utils.py deleted file mode 100644 index e61a2c5b7..000000000 --- a/deep_ep/utils.py +++ /dev/null @@ -1,101 +0,0 @@ -import os -import torch -import torch.distributed as dist -from typing import Any, Optional, Tuple - -# noinspection PyUnresolvedReferences -from deep_ep_cpp import EventHandle - - -class EventOverlap: - """ - A wrapper class to manage CUDA events, also for better overlapping convenience. - - Attributes: - event: the CUDA event captured. - extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph. - """ - - def __init__(self, event: Optional[EventHandle] = None, extra_tensors: Optional[Tuple[torch.Tensor]] = None) -> None: - """ - Initialize the class. - - Arguments: - event: the CUDA event captured. - extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph. - """ - self.event = event - - # NOTES: we use extra tensors to achieve stream recording, otherwise, - # stream recording will be incompatible with CUDA graph. - self.extra_tensors = extra_tensors - - def current_stream_wait(self) -> None: - """ - The current stream `torch.cuda.current_stream()` waits for the event to be finished. - """ - assert self.event is not None - self.event.current_stream_wait() - - def __enter__(self) -> Any: - """ - Utility for overlapping and Python `with` syntax. - - You can overlap the kernels on the current stream with the following example: - ```python - event_overlap = event_after_all_to_all_kernels() - with event_overlap(): - do_something_on_current_stream() - # After exiting the `with` scope, the current stream with wait the event to be finished. - ``` - """ - return self - - def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: - """ - Utility for overlapping and Python `with` syntax. - - Please follow the example in the `__enter__` function. - """ - if self.event is not None: - self.event.current_stream_wait() - - -def check_nvlink_connections(group: dist.ProcessGroup): - """ - Check NVLink connection between every pair of GPUs. - - Arguments: - group: the communication group. - """ - # Check NVLink connection - # NOTES: some A100 PCIE GPUs only have pairwise NVLink connection, so that we can only use EP2 - # TODO: check all cases, all local-node GPUs in the group should be connected via NVLink - if 'PCIE' in torch.cuda.get_device_name(): - assert group.size() <= 2, 'PCIe GPUs only have pairwise NVLink connections' - - # noinspection PyUnresolvedReferences - import pynvml - pynvml.nvmlInit() - - # noinspection PyTypeChecker - devices = os.environ.get('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7').strip(',').split(',') - physical_device_idx = int(devices[torch.cuda.current_device()]) - physical_device_indices = [ - 0, - ] * group.size() - dist.all_gather_object(physical_device_indices, physical_device_idx, group) - - # Check whether they are all connected via NVLink - # Reference: https://github.com/vllm-project/vllm/blob/b8e809a057765c574726a6077fd124db5077ce1f/vllm/platforms/cuda.py#L438 - handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_indices] - for i, handle in enumerate(handles): - for j, peer_handle in enumerate(handles): - if i >= j: - continue - status = pynvml.nvmlDeviceGetP2PStatus(handle, peer_handle, pynvml.NVML_P2P_CAPS_INDEX_NVLINK) - assert status == pynvml.NVML_P2P_STATUS_OK,\ - f'GPU {physical_device_indices[i]} and GPU {physical_device_indices[j]} are not connected via NVLink' - - # Close NVML - pynvml.nvmlShutdown() diff --git a/deep_ep/utils/__init__.py b/deep_ep/utils/__init__.py new file mode 100644 index 000000000..9d4b6fb9f --- /dev/null +++ b/deep_ep/utils/__init__.py @@ -0,0 +1,3 @@ +# For forward compatibility +# noinspection PyUnresolvedReferences +from .event import EventHandle diff --git a/deep_ep/utils/comm.py b/deep_ep/utils/comm.py new file mode 100644 index 000000000..6cec159d1 --- /dev/null +++ b/deep_ep/utils/comm.py @@ -0,0 +1,79 @@ +import os + +import torch +import torch.distributed as dist + +# noinspection PyUnresolvedReferences +import deep_ep._C as _C + + +class NCCLCommHandle: + """ + A wrapper around a raw NCCL communicator. Manages the lifecycle of the communicator if created by DeepEP, + or simply wraps an existing one if obtained from PyTorch. + + Attributes: + nccl_comm: the raw NCCL communicator. + managed: whether the communicator was created by DeepEP and should be destroyed when this handle is dropped. + """ + + def __init__(self, nccl_comm: int, managed: bool): + self.nccl_comm = nccl_comm + self.managed = managed + self.destroy = _C.destroy_nccl_comm + + def __del__(self): + if self.managed: + self.destroy(self.nccl_comm) + + def get(self) -> int: + """ + Get the raw NCCL communicator. + + Returns: + nccl_comm: the raw NCCL communicator. + """ + return self.nccl_comm + +_storage = dict() + + +def get_nccl_comm_handle(group: dist.ProcessGroup) -> NCCLCommHandle: + """ + Get or create an NCCL communicator handle for the given process group. + Results are cached, so subsequent calls with the same group return the same handle. + + Arguments: + group: the communication group. + + Returns: + handle: the NCCL communicator handle. + """ + # Check cache hit + global _storage + if group in _storage: + return _storage[group] + + # New PyTorch has such API + backend = group._get_backend(torch.device('cuda')) + if hasattr(backend, '_comm_ptr') and int(os.getenv('EP_REUSE_NCCL_COMM', '0')): + _storage[group] = NCCLCommHandle(backend._comm_ptr(), False) + return _storage[group] + + # For old PyTorch, we have to recreate a NCCL comm + nccl_unique_ids = [None, ] * group.size() + dist.all_gather_object(nccl_unique_ids, _C.get_local_nccl_unique_id(), group) + root_unique_id = nccl_unique_ids[0] + + # Create a new communicator + _storage[group] = NCCLCommHandle( + _C.create_nccl_comm(root_unique_id, group.size(), group.rank()), True) + return _storage[group] + + +def destroy_all_managed_nccl_comm() -> None: + """ + Destroy all cached NCCL communicator handles and clear the cache. + + """ + _storage.clear() diff --git a/deep_ep/utils/envs.py b/deep_ep/utils/envs.py new file mode 100644 index 000000000..f6e34d988 --- /dev/null +++ b/deep_ep/utils/envs.py @@ -0,0 +1,268 @@ +import functools +import inspect +import os +import random +import re +import subprocess +import torch +import torch.distributed as dist +from typing import Tuple + +# noinspection PyUnresolvedReferences +import deep_ep._C as _C + +from .comm import get_nccl_comm_handle + +_local_rank = None +_local_seed = 0 +_global_seed = 0 + +# Default NIC name for RDMA operations, configurable via environment variable +_DEFAULT_NIC_NAME = os.getenv('EP_NIC_NAME', 'mlx5_0') + + +def init_seed(global_seed: int) -> None: + """ + Initialize the random seed for reproducibility. The local seed is derived from the global seed plus rank. + + Arguments: + global_seed: the global random seed. + """ + global _local_seed, _global_seed + _local_seed = global_seed + dist.get_rank() + _global_seed = global_seed + torch.manual_seed(_local_seed) + random.seed(_local_seed) + + +def get_local_seed() -> int: + """ + Get the local random seed. + + Returns: + seed: the local random seed. + """ + return _local_seed + + +def get_global_seed() -> int: + """ + Get the global random seed. + + Returns: + seed: the global random seed. + """ + return _global_seed + + +def dist_print(s: str = '', once_in_node: bool = False) -> None: + """ + Print a message from all ranks, or only from rank 0 of each node, followed by a barrier. + + Arguments: + s: the message to print. + once_in_node: if `True`, only the first local rank in each node prints. + """ + global _local_rank + assert _local_rank is not None + if not once_in_node or _local_rank == 0: + print(s, flush=True) + dist.barrier() + + +def init_dist(local_rank: int, num_local_ranks: int, seed: int = 0) -> Tuple[int, int, dist.ProcessGroup]: + """ + Initialize the distributed environment with NCCL backend. + + Arguments: + local_rank: the local rank index. + num_local_ranks: the number of local ranks. + seed: the global random seed. + + Returns: + rank: the global rank index. + world_size: the total number of ranks. + group: the communication group. + """ + # NOTES: you may rewrite this function with your own cluster settings + ip = os.getenv('MASTER_ADDR', '127.0.0.1') + port = int(os.getenv('MASTER_PORT', '8361')) + num_nodes = int(os.getenv('WORLD_SIZE', 1)) + node_rank = int(os.getenv('RANK', 0)) + + # Set local rank + global _local_rank + _local_rank = local_rank + + sig = inspect.signature(dist.init_process_group) + params = { + 'backend': 'nccl', + 'init_method': f'tcp://{ip}:{port}', + 'world_size': num_nodes * num_local_ranks, + 'rank': node_rank * num_local_ranks + local_rank, + } + if 'device_id' in sig.parameters: + # noinspection PyTypeChecker + params['device_id'] = torch.device(f'cuda:{local_rank}') + dist.init_process_group(**params) + torch.set_default_dtype(torch.bfloat16) + torch.set_default_device('cuda') + torch.cuda.set_device(local_rank) + + init_seed(seed) + return dist.get_rank(), dist.get_world_size(), dist.new_group(list(range(num_local_ranks * num_nodes))) + + +def get_physical_domain_size(group: dist.ProcessGroup) -> Tuple[int, int]: + """ + Get the physical domain sizes (RDMA ranks and NVLink ranks). + + Arguments: + group: the communication group. + + Returns: + num_rdma_ranks: the number of physical RDMA ranks. + num_nvlink_ranks: the number of physical NVLink ranks. + """ + return _C.get_physical_domain_size(get_nccl_comm_handle(group).get()) + + +def get_logical_domain_size(group: dist.ProcessGroup, allow_hybrid_mode: bool = True) -> Tuple[int, int]: + """ + Get the logical domain sizes (scaleout ranks and scaleup ranks). + + Arguments: + group: the communication group. + allow_hybrid_mode: whether to enable hybrid mode. + + Returns: + num_scaleout_ranks: the number of logical scaleout ranks. + num_scaleup_ranks: the number of logical scaleup ranks. + """ + return _C.get_logical_domain_size(get_nccl_comm_handle(group).get(), allow_hybrid_mode) + + +def check_nvlink_connections(group: dist.ProcessGroup) -> None: + """ + Check NVLink connection between every pair of GPUs. + + Arguments: + group: the communication group. + """ + # Check NVLink connection + # NOTES: some A100 PCIE GPUs only have pairwise NVLink connection, so that we can only use EP2 + # TODO: check all cases, all local-node GPUs in the group should be connected via NVLink + if 'PCIE' in torch.cuda.get_device_name(): + assert group.size() <= 2, 'PCIe GPUs only have pairwise NVLink connections' + + # noinspection PyUnresolvedReferences + import pynvml + pynvml.nvmlInit() + + # noinspection PyTypeChecker + devices = os.environ.get('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7').strip(',').split(',') + physical_device_idx = int(devices[torch.cuda.current_device()]) + physical_device_indices = [0, ] * group.size() + dist.all_gather_object(physical_device_indices, physical_device_idx, group) + + # Check whether they are all connected via NVLink + # Reference: https://github.com/vllm-project/vllm/blob/b8e809a057765c574726a6077fd124db5077ce1f/vllm/platforms/cuda.py#L438 + handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_indices] + for i, handle in enumerate(handles): + for j, peer_handle in enumerate(handles): + if i >= j: + continue + status = pynvml.nvmlDeviceGetP2PStatus(handle, peer_handle, pynvml.NVML_P2P_CAPS_INDEX_NVLINK) + assert status == pynvml.NVML_P2P_STATUS_OK, \ + f'GPU {physical_device_indices[i]} and GPU {physical_device_indices[j]} are not connected via NVLink' + + # Close NVML + pynvml.nvmlShutdown() + + +def check_torch_deterministic() -> None: + """ + Ensure PyTorch deterministic algorithms and fill_uninitialized_memory are not both enabled. + When both are on, `torch.empty()` calls an initialization kernel that may overlap with communication streams, + causing errors. + """ + assert not (torch.are_deterministic_algorithms_enabled() and torch.utils.deterministic.fill_uninitialized_memory) + + +@functools.lru_cache() +def get_nvlink_gbs(factor: float = 0.9) -> float: + """ + Get the total NVLink bandwidth in GB/s, cached. + + Arguments: + factor: the bandwidth efficiency factor. + + Returns: + gbs: the total NVLink bandwidth in GB/s (0 if detection fails). + """ + # noinspection PyBroadException + try: + result = subprocess.run(['nvidia-smi', 'nvlink', '-s'], + capture_output=True, text=True, check=True) + output = result.stdout + pattern = r'GPU \d+:.*?(?=^GPU \d+:|^$)' + match = re.search(pattern, output, re.MULTILINE | re.DOTALL) + assert match + + gpu_block = match.group(0) + link_pattern = r'Link \d+:\s*([\d\.]+) GB/s' + link_matches = re.findall(link_pattern, gpu_block) + assert link_matches + return sum(float(bw) for bw in link_matches) * factor + except Exception as e: + print(f'Failed to get NVLink connection speed: {e}') + return 0 + + +@functools.lru_cache() +def check_fast_rdma_atomic_support(nic_name: str = _DEFAULT_NIC_NAME) -> bool: + """ + Check whether the NIC supports fast RDMA atomic operations (MT4131 or newer). + + Arguments: + nic_name: the NIC device name. + + Returns: + supported: `True` if fast RDMA atomics are supported. + """ + # noinspection PyBroadException + try: + result = subprocess.run(['ibstat'], capture_output=True, text=True, check=True) + output = result.stdout + pattern = rf"CA '{nic_name}'.*?CA type:\s*(\S+)" + match = re.search(pattern, output, re.DOTALL) + assert match + return match.group(1) == 'MT4131' + except Exception: + return False + + +@functools.lru_cache() +def get_rdma_gbs(nic_name: str = _DEFAULT_NIC_NAME) -> float: + """ + Get the RDMA bandwidth in GB/s, cached. + + Arguments: + nic_name: the NIC device name. + + Returns: + gbs: the RDMA bandwidth in GB/s (0 if detection fails). + """ + # noinspection PyBroadException + try: + result = subprocess.run(['ibstat'], capture_output=True, text=True, check=True) + output = result.stdout + + pattern = rf"CA '{nic_name}'.*?Port \d+:\s*.*?Rate:\s*(\d+)" + match = re.search(pattern, output, re.DOTALL) + assert match + rate = int(match.group(1)) + return rate / 8 + except Exception as e: + print(f'Failed to get RDMA connection speed: {e}') + return 0 diff --git a/deep_ep/utils/event.py b/deep_ep/utils/event.py new file mode 100644 index 000000000..9e9334b13 --- /dev/null +++ b/deep_ep/utils/event.py @@ -0,0 +1,81 @@ +import torch +from typing import Any, Optional, Tuple + +# noinspection PyUnresolvedReferences +from deep_ep._C import EventHandle + + +class EventOverlap: + """ + A wrapper class to manage CUDA events, also for better overlapping convenience. + + Attributes: + event: the CUDA event captured. + extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph. + """ + + def __init__(self, event: Optional[EventHandle] = None, extra_tensors: Optional[Tuple[torch.Tensor]] = None) -> None: + """ + Initialize the class. + + Arguments: + event: the CUDA event captured. + extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph. + """ + self.event = event + + # NOTES: we use extra tensors to achieve stream recording, otherwise, + # stream recording will be incompatible with CUDA graph. + # TODO: `extra_tensors` is not longer useful for EPv2, as objects are stored in `self.event` + self.extra_tensors = extra_tensors + + # A wrapper for `with event_overlap(release_handle=True)` + self._release_handle_by_call = False + + def current_stream_wait(self, release_handle: bool = False) -> None: + """ + The current stream `torch.cuda.current_stream()` waits for the event to be finished. + """ + assert self.event is not None + self.event.current_stream_wait() + + # In `self.event`, we also have some V2 APIs storing tensors to record in it, + # So, after waiting the current stream, those tensors can be released by deleting `self.event` + # However, you better do it by yourself (to be compatible with multi-stream waits) + if release_handle: + self.event = None + + def __call__(self, release_handle: bool = False) -> "EventOverlap": + """ + Configures the 'release_handle' behavior for the upcoming context manager usage. + Usage: + with event_overlap(release_handle=True): + ... + Returns `self` to ensure no new wrapper object is created, keeping the reference count of the underlying event unchanged/managed solely by this instance. + """ + self._release_handle_by_call = release_handle + return self + + def __enter__(self) -> Any: + """ + Utility for overlapping and Python `with` syntax. + + You can overlap the kernels on the current stream with the following example: + ```python + event_overlap = event_after_all_to_all_kernels() + with event_overlap: + do_something_on_current_stream() + # After exiting the `with` scope, the current stream with wait the event to be finished. + ``` + """ + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + """ + Utility for overlapping and Python `with` syntax. + + Please follow the example in the `__enter__` function. + """ + if self.event is not None: + self.current_stream_wait(release_handle=self._release_handle_by_call) + self._release_handle_by_call = False diff --git a/deep_ep/utils/find_pkgs.py b/deep_ep/utils/find_pkgs.py new file mode 100644 index 000000000..c68df972d --- /dev/null +++ b/deep_ep/utils/find_pkgs.py @@ -0,0 +1,82 @@ +import functools +import os +import sys +from importlib.metadata import distributions +from typing import Optional + + +def find_pkg_root(name: str, lib_name: Optional[str] = None, optional: bool = False): + """ + Find the root directory of an installed NVIDIA package by inspecting Python package metadata. + Checks environment variables `EP_{NAME}_ROOT_DIR` and `{NAME}_DIR` first. + + Arguments: + name: the package name (e.g., `'nccl'`, `'nvshmem'`). + lib_name: the library filename to search for within the package files. + optional: if ``False``, raises an assertion error when the package is not found. + + Returns: + root: the package root directory, or `None` if not found and optional. + """ + upper = name.upper() + for env_name in (f'EP_{upper}_ROOT_DIR', f'{upper}_DIR'): + if env_name in os.environ: + return os.environ[env_name] + + path_priority = {p: i for i, p in enumerate(sys.path)} + best, best_priority = None, len(sys.path) + + for dist in distributions(): + dist_name = (dist.metadata['Name'] or '').lower() + if f'nvidia-{name}' not in dist_name and f'nvidia_{name}' not in dist_name: + continue + + dist_site = str(dist._path.parent) + priority = path_priority.get(dist_site, len(sys.path)) + if priority > best_priority: + continue + + if lib_name is not None: + for f in (dist.files or []): + if lib_name in str(f): + lib_dir = os.path.dirname(str(f.locate())) + root = os.path.dirname(lib_dir) if os.path.basename(lib_dir) == 'lib' else lib_dir + best, best_priority = root, priority + break + else: + pkg_dir = os.path.join(dist_site, 'nvidia', name) + if os.path.isdir(pkg_dir): + best, best_priority = pkg_dir, priority + + # Raise error if not optional + if not optional: + assert best is not None, f'Cannot find package: {name}' + return best + + +@functools.lru_cache() +def find_nccl_root(optional: bool = False): + """ + Find the NCCL installation root directory, cached. + + Arguments: + optional: if `False`, raises an assertion error when NCCL is not found. + + Returns: + root: the NCCL root directory. + """ + return find_pkg_root('nccl', lib_name='libnccl.so', optional=optional) + + +@functools.lru_cache() +def find_nvshmem_root(optional: bool = False): + """ + Find the NVSHMEM installation root directory, cached. + + Arguments: + optional: if `False`, raises an assertion error when NVSHMEM is not found. + + Returns: + root: the NVSHMEM root directory. + """ + return find_pkg_root('nvshmem', optional=optional) diff --git a/deep_ep/utils/gate.py b/deep_ep/utils/gate.py new file mode 100644 index 000000000..e796f6253 --- /dev/null +++ b/deep_ep/utils/gate.py @@ -0,0 +1,180 @@ +import torch + + +def generate_topk_idx(rank_count: torch.Tensor, num_tokens: int, num_experts: int, num_ranks: int, num_topk: int) -> torch.Tensor: + """ + Map rank count to expert indices + """ + assert torch.equal(torch.sum(rank_count, dim=1), torch.ones(num_tokens, dtype=torch.int, device='cuda') * num_topk) + assert (num_tokens, num_ranks) == rank_count.shape + num_experts_per_rank = num_experts // num_ranks + + # Generate base value + base_vals = torch.arange(num_experts, device='cuda').view(1, num_ranks, num_experts_per_rank).expand(num_tokens, num_ranks, num_experts_per_rank) + + # Randomize the ordering within each row + rand_vals = torch.rand(num_tokens, num_ranks, num_experts_per_rank, device='cuda') + perm_indices = torch.argsort(rand_vals, dim=-1) + permuted = torch.gather(base_vals, 2, perm_indices) + + # Create the mask + k_idx = torch.arange(num_experts_per_rank, device='cuda').view(1, 1, num_experts_per_rank).expand(num_tokens, num_ranks, num_experts_per_rank) + rank_count_expanded = rank_count.unsqueeze(2).expand(num_tokens, num_ranks, num_experts_per_rank) + mask = k_idx < rank_count_expanded + + # Get the final indices by masking and reshaping + selected = permuted[mask] # (num_tokens * num_topk,) + topk_idx = selected.view(num_tokens, num_topk) + + return topk_idx + + +def generate_rank_count(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float) -> torch.Tensor: + """ + Generate rank count tensor for a given number of tokens, experts, ranks, and top-k. + + This function generates a tensor of shape `(num_tokens, num_ranks)` where each element `[i, j]` represents + the number of topk experts that token `i` have on rank `j`. The distribution is such that + one special rank gets `ratio` times more traffic than the others. + """ + num_experts_per_rank = num_experts // num_ranks + num_normal_ranks = num_ranks - 1 + + assert ratio >= 1.0, 'ratio must be no less than 1.0' + + # Generate rank count of each token from random distribution + random_scores = torch.rand(num_tokens, num_experts, device='cuda') + topk_weights_, topk_indices = torch.topk(random_scores, num_topk, dim=1, largest=True, sorted=False) + topk_indices //= num_experts_per_rank + sorted_topk_indices = torch.sort(topk_indices, dim=1)[0] + topk_indices_diff_mask = sorted_topk_indices[:, 1:] != sorted_topk_indices[:, :-1] + a = topk_indices_diff_mask.sum(dim=1) + 1 + + # Upper bound for this generating algorithm + upper_bound_per_token = int(num_normal_ranks / ratio) + 1 + + # Clamp the value in range [1, upper_bound_per_token] for each token + a = torch.clamp(a, None, upper_bound_per_token) + + # Consider the special rank + sum_a = torch.sum(a).item() + normal_token_count = int(sum_a / (num_normal_ranks + ratio)) + special_token_count = sum_a - normal_token_count * num_normal_ranks + special_token_count = min(special_token_count, int(normal_token_count * ratio) + 1) + + # Tokens that the special rank must be in topk + must_mask = (a == num_ranks) + must_count = int(must_mask.sum().item()) + special_token_count = max(must_count, special_token_count) + assert must_count <= special_token_count, 'Too many tokens with full rank assignment' + + # Tokens that the special rank can optionally be in topk + optional_token_indices = torch.where(must_mask == 0)[0] + optional_token_indices = optional_token_indices[torch.randperm(num_tokens - must_count, device='cuda')][:special_token_count - must_count] + must_token_indices = torch.where(must_mask != 0)[0] + special_token_row_index = torch.cat(([must_token_indices, optional_token_indices])) + + # Generate permutations for normal ranks + rank_perm = (torch.randperm(num_normal_ranks, device='cuda') + 1).repeat(num_tokens * num_topk // num_normal_ranks + 1) + + # Compute cumulative sum of a to get starting indices in b for each row + a_cumsum = torch.cumsum(torch.cat((torch.tensor([0], device='cuda'), a)), dim=0) + row_starts = a_cumsum[:-1] # Starting indices for each row in b, shape (n,) + + # Insert special rank index into the permutation for special tokens + rank_perm_with_special_rank = torch.zeros(num_tokens * num_topk, dtype=torch.long, device='cuda') # (n * k,) + special_token_mask = torch.zeros(num_tokens * num_topk, dtype=torch.bool, device='cuda') + special_token_flattened_row_index = row_starts[special_token_row_index] + special_token_mask[special_token_flattened_row_index] = 1 + all_indices = torch.arange(num_tokens * num_topk, device='cuda') + non_special_indices = all_indices[special_token_mask != True] + rank_perm_with_special_rank[non_special_indices] = rank_perm[:len(non_special_indices)] + + # Create column index grids + col_idx = torch.arange(num_topk, device='cuda').view(1, num_topk) # (1, num_topk) + + # Compute modulo indices: col_idx % a[i] for each row + # torch.max is used to avoid zeros in case a[i] = 0 (which happens when the only topk rank is the special rank) + mod_idx = col_idx % a.view(num_tokens, 1) # (n, num_topk) + + # Compute indices in b: row_start + (col % a[i]) + b_indices = row_starts.view(num_tokens, 1) + mod_idx # (n, k) + + # Gather values from b using computed indices + result = rank_perm_with_special_rank[b_indices] + + # Shuffle rows randomly to avoid any pattern + shuffle_indices = torch.randperm(num_tokens, device='cuda') + result = result[shuffle_indices] # Shuffle rows + + # Create rank count tensor + rank_count = torch.zeros((num_tokens, num_ranks), dtype=torch.int32, device='cuda') + rank_count.scatter_add_(dim=1, index=result, src=torch.ones_like(result, dtype=torch.int32)) + return rank_count + + +def get_precise_unbalanced_scores(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float): + """ + Generate precise unbalanced scores for testing. + + Note that this function generates a distribution with precise unbalanced distribution, + which **differs from real distribution**. + """ + # Generate num topk experts for each rank + rank_count = generate_rank_count(num_tokens, num_experts, num_ranks, num_topk, ratio) + + # Generate scores in a low distribution + threshold = 0.9 + scores = torch.empty((num_tokens, num_experts), dtype=torch.float32, device='cuda') + scores.uniform_(to=threshold) + + # Generate topk indices and change their scores to a high distribution + topk_idx = generate_topk_idx(rank_count, num_tokens, num_experts, num_ranks, num_topk) + topk_scores = torch.empty((num_tokens, num_topk), dtype=torch.float32, device='cuda') + topk_scores.uniform_(threshold + 1e-6, 1.0) + row_idx = torch.arange(num_tokens).unsqueeze(1).expand(num_tokens, num_topk) + scores[row_idx, topk_idx] = topk_scores + return scores + + +def get_scores_by_factor(num_tokens: int, num_experts: int, num_ranks: int, factor: float) -> torch.Tensor: + num_experts_per_rank = num_experts // num_ranks + scores = torch.empty((num_tokens, num_experts), dtype=torch.float32, device='cuda') + scores[:, :num_experts_per_rank].uniform_(to=factor) + scores[:, num_experts_per_rank:].uniform_(to=1) + return scores + + +def map_unbalanced_ratio_to_factor(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float) -> float: + num_iterations = 20 + factor_l, factor_r = 1.0, 100.0 + + num_experts_per_rank = num_experts // num_ranks + for _i in range(num_iterations): + factor_mid = (factor_l + factor_r) / 2 + scores = get_scores_by_factor(num_tokens, num_experts, num_ranks, factor_mid) + _, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False) + rank_idx = topk_idx // num_experts_per_rank + one_hot = torch.nn.functional.one_hot(rank_idx, num_ranks) + counts = one_hot.any(dim=1).to(torch.float).sum(dim=0) + if counts[0].item() > counts[1:].mean().item() * ratio: + factor_r = factor_mid + else: + factor_l = factor_mid + return factor_l + + +def get_random_unbalanced_scores(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float): + """Generate unbalanced scores with a given ratio. + """ + factor = 1.0 + if ratio != 1.0: + factor = map_unbalanced_ratio_to_factor(num_tokens, num_experts, num_ranks, num_topk, ratio) + return get_scores_by_factor(num_tokens, num_experts, num_ranks, factor) + + +def get_unbalanced_scores(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float, precise: bool): + if precise: + return get_precise_unbalanced_scores(num_tokens, num_experts, num_ranks, num_topk, ratio) + else: + return get_random_unbalanced_scores(num_tokens, num_experts, num_ranks, num_topk, ratio) diff --git a/deep_ep/utils/math.py b/deep_ep/utils/math.py new file mode 100644 index 000000000..199418968 --- /dev/null +++ b/deep_ep/utils/math.py @@ -0,0 +1,103 @@ +import torch +from typing import Tuple + + +def calc_diff(x: torch.Tensor, y: torch.Tensor) -> float: + x, y = x.double() + 1, y.double() + 1 + denominator = (x * x + y * y).sum() + sim = 2 * (x * y).sum() / denominator + return (1 - sim).item() + + +def safe_div(a, b) -> float: + try: + return a / b + except ZeroDivisionError as e: + if a == 0: + return 0 + else: + raise + + +def ceil_div(x: int, y: int) -> int: + return (x + y - 1) // y + + +def align(x: int, y: int) -> int: + return ceil_div(x, y) * y + + +@torch.compile(dynamic=True) +def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 + m, n = x.shape + aligned_n = align(n, 128) + x_padded = torch.nn.functional.pad(x, (0, aligned_n - n), mode='constant', value=0) + x_padded_view = x_padded.view(m, -1, 128) + x_amax = x_padded_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) + return (x_padded_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view( + m, aligned_n)[:, :n].contiguous(), (x_amax / 448.0).view(m, -1) + + +@torch.compile(dynamic=True) +def per_token_cast_back(x_fp8: torch.Tensor, x_scales: torch.Tensor) -> torch.Tensor: + if x_fp8.numel() == 0: + return x_fp8.to(torch.bfloat16) + + assert x_fp8.dim() == 2 + m, n = x_fp8.shape + aligned_n = align(n, 128) + x_fp8_padded = torch.nn.functional.pad(x_fp8, (0, aligned_n - n), mode='constant', value=0) + if x_scales.dtype == torch.int: + x_scales = x_scales.view(dtype=torch.uint8).to(torch.int) << 23 + x_scales = x_scales.view(dtype=torch.float) + x_fp32_padded = x_fp8_padded.to(torch.float32).view(x_fp8.shape[0], -1, 128) + x_scales = x_scales.view(x_fp8.shape[0], -1, 1) + return (x_fp32_padded * x_scales).view(x_fp8_padded.shape).to(torch.bfloat16)[:, :n].contiguous() + + +def inplace_unique(x: torch.Tensor, num_slots: int) -> None: + assert x.dim() == 2 + mask = x < 0 + x_padded = x.masked_fill(mask, num_slots) + bin_count = torch.zeros((x.size(0), num_slots + 1), dtype=x.dtype, device=x.device) + bin_count.scatter_add_(1, x_padded, torch.ones_like(x_padded)) + bin_count = bin_count[:, :num_slots] + sorted_bin_count, sorted_bin_idx = torch.sort(bin_count, dim=-1, descending=True) + sorted_bin_idx.masked_fill_(sorted_bin_count == 0, -1) + sorted_bin_idx = torch.sort(sorted_bin_idx, descending=True, dim=-1).values + x[:, :].fill_(-1) + valid_len = min(num_slots, x.size(1)) + x[:, :valid_len] = sorted_bin_idx[:, :valid_len] + + +def create_grouped_scores(scores: torch.Tensor, group_idx: torch.Tensor, num_groups: int) -> torch.Tensor: + num_tokens, num_experts = scores.shape + scores = scores.view(num_tokens, num_groups, -1) + mask = torch.zeros((num_tokens, num_groups), dtype=torch.bool, device=scores.device) + mask = mask.scatter_(1, group_idx, True).unsqueeze(-1).expand_as(scores) + return (scores * mask).view(num_tokens, num_experts) + + +def hash_tensor(t: torch.Tensor) -> int: + return t.view(torch.int).sum().item() + + +def hash_tensors(*tensors) -> int: + value = 0 + for t in tensors: + if isinstance(t, (tuple, list)): + value ^= hash_tensors(*t) + elif t is not None and isinstance(t, torch.Tensor): + value ^= hash_tensor(t) + return value + + +def count_bytes(*tensors) -> int: + total = 0 + for t in tensors: + if isinstance(t, (tuple, list)): + total += count_bytes(*t) + elif t is not None: + total += t.numel() * t.element_size() + return total diff --git a/deep_ep/utils/refs.py b/deep_ep/utils/refs.py new file mode 100644 index 000000000..c3ab3cbcb --- /dev/null +++ b/deep_ep/utils/refs.py @@ -0,0 +1,243 @@ +import math +import torch +import torch.distributed as dist +from typing import Optional, Tuple, Union + +from .envs import get_global_seed +from .math import ceil_div + + +def dispatch(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + topk_idx: torch.Tensor, topk_weights: Optional[torch.Tensor], + num_max_tokens_per_rank: int, num_experts: int): + """ + The reference implementation of dispatching tokens to experts across multiple ranks. + + Not expanded. Sorted by rank and then by token within each rank (i.e. sorted by `src_token_global_idx`). + + Arguments: + - `x`: Input tokens, `[num_tokens, hidden]` or (`[num_tokens, hidden], [num_tokens, hidden_sf]`) + - `topk_idx`: Top-k expert indices for each token, `[num_tokens, num_topk]` + - `topk_weights`: Top-k weights for each token, can be None, `[num_tokens, num_topk]` + - `num_max_tokens_per_rank`: Maximum number of tokens per rank, must >= actual number of tokens per rank and aligned with DeepEP's `num_max_tokens_per_rank` since we're going to calculate `src_token_global_idx = src_rank_idx*num_max_tokens_per_rank + src_token_local_idx` + - `num_experts`: Total number of experts across all ranks + + Returns: + - `recv_x`, `recv_topk_idx`, and `recv_topk_weights`: Received tokens, top-k indices, and top-k weights after dispatching. Out of range `recv_topk_idx` (i.e. that expert is not on the current rank) are set to -1. + - `recv_src_token_idx`: Received `src_token_global_idx` for each received token + - `num_recv_tokens_per_rank`: Number of received tokens from each rank, `[num_ranks]` + """ + # TODO: support forwarding + # TODO: make top-k weight fully optional + rank_idx = dist.get_rank() + num_ranks = dist.get_world_size() + + assert num_experts % num_ranks == 0 + num_experts_per_rank = num_experts // num_ranks + + # Unpack SF + use_fp8 = isinstance(x, tuple) + x, sf = x if use_fp8 else (x, None) + + # TODO: use SF bytes instead of hardcoded recipe + num_tokens, hidden = x.size() + num_tokens_, num_topk = topk_idx.size() + assert num_tokens == num_tokens_ + if sf is not None: + num_tokens__, hidden_sf = sf.size() + assert num_tokens == num_tokens__ + assert hidden_sf == ceil_div(hidden, 128) + if topk_weights is not None: + num_tokens__, num_topk_ = topk_weights.size() + assert num_tokens == num_tokens__ + assert num_topk == num_topk_ + + # Prepare per-peer send buffers + send_x_list = [] + send_sf_list = [] + send_topk_idx_list = [] + send_topk_weights_list = [] + send_src_token_idx_list = [] + num_send_tokens_per_rank = torch.zeros((num_ranks, ), dtype=torch.int, device=x.device) + for dst_rank_idx in range(num_ranks): + expert_start_idx = dst_rank_idx * num_experts_per_rank + expert_end_idx = expert_start_idx + num_experts_per_rank + + # Get the indices of tokens + mask_to_send = ((expert_start_idx <= topk_idx) & (topk_idx < expert_end_idx)).any(dim=1) + indices_to_send = mask_to_send.nonzero(as_tuple=True)[0] + num_send_tokens_per_rank[dst_rank_idx] = indices_to_send.numel() + + # Select the data for tokens + x_to_send = x[indices_to_send] + sf_to_send = sf[indices_to_send] if use_fp8 else None + topk_idx_to_send = topk_idx[indices_to_send] + topk_weights_to_send = topk_weights[indices_to_send] + masked_topk_idx = torch.where((expert_start_idx <= topk_idx_to_send) & (topk_idx_to_send < expert_end_idx), + topk_idx_to_send, torch.full_like(topk_idx_to_send, -1)) + + send_x_list.append(x_to_send) + send_sf_list.append(sf_to_send) + send_topk_idx_list.append(masked_topk_idx) + send_topk_weights_list.append(topk_weights_to_send) + send_src_token_idx_list.append(indices_to_send) + + send_x = torch.cat(send_x_list, dim=0) + send_sf = torch.cat(send_sf_list, dim=0) if use_fp8 else None + send_topk_idx = torch.cat(send_topk_idx_list, dim=0) + send_topk_weights = torch.cat(send_topk_weights_list, dim=0) + send_src_token_idx = torch.cat(send_src_token_idx_list, dim=0).to(torch.int) + send_src_token_idx += rank_idx * num_max_tokens_per_rank + + # Exchange size + num_recv_tokens_per_rank = torch.empty((num_ranks, ), dtype=torch.int, device=x.device) + dist.all_to_all_single(num_recv_tokens_per_rank, num_send_tokens_per_rank) + num_recv_tokens = int(num_recv_tokens_per_rank.sum().item()) + + # Exchange main data + num_send_tokens_per_rank = num_send_tokens_per_rank.tolist() + num_recv_tokens_per_rank = num_recv_tokens_per_rank.tolist() + recv_x = torch.empty((num_recv_tokens, hidden), dtype=x.dtype, device=x.device) + recv_sf = torch.empty((num_recv_tokens, hidden_sf), dtype=sf.dtype, device=x.device) if use_fp8 else None + recv_topk_idx = torch.empty((num_recv_tokens, num_topk), dtype=topk_idx.dtype, device=x.device) + recv_topk_weights = torch.empty((num_recv_tokens, num_topk), dtype=topk_weights.dtype, device=x.device) + recv_src_token_idx = torch.empty((num_recv_tokens, ), dtype=torch.int, device=x.device) + dist.all_to_all_single(recv_x, send_x, num_recv_tokens_per_rank, num_send_tokens_per_rank) + if use_fp8: + dist.all_to_all_single(recv_sf, send_sf, num_recv_tokens_per_rank, num_send_tokens_per_rank) + dist.all_to_all_single(recv_topk_idx, send_topk_idx, num_recv_tokens_per_rank, num_send_tokens_per_rank) + dist.all_to_all_single(recv_topk_weights, send_topk_weights, num_recv_tokens_per_rank, num_send_tokens_per_rank) + dist.all_to_all_single(recv_src_token_idx, send_src_token_idx, num_recv_tokens_per_rank, num_send_tokens_per_rank) + + # Mask top-k indices + expert_start_idx = rank_idx * num_experts_per_rank + expert_end_idx = expert_start_idx + num_experts_per_rank + mask = (expert_start_idx <= recv_topk_idx) & (recv_topk_idx < expert_end_idx) + recv_topk_idx = recv_topk_idx - expert_start_idx + recv_topk_idx.masked_fill_(~mask, -1) + + # Pack SF + recv_x = (recv_x, recv_sf) if use_fp8 else recv_x + + return (recv_x, recv_topk_idx, recv_topk_weights, + recv_src_token_idx, torch.tensor(num_recv_tokens_per_rank, dtype=torch.int)) + + +def generate_pre_combine_data(src_token_global_idx: torch.Tensor, + num_max_tokens_per_rank: int, num_topk: int, hidden: int) -> torch.Tensor: + """ + Generate data needed for combine from `src_token_global_idx`. + Recall that `src_token_global_idx = src_rank_idx * num_max_tokens_per_rank + src_token_local_idx`. + The generated data (denoted as `y`) of the i-th token has a shape of [num_topk, hidden], with + + `y[j, k] = sin((token_seeds * P % max_seed + 1) / max_seed * (k + 1) + sin(seed))` + + where `P=131071` is a large prime, `token_seeds` is calculated via `token_seeds = src_token_global_idx[i] * num_topk + j`, + and `max_seed = num_ranks * num_max_tokens_per_rank * num_topk`. + + Arguments: + - `src_token_global_idx`: Source token global indices, `[num_tokens]` + + Returns: + - Generated data, `[num_tokens, num_topk, hidden]` + """ + num_ranks = dist.get_world_size() + token_seeds = (src_token_global_idx.unsqueeze(1) * num_topk + + torch.arange(num_topk, device=src_token_global_idx.device).unsqueeze(0)) # [num_tokens, num_topk] + max_seed = num_ranks * num_max_tokens_per_rank * num_topk + result = torch.sin( + (((token_seeds * 131071 % max_seed).float() + 1) / max_seed).unsqueeze(-1) * + torch.arange(1, hidden + 1, device=src_token_global_idx.device, dtype=torch.float32).broadcast_to(1, 1, hidden) + + math.sin(float(get_global_seed())) + ) + return result.to(torch.bfloat16) + + +def ordered_accumulate(data: torch.Tensor, initial_value: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Accumulate `data` in order along the num_topk dimension. + + Arguments: + - `data`: Data to be accumulated, `[num_tokens, num_topk, hidden]` + - `initial_value`: Initial value for accumulation, `[num_tokens, hidden]` + + Returns: + - Result, `[num_tokens, hidden]` + """ + num_topk = data.shape[1] + if initial_value is None: + result = torch.zeros((data.shape[0], data.shape[2]), dtype=torch.float32, device=data.device) + else: + result = initial_value.clone() + for i in range(num_topk): + result += data[:, i, :].float() + return result.to(data.dtype) + + +def combine(y: torch.Tensor, topk_idx: torch.Tensor, + num_scaleout_ranks: int, num_scaleup_ranks: int, num_experts: int, + bias: Optional[Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]], + reduce_in_local: bool, reduce_in_scaleup: bool) -> torch.Tensor: + """ + The reference implementation of (possibly multi-level reduction) combining tokens. + + Arguments: + - `y`: Input tokens to be combined, `[num_tokens, num_topk, hidden]` + - `topk_idx`: `[num_tokens, num_topk]` + - `reduce_in_local` and `reduce_in_scaleup`: Whether to do reduction within rank and within scale-up group. + - `(True, True)` -> Hybrid combine + - `(True, False)` -> Non-hybrid combine + - `(False, False)` -> Equivalent to `allow_multiple_reduction` is False + Pay attention that `reduce_in_scaleup` = `True` or `False` is NOT equivalent even if `num_scaleout_ranks == 1` due to `bias` handling. + + Returns: + - Combined result, `[num_tokens, hidden]` + """ + num_ranks = num_scaleout_ranks * num_scaleup_ranks + num_tokens, hidden = y.shape[0], y.shape[2] + num_topk = y.shape[1] + assert not (not reduce_in_local and reduce_in_scaleup), 'Invalid reduction configuration' + + def grouped_reduce(data_to_reduce: torch.Tensor, group_id: torch.Tensor) -> torch.Tensor: + """ + Perform in-place grouped reduction on `data_to_reduce` according to `group_id`. + The summation within each group are performed in strict order along the `num_topk` dimension. + The result for each group is stored at the rightmost token of that group, and other tokens are set to zero. + + Arguments: + - `data_to_reduce`: Data to be reduced, `[num_tokens, num_topk, hidden]` + - `group_id`: group IDs for each token, `[num_tokens, num_topk]` + """ + # Shuffle to make tokens with the same group_id contiguous + group_id, src_indices = torch.sort(group_id, dim=-1, stable=True) + # transformed_src_indices[i, j] = i * num_topk + src_indices[i, j] + transformed_src_indices = ( + (src_indices + torch.arange(0, num_tokens, device=y.device).unsqueeze(-1) * num_topk).flatten()) + data_to_reduce = data_to_reduce.view(-1, hidden)[transformed_src_indices].view(num_tokens, num_topk, hidden) + # Perform segmented reduce within each group + cur_accum_buf = torch.zeros((num_tokens, hidden), dtype=torch.float32, device=y.device) + for i in range(num_topk): + is_segment_break = torch.full((num_tokens, ), True, dtype=torch.bool, device=y.device) \ + if i == num_topk - 1 else group_id[:, i] != group_id[:, i + 1] + cur_accum_buf += data_to_reduce[:, i, :].float() + # For one token, if `is_segment_break` is True, + # save the accumulated value and clear the buffer, otherwise, clear `data_to_reduce[:, i, :]` + segment_break_token_indices = torch.where(is_segment_break)[0] + data_to_reduce[segment_break_token_indices, i] = cur_accum_buf[segment_break_token_indices].to(data_to_reduce.dtype) + cur_accum_buf[segment_break_token_indices] = 0.0 + non_segment_break_token_indices = torch.where(~is_segment_break)[0] + data_to_reduce[non_segment_break_token_indices, i] = 0.0 + # Unshuffle + # noinspection PyShadowingNames + result = torch.empty_like(data_to_reduce) + result.view(-1, hidden)[transformed_src_indices] = data_to_reduce.view(-1, hidden) + return result.view(num_tokens, num_topk, hidden) + + num_experts_per_rank = num_experts // num_ranks + if reduce_in_local: + y = grouped_reduce(y, topk_idx // num_experts_per_rank) + if reduce_in_scaleup: + y = grouped_reduce(y, topk_idx // (num_experts_per_rank * num_scaleup_ranks)) + bias_sum = bias[0].float() + bias[1].float() if isinstance(bias, tuple) else bias.float() if bias is not None else None + result = ordered_accumulate(y, bias_sum) + return result diff --git a/deep_ep/utils/semantic.py b/deep_ep/utils/semantic.py new file mode 100644 index 000000000..bd8af6ee0 --- /dev/null +++ b/deep_ep/utils/semantic.py @@ -0,0 +1,27 @@ +from typing import Any, Optional +import weakref +import functools + +def value_or(value: Optional[Any], default: Any) -> Any: + return default if value is None else value + + +def weak_lru(maxsize: Optional[int] = 128, typed: bool = False): + """ + LRU Cache decorator that keeps a weak reference to `self` + Useful for caching methods in classes that may cause memory leaks if `functools.lru_cache` is used directly. + From https://stackoverflow.com/a/68052994/16569836 + """ + def wrapper(func): + + @functools.lru_cache(maxsize, typed) + def _func(_self, *args, **kwargs): + return func(_self(), *args, **kwargs) + + @functools.wraps(func) + def inner(self, *args, **kwargs): + return _func(weakref.ref(self), *args, **kwargs) + + return inner + + return wrapper diff --git a/deep_ep/utils/testing.py b/deep_ep/utils/testing.py new file mode 100644 index 000000000..cc005bb89 --- /dev/null +++ b/deep_ep/utils/testing.py @@ -0,0 +1,219 @@ +import json +import os +import sys +import torch +import numpy as np +import tempfile +import torch.distributed as dist +from pathlib import Path +from typing import Callable, Optional, Union + + +def flush_l2_cache(enabled: bool = True): + """ + Flush the GPU L2 cache by writing a large zero-initialized tensor. + + Arguments: + enabled: if `False`, does nothing. + """ + l2_flush_cache_size = 256e6 + if enabled: + torch.empty(int(l2_flush_cache_size // 4), dtype=torch.int, device='cuda').zero_() + + +def bench(fn, num_warmups: int = 50, num_tests: int = 50, + post_fn: Optional[Callable] = None, flush_l2: bool = True): + """ + Benchmark a function using CUDA events. + + Arguments: + fn: the function to benchmark. + num_warmups: the number of warmup iterations. + num_tests: the number of measurement iterations. + post_fn: an optional function to call after each test iteration. + flush_l2: whether to flush the L2 cache before each iteration. + + Returns: + avg: the average execution time in seconds. + min: the minimum execution time in seconds. + max: the maximum execution time in seconds. + """ + torch.cuda.synchronize() + + # Warmup + for _ in range(num_warmups): + fn() + + # Testing + start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] + end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] + for i in range(num_tests): + flush_l2_cache(flush_l2) + start_events[i].record() + fn() + end_events[i].record() + if post_fn is not None: + post_fn() + torch.cuda.synchronize() + + times = np.array([s.elapsed_time(e) / 1e3 for s, e in zip(start_events, end_events)])[1:] + return np.average(times), np.min(times), np.max(times) + + +class empty_suppress: + + def __enter__(self): + return self + + def __exit__(self, *_): + pass + + +class suppress_stdout_stderr: + """ + Context manager to suppress stdout and stderr output. + """ + + def __enter__(self): + self.outnull_file = open(os.devnull, 'w') + self.errnull_file = open(os.devnull, 'w') + + self.old_stdout_fileno_undup = sys.stdout.fileno() + self.old_stderr_fileno_undup = sys.stderr.fileno() + + self.old_stdout_fileno = os.dup(sys.stdout.fileno()) + self.old_stderr_fileno = os.dup(sys.stderr.fileno()) + + self.old_stdout = sys.stdout + self.old_stderr = sys.stderr + + os.dup2(self.outnull_file.fileno(), self.old_stdout_fileno_undup) + os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup) + + sys.stdout = self.outnull_file + sys.stderr = self.errnull_file + return self + + def __exit__(self, *_): + sys.stdout = self.old_stdout + sys.stderr = self.old_stderr + + os.dup2(self.old_stdout_fileno, self.old_stdout_fileno_undup) + os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup) + + os.close(self.old_stdout_fileno) + os.close(self.old_stderr_fileno) + + self.outnull_file.close() + self.errnull_file.close() + + +def bench_kineto(fn, + kernel_names: Union[str, tuple], + num_tests: int = 30, + suppress_kineto_output: bool = False, + trace_path: Optional[str] = None, + flush_l2: bool = True, + barrier_comm_profiling: bool = False, + num_kernels_per_period: int = 1, + barrier: Optional[Callable] = None): + """ + Benchmark a function using the PyTorch profiler (kineto) to get per-kernel timing. + + Arguments: + fn: the function to benchmark. + kernel_names: the CUDA kernel name(s) to profile. + num_tests: the number of test iterations. + suppress_kineto_output: whether to suppress profiler output. + trace_path: the path to save the Chrome trace (`None` to skip). + flush_l2: whether to flush the L2 cache before each iteration. + barrier_comm_profiling: whether to insert a barrier before each iteration to reduce + unbalanced CPU launch overhead. + num_kernels_per_period: the number of kernels launched per test period. + barrier: a custom barrier function to use instead of `dist.all_reduce`. + + Returns: + durations: the average kernel duration(s) in seconds. + """ + assert isinstance(kernel_names, (str, tuple)) + is_tuple = isinstance(kernel_names, tuple) + + # Skip profiling + # Conflict with Nsight Systems, Nsight Compute and Compute Sanitizer + if int(os.environ.get('EP_USE_NVIDIA_TOOLS', 0)): + return (1, ) * len(kernel_names) if is_tuple else 1 + + # For some auto-tuning kernels with prints + fn() + torch.cuda.synchronize() + + # Profile + suppress = suppress_stdout_stderr if suppress_kineto_output else empty_suppress + barrier_comm_profiling &= int(os.environ.get('EP_DISABLE_BARRIER_PROFILING', 0)) == 0 + with suppress(): + schedule = torch.profiler.schedule(wait=0, warmup=1, active=1, repeat=1) + profiler = torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA], schedule=schedule, acc_events=True) + dummy = torch.ones(1, dtype=torch.float, device='cuda') + with profiler: + for i in range(2): + for _ in range(num_tests): + # Flush L2 cache + flush_l2_cache(flush_l2) + + # NOTES: use a large kernel and a barrier to eliminate the unbalanced CPU launch overhead + if barrier_comm_profiling: + torch.cuda._sleep(int(2e7)) # ~10ms + + # Some network may have ring-based implement, so be careful to use `all_reduce` + if barrier is None: + dist.all_reduce(dummy) + else: + barrier() + fn() + torch.cuda.synchronize() + profiler.step() + + # Parse the profiling table + prof_lines = profiler.key_averages().table(sort_by='cuda_time_total', max_name_column_width=100).split('\n') + kernel_names = (kernel_names, ) if isinstance(kernel_names, str) else kernel_names + assert all([isinstance(name, str) for name in kernel_names]) + for name in kernel_names: + assert sum([name in line for line in prof_lines]) <= 1, f'Errors of the kernel {name} in the profiling table: {prof_lines}' + + # Save chrome traces + if trace_path is not None: + profiler.export_chrome_trace(trace_path) + + # Return average kernel durations + units = {'ms': 1e3, 'us': 1e6} + kernel_durations = [] + for name in kernel_names: + total_time = 0 + total_num = 0 + for line in prof_lines: + if name in line: + time_str = line.split()[-2] + num_str = line.split()[-1] + for unit, scale in units.items(): + if unit in time_str: + total_time += float(time_str.replace(unit, '')) / scale * int(num_str) + total_num += int(num_str) + break + kernel_durations.append(total_time / total_num if total_num > 0 else 0) + + # Expand the kernels by periods + if num_kernels_per_period > 1: + with tempfile.NamedTemporaryFile(suffix='.json') as tmp: + profiler.export_chrome_trace(tmp.name) + profile_data = json.loads(Path(tmp.name).read_text()) + + for i, kernel_name in enumerate(kernel_names): + events = [event for event in profile_data['traceEvents'] if f'::{kernel_name}' in event['name']] + events = sorted(events, key=lambda event: event['ts']) + durations = [event['dur'] / 1e6 for event in events] + assert len(durations) % num_kernels_per_period == 0 + num_kernel_patterns = len(durations) // num_kernels_per_period + kernel_durations[i] = [sum(durations[j::num_kernels_per_period]) / num_kernel_patterns for j in range(num_kernels_per_period)] + + # Return execution durations + return kernel_durations if is_tuple else kernel_durations[0] diff --git a/develop.sh b/develop.sh new file mode 100755 index 000000000..9c5f85a21 --- /dev/null +++ b/develop.sh @@ -0,0 +1,21 @@ +# Change current directory into project root +original_dir=$(pwd) +script_dir=$(realpath "$(dirname "$0")") +cd "$script_dir" + +# Remove old dist file, build files, and build +rm -rf build dist +rm -rf *.egg-info +python setup.py build + +# Find the .so file in build directory and create symlink in current directory +so_file=$(find build -name "*.so" -type f | head -n 1) +if [ -n "$so_file" ]; then + ln -sf "../$so_file" deep_ep/ +else + echo "Error: No SO file found in build directory" >&2 + exit 1 +fi + +# Open users' original directory +cd "$original_dir" diff --git a/docs/legacy.md b/docs/legacy.md new file mode 100644 index 000000000..2aa77061d --- /dev/null +++ b/docs/legacy.md @@ -0,0 +1,320 @@ +# DeepEP V1 (Legacy) + +> **Note:** This is the archived documentation for DeepEP V1 (NVSHMEM-based). For the latest V2 documentation, see the [main README](../README.md). + +--- + +DeepEP (DeepEveryParallel) V1 is the original high-performance communication library for modern machine learning, focused on expert parallelism (EP). It provides high-throughput and low-latency all-to-all GPU kernels, which are also known as MoE dispatch and combine. The library also supports low-precision operations, including FP8. + +To align with the group-limited gating algorithm proposed in the [DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3) paper, DeepEP V1 offers a set of kernels optimized for asymmetric-domain bandwidth forwarding, such as forwarding data from NVLink domain to RDMA domain. These kernels deliver high throughput, making them suitable for both training and inference prefilling tasks. Additionally, they support SM (Streaming Multiprocessors) number control. + +For latency-sensitive inference decoding, DeepEP V1 includes a set of low-latency kernels with pure RDMA to minimize delays. The library also introduces a hook-based communication-computation overlapping method that does not occupy any SM resource. + +Notice: the implementation in this library may have some slight differences from the [DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3) paper. + +## Performance + +### Normal kernels with NVLink and RDMA forwarding + +We test normal kernels on H800 (~160 GB/s NVLink maximum bandwidth), with each connected to a CX7 InfiniBand 400 Gb/s RDMA network card (~50 GB/s maximum bandwidth). And we follow the DeepSeek-V3/R1 pretraining setting (4096 tokens per batch, 7168 hidden, top-4 groups, top-8 experts, FP8 dispatching and BF16 combining). + +| Type | Dispatch #EP | Bottleneck bandwidth | Combine #EP | Bottleneck bandwidth | +|:---------:|:------------:|:--------------------:|:-----------:|:--------------------:| +| Intranode | 8 | 153 GB/s (NVLink) | 8 | 158 GB/s (NVLink) | +| Internode | 16 | 43 GB/s (RDMA) | 16 | 43 GB/s (RDMA) | +| Internode | 32 | 58 GB/s (RDMA) | 32 | 57 GB/s (RDMA) | +| Internode | 64 | 51 GB/s (RDMA) | 64 | 50 GB/s (RDMA) | + +### Low-latency kernels with pure RDMA + +We test low-latency kernels on H800 with each connected to a CX7 InfiniBand 400 Gb/s RDMA network card (~50 GB/s maximum bandwidth). And we follow a typical DeepSeek-V3/R1 production setting (128 tokens per batch, 7168 hidden, top-8 experts, FP8 dispatching and BF16 combining). + +| Dispatch #EP | Latency | RDMA bandwidth | Combine #EP | Latency | RDMA bandwidth | +|:------------:|:-------:|:--------------:|:-----------:|:-------:|:--------------:| +| 8 | 77 us | 98 GB/s | 8 | 114 us | 127 GB/s | +| 16 | 118 us | 63 GB/s | 16 | 195 us | 74 GB/s | +| 32 | 155 us | 48 GB/s | 32 | 273 us | 53 GB/s | +| 64 | 173 us | 43 GB/s | 64 | 314 us | 46 GB/s | +| 128 | 192 us | 39 GB/s | 128 | 369 us | 39 GB/s | +| 256 | 194 us | 39 GB/s | 256 | 360 us | 40 GB/s | + +## Quick Start + +### Requirements + +- Ampere (SM80), Hopper (SM90) GPUs, or other architectures with SM90 PTX ISA support +- Python 3.8 and above +- CUDA version + - CUDA 11.0 and above for SM80 GPUs + - CUDA 12.3 and above for SM90 GPUs +- PyTorch 2.1 and above +- NVLink for intranode communication +- RDMA network for internode communication + +### Download and install NVSHMEM dependency + +DeepEP V1 depends on NVSHMEM. Please refer to the NVSHMEM Installation Guide for instructions. + +### Development + +```bash +# Build and make symbolic links for SO files +NVSHMEM_DIR=/path/to/installed/nvshmem 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.py` +# according to your own cluster settings, and launch into multiple nodes +python tests/test_intranode.py +python tests/test_internode.py +python tests/test_low_latency.py +``` + +### Installation + +```bash +NVSHMEM_DIR=/path/to/installed/nvshmem python setup.py install +``` + +#### Installation environment variables + +- `NVSHMEM_DIR`: the path to the NVSHMEM directory, disable all internode and low-latency features if not specified +- `DISABLE_SM90_FEATURES`: 0 or 1, whether to disable SM90 features, it is required for SM90 devices or CUDA 11 +- `TORCH_CUDA_ARCH_LIST`: the list of target architectures, e.g. `TORCH_CUDA_ARCH_LIST="9.0"` +- `DISABLE_AGGRESSIVE_PTX_INSTRS`: 0 or 1, whether to disable aggressive load/store instructions, see [Undefined-behavior PTX usage](#undefined-behavior-ptx-usage) for more details + +## 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: + +- workloads using normal kernels +- workloads using low-latency kernels +- other workloads + +For DeepEP V1, you can control the virtual lane assignment by setting the `NVSHMEM_IB_SL` environment variable. + +### Adaptive routing + +Adaptive routing is an advanced routing feature provided by InfiniBand switches that can evenly distribute traffic across multiple paths. Enabling adaptive routing can completely eliminate network congestion caused by routing conflicts, but it also introduces additional latency. We recommend the following configuration for optimal performance: + +- enable adaptive routing in environments with heavy network loads +- use static routing in environments with light network loads + +### Congestion control + +Congestion control is disabled as we have not observed significant congestion in our production environment. + +## Interfaces and Examples + +### Example use in model training or inference prefilling + +The normal kernels can be used in model training or the inference prefilling phase (without the backward part) as the below example code shows. + +```python +import torch +import torch.distributed as dist +from typing import List, Tuple, Optional, Union + +from deep_ep import Buffer, EventOverlap + +# Communication buffer (will allocate at runtime) +_buffer: Optional[Buffer] = None + +# Set the number of SMs to use +# NOTES: this is a static variable +Buffer.set_num_sms(24) + + +# You may call this function at the framework initialization +def get_buffer(group: dist.ProcessGroup, hidden_bytes: int) -> Buffer: + global _buffer + + # NOTES: you may also replace `get_*_config` with your auto-tuned results via all the tests + num_nvl_bytes, num_rdma_bytes = 0, 0 + for config in (Buffer.get_dispatch_config(group.size()), Buffer.get_combine_config(group.size())): + num_nvl_bytes = max(config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes) + num_rdma_bytes = max(config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes) + + # Allocate a buffer if not existed or not enough buffer size + if _buffer is None or _buffer.group != group or _buffer.num_nvl_bytes < num_nvl_bytes or _buffer.num_rdma_bytes < num_rdma_bytes: + _buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes) + return _buffer + + +def get_hidden_bytes(x: torch.Tensor) -> int: + t = x[0] if isinstance(x, tuple) else x + return t.size(1) * max(t.element_size(), 2) + + +def dispatch_forward(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + topk_idx: torch.Tensor, topk_weights: torch.Tensor, + num_experts: int, previous_event: Optional[EventOverlap] = None) -> \ + Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], torch.Tensor, torch.Tensor, List, Tuple, EventOverlap]: + # NOTES: an optional `previous_event` means a CUDA event captured that you want to make it as a dependency + # of the dispatch kernel, it may be useful with communication-computation overlap. For more information, please + # refer to the docs of `Buffer.dispatch` + global _buffer + + # Calculate layout before actual dispatch + num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, previous_event = \ + _buffer.get_dispatch_layout(topk_idx, num_experts, + previous_event=previous_event, async_finish=True, + allocate_on_comm_stream=previous_event is not None) + # Do MoE dispatch + # NOTES: the CPU will wait for GPU's signal to arrive, so this is not compatible with CUDA graph + # Unless you specify `num_worst_tokens`, but this flag is for intranode only + # For more advanced usages, please refer to the docs of the `dispatch` function + recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event = \ + _buffer.dispatch(x, topk_idx=topk_idx, topk_weights=topk_weights, + num_tokens_per_rank=num_tokens_per_rank, num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, num_tokens_per_expert=num_tokens_per_expert, + previous_event=previous_event, async_finish=True, + allocate_on_comm_stream=True) + # For event management, please refer to the docs of the `EventOverlap` class + return recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event + + +def dispatch_backward(grad_recv_x: torch.Tensor, grad_recv_topk_weights: torch.Tensor, handle: Tuple) -> \ + Tuple[torch.Tensor, torch.Tensor, EventOverlap]: + global _buffer + + # The backward process of MoE dispatch is actually a combine + # For more advanced usages, please refer to the docs of the `combine` function + combined_grad_x, combined_grad_recv_topk_weights, event = \ + _buffer.combine(grad_recv_x, handle, topk_weights=grad_recv_topk_weights, async_finish=True) + + # For event management, please refer to the docs of the `EventOverlap` class + return combined_grad_x, combined_grad_recv_topk_weights, event + + +def combine_forward(x: torch.Tensor, handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \ + Tuple[torch.Tensor, EventOverlap]: + global _buffer + + # Do MoE combine + # For more advanced usages, please refer to the docs of the `combine` function + combined_x, _, event = _buffer.combine(x, handle, async_finish=True, previous_event=previous_event, + allocate_on_comm_stream=previous_event is not None) + + # For event management, please refer to the docs of the `EventOverlap` class + return combined_x, event + + +def combine_backward(grad_combined_x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \ + Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], EventOverlap]: + global _buffer + + # The backward process of MoE combine is actually a dispatch + # For more advanced usages, please refer to the docs of the `dispatch` function + grad_x, _, _, _, _, event = _buffer.dispatch(grad_combined_x, handle=handle, async_finish=True, + previous_event=previous_event, + allocate_on_comm_stream=previous_event is not None) + + # For event management, please refer to the docs of the `EventOverlap` class + return grad_x, event +``` + +Moreover, inside the dispatch function, we may not know how many tokens to receive for the current rank. So an implicit CPU wait for GPU received count signal will be involved, as the following figure shows. + +![normal](../figures/normal.png) + +### Example use in inference decoding + +The low latency kernels can be used in the inference decoding phase as the below example code shows. + +```python +import torch +import torch.distributed as dist +from typing import Tuple, Optional + +from deep_ep import Buffer + +# Communication buffer (will allocate at runtime) +# NOTES: there is no SM control API for the low-latency kernels +_buffer: Optional[Buffer] = None + + +# You may call this function at the framework initialization +def get_buffer(group: dist.ProcessGroup, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int) -> Buffer: + # NOTES: the low-latency mode will consume much more space than the normal mode + # So we recommend that `num_max_dispatch_tokens_per_rank` (the actual batch size in the decoding engine) should be less than 256 + global _buffer + num_rdma_bytes = Buffer.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, group.size(), num_experts) + + # Allocate a buffer if not existed or not enough buffer size + if _buffer is None or _buffer.group != group or not _buffer.low_latency_mode or _buffer.num_rdma_bytes < num_rdma_bytes: + # NOTES: for the best performance, the QP number **must** be equal to the number of the local experts + assert num_experts % group.size() == 0 + _buffer = Buffer(group, 0, num_rdma_bytes, low_latency_mode=True, num_qps_per_rank=num_experts // group.size()) + return _buffer + + +def low_latency_dispatch(hidden_states: torch.Tensor, topk_idx: torch.Tensor, num_max_dispatch_tokens_per_rank: int, num_experts: int): + global _buffer + + # Do MoE dispatch, compatible with CUDA graph (but you may restore some buffer status once you replay) + recv_hidden_states, recv_expert_count, handle, event, hook = \ + _buffer.low_latency_dispatch(hidden_states, topk_idx, num_max_dispatch_tokens_per_rank, num_experts, + async_finish=False, return_recv_hook=True) + + # NOTES: the actual tensor will not be received only if you call `hook()`, + # it is useful for double-batch overlapping, but **without any SM occupation** + # If you don't want to overlap, please set `return_recv_hook=False` + # Later, you can use our GEMM library to do the computation with this specific format + return recv_hidden_states, recv_expert_count, handle, event, hook + + +def low_latency_combine(hidden_states: torch.Tensor, + topk_idx: torch.Tensor, topk_weights: torch.Tensor, handle: Tuple): + global _buffer + + # Do MoE combine, compatible with CUDA graph (but you may restore some buffer status once you replay) + combined_hidden_states, event_overlap, hook = \ + _buffer.low_latency_combine(hidden_states, topk_idx, topk_weights, handle, + async_finish=False, return_recv_hook=True) + + # NOTES: the same behavior as described in the dispatch kernel + return combined_hidden_states, event_overlap, hook +``` + +For two-micro-batch overlapping, you can refer to the following figure. With our receiving hook interface, the RDMA network traffic is happening in the background, without costing any GPU SMs from the computation part. But notice, the overlapped parts can be adjusted, i.e., the 4 parts of attention/dispatch/MoE/combine may not have the exact same execution time. You may adjust the stage settings according to your workload. + +![low-latency](../figures/low-latency.png) + +## Roadmap (V1) + +- [x] AR support +- [x] Refactor low-latency mode AR code +- [x] A100 support (intranode only) +- [x] Support BF16 for the low-latency dispatch kernel +- [x] Support NVLink protocol for intranode low-latency kernels +- [ ] TMA copy instead of LD/ST + - [x] Intranode kernels + - [ ] Internode kernels + - [ ] Low-latency kernels +- [ ] SM-free kernels and refactors +- [ ] Fully remove undefined-behavior PTX instructions + +## Notices + +#### Easier potential overall design + +The V1 implementation uses queues for communication buffers which save memory but introduce complexity and potential deadlocks. If you're implementing your own version based on DeepEP V1, consider using fixed-size buffers allocated to maximum capacity for simplicity and better performance. For a detailed discussion of this alternative approach, see https://github.com/deepseek-ai/DeepEP/issues/39. + +#### Undefined-behavior PTX usage + +- For extreme performance, we discover and use an undefined-behavior PTX usage: using read-only PTX `ld.global.nc.L1::no_allocate.L2::256B` to **read volatile data**. The PTX modifier `.nc` indicates that a non-coherent cache is used. But the correctness is tested to be guaranteed with `.L1::no_allocate` on Hopper architectures, and performance will be much better. The reason we guess may be: the non-coherent cache is unified with L1, and the L1 modifier is not just a hint but a strong option, so that the correctness can be guaranteed by no dirty data in L1. +- Initially, because NVCC could not automatically unroll volatile read PTX, we tried using `__ldg` (i.e., `ld.nc`). Even compared to manually unrolled volatile reads, it was significantly faster (likely due to additional compiler optimizations). However, the results could be incorrect or dirty. After consulting the PTX documentation, we discovered that L1 and non-coherent cache are unified on Hopper architectures. We speculated that `.L1::no_allocate` might resolve the issue, leading to this discovery. +- If you find kernels not working on some other platforms, you may add `DISABLE_AGGRESSIVE_PTX_INSTRS=1` to `setup.py` and disable this, or file an issue. + +#### Auto-tuning on your cluster + +For better performance on your cluster, we recommend to run all the tests and use the best auto-tuned configuration. The default configurations are optimized on the DeepSeek's internal cluster. diff --git a/third-party/README.md b/docs/nvshmem.md similarity index 99% rename from third-party/README.md rename to docs/nvshmem.md index 63d855777..d1aeceeb3 100644 --- a/third-party/README.md +++ b/docs/nvshmem.md @@ -74,4 +74,4 @@ export PATH="${NVSHMEM_DIR}/bin:$PATH" ```bash nvshmem-info -a # Should display details of nvshmem -``` +``` \ No newline at end of file diff --git a/requirements-lint.txt b/requirements-lint.txt index 23419552b..eca1cd620 100644 --- a/requirements-lint.txt +++ b/requirements-lint.txt @@ -1,3 +1,3 @@ clang-format==15.0.7 yapf==0.40.2 -ruff==0.6.5 \ No newline at end of file +ruff==0.6.5 diff --git a/setup.py b/setup.py index e89a9dbb4..9b30f9955 100644 --- a/setup.py +++ b/setup.py @@ -1,13 +1,24 @@ +import ast +import re import os import subprocess import setuptools import importlib from pathlib import Path +from setuptools.command.build_py import build_py from torch.utils.cpp_extension import BuildExtension, CUDAExtension +current_dir = os.path.dirname(os.path.realpath(__file__)) +persistent_env_names = ('EP_JIT_CACHE_DIR', 'EP_JIT_PRINT_COMPILER_COMMAND', 'EP_NUM_TOPK_IDX_BITS', 'EP_NCCL_ROOT_DIR') -# Wheel specific: the wheels only include the soname of the host library `libnvshmem_host.so.X` +# Load discover module without triggering `deep_ep.__init__` +find_pkgs_spec = importlib.util.spec_from_file_location('find_pkgs', os.path.join(current_dir, 'deep_ep', 'utils', 'find_pkgs.py')) +find_pkgs = importlib.util.module_from_spec(find_pkgs_spec) +find_pkgs_spec.loader.exec_module(find_pkgs) + + +# Wheel specific: the wheels only include the SO name of the host library `libnvshmem_host.so.X` def get_nvshmem_host_lib_name(base_dir): path = Path(base_dir).joinpath('lib') for file in path.rglob('libnvshmem_host.so.*'): @@ -15,45 +26,80 @@ def get_nvshmem_host_lib_name(base_dir): raise ModuleNotFoundError('libnvshmem_host.so not found') -if __name__ == '__main__': - disable_nvshmem = False - nvshmem_dir = os.getenv('NVSHMEM_DIR', None) - nvshmem_host_lib = 'libnvshmem_host.so' - if nvshmem_dir is None: - try: - nvshmem_dir = importlib.util.find_spec("nvidia.nvshmem").submodule_search_locations[0] - nvshmem_host_lib = get_nvshmem_host_lib_name(nvshmem_dir) - import nvidia.nvshmem as nvshmem # noqa: F401 - except (ModuleNotFoundError, AttributeError, IndexError): - print( - 'Warning: `NVSHMEM_DIR` is not specified, and the NVSHMEM module is not installed. All internode and low-latency features are disabled\n' - ) - disable_nvshmem = True - else: - disable_nvshmem = False +def get_package_version(): + with open(Path(current_dir) / 'deep_ep' / '__init__.py', 'r') as f: + version_match = re.search(r'^__version__\s*=\s*(.*)$', f.read(), re.MULTILINE) + public_version = ast.literal_eval(version_match.group(1)) + + # noinspection PyBroadException + try: + status_cmd = ['git', 'status', '--porcelain'] + status_output = subprocess.check_output(status_cmd).decode('ascii').strip() + if status_output: + print(f'Warning: Git working directory is not clean. Uncommitted changes:\n{status_output}') + assert False, 'Git working directory is not clean' + + cmd = ['git', 'rev-parse', '--short', 'HEAD'] + revision = '+' + subprocess.check_output(cmd).decode('ascii').rstrip() + except: + revision = '+local' + return f'{public_version}{revision}' + + +class CustomBuildPy(build_py): + def run(self): + # Make clusters' cache setting default into `envs.py` + self.generate_default_envs() + + # Finally, run the regular build + build_py.run(self) - if not disable_nvshmem: - assert os.path.exists(nvshmem_dir), f'The specified NVSHMEM directory does not exist: {nvshmem_dir}' + def generate_default_envs(self): + code = '# Pre-installed environment variables\n' + code += 'persistent_envs = dict()\n' + # noinspection PyShadowingNames + for name in persistent_env_names: + code += f"persistent_envs['{name}'] = '{os.environ[name]}'\n" if name in os.environ else '' + # Create temporary build directory + build_include_dir = os.path.join(self.build_lib, 'deep_ep') + os.makedirs(build_include_dir, exist_ok=True) + with open(os.path.join(self.build_lib, 'deep_ep', 'envs.py'), 'w') as f: + f.write(code) + + +if __name__ == '__main__': + # TODO: make NVSHMEM and legacy optional + nvshmem_root_dir = find_pkgs.find_nvshmem_root() + nccl_root_dir = find_pkgs.find_nccl_root() + + # `128,2417` is used to suppress warnings of `fmt` cxx_flags = ['-O3', '-Wno-deprecated-declarations', '-Wno-unused-variable', '-Wno-sign-compare', '-Wno-reorder', '-Wno-attributes'] - nvcc_flags = ['-O3', '-Xcompiler', '-O3'] - sources = ['csrc/deep_ep.cpp', 'csrc/kernels/runtime.cu', 'csrc/kernels/layout.cu', 'csrc/kernels/intranode.cu'] - include_dirs = ['csrc/'] + nvcc_flags = ['-O3', '-Xcompiler', '-O3', '--extended-lambda', '--diag-suppress=128,2417'] + sources = ['csrc/python_api.cpp', 'csrc/kernels/legacy/layout.cu', 'csrc/kernels/legacy/intranode.cu'] + include_dirs = [f'{current_dir}/deep_ep/include', + f'{current_dir}/third-party/fmt/include', + '/usr/local/cuda/include/cccl'] library_dirs = [] nvcc_dlink = [] extra_link_args = ['-lcuda'] # NVSHMEM flags - if disable_nvshmem: - cxx_flags.append('-DDISABLE_NVSHMEM') - nvcc_flags.append('-DDISABLE_NVSHMEM') - else: - sources.extend(['csrc/kernels/internode.cu', 'csrc/kernels/internode_ll.cu']) - include_dirs.extend([f'{nvshmem_dir}/include']) - library_dirs.extend([f'{nvshmem_dir}/lib']) - nvcc_dlink.extend(['-dlink', f'-L{nvshmem_dir}/lib', '-lnvshmem_device']) - extra_link_args.extend([f'-l:{nvshmem_host_lib}', '-l:libnvshmem_device.a', f'-Wl,-rpath,{nvshmem_dir}/lib']) + sources.extend(['csrc/kernels/legacy/internode.cu', 'csrc/kernels/legacy/internode_ll.cu', 'csrc/kernels/backend/nvshmem.cu']) + include_dirs.extend([f'{nvshmem_root_dir}/include']) + library_dirs.extend([f'{nvshmem_root_dir}/lib']) + nvcc_dlink.extend(['-dlink', f'-L{nvshmem_root_dir}/lib', '-lnvshmem_device']) + extra_link_args.extend([f'-l:libnvshmem_host.so', '-l:libnvshmem_device.a', f'-Wl,-rpath,{nvshmem_root_dir}/lib']) + + # NCCL flags + sources.extend(['csrc/kernels/backend/nccl.cu']) + include_dirs.extend([f'{nccl_root_dir}/include']) + extra_link_args.extend([f'-l:libnccl.so', f'-Wl,-rpath,{nccl_root_dir}/lib']) + # CUDA driver sources + sources.extend(['csrc/kernels/backend/cuda_driver.cu']) + + # TODO: remove these if int(os.getenv('DISABLE_SM90_FEATURES', 0)): # Prefer A100 os.environ['TORCH_CUDA_ARCH_LIST'] = os.getenv('TORCH_CUDA_ARCH_LIST', '8.0') @@ -63,7 +109,7 @@ def get_nvshmem_host_lib_name(base_dir): nvcc_flags.append('-DDISABLE_SM90_FEATURES') # Disable internode and low-latency kernels - assert disable_nvshmem + assert False, 'Not implemented' else: # Prefer H800 series os.environ['TORCH_CUDA_ARCH_LIST'] = os.getenv('TORCH_CUDA_ARCH_LIST', '9.0') @@ -81,11 +127,16 @@ def get_nvshmem_host_lib_name(base_dir): cxx_flags.append('-DDISABLE_AGGRESSIVE_PTX_INSTRS') nvcc_flags.append('-DDISABLE_AGGRESSIVE_PTX_INSTRS') + # Legacy environment name + if 'TOPK_IDX_BITS' in os.environ: + assert 'EP_NUM_TOPK_IDX_BITS' not in os.environ + os.environ['EP_NUM_TOPK_IDX_BITS'] = os.environ['TOPK_IDX_BITS'] + # Bits of `topk_idx.dtype`, choices are 32 and 64 - if "TOPK_IDX_BITS" in os.environ: - topk_idx_bits = int(os.environ['TOPK_IDX_BITS']) - cxx_flags.append(f'-DTOPK_IDX_BITS={topk_idx_bits}') - nvcc_flags.append(f'-DTOPK_IDX_BITS={topk_idx_bits}') + if 'EP_NUM_TOPK_IDX_BITS' in os.environ: + num_topk_idx_bits = int(os.environ['EP_NUM_TOPK_IDX_BITS']) + cxx_flags.append(f'-DEP_NUM_TOPK_IDX_BITS={num_topk_idx_bits}') + nvcc_flags.append(f'-DEP_NUM_TOPK_IDX_BITS={num_topk_idx_bits}') # Put them together extra_compile_args = { @@ -103,25 +154,38 @@ def get_nvshmem_host_lib_name(base_dir): print(f' > Compilation flags: {extra_compile_args}') print(f' > Link flags: {extra_link_args}') print(f' > Arch list: {os.environ["TORCH_CUDA_ARCH_LIST"]}') - print(f' > NVSHMEM path: {nvshmem_dir}') + print(f' > NVSHMEM path: {nvshmem_root_dir}') + print(f' > NCCL path: {nccl_root_dir}') + # Print persistent env variables + persistent_envs = [] + for name in persistent_env_names: + if name in os.environ: + persistent_envs.append((name, os.environ[name])) + if len(persistent_envs) > 0: + print(f' > Persistent envs:') + for k, v in persistent_envs: + print(f' > {k}: {v}') print() - # noinspection PyBroadException - try: - cmd = ['git', 'rev-parse', '--short', 'HEAD'] - revision = '+' + subprocess.check_output(cmd).decode('ascii').rstrip() - except Exception as _: - revision = '' - - setuptools.setup(name='deep_ep', - version='1.2.1' + revision, - packages=setuptools.find_packages(include=['deep_ep']), - ext_modules=[ - CUDAExtension(name='deep_ep_cpp', - include_dirs=include_dirs, - library_dirs=library_dirs, - sources=sources, - extra_compile_args=extra_compile_args, - extra_link_args=extra_link_args) - ], - cmdclass={'build_ext': BuildExtension}) + setuptools.setup( + name='deep_ep', + version=get_package_version(), + packages=setuptools.find_packages(include=['deep_ep', 'deep_ep.*']), + package_data={ + 'deep_ep': [ + 'include/deep_ep/**/*', + ] + }, + ext_modules=[ + CUDAExtension(name='deep_ep._C', + include_dirs=include_dirs, + library_dirs=library_dirs, + sources=sources, + extra_compile_args=extra_compile_args, + extra_link_args=extra_link_args) + ], + cmdclass={ + 'build_ext': BuildExtension, + 'build_py': CustomBuildPy + } + ) diff --git a/tests/elastic/test_agrs.py b/tests/elastic/test_agrs.py new file mode 100644 index 000000000..4d416cd35 --- /dev/null +++ b/tests/elastic/test_agrs.py @@ -0,0 +1,193 @@ +import argparse +import math +import random +import torch +import torch.distributed as dist +import numpy as np + +import deep_ep +from deep_ep.utils.envs import init_dist, dist_print + + +def all_gather_ref(shape: tuple, rank_idx: int, num_ranks: int, round_idx: int = 0): + ref_list = [] + for i in range(num_ranks): + torch.manual_seed(42 + round_idx * 43 + i) + ref_list.append(torch.randn(shape, dtype=torch.bfloat16, device='cuda')) + return ref_list[rank_idx], torch.stack(ref_list, dim=0) + + +def generate_stress_ops( + num_ops: int, + num_max_inflight_agrs: int, + shape: tuple, + rank_idx: int, + num_ranks: int, +) -> tuple[list[tuple], tuple[torch.Tensor], tuple[torch.Tensor]]: + tensors, refs = zip(*(all_gather_ref(shape, rank_idx, num_ranks, round_idx=i) for i in range(num_ops)), strict=True) + unprocessed = random.sample(range(num_ops), num_ops) + inflight, ops = [], [('create_session', (-1,))] + limit = num_max_inflight_agrs + while unprocessed or inflight: + max_g = min(len(unprocessed), limit) + choices = [] + if max_g > 0: + choices.append('ag') + if inflight: + choices.append('fetch') + else: + choices.append('destroy') + op = random.choice(choices) + if op == 'ag': + b = tuple(unprocessed[-random.randint(1, max_g):]) + limit -= len(b) + del unprocessed[-len(b):] + inflight.append(b) + ops.append(('ag', b)) + elif op == 'fetch': + ops.append(('fetch', inflight.pop(random.randrange(len(inflight))))) + else: + ops.extend([('destroy_session', (-1,)), ('create_session', (-1,))]) + limit = num_max_inflight_agrs + + ops.append(('destroy_session', (-1,))) + return ops, tensors, refs + + +def do_all_gather(buffer: deep_ep.ElasticBuffer, + is_inplace: bool, is_batched: bool, + tensors: tuple[torch.Tensor, ...], + start_event: torch.cuda.Event | None = None): + # Copy into buffer if inplace + if is_inplace: + ag_tensors = buffer.agrs_get_inplace_tensor(tuple(t.shape for t in tensors), torch.bfloat16) + for x, y in zip(ag_tensors, tensors, strict=True): + x.copy_(y) + else: + ag_tensors = tensors + + # Record event + if start_event is not None: + torch.zeros(int(256e6 // 4), dtype=torch.int, device='cuda') # flush L2 cache + start_event.record() + + # Do all-gather + if is_batched: + *out_tensors, handle = buffer.all_gather(ag_tensors) + return out_tensors, [handle] + else: + out_tensors, handles = [], [] + for t in ag_tensors: + out_tensor, handle = buffer.all_gather(t) + out_tensors.append(out_tensor) + handles.append(handle) + return out_tensors, handles + + +# noinspection PyTypeChecker,PyCallingNonCallable,PyShadowingNames +@torch.inference_mode() +def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + + # Print configs + shape = (32, 64, 2048) + num_bytes_per_tensor = math.prod(shape) * 2 + num_max_inflight_agrs = args.num_max_inflight_agrs + num_max_session_bytes = deep_ep.ElasticBuffer.get_agrs_buffer_size_hint( + group, num_bytes_per_tensor * group.size() * num_max_inflight_agrs) + dist_print(f'Config:\n' + f' > Ranks: {num_ranks}\n' + f' > Shape: {shape}\n' + f' > Max inflight AGRS: {num_max_inflight_agrs}\n', + once_in_node=True) + + # Create buffer + buffer = deep_ep.ElasticBuffer(group, explicitly_destroy=True, num_bytes=num_max_session_bytes) + buffer.agrs_set_config(num_max_session_bytes, num_max_inflight_agrs) + + # Run stress tests + dist_print('Running stress tests:', once_in_node=True) + for seed in range(args.num_stress_iterations): + random.seed(42 + seed) + num_ops = 128 + ops, tensors, refs = generate_stress_ops(num_ops, num_max_inflight_agrs, shape, rank_idx, num_ranks) + results = [None] * num_ops + handles = dict() + torch.cuda.synchronize() + for op, indices in ops: + if op == 'create_session': + buffer.create_agrs_session() + elif op == 'destroy_session': + buffer.destroy_agrs_session() + elif op == 'ag': + is_inplace, is_batched = random.random() < 0.5, random.random() < 0.8 + handles[indices] = do_all_gather(buffer, is_inplace, is_batched, tuple(tensors[i] for i in indices)) + elif op == 'fetch': + out_tensors, wait_handles = handles[indices] + for h in wait_handles: + h() + for out, idx in zip(out_tensors, indices, strict=True): + results[idx] = out.clone() + + for i in range(num_ops): + assert results[i] is not None and torch.equal(results[i], refs[i]), \ + f'Rank {rank_idx}: stress mismatch at seed={seed}, op={i}' + dist_print(f' > Seed {seed} passed ({num_ops} ops)', once_in_node=True) + dist_print(once_in_node=True) + + # Destroy the buffer + dist_print(f'Profiling all-gather:', once_in_node=True) + buffer.destroy() + + # Profiling + num_max_session_bytes = deep_ep.ElasticBuffer.get_agrs_buffer_size_hint( + group, (2 ** 26) * group.size() * num_max_inflight_agrs) + buffer = deep_ep.ElasticBuffer(group, explicitly_destroy=True, num_bytes=num_max_session_bytes) + buffer.agrs_set_config(num_max_session_bytes, num_max_inflight_agrs) + for num_bytes in (2 ** p for p in range(20, 27)): + # Create tensors + shape = (num_bytes // 2, ) + tensors = tuple(torch.randn(shape, dtype=torch.bfloat16, device='cuda') for _ in range(num_max_inflight_agrs)) + + # Tests + for is_inplace in (False, True): + for is_batched in (False, True): + num_tests = 50 + start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] + end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] + torch.cuda.synchronize() + + for i in range(num_tests): + with buffer.agrs_new_session(): + _, wait_handles = do_all_gather(buffer, is_inplace, is_batched, tensors, start_event=start_events[i]) + for h in wait_handles: + h() + end_events[i].record() + torch.cuda.synchronize() + + times = np.array([s.elapsed_time(e) / 1e3 for s, e in zip(start_events, end_events, strict=True)])[1:] + avg_t = np.average(times) + unit = ('MB', 1e6) if num_bytes >= 1e6 else ('KB', 1e3) + bandwidth_info = f', {num_bytes * num_ranks * num_max_inflight_agrs / avg_t / 1e9:.3f} GB/s' if num_ranks > 1 else '' + + dist_print( + f' > Rank: {rank_idx:3}/{num_ranks:3} | ' + f'{num_ranks} x {(num_bytes / unit[1]):.0f} {unit[0]} | ' + f'avg: {avg_t / num_max_inflight_agrs * 1e6:.3f} us' + f'{bandwidth_info}' + f' (inplace={int(is_inplace)}, batched={int(is_batched)})') + dist_print(once_in_node=True) + + # Destroy the runtime and communication group + buffer.destroy() + dist.destroy_process_group() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Test all_gather kernels') + parser.add_argument('--num-processes', type=int, default=8) + parser.add_argument('--num-max-inflight-agrs', type=int, default=4) + parser.add_argument('--num-stress-iterations', type=int, default=4) + args = parser.parse_args() + + torch.multiprocessing.spawn(test, args=(args.num_processes, args), nprocs=args.num_processes) diff --git a/tests/elastic/test_barrier.py b/tests/elastic/test_barrier.py new file mode 100644 index 000000000..44fd51936 --- /dev/null +++ b/tests/elastic/test_barrier.py @@ -0,0 +1,62 @@ +import argparse +import torch +import torch.distributed as dist + +import deep_ep +from deep_ep.utils.envs import init_dist, dist_print +from deep_ep.utils.testing import bench_kineto + + +def test_barrier(buffer: deep_ep.ElasticBuffer, args: argparse.Namespace): + dist_print('Profiling barrier:', once_in_node=True) + num_scaleout_ranks, num_scaleup_ranks = buffer.get_logical_domain_size() + dist_print(f'Config:\n' + f' > Ranks: {num_scaleout_ranks} x {num_scaleup_ranks}\n' + f' > #QPs: {buffer.num_allocated_qps}\n', + once_in_node=True) + + # Test barrier time + def loop_barrier(num_tests=1000): + for i in range(num_tests): + buffer.barrier() + + t = bench_kineto(lambda: loop_barrier(), 'barrier', barrier_comm_profiling=True, barrier=buffer.barrier) + dist_print(f' > EP: {buffer.rank_idx:3}/{buffer.num_ranks:3}, ' + f'barrier time: {t * 1e6:.3f} us') + dist_print(once_in_node=True) + + +# noinspection PyShadowingNames +@torch.inference_mode() +def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + + do_pressure_test = args.do_pressure_test + for i in range(int(1e9) if do_pressure_test else 1): + buffer = deep_ep.ElasticBuffer( + group, num_bytes=2 ** 30, + allow_hybrid_mode=args.allow_hybrid_mode, + num_allocated_qps=args.num_allocated_qps, + explicitly_destroy=True + ) + + # Test barrier + test_barrier(buffer, args) + + # Destroy the runtime and communication group + buffer.destroy() + dist.destroy_process_group() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Test elastic EP barrier performance') + + parser.add_argument('--num-processes', type=int, default=8, help='Number of processes to spawn (default: 8)') + parser.add_argument('--allow-hybrid-mode', type=int, default=1, help='Whether to allow hybrid mode') + parser.add_argument('--num-allocated-qps', type=int, default=8, help='Number of QPs to use (0 means auto)') + parser.add_argument('--do-pressure-test', action='store_true', help='Whether to do pressure test') + args = parser.parse_args() + + # Launch test processes + num_processes = args.num_processes + torch.multiprocessing.spawn(test_loop, args=(num_processes, args), nprocs=num_processes) diff --git a/tests/elastic/test_engram.py b/tests/elastic/test_engram.py new file mode 100644 index 000000000..20ea1784a --- /dev/null +++ b/tests/elastic/test_engram.py @@ -0,0 +1,102 @@ +import argparse +import os +import torch +import torch.distributed as dist + +import deep_ep +from deep_ep.utils.envs import init_dist, dist_print +from deep_ep.utils.testing import bench_kineto + + +# noinspection PyUnboundLocalVariable,PyShadowingNames +@torch.inference_mode() +def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank, num_ranks, group = init_dist(local_rank, num_local_ranks) + num_bytes = deep_ep.ElasticBuffer.get_engram_storage_size_hint( + args.num_entries, args.hidden, args.num_tokens, torch.bfloat16) + + # 1 QP uses 1 SM + num_qps = args.num_qps + if num_qps == 0: + num_qps = torch.cuda.get_device_properties('cuda').multi_processor_count + + # Allocate buffer + dist_print(f'Config:\n' + f' > Ranks: {num_ranks}\n' + f' > QPs: {num_qps}\n' + f' > Entries per rank: {args.num_entries}, hidden: {args.hidden}\n' + f' > Tokens to fetch: {args.num_tokens}\n' + f' > Storage per rank: {args.num_entries * args.hidden * 2 / 1024 / 1024:.1f} MB\n', + once_in_node=True) + buffer = deep_ep.ElasticBuffer( + group, + num_bytes=num_bytes, explicitly_destroy=True, num_allocated_qps=num_qps, + allow_hybrid_mode=False, allow_multiple_reduction=False) + + # Write buffer: each rank writes its own local storage into the NCCL window + local_storage = torch.randn((args.num_entries, args.hidden), dtype=torch.bfloat16, device='cuda') + global_storage = torch.empty((num_ranks * args.num_entries, args.hidden), dtype=torch.bfloat16, device='cuda') + dist.all_gather_into_tensor(global_storage, local_storage, group) + buffer.engram_write(local_storage) + + # Generate random indices to fetch + indices = torch.randint(0, num_ranks * args.num_entries, (args.num_tokens, ), device='cuda', dtype=torch.int) + + # Correctness check + ref_fetched = global_storage[indices] + hook = buffer.engram_fetch(indices) + fetched = hook() + if not args.skip_check: + assert torch.equal(ref_fetched, fetched), f'{(ref_fetched - fetched).abs().max().item()}' + + # Performance test + dist_print('Running performance test ...', once_in_node=True) + msg_bytes = args.hidden * 2 # bfloat16 + num_fetched_bytes = args.num_tokens * msg_bytes + + # Measure fetch + wait (end-to-end) + def fetch_and_wait(): + # noinspection PyShadowingNames + hook = buffer.engram_fetch(indices) + hook() + + issue_t, wait_t = bench_kineto( + fetch_and_wait, + kernel_names=('engram_fetch_impl', 'engram_fetch_wait_impl'), + barrier_comm_profiling=True, + barrier=buffer.barrier, + trace_path=f'{args.dump_profile_traces}/engram_fetch_rank{buffer.rank_idx}.json' if args.dump_profile_traces else None) + mpps = args.num_tokens / (issue_t + wait_t) / 1e6 + + dist_print(f' > Rank {rank:3}/{num_ranks} | ' + f'issue: {issue_t * 1e6:.1f} us, ' + f'wait: {wait_t * 1e6:.1f} us, ' + f'{num_fetched_bytes / (issue_t + wait_t) / 1e9:.1f} GB/s, ' + f'bytes: {num_fetched_bytes / 1024 / 1024:.1f} MB, ' + f'{mpps:.2f} MPPS ({msg_bytes} B/msg)') + + dist_print('', once_in_node=True) + + # Destroy the runtime and communication group + buffer.destroy() + dist.destroy_process_group() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Test engram fetch kernels') + parser.add_argument('--num-processes', type=int, default=4, help='Number of processes to spawn') + parser.add_argument('--num-qps', type=int, default=0, help='Number of QPs used (0 for maximum)') + parser.add_argument('--num-entries', type=int, default=524288, help='Number of entries per rank') + parser.add_argument('--hidden', type=int, default=128, help='Hidden dimension size') + parser.add_argument('--num-tokens', type=int, default=4096, help='Number of tokens to fetch') + parser.add_argument('--skip-check', action='store_true', help='Skip correctness check') + parser.add_argument('--dump-profile-traces', type=str, default='', help='Dump profiling trace JSONs') + args = parser.parse_args() + + # Create dump trace directories + if args.dump_profile_traces: + os.makedirs(args.dump_profile_traces, exist_ok=True) + + # Launch + num_processes = args.num_processes + torch.multiprocessing.spawn(test, args=(num_processes, args), nprocs=num_processes) diff --git a/tests/elastic/test_ep.py b/tests/elastic/test_ep.py new file mode 100644 index 000000000..f749d6939 --- /dev/null +++ b/tests/elastic/test_ep.py @@ -0,0 +1,563 @@ +import argparse +import os +import torch +import torch.distributed as dist +from typing import Union, Tuple, Optional + +import deep_ep +from deep_ep.utils.math import ( + align, count_bytes, calc_diff, + per_token_cast_back, per_token_cast_to_fp8, + safe_div +) +from deep_ep.utils.gate import get_unbalanced_scores +from deep_ep.utils.envs import init_dist, init_seed, dist_print +from deep_ep.utils.refs import dispatch as ref_dispatch +from deep_ep.utils.refs import combine as ref_combine +from deep_ep.utils.refs import generate_pre_combine_data, ordered_accumulate +from deep_ep.utils.testing import bench_kineto + + +# noinspection PyUnusedLocal,PyShadowingNames +def enumerate_ep_modes(): + for do_handle_copy in (1, 0): + for expert_alignment in (128, 1): + for use_fp8_dispatch in (1, 0): + for num_bias in (0, 1, 2): + for with_previous_event in (0, 1): + for async_with_compute_stream in (0, 1): + for allocate_on_comm_stream in ((1, ) if with_previous_event else (0, 1)): + yield (do_handle_copy, expert_alignment, use_fp8_dispatch, num_bias, + with_previous_event, async_with_compute_stream, allocate_on_comm_stream) + + +def launch(buffer: deep_ep.ElasticBuffer, name: str, + with_previous_event: int, async_with_compute_stream: int, + params: dict): + if with_previous_event: + params.update(previous_event=buffer.capture()) + values = getattr(buffer, name)(**params) + values[-1].current_stream_wait() if async_with_compute_stream else () + return values + + +def fold_expanded(expanded: Union[Tuple[torch.Tensor], torch.Tensor], + indices: torch.Tensor, valid_mask: torch.Tensor): + if not isinstance(expanded, torch.Tensor): + return tuple(fold_expanded(t, indices, valid_mask) for t in expanded) + + gathered = expanded[indices] + first_valid_idx = valid_mask.to(torch.int).argmax(dim=1) + folded = gathered[torch.arange(gathered.shape[0], device='cuda'), first_valid_idx] + result = (gathered == folded.unsqueeze(1)).all(dim=-1) + result = result | (~valid_mask) + assert result.all() + return folded + + +# noinspection PyUnboundLocalVariable,PyShadowingNames +def test_dispatch_combine(buffer: deep_ep.ElasticBuffer, args: argparse.Namespace): + # Settings + num_scaleout_ranks, num_scaleup_ranks = buffer.get_logical_domain_size() + num_max_tokens_per_rank, num_tokens, hidden = args.num_tokens, max(1, args.num_tokens - dist.get_rank()), args.hidden + num_topk, num_experts = args.num_topk, args.num_experts + num_local_experts = num_experts // buffer.num_ranks + num_sms = buffer.get_theoretical_num_sms(num_experts, num_topk) if args.num_sms == 0 else args.num_sms + num_qps = buffer.get_theoretical_num_qps(num_sms) if args.num_qps == 0 else args.num_qps + dist_print(f'Config:\n' + f' > Ranks: {num_scaleout_ranks} x {num_scaleup_ranks}\n' + f' > Experts: {num_topk}/{num_experts}\n' + f' > Tokens: {num_tokens} (max: {num_max_tokens_per_rank}), hidden: {hidden}\n' + f' > #SM: {num_sms}, #QPs: {num_qps}/{buffer.num_allocated_qps}\n', + once_in_node=True) + + # Construct expert selections first (may have an unbalanced ratio here) + scores = get_unbalanced_scores(num_tokens, num_experts, buffer.num_ranks, num_topk, args.unbalanced_ratio, args.precise_unbalanced_ratio) + topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False) + topk_idx = topk_idx.to(deep_ep.topk_idx_t) + if args.masked_ratio > 0: + rand_mask = torch.rand_like(topk_idx, dtype=torch.float) + topk_idx.masked_fill_(rand_mask < args.masked_ratio, -1) + topk_weights.masked_fill_(topk_idx < 0, 0) + + # Run all tests + dist_print('Running all test cases:', once_in_node=True) + for (do_handle_copy, expert_alignment, use_fp8_dispatch, num_bias, + with_previous_event, async_with_compute_stream, allocate_on_comm_stream) in enumerate_ep_modes(): + dist_print(f' > Testing with ' + f'{do_handle_copy=}, {expert_alignment=}, {use_fp8_dispatch=}, {num_bias=}, ' + f'{with_previous_event=}, {async_with_compute_stream=}, {allocate_on_comm_stream=} ...', + once_in_node=True) + + # Random data + # TODO: support top-k groups + x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + x = per_token_cast_to_fp8(x) if use_fp8_dispatch else x + bias = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') if num_bias == 1 else None + if num_bias == 2: + bias = tuple(torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') for _ in range(num_bias)) + assert len(bias) == 2 # To prevent linter warning + + # Test correctness with NCCL reference + if not args.skip_check: + ref_recv_x, ref_recv_topk_idx, ref_recv_topk_weights, \ + ref_recv_src_token_idx, ref_num_recv_tokens_per_rank = \ + ref_dispatch(x, topk_idx, topk_weights, num_max_tokens_per_rank, num_experts) + ref_recv_x_bf16 = per_token_cast_back(ref_recv_x[0], ref_recv_x[1]) if use_fp8_dispatch else ref_recv_x + + if args.allow_multiple_reduction: + # Should be the same as the trigger condition of DeepEP's hybrid combine, which performs intra-scaleup reduction first + if args.allow_hybrid_mode and num_scaleout_ranks > 1: + reduced_combine_recipe = (True, True) + combine_recipe = (True, True) + else: + reduced_combine_recipe = (True, False) + combine_recipe = (True, False) + else: + reduced_combine_recipe = (False, False) + combine_recipe = (True, False) + ref_y = generate_pre_combine_data( + dist.get_rank() * num_max_tokens_per_rank + torch.arange(num_tokens, device='cuda'), + num_max_tokens_per_rank, num_topk, hidden) + ref_y[topk_idx == -1] = 0 + ref_reduced_combined_y = ref_combine( + ref_y, topk_idx, + num_scaleout_ranks, num_scaleup_ranks, num_experts, + bias, + *reduced_combine_recipe + ) + ref_combined_y = ref_combine( + ref_y, topk_idx, + num_scaleout_ranks, num_scaleup_ranks, + num_experts, bias, + *combine_recipe + ) # Reduce within rank, then globally, for non-expand combine mode + torch.cuda.synchronize() + + # Do dispatch + dispatch_args = dict( + x=x, topk_idx=topk_idx, topk_weights=topk_weights, + num_sms=num_sms, num_qps=num_qps, + num_max_tokens_per_rank=num_max_tokens_per_rank, num_experts=num_experts, + expert_alignment=expert_alignment, + async_with_compute_stream=async_with_compute_stream, + allocate_on_comm_stream=allocate_on_comm_stream, + do_handle_copy=do_handle_copy, do_cpu_sync=args.do_cpu_sync) + recv_x, recv_topk_idx, recv_topk_weights, handle, dispatch_event = \ + launch(buffer, 'dispatch', with_previous_event, async_with_compute_stream, dispatch_args) + recv_x_bf16 = per_token_cast_back(recv_x[0], recv_x[1]) if use_fp8_dispatch else recv_x + + # Expanding mode + expanded_dispatch_args = dispatch_args | dict(do_expand=True, use_tma_aligned_col_major_sf=True) + expanded_recv_x, expanded_recv_topk_idx, expanded_recv_topk_weights, expanded_handle, expanded_dispatch_event = \ + launch(buffer, 'dispatch', with_previous_event, async_with_compute_stream, expanded_dispatch_args) + expanded_recv_x_bf16 = per_token_cast_back(expanded_recv_x[0], expanded_recv_x[1]) if use_fp8_dispatch else expanded_recv_x + + # Cached mode + cached_dispatch_args = dict( + x=x, + num_sms=num_sms, num_qps=num_qps, + async_with_compute_stream=async_with_compute_stream, + allocate_on_comm_stream=allocate_on_comm_stream, + handle=handle) + cached_recv_x, cached_recv_topk_idx, cached_recv_topk_weights, cached_handle, cached_dispatch_event = \ + launch(buffer, 'dispatch', with_previous_event, async_with_compute_stream, cached_dispatch_args) + + # Count the number of received tokens + num_recv_tokens = handle.psum_num_recv_tokens_per_scaleup_rank[-1].item() + assert num_recv_tokens == expanded_handle.psum_num_recv_tokens_per_scaleup_rank[-1].item(), \ + 'Expand should not affect the number of received tokens.' + num_expanded_tokens = expanded_handle.psum_num_recv_tokens_per_expert[-1].item() + + # Construction the input data for DeepEP combine + src_token_global_idx = handle.recv_src_metadata[:num_recv_tokens, 0] + if not args.skip_check: + sorted_src_token_global_idx = torch.sort(src_token_global_idx).values + assert torch.equal(ref_recv_src_token_idx, sorted_src_token_global_idx), \ + f'{ref_recv_src_token_idx=}, {sorted_src_token_global_idx=}' + local_y = generate_pre_combine_data(src_token_global_idx, num_max_tokens_per_rank, num_topk, hidden) # [num_recv_tokens, topk, hidden] + local_y[recv_topk_idx[:num_recv_tokens] == -1] = 0 + local_reduced_y = ordered_accumulate(local_y) + input_for_combine = torch.empty_like(recv_x_bf16, dtype=torch.bfloat16, device='cuda') + input_for_combine[:num_recv_tokens] = local_reduced_y + + expanded_src_token_global_idx = expanded_handle.recv_src_metadata[:num_recv_tokens, 0] + if not args.skip_check: + sorted_expanded_src_token_global_idx = torch.sort(expanded_src_token_global_idx).values + assert torch.equal(ref_recv_src_token_idx, sorted_expanded_src_token_global_idx), \ + f'{ref_recv_src_token_idx=}, {sorted_expanded_src_token_global_idx=}' + local_y_expand = generate_pre_combine_data(expanded_src_token_global_idx, num_max_tokens_per_rank, num_topk, hidden) # [num_recv_tokens, topk, hidden] + # We put an extra row to conveniently handle the -1 index + input_for_expand_combine = torch.empty((expanded_recv_x_bf16.shape[0] + 1, hidden), dtype=torch.bfloat16, device='cuda') + input_for_expand_combine[expanded_handle.recv_src_metadata[:num_recv_tokens, 2:].flatten()] = local_y_expand.view(-1, hidden) + input_for_expand_combine = input_for_expand_combine[:-1, ...] + + # Do combine + combine_args = dict( + x=input_for_combine, topk_weights=recv_topk_weights, bias=bias, + handle=handle, + num_sms=num_sms, num_qps=num_qps, + async_with_compute_stream=async_with_compute_stream, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + combined_x, combined_topk_weights, combine_event = \ + launch(buffer, 'combine', with_previous_event, async_with_compute_stream, combine_args) + + # Reduced combine + reduced_combine_args = dict( + x=input_for_expand_combine, bias=bias, + handle=expanded_handle, + num_sms=num_sms, num_qps=num_qps, + async_with_compute_stream=async_with_compute_stream, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + reduced_combined_x, reduced_combined_topk_weights, reduced_combine_event = \ + launch(buffer, 'combine', with_previous_event, async_with_compute_stream, reduced_combine_args) + + assert not (args.dump_profile_traces and args.skip_perf_test), '`--skip-perf-test` should not be specified when `--dump-profile-traces` is provided' + if not args.skip_perf_test: + # Profiling + def get_trace_path(prefix: str): + return None if not args.dump_profile_traces else f'{args.dump_profile_traces}/{prefix}_rank{buffer.rank_idx}.json' + + # Calculate the number of tokens that are sent to the other scaleout peers + dst_scaleout_rank_idx = topk_idx // (num_experts // num_scaleout_ranks) + num_scaleout_send_tokens = 0 + for i in range(num_scaleout_ranks if num_scaleout_ranks > 1 else 0): + if args.ignore_local_traffic and i == dist.get_rank() // num_scaleup_ranks: + continue + num_scaleout_send_tokens += (dst_scaleout_rank_idx == i).any(dim=1).sum().item() + + # Calculate the number of tokens that are received via the other scaleup peers + num_scaleup_recv_tokens = num_recv_tokens + if args.ignore_local_traffic: + num_scaleup_recv_tokens -= (src_token_global_idx // num_max_tokens_per_rank % num_scaleup_ranks == dist.get_rank() % num_scaleup_ranks).sum().item() + + # Test dispatch performance + num_bytes_per_dispatch_token = safe_div(count_bytes(recv_x, recv_topk_idx, recv_topk_weights), recv_topk_idx.size(0)) + num_scaleup_bytes = num_bytes_per_dispatch_token * num_scaleup_recv_tokens # Received via scaleup + num_scaleout_bytes = num_bytes_per_dispatch_token * num_scaleout_send_tokens # Send via scaleout + t, copy_t = bench_kineto(lambda: buffer.dispatch(**dispatch_args), + kernel_names=('dispatch_impl', 'dispatch_copy_epilogue_impl'), + barrier_comm_profiling=True, barrier=buffer.barrier, trace_path=get_trace_path('dispatch')) + dist_print(f' * EP: {buffer.rank_idx:3}/{buffer.num_ranks} | ' + f'dispatch: ' + f'{num_scaleout_bytes / t / 1e9:.0f} GB/s (SO), ' + f'{num_scaleup_bytes / t / 1e9:.0f} GB/s (SU), {t * 1e6:.3f} us, {num_scaleup_bytes:.0f} bytes | ' + f'copy: {2 * num_recv_tokens * num_bytes_per_dispatch_token / copy_t / 1e9:.0f} GB/s, {copy_t * 1e6:.3f} us') + + # Test expanded dispatch performance + num_bytes_per_dispatch_token_meta = safe_div(count_bytes(expanded_handle.recv_src_metadata), expanded_handle.recv_src_metadata.size(0)) + t, copy_t = bench_kineto(lambda: buffer.dispatch(**expanded_dispatch_args), + kernel_names=('dispatch_impl', 'dispatch_copy_epilogue_impl'), + barrier_comm_profiling=True, barrier=buffer.barrier, trace_path=get_trace_path('expanded_dispatch')) + dist_print(f' - EP: {buffer.rank_idx:3}/{buffer.num_ranks} | ' + f'expanded dispatch: ' + f'{num_scaleout_bytes / t / 1e9:.0f} GB/s (SO), ' + f'{num_scaleup_bytes / t / 1e9:.0f} GB/s (SU), {t * 1e6:.3f} us, {num_scaleup_bytes:.0f} bytes | ' + f'copy: {(num_recv_tokens * (num_bytes_per_dispatch_token_meta + num_bytes_per_dispatch_token) + num_expanded_tokens * num_bytes_per_dispatch_token) / copy_t / 1e9:.0f} GB/s, {copy_t * 1e6:.3f} us') + + # Test cached dispatch performance + t, copy_t = bench_kineto(lambda: buffer.dispatch(**cached_dispatch_args), + kernel_names=('dispatch_impl', 'dispatch_copy_epilogue_impl'), + barrier_comm_profiling=True, barrier=buffer.barrier, trace_path=get_trace_path('cached_dispatch')) + dist_print(f' # EP: {buffer.rank_idx:3}/{buffer.num_ranks} | ' + f'cached dispatch: ' + f'{num_scaleout_bytes / t / 1e9:.0f} GB/s (SO), ' + f'{num_scaleup_bytes / t / 1e9:.0f} GB/s (SU), {t * 1e6:.3f} us, {num_scaleup_bytes:.0f} bytes | ' + f'copy: {2 * num_scaleup_bytes / copy_t / 1e9:.0f} GB/s, {copy_t * 1e6:.3f} us') + + # Test combine performance + num_bytes_per_combine_token = safe_div(count_bytes(recv_x_bf16, recv_topk_weights), recv_x_bf16.size(0)) + num_bias_bytes = count_bytes(bias) + num_reduction_write_bytes = count_bytes(combined_x, combined_topk_weights) + + def get_combine_bytes(is_expand_mode: bool) -> Tuple[float, float, float]: + num_experts_per_rank = num_experts // (num_scaleup_ranks * num_scaleout_ranks) + num_experts_per_scaleout_rank = num_experts_per_rank * num_scaleup_ranks + + def get_unique_and_valid_dst_count(dst_idx: torch.Tensor, + ignored_nums_l: Optional[int] = None, ignored_nums_r: Optional[int] = None, + max_num_in_dst_idx: int = num_experts - 1) -> int: + """ + Get the number of valid destinations, with deduplication within each token and numbers within `[ignored_nums_l, ignored_nums_r)` being ignored + """ + dst_idx = dst_idx.clone() + ignore_mask = dst_idx == -1 + if args.ignore_local_traffic and ignored_nums_l is not None: + assert ignored_nums_r is not None + ignore_mask |= ((dst_idx >= ignored_nums_l) & (dst_idx < ignored_nums_r)) + dst_idx = dst_idx + torch.arange(0, dst_idx.shape[0], dtype=dst_idx.dtype, device=dst_idx.device).unsqueeze(-1) * (max_num_in_dst_idx + 1) # So that different rows will have different values + dst_idx[ignore_mask] = dst_idx[0][0].item() # So that these `-1`s won't affect the count of unique numbers + return torch.unique(dst_idx, sorted=False).numel() + + if not args.allow_multiple_reduction: + # No multiple reduction + if not is_expand_mode: + num_scaleup_tokens = num_scaleup_recv_tokens + num_scaleout_tokens = get_unique_and_valid_dst_count( + topk_idx // num_experts_per_rank, buffer.scaleout_rank_idx * num_scaleup_ranks, (buffer.scaleout_rank_idx + 1) * num_scaleup_ranks) + num_reduction_read_tokens = get_unique_and_valid_dst_count(topk_idx // num_experts_per_rank) + else: + tokens_src_rank_idx = src_token_global_idx//num_max_tokens_per_rank + if args.ignore_local_traffic: + num_scaleup_tokens = (recv_topk_idx[:num_recv_tokens] != -1)[tokens_src_rank_idx % num_scaleup_ranks != buffer.scaleup_rank_idx].sum().item() + else: + num_scaleup_tokens = (recv_topk_idx[:num_recv_tokens] != -1).sum().item() + num_scaleout_tokens = get_unique_and_valid_dst_count( + topk_idx, buffer.scaleout_rank_idx * num_experts_per_scaleout_rank, (buffer.scaleout_rank_idx + 1) * num_experts_per_scaleout_rank) + num_reduction_read_tokens = get_unique_and_valid_dst_count(topk_idx) + else: + # With `allow_multiple_reduction`, "combine" has exactly the same number of tokens as "dispatch" + num_scaleup_tokens = num_scaleup_recv_tokens + num_scaleout_tokens = num_scaleout_send_tokens + if args.allow_hybrid_mode: + num_reduction_read_tokens = get_unique_and_valid_dst_count(topk_idx // num_experts_per_scaleout_rank) + else: + num_reduction_read_tokens = get_unique_and_valid_dst_count(topk_idx // num_experts_per_rank) + if not args.ignore_local_traffic and num_scaleout_ranks == 1: + num_scaleout_tokens = 0 + return num_scaleout_tokens * num_bytes_per_combine_token, num_scaleup_tokens * num_bytes_per_combine_token, num_reduction_read_tokens * num_bytes_per_combine_token + + num_scaleout_bytes, num_scaleup_bytes, num_reduction_read_bytes = get_combine_bytes(False) + t, copy_t = bench_kineto(lambda: buffer.combine(**combine_args), + kernel_names=('combine_impl', 'combine_reduce_epilogue_impl'), + barrier_comm_profiling=True, barrier=buffer.barrier, trace_path=get_trace_path('combine')) + dist_print(f' @ EP: {buffer.rank_idx:3}/{buffer.num_ranks} | ' + f'combine: ' + f'{num_scaleout_bytes / t / 1e9:.0f} GB/s (SO), ' + f'{num_scaleup_bytes / t / 1e9:.0f} GB/s (SU), {t * 1e6:.3f} us, {num_scaleup_bytes:.0f} bytes | ' + f'reduce: {(num_bias_bytes + num_reduction_read_bytes + num_reduction_write_bytes) / copy_t / 1e9:.0f} GB/s, {copy_t * 1e6:.3f} us') + + # Test reduced combine performance + num_scaleout_bytes, num_scaleup_bytes, num_reduction_read_bytes = get_combine_bytes(True) + t, copy_t = bench_kineto(lambda: buffer.combine(**reduced_combine_args), + kernel_names=('combine_impl', 'combine_reduce_epilogue_impl'), + barrier_comm_profiling=True, barrier=buffer.barrier, trace_path=get_trace_path('reduced_combine')) + dist_print(f' + EP: {buffer.rank_idx:3}/{buffer.num_ranks} | ' + f'reduced combine: ' + f'{num_scaleout_bytes / t / 1e9:.0f} GB/s (SO), ' + f'{num_scaleup_bytes / t / 1e9:.0f} GB/s (SU), {t * 1e6:.3f} us, {num_scaleup_bytes:.0f} bytes | ' + f'reduce: {(num_bias_bytes + num_reduction_read_bytes + num_reduction_write_bytes) / copy_t / 1e9:.0f} GB/s, {copy_t * 1e6:.3f} us') + dist_print(once_in_node=True) + + # Checks + # NOTES: we do checks after the performance tests, as we may modify some tensors + if not args.skip_check: + # Handle copy checks + assert (topk_idx.data_ptr() != handle.topk_idx.data_ptr()) == do_handle_copy + assert (topk_idx.data_ptr() != cached_handle.topk_idx.data_ptr()) == do_handle_copy + assert handle.topk_idx.data_ptr() == cached_handle.topk_idx.data_ptr() + + # Make the valid part of the whole tensor for no CPU sync mode + if not args.do_cpu_sync: + if use_fp8_dispatch: + recv_x = (recv_x[0][:num_recv_tokens], recv_x[1][:num_recv_tokens]) + cached_recv_x = (cached_recv_x[0][:num_recv_tokens], cached_recv_x[1][:num_recv_tokens]) + else: + recv_x = recv_x[:num_recv_tokens] + cached_recv_x = cached_recv_x[:num_recv_tokens] + recv_x_bf16 = recv_x_bf16[:num_recv_tokens] + recv_topk_idx = recv_topk_idx[:num_recv_tokens] + recv_topk_weights = recv_topk_weights[:num_recv_tokens] + cached_recv_topk_idx = cached_recv_topk_idx[:num_recv_tokens] + handle.recv_src_metadata = handle.recv_src_metadata[:num_recv_tokens] + expanded_handle.recv_src_metadata = expanded_handle.recv_src_metadata[:num_recv_tokens] + + # Make sure deterministic mode works by doing the dispatch twice + if args.deterministic: + recv_x_twice, recv_topk_idx_twice, recv_topk_weights_twice, handle_twice, dispatch_event_twice = \ + launch(buffer, 'dispatch', with_previous_event, async_with_compute_stream, dispatch_args) + if not args.do_cpu_sync: + assert num_recv_tokens == handle_twice.psum_num_recv_tokens_per_scaleup_rank[-1].item() + handle_twice.recv_src_metadata = handle_twice.recv_src_metadata[:num_recv_tokens] + assert torch.equal(handle.recv_src_metadata[:, :2], handle_twice.recv_src_metadata[:, :2]) + + # Test cumulative stats counter + cumulative_local_expert_recv_stats = torch.zeros((num_local_experts, ), dtype=torch.int, device='cuda') + dispatch_args['cumulative_local_expert_recv_stats'] = cumulative_local_expert_recv_stats + launch(buffer, 'dispatch', with_previous_event, async_with_compute_stream, dispatch_args) + + # Expanded checks + assert expanded_recv_topk_idx is None + assert expanded_handle.recv_src_metadata.size(0) == num_recv_tokens + expanded_indices = expanded_handle.recv_src_metadata[:, 2:] + expanded_mask = expanded_indices >= 0 + expanded_safe_indices = expanded_indices.clone() + expanded_safe_indices[~expanded_mask] = 0 + expanded_recv_x = fold_expanded(expanded_recv_x, expanded_safe_indices, expanded_mask) + expanded_recv_topk_weights = expanded_recv_topk_weights[expanded_safe_indices] + + # Cached checks + if use_fp8_dispatch: + assert torch.equal(recv_x[0], cached_recv_x[0]) + assert torch.equal(recv_x[1], cached_recv_x[1]) + else: + assert torch.equal(recv_x, cached_recv_x) + assert torch.equal(recv_topk_idx, cached_recv_topk_idx) + assert torch.equal(handle.dst_buffer_slot_idx, cached_handle.dst_buffer_slot_idx) + assert torch.equal(handle.psum_num_recv_tokens_per_scaleup_rank, cached_handle.psum_num_recv_tokens_per_scaleup_rank) + assert handle.num_recv_tokens_per_expert_list == cached_handle.num_recv_tokens_per_expert_list + + # Check dispatch expert count + assert recv_x_bf16.size() == ref_recv_x_bf16.size(), f'{recv_x_bf16.size()=}, {ref_recv_x_bf16.size()=}' + assert recv_x_bf16.size(0) == num_recv_tokens + for i in range(num_local_experts if args.do_cpu_sync else 0): + ref_count = (ref_recv_topk_idx == i).sum().item() + aligned_ref_count = align(ref_count, expert_alignment) + assert ref_count == cumulative_local_expert_recv_stats[i].item(),\ + f'{i}, {ref_count}, {cumulative_local_expert_recv_stats[i].item()}' + assert aligned_ref_count == handle.num_recv_tokens_per_expert_list[i] + psum_num_recv_tokens_per_expert_list = [0] + handle.psum_num_recv_tokens_per_expert.tolist() + expanded_psum_num_recv_tokens_per_expert_list = [0] + expanded_handle.psum_num_recv_tokens_per_expert.tolist() + for i in range(num_local_experts): + ref_count = (ref_recv_topk_idx == i).sum().item() + count = psum_num_recv_tokens_per_expert_list[i + 1] - psum_num_recv_tokens_per_expert_list[i] + expanded_count = (expanded_psum_num_recv_tokens_per_expert_list[i + 1] - + align(expanded_psum_num_recv_tokens_per_expert_list[i], expert_alignment)) + assert align(ref_count, expert_alignment) == count, f'{buffer.rank_idx=}, {i=}, {ref_count=}, {count=}' + assert ref_count == expanded_count, f'{ref_count=}, {expanded_count=}' + + # Check dispatch scale-up received token psum + psum_num_recv_tokens_per_scaleup_rank_list = [0] + handle.psum_num_recv_tokens_per_scaleup_rank.tolist() + for i in range(num_scaleup_ranks): + count = psum_num_recv_tokens_per_scaleup_rank_list[i + 1] - psum_num_recv_tokens_per_scaleup_rank_list[i] + ref_count = sum(ref_num_recv_tokens_per_rank[i::num_scaleup_ranks]) + assert count == ref_count, f'{ref_count=}, {count=}' + + # Check dispatch data + for check_recv_x, check_recv_topk_idx, check_recv_topk_weights, check_handle in ( + (expanded_recv_x, None, expanded_recv_topk_weights, expanded_handle), # Expanded + (recv_x, recv_topk_idx, recv_topk_weights, handle), # Unexpanded + ): + for i in range(buffer.num_ranks): + rank_start_idx = sum(ref_num_recv_tokens_per_rank[:i]) + rank_end_idx = rank_start_idx + ref_num_recv_tokens_per_rank[i] + sorted_metadata = torch.sort(check_handle.recv_src_metadata[:, 0]) + sorted_indices = sorted_metadata.indices[rank_start_idx:rank_end_idx] + sorted_values = sorted_metadata.values[rank_start_idx:rank_end_idx] + assert torch.equal(ref_recv_src_token_idx[rank_start_idx:rank_end_idx], sorted_values) + + # Data should be bitwise identical + check_list = [(ref_recv_topk_weights, check_recv_topk_weights, True)] + if check_recv_topk_idx is not None: + check_list.append((ref_recv_topk_idx, check_recv_topk_idx, False)) + if use_fp8_dispatch: + check_list.append((ref_recv_x[0], check_recv_x[0], False)) + check_list.append((ref_recv_x[1], check_recv_x[1], False)) + else: + check_list.append((ref_recv_x, check_recv_x, False)) + ref_mask = ref_recv_topk_idx[rank_start_idx:rank_end_idx] < 0 + for ref_t, t, do_mask in check_list: + ref_t = ref_t[rank_start_idx:rank_end_idx] + t = t[sorted_indices] + if do_mask: + ref_t = ref_t.masked_fill(ref_mask, 0) + t = t.masked_fill(ref_mask, 0) + assert torch.equal(ref_t, t), f'{ref_t=}, {t=}' + + # Combined data should also be bitwise-identical + assert torch.equal(combined_x, ref_combined_y), \ + f'Diff: {calc_diff(combined_x, ref_combined_y)}' + assert torch.equal(reduced_combined_x, ref_reduced_combined_y), \ + f'Diff: {calc_diff(reduced_combined_x, ref_reduced_combined_y)}' + assert torch.equal(combined_topk_weights, topk_weights), \ + f'{calc_diff(combined_topk_weights, topk_weights)}' + + # Break on the first test case + if args.test_first_only: + break + dist_print('', once_in_node=True) + + +# noinspection PyUnboundLocalVariable,PyShadowingNames +@torch.inference_mode() +def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks, seed=args.seed) + def construct_elastic_buffer(): + return deep_ep.ElasticBuffer(group, + num_max_tokens_per_rank=args.num_tokens, hidden=args.hidden, + deterministic=args.deterministic, + allow_hybrid_mode=args.allow_hybrid_mode, + allow_multiple_reduction=args.allow_multiple_reduction, + prefer_overlap_with_compute=bool(args.prefer_overlap_with_compute), + sl_idx=args.sl_idx, + num_allocated_qps=max(args.num_allocated_qps, args.num_qps), + explicitly_destroy=True, + num_gpu_timeout_secs=args.num_gpu_timeout_secs, + num_cpu_timeout_secs=args.num_cpu_timeout_secs) + + buffer = construct_elastic_buffer() + + # Warning in case of precise unbalanced ratio + if args.precise_unbalanced_ratio: + dist_print('\033[33mWarning: Using precise unbalanced ratio mode. ' + 'Test data is manually constructed and may differ from real world distribution.\033[0m', + once_in_node=True) + + # Test MoE kernels + test_dispatch_combine(buffer, args) + + # Pressure tests + for seed in range(int(1e9) if args.do_pressure_test else 0): + if not args.reuse_elastic_buffer: + # Recreate elastic buffer + buffer.destroy() + buffer = construct_elastic_buffer() + + assert not args.skip_check + dist_print(f'Testing with {seed=} ...', once_in_node=True) + init_seed(seed) + test_dispatch_combine(buffer, args) + + # Destroy the runtime and communication group + buffer.destroy() + dist.destroy_process_group() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Test elastic EP kernels') + + # Resource settings + parser.add_argument('--num-processes', type=int, default=8, help='Number of processes to spawn (default: 8)') + parser.add_argument('--num-sms', type=int, default=0, help='Number of SMs to use (0 means auto)') + parser.add_argument('--num-qps', type=int, default=0, help='Number of QPs to use (0 means auto)') + parser.add_argument('--num-allocated-qps', type=int, default=0, help='Number of QPs to allocate (0 means auto)') + parser.add_argument('--num-gpu-timeout-secs', type=int, default=100, help='Timeout in seconds (GPU side)') + parser.add_argument('--num-cpu-timeout-secs', type=int, default=100, help='Timeout in seconds (CPU side)') + parser.add_argument('--sl-idx', type=int, default=0, help='SL index') + + # Model settings + parser.add_argument('--num-tokens', type=int, default=4096, help='Number of tokens') + parser.add_argument('--hidden', type=int, default=7168, help='Hidden dimension size') + parser.add_argument('--num-topk', type=int, default=6, help='Number of top-k experts') + parser.add_argument('--num-experts', type=int, default=256, help='Number of experts') + + # Scenario settings + parser.add_argument('--do-cpu-sync', type=int, default=1, help='Whether to do CPU sync') + parser.add_argument('--allow-hybrid-mode', type=int, default=1, help='Whether to allow hybrid mode') + parser.add_argument('--allow-multiple-reduction', type=int, default=1, help='Whether to allow multiple reductions') + parser.add_argument('--prefer-overlap-with-compute', type=int, default=0, help='Whether to prefer overlap with compute') + parser.add_argument('--deterministic', action='store_true', help='Use deterministic algorithm') + + # Test settings + parser.add_argument('--seed', type=int, default=0, help='Default seed for pressure tests') + parser.add_argument('--skip-check', action='store_true', help='Whether to skip correctness checks') + parser.add_argument('--skip-perf-test', action='store_true', help='Whether to skip performance tests') + parser.add_argument('--do-pressure-test', action='store_true', help='Whether to do pressure test') + parser.add_argument('--reuse-elastic-buffer', action='store_true', help='Whether to reuse elastic buffer for each test') + parser.add_argument('--test-first-only', action='store_true', help='Only test the first case') + parser.add_argument('--unbalanced-ratio', type=float, default=1.0, help='The MoE unbalanced ratio') + parser.add_argument('--precise-unbalanced-ratio', action='store_true', help='Generate topk index with precise unbalanced ratio') + parser.add_argument('--masked-ratio', type=float, default=0.0, help='Mask some expert selections') + parser.add_argument('--dump-profile-traces', type=str, default='', help='Dump profiling trace JSONs') + parser.add_argument('--ignore-local-traffic', action='store_true', help='Whether to ignore local traffic during bandwidth calculation') + args = parser.parse_args() + + # Create dump trace directories + if args.dump_profile_traces: + os.makedirs(args.dump_profile_traces, exist_ok=True) + + # Launch test processes + num_processes = args.num_processes + torch.multiprocessing.spawn(test_loop, args=(num_processes, args), nprocs=num_processes) diff --git a/tests/elastic/test_pp.py b/tests/elastic/test_pp.py new file mode 100644 index 000000000..e572e1027 --- /dev/null +++ b/tests/elastic/test_pp.py @@ -0,0 +1,139 @@ +import argparse +import math +import os +import random +import torch +import torch.distributed as dist + +import deep_ep +from deep_ep.utils.envs import init_dist, dist_print, get_rdma_gbs +from deep_ep.utils.testing import bench_kineto + + +def generate_stress_ops(rank_idx: int, num_ranks: int, num_sends: int, shape: tuple): + send_times = {(s, d): [] for s in range(num_ranks) for d in range(num_ranks) if s != d} + recv_times = {(s, d): [] for s in range(num_ranks) for d in range(num_ranks) if s != d} + + for _ in range(num_sends): + src_rank_idx = random.randint(0, num_ranks - 1) + dst_rank_idx = (src_rank_idx + (1 if random.randint(0, 1) else -1)) % num_ranks + st = random.randint(0, 10 ** 8) + rt = st + random.randint(1, 3 * 10 ** 6) + send_times[(src_rank_idx, dst_rank_idx)].append(st) + recv_times[(src_rank_idx, dst_rank_idx)].append(rt) + + ops = [] + for (src_rank_idx, dst_rank_idx) in send_times: + n = len(send_times[(src_rank_idx, dst_rank_idx)]) + sorted_send = sorted(send_times[(src_rank_idx, dst_rank_idx)]) + sorted_recv = sorted(recv_times[(src_rank_idx, dst_rank_idx)]) + for i in range(n): + tensor = torch.randn(shape, dtype=torch.bfloat16, device='cuda') + if src_rank_idx == rank_idx: + ops.append(('send', sorted_send[i], dst_rank_idx, i, tensor)) + if dst_rank_idx == rank_idx: + ops.append(('recv', sorted_recv[i], src_rank_idx, i, tensor)) + ops.sort(key=lambda x: (x[1], x[3])) + return ops + + +# noinspection PyShadowingNames +@torch.inference_mode() +def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + shape = (args.num_tokens, args.hidden) + num_max_tensor_bytes = math.prod(shape) * 2 + num_max_inflight_tensors = args.num_max_inflight_tensors + buffer = deep_ep.ElasticBuffer( + group, explicitly_destroy=True, allow_hybrid_mode=False, + num_bytes=deep_ep.ElasticBuffer.get_pp_buffer_size_hint( + num_max_tensor_bytes, num_max_inflight_tensors)) + buffer.pp_set_config(num_max_tensor_bytes, num_max_inflight_tensors) + + # Print configs + assert num_ranks > 1 + dist_print(f'Config:\n' + f' > Ranks: {num_ranks}\n' + f' > Shape: {shape}\n' + f' > Max inflight tensors: {num_max_inflight_tensors}\n', + once_in_node=True) + + # Run stress tests + dist_print('Running stress tests:', once_in_node=True) + for seed in range(args.num_stress_iterations): + dist_print(f' > Testing with {seed=} ...', once_in_node=True) + torch.manual_seed(42 + seed) + random.seed(42 + seed) + ops = generate_stress_ops(rank_idx, num_ranks, args.num_sends, shape) + + prev = 0 + for j, (op, timestamp, peer, _, tensor) in enumerate(ops): + if op == 'send': + buffer.pp_send(tensor, peer) + else: + result = torch.empty_like(tensor) + buffer.pp_recv(result, peer) + assert torch.equal(result, tensor), \ + f'Rank {rank_idx}: mismatch at op {j}' + if timestamp > prev: + torch.cuda._sleep(int((timestamp - prev) / 10 ** 8 * args.num_sleep_cycles)) + prev = timestamp + dist_print(' > All stress tests passed', once_in_node=True) + dist_print(once_in_node=True) + + # Profiling + dist_print('Profiling PP send and recv:', once_in_node=True) + num_approx_rdma_cycles = int(num_max_tensor_bytes * 2 / get_rdma_gbs() * 1.5) + + def get_trace_path(prefix: str): + return (None if not args.dump_profile_traces + else f'{args.dump_profile_traces}/{prefix}_rank{rank_idx}.json') + + for hide_rdma_latency in (True, False): + for num_concurrent in (1, 2, 3): + send_tensors = [torch.randn(shape, dtype=torch.bfloat16, device='cuda') for _ in range(num_concurrent)] + recv_tensors = [torch.empty(shape, dtype=torch.bfloat16, device='cuda') for _ in range(num_concurrent)] + + def loop(_hide_rdma_latency=hide_rdma_latency): + torch.zeros((131072, 32768), dtype=torch.int, device='cuda') + for t in send_tensors: + buffer.pp_send(t, (rank_idx + 1) % num_ranks) + if _hide_rdma_latency: + torch.cuda._sleep(num_approx_rdma_cycles * num_concurrent) + for t in recv_tensors: + buffer.pp_recv(t, (rank_idx - 1) % num_ranks) + + send_t, recv_t = bench_kineto( + loop, kernel_names=('send_impl', 'recv_impl'), + barrier_comm_profiling=True, barrier=buffer.barrier, + trace_path=get_trace_path(f'pp_{num_concurrent}_{hide_rdma_latency}')) + dist_print( + f' > EP: {rank_idx:3}/{num_ranks:3} | ' + f'hide={int(hide_rdma_latency)}, concurrent={num_concurrent} | ' + f'send: {send_t * 1e6:.3f} us, ' + f'{2 * num_max_tensor_bytes / send_t / 1e9:.3f} GB/s | ' + f'recv: {recv_t * 1e6:.3f} us, ' + f'{(2 if hide_rdma_latency else 1) * num_max_tensor_bytes / recv_t / 1e9:.3f} GB/s') + dist_print(once_in_node=True) + + # Destroy the runtime and communication group + buffer.destroy() + dist.destroy_process_group() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Test PP send/recv kernels') + parser.add_argument('--num-processes', type=int, default=4) + parser.add_argument('--num-tokens', type=int, default=4096) + parser.add_argument('--hidden', type=int, default=7168) + parser.add_argument('--num-max-inflight-tensors', type=int, default=4) + parser.add_argument('--num-stress-iterations', type=int, default=4) + parser.add_argument('--num-sends', type=int, default=128) + parser.add_argument('--num-sleep-cycles', type=int, default=10 ** 7) + parser.add_argument('--dump-profile-traces', type=str, default='') + args = parser.parse_args() + + if args.dump_profile_traces: + os.makedirs(args.dump_profile_traces, exist_ok=True) + + torch.multiprocessing.spawn(test, args=(args.num_processes, args), nprocs=args.num_processes) diff --git a/tests/test_internode.py b/tests/legacy/test_internode.py similarity index 98% rename from tests/test_internode.py rename to tests/legacy/test_internode.py index 6530669da..6c4a9b795 100644 --- a/tests/test_internode.py +++ b/tests/legacy/test_internode.py @@ -6,7 +6,9 @@ # noinspection PyUnresolvedReferences import deep_ep -from utils import init_dist, bench, bench_kineto, calc_diff, create_grouped_scores, inplace_unique, per_token_cast_to_fp8, per_token_cast_back, hash_tensor +from deep_ep.utils.envs import init_dist +from deep_ep.utils.math import calc_diff, create_grouped_scores, inplace_unique, per_token_cast_to_fp8, per_token_cast_back, hash_tensor +from deep_ep.utils.testing import bench, bench_kineto # Test compatibility with low latency functions import test_low_latency diff --git a/tests/test_intranode.py b/tests/legacy/test_intranode.py similarity index 98% rename from tests/test_intranode.py rename to tests/legacy/test_intranode.py index 487491372..5a1730fbf 100644 --- a/tests/test_intranode.py +++ b/tests/legacy/test_intranode.py @@ -5,7 +5,9 @@ # noinspection PyUnresolvedReferences import deep_ep -from utils import init_dist, bench, calc_diff, inplace_unique, per_token_cast_to_fp8, per_token_cast_back +from deep_ep.utils.envs import init_dist +from deep_ep.utils.math import calc_diff, inplace_unique, per_token_cast_to_fp8, per_token_cast_back +from deep_ep.utils.testing import bench # Test compatibility with low latency functions import test_low_latency @@ -276,8 +278,7 @@ def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Namespace): low_latency_mode=test_ll_compatibility, num_qps_per_rank=(ll_num_experts // num_ranks if test_ll_compatibility else 1), explicitly_destroy=True, - allow_mnnvl=args.allow_mnnvl, - use_fabric=args.use_fabric) + allow_mnnvl=args.allow_mnnvl) torch.manual_seed(rank) for i in (24, ): @@ -304,7 +305,6 @@ def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Namespace): parser.add_argument('--num-topk', type=int, default=8, help='Number of top-k experts (default: 8)') parser.add_argument('--num-experts', type=int, default=256, help='Number of experts (default: 256)') parser.add_argument('--allow-mnnvl', action="store_true", help='Enable MNNVL support') - parser.add_argument('--use-fabric', action="store_true", help='Enable fabric mode') args = parser.parse_args() num_processes = args.num_processes diff --git a/tests/test_low_latency.py b/tests/legacy/test_low_latency.py similarity index 99% rename from tests/test_low_latency.py rename to tests/legacy/test_low_latency.py index 456dcf272..99a78bc8a 100644 --- a/tests/test_low_latency.py +++ b/tests/legacy/test_low_latency.py @@ -6,7 +6,9 @@ from typing import Literal, Set import deep_ep -from utils import init_dist, bench, bench_kineto, calc_diff, hash_tensor, per_token_cast_back +from deep_ep.utils.envs import init_dist +from deep_ep.utils.math import calc_diff, per_token_cast_back, hash_tensor +from deep_ep.utils.testing import bench, bench_kineto def simulate_failure_and_skip(rank: int, api: Literal["dispatch", "combine", "clean"], expected_masked_ranks: Set[int]): diff --git a/tests/utils.py b/tests/utils.py deleted file mode 100644 index 1390b2b9a..000000000 --- a/tests/utils.py +++ /dev/null @@ -1,242 +0,0 @@ -import inspect -import json -import tempfile -from pathlib import Path - -import numpy as np -import os -import sys -import torch -import torch.distributed as dist -from typing import Optional, Union - - -def init_dist(local_rank: int, num_local_ranks: int): - # NOTES: you may rewrite this function with your own cluster settings - ip = os.getenv('MASTER_ADDR', '127.0.0.1') - port = int(os.getenv('MASTER_PORT', '8361')) - num_nodes = int(os.getenv('WORLD_SIZE', 1)) - node_rank = int(os.getenv('RANK', 0)) - - sig = inspect.signature(dist.init_process_group) - params = { - 'backend': 'nccl', - 'init_method': f'tcp://{ip}:{port}', - 'world_size': num_nodes * num_local_ranks, - 'rank': node_rank * num_local_ranks + local_rank, - } - if 'device_id' in sig.parameters: - # noinspection PyTypeChecker - params['device_id'] = torch.device(f'cuda:{local_rank}') - dist.init_process_group(**params) - torch.set_default_dtype(torch.bfloat16) - torch.set_default_device('cuda') - torch.cuda.set_device(local_rank) - - return dist.get_rank(), dist.get_world_size(), dist.new_group(list(range(num_local_ranks * num_nodes))) - - -def calc_diff(x: torch.Tensor, y: torch.Tensor): - x, y = x.double() + 1, y.double() + 1 - denominator = (x * x + y * y).sum() - sim = 2 * (x * y).sum() / denominator - return (1 - sim).item() - - -def align_up(x, y): - return (x + y - 1) // y * y - - -def per_token_cast_to_fp8(x: torch.Tensor): - assert x.dim() == 2 - m, n = x.shape - aligned_n = align_up(n, 128) - x_padded = torch.nn.functional.pad(x, (0, aligned_n - n), mode='constant', value=0) - x_padded_view = x_padded.view(m, -1, 128) - x_amax = x_padded_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) - return (x_padded_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view( - m, aligned_n)[:, :n].contiguous(), (x_amax / 448.0).view(m, -1) - - -def per_token_cast_back(x_fp8: torch.Tensor, x_scales: torch.Tensor): - if x_fp8.numel() == 0: - return x_fp8.to(torch.bfloat16) - - assert x_fp8.dim() == 2 - m, n = x_fp8.shape - aligned_n = align_up(n, 128) - x_fp8_padded = torch.nn.functional.pad(x_fp8, (0, aligned_n - n), mode='constant', value=0) - if x_scales.dtype == torch.int: - x_scales = x_scales.view(dtype=torch.uint8).to(torch.int) << 23 - x_scales = x_scales.view(dtype=torch.float) - x_fp32_padded = x_fp8_padded.to(torch.float32).view(x_fp8.size(0), -1, 128) - x_scales = x_scales.view(x_fp8.size(0), -1, 1) - return (x_fp32_padded * x_scales).view(x_fp8_padded.shape).to(torch.bfloat16)[:, :n].contiguous() - - -def inplace_unique(x: torch.Tensor, num_slots: int): - assert x.dim() == 2 - mask = x < 0 - x_padded = x.masked_fill(mask, num_slots) - bin_count = torch.zeros((x.size(0), num_slots + 1), dtype=x.dtype, device=x.device) - bin_count.scatter_add_(1, x_padded, torch.ones_like(x_padded)) - bin_count = bin_count[:, :num_slots] - sorted_bin_count, sorted_bin_idx = torch.sort(bin_count, dim=-1, descending=True) - sorted_bin_idx.masked_fill_(sorted_bin_count == 0, -1) - sorted_bin_idx = torch.sort(sorted_bin_idx, descending=True, dim=-1).values - x[:, :].fill_(-1) - valid_len = min(num_slots, x.size(1)) - x[:, :valid_len] = sorted_bin_idx[:, :valid_len] - - -def create_grouped_scores(scores: torch.Tensor, group_idx: torch.Tensor, num_groups: int): - num_tokens, num_experts = scores.shape - scores = scores.view(num_tokens, num_groups, -1) - mask = torch.zeros((num_tokens, num_groups), dtype=torch.bool, device=scores.device) - mask = mask.scatter_(1, group_idx, True).unsqueeze(-1).expand_as(scores) - return (scores * mask).view(num_tokens, num_experts) - - -def bench(fn, num_warmups: int = 50, num_tests: int = 50, post_fn=None): - # Flush L2 cache with 256 MB data - torch.cuda.synchronize() - cache = torch.empty(int(256e6 // 4), dtype=torch.int, device='cuda') - - # Warmup - for _ in range(num_warmups): - fn() - - # Flush L2 - cache.zero_() - - # Testing - start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] - end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] - for i in range(num_tests): - # Record - start_events[i].record() - fn() - end_events[i].record() - if post_fn is not None: - post_fn() - torch.cuda.synchronize() - - times = np.array([s.elapsed_time(e) / 1e3 for s, e in zip(start_events, end_events)])[1:] - return np.average(times), np.min(times), np.max(times) - - -class empty_suppress: - - def __enter__(self): - return self - - def __exit__(self, *_): - pass - - -class suppress_stdout_stderr: - - def __enter__(self): - self.outnull_file = open(os.devnull, 'w') - self.errnull_file = open(os.devnull, 'w') - - self.old_stdout_fileno_undup = sys.stdout.fileno() - self.old_stderr_fileno_undup = sys.stderr.fileno() - - self.old_stdout_fileno = os.dup(sys.stdout.fileno()) - self.old_stderr_fileno = os.dup(sys.stderr.fileno()) - - self.old_stdout = sys.stdout - self.old_stderr = sys.stderr - - os.dup2(self.outnull_file.fileno(), self.old_stdout_fileno_undup) - os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup) - - sys.stdout = self.outnull_file - sys.stderr = self.errnull_file - return self - - def __exit__(self, *_): - sys.stdout = self.old_stdout - sys.stderr = self.old_stderr - - os.dup2(self.old_stdout_fileno, self.old_stdout_fileno_undup) - os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup) - - os.close(self.old_stdout_fileno) - os.close(self.old_stderr_fileno) - - self.outnull_file.close() - self.errnull_file.close() - - -def bench_kineto(fn, - kernel_names: Union[str, tuple], - num_tests: int = 30, - suppress_kineto_output: bool = False, - trace_path: Optional[str] = None, - barrier_comm_profiling: bool = False, - num_kernels_per_period: int = 1): - # Profile - suppress = suppress_stdout_stderr if suppress_kineto_output else empty_suppress - with suppress(): - schedule = torch.profiler.schedule(wait=1, warmup=0, active=1, repeat=1) - with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA], schedule=schedule) as prof: - for _ in range(2): - # NOTES: use a large kernel and a barrier to eliminate the unbalanced CPU launch overhead - if barrier_comm_profiling: - lhs = torch.randn((8192, 8192), dtype=torch.float, device='cuda') - rhs = torch.randn((8192, 8192), dtype=torch.float, device='cuda') - lhs @ rhs - dist.all_reduce(torch.ones(1, dtype=torch.float, device='cuda')) - for _ in range(num_tests): - fn() - torch.cuda.synchronize() - prof.step() - - # Parse the profiling table - assert isinstance(kernel_names, (str, tuple)) - is_tuple = isinstance(kernel_names, tuple) - prof_lines = prof.key_averages().table(sort_by='cuda_time_total', max_name_column_width=100).split('\n') - kernel_names = (kernel_names, ) if isinstance(kernel_names, str) else kernel_names - assert all([isinstance(name, str) for name in kernel_names]) - for name in kernel_names: - assert sum([name in line for line in prof_lines]) == 1, f'Errors of the kernel {name} in the profiling table' - - # Save chrome traces - if trace_path is not None: - prof.export_chrome_trace(trace_path) - - # Return average kernel durations - units = {'ms': 1e3, 'us': 1e6} - kernel_durations = [] - for name in kernel_names: - for line in prof_lines: - if name in line: - time_str = line.split()[-2] - for unit, scale in units.items(): - if unit in time_str: - kernel_durations.append(float(time_str.replace(unit, '')) / scale) - break - break - - # Expand the kernels by periods - if num_kernels_per_period > 1: - with tempfile.NamedTemporaryFile(suffix='.json') as tmp: - prof.export_chrome_trace(tmp.name) - profile_data = json.loads(Path(tmp.name).read_text()) - - for i, kernel_name in enumerate(kernel_names): - events = [event for event in profile_data['traceEvents'] if f'::{kernel_name}' in event['name']] - events = sorted(events, key=lambda event: event['ts']) - durations = [event['dur'] / 1e6 for event in events] - assert len(durations) % num_kernels_per_period == 0 - num_kernel_patterns = len(durations) // num_kernels_per_period - kernel_durations[i] = [sum(durations[j::num_kernels_per_period]) / num_kernel_patterns for j in range(num_kernels_per_period)] - - # Return execution durations - return kernel_durations if is_tuple else kernel_durations[0] - - -def hash_tensor(t: torch.Tensor): - return t.view(torch.int).sum().item() diff --git a/tests/utils/test_gate.py b/tests/utils/test_gate.py new file mode 100644 index 000000000..fa46c5c0a --- /dev/null +++ b/tests/utils/test_gate.py @@ -0,0 +1,57 @@ +import torch + +from deep_ep.utils.math import ceil_div +from deep_ep.utils.gate import get_unbalanced_scores + + +def test_unbalanced_scores(): + print('Testing gate score generation (Output with num_tokens = 4096, num_experts = 512):') + for num_tokens in [1, 4096]: + for num_experts_per_rank in [1, 4, 8, 16]: + for num_ranks in [2, 4, 8, 16, 64, 72]: + num_experts = num_experts_per_rank * num_ranks + for num_topk in [1, 2, 4, 6, 8, 9]: + if num_topk > num_experts: + continue + for ratio in [1.0, 2.0, 4.0]: + for precise in [1, 0]: + total_rank_count = torch.zeros(num_ranks, device='cuda') + + # This is the requirement from precise generation algorithm + lower_bound_per_token = max(1, ceil_div(num_topk, num_experts_per_rank)) + upper_bound_per_token = min(min(num_topk, num_ranks), int((num_ranks - 1) / ratio) + 1) + if lower_bound_per_token > upper_bound_per_token: + continue + + # Repeat for each rank + for rank_idx in range(num_ranks): + scores = get_unbalanced_scores(num_tokens, num_experts, num_ranks, num_topk, ratio, precise) + _topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False) + topk_idx = topk_idx // num_experts_per_rank + row_indices = torch.arange(num_tokens).unsqueeze(1).expand(num_tokens, num_topk).flatten() + topk_idx = topk_idx.flatten() + rank_count = torch.zeros((num_tokens, num_ranks), device='cuda') + rank_count[row_indices, topk_idx] = 1 + rank_count = rank_count.sum(dim=0) + total_rank_count += rank_count + + # Calculate the actual ratio and inequality + practical_ratio = total_rank_count[0].item() / max(total_rank_count[1:].min().item(), 1) + inequality = total_rank_count[1:].max().item() / max(total_rank_count[1:].min().item(), 1) + total_sent_tokens = int(total_rank_count.sum().item()) + if num_tokens > 1000: + if num_ranks in [8, 64] and num_experts_per_rank == 8: + print(f' > {precise=}, {num_ranks=:2d}, {num_topk=}, expected_ratio={ratio} | ' + f'ratio={practical_ratio:6.3f}, {inequality=:6.3f}, {total_sent_tokens=:7d}') + + # Only check the ratio and inequality in precise mode + if precise: + assert abs(practical_ratio - ratio) / ratio < 0.1 and inequality < 1.02, \ + f'Failed to generate unbalanced scores with following config: \n' \ + f'{precise=}, {num_tokens=}, {num_experts=:3d}, {num_ranks=:2d}, {num_topk=}, expected_ratio={ratio} | ' \ + f'ratio={practical_ratio:6.3f}, {inequality=:6.3f}, {total_sent_tokens=:7d}' + print() + + +if __name__ == '__main__': + test_unbalanced_scores() diff --git a/third-party/fmt b/third-party/fmt new file mode 160000 index 000000000..a4c7e1713 --- /dev/null +++ b/third-party/fmt @@ -0,0 +1 @@ +Subproject commit a4c7e17133ee9cb6a2f45545f6e974dd3c393efa diff --git a/third-party/nvshmem.patch b/third-party/nvshmem.patch deleted file mode 100644 index 6ec064adc..000000000 --- a/third-party/nvshmem.patch +++ /dev/null @@ -1,474 +0,0 @@ -From 9e6cc27cceb3130784e4ea7b61ea3171156365fd Mon Sep 17 00:00:00 2001 -From: Shangyan Zhou -Date: Fri, 20 Dec 2024 10:57:12 +0800 -Subject: [PATCH 1/4] Change QP creating order. - ---- - src/modules/transport/ibgda/ibgda.cpp | 13 ++++++++----- - 1 file changed, 8 insertions(+), 5 deletions(-) - -diff --git a/src/modules/transport/ibgda/ibgda.cpp b/src/modules/transport/ibgda/ibgda.cpp -index ef325cd..286132e 100644 ---- a/src/modules/transport/ibgda/ibgda.cpp -+++ b/src/modules/transport/ibgda/ibgda.cpp -@@ -2936,17 +2936,20 @@ int nvshmemt_ibgda_connect_endpoints(nvshmem_transport_t t, int *selected_dev_id - INFO(ibgda_state->log_level, "Creating %d RC QPs", device->rc.num_eps_per_pe); - for (int i = 0; i < num_rc_eps; ++i) { - // Do not create loopback to self -- if (i / device->rc.num_eps_per_pe == mype) { -+ int dst_pe = (i + 1 + mype) % n_pes; -+ int offset = i / n_pes; -+ int mapped_i = dst_pe * device->rc.num_eps_per_pe + offset; -+ if (dst_pe == mype) { - continue; - } -- status = ibgda_create_qp(&device->rc.eps[i], device, portid, i, -+ status = ibgda_create_qp(&device->rc.eps[mapped_i], device, portid, mapped_i, - NVSHMEMI_IBGDA_DEVICE_QP_TYPE_RC); - NVSHMEMI_NZ_ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out, -- "ibgda_create_dci failed on RC #%d.", i); -+ "ibgda_create_dci failed on RC #%d.", mapped_i); - -- status = ibgda_get_rc_handle(&local_rc_handles[i], device->rc.eps[i], device); -+ status = ibgda_get_rc_handle(&local_rc_handles[mapped_i], device->rc.eps[mapped_i], device); - NVSHMEMI_NZ_ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out, -- "ibgda_get_rc_handle failed on RC #%d.", i); -+ "ibgda_get_rc_handle failed on RC #%d.", mapped_i); - } - - if (num_rc_eps) { --- -2.25.1 - - -From b11d41e4f3727f2f6ccc00a8c852e59e2ee33c8a Mon Sep 17 00:00:00 2001 -From: Shangyan Zhou -Date: Fri, 10 Jan 2025 11:53:38 +0800 -Subject: [PATCH 2/4] Add recv queue and recv cq for rc qps. - -Let the ibgda rc qps use regular recv queue. - -Add recv queue to ibgda dev qp. - -IBGDA create recv cq - -Setup recv cq. - -fix recv queue. - -Remove some useless idx. - -Longer recv queue. ---- - .../nvshmem_common_ibgda.h | 19 +++++- - src/modules/transport/ibgda/ibgda.cpp | 65 ++++++++++++++++--- - 2 files changed, 71 insertions(+), 13 deletions(-) - -diff --git a/src/include/device_host_transport/nvshmem_common_ibgda.h b/src/include/device_host_transport/nvshmem_common_ibgda.h -index 8b8a263..1be3dec 100644 ---- a/src/include/device_host_transport/nvshmem_common_ibgda.h -+++ b/src/include/device_host_transport/nvshmem_common_ibgda.h -@@ -168,14 +168,17 @@ typedef struct { - uint64_t get_head; // last wqe idx + 1 with a "fetch" operation (g, get, amo_fetch) - uint64_t get_tail; // last wqe idx + 1 polled with cst; get_tail > get_head is possible - } tx_wq; -+ struct { -+ uint64_t resv_head; // last reserved wqe idx + 1 -+ } rx_wq; - struct { - uint64_t head; - uint64_t tail; - } ibuf; - char padding[NVSHMEMI_IBGDA_QP_MANAGEMENT_PADDING]; - } __attribute__((__aligned__(8))) nvshmemi_ibgda_device_qp_management_v1; --static_assert(sizeof(nvshmemi_ibgda_device_qp_management_v1) == 96, -- "ibgda_device_qp_management_v1 must be 96 bytes."); -+static_assert(sizeof(nvshmemi_ibgda_device_qp_management_v1) == 104, -+ "ibgda_device_qp_management_v1 must be 104 bytes."); - - typedef nvshmemi_ibgda_device_qp_management_v1 nvshmemi_ibgda_device_qp_management_t; - -@@ -199,9 +202,19 @@ typedef struct nvshmemi_ibgda_device_qp { - // May point to mvars.prod_idx or internal prod_idx - uint64_t *prod_idx; - } tx_wq; -+ struct { -+ uint16_t nwqes; -+ uint64_t tail; -+ void *wqe; -+ __be32 *dbrec; -+ void *bf; -+ nvshmemi_ibgda_device_cq_t *cq; -+ // May point to mvars.prod_idx or internal prod_idx -+ uint64_t *prod_idx; -+ } rx_wq; - nvshmemi_ibgda_device_qp_management_v1 mvars; // management variables - } nvshmemi_ibgda_device_qp_v1; --static_assert(sizeof(nvshmemi_ibgda_device_qp_v1) == 184, "ibgda_device_qp_v1 must be 184 bytes."); -+static_assert(sizeof(nvshmemi_ibgda_device_qp_v1) == 248, "ibgda_device_qp_v1 must be 248 bytes."); - - typedef nvshmemi_ibgda_device_qp_v1 nvshmemi_ibgda_device_qp_t; - -diff --git a/src/modules/transport/ibgda/ibgda.cpp b/src/modules/transport/ibgda/ibgda.cpp -index 286132e..e0b2d5c 100644 ---- a/src/modules/transport/ibgda/ibgda.cpp -+++ b/src/modules/transport/ibgda/ibgda.cpp -@@ -198,6 +198,7 @@ struct ibgda_ep { - off_t dbr_offset; - - struct ibgda_cq *send_cq; -+ struct ibgda_cq *recv_cq; - struct ibv_ah *ah; - - uint32_t user_index; -@@ -1538,7 +1539,8 @@ static int ibgda_create_cq_shared_objects(nvshmemt_ibgda_state_t *ibgda_state, - - struct ibv_context *context = device->context; - -- unsigned int num_cqs = device->dci.num_eps + device->rc.num_eps_per_pe * n_pes; -+ // Each RC qp has one send CQ and one recv CQ. -+ unsigned int num_cqs = device->dci.num_eps + device->rc.num_eps_per_pe * n_pes * 2; - - assert(ibgda_qp_depth > 0); - size_t num_cqe = IBGDA_ROUND_UP_POW2_OR_0(ibgda_qp_depth); -@@ -1701,7 +1703,8 @@ static int ibgda_create_qp_shared_objects(nvshmemt_ibgda_state_t *ibgda_state, - } - - // Allocate and map WQ buffer for all QPs. -- wq_buf_size_per_qp = num_wqebb * MLX5_SEND_WQE_BB; // num_wqebb is always a power of 2 -+ // Todo: reduce the size of wq buffer. -+ wq_buf_size_per_qp = num_wqebb * MLX5_SEND_WQE_BB * 2; // num_wqebb is always a power of 2 - wq_buf_size = wq_buf_size_per_qp * num_eps; - status = ibgda_nic_control_alloc(&wq_mobject, wq_buf_size, IBGDA_GPAGE_SIZE); - NVSHMEMI_NZ_ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out, "cannot allocate wq buf.\n"); -@@ -1882,8 +1885,11 @@ static int ibgda_create_qp(struct ibgda_ep **ep_ptr, struct ibgda_device *device - int cqe_version = 0; - - struct ibgda_cq *send_cq = NULL; -+ struct ibgda_cq *recv_cq = NULL; - - size_t num_wqebb = IBGDA_ROUND_UP_POW2_OR_0(ibgda_qp_depth); -+ size_t num_recv_wqe = ibgda_qp_depth; -+ size_t recv_wqe_size = 16; - - int status = 0; - -@@ -1911,6 +1917,11 @@ static int ibgda_create_qp(struct ibgda_ep **ep_ptr, struct ibgda_device *device - status = ibgda_create_cq(&send_cq, device); - NVSHMEMI_NZ_ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out, "ibgda_create_cq failed.\n"); - -+ if (qp_type == NVSHMEMI_IBGDA_DEVICE_QP_TYPE_RC) { -+ status = ibgda_create_cq(&recv_cq, device); -+ NVSHMEMI_NZ_ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out, "ibgda_create_cq failed.\n"); -+ } -+ - ep = (struct ibgda_ep *)calloc(1, sizeof(struct ibgda_ep)); - NVSHMEMI_NULL_ERROR_JMP(ep, status, NVSHMEMX_ERROR_OUT_OF_MEMORY, out, - "Unable to allocate mem for ep.\n"); -@@ -1939,12 +1950,9 @@ static int ibgda_create_qp(struct ibgda_ep **ep_ptr, struct ibgda_device *device - DEVX_SET(qpc, qp_context, pm_state, MLX5_QPC_PM_STATE_MIGRATED); - DEVX_SET(qpc, qp_context, pd, device->qp_shared_object.pdn); - DEVX_SET(qpc, qp_context, uar_page, uar_mobject->uar->page_id); // BF register -- DEVX_SET(qpc, qp_context, rq_type, IBGDA_SRQ_TYPE_VALUE); // Shared Receive Queue -- DEVX_SET(qpc, qp_context, srqn_rmpn_xrqn, device->qp_shared_object.srqn); - DEVX_SET(qpc, qp_context, cqn_snd, send_cq->cqn); -- DEVX_SET(qpc, qp_context, cqn_rcv, device->qp_shared_object.rcqn); -+ DEVX_SET(qpc, qp_context, cqn_rcv, qp_type == NVSHMEMI_IBGDA_DEVICE_QP_TYPE_RC ? recv_cq->cqn : device->qp_shared_object.rcqn); - DEVX_SET(qpc, qp_context, log_sq_size, IBGDA_ILOG2_OR0(num_wqebb)); -- DEVX_SET(qpc, qp_context, log_rq_size, 0); - DEVX_SET(qpc, qp_context, cs_req, 0); // Disable CS Request - DEVX_SET(qpc, qp_context, cs_res, 0); // Disable CS Response - DEVX_SET(qpc, qp_context, dbr_umem_valid, IBGDA_MLX5_UMEM_VALID_ENABLE); // Enable dbr_umem_id -@@ -1953,6 +1961,15 @@ static int ibgda_create_qp(struct ibgda_ep **ep_ptr, struct ibgda_device *device - DEVX_SET(qpc, qp_context, dbr_umem_id, dbr_umem->umem_id); // DBR buffer - DEVX_SET(qpc, qp_context, user_index, qp_idx); - DEVX_SET(qpc, qp_context, page_offset, 0); -+ if (qp_type == NVSHMEMI_IBGDA_DEVICE_QP_TYPE_RC){ -+ DEVX_SET(qpc, qp_context, rq_type, 0); // Regular recv queue -+ DEVX_SET(qpc, qp_context, log_rq_size, IBGDA_ILOG2(num_recv_wqe)); // 4 wqe -+ DEVX_SET(qpc, qp_context, log_rq_stride, IBGDA_ILOG2(recv_wqe_size) - 4); // max recv wqe size = 16B -+ } else { -+ DEVX_SET(qpc, qp_context, rq_type, IBGDA_SRQ_TYPE_VALUE); // Shared Receive Queue, DC must use this. -+ DEVX_SET(qpc, qp_context, srqn_rmpn_xrqn, device->qp_shared_object.srqn); -+ DEVX_SET(qpc, qp_context, log_rq_size, 0); -+ } - - ep->devx_qp = mlx5dv_devx_obj_create(context, cmd_in, sizeof(cmd_in), cmd_out, sizeof(cmd_out)); - NVSHMEMI_NULL_ERROR_JMP(ep->devx_qp, status, NVSHMEMX_ERROR_INTERNAL, out, -@@ -1962,9 +1979,9 @@ static int ibgda_create_qp(struct ibgda_ep **ep_ptr, struct ibgda_device *device - ep->portid = portid; - - ep->sq_cnt = num_wqebb; -- ep->sq_buf_offset = 0; -+ ep->sq_buf_offset = num_recv_wqe * recv_wqe_size; - -- ep->rq_cnt = 0; -+ ep->rq_cnt = num_recv_wqe; - ep->rq_buf_offset = 0; - - ep->wq_mobject = device->qp_shared_object.wq_mobject; -@@ -1978,6 +1995,7 @@ static int ibgda_create_qp(struct ibgda_ep **ep_ptr, struct ibgda_device *device - ep->uar_mobject = uar_mobject; - - ep->send_cq = send_cq; -+ ep->recv_cq = recv_cq; - - ep->qp_type = qp_type; - -@@ -1989,6 +2007,7 @@ out: - if (status) { - if (uar_mobject) ibgda_unmap_and_free_qp_uar(uar_mobject); - if (send_cq) ibgda_destroy_cq(send_cq); -+ if (recv_cq) ibgda_destroy_cq(recv_cq); - if (ep) free(ep); - } - -@@ -2287,6 +2306,10 @@ static int ibgda_destroy_ep(struct ibgda_ep *ep) { - ibgda_destroy_cq(ep->send_cq); - } - -+ if (ep->recv_cq) { -+ ibgda_destroy_cq(ep->recv_cq); -+ } -+ - if (ep->ah) { - ftable.destroy_ah(ep->ah); - } -@@ -2318,7 +2341,7 @@ static void ibgda_get_device_qp(nvshmemi_ibgda_device_qp_t *dev_qp, struct ibgda - dev_qp->qpn = ep->qpn; - - assert(ep->wq_mobject->has_gpu_mapping); -- dev_qp->tx_wq.wqe = (void *)((uintptr_t)ep->wq_mobject->aligned.gpu_ptr + ep->wq_offset); -+ dev_qp->tx_wq.wqe = (void *)((uintptr_t)ep->wq_mobject->aligned.gpu_ptr + ep->wq_offset + ep->sq_buf_offset); - - if (ibgda_nic_handler == IBGDA_NIC_HANDLER_GPU) { - assert(ep->dbr_mobject->has_gpu_mapping); -@@ -2330,6 +2353,12 @@ static void ibgda_get_device_qp(nvshmemi_ibgda_device_qp_t *dev_qp, struct ibgda - } - - dev_qp->tx_wq.nwqes = ep->sq_cnt; -+ if (ep->qp_type == NVSHMEMI_IBGDA_DEVICE_QP_TYPE_RC) { -+ dev_qp->rx_wq.nwqes = ep->rq_cnt; -+ dev_qp->rx_wq.wqe = (void *)((uintptr_t)ep->wq_mobject->aligned.gpu_ptr + ep->wq_offset + ep->rq_buf_offset); -+ dev_qp->rx_wq.dbrec = (__be32 *)((uintptr_t)ep->dbr_mobject->aligned.gpu_ptr + ep->dbr_offset); -+ dev_qp->rx_wq.bf = (void *)ep->uar_mobject->aligned.gpu_ptr; -+ } - - ibuf_dci_start = (uintptr_t)device->qp_shared_object.internal_buf.mem_object->aligned.gpu_ptr; - ibuf_rc_start = ibuf_dci_start + (size_per_dci * device->dci.num_eps); -@@ -2379,6 +2408,9 @@ static int ibgda_setup_gpu_state(nvshmem_transport_t t) { - nvshmemi_ibgda_device_cq_t *cq_d = NULL; - nvshmemi_ibgda_device_cq_t *cq_h = NULL; - -+ nvshmemi_ibgda_device_cq_t *recv_cq_d = NULL; -+ nvshmemi_ibgda_device_cq_t *recv_cq_h = NULL; -+ - uint8_t *qp_group_switches_d = NULL; - - const size_t mvars_offset = offsetof(nvshmemi_ibgda_device_qp_t, mvars); -@@ -2386,6 +2418,7 @@ static int ibgda_setup_gpu_state(nvshmem_transport_t t) { - const size_t cons_t_offset = offsetof(nvshmemi_ibgda_device_qp_management_t, tx_wq.cons_idx); - const size_t wqe_h_offset = offsetof(nvshmemi_ibgda_device_qp_management_t, tx_wq.resv_head); - const size_t wqe_t_offset = offsetof(nvshmemi_ibgda_device_qp_management_t, tx_wq.ready_head); -+ const size_t rx_resv_head_offset = offsetof(nvshmemi_ibgda_device_qp_management_t, rx_wq.resv_head); - - nvshmemi_ibgda_device_qp_map_type_t rc_map_type = NVSHMEMI_IBGDA_DEVICE_QP_MAP_TYPE_INVALID; - nvshmemi_ibgda_device_qp_map_type_t dc_map_type = NVSHMEMI_IBGDA_DEVICE_QP_MAP_TYPE_INVALID; -@@ -2421,7 +2454,7 @@ static int ibgda_setup_gpu_state(nvshmem_transport_t t) { - num_dct_handles += device->dct.num_eps * n_pes; - num_dci_handles += device->dci.num_eps; - num_rc_handles += device->rc.num_eps_per_pe * n_pes; -- num_cq_handles += device->dci.num_eps + (device->rc.num_eps_per_pe * (n_pes - 1)); -+ num_cq_handles += device->dci.num_eps + (device->rc.num_eps_per_pe * (n_pes - 1) * 2); - num_shared_dci_handles += device->dci.num_shared_eps; - } - assert(num_dci_handles - num_shared_dci_handles >= 0); -@@ -2456,6 +2489,10 @@ static int ibgda_setup_gpu_state(nvshmem_transport_t t) { - for (int i = 0; i < num_cq_handles; i++) { - nvshmemi_init_ibgda_device_cq(cq_h[i]); - } -+ -+ recv_cq_h = (nvshmemi_ibgda_device_cq_t *)calloc(1, sizeof(*recv_cq_h)); -+ NVSHMEMI_NULL_ERROR_JMP(recv_cq_h, status, NVSHMEMX_ERROR_OUT_OF_MEMORY, out, "recv_cq calloc err."); -+ nvshmemi_init_ibgda_device_cq(recv_cq_h[0]); - /* allocate host memory for dct, rc, cq, dci end */ - - /* allocate device memory for dct, rc, cq, dci start */ -@@ -2559,6 +2596,14 @@ static int ibgda_setup_gpu_state(nvshmem_transport_t t) { - } - - ++cq_idx; -+ -+ rc_h[arr_idx].rx_wq.cq = &cq_d[cq_idx]; -+ -+ ibgda_get_device_cq(&cq_h[cq_idx], device->rc.eps[i]->recv_cq); -+ cq_h[cq_idx].resv_head = (uint64_t *)(base_mvars_d_addr + rx_resv_head_offset); -+ cq_h[cq_idx].qpn = rc_h[arr_idx].qpn; -+ cq_h[cq_idx].qp_type = rc_h[arr_idx].qp_type; -+ ++cq_idx; - } - } - } --- -2.25.1 - - -From af479f9f23103d4a1579fae38676d6b3022df887 Mon Sep 17 00:00:00 2001 -From: Shangyan Zhou -Date: Sat, 8 Feb 2025 18:02:39 +0800 -Subject: [PATCH 3/4] Maintain recv queue's cons_idx. - ---- - src/include/device_host_transport/nvshmem_common_ibgda.h | 5 +++-- - src/modules/transport/ibgda/ibgda.cpp | 6 ++++-- - 2 files changed, 7 insertions(+), 4 deletions(-) - -diff --git a/src/include/device_host_transport/nvshmem_common_ibgda.h b/src/include/device_host_transport/nvshmem_common_ibgda.h -index 1be3dec..ea1e284 100644 ---- a/src/include/device_host_transport/nvshmem_common_ibgda.h -+++ b/src/include/device_host_transport/nvshmem_common_ibgda.h -@@ -170,6 +170,7 @@ typedef struct { - } tx_wq; - struct { - uint64_t resv_head; // last reserved wqe idx + 1 -+ uint64_t cons_idx; // polled wqe idx + 1 (consumer index + 1) - } rx_wq; - struct { - uint64_t head; -@@ -177,7 +178,7 @@ typedef struct { - } ibuf; - char padding[NVSHMEMI_IBGDA_QP_MANAGEMENT_PADDING]; - } __attribute__((__aligned__(8))) nvshmemi_ibgda_device_qp_management_v1; --static_assert(sizeof(nvshmemi_ibgda_device_qp_management_v1) == 104, -- "ibgda_device_qp_management_v1 must be 104 bytes."); -+static_assert(sizeof(nvshmemi_ibgda_device_qp_management_v1) == 112, -+ "ibgda_device_qp_management_v1 must be 112 bytes."); - - typedef nvshmemi_ibgda_device_qp_management_v1 nvshmemi_ibgda_device_qp_management_t; -@@ -214,7 +215,7 @@ typedef struct nvshmemi_ibgda_device_qp { - } rx_wq; - nvshmemi_ibgda_device_qp_management_v1 mvars; // management variables - } nvshmemi_ibgda_device_qp_v1; --static_assert(sizeof(nvshmemi_ibgda_device_qp_v1) == 248, "ibgda_device_qp_v1 must be 248 bytes."); -+static_assert(sizeof(nvshmemi_ibgda_device_qp_v1) == 256, "ibgda_device_qp_v1 must be 256 bytes."); - - typedef nvshmemi_ibgda_device_qp_v1 nvshmemi_ibgda_device_qp_t; - -diff --git a/src/modules/transport/ibgda/ibgda.cpp b/src/modules/transport/ibgda/ibgda.cpp -index e0b2d5c..bc339c5 100644 ---- a/src/modules/transport/ibgda/ibgda.cpp -+++ b/src/modules/transport/ibgda/ibgda.cpp -@@ -1067,7 +1067,7 @@ static inline void ibgda_nic_control_free(struct ibgda_mem_object *mobject) { - ibgda_host_mem_free(mobject); - } - --static int ibgda_create_cq(struct ibgda_cq **pgcq, struct ibgda_device *device) { -+static int ibgda_create_cq(struct ibgda_cq **pgcq, struct ibgda_device *device, int cc = 1) { - int status = 0; - - struct ibgda_cq *gcq = NULL; -@@ -1118,7 +1118,7 @@ static int ibgda_create_cq(struct ibgda_cq **pgcq, struct ibgda_device *device) - cq_context = DEVX_ADDR_OF(create_cq_in, cmd_in, cq_context); - DEVX_SET(cqc, cq_context, dbr_umem_valid, IBGDA_MLX5_UMEM_VALID_ENABLE); - DEVX_SET(cqc, cq_context, cqe_sz, MLX5_CQE_SIZE_64B); -- DEVX_SET(cqc, cq_context, cc, 0x1); // Use collapsed CQ -+ DEVX_SET(cqc, cq_context, cc, cc); // Use collapsed CQ - DEVX_SET(cqc, cq_context, oi, 0x1); // Allow overrun - DEVX_SET(cqc, cq_context, dbr_umem_id, dbr_umem->umem_id); - DEVX_SET(cqc, cq_context, log_cq_size, IBGDA_ILOG2_OR0(num_cqe)); -@@ -2419,6 +2419,7 @@ static int ibgda_setup_gpu_state(nvshmem_transport_t t) { - const size_t wqe_h_offset = offsetof(nvshmemi_ibgda_device_qp_management_t, tx_wq.resv_head); - const size_t wqe_t_offset = offsetof(nvshmemi_ibgda_device_qp_management_t, tx_wq.ready_head); - const size_t rx_resv_head_offset = offsetof(nvshmemi_ibgda_device_qp_management_t, rx_wq.resv_head); -+ const size_t rx_cons_offset = offsetof(nvshmemi_ibgda_device_qp_management_t, rx_wq.cons_idx); - - nvshmemi_ibgda_device_qp_map_type_t rc_map_type = NVSHMEMI_IBGDA_DEVICE_QP_MAP_TYPE_INVALID; - nvshmemi_ibgda_device_qp_map_type_t dc_map_type = NVSHMEMI_IBGDA_DEVICE_QP_MAP_TYPE_INVALID; -@@ -2601,6 +2602,7 @@ static int ibgda_setup_gpu_state(nvshmem_transport_t t) { - - ibgda_get_device_cq(&cq_h[cq_idx], device->rc.eps[i]->recv_cq); - cq_h[cq_idx].resv_head = (uint64_t *)(base_mvars_d_addr + rx_resv_head_offset); -+ cq_h[cq_idx].cons_idx = (uint64_t *)(base_mvars_d_addr + rx_cons_offset); - cq_h[cq_idx].qpn = rc_h[arr_idx].qpn; - cq_h[cq_idx].qp_type = rc_h[arr_idx].qp_type; - ++cq_idx; --- -2.25.1 - - -From e0ba3fa21b4b633b481c6684c3ad04f2670c8df4 Mon Sep 17 00:00:00 2001 -From: Shangyan Zhou -Date: Tue, 11 Feb 2025 11:00:57 +0800 -Subject: [PATCH 4/4] Init rx_wq counters. - ---- - src/include/device_host_transport/nvshmem_common_ibgda.h | 2 ++ - 1 file changed, 2 insertions(+) - -diff --git a/src/include/device_host_transport/nvshmem_common_ibgda.h b/src/include/device_host_transport/nvshmem_common_ibgda.h -index ea1e284..e6640d6 100644 ---- a/src/include/device_host_transport/nvshmem_common_ibgda.h -+++ b/src/include/device_host_transport/nvshmem_common_ibgda.h -@@ -46,6 +46,8 @@ - qp_man.tx_wq.cons_idx = NVSHMEMI_IBGDA_ULSCALAR_INVALID; \ - qp_man.tx_wq.get_head = NVSHMEMI_IBGDA_ULSCALAR_INVALID; \ - qp_man.tx_wq.get_tail = NVSHMEMI_IBGDA_ULSCALAR_INVALID; \ -+ qp_man.rx_wq.resv_head = NVSHMEMI_IBGDA_ULSCALAR_INVALID; \ -+ qp_man.rx_wq.cons_idx = NVSHMEMI_IBGDA_ULSCALAR_INVALID; \ - qp_man.ibuf.head = NVSHMEMI_IBGDA_ULSCALAR_INVALID; \ - qp_man.ibuf.tail = NVSHMEMI_IBGDA_ULSCALAR_INVALID; \ - } while (0); --- -2.25.1 - -diff --git a/src/modules/transport/common/transport_ib_common.cpp b/src/modules/transport/common/transport_ib_common.cpp -index c89f408..f99018a 100644 ---- a/src/modules/transport/common/transport_ib_common.cpp -+++ b/src/modules/transport/common/transport_ib_common.cpp -@@ -26,6 +26,9 @@ int nvshmemt_ib_common_nv_peer_mem_available() { - if (access("/sys/kernel/mm/memory_peers/nvidia-peermem/version", F_OK) == 0) { - return NVSHMEMX_SUCCESS; - } -+ if (access("/sys/module/nvidia_peermem/version", F_OK) == 0) { -+ return NVSHMEMX_SUCCESS; -+ } - - return NVSHMEMX_ERROR_INTERNAL; - } - - -From 099f608fcd9a1d34c866ad75d0af5d02d2020374 Mon Sep 17 00:00:00 2001 -From: Kaichao You -Date: Tue, 10 Jun 2025 00:35:03 -0700 -Subject: [PATCH] remove gdrcopy dependency - ---- - src/modules/transport/ibgda/ibgda.cpp | 6 ++++++ - 1 file changed, 6 insertions(+) - -diff --git a/src/modules/transport/ibgda/ibgda.cpp b/src/modules/transport/ibgda/ibgda.cpp -index ef325cd..16ee09c 100644 ---- a/src/modules/transport/ibgda/ibgda.cpp -+++ b/src/modules/transport/ibgda/ibgda.cpp -@@ -406,6 +406,7 @@ static size_t ibgda_get_host_page_size() { - return host_page_size; - } - -+#ifdef NVSHMEM_USE_GDRCOPY - int nvshmemt_ibgda_progress(nvshmem_transport_t t) { - nvshmemt_ibgda_state_t *ibgda_state = (nvshmemt_ibgda_state_t *)t->state; - int n_devs_selected = ibgda_state->n_devs_selected; -@@ -459,6 +460,11 @@ int nvshmemt_ibgda_progress(nvshmem_transport_t t) { - } - return 0; - } -+#else -+int nvshmemt_ibgda_progress(nvshmem_transport_t t) { -+ return NVSHMEMX_ERROR_NOT_SUPPORTED; -+} -+#endif - - int nvshmemt_ibgda_show_info(struct nvshmem_transport *transport, int style) { - NVSHMEMI_ERROR_PRINT("ibgda show info not implemented"); --- -2.34.1