Showing preview only (640K chars total). Download the full file or copy to clipboard to get everything.
Repository: deepseek-ai/DeepEP
Branch: main
Commit: 567632dd5981
Files: 37
Total size: 619.5 KB
Directory structure:
gitextract_l_jv2phj/
├── .clang-format
├── .github/
│ └── workflows/
│ └── format.yml
├── .gitignore
├── LICENSE
├── README.md
├── csrc/
│ ├── CMakeLists.txt
│ ├── config.hpp
│ ├── deep_ep.cpp
│ ├── deep_ep.hpp
│ ├── event.hpp
│ └── kernels/
│ ├── CMakeLists.txt
│ ├── api.cuh
│ ├── buffer.cuh
│ ├── configs.cuh
│ ├── exception.cuh
│ ├── ibgda_device.cuh
│ ├── internode.cu
│ ├── internode_ll.cu
│ ├── intranode.cu
│ ├── launch.cuh
│ ├── layout.cu
│ ├── runtime.cu
│ └── utils.cuh
├── deep_ep/
│ ├── __init__.py
│ ├── buffer.py
│ └── utils.py
├── format.sh
├── install.sh
├── pyproject.toml
├── requirements-lint.txt
├── setup.py
├── tests/
│ ├── test_internode.py
│ ├── test_intranode.py
│ ├── test_low_latency.py
│ └── utils.py
└── third-party/
├── README.md
└── nvshmem.patch
================================================
FILE CONTENTS
================================================
================================================
FILE: .clang-format
================================================
BasedOnStyle: Google
UseTab: Never
IndentWidth: 4
ColumnLimit: 140
AccessModifierOffset: -4
# Force pointers to the type for C++.
DerivePointerAlignment: false
PointerAlignment: Left
ReferenceAlignment: Left
AllowShortFunctionsOnASingleLine: Inline
AllowShortIfStatementsOnASingleLine: false
AllowShortLoopsOnASingleLine: false
AlignOperands: false
BreakBeforeBinaryOperators: None
Cpp11BracedListStyle: true
ContinuationIndentWidth: 4
BinPackArguments: false
BinPackParameters: false
================================================
FILE: .github/workflows/format.yml
================================================
name: Code Format Check
on:
push:
branches: [ main ]
pull_request:
branches: [ main ]
jobs:
format-check:
runs-on: ubuntu-latest
steps:
- name: Checkout source
uses: actions/checkout@v6
with:
fetch-depth: 0
- name: Setup environment
run: |
sudo apt-get update
sudo apt-get install -y bash
- name: Run format.sh
run: |
bash ./format.sh
# If format.sh return non-zero, GitHub Actions will mark it as failure.
================================================
FILE: .gitignore
================================================
compile_commands.json
.idea
.DS_Store
*.pyc
build/
.cache/
.vscode/
*/cmake-build-*/
================================================
FILE: LICENSE
================================================
MIT License
Copyright (c) 2025 DeepSeek
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
================================================
FILE: README.md
================================================
# 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.
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.
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.
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) |
**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!
### 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 |
**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!
## 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 also depends on NVSHMEM. Please refer to our [NVSHMEM Installation Guide](third-party/README.md) 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
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
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.
### 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.
## Roadmap
- [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 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.
#### 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.
## License
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).
## Experimental Branches
- [Zero-copy](https://github.com/deepseek-ai/DeepEP/pull/453)
- Removing the copy between PyTorch tensors and communication buffers, which reduces the SM usages significantly for normal kernels
- This PR is authored by **Tencent Network Platform Department**
- [Eager](https://github.com/deepseek-ai/DeepEP/pull/437)
- Using a low-latency protocol removes the extra RTT latency introduced by RDMA atomic OPs
- [Hybrid-EP](https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep)
- A new backend implementation using TMA instructions for minimal SM usage and larger NVLink domain support
- Fine-grained communication-computation overlap for single-batch scenarios
- PCIe kernel support for non-NVLink environments
- NVFP4 data type support
- [AntGroup-Opt](https://github.com/deepseek-ai/DeepEP/tree/antgroup-opt)
- This optimization series is authored by **AntGroup Network Platform Department**
- [Normal-SMFree](https://github.com/deepseek-ai/DeepEP/pull/347) Eliminating SM from RDMA path by decoupling comm-kernel execution from NIC token transfer, freeing SMs for compute
- [LL-SBO](https://github.com/deepseek-ai/DeepEP/pull/483) Overlapping Down GEMM computation with Combine Send communication via signaling mechanism to reduce end-to-end latency
- [LL-Layered](https://github.com/deepseek-ai/DeepEP/pull/500) Optimizing cross-node LL operator communication using rail-optimized forwarding and data merging to reduce latency
- [Mori-EP](https://github.com/deepseek-ai/DeepEP/tree/mori-ep)
- ROCm/AMD GPU support powered by [MORI](https://github.com/ROCm/mori) backend (low-latency mode)
## 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
If you use this codebase or otherwise find our work valuable, please cite:
```bibtex
@misc{deepep2025,
title={DeepEP: an efficient expert-parallel communication library},
author={Chenggang Zhao and Shangyan Zhou and Liyue Zhang and Chengqi Deng and Zhean Xu and Yuxuan Liu and Kuai Yu and Jiashi Li and Liang Zhao},
year={2025},
publisher = {GitHub},
howpublished = {\url{https://github.com/deepseek-ai/DeepEP}},
}
```
================================================
FILE: csrc/CMakeLists.txt
================================================
# 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)
================================================
FILE: csrc/config.hpp
================================================
#pragma once
#include "kernels/api.cuh"
#include "kernels/exception.cuh"
namespace deep_ep {
template <typename dtype_t>
dtype_t ceil_div(dtype_t a, dtype_t b) {
return (a + b - 1) / b;
}
template <typename dtype_t>
dtype_t align_up(dtype_t a, dtype_t b) {
return ceil_div<dtype_t>(a, b) * b;
}
template <typename dtype_t>
dtype_t align_down(dtype_t a, dtype_t b) {
return a / b * b;
}
struct Config {
int num_sms;
int num_max_nvl_chunked_send_tokens;
int num_max_nvl_chunked_recv_tokens;
int num_max_rdma_chunked_send_tokens;
int num_max_rdma_chunked_recv_tokens;
Config(int num_sms,
int num_max_nvl_chunked_send_tokens,
int num_max_nvl_chunked_recv_tokens,
int num_max_rdma_chunked_send_tokens,
int num_max_rdma_chunked_recv_tokens)
: num_sms(num_sms),
num_max_nvl_chunked_send_tokens(num_max_nvl_chunked_send_tokens),
num_max_nvl_chunked_recv_tokens(num_max_nvl_chunked_recv_tokens),
num_max_rdma_chunked_send_tokens(num_max_rdma_chunked_send_tokens),
num_max_rdma_chunked_recv_tokens(num_max_rdma_chunked_recv_tokens) {
EP_HOST_ASSERT(num_sms >= 0);
EP_HOST_ASSERT(num_max_nvl_chunked_send_tokens > 0 and num_max_nvl_chunked_recv_tokens > 0);
EP_HOST_ASSERT(num_max_nvl_chunked_send_tokens < num_max_nvl_chunked_recv_tokens);
EP_HOST_ASSERT(num_max_rdma_chunked_send_tokens > 0 and num_max_rdma_chunked_recv_tokens > 0);
// Ceil up RDMA buffer size
this->num_max_rdma_chunked_recv_tokens = align_up<int>(num_max_rdma_chunked_recv_tokens, num_max_rdma_chunked_send_tokens);
EP_HOST_ASSERT(num_max_rdma_chunked_send_tokens < num_max_rdma_chunked_recv_tokens);
// NOTES: this assertion is related to RDMA lazy head update, we must ensure senders always have space to push
EP_HOST_ASSERT(num_max_rdma_chunked_send_tokens <= num_max_rdma_chunked_recv_tokens / 2);
}
size_t get_nvl_buffer_size_hint(size_t hidden_bytes, int num_ranks) const {
// Below are some assumptions
// TODO: add assertions
constexpr int kNumMaxTopK = 128;
constexpr int kNumMaxScales = 128;
EP_HOST_ASSERT(num_ranks < NUM_MAX_NVL_PEERS or num_ranks % NUM_MAX_NVL_PEERS == 0);
EP_HOST_ASSERT(num_ranks <= NUM_MAX_NVL_PEERS or num_sms % 2 == 0);
const auto num_rdma_ranks = std::max(num_ranks / NUM_MAX_NVL_PEERS, 1);
const auto num_nvl_ranks = std::min(num_ranks, NUM_MAX_NVL_PEERS);
const int num_channels = num_sms / 2;
size_t num_bytes = 0;
num_bytes += num_channels * num_nvl_ranks * (2 * num_rdma_ranks + 3) * sizeof(int);
num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * hidden_bytes;
#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);
num_bytes = ((num_bytes + 127) / 128) * 128;
return num_bytes;
}
size_t get_rdma_buffer_size_hint(int64_t hidden_bytes, int num_ranks) const {
#ifndef DISABLE_NVSHMEM
// Legacy mode
if (num_ranks <= NUM_MAX_NVL_PEERS)
return 0;
// Below are some assumptions
// TODO: add assertions
constexpr int kNumMaxTopK = 128;
constexpr int kNumMaxScales = 128;
EP_HOST_ASSERT(num_ranks % NUM_MAX_NVL_PEERS == 0);
EP_HOST_ASSERT(num_sms % 2 == 0);
const int num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS;
const int num_channels = num_sms / 2;
size_t num_bytes = 0;
num_bytes += num_channels * num_rdma_ranks * (NUM_MAX_NVL_PEERS * 2 + 2) * 2 * sizeof(int);
num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * hidden_bytes * 2;
num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * 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;
num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * kNumMaxTopK * sizeof(float) * 2;
num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * kNumMaxScales * sizeof(float) * 2;
num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * sizeof(int4) * 2;
num_bytes = ((num_bytes + 127) / 128) * 128;
return num_bytes;
#else
EP_HOST_ASSERT(false and "NVSHMEM is disable during compilation");
#endif
}
};
struct LowLatencyBuffer {
int num_clean_int = 0;
void* dispatch_rdma_send_buffer = nullptr;
void* dispatch_rdma_recv_data_buffer = nullptr;
int* dispatch_rdma_recv_count_buffer = nullptr;
void* combine_rdma_send_buffer = nullptr;
void* combine_rdma_recv_data_buffer = nullptr;
int* combine_rdma_recv_flag_buffer = nullptr;
void* combine_rdma_send_buffer_data_start = nullptr;
size_t num_bytes_per_combine_msg = 0;
std::pair<int*, int> clean_meta() {
EP_HOST_ASSERT(dispatch_rdma_recv_count_buffer == combine_rdma_recv_flag_buffer);
return {dispatch_rdma_recv_count_buffer, num_clean_int};
}
};
struct LowLatencyLayout {
size_t total_bytes = 0;
LowLatencyBuffer buffers[2];
template <typename out_ptr_t = void*, typename count_ptr_t = uint8_t*, typename in_ptr_t = void*>
out_ptr_t advance(const in_ptr_t& ptr, size_t count) {
return reinterpret_cast<out_ptr_t>(reinterpret_cast<count_ptr_t>(ptr) + count);
}
LowLatencyLayout(void* rdma_buffer, int num_max_dispatch_tokens_per_rank, int hidden, int num_ranks, int num_experts) {
const int num_scales = hidden / 128;
// Dispatch and combine layout:
// - 2 symmetric odd/even send buffer
// - 2 symmetric odd/even receive buffers
// - 2 symmetric odd/even signaling buffers
// Message sizes
// NOTES: you should add a control `int4` for combine messages if you want to do data transformation
// NOTES: `num_scales * sizeof(nv_bfloat162)` means the per-128-channel min/max
EP_HOST_ASSERT(num_scales * sizeof(float) <= hidden);
size_t num_bytes_per_dispatch_msg = sizeof(int4) + std::max(hidden * sizeof(nv_bfloat16), hidden + num_scales * sizeof(float));
size_t num_bytes_per_combine_msg = num_scales * sizeof(nv_bfloat162) + hidden * sizeof(nv_bfloat16);
// Send buffer
size_t dispatch_send_buffer_bytes = num_max_dispatch_tokens_per_rank * num_bytes_per_dispatch_msg;
size_t combine_send_buffer_bytes = num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_combine_msg;
size_t send_buffer_bytes = std::max(dispatch_send_buffer_bytes, combine_send_buffer_bytes);
EP_HOST_ASSERT(send_buffer_bytes % sizeof(int4) == 0);
total_bytes += send_buffer_bytes * 2;
// Symmetric receive buffers
// TODO: optimize memory usages
size_t dispatch_recv_data_buffer_bytes = num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_dispatch_msg;
size_t combine_recv_buffer_bytes = num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_combine_msg;
size_t recv_buffer_bytes = std::max(dispatch_recv_data_buffer_bytes, combine_recv_buffer_bytes);
EP_HOST_ASSERT(recv_buffer_bytes % sizeof(int4) == 0);
total_bytes += recv_buffer_bytes * 2;
// Symmetric signaling buffers
size_t dispatch_recv_count_buffer_bytes = num_experts * sizeof(int);
size_t combine_recv_flag_buffer_bytes = dispatch_recv_count_buffer_bytes;
size_t signaling_buffer_bytes = std::max(dispatch_recv_count_buffer_bytes, combine_recv_flag_buffer_bytes);
size_t signaling_buffer_bytes_aligned = align_up<size_t>(signaling_buffer_bytes, 128);
total_bytes += signaling_buffer_bytes_aligned * 2;
// Assign pointers
// NOTES: we still leave some space for distinguishing dispatch/combine buffer,
// so you may see some parameters are duplicated
for (int i = 0; i < 2; ++i) {
buffers[i] = {static_cast<int>(signaling_buffer_bytes / sizeof(int)),
advance(rdma_buffer, signaling_buffer_bytes_aligned * 2 + send_buffer_bytes * i),
advance(rdma_buffer, signaling_buffer_bytes_aligned * 2 + send_buffer_bytes * 2 + recv_buffer_bytes * i),
advance<int*>(rdma_buffer, signaling_buffer_bytes_aligned * i),
advance(rdma_buffer, signaling_buffer_bytes_aligned * 2 + send_buffer_bytes * i),
advance(rdma_buffer, signaling_buffer_bytes_aligned * 2 + send_buffer_bytes * 2 + recv_buffer_bytes * i),
advance<int*>(rdma_buffer, signaling_buffer_bytes_aligned * i),
advance(rdma_buffer, signaling_buffer_bytes_aligned * 2 + send_buffer_bytes * i),
num_bytes_per_combine_msg};
}
}
};
size_t get_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;
}
} // namespace deep_ep
================================================
FILE: csrc/deep_ep.cpp
================================================
#include "deep_ep.hpp"
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/CUDADataType.h>
#include <cuda_runtime.h>
#include <pybind11/functional.h>
#include <torch/python.h>
#include <chrono>
#include <memory>
#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<int>::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<int>::max()));
EP_HOST_ASSERT(num_nvl_bytes / sizeof(int4) < std::numeric_limits<int>::max());
EP_HOST_ASSERT(num_rdma_bytes / sizeof(int4) < std::numeric_limits<int>::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<int64_t>(num_nvl_bytes, num_device_sms / 2) < std::numeric_limits<int>::max());
EP_HOST_ASSERT(ceil_div<int64_t>(num_rdma_bytes, num_device_sms / 2) < std::numeric_limits<int>::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<void**>(static_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes + barrier_signal_bytes);
// Set barrier signals
barrier_signal_ptrs[nvl_rank] = reinterpret_cast<int*>(static_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes);
barrier_signal_ptrs_gpu =
reinterpret_cast<int**>(static_cast<uint8_t*>(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<int*>(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<int*>(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<int*>(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<const char*>(&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<const char*>(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<int64_t>(elementSize(casted_dtype));
auto base_ptr = static_cast<uint8_t*>(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<int*>(moe_recv_counter)));
// Free chunked mode staffs
CUDA_CHECK(cudaFreeHost(const_cast<int*>(moe_recv_expert_counter)));
destroyed = true;
available = false;
}
void Buffer::sync(const std::vector<int>& device_ids,
const std::vector<std::optional<pybind11::bytearray>>& all_gathered_handles,
const std::optional<pybind11::bytearray>& 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<int*>(static_cast<uint8_t*>(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<uint8_t> root_unique_id(root_unique_id_opt->size());
auto root_unique_id_str = root_unique_id_opt->cast<std::string>();
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<int*>(internode::alloc(num_mask_buffer_bytes, NUM_BUFFER_ALIGNMENT_BYTES));
sync_buffer_ptr = reinterpret_cast<int*>(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, std::optional<torch::Tensor>, torch::Tensor, torch::Tensor, std::optional<EventHandle>>
Buffer::get_dispatch_layout(
const torch::Tensor& topk_idx, int num_experts, std::optional<EventHandle>& 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<int>(topk_idx.size(0)), num_topk = static_cast<int>(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<torch::Tensor>();
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<topk_idx_t>(),
num_tokens_per_rank.data_ptr<int>(),
num_tokens_per_rdma_rank.has_value() ? num_tokens_per_rdma_rank.value().data_ptr<int>() : nullptr,
num_tokens_per_expert.data_ptr<int>(),
is_token_in_rank.data_ptr<bool>(),
num_tokens,
num_topk,
num_ranks,
num_experts,
comm_stream);
// Wait streams
std::optional<EventHandle> 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<torch::Tensor,
std::optional<torch::Tensor>,
std::optional<torch::Tensor>,
std::optional<torch::Tensor>,
std::vector<int>,
torch::Tensor,
torch::Tensor,
torch::Tensor,
torch::Tensor,
torch::Tensor,
std::optional<EventHandle>>
Buffer::intranode_dispatch(const torch::Tensor& x,
const std::optional<torch::Tensor>& x_scales,
const std::optional<torch::Tensor>& topk_idx,
const std::optional<torch::Tensor>& topk_weights,
const std::optional<torch::Tensor>& num_tokens_per_rank,
const torch::Tensor& is_token_in_rank,
const std::optional<torch::Tensor>& num_tokens_per_expert,
int cached_num_recv_tokens,
const std::optional<torch::Tensor>& cached_rank_prefix_matrix,
const std::optional<torch::Tensor>& cached_channel_prefix_matrix,
int expert_alignment,
int num_worst_tokens,
const Config& config,
std::optional<EventHandle>& 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<int>(x.size(0)), hidden = static_cast<int>(x.size(1));
auto num_experts = cached_mode ? 0 : static_cast<int>(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<int>(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_idx_t>();
topk_weights_ptr = topk_weights->data_ptr<float>();
}
// 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<int>(x_scales->size(1));
x_scales_ptr = static_cast<float*>(x_scales->data_ptr());
scale_token_stride = static_cast<int>(x_scales->stride(0));
scale_hidden_stride = static_cast<int>(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<int> 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<int>(), 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<int>(),
moe_recv_counter_mapped,
num_ranks,
num_tokens_per_expert->data_ptr<int>(),
moe_recv_expert_counter_mapped,
num_experts,
num_tokens,
is_token_in_rank.data_ptr<bool>(),
channel_prefix_matrix.data_ptr<int>(),
rank_prefix_matrix.data_ptr<int>(),
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<int>(*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::seconds>(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<int>(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<torch::Tensor>(), recv_topk_weights = std::optional<torch::Tensor>(),
recv_x_scales = std::optional<torch::Tensor>();
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<topk_idx_t>();
recv_topk_weights_ptr = recv_topk_weights->data_ptr<float>();
}
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<float*>(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<int>(),
recv_topk_idx_ptr,
recv_topk_weights_ptr,
recv_channel_prefix_matrix.data_ptr<int>(),
send_head.data_ptr<int>(),
x.data_ptr(),
x_scales_ptr,
topk_idx_ptr,
topk_weights_ptr,
is_token_in_rank.data_ptr<bool>(),
channel_prefix_matrix.data_ptr<int>(),
num_tokens,
num_worst_tokens,
static_cast<int>(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<EventHandle> 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<torch::Tensor, std::optional<torch::Tensor>, std::optional<EventHandle>> Buffer::intranode_combine(
const torch::Tensor& x,
const std::optional<torch::Tensor>& topk_weights,
const std::optional<torch::Tensor>& bias_0,
const std::optional<torch::Tensor>& 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<EventHandle>& 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<int>(x.size(0)), hidden = static_cast<int>(x.size(1));
auto num_recv_tokens = static_cast<int>(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<torch::Tensor>();
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<int>(topk_weights->size(1));
topk_weights_ptr = topk_weights->data_ptr<float>();
recv_topk_weights = torch::empty({num_recv_tokens, num_topk}, topk_weights->options());
recv_topk_weights_ptr = recv_topk_weights->data_ptr<float>();
}
// 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<int>(),
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<std::optional<torch::Tensor>>({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<int>(),
rank_prefix_matrix.data_ptr<int>(),
channel_prefix_matrix.data_ptr<int>(),
send_head.data_ptr<int>(),
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<EventHandle> 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<torch::Tensor,
std::optional<torch::Tensor>,
std::optional<torch::Tensor>,
std::optional<torch::Tensor>,
std::vector<int>,
torch::Tensor,
torch::Tensor,
std::optional<torch::Tensor>,
torch::Tensor,
std::optional<torch::Tensor>,
torch::Tensor,
std::optional<torch::Tensor>,
std::optional<torch::Tensor>,
std::optional<torch::Tensor>,
std::optional<EventHandle>>
Buffer::internode_dispatch(const torch::Tensor& x,
const std::optional<torch::Tensor>& x_scales,
const std::optional<torch::Tensor>& topk_idx,
const std::optional<torch::Tensor>& topk_weights,
const std::optional<torch::Tensor>& num_tokens_per_rank,
const std::optional<torch::Tensor>& num_tokens_per_rdma_rank,
const torch::Tensor& is_token_in_rank,
const std::optional<torch::Tensor>& num_tokens_per_expert,
int cached_num_recv_tokens,
int cached_num_rdma_recv_tokens,
const std::optional<torch::Tensor>& cached_rdma_channel_prefix_matrix,
const std::optional<torch::Tensor>& cached_recv_rdma_rank_prefix_sum,
const std::optional<torch::Tensor>& cached_gbl_channel_prefix_matrix,
const std::optional<torch::Tensor>& cached_recv_gbl_rank_prefix_sum,
int expert_alignment,
int num_worst_tokens,
const Config& config,
std::optional<EventHandle>& 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<int>(x.size(0)), hidden = static_cast<int>(x.size(1)),
hidden_int4 = static_cast<int>(x.size(1) * x.element_size() / sizeof(int4));
auto num_experts = cached_mode ? 0 : static_cast<int>(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<int>(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_idx_t>();
topk_weights_ptr = topk_weights->data_ptr<float>();
}
// 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<int>(x_scales->size(1));
x_scales_ptr = static_cast<float*>(x_scales->data_ptr());
scale_token_stride = static_cast<int>(x_scales->stride(0));
scale_hidden_stride = static_cast<int>(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<int> 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<int>(),
moe_recv_counter_mapped,
num_ranks,
num_tokens_per_rdma_rank->data_ptr<int>(),
moe_recv_rdma_counter_mapped,
num_tokens_per_expert->data_ptr<int>(),
moe_recv_expert_counter_mapped,
num_experts,
is_token_in_rank.data_ptr<bool>(),
num_tokens,
num_worst_tokens,
num_channels,
hidden_int4,
num_scales,
num_topk,
expert_alignment,
rdma_channel_prefix_matrix.data_ptr<int>(),
recv_rdma_rank_prefix_sum.data_ptr<int>(),
gbl_channel_prefix_matrix.data_ptr<int>(),
recv_gbl_rank_prefix_sum.data_ptr<int>(),
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<int>(*moe_recv_counter);
num_rdma_recv_tokens = static_cast<int>(*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::seconds>(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<int>(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<torch::Tensor>(), recv_topk_weights = std::optional<torch::Tensor>(),
recv_x_scales = std::optional<torch::Tensor>();
auto recv_src_meta = std::optional<torch::Tensor>();
auto recv_rdma_channel_prefix_matrix = std::optional<torch::Tensor>();
auto recv_gbl_channel_prefix_matrix = std::optional<torch::Tensor>();
auto send_rdma_head = std::optional<torch::Tensor>();
auto send_nvl_head = std::optional<torch::Tensor>();
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<topk_idx_t>();
recv_topk_weights_ptr = recv_topk_weights->data_ptr<float>();
}
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<float*>(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<int>(),
cached_mode ? nullptr : send_nvl_head->data_ptr<int>(),
cached_mode ? nullptr : recv_rdma_channel_prefix_matrix->data_ptr<int>(),
cached_mode ? nullptr : recv_gbl_channel_prefix_matrix->data_ptr<int>(),
rdma_channel_prefix_matrix.data_ptr<int>(),
recv_rdma_rank_prefix_sum.data_ptr<int>(),
gbl_channel_prefix_matrix.data_ptr<int>(),
recv_gbl_rank_prefix_sum.data_ptr<int>(),
is_token_in_rank.data_ptr<bool>(),
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<EventHandle> 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<torch::Tensor, std::optional<torch::Tensor>, std::optional<EventHandle>> Buffer::internode_combine(
const torch::Tensor& x,
const std::optional<torch::Tensor>& topk_weights,
const std::optional<torch::Tensor>& bias_0,
const std::optional<torch::Tensor>& 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<EventHandle>& 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<int>(x.size(0)), hidden = static_cast<int>(x.size(1)),
hidden_int4 = static_cast<int>(x.size(1) * x.element_size() / sizeof(int4));
auto num_combined_tokens = static_cast<int>(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<torch::Tensor>();
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<int>(topk_weights->size(1));
topk_weights_ptr = topk_weights->data_ptr<float>();
combined_topk_weights = torch::empty({num_combined_tokens, num_topk}, topk_weights->options());
combined_topk_weights_ptr = combined_topk_weights->data_ptr<float>();
}
// 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<int>(),
rdma_channel_prefix_matrix.data_ptr<int>(),
rdma_rank_prefix_sum.data_ptr<int>(),
combined_nvl_head.data_ptr<int>(),
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<std::optional<torch::Tensor>>({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<bool>(),
x.data_ptr(),
topk_weights_ptr,
bias_ptrs[0],
bias_ptrs[1],
combined_rdma_head.data_ptr<int>(),
combined_nvl_head.data_ptr<int>(),
src_meta.data_ptr(),
rdma_channel_prefix_matrix.data_ptr<int>(),
rdma_rank_prefix_sum.data_ptr<int>(),
gbl_channel_prefix_matrix.data_ptr<int>(),
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<EventHandle> 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<int64_t>(ptr) - reinterpret_cast<int64_t>(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,
std::optional<torch::Tensor>,
torch::Tensor,
torch::Tensor,
torch::Tensor,
std::optional<EventHandle>,
std::optional<std::function<void()>>>
Buffer::low_latency_dispatch(const torch::Tensor& x,
const torch::Tensor& topk_idx,
const std::optional<torch::Tensor>& cumulative_local_expert_recv_stats,
const std::optional<torch::Tensor>& dispatch_wait_recv_cost_stats,
int num_max_dispatch_tokens_per_rank,
int num_experts,
bool use_fp8,
bool round_scale,
bool use_ue8m0,
bool async,
bool return_recv_hook) {
#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<topk_idx_t>::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<int>(x.size(0)), hidden = static_cast<int>(x.size(1));
auto num_topk = static_cast<int>(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<torch::Tensor>();
void* packed_recv_x_scales_ptr = nullptr;
EP_HOST_ASSERT((num_ranks * num_max_dispatch_tokens_per_rank) % 4 == 0 and "TMA requires the number of tokens to be multiple of 4");
if (use_fp8) {
// TODO: support unaligned cases
EP_HOST_ASSERT(hidden % 512 == 0);
if (not use_ue8m0) {
packed_recv_x_scales = torch::empty({num_local_experts, hidden / 128, num_ranks * num_max_dispatch_tokens_per_rank},
torch::dtype(torch::kFloat32).device(torch::kCUDA));
} else {
EP_HOST_ASSERT(round_scale);
packed_recv_x_scales = torch::empty({num_local_experts, hidden / 512, num_ranks * num_max_dispatch_tokens_per_rank},
torch::dtype(torch::kInt).device(torch::kCUDA));
}
packed_recv_x_scales = torch::transpose(packed_recv_x_scales.value(), 1, 2);
packed_recv_x_scales_ptr = packed_recv_x_scales->data_ptr();
}
// Kernel launch
auto next_clean_meta = next_buffer.clean_meta();
auto launcher = [=](int phases) {
internode_ll::dispatch(
packed_recv_x.data_ptr(),
packed_recv_x_scales_ptr,
packed_recv_src_info.data_ptr<int>(),
packed_recv_layout_range.data_ptr<int64_t>(),
packed_recv_count.data_ptr<int>(),
mask_buffer_ptr,
cumulative_local_expert_recv_stats.has_value() ? cumulative_local_expert_recv_stats->data_ptr<int>() : nullptr,
dispatch_wait_recv_cost_stats.has_value() ? dispatch_wait_recv_cost_stats->data_ptr<int64_t>() : nullptr,
buffer.dispatch_rdma_recv_data_buffer,
buffer.dispatch_rdma_recv_count_buffer,
buffer.dispatch_rdma_send_buffer,
x.data_ptr(),
topk_idx.data_ptr<topk_idx_t>(),
next_clean_meta.first,
next_clean_meta.second,
num_tokens,
hidden,
num_max_dispatch_tokens_per_rank,
num_topk,
num_experts,
rank,
num_ranks,
use_fp8,
round_scale,
use_ue8m0,
workspace,
num_device_sms,
launch_stream,
phases);
};
launcher(return_recv_hook ? LOW_LATENCY_SEND_PHASE : (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE));
// Wait streams
std::optional<EventHandle> event;
if (async) {
// NOTES: we must ensure the all tensors will not be deallocated before the stream-wait happens,
// so in Python API, we must wrap all tensors into the event handle.
event = EventHandle(launch_stream);
} else if (not return_recv_hook) {
stream_wait(compute_stream, launch_stream);
}
// Receiver callback
std::optional<std::function<void()>> recv_hook = std::nullopt;
if (return_recv_hook)
recv_hook = [=]() { launcher(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<torch::Tensor, std::optional<EventHandle>, std::optional<std::function<void()>>> 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<torch::Tensor>& combine_wait_recv_cost_stats,
int num_max_dispatch_tokens_per_rank,
int num_experts,
bool use_logfmt,
bool zero_copy,
bool async,
bool return_recv_hook,
const std::optional<torch::Tensor>& out) {
#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<topk_idx_t>::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<int>(x.size(2));
auto num_topk = static_cast<int>(topk_weights.size(1));
auto num_combined_tokens = static_cast<int>(topk_weights.size(0));
// Buffer control
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_idx_t>(),
topk_weights.data_ptr<float>(),
src_info.data_ptr<int>(),
layout_range.data_ptr<int64_t>(),
mask_buffer_ptr,
combine_wait_recv_cost_stats.has_value() ? combine_wait_recv_cost_stats->data_ptr<int64_t>() : nullptr,
next_clean_meta.first,
next_clean_meta.second,
num_combined_tokens,
hidden,
num_max_dispatch_tokens_per_rank,
num_topk,
num_experts,
rank,
num_ranks,
use_logfmt,
workspace,
num_device_sms,
launch_stream,
phases,
zero_copy);
};
launcher(return_recv_hook ? LOW_LATENCY_SEND_PHASE : (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE));
// Wait streams
std::optional<EventHandle> event;
if (async) {
// NOTES: we must ensure the all tensors will not be deallocated before the stream-wait happens,
// so in Python API, we must wrap all tensors into the event handle.
event = EventHandle(launch_stream);
} else if (not return_recv_hook) {
stream_wait(compute_stream, launch_stream);
}
// Receiver callback
std::optional<std::function<void()>> recv_hook = std::nullopt;
if (return_recv_hook)
recv_hook = [=]() { launcher(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<int>(buffer.num_bytes_per_combine_msg / elementSize(torch::kBFloat16));
EP_HOST_ASSERT(buffer.num_bytes_per_combine_msg % elementSize(torch::kBFloat16) == 0);
return torch::from_blob(buffer.combine_rdma_send_buffer_data_start,
{num_experts / num_ranks, num_ranks * num_max_dispatch_tokens_per_rank, hidden},
{num_ranks * num_max_dispatch_tokens_per_rank * num_msg_elems, num_msg_elems, 1},
torch::TensorOptions().dtype(dtype).device(torch::kCUDA));
#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<int*>(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_<deep_ep::Config>(m, "Config")
.def(pybind11::init<int, int, int, int, int>(),
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_<deep_ep::EventHandle>(m, "EventHandle")
.def(pybind11::init<>())
.def("current_stream_wait", &deep_ep::EventHandle::current_stream_wait);
pybind11::class_<deep_ep::Buffer>(m, "Buffer")
.def(pybind11::init<int, int, int64_t, int64_t, bool, bool, bool, bool>())
.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<py::object>((PyObject*)torch::getTHPDtype(c10::CppTypeToScalarType<deep_ep::topk_idx_t>::value));
}
================================================
FILE: csrc/deep_ep.hpp
================================================
#pragma once
// Forcibly disable NDEBUG
#ifdef NDEBUG
#undef NDEBUG
#endif
#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
#include <torch/types.h>
#include <tuple>
#include <vector>
#include "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<int>& device_ids,
const std::vector<std::optional<pybind11::bytearray>>& all_gathered_handles,
const std::optional<pybind11::bytearray>& root_unique_id_opt);
void destroy();
std::tuple<torch::Tensor, std::optional<torch::Tensor>, torch::Tensor, torch::Tensor, std::optional<EventHandle>> get_dispatch_layout(
const torch::Tensor& topk_idx,
int num_experts,
std::optional<EventHandle>& previous_event,
bool async,
bool allocate_on_comm_stream);
std::tuple<torch::Tensor,
std::optional<torch::Tensor>,
std::optional<torch::Tensor>,
std::optional<torch::Tensor>,
std::vector<int>,
torch::Tensor,
torch::Tensor,
torch::Tensor,
torch::Tensor,
torch::Tensor,
std::optional<EventHandle>>
intranode_dispatch(const torch::Tensor& x,
const std::optional<torch::Tensor>& x_scales,
const std::optional<torch::Tensor>& topk_idx,
const std::optional<torch::Tensor>& topk_weights,
const std::optional<torch::Tensor>& num_tokens_per_rank,
const torch::Tensor& is_token_in_rank,
const std::optional<torch::Tensor>& num_tokens_per_expert,
int cached_num_recv_tokens,
const std::optional<torch::Tensor>& cached_rank_prefix_matrix,
const std::optional<torch::Tensor>& cached_channel_prefix_matrix,
int expert_alignment,
int num_worst_tokens,
const Config& config,
std::optional<EventHandle>& previous_event,
bool async,
bool allocate_on_comm_stream);
std::tuple<torch::Tensor, std::optional<torch::Tensor>, std::optional<EventHandle>> intranode_combine(
const torch::Tensor& x,
const std::optional<torch::Tensor>& topk_weights,
const std::optional<torch::Tensor>& bias_0,
const std::optional<torch::Tensor>& 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<EventHandle>& previous_event,
bool async,
bool allocate_on_comm_stream);
std::tuple<torch::Tensor,
std::optional<torch::Tensor>,
std::optional<torch::Tensor>,
std::optional<torch::Tensor>,
std::vector<int>,
torch::Tensor,
torch::Tensor,
std::optional<torch::Tensor>,
torch::Tensor,
std::optional<torch::Tensor>,
torch::Tensor,
std::optional<torch::Tensor>,
std::optional<torch::Tensor>,
std::optional<torch::Tensor>,
std::optional<EventHandle>>
internode_dispatch(const torch::Tensor& x,
const std::optional<torch::Tensor>& x_scales,
const std::optional<torch::Tensor>& topk_idx,
const std::optional<torch::Tensor>& topk_weights,
const std::optional<torch::Tensor>& num_tokens_per_rank,
const std::optional<torch::Tensor>& num_tokens_per_rdma_rank,
const torch::Tensor& is_token_in_rank,
const std::optional<torch::Tensor>& num_tokens_per_expert,
int cached_num_recv_tokens,
int cached_num_rdma_recv_tokens,
const std::optional<torch::Tensor>& cached_rdma_channel_prefix_matrix,
const std::optional<torch::Tensor>& cached_recv_rdma_rank_prefix_sum,
const std::optional<torch::Tensor>& cached_gbl_channel_prefix_matrix,
const std::optional<torch::Tensor>& cached_recv_gbl_rank_prefix_sum,
int expert_alignment,
int num_worst_tokens,
const Config& config,
std::optional<EventHandle>& previous_event,
bool async,
bool allocate_on_comm_stream);
std::tuple<torch::Tensor, std::optional<torch::Tensor>, std::optional<EventHandle>> internode_combine(
const torch::Tensor& x,
const std::optional<torch::Tensor>& topk_weights,
const std::optional<torch::Tensor>& bias_0,
const std::optional<torch::Tensor>& 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<EventHandle>& 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,
std::optional<torch::Tensor>,
torch::Tensor,
torch::Tensor,
torch::Tensor,
std::optional<EventHandle>,
std::optional<std::function<void()>>>
low_latency_dispatch(const torch::Tensor& x,
const torch::Tensor& topk_idx,
const std::optional<torch::Tensor>& cumulative_local_expert_recv_stats,
const std::optional<torch::Tensor>& dispatch_wait_recv_cost_stats,
int num_max_dispatch_tokens_per_rank,
int num_experts,
bool use_fp8,
bool round_scale,
bool use_ue8m0,
bool async,
bool return_recv_hook);
std::tuple<torch::Tensor, std::optional<EventHandle>, std::optional<std::function<void()>>> 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<torch::Tensor>& combine_wait_recv_cost_stats,
int num_max_dispatch_tokens_per_rank,
int num_experts,
bool use_logfmt,
bool zero_copy,
bool async,
bool return_recv_hook,
const std::optional<torch::Tensor>& out = std::nullopt);
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
================================================
FILE: csrc/event.hpp
================================================
#include <ATen/cuda/CUDAContext.h>
#include <memory>
#include "kernels/exception.cuh"
namespace deep_ep {
struct EventHandle {
std::shared_ptr<torch::Event> event;
EventHandle() {
event = std::make_shared<torch::Event>(torch::kCUDA);
event->record(at::cuda::getCurrentCUDAStream());
}
explicit EventHandle(const at::cuda::CUDAStream& stream) {
event = std::make_shared<torch::Event>(torch::kCUDA);
event->record(stream);
}
EventHandle(const EventHandle& other) = default;
void current_stream_wait() const { at::cuda::getCurrentCUDAStream().unwrap().wait(*event); }
};
torch::Event create_event(const at::cuda::CUDAStream& s) {
auto event = torch::Event(torch::kCUDA);
event.record(s);
return event;
}
void stream_wait(const at::cuda::CUDAStream& s_0, const at::cuda::CUDAStream& s_1) {
EP_HOST_ASSERT(s_0.id() != s_1.id());
s_0.unwrap().wait(create_event(s_1));
}
void stream_wait(const at::cuda::CUDAStream& s, const EventHandle& event) {
s.unwrap().wait(*event.event);
}
} // namespace deep_ep
================================================
FILE: csrc/kernels/CMakeLists.txt
================================================
function(add_deep_ep_library target_name source_file)
add_library(${target_name} STATIC ${source_file})
set_target_properties(${target_name} PROPERTIES
POSITION_INDEPENDENT_CODE ON
CXX_STANDARD_REQUIRED ON
CUDA_STANDARD_REQUIRED ON
CXX_STANDARD 17
CUDA_STANDARD 17
CUDA_SEPARABLE_COMPILATION ON
)
target_link_libraries(${target_name} PUBLIC nvshmem 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)
================================================
FILE: csrc/kernels/api.cuh
================================================
#pragma once
#include <vector>
#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<uint8_t> get_unique_id();
int init(const std::vector<uint8_t>& 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
// Layout kernels
namespace layout {
void get_dispatch_layout(const topk_idx_t* topk_idx,
int* num_tokens_per_rank,
int* num_tokens_per_rdma_rank,
int* num_tokens_per_expert,
bool* is_token_in_rank,
int num_tokens,
int num_topk,
int num_ranks,
int num_experts,
cudaStream_t stream);
} // namespace layout
// Intranode kernels
namespace intranode {
void notify_dispatch(const int* num_tokens_per_rank,
int* moe_recv_counter_mapped,
int num_ranks,
const int* num_tokens_per_expert,
int* moe_recv_expert_counter_mapped,
int num_experts,
int num_tokens,
const bool* is_token_in_rank,
int* channel_prefix_matrix,
int* rank_prefix_matrix_copy,
int num_memset_int,
int expert_alignment,
void** buffer_ptrs,
int** barrier_signal_ptrs,
int rank,
cudaStream_t stream,
int num_sms);
void cached_notify_dispatch(const int* rank_prefix_matrix,
int num_memset_int,
void** buffer_ptrs,
int** barrier_signal_ptrs,
int rank,
int num_ranks,
cudaStream_t stream);
void dispatch(void* recv_x,
float* recv_x_scales,
int* recv_src_idx,
topk_idx_t* recv_topk_idx,
float* recv_topk_weights,
int* recv_channel_offset,
int* send_head,
const void* x,
const float* x_scales,
const topk_idx_t* topk_idx,
const float* topk_weights,
const bool* is_token_in_rank,
const int* channel_prefix_matrix,
int num_tokens,
int num_worst_tokens,
int hidden_int4,
int num_topk,
int num_experts,
int num_scales,
int scale_token_stride,
int scale_hidden_stride,
void** buffer_ptrs,
int rank,
int num_ranks,
cudaStream_t stream,
int num_sms,
int num_max_send_tokens,
int num_recv_buffer_tokens);
void cached_notify_combine(void** buffer_ptrs,
int* send_head,
int num_channels,
int num_recv_tokens,
int num_memset_int,
int** barrier_signal_ptrs,
int rank,
int num_ranks,
cudaStream_t stream);
void combine(cudaDataType_t type,
void* recv_x,
float* recv_topk_weights,
const void* x,
const float* topk_weights,
const void* bias_0,
const void* bias_1,
const int* src_idx,
const int* rank_prefix_matrix,
const int* channel_prefix_matrix,
int* send_head,
int num_tokens,
int num_recv_tokens,
int hidden,
int num_topk,
void** buffer_ptrs,
int rank,
int num_ranks,
cudaStream_t stream,
int num_sms,
int num_max_send_tokens,
int num_recv_buffer_tokens);
} // namespace intranode
// Internode kernels
namespace internode {
int get_source_meta_bytes();
void notify_dispatch(const int* num_tokens_per_rank,
int* moe_recv_counter_mapped,
int num_ranks,
const int* num_tokens_per_rdma_rank,
int* moe_recv_rdma_counter_mapped,
const int* num_tokens_per_expert,
int* moe_recv_expert_counter_mapped,
int num_experts,
const bool* is_token_in_rank,
int num_tokens,
int num_worst_tokens,
int num_channels,
int hidden_int4,
int num_scales,
int num_topk,
int expert_alignment,
int* rdma_channel_prefix_matrix,
int* recv_rdma_rank_prefix_sum,
int* gbl_channel_prefix_matrix,
int* recv_gbl_rank_prefix_sum,
void* rdma_buffer_ptr,
int num_max_rdma_chunked_recv_tokens,
void** buffer_ptrs,
int num_max_nvl_chunked_recv_tokens,
int** barrier_signal_ptrs,
int rank,
cudaStream_t stream,
int64_t num_rdma_bytes,
int64_t num_nvl_bytes,
bool low_latency_mode);
void dispatch(void* recv_x,
float* recv_x_scales,
topk_idx_t* recv_topk_idx,
float* recv_topk_weights,
void* recv_src_meta,
const void* x,
const float* x_scales,
const topk_idx_t* topk_idx,
const float* topk_weights,
int* send_rdma_head,
int* send_nvl_head,
int* recv_rdma_channel_prefix_matrix,
int* recv_gbl_channel_prefix_matrix,
const int* rdma_channel_prefix_matrix,
const int* recv_rdma_rank_prefix_sum,
const int* gbl_channel_prefix_matrix,
const int* recv_gbl_rank_prefix_sum,
const bool* is_token_in_rank,
int num_tokens,
int num_worst_tokens,
int hidden_int4,
int num_scales,
int num_topk,
int num_experts,
int scale_token_stride,
int scale_hidden_stride,
void* rdma_buffer_ptr,
int num_max_rdma_chunked_send_tokens,
int num_max_rdma_chunked_recv_tokens,
void** buffer_ptrs,
int num_max_nvl_chunked_send_tokens,
int num_max_nvl_chunked_recv_tokens,
int rank,
int num_ranks,
bool is_cached_dispatch,
cudaStream_t stream,
int num_channels,
bool low_latency_mode);
void cached_notify(int hidden_int4,
int num_scales,
int num_topk_idx,
int num_topk_weights,
int num_ranks,
int num_channels,
int num_combined_tokens,
int* combined_rdma_head,
const int* rdma_channel_prefix_matrix,
const int* rdma_rank_prefix_sum,
int* combined_nvl_head,
void* rdma_buffer_ptr,
int num_max_rdma_chunked_recv_tokens,
void** buffer_ptrs,
int num_max_nvl_chunked_recv_tokens,
int** barrier_signal_ptrs,
int rank,
cudaStream_t stream,
int64_t num_rdma_bytes,
int64_t num_nvl_bytes,
bool is_cached_dispatch,
bool low_latency_mode);
void combine(cudaDataType_t type,
void* combined_x,
float* combined_topk_weights,
const bool* is_combined_token_in_rank,
const void* x,
const float* topk_weights,
const void* bias_0,
const void* bias_1,
const int* combined_rdma_head,
const int* combined_nvl_head,
const void* src_meta,
const int* rdma_channel_prefix_matrix,
const int* rdma_rank_prefix_sum,
const int* gbl_channel_prefix_matrix,
int num_tokens,
int num_combined_tokens,
int hidden,
int num_topk,
void* rdma_buffer_ptr,
int num_max_rdma_chunked_send_tokens,
int num_max_rdma_chunked_recv_tokens,
void** buffer_ptrs,
int num_max_nvl_chunked_send_tokens,
int num_max_nvl_chunked_recv_tokens,
int rank,
int num_ranks,
cudaStream_t stream,
int num_channels,
bool low_latency_mode);
} // namespace internode
// Internode low-latency kernels
namespace internode_ll {
void clean_low_latency_buffer(int* clean_0,
int num_clean_int_0,
int* clean_1,
int num_clean_int_1,
int rank,
int num_ranks,
int* mask_buffer,
int* sync_buffer,
cudaStream_t stream);
void dispatch(void* packed_recv_x,
void* packed_recv_x_scales,
int* packed_recv_src_info,
int64_t* packed_recv_layout_range,
int* packed_recv_count,
int* mask_buffer,
int* cumulative_local_expert_recv_stats,
int64_t* dispatch_wait_recv_cost_stats,
void* rdma_recv_x,
int* rdma_recv_count,
void* rdma_x,
const void* x,
const topk_idx_t* topk_idx,
int* next_clean,
int num_next_clean_int,
int num_tokens,
int hidden,
int num_max_dispatch_tokens_per_rank,
int num_topk,
int num_experts,
int rank,
int num_ranks,
bool use_fp8,
bool round_scale,
bool use_ue8m0,
void* workspace,
int num_device_sms,
cudaStream_t stream,
int phases);
void combine(void* combined_x,
void* rdma_recv_x,
int* rdma_recv_flag,
void* rdma_send_x,
const void* x,
const topk_idx_t* topk_idx,
const float* topk_weights,
const int* src_info,
const int64_t* layout_range,
int* mask_buffer,
int64_t* combine_wait_recv_cost_stats,
int* next_clean,
int num_next_clean_int,
int num_combined_tokens,
int hidden,
int num_max_dispatch_tokens_per_rank,
int num_topk,
int num_experts,
int rank,
int num_ranks,
bool use_logfmt,
void* workspace,
int num_device_sms,
cudaStream_t stream,
int phases,
bool zero_copy);
void query_mask_buffer(int* mask_buffer_ptr, int num_ranks, int* output_mask_tensor, cudaStream_t stream);
void update_mask_buffer(int* mask_buffer_ptr, int rank_to_mask, bool mask, cudaStream_t stream);
void clean_mask_buffer(int* mask_buffer_ptr, int num_ranks, cudaStream_t stream);
} // namespace internode_ll
} // namespace deep_ep
================================================
FILE: csrc/kernels/buffer.cuh
================================================
#pragma once
#include "configs.cuh"
#include "exception.cuh"
namespace deep_ep {
template <typename dtype_t>
struct Buffer {
private:
uint8_t* ptr;
public:
int64_t total_bytes;
__device__ __forceinline__ Buffer() : ptr(nullptr), total_bytes(0) {}
__device__ __forceinline__ Buffer(void*& gbl_ptr, int num_elems, int offset = 0) {
total_bytes = num_elems * sizeof(dtype_t);
ptr = static_cast<uint8_t*>(gbl_ptr) + offset * sizeof(dtype_t);
gbl_ptr = static_cast<uint8_t*>(gbl_ptr) + total_bytes;
}
__device__ __forceinline__ Buffer advance_also(void*& gbl_ptr) {
gbl_ptr = static_cast<uint8_t*>(gbl_ptr) + total_bytes;
return *this;
}
__device__ __forceinline__ dtype_t* buffer() { return reinterpret_cast<dtype_t*>(ptr); }
__device__ __forceinline__ dtype_t& operator[](int idx) { return buffer()[idx]; }
};
template <typename dtype_t, int kNumRanks = 1>
struct AsymBuffer {
private:
uint8_t* ptrs[kNumRanks];
int64_t num_bytes;
public:
int64_t total_bytes;
__device__ __forceinline__ AsymBuffer(void*& gbl_ptr, int num_elems, int num_ranks, int sm_id = 0, int num_sms = 1, int offset = 0) {
EP_STATIC_ASSERT(kNumRanks == 1, "");
num_bytes = num_elems * sizeof(dtype_t);
int64_t per_channel_bytes = num_bytes * num_ranks;
total_bytes = per_channel_bytes * num_sms;
ptrs[0] = static_cast<uint8_t*>(gbl_ptr) + per_channel_bytes * sm_id + num_bytes * offset;
gbl_ptr = static_cast<uint8_t*>(gbl_ptr) + total_bytes;
}
__device__ __forceinline__ AsymBuffer(void** gbl_ptrs, int num_elems, int num_ranks, int sm_id = 0, int num_sms = 1, int offset = 0) {
EP_STATIC_ASSERT(kNumRanks > 1, "");
num_bytes = num_elems * sizeof(dtype_t);
int64_t per_channel_bytes = num_bytes * num_ranks;
total_bytes = per_channel_bytes * num_sms;
for (int i = 0; i < kNumRanks; ++i) {
ptrs[i] = static_cast<uint8_t*>(gbl_ptrs[i]) + per_channel_bytes * sm_id + num_bytes * offset;
gbl_ptrs[i] = static_cast<uint8_t*>(gbl_ptrs[i]) + total_bytes;
}
}
__device__ __forceinline__ void advance(int shift) {
#pragma unroll
for (int i = 0; i < kNumRanks; ++i)
ptrs[i] = ptrs[i] + shift * sizeof(dtype_t);
}
__device__ __forceinline__ AsymBuffer advance_also(void*& gbl_ptr) {
gbl_ptr = static_cast<uint8_t*>(gbl_ptr) + total_bytes;
return *this;
}
template <int kNumAlsoRanks>
__device__ __forceinline__ AsymBuffer advance_also(void** gbl_ptrs) {
for (int i = 0; i < kNumAlsoRanks; ++i)
gbl_ptrs[i] = static_cast<uint8_t*>(gbl_ptrs[i]) + total_bytes;
return *this;
}
__device__ __forceinline__ dtype_t* buffer(int idx = 0) {
EP_STATIC_ASSERT(kNumRanks == 1, "`buffer` is only available for single rank case");
return reinterpret_cast<dtype_t*>(ptrs[0] + num_bytes * idx);
}
__device__ __forceinline__ dtype_t* buffer_by(int rank_idx, int idx = 0) {
EP_STATIC_ASSERT(kNumRanks > 1, "`buffer` is only available for single rank case");
return reinterpret_cast<dtype_t*>(ptrs[rank_idx] + num_bytes * idx);
}
};
template <typename dtype_t, bool kDecoupled = true>
struct SymBuffer {
private:
// NOTES: for non-decoupled case, `recv_ptr` is not used
uint8_t* send_ptr;
uint8_t* recv_ptr;
int64_t num_bytes;
public:
int64_t total_bytes;
__device__ __forceinline__ SymBuffer(void*& gbl_ptr, int num_elems, int num_ranks, int sm_id = 0, int num_sms = 1) {
num_bytes = num_elems * sizeof(dtype_t);
int64_t per_channel_bytes = num_bytes * num_ranks;
total_bytes = per_channel_bytes * num_sms * (static_cast<int>(kDecoupled) + 1);
send_ptr = static_cast<uint8_t*>(gbl_ptr) + per_channel_bytes * sm_id;
recv_ptr = static_cast<uint8_t*>(gbl_ptr) + per_channel_bytes * (sm_id + num_sms);
gbl_ptr = static_cast<uint8_t*>(gbl_ptr) + total_bytes;
}
__device__ __forceinline__ dtype_t* send_buffer(int idx = 0) {
EP_STATIC_ASSERT(kDecoupled, "`send_buffer` is only available for non-decoupled case");
return reinterpret_cast<dtype_t*>(send_ptr + num_bytes * idx);
}
__device__ __forceinline__ dtype_t* recv_buffer(int idx = 0) {
EP_STATIC_ASSERT(kDecoupled, "`recv_buffer` is only available for non-decoupled case");
return reinterpret_cast<dtype_t*>(recv_ptr + num_bytes * idx);
}
__device__ __forceinline__ dtype_t* buffer(int idx = 0) {
EP_STATIC_ASSERT(not kDecoupled, "`buffer` is only available for decoupled case");
return reinterpret_cast<dtype_t*>(send_ptr + num_bytes * idx);
}
};
} // namespace deep_ep
================================================
FILE: csrc/kernels/configs.cuh
================================================
#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 <cuda_bf16.h>
#include <cuda_runtime.h>
#include <cstdint>
#ifndef DISABLE_SM90_FEATURES
#include <cuda_fp8.h>
#else
// Ampere does not support FP8 features
#define __NV_E4M3 0
#define __NV_E5M2 1
typedef int __nv_fp8_interpretation_t;
typedef int __nv_fp8x4_e4m3;
typedef uint8_t __nv_fp8_storage_t;
#endif
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 <device_host_transport/nvshmem_common_ibgda.h>
#include <infiniband/mlx5dv.h>
#include <nvshmem.h>
#include <nvshmemx.h>
#include <non_abi/device/threadgroup/nvshmemi_common_device_defines.cuh>
#endif
================================================
FILE: csrc/kernels/exception.cuh
================================================
#pragma once
#include <exception>
#include <string>
#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
================================================
FILE: csrc/kernels/ibgda_device.cuh
================================================
// Portions derived from NVSHMEM (https://developer.nvidia.com/nvshmem)
// Copyright (c) NVIDIA Corporation.
// Licensed under the NVSHMEM Software License Agreement (version: September 3, 2019).
// See full license at: https://docs.nvidia.com/nvshmem/api/sla.html
//
// Modified from original source:
// - nvshmem/src/include/non_abi/device/pt-to-pt/ibgda_device.cuh
#pragma once
#include <type_traits>
#include "configs.cuh"
#include "exception.cuh"
#include "utils.cuh"
namespace deep_ep {
EP_STATIC_ASSERT(NVSHMEMI_IBGDA_MIN_QP_DEPTH >= 64, "Invalid QP minimum depth");
__device__ static __forceinline__ uint64_t HtoBE64(uint64_t x) {
uint64_t ret;
asm("{\n\t"
".reg .b32 ign;\n\t"
".reg .b32 lo;\n\t"
".reg .b32 hi;\n\t"
".reg .b32 new_lo;\n\t"
".reg .b32 new_hi;\n\t"
"mov.b64 {lo,hi}, %1;\n\t"
"prmt.b32 new_hi, lo, ign, 0x0123;\n\t"
"prmt.b32 new_lo, hi, ign, 0x0123;\n\t"
"mov.b64 %0, {new_lo,new_hi};\n\t"
"}"
: "=l"(ret)
: "l"(x));
return ret;
}
__device__ static __forceinline__ uint32_t HtoBE32(uint32_t x) {
uint32_t ret;
asm("{\n\t"
".reg .b32 ign;\n\t"
"prmt.b32 %0, %1, ign, 0x0123;\n\t"
"}"
: "=r"(ret)
: "r"(x));
return ret;
}
__device__ static __forceinline__ uint16_t HtoBE16(uint16_t x) {
// TODO: simplify PTX using 16-bit instructions
auto a = static_cast<uint32_t>(x);
uint32_t d;
asm volatile(
"{\n\t"
".reg .b32 mask;\n\t"
".reg .b32 ign;\n\t"
"mov.b32 mask, 0x4401;\n\t"
"mov.b32 ign, 0x0;\n\t"
"prmt.b32 %0, %1, ign, mask;\n\t"
"}"
: "=r"(d)
: "r"(a));
return static_cast<uint16_t>(d);
}
typedef struct mlx5_wqe_ctrl_seg __attribute__((__aligned__(8))) ibgda_ctrl_seg_t;
typedef struct {
uint32_t add_data;
uint32_t field_boundary;
uint64_t reserved;
} __attribute__((__packed__)) ibgda_atomic_32_masked_fa_seg_t;
__device__ static __forceinline__ nvshmemi_ibgda_device_state_t* ibgda_get_state() {
return &nvshmemi_ibgda_device_state_d;
}
// Template helper to get RC - uses compile-time type checking with if constexpr (C++17)
template <typename StateType>
__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<StateType, nvshmemi_ibgda_device_state_v1>) {
// 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);
}
__device__ static __forceinline__ void ibgda_lock_acquire(int* lock) {
while (atomicCAS(lock, 0, 1) == 1)
;
// Prevent reordering before the lock is acquired
memory_fence_cta();
}
__device__ static __forceinline__ void ibgda_lock_release(int* lock) {
memory_fence_cta();
// Prevent reordering before lock is released
st_na_relaxed(lock, 0);
}
__device__ static __forceinline__ void ibgda_update_dbr(nvshmemi_ibgda_device_qp_t* qp, uint32_t dbrec_head) {
// `DBREC` contains the index of the next empty `WQEBB`
__be32 dbrec_val;
__be32* dbrec_ptr = qp->tx_wq.dbrec;
// This is equivalent to `WRITE_ONCE(dbrec_ptr, HtoBE32(dbrec_head & 0xffff))`
asm("{\n\t"
".reg .b32 dbrec_head_16b;\n\t"
".reg .b32 ign;\n\t"
"and.b32 dbrec_head_16b, %1, 0xffff;\n\t"
"prmt.b32 %0, dbrec_head_16b, ign, 0x123;\n\t"
"}"
: "=r"(dbrec_val)
: "r"(dbrec_head));
st_na_release(dbrec_ptr, dbrec_val);
}
__device__ static __forceinline__ void ibgda_ring_db(nvshmemi_ibgda_device_qp_t* qp, uint16_t prod_idx) {
auto bf_ptr = reinterpret_cast<uint64_t*>(qp->tx_wq.bf);
ibgda_ctrl_seg_t ctrl_seg = {.opmod_idx_opcode = HtoBE32(prod_idx << 8), .qpn_ds = HtoBE32(qp->qpn << 8)};
EP_STATIC_ASSERT(sizeof(decltype(&ctrl_seg)) == sizeof(uint64_t), "");
st_na_release(bf_ptr, *(reinterpret_cast<uint64_t*>(&ctrl_seg)));
}
__device__ static __forceinline__ void ibgda_post_send(nvshmemi_ibgda_device_qp_t* qp, uint64_t new_prod_idx) {
nvshmemi_ibgda_device_qp_management_t* mvars = &qp->mvars;
uint64_t old_prod_idx;
// Update `prod_idx` before ringing the doorbell, so that we know which index is needed in quiet/fence
ibgda_lock_acquire(&mvars->post_send_lock);
old_prod_idx = atomicMax(reinterpret_cast<unsigned long long int*>(&mvars->tx_wq.prod_idx), new_prod_idx);
if (new_prod_idx > old_prod_idx) {
ibgda_update_dbr(qp, new_prod_idx);
ibgda_ring_db(qp, new_prod_idx);
}
ibgda_lock_release(&mvars->post_send_lock);
}
template <bool kAlwaysDoPostSend>
__device__ static __forceinline__ void ibgda_submit_requests(nvshmemi_ibgda_device_qp_t* qp,
uint64_t base_wqe_idx,
uint32_t num_wqes,
int message_idx = 0) {
auto state = ibgda_get_state();
nvshmemi_ibgda_device_qp_management_t* mvars = &qp->mvars;
uint64_t new_wqe_idx = base_wqe_idx + num_wqes;
// WQE writes must be finished first
__threadfence();
unsigned long long int* ready_idx =
(unsigned long long int*)(state->use_async_postsend ? qp->tx_wq.prod_idx : &mvars->tx_wq.ready_head);
// Wait for prior WQE slots to be filled first
while (atomicCAS(ready_idx, base_wqe_idx, new_wqe_idx) != base_wqe_idx)
;
// Always post, not in batch
if (!state->use_async_postsend) {
constexpr int kNumRequestInBatch = 4;
if (kAlwaysDoPostSend or (message_idx + 1) % kNumRequestInBatch == 0)
ibgda_post_send(qp, new_wqe_idx);
}
}
__device__ static __forceinline__ void ibgda_write_rdma_write_inl_wqe(
nvshmemi_ibgda_device_qp_t* qp, const uint32_t* val, uint64_t raddr, __be32 rkey, uint16_t wqe_idx, void** out_wqes, uint32_t imm) {
ibgda_ctrl_seg_t ctrl_seg;
struct mlx5_wqe_raddr_seg raddr_seg;
struct mlx5_wqe_inl_data_seg inl_seg;
auto* ctrl_seg_ptr = reinterpret_cast<ibgda_ctrl_seg_t*>(out_wqes[0]);
auto* raddr_seg_ptr = reinterpret_cast<mlx5_wqe_raddr_seg*>(reinterpret_cast<uintptr_t>(ctrl_seg_ptr) + sizeof(*ctrl_seg_ptr));
auto* inl_seg_ptr = reinterpret_cast<mlx5_wqe_inl_data_seg*>(reinterpret_cast<uintptr_t>(raddr_seg_ptr) + sizeof(*raddr_seg_ptr));
auto* wqe_data_ptr = reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(inl_seg_ptr) + sizeof(*inl_seg_ptr));
raddr_seg.raddr = HtoBE64(raddr);
raddr_seg.rkey = rkey;
raddr_seg.reserved = 0;
inl_seg.byte_count = HtoBE32(4 | MLX5_INLINE_SEG);
// `imm == std::numeric_limits<uint32_t>::max()` means no imm writes
ctrl_seg = {0};
ctrl_seg.qpn_ds = HtoBE32((qp->qpn << 8) | 3);
ctrl_seg.fm_ce_se = MLX5_WQE_CTRL_CQ_UPDATE;
ctrl_seg.opmod_idx_opcode =
HtoBE32((wqe_idx << 8) | (imm != std::numeric_limits<uint32_t>::max() ? MLX5_OPCODE_RDMA_WRITE_IMM : MLX5_OPCODE_RDMA_WRITE));
if (imm != std::numeric_limits<uint32_t>::max())
ctrl_seg.imm = HtoBE32(imm);
EP_STATIC_ASSERT(sizeof(*ctrl_seg_ptr) == 16, "sizeof(*ctrl_seg_ptr) == 16");
EP_STATIC_ASSERT(sizeof(*raddr_seg_ptr) == 16, "sizeof(*raddr_seg_ptr) == 16");
EP_STATIC_ASSERT(sizeof(*inl_seg_ptr) == 4, "sizeof(*inl_seg_ptr) == 4");
st_na_relaxed(reinterpret_cast<int4*>(ctrl_seg_ptr), *reinterpret_cast<const int4*>(&ctrl_seg));
st_na_relaxed(reinterpret_cast<int4*>(raddr_seg_ptr), *reinterpret_cast<const int4*>(&raddr_seg));
st_na_relaxed(reinterpret_cast<uint32_t*>(inl_seg_ptr), *reinterpret_cast<const uint32_t*>(&inl_seg));
st_na_relaxed(reinterpret_cast<uint32_t*>(wqe_data_ptr), *reinterpret_cast<const uint32_t*>(val));
}
__device__ static __forceinline__ uint64_t
ibgda_get_lkey_and_rkey(uint64_t laddr, __be32* lkey, uint64_t raddr, int dst_pe, uint64_t* out_raddr, __be32* out_rkey, uint32_t dev_idx) {
auto state = ibgda_get_state();
auto heap_start = reinterpret_cast<uint64_t>(nvshmemi_device_state_d.heap_base);
auto log2_cumem_granularity = state->log2_cumem_granularity;
// Local key
uint64_t idx = ((laddr - heap_start) >> log2_cumem_granularity) * state->num_devices_initialized + dev_idx;
auto device_key = state->constmem.lkeys[idx];
auto lchunk_size = device_key.next_addr - laddr;
*lkey = device_key.key;
// Remote key
uint64_t roffset = raddr - heap_start;
idx = ((roffset >> log2_cumem_granularity) * nvshmemi_device_state_d.npes) * state->num_devices_initialized +
dst_pe * state->num_devices_initialized + dev_idx;
if (idx < NVSHMEMI_IBGDA_MAX_CONST_RKEYS) {
device_key = state->constmem.rkeys[idx];
} else {
device_key = state->globalmem.rkeys[idx - NVSHMEMI_IBGDA_MAX_CONST_RKEYS];
}
*out_raddr = reinterpret_cast<uint64_t>(nvshmemi_device_state_d.peer_heap_base_remote[dst_pe]) + roffset;
*out_rkey = device_key.key;
// Return the minimum of local and remote chunk sizes
auto rchunk_size = device_key.next_addr - roffset;
return min(lchunk_size, rchunk_size);
}
__device__ static __forceinline__ void ibgda_get_rkey(uint64_t addr, int dst_pe, uint64_t* out_raddr, __be32* out_rkey, uint32_t dev_idx) {
auto state = ibgda_get_state();
auto heap_start = reinterpret_cast<uint64_t>(nvshmemi_device_state_d.heap_base);
uint64_t roffset = addr - heap_start;
uint64_t idx = ((roffset >> state->log2_cumem_granularity) * nvshmemi_device_state_d.npes * state->num_devices_initialized) +
dst_pe * state->num_devices_initialized + dev_idx;
nvshmemi_ibgda_device_key_t device_key;
if (idx < NVSHMEMI_IBGDA_MAX_CONST_RKEYS)
device_key = state->constmem.rkeys[idx];
else
device_key = state->globalmem.rkeys[idx - NVSHMEMI_IBGDA_MAX_CONST_RKEYS];
*out_raddr = reinterpret_cast<uint64_t>(nvshmemi_device_state_d.peer_heap_base_remote[dst_pe]) + roffset;
*out_rkey = device_key.key;
}
__device__ static __forceinline__ uint64_t ibgda_reserve_wqe_slots(nvshmemi_ibgda_device_qp_t* qp, uint32_t num_wqes) {
auto mvars = &qp->mvars;
return atomicAdd(reinterpret_cast<unsigned long long*>(&mvars->tx_wq.resv_head), static_cast<unsigned long long>(num_wqes));
}
__device__ static __forceinline__ void* ibgda_get_wqe_ptr(nvshmemi_ibgda_device_qp_t* qp, uint16_t wqe_idx) {
uint16_t cnt = qp->tx_wq.nwqes;
uint16_t idx = wqe_idx & (cnt - 1);
return reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(qp->tx_wq.wqe) + (idx << MLX5_SEND_WQE_SHIFT));
}
__device__ static __forceinline__ void nvshmemi_ibgda_rma_p(
int* rptr, const int value, int dst_pe, int qp_id, uint32_t imm = std::numeric_limits<uint32_t>::max()) {
// Get rkey
// NOTES: the `p` operation will not cross multiple remote chunks
__be32 rkey;
uint64_t raddr;
auto qp = ibgda_get_rc(dst_pe, qp_id);
ibgda_get_rkey(reinterpret_cast<uint64_t>(rptr), dst_pe, &raddr, &rkey, qp->dev_idx);
// Write WQEs
uint64_t base_wqe_idx = ibgda_reserve_wqe_slots(qp, 1);
void* wqe_ptrs;
wqe_ptrs = ibgda_get_wqe_ptr(qp, base_wqe_idx);
ibgda_write_rdma_write_inl_wqe(qp, reinterpret_cast<const uint32_t*>(&value), raddr, rkey, base_wqe_idx, &wqe_ptrs, imm);
// Submit requests
ibgda_submit_requests<true>(qp, base_wqe_idx, 1);
}
__device__ static __forceinline__ void ibgda_write_rdma_write_wqe(nvshmemi_ibgda_device_qp_t* qp,
uint64_t laddr,
__be32 lkey,
uint64_t raddr,
__be32 rkey,
uint32_t bytes,
uint16_t wqe_idx,
void** out_wqes) {
ibgda_ctrl_seg_t ctrl_seg;
struct mlx5_wqe_raddr_seg raddr_seg;
struct mlx5_wqe_data_seg data_seg;
auto* ctrl_seg_ptr = reinterpret_cast<ibgda_ctrl_seg_t*>(out_wqes[0]);
void* av_seg_ptr = reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(ctrl_seg_ptr) + sizeof(*ctrl_seg_ptr));
struct mlx5_wqe_raddr_seg* raddr_seg_ptr;
struct mlx5_wqe_data_seg* data_seg_ptr;
raddr_seg_ptr = reinterpret_cast<mlx5_wqe_raddr_seg*>(reinterpret_cast<uintptr_t>(av_seg_ptr));
data_seg_ptr = reinterpret_cast<mlx5_wqe_data_seg*>(reinterpret_cast<uintptr_t>(raddr_seg_ptr) + sizeof(*raddr_seg_ptr));
raddr_seg.raddr = HtoBE64(raddr);
raddr_seg.rkey = rkey;
raddr_seg.reserved = 0;
data_seg.byte_count = HtoBE32(bytes);
data_seg.lkey = lkey;
data_seg.addr = HtoBE64(laddr);
ctrl_seg = {0};
ctrl_seg.qpn_ds = HtoBE32((qp->qpn << 8) | 3);
ctrl_seg.fm_ce_se = MLX5_WQE_CTRL_CQ_UPDATE;
ctrl_seg.opmod_idx_opcode = HtoBE32((wqe_idx << 8) | MLX5_OPCODE_RDMA_WRITE);
EP_STATIC_ASSERT(sizeof(*ctrl_seg_ptr) == 16, "sizeof(*ctrl_seg_ptr) == 16");
EP_STATIC_ASSERT(sizeof(*raddr_seg_ptr) == 16, "sizeof(*raddr_seg_ptr) == 16");
EP_STATIC_ASSERT(sizeof(*data_seg_ptr) == 16, "sizeof(*data_seg_ptr) == 16");
st_na_relaxed(reinterpret_cast<int4*>(ctrl_seg_ptr), *reinterpret_cast<const int4*>(&ctrl_seg));
st_na_relaxed(reinterpret_cast<int4*>(raddr_seg_ptr), *reinterpret_cast<const int4*>(&raddr_seg));
st_na_relaxed(reinterpret_cast<int4*>(data_seg_ptr), *reinterpret_cast<const int4*>(&data_seg));
}
__device__ static __forceinline__ void ibgda_write_empty_recv_wqe(void* out_wqe) {
auto* data_seg_ptr = reinterpret_cast<struct mlx5_wqe_data_seg*>(out_wqe);
struct mlx5_wqe_data_seg data_seg;
// Make the first segment in the WQE invalid, then the entire list will be invalid
data_seg.byte_count = 0;
data_seg.lkey = HtoBE64(MLX5_INVALID_LKEY);
data_seg.addr = 0;
EP_STATIC_ASSERT(sizeof(mlx5_wqe_data_seg) == sizeof(int4), "Invalid data type length");
st_na_relaxed(reinterpret_cast<int4*>(data_seg_ptr), *reinterpret_cast<const int4*>(&data_seg));
}
template <bool kAlwaysDoPostSend = false>
__device__ static __forceinline__ void nvshmemi_ibgda_put_nbi_warp(
uint64_t req_rptr, uint64_t req_lptr, size_t bytes, int dst_pe, int qp_id, int lane_id, int message_idx) {
// Get lkey and rkey, store them into lanes
uint32_t num_wqes = 0;
__be32 my_lkey = 0;
uint64_t my_laddr = 0;
__be32 my_rkey = 0;
uint64_t my_raddr = 0;
uint64_t my_chunk_size = 0;
auto qp = ibgda_get_rc(dst_pe, qp_id);
// Decide how many messages (theoretically 3 for maximum)
auto remaining_bytes = bytes;
while (remaining_bytes > 0) {
if (lane_id == num_wqes) {
my_chunk_size = min(remaining_bytes,
ibgda_get_lkey_and_rkey(my_laddr = req_lptr, &my_lkey, req_rptr, dst_pe, &my_raddr, &my_rkey, qp->dev_idx));
}
// Move one more message
auto chunk_size = __shfl_sync(0xffffffff, my_chunk_size, static_cast<int>(num_wqes));
remaining_bytes -= chunk_size;
req_lptr += chunk_size;
req_rptr += chunk_size;
++num_wqes;
}
EP_DEVICE_ASSERT(num_wqes <= 32);
// Process WQE
uint64_t base_wqe_idx = 0;
if (lane_id == 0)
base_wqe_idx = ibgda_reserve_wqe_slots(qp, num_wqes);
base_wqe_idx = __shfl_sync(0xffffffff, base_wqe_idx, 0);
if (lane_id < num_wqes) {
auto wqe_idx = base_wqe_idx + lane_id;
auto wqe_ptr = ibgda_get_wqe_ptr(qp, wqe_idx);
ibgda_write_rdma_write_wqe(qp, my_laddr, my_lkey, my_raddr, my_rkey, my_chunk_size, wqe_idx, &wqe_ptr);
}
__syncwarp();
// Submit
if (lane_id == 0)
ibgda_submit_requests<kAlwaysDoPostSend>(qp, base_wqe_idx, num_wqes, message_idx);
__syncwarp();
}
__device__ static __forceinline__ void ibgda_write_amo_add_wqe(nvshmemi_ibgda_device_qp_t* qp,
const int& value,
uint64_t laddr,
__be32 lkey,
uint64_t raddr,
__be32 rkey,
uint16_t wqe_idx,
void** out_wqes) {
ibgda_ctrl_seg_t ctrl_seg = {0};
struct mlx5_wqe_raddr_seg raddr_seg;
struct mlx5_wqe_atomic_seg atomic_seg_1;
struct mlx5_wqe_data_seg data_seg;
auto ctrl_seg_ptr = reinterpret_cast<ibgda_ctrl_seg_t*>(out_wqes[0]);
auto raddr_seg_ptr = reinterpret_cast<mlx5_wqe_raddr_seg*>(reinterpret_cast<uintptr_t>(ctrl_seg_ptr) + sizeof(*ctrl_seg_ptr));
auto atomic_seg_ptr = reinterpret_cast<mlx5_wqe_atomic_seg*>(reinterpret_cast<uintptr_t>(raddr_seg_ptr) + sizeof(*raddr_seg_ptr));
auto data_seg_ptr = reinterpret_cast<mlx5_wqe_data_seg*>(reinterpret_cast<uintptr_t>(atomic_seg_ptr) + sizeof(*atomic_seg_ptr));
raddr_seg.raddr = HtoBE64(raddr);
raddr_seg.rkey = rkey;
raddr_seg.reserved = 0;
// NOTES: `0x08000000` means `IBGDA_4_BYTE_EXT_AMO_OPMOD`
ctrl_seg.opmod_idx_opcode = HtoBE32(MLX5_OPCODE_ATOMIC_MASKED_FA | (wqe_idx << 8) | 0x08000000);
auto atomic_32_masked_fa_seg = reinterpret_cast<ibgda_atomic_32_masked_fa_seg_t*>(&atomic_seg_1);
atomic_32_masked_fa_seg->add_data = HtoBE32(value);
atomic_32_masked_fa_seg->field_boundary = 0;
ctrl_seg.qpn_ds = HtoBE32((qp->qpn << 8) | 4);
ctrl_seg.fm_ce_se = MLX5_WQE_CTRL_CQ_UPDATE;
data_seg.byte_count = HtoBE32(sizeof(int));
data_seg.lkey = lkey;
data_seg.addr = HtoBE64(laddr);
EP_STATIC_ASSERT(sizeof(*ctrl_seg_ptr) == sizeof(int4), "Invalid vectorization");
EP_STATIC_ASSERT(sizeof(*raddr_seg_ptr) == sizeof(int4), "Invalid vectorization");
EP_STATIC_ASSERT(sizeof(*atomic_seg_ptr) == sizeof(int4), "Invalid vectorization");
EP_STATIC_ASSERT(sizeof(*data_seg_ptr) == sizeof(int4), "Invalid vectorization");
st_na_relaxed(reinterpret_cast<int4*>(ctrl_seg_ptr), *reinterpret_cast<int4*>(&ctrl_seg));
st_na_relaxed(reinterpret_cast<int4*>(raddr_seg_ptr), *reinterpret_cast<int4*>(&raddr_seg));
st_na_relaxed(reinterpret_cast<int4*>(atomic_seg_ptr), *reinterpret_cast<int4*>(&atomic_seg_1));
st_na_relaxed(reinterpret_cast<int4*>(data_seg_ptr), *reinterpret_cast<int4*>(&data_seg));
}
__device__ __forceinline__ void nvshmemi_ibgda_amo_nonfetch_add(
void* rptr, const int& value, int pe, int qp_id, bool is_local_copy = false) {
if (is_local_copy) {
atomicAdd(static_cast<unsigned long long*>(rptr), value);
} else {
nvshmemi_ibgda_device_qp_t* qp = ibgda_get_rc(pe, qp_id);
__be32 rkey;
uint64_t raddr;
ibgda_get_rkey(reinterpret_cast<uint64_t>(rptr), pe, &raddr, &rkey, qp->dev_idx);
uint64_t my_wqe_idx = ibgda_reserve_wqe_slots(qp, 1);
void* wqe_ptrs = ibgda_get_wqe_ptr(qp, my_wqe_idx);
ibgda_write_amo_add_wqe(qp, value, reinterpret_cast<uint64_t>(qp->ibuf.buf), qp->ibuf.lkey, raddr, rkey, my_wqe_idx, &wqe_ptrs);
ibgda_submit_requests<true>(qp, my_wqe_idx, 1);
}
}
__device__ __forceinline__ uint64_t nvshmemi_get_p2p_ptr(const uint64_t& ptr, const int& rank, const int& dst_rank) {
// Local rank, no need for mapping
if (rank == dst_rank)
return ptr;
auto peer_base = __ldg(reinterpret_cast<uint64_t*>(nvshmemi_device_state_d.peer_heap_base_p2p) + dst_rank);
// RDMA connected
if (peer_base == 0)
return 0;
// NVLink P2P is enabled
return peer_base + (ptr - reinterpret_cast<uint64_t>(nvshmemi_device_state_d.heap_base));
}
// This is a simplified version of NVSHMEM's `ibgda_poll_cq`.
// Note that this implementation does not guarantee thread safety,
// so we must ensure that no other threads are concurrently using the same QP.
__device__ static __forceinline__ void ibgda_poll_cq(nvshmemi_ibgda_device_cq_t* cq, uint64_t idx) {
const auto cqe64 = static_cast<mlx5_cqe64*>(cq->cqe);
const uint32_t ncqes = cq->ncqes;
memory_fence_cta();
if (*cq->cons_idx >= idx)
return;
// NOTES: this while loop is part of do-while below.
// `wqe_counter` is the HW consumer index. However, we always maintain `index + 1`.
// To be able to compare with the index, we need to use `wqe_counter + 1`.
// Because `wqe_counter` is `uint16_t`, it may be overflow. Still, we know for
// sure that if `idx - wqe_counter - 1 < ncqes`, `wqe_counter + 1 is less than
// idx, and thus we need to wait. We don't need to wait when `idx == wqe_counter + 1`
// That's why we use `- 2` here to make this case overflow.
uint16_t wqe_counter;
do {
wqe_counter = HtoBE16(ld_na_relaxed(&cqe64->wqe_counter));
} while ((static_cast<uint16_t>(static_cast<uint16_t>(idx) - wqe_counter - static_cast<uint16_t>(2)) < ncqes));
*cq->cons_idx = idx;
// Prevent reordering of this function and later instructions
memory_fence_cta();
}
// Wait until wqe `idx - 1` is completed.
__device__ static __forceinline__ void nvshmemi_ibgda_quiet(int dst_pe, int qp_id) {
auto qp = ibgda_get_rc(dst_pe, qp_id);
auto state = ibgda_get_state();
uint64_t prod_idx = state->use_async_postsend ? ld_na_relaxed(qp->tx_wq.prod_idx) : ld_na_relaxed(&qp->mvars.tx_wq.ready_head);
ibgda_poll_cq(qp->tx_wq.cq, prod_idx);
}
} // namespace deep_ep
================================================
FILE: csrc/kernels/internode.cu
================================================
#include <functional>
#include <optional>
#include "buffer.cuh"
#include "configs.cuh"
#include "exception.cuh"
#include "ibgda_device.cuh"
#include "launch.cuh"
#include "utils.cuh"
namespace deep_ep {
namespace internode {
extern nvshmem_team_t cpu_rdma_team;
struct SourceMeta {
int src_rdma_rank, is_token_in_nvl_rank_bits;
EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS == 8, "Invalid number of maximum NVL peers");
__forceinline__ SourceMeta() = default;
// TODO: faster encoding
__device__ __forceinline__ SourceMeta(int rdma_rank, const bool* is_token_in_nvl_ranks) {
src_rdma_rank = rdma_rank;
is_token_in_nvl_rank_bits = is_token_in_nvl_ranks[0];
#pragma unroll
for (int i = 1; i < NUM_MAX_NVL_PEERS; ++i)
is_token_in_nvl_rank_bits |= is_token_in_nvl_ranks[i] << i;
}
__device__ __forceinline__ bool is_token_in_nvl_rank(int nvl_rank) const { return (is_token_in_nvl_rank_bits >> nvl_rank) & 1; }
};
EP_STATIC_ASSERT(sizeof(SourceMeta) % sizeof(int) == 0, "Invalid size of `SourceMeta`");
int get_source_meta_bytes() {
return sizeof(SourceMeta);
}
__host__ __device__ __forceinline__ int get_num_bytes_per_token(int hidden_int4, int num_scales, int num_topk_idx, int num_topk_weights) {
return static_cast<int>(align_up(hidden_int4 * sizeof(int4) + sizeof(SourceMeta) + num_scales * sizeof(float) +
num_topk_idx * sizeof(int) + num_topk_weights * sizeof(float),
sizeof(int4)));
}
__host__ __device__ __forceinline__ std::pair<int, int> get_rdma_clean_meta(int hidden_int4,
int num_scales,
int num_topk_idx,
int num_topk_weights,
int num_rdma_ranks,
int num_rdma_recv_buffer_tokens,
int num_channels) {
// Return `int32_t` offset and count to clean
return {(get_num_bytes_per_token(hidden_int4, num_scales, num_topk_idx, num_topk_weights) * num_rdma_recv_buffer_tokens *
num_rdma_ranks * 2 * num_channels) /
sizeof(int),
(NUM_MAX_NVL_PEERS * 2 + 4) * num_rdma_ranks * 2 * num_channels};
}
__host__ __device__ __forceinline__ std::pair<int, int> get_nvl_clean_meta(int hidden_int4,
int num_scales,
int num_topk_idx,
int num_topk_weights,
int num_rdma_ranks,
int num_nvl_ranks,
int num_nvl_recv_buffer_tokens,
int num_channels,
bool is_dispatch) {
// Return `int32_t` offset and to clean
EP_STATIC_ASSERT(sizeof(SourceMeta) % sizeof(int) == 0, "Invalid size of `SourceMeta`");
return {
(num_nvl_recv_buffer_tokens * get_num_bytes_per_token(hidden_int4, num_scales, num_topk_idx, num_topk_weights) * num_nvl_ranks *
num_channels) /
sizeof(int),
num_nvl_ranks * (2 * num_rdma_ranks + 2) * num_channels,
};
}
template <bool kLowLatencyMode>
__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;
}
template <bool kLowLatencyMode>
__forceinline__ __device__ void nvshmem_sync_with_same_gpu_idx(const nvshmem_team_t& rdma_team) {
kLowLatencyMode ? void(nvshmem_sync(rdma_team)) : nvshmem_sync_all();
}
template <bool kLowLatencyMode, int kNumRDMARanks>
__global__ void notify_dispatch(const int* num_tokens_per_rank,
int* moe_recv_counter_mapped,
int num_ranks,
const int* num_tokens_per_rdma_rank,
int* moe_recv_rdma_counter_mapped,
const int* num_tokens_per_expert,
int* moe_recv_expert_counter_mapped,
int num_experts,
const bool* is_token_in_rank,
int num_tokens,
int num_worst_tokens,
int num_channels,
int expert_alignment,
const int rdma_clean_offset,
const int rdma_num_int_clean,
const int nvl_clean_offset,
const int nvl_num_int_clean,
int* rdma_channel_prefix_matrix,
int* recv_rdma_rank_prefix_sum,
int* gbl_channel_prefix_matrix,
int* recv_gbl_rank_prefix_sum,
void* rdma_buffer_ptr,
void** buffer_ptrs,
int** barrier_signal_ptrs,
int rank,
const nvshmem_team_t rdma_team) {
auto sm_id = static_cast<int>(blockIdx.x);
auto thread_id = static_cast<int>(threadIdx.x), warp_id = thread_id / 32, lane_id = get_lane_id();
auto num_threads = static_cast<int>(blockDim.x), num_warps = num_threads / 32;
auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS;
auto num_rdma_experts = num_experts / kNumRDMARanks, num_nvl_experts = num_rdma_experts / NUM_MAX_NVL_PEERS;
if (sm_id == 0) {
// Communication with others
// Global barrier: the first warp does intra-node sync, the second warp does internode sync
EP_DEVICE_ASSERT(num_warps > 1);
EP_DEVICE_ASSERT(kNumRDMARanks <= num_threads);
// waiting for all previous inflight wrs to complete,
// in case of rewriting cleared rdma_buffer
auto qps_per_rdma_rank = ibgda_get_state()->num_rc_per_pe * ibgda_get_state()->num_devices_initialized;
for (int i = thread_id; i < qps_per_rdma_rank * (kNumRDMARanks - 1); i += num_threads) {
auto dst_rdma_rank = (i / qps_per_rdma_rank + rdma_rank + 1) % kNumRDMARanks;
auto qp_id = i % qps_per_rdma_rank;
nvshmemi_ibgda_quiet(translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank), qp_id);
}
__syncthreads();
if (thread_id == 32)
nvshmem_sync_with_same_gpu_idx<kLowLatencyMode>(rdma_team);
barrier_block<NUM_MAX_NVL_PEERS, true>(barrier_signal_ptrs, nvl_rank);
// Send numbers of tokens per rank/expert to RDMA ranks
auto rdma_buffer_ptr_int = static_cast<int*>(rdma_buffer_ptr);
auto rdma_recv_num_tokens_mixed = SymBuffer<int>(rdma_buffer_ptr, NUM_MAX_NVL_PEERS + num_rdma_experts + 1, kNumRDMARanks);
// Clean up for later data dispatch
EP_DEVICE_ASSERT(rdma_recv_num_tokens_mixed.total_bytes <= rdma_clean_offset * sizeof(int));
#pragma unroll
for (int i = thread_id; i < rdma_num_int_clean; i += num_threads)
rdma_buffer_ptr_int[rdma_clean_offset + i] = 0;
// Copy to send buffer
#pragma unroll
for (int i = thread_id; i < num_ranks; i += num_threads)
rdma_recv_num_tokens_mixed.send_buffer(i / NUM_MAX_NVL_PEERS)[i % NUM_MAX_NVL_PEERS] = num_tokens_per_rank[i];
#pragma unroll
for (int i = thread_id; i < num_experts; i += num_threads)
rdma_recv_num_tokens_mixed.send_buffer(i / num_rdma_experts)[NUM_MAX_NVL_PEERS + i % num_rdma_experts] =
num_tokens_per_expert[i];
if (thread_id < kNumRDMARanks)
rdma_recv_num_tokens_mixed.send_buffer(thread_id)[NUM_MAX_NVL_PEERS + num_rdma_experts] = num_tokens_per_rdma_rank[thread_id];
__syncthreads();
// Issue send
// TODO: more light fence or barrier or signaling
// TODO: overlap EP barrier and NVL cleaning
for (int i = warp_id; i < kNumRDMARanks; i += num_warps) {
if (i != rdma_rank) {
nvshmemi_ibgda_put_nbi_warp<true>(reinterpret_cast<uint64_t>(rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank)),
reinterpret_cast<uint64_t>(rdma_recv_num_tokens_mixed.send_buffer(i)),
(NUM_MAX_NVL_PEERS + num_rdma_experts + 1) * sizeof(int),
translate_dst_rdma_rank<kLowLatencyMode>(i, nvl_rank),
0,
lane_id,
0);
} else {
UNROLLED_WARP_COPY(1,
lane_id,
NUM_MAX_NVL_PEERS + num_rdma_experts + 1,
rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank),
rdma_recv_num_tokens_mixed.send_buffer(i),
ld_volatile_global,
st_na_global);
}
}
__syncthreads();
// Wait previous operations to be finished
if (thread_id < kNumRDMARanks and thread_id != rdma_rank)
nvshmemi_ibgda_quiet(translate_dst_rdma_rank<kLowLatencyMode>(thread_id, nvl_rank), 0);
__syncthreads();
// Barrier
if (thread_id == 0)
nvshmem_sync_with_same_gpu_idx<kLowLatencyMode>(rdma_team);
__syncthreads();
// NVL buffers
auto nvl_send_buffer = thread_id < NUM_MAX_NVL_PEERS ? buffer_ptrs[thread_id] : nullptr;
auto nvl_recv_buffer = buffer_ptrs[nvl_rank];
auto nvl_reduced_num_tokens_per_expert = Buffer<int>(nvl_recv_buffer, num_rdma_experts).advance_also(nvl_send_buffer);
auto nvl_send_num_tokens_per_rank = AsymBuffer<int>(nvl_send_buffer, kNumRDMARanks, NUM_MAX_NVL_PEERS);
auto nvl_send_num_tokens_per_expert = AsymBuffer<int>(nvl_send_buffer, num_nvl_experts, NUM_MAX_NVL_PEERS);
auto nvl_recv_num_tokens_per_rank = AsymBuffer<int>(nvl_recv_buffer, kNumRDMARanks, NUM_MAX_NVL_PEERS);
auto nvl_recv_num_tokens_per_expert = AsymBuffer<int>(nvl_recv_buffer, num_nvl_experts, NUM_MAX_NVL_PEERS);
// Clean up for later data dispatch
auto nvl_buffer_ptr_int = static_cast<int*>(buffer_ptrs[nvl_rank]);
EP_DEVICE_ASSERT(nvl_reduced_num_tokens_per_expert.total_bytes + nvl_send_num_tokens_per_rank.total_bytes +
nvl_send_num_tokens_per_expert.total_bytes <=
nvl_clean_offset * sizeof(int));
#pragma unroll
for (int i = thread_id; i < nvl_num_int_clean; i += num_threads)
nvl_buffer_ptr_int[nvl_clean_offset + i] = 0;
// Reduce number of tokens per expert into the NVL send buffer
// TODO: may use NVSHMEM reduction
EP_DEVICE_ASSERT(num_rdma_experts <= num_threads);
if (thread_id < num_rdma_experts) {
int sum = 0;
#pragma unroll
for (int i = 0; i < kNumRDMARanks; ++i)
sum += rdma_recv_num_tokens_mixed.recv_buffer(i)[NUM_MAX_NVL_PEERS + thread_id];
nvl_reduced_num_tokens_per_expert[thread_id] = sum;
}
__syncthreads();
// Reduce RDMA received tokens
if (thread_id == 0) {
int sum = 0;
#pragma unroll
for (int i = 0; i < kNumRDMARanks; ++i) {
sum += rdma_recv_num_tokens_mixed.recv_buffer(i)[NUM_MAX_NVL_PEERS + num_rdma_experts];
recv_rdma_rank_prefix_sum[i] = sum;
}
if (num_worst_tokens == 0) {
while (ld_volatile_global(moe_recv_rdma_counter_mapped) != -1)
;
*moe_recv_rdma_counter_mapped = sum;
}
}
// Send numbers of tokens per rank/expert to NVL ranks
EP_DEVICE_ASSERT(NUM_MAX_NVL_PEERS <= num_threads);
if (thread_id < NUM_MAX_NVL_PEERS) {
#pragma unroll
for (int i = 0; i < kNumRDMARanks; ++i)
nvl_send_num_tokens_per_rank.buffer(nvl_rank)[i] = rdma_recv_num_tokens_mixed.recv_buffer(i)[thread_id];
#pragma unroll
for (int i = 0; i < num_nvl_experts; ++i)
nvl_send_num_tokens_per_expert.buffer(nvl_rank)[i] = nvl_reduced_num_tokens_per_expert[thread_id * num_nvl_experts + i];
}
barrier_block<NUM_MAX_NVL_PEERS>(barrier_signal_ptrs, nvl_rank);
// Reduce the number of tokens per rank/expert
EP_DEVICE_ASSERT(num_nvl_experts <= num_threads);
if (thread_id == 0) {
int sum = 0;
#pragma unroll
for (int i = 0; i < num_ranks; ++i) {
int src_rdma_rank = i / NUM_MAX_NVL_PEERS, src_nvl_rank = i % NUM_MAX_NVL_PEERS;
sum += nvl_recv_num_tokens_per_rank.buffer(src_nvl_rank)[src_rdma_rank];
recv_gbl_rank_prefix_sum[i] = sum;
}
if (num_worst_tokens == 0) {
while (ld_volatile_global(moe_recv_counter_mapped) != -1)
;
*moe_recv_counter_mapped = sum;
}
}
if (thread_id < num_nvl_experts) {
int sum = 0;
#pragma unroll
for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i)
sum += nvl_recv_num_tokens_per_expert.buffer(i)[thread_id];
sum = (sum + expert_alignment - 1) / expert_alignment * expert_alignment;
if (num_worst_tokens == 0) {
while (ld_volatile_global(moe_recv_expert_counter_mapped + thread_id) != -1)
;
moe_recv_expert_counter_mapped[thread_id] = sum;
}
}
// Finally barrier
if (thread_id == 32)
nvshmem_sync_with_same_gpu_idx<kLowLatencyMode>(rdma_team);
barrier_block<NUM_MAX_NVL_PEERS>(barrier_signal_ptrs, nvl_rank);
} else {
// Calculate meta data
int dst_rdma_rank = sm_id - 1;
for (int channel_id = warp_id; channel_id < num_channels; channel_id += num_warps) {
int token_start_idx, token_end_idx;
get_channel_task_range(num_tokens, num_channels, channel_id, token_start_idx, token_end_idx);
// Iterate over tokens
int total_count = 0, per_nvl_rank_count
gitextract_l_jv2phj/
├── .clang-format
├── .github/
│ └── workflows/
│ └── format.yml
├── .gitignore
├── LICENSE
├── README.md
├── csrc/
│ ├── CMakeLists.txt
│ ├── config.hpp
│ ├── deep_ep.cpp
│ ├── deep_ep.hpp
│ ├── event.hpp
│ └── kernels/
│ ├── CMakeLists.txt
│ ├── api.cuh
│ ├── buffer.cuh
│ ├── configs.cuh
│ ├── exception.cuh
│ ├── ibgda_device.cuh
│ ├── internode.cu
│ ├── internode_ll.cu
│ ├── intranode.cu
│ ├── launch.cuh
│ ├── layout.cu
│ ├── runtime.cu
│ └── utils.cuh
├── deep_ep/
│ ├── __init__.py
│ ├── buffer.py
│ └── utils.py
├── format.sh
├── install.sh
├── pyproject.toml
├── requirements-lint.txt
├── setup.py
├── tests/
│ ├── test_internode.py
│ ├── test_intranode.py
│ ├── test_low_latency.py
│ └── utils.py
└── third-party/
├── README.md
└── nvshmem.patch
SYMBOL INDEX (89 symbols across 11 files)
FILE: csrc/config.hpp
type deep_ep (line 6) | namespace deep_ep {
function dtype_t (line 9) | dtype_t ceil_div(dtype_t a, dtype_t b) {
function dtype_t (line 14) | dtype_t align_up(dtype_t a, dtype_t b) {
function dtype_t (line 19) | dtype_t align_down(dtype_t a, dtype_t b) {
type Config (line 23) | struct Config {
method Config (line 30) | Config(int num_sms,
method get_nvl_buffer_size_hint (line 52) | size_t get_nvl_buffer_size_hint(size_t hidden_bytes, int num_ranks) ...
method get_rdma_buffer_size_hint (line 76) | size_t get_rdma_buffer_size_hint(int64_t hidden_bytes, int num_ranks...
type LowLatencyBuffer (line 107) | struct LowLatencyBuffer {
method clean_meta (line 121) | std::pair<int*, int> clean_meta() {
type LowLatencyLayout (line 127) | struct LowLatencyLayout {
method out_ptr_t (line 132) | out_ptr_t advance(const in_ptr_t& ptr, size_t count) {
method LowLatencyLayout (line 136) | LowLatencyLayout(void* rdma_buffer, int num_max_dispatch_tokens_per_...
function get_low_latency_rdma_size_hint (line 190) | size_t get_low_latency_rdma_size_hint(int num_max_dispatch_tokens_per_...
FILE: csrc/deep_ep.cpp
type shared_memory (line 15) | namespace shared_memory {
function cu_mem_set_access_all (line 16) | void cu_mem_set_access_all(void* ptr, size_t size) {
function cu_mem_free (line 30) | void cu_mem_free(void* ptr) {
function get_size_align_to_granularity (line 42) | size_t get_size_align_to_granularity(size_t size_raw, size_t granulari...
type deep_ep (line 126) | namespace deep_ep {
function is_sm90_compiled (line 1817) | bool is_sm90_compiled() {
function PYBIND11_MODULE (line 1846) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
FILE: csrc/deep_ep.hpp
type shared_memory (line 24) | namespace shared_memory {
type MemHandle (line 31) | struct MemHandle {
class SharedMemoryAllocator (line 38) | class SharedMemoryAllocator {
type deep_ep (line 52) | namespace deep_ep {
FILE: csrc/event.hpp
type deep_ep (line 7) | namespace deep_ep {
type EventHandle (line 9) | struct EventHandle {
method EventHandle (line 12) | EventHandle() {
method EventHandle (line 17) | explicit EventHandle(const at::cuda::CUDAStream& stream) {
method EventHandle (line 22) | EventHandle(const EventHandle& other) = default;
method current_stream_wait (line 24) | void current_stream_wait() const { at::cuda::getCurrentCUDAStream()....
function create_event (line 27) | torch::Event create_event(const at::cuda::CUDAStream& s) {
function stream_wait (line 33) | void stream_wait(const at::cuda::CUDAStream& s_0, const at::cuda::CUDA...
function stream_wait (line 38) | void stream_wait(const at::cuda::CUDAStream& s, const EventHandle& eve...
FILE: deep_ep/buffer.py
class Buffer (line 13) | class Buffer:
method __init__ (line 32) | def __init__(self,
method destroy (line 138) | def destroy(self):
method is_sm90_compiled (line 150) | def is_sm90_compiled():
method set_num_sms (line 154) | def set_num_sms(new_num_sms: int) -> None:
method capture (line 166) | def capture() -> EventOverlap:
method get_low_latency_rdma_size_hint (line 176) | def get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank: i...
method get_comm_stream (line 191) | def get_comm_stream(self) -> torch.Stream:
method get_local_buffer_tensor (line 201) | def get_local_buffer_tensor(self,
method _unpack_bias (line 223) | def _unpack_bias(bias: Union[torch.Tensor, Tuple[torch.Tensor, torch.T...
method get_dispatch_config (line 233) | def get_dispatch_config(num_ranks: int) -> Config:
method get_combine_config (line 263) | def get_combine_config(num_ranks: int) -> Config:
method get_dispatch_layout (line 293) | def get_dispatch_layout(self, topk_idx: torch.Tensor, num_experts: int,
method dispatch (line 322) | def dispatch(self, x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Te...
method combine (line 405) | def combine(self, x: torch.Tensor, handle: Tuple,
method internode_dispatch (line 453) | def internode_dispatch(self, x: Union[torch.Tensor, Tuple[torch.Tensor...
method internode_combine (line 504) | def internode_combine(self, x: torch.Tensor, handle: Union[tuple, list],
method clean_low_latency_buffer (line 533) | def clean_low_latency_buffer(self, num_max_dispatch_tokens_per_rank: i...
method low_latency_dispatch (line 548) | def low_latency_dispatch(self, x: torch.Tensor, topk_idx: torch.Tensor,
method low_latency_combine (line 617) | def low_latency_combine(self, x: torch.Tensor, topk_idx: torch.Tensor,...
method low_latency_update_mask_buffer (line 663) | def low_latency_update_mask_buffer(self, rank_to_mask: int, mask: bool...
method low_latency_query_mask_buffer (line 674) | def low_latency_query_mask_buffer(self, mask_status: torch.Tensor):
method low_latency_clean_mask_buffer (line 684) | def low_latency_clean_mask_buffer(self):
method get_next_low_latency_combine_buffer (line 691) | def get_next_low_latency_combine_buffer(self, handle: object):
FILE: deep_ep/utils.py
class EventOverlap (line 10) | class EventOverlap:
method __init__ (line 19) | def __init__(self, event: Optional[EventHandle] = None, extra_tensors:...
method current_stream_wait (line 33) | def current_stream_wait(self) -> None:
method __enter__ (line 40) | def __enter__(self) -> Any:
method __exit__ (line 54) | def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
function check_nvlink_connections (line 64) | def check_nvlink_connections(group: dist.ProcessGroup):
FILE: setup.py
function get_nvshmem_host_lib_name (line 11) | def get_nvshmem_host_lib_name(base_dir):
FILE: tests/test_internode.py
function test_main (line 16) | def test_main(args: argparse.Namespace,
function test_loop (line 316) | def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Name...
FILE: tests/test_intranode.py
function test_main (line 15) | def test_main(args: argparse.Namespace, num_sms: int, local_rank: int, n...
function test_loop (line 266) | def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Name...
FILE: tests/test_low_latency.py
function simulate_failure_and_skip (line 12) | def simulate_failure_and_skip(rank: int, api: Literal["dispatch", "combi...
function query_mask_buffer_and_check (line 31) | def query_mask_buffer_and_check(api: Literal["dispatch", "combine", "cle...
function test_main (line 37) | def test_main(num_tokens: int,
function test_loop (line 253) | def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Name...
FILE: tests/utils.py
function init_dist (line 14) | def init_dist(local_rank: int, num_local_ranks: int):
function calc_diff (line 39) | def calc_diff(x: torch.Tensor, y: torch.Tensor):
function align_up (line 46) | def align_up(x, y):
function per_token_cast_to_fp8 (line 50) | def per_token_cast_to_fp8(x: torch.Tensor):
function per_token_cast_back (line 61) | def per_token_cast_back(x_fp8: torch.Tensor, x_scales: torch.Tensor):
function inplace_unique (line 77) | def inplace_unique(x: torch.Tensor, num_slots: int):
function create_grouped_scores (line 92) | def create_grouped_scores(scores: torch.Tensor, group_idx: torch.Tensor,...
function bench (line 100) | def bench(fn, num_warmups: int = 50, num_tests: int = 50, post_fn=None):
class empty_suppress (line 128) | class empty_suppress:
method __enter__ (line 130) | def __enter__(self):
method __exit__ (line 133) | def __exit__(self, *_):
class suppress_stdout_stderr (line 137) | class suppress_stdout_stderr:
method __enter__ (line 139) | def __enter__(self):
method __exit__ (line 159) | def __exit__(self, *_):
function bench_kineto (line 173) | def bench_kineto(fn,
function hash_tensor (line 241) | def hash_tensor(t: torch.Tensor):
Condensed preview — 37 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (651K chars).
[
{
"path": ".clang-format",
"chars": 488,
"preview": "BasedOnStyle: Google\nUseTab: Never\nIndentWidth: 4\nColumnLimit: 140\nAccessModifierOffset: -4\n\n# Force pointers to the typ"
},
{
"path": ".github/workflows/format.yml",
"chars": 507,
"preview": "name: Code Format Check\n\non:\n push:\n branches: [ main ]\n pull_request:\n branches: [ main ]\n\njobs:\n format-check"
},
{
"path": ".gitignore",
"chars": 85,
"preview": "compile_commands.json\n.idea\n.DS_Store\n*.pyc\nbuild/\n.cache/\n.vscode/\n*/cmake-build-*/\n"
},
{
"path": "LICENSE",
"chars": 1065,
"preview": "MIT License\n\nCopyright (c) 2025 DeepSeek\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\no"
},
{
"path": "README.md",
"chars": 21309,
"preview": "# DeepEP\n\nDeepEP is a communication library tailored for Mixture-of-Experts (MoE) and expert parallelism (EP). It provid"
},
{
"path": "csrc/CMakeLists.txt",
"chars": 1469,
"preview": "# NOTES: this CMake is only for debugging; for setup, please use Torch extension\ncmake_minimum_required(VERSION 3.10)\npr"
},
{
"path": "csrc/config.hpp",
"chars": 9860,
"preview": "#pragma once\n\n#include \"kernels/api.cuh\"\n#include \"kernels/exception.cuh\"\n\nnamespace deep_ep {\n\ntemplate <typename dtype"
},
{
"path": "csrc/deep_ep.cpp",
"chars": 90719,
"preview": "#include \"deep_ep.hpp\"\n\n#include <ATen/cuda/CUDAContext.h>\n#include <ATen/cuda/CUDADataType.h>\n#include <cuda_runtime.h>"
},
{
"path": "csrc/deep_ep.hpp",
"chars": 10854,
"preview": "#pragma once\n\n// Forcibly disable NDEBUG\n#ifdef NDEBUG\n#undef NDEBUG\n#endif\n\n#include <pybind11/pybind11.h>\n#include <py"
},
{
"path": "csrc/event.hpp",
"chars": 1093,
"preview": "#include <ATen/cuda/CUDAContext.h>\n\n#include <memory>\n\n#include \"kernels/exception.cuh\"\n\nnamespace deep_ep {\n\nstruct Eve"
},
{
"path": "csrc/kernels/CMakeLists.txt",
"chars": 887,
"preview": "function(add_deep_ep_library target_name source_file)\n add_library(${target_name} STATIC ${source_file})\n set_targ"
},
{
"path": "csrc/kernels/api.cuh",
"chars": 12147,
"preview": "#pragma once\n\n#include <vector>\n\n#include \"configs.cuh\"\n\nnamespace deep_ep {\n\n// Intranode runtime\nnamespace intranode {"
},
{
"path": "csrc/kernels/buffer.cuh",
"chars": 4809,
"preview": "#pragma once\n\n#include \"configs.cuh\"\n#include \"exception.cuh\"\n\nnamespace deep_ep {\n\ntemplate <typename dtype_t>\nstruct B"
},
{
"path": "csrc/kernels/configs.cuh",
"chars": 2230,
"preview": "#pragma once\n\n#define NUM_MAX_NVL_PEERS 8\n#define NUM_MAX_RDMA_PEERS 20\n#define NUM_WORKSPACE_BYTES (32 * 1024 * 1024)\n#"
},
{
"path": "csrc/kernels/exception.cuh",
"chars": 2850,
"preview": "#pragma once\n\n#include <exception>\n#include <string>\n\n#include \"configs.cuh\"\n\n#ifndef EP_STATIC_ASSERT\n#define EP_STATIC"
},
{
"path": "csrc/kernels/ibgda_device.cuh",
"chars": 22226,
"preview": "// Portions derived from NVSHMEM (https://developer.nvidia.com/nvshmem)\n// Copyright (c) NVIDIA Corporation.\n// Licensed"
},
{
"path": "csrc/kernels/internode.cu",
"chars": 130621,
"preview": "#include <functional>\n#include <optional>\n\n#include \"buffer.cuh\"\n#include \"configs.cuh\"\n#include \"exception.cuh\"\n#includ"
},
{
"path": "csrc/kernels/internode_ll.cu",
"chars": 69550,
"preview": "#include \"configs.cuh\"\n#include \"exception.cuh\"\n#include \"ibgda_device.cuh\"\n#include \"launch.cuh\"\n\nnamespace deep_ep {\n\n"
},
{
"path": "csrc/kernels/intranode.cu",
"chars": 57364,
"preview": "#include \"buffer.cuh\"\n#include \"configs.cuh\"\n#include \"exception.cuh\"\n#include \"launch.cuh\"\n#include \"utils.cuh\"\n\nnamesp"
},
{
"path": "csrc/kernels/launch.cuh",
"chars": 7293,
"preview": "#pragma once\n\n#include \"configs.cuh\"\n#include \"exception.cuh\"\n\n#ifndef SETUP_LAUNCH_CONFIG\n#ifndef DISABLE_SM90_FEATURES"
},
{
"path": "csrc/kernels/layout.cu",
"chars": 6939,
"preview": "#include \"configs.cuh\"\n#include \"exception.cuh\"\n#include \"launch.cuh\"\n\nnamespace deep_ep {\n\nnamespace layout {\n\ntemplate"
},
{
"path": "csrc/kernels/runtime.cu",
"chars": 2962,
"preview": "#include <cstring>\n#include <vector>\n\n#include \"configs.cuh\"\n#include \"exception.cuh\"\n#include \"launch.cuh\"\n#include \"ut"
},
{
"path": "csrc/kernels/utils.cuh",
"chars": 25194,
"preview": "#pragma once\n\n#include \"exception.cuh\"\n\n#define UNROLLED_WARP_COPY(UNROLL_FACTOR, LANE_ID, N, DST, SRC, LD_FUNC, ST_FUNC"
},
{
"path": "deep_ep/__init__.py",
"chars": 155,
"preview": "import torch\n\nfrom .utils import EventOverlap\nfrom .buffer import Buffer\n\n# noinspection PyUnresolvedReferences\nfrom dee"
},
{
"path": "deep_ep/buffer.py",
"chars": 42754,
"preview": "import os\nimport torch\nimport torch.distributed as dist\nfrom typing import Callable, List, Tuple, Optional, Union\n\n# noi"
},
{
"path": "deep_ep/utils.py",
"chars": 3839,
"preview": "import os\nimport torch\nimport torch.distributed as dist\nfrom typing import Any, Optional, Tuple\n\n# noinspection PyUnreso"
},
{
"path": "format.sh",
"chars": 6565,
"preview": "#!/usr/bin/env bash\n# Usage:\n# # Do work and commit your work.\n\n# # Format files that differ from origin/main.\n# "
},
{
"path": "install.sh",
"chars": 269,
"preview": "# Change current directory into project root\noriginal_dir=$(pwd)\nscript_dir=$(dirname \"$0\")\ncd \"$script_dir\"\n\n# Remove o"
},
{
"path": "pyproject.toml",
"chars": 600,
"preview": "[tool.yapf]\nbased_on_style = \"pep8\"\ncolumn_limit = 140\nindent_width = 4\n\n[tool.ruff.lint]\nselect = [\n # pycodestyle\n "
},
{
"path": "requirements-lint.txt",
"chars": 45,
"preview": "clang-format==15.0.7\nyapf==0.40.2\nruff==0.6.5"
},
{
"path": "setup.py",
"chars": 5229,
"preview": "import os\nimport subprocess\nimport setuptools\nimport importlib\n\nfrom pathlib import Path\nfrom torch.utils.cpp_extension "
},
{
"path": "tests/test_internode.py",
"chars": 21468,
"preview": "import argparse\nimport os\nimport time\nimport torch\nimport torch.distributed as dist\n\n# noinspection PyUnresolvedReferenc"
},
{
"path": "tests/test_intranode.py",
"chars": 17247,
"preview": "import argparse\nimport time\nimport torch\nimport torch.distributed as dist\n\n# noinspection PyUnresolvedReferences\nimport "
},
{
"path": "tests/test_low_latency.py",
"chars": 19180,
"preview": "import argparse\nimport random\nimport torch\nimport torch.distributed as dist\nfrom functools import partial\nfrom typing im"
},
{
"path": "tests/utils.py",
"chars": 9161,
"preview": "import inspect\nimport json\nimport tempfile\nfrom pathlib import Path\n\nimport numpy as np\nimport os\nimport sys\nimport torc"
},
{
"path": "third-party/README.md",
"chars": 3508,
"preview": "# Install NVSHMEM\n\n## Important notices\n\n**This project is neither sponsored nor supported by NVIDIA.**\n\n**Use of NVIDIA"
},
{
"path": "third-party/nvshmem.patch",
"chars": 21868,
"preview": "From 9e6cc27cceb3130784e4ea7b61ea3171156365fd Mon Sep 17 00:00:00 2001\nFrom: Shangyan Zhou <sy.zhou@deepseek.com>\nDate: "
}
]
About this extraction
This page contains the full source code of the deepseek-ai/DeepEP GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 37 files (619.5 KB), approximately 148.8k tokens, and a symbol index with 89 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.