mirrors/deepseek-ai--deepep
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Code Format Check / format-check (push) Failing after 1s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:24:32 +08:00
commit b5ecf06f65
107 changed files with 24727 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.clang-format
+21
View File
@@ -0,0 +1,21 @@
BasedOnStyle: Google
UseTab: Never
IndentWidth: 4
ColumnLimit: 140
AccessModifierOffset: -4
# Force pointers to the type for C++.
DerivePointerAlignment: false
PointerAlignment: Left
ReferenceAlignment: Left
AllowShortFunctionsOnASingleLine: Inline
AllowShortIfStatementsOnASingleLine: false
AllowShortLoopsOnASingleLine: false
AlignOperands: false
BreakBeforeBinaryOperators: None
Cpp11BracedListStyle: true
ContinuationIndentWidth: 4
BinPackArguments: false
BinPackParameters: false
.editorconfig
+58
View File
@@ -0,0 +1,58 @@
# https://editorconfig.org/
root = true
[*]
charset = utf-8
end_of_line = lf
indent_style = space
indent_size = 4
trim_trailing_whitespace = true
insert_final_newline = true
[*.{py,pyi}]
indent_size = 4
[*.{cpp,hpp,cxx,cc,c,h,cu,cuh}]
indent_size = 4
[*.rs]
indent_size = 4
[*.go]
indent_style = tab
[*.{yaml,yml}]
indent_size = 2
[.clang-{format,tidy}]
indent_size = 2
[Makefile]
indent_style = tab
[*.sh]
indent_size = 4
[*.bat]
indent_size = 4
end_of_line = crlf
[*.md]
indent_size = 2
x-soft-wrap-text = true
[*.rst]
indent_size = 4
x-soft-wrap-text = true
[*.{html,xml,css,scss,js,jsx,ts,tsx,vue}]
indent_size = 2
[**/test{,s,ing}/**/*.txt]
trim_trailing_whitespace = false
insert_final_newline = false
[**/example{,s}/**/*.txt]
trim_trailing_whitespace = false
insert_final_newline = false
.github/workflows/format.yml
+28
View File
@@ -0,0 +1,28 @@
name: Code Format Check
on:
push:
branches: [ main ]
pull_request:
branches: [ main ]
jobs:
format-check:
runs-on: ubuntu-latest
steps:
- name: Checkout source
uses: actions/checkout@v6
with:
fetch-depth: 0
- name: Setup environment
run: |
sudo apt-get update
sudo apt-get install -y bash
- name: Run format.sh
run: |
bash ./format.sh
# If format.sh return non-zero, GitHub Actions will mark it as failure.
.gitignore
+28
View File
@@ -0,0 +1,28 @@
# VS Code
compile_commands.json
.clangd
.cache/
.vscode/
# macOS
.DS_Store
# JetBrains
.idea
*/cmake-build-*/
# Python
*.pyc
# Build
build/
*.so
deep_ep.egg-info/
dist/
# Coredump
coredump/
# OpenCode
.sisyphus/
AGENTS.md
.gitmodules
+3
View File
@@ -0,0 +1,3 @@
[submodule "third-party/fmt"]
path = third-party/fmt
url = https://github.com/fmtlib/fmt.git
CMakeLists.txt
+52
View File
@@ -0,0 +1,52 @@
# NOTES: this CMake is only for debugging; for setup, please use Torch extension
cmake_minimum_required(VERSION 3.10)
project(deep_ep LANGUAGES CUDA CXX)
set(CMAKE_VERBOSE_MAKEFILE ON)
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -O3 -fPIC")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC")
set(CUDA_SEPARABLE_COMPILATION ON)
list(APPEND CUDA_NVCC_FLAGS "-DENABLE_FAST_DEBUG")
list(APPEND CUDA_NVCC_FLAGS "-O3")
list(APPEND CUDA_NVCC_FLAGS "--ptxas-options=--verbose,--register-usage-level=10,--warn-on-local-memory-usage")
list(APPEND CUDA_NVCC_FLAGS "-Xcompiler=-rdynamic")
list(APPEND CUDA_NVCC_FLAGS "-lineinfo")
# Suppress warnings for the `fmt` library
list(APPEND CUDA_NVCC_FLAGS "--diag-suppress=128,2417")
set(USE_SYSTEM_NVTX on)
set(CUDA_ARCH_LIST "9.0" CACHE STRING "List of CUDA architectures to compile")
set(TORCH_CUDA_ARCH_LIST "${CUDA_ARCH_LIST}")
find_package(CUDAToolkit REQUIRED)
find_package(pybind11 REQUIRED)
find_package(Torch REQUIRED)
# NVSHMEM
find_package(NVSHMEM REQUIRED HINTS ${NVSHMEM_ROOT_DIR}/lib/cmake/nvshmem)
add_library(nvshmem ALIAS nvshmem::nvshmem)
add_library(nvshmem_host ALIAS nvshmem::nvshmem_host)
add_library(nvshmem_device ALIAS nvshmem::nvshmem_device)
# NCCL
# TODO: use `find_package` instead of manual checks
if (NOT NCCL_ROOT_DIR)
message(FATAL_ERROR "NCCL_ROOT_DIR is not set.")
endif()
set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CUDA_STANDARD 20)
include_directories(deep_ep/include third-party/fmt/include)
include_directories(${CUDA_TOOLKIT_ROOT_DIR}/include ${CUDA_TOOLKIT_ROOT_DIR}/include/cccl ${TORCH_INCLUDE_DIRS} ${PYTHON_INCLUDE_DIRS} ${NVSHMEM_INCLUDE_DIR} ${NCCL_ROOT_DIR}/include .)
link_directories(${TORCH_INSTALL_PREFIX}/lib ${CUDA_TOOLKIT_ROOT_DIR}/lib ${NVSHMEM_LIB_DIR} ${NCCL_ROOT_DIR}/lib)
# Add kernels
add_subdirectory(csrc)
# Link CPP and CUDA together
pybind11_add_module(_C csrc/python_api.cpp)
target_link_libraries(_C PRIVATE nccl ${RUNTIME_CUDA_LIBRARIES} ${LEGACY_CUDA_LIBRARIES} ${TORCH_LIBRARIES} torch_python)
# Enable kernel code indexing with CMake-based IDEs
cuda_add_library(deep_ep_indexing_cuda STATIC csrc/indexing/main.cu)
LICENSE
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2025 DeepSeek
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
README.md
+449
View File
@@ -0,0 +1,449 @@
# DeepEP
DeepEP (DeepEveryParallel) is a high-performance communication library for modern machine learning training and inference. The library currently focuses on expert parallelism (EP) — providing high-throughput and low-latency all-to-all GPU kernels (MoE dispatch and combine) with low-precision support including FP8 — while also offering experimental primitives for pipeline parallelism (PP), context parallelism (CP), and remote memory access (Engram), all designed for zero or minimal SM occupation. All kernels are compiled at runtime via a lightweight Just-In-Time (JIT) module, requiring no CUDA compilation during installation.
Despite its lightweight design, DeepEP's performance matches or exceeds hardware bandwidth limits across various configurations.
## News
- **V2 release**: A complete refactoring of Expert Parallelism — achieving extreme performance with several times fewer SM resources compared to V1, while supporting significantly larger scale-up and scale-out domains. V2 has also switched from the NVSHMEM backend to the more lightweight **NCCL Gin backend**.
### New features
- **Fully JIT** (Just-In-Time compilation)
- **NCCL Gin backend**
- Header-only & lightweight
- Able to reuse existing NCCL communicators
- **EPv2**
- High-throughput and low-latency APIs unified into a single `ElasticBuffer` interface, with a new GEMM layout
- Larger scale-up & scale-out domain support (up to EP2048)
- Analytical SM & QP count calculation — no more auto-tuning needed
- Both hybrid & direct modes remain supported
- For V3-like legacy training, SM usage reduced from 24 to 4 - 6 while maintaining equivalent or better performance
- **0 SM Engram** (with RDMA)
- **0 SM PP** (with RDMA)
- **0 SM CP** (with Copy Engine)
### Notes
- Buffer size consumption is larger than V1
- 0 SM RDMA low-latency EP is no longer supported
- Engram, PP, and CP are experimental features
### Still on-going features
- **Elastic GPU & CPU buffers**: A contiguous virtual address space that maps to a hybrid of GPU and CPU physical memory under the hood, enabling fully automatic and transparent Engram or imbalanced EP
- Reducing intermediate buffer sizes by leveraging EP replay to handle load imbalance
- All-gather updates and reduce-scatter implementations for DP & TP
For the legacy V1 documentation (NVSHMEM-based), see [docs/legacy.md](docs/legacy.md).
## Performance
Following V3's configuration, we tested with 8K tokens per batch, 7168 hidden dimensions, top 8 experts, FP8 dispatching, and BF16 combining, and obtained the following results:
| Arch | NIC type | Topo | Dispatch Bottleneck Bandwidth | Combine Bottleneck Bandwidth | #SMs |
|--|--|--|--|--|--|
| SM90 | CX7 | EP 8 x 2 | 90 GB/s (RDMA) | 81 GB/s (RDMA) | 12 |
| SM90 | CX7 | EP 8 x 4 | 61 GB/s (RDMA) | 61 GB/s (RDMA) | 6 |
| SM100 | CX7 | EP 8 x 2 | 90 GB/s (RDMA) | 91 GB/s (RDMA) | 12 |
| SM100 | N/A | EP 8 | 726 GB/s (NVLink) | 740 GB/s (NVLink) | 64 (Max perf) |
| SM100 | N/A | EP 8 | 643 GB/s (NVLink) | 675 GB/s (NVLink) | 24 (Min #SM) |
Notes: the results are logical bandwidth. For example, under the `EP 8 x 2` case, 90 GB/s actually contains local rank traffic.
Comparing with V1, **V2 achieves up to 1.3x peak performance, while saving up to 4x SM count**.
We omit results for larger EP configurations for the time being, but encourage interested users to benchmark them directly. Based on our internal experience, we expect the kernel to continue saturating hardware bandwidth at scale.
For V1 performance data, see [docs/legacy.md](docs/legacy.md#performance).
## Quick start
### Requirements
- Hopper (SM90) GPUs, or other architectures with SM90 PTX ISA support
- Python 3.8 and above
- CUDA version
- CUDA 12.3 and above for SM90 GPUs
- PyTorch 2.10 and above
- NCCL 2.30.4 and above
- NVLink for intranode communication
- RDMA network for internode communication
### Install NCCL dependency
We recommend using pip to install NCCL so that DeepEP can automatically locate it within the Python environment. You can install it using the following command:
```bash
pip install "nvidia-nccl-cu13>=2.30.4" --no-deps
```
### Install NVSHMEM dependency
DeepEP also depends on NVSHMEM to provide support for legacy methods. Please refer to our [NVSHMEM Installation Guide](docs/nvshmem.md) for instructions.
### Development
```bash
# Build and make symbolic links for SO files
python setup.py build
# You may modify the specific SO names according to your own platform
ln -s build/lib.linux-x86_64-cpython-38/deep_ep_cpp.cpython-38-x86_64-linux-gnu.so
# Run test cases
# NOTES: you may modify the `init_dist` function in `tests/utils/envs.py`
# according to your own cluster settings, and launch into multiple nodes
python tests/elastic/test_ep.py
python tests/elastic/test_agrs.py
python tests/elastic/test_engram.py
python tests/elastic/test_pp.py
```
### Installation
```bash
python setup.py install
```
Then, import `deep_ep` in your Python project, and enjoy!
## Interfaces and examples
### Buffer initialization
In V2, all EP operations — high-throughput and low-latency — are unified under a single `ElasticBuffer` interface. The buffer can be initialized by specifying MoE settings directly, and the optimal SM and QP counts are calculated analytically.
```python
import torch
import torch.distributed as dist
from typing import Optional
from deep_ep import ElasticBuffer
# Communication buffer (will allocate at runtime)
_buffer: Optional[ElasticBuffer] = None
# Number of SMs to use for communication kernels (will be set at buffer creation)
_num_comm_sms: int = 0
def get_buffer(group: dist.ProcessGroup,
num_max_tokens_per_rank: int,
hidden: int,
num_topk: int,
num_experts: int,
use_fp8_dispatch: bool = False) -> ElasticBuffer:
"""Initialize or retrieve the ElasticBuffer for EP communication."""
global _buffer, _num_comm_sms
# Check if we can reuse the existing buffer
required_bytes = ElasticBuffer.get_buffer_size_hint(
group, num_max_tokens_per_rank, hidden,
num_topk=num_topk, use_fp8_dispatch=use_fp8_dispatch,
)
if _buffer is not None and _buffer.group == group and _buffer.num_bytes >= required_bytes:
return _buffer
# Allocate a new buffer with MoE settings
# NOTES: V2 buffer size consumption is larger than V1
_buffer = ElasticBuffer(
group,
num_max_tokens_per_rank=num_max_tokens_per_rank,
hidden=hidden,
num_topk=num_topk,
use_fp8_dispatch=use_fp8_dispatch,
)
# V2 analytically calculates the optimal SM count — no more auto-tuning needed
# You may also specify `num_sms` manually in dispatch/combine calls to override
_num_comm_sms = _buffer.get_theoretical_num_sms(num_experts, num_topk)
return _buffer
```
### Example use in model training or inference prefilling
V2 unifies the dispatch and combine APIs into a single `ElasticBuffer` interface. The example below shows how to use them for training (with backward passes) or inference prefilling.
```python
import torch
import torch.distributed as dist
from typing import Tuple, Union
from deep_ep import ElasticBuffer, EPHandle, EventOverlap
def dispatch_forward(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
topk_idx: torch.Tensor, topk_weights: torch.Tensor,
num_experts: int,
num_max_tokens_per_rank: int,
expert_alignment: int = 1) -> \
Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
torch.Tensor, torch.Tensor, EPHandle, EventOverlap]:
"""
MoE dispatch: route tokens to the corresponding experts across all ranks.
Supports both BF16 and FP8 (x as a tuple of [data, scale_factors]) inputs.
"""
global _buffer, _num_comm_sms
recv_x, recv_topk_idx, recv_topk_weights, handle, event = _buffer.dispatch(
x,
topk_idx=topk_idx,
topk_weights=topk_weights,
num_experts=num_experts,
num_max_tokens_per_rank=num_max_tokens_per_rank,
expert_alignment=expert_alignment,
num_sms=_num_comm_sms,
async_with_compute_stream=True,
)
# `handle` contains routing metadata for the subsequent combine call
# `handle.num_recv_tokens_per_expert_list` provides per-expert token counts for GEMM
# Use `event.current_stream_wait()` to synchronize the compute stream before using results
return recv_x, recv_topk_idx, recv_topk_weights, handle, event
def dispatch_backward(grad_recv_x: torch.Tensor,
grad_recv_topk_weights: torch.Tensor,
handle: EPHandle) -> Tuple[torch.Tensor, torch.Tensor, EventOverlap]:
"""The backward pass of MoE dispatch is actually a combine."""
global _buffer, _num_comm_sms
combined_grad_x, combined_grad_topk_weights, event = _buffer.combine(
grad_recv_x,
handle=handle,
topk_weights=grad_recv_topk_weights,
num_sms=_num_comm_sms,
async_with_compute_stream=True,
)
return combined_grad_x, combined_grad_topk_weights, event
def combine_forward(x: torch.Tensor,
handle: EPHandle) -> Tuple[torch.Tensor, EventOverlap]:
"""MoE combine: reduce expert outputs back to their original ranks."""
global _buffer, _num_comm_sms
combined_x, _, event = _buffer.combine(
x,
handle=handle,
num_sms=_num_comm_sms,
async_with_compute_stream=True,
)
return combined_x, event
def combine_backward(grad_combined_x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
handle: EPHandle) -> \
Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], EventOverlap]:
"""The backward pass of MoE combine is actually a dispatch."""
global _buffer, _num_comm_sms
grad_x, _, _, _, event = _buffer.dispatch(
grad_combined_x,
handle=handle,
num_sms=_num_comm_sms,
async_with_compute_stream=True,
)
return grad_x, event
```
For communication-computation overlap, use the `EventOverlap` interface to manage dependencies between the communication stream and the compute stream:
```python
# After dispatch, overlap computation while communication is in-flight
recv_x, recv_topk_idx, recv_topk_weights, handle, event = dispatch_forward(...)
# ... do some independent computation here ...
# Wait for communication to finish before using results
event.current_stream_wait()
# Now safe to use recv_x, recv_topk_idx, recv_topk_weights
```
### Example use in inference decoding
For inference decoding, the same `ElasticBuffer` is used. The handle-caching pattern allows reusing routing metadata across iterations when the gating decisions remain unchanged, avoiding redundant CPU synchronization.
```python
import torch
from typing import Tuple, Optional, Union
from deep_ep import ElasticBuffer, EPHandle, EventOverlap
def decode_dispatch(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
topk_idx: torch.Tensor, topk_weights: torch.Tensor,
num_experts: int,
num_max_tokens_per_rank: int,
cached_handle: Optional[EPHandle] = None) -> \
Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
torch.Tensor, torch.Tensor, EPHandle, EventOverlap]:
"""
MoE dispatch for inference decoding.
If `cached_handle` is provided, the layout is reused without CPU synchronization.
"""
global _buffer, _num_comm_sms
if cached_handle is not None:
# Reuse cached handle: skip layout recomputation and CPU sync
recv_x, _, _, handle, event = _buffer.dispatch(
x,
handle=cached_handle,
num_sms=_num_comm_sms,
async_with_compute_stream=True,
)
return recv_x, cached_handle.topk_idx, None, handle, event
recv_x, recv_topk_idx, recv_topk_weights, handle, event = _buffer.dispatch(
x,
topk_idx=topk_idx,
topk_weights=topk_weights,
num_experts=num_experts,
num_max_tokens_per_rank=num_max_tokens_per_rank,
num_sms=_num_comm_sms,
async_with_compute_stream=True,
)
return recv_x, recv_topk_idx, recv_topk_weights, handle, event
def decode_combine(x: torch.Tensor,
handle: EPHandle) -> Tuple[torch.Tensor, EventOverlap]:
"""MoE combine for inference decoding."""
global _buffer, _num_comm_sms
combined_x, _, event = _buffer.combine(
x,
handle=handle,
num_sms=_num_comm_sms,
async_with_compute_stream=True,
)
return combined_x, event
```
### Environment variables
The library provides some environment variables, which may be useful:
- General
- `EP_BUFFER_DEBUG`: `0` or `1`, print buffer initialization, SM approximation, and backend debugging information, `0` by default
- `EP_SUPPRESS_NCCL_CHECK`: `0` or `1`, suppress NCCL version mismatch checking, `0` by default
- `EP_AVOID_RECORD_STREAM`: `0` or `1`, avoid `record_stream` on output tensors, `0` by default
- `EP_NUM_TOPK_IDX_BITS`: integer, override the number of bits for top-k index encoding, `0` (auto) by default
- Networking
- `EP_NIC_NAME`: string, the default NIC name used to query NIC properties, `mlx5_0` by default
- `EP_OVERRIDE_RDMA_SL`: integer, override the RDMA service level index for traffic isolation
- `EP_DISABLE_GIN`: `0` or `1`, disable the NCCL Gin backend (fall back to non-Gin path), `0` by default
- JIT
- `EP_JIT_DEBUG`: `0` or `1`, print JIT debugging information, `0` by default
- `EP_JIT_CACHE_DIR`: string, cache directory for compiled kernels, `$HOME/.deep_ep` by default
- `EP_JIT_NVCC_COMPILER`: string, NVCC compiler path; defaults to `torch.utils.cpp_extension.CUDA_HOME`
- `EP_JIT_CPP_STANDARD`: integer, C++ standard version, `20` by default
- `EP_JIT_PRINT_COMPILER_COMMAND`: `0` or `1`, print compilation commands, `0` by default
- `EP_JIT_PTXAS_VERBOSE`: `0` or `1`, show detailed PTXAS output, `0` by default
- `EP_JIT_PTXAS_CHECK`: `0` or `1`, assert no local memory usage in compiled kernels, `0` by default
- `EP_JIT_WITH_LINEINFO`: `0` or `1`, embed source line info for profiling tools, `0` by default
- `EP_JIT_DUMP_ASM`: `0` or `1`, dump both PTX and SASS, `0` by default
- `EP_JIT_DUMP_PTX`: `0` or `1`, dump PTX output, `0` by default
- `EP_JIT_DUMP_SASS`: `0` or `1`, dump SASS output, `0` by default
- Debug and profiling
- `EP_GIN_GDAKI_DEBUG`: `0` or `1`, enable NCCL Gin GDAKI debugging output, `0` by default
- `EP_USE_NVIDIA_TOOLS`: `0` or `1`, skip internal profiling when running under external NVIDIA tools, `0` by default
- `EP_DISABLE_BARRIER_PROFILING`: `0` or `1`, disable barrier-based communication profiling in benchmarks, `0` by default
- Build
- `EP_NCCL_ROOT_DIR`: string, path to the NCCL installation directory; auto-detected from the Python environment if not set
- `EP_NVSHMEM_ROOT_DIR`: string, path to the NVSHMEM installation directory; auto-detected from the Python environment if not set
- `TORCH_CUDA_ARCH_LIST`: string, list of target CUDA architectures, e.g. `"9.0"`
- `DISABLE_SM90_FEATURES`: `0` or `1`, disable SM90 features for legacy methods, `0` by default
- `DISABLE_AGGRESSIVE_PTX_INSTRS`: `0` or `1`, disable aggressive load/store instructions in legacy methods, `0` by default
Some environment variables are **persistent**: they are captured at build time and baked into the installed package as default values. At import time, these defaults are applied automatically unless overridden by current environment variables. The persistent variables are: `EP_JIT_CACHE_DIR`, `EP_JIT_PRINT_COMPILER_COMMAND`, `EP_NUM_TOPK_IDX_BITS`, `EP_NCCL_ROOT_DIR`.
For additional details, please refer to [the test code](tests/elastic/test_ep.py) or review the corresponding Python documentation.
## Network configurations
DeepEP is fully tested with InfiniBand networks. However, it is theoretically compatible with RDMA over Converged Ethernet (RoCE) as well.
### Traffic isolation
Traffic isolation is supported by InfiniBand through Virtual Lanes (VL).
To prevent interference between different types of traffic, we recommend segregating workloads across different virtual lanes as follows:
- expert-parallel workloads
- other workloads
For DeepEP V2, you can control the virtual lane assignment by setting the `sl_idx` argument or the `EP_OVERRIDE_RDMA_SL` environment variable.
### Adaptive routing
Adaptive routing is an advanced routing feature provided by InfiniBand switches that can evenly distribute traffic across multiple paths. Even though adaptive routing introduces additional latency, we still recommend enabling it under all network load conditions.
### Congestion control
Congestion control is disabled because it hurts maximum bandwidth. If congestion is unavoidable in some scenarios, we recommend assigning those workloads to low-priority virtual lanes.
### PCI atomic mode
If the hardware supports it, we recommend using the following command to set the NIC's `PCI_ATOMIC_MODE` to improve RDMA atomic operation performance:
```bash
sudo mlxconfig -y -d mlx5_$i set PCI_ATOMIC_MODE=4
```
## Experimental branches
- [Zero-copy](https://github.com/deepseek-ai/DeepEP/pull/453)
- Removing the copy between PyTorch tensors and communication buffers, which reduces the SM usages significantly for normal kernels
- This PR is authored by **Tencent Network Platform Department**
- [Eager](https://github.com/deepseek-ai/DeepEP/pull/437)
- Using a low-latency protocol removes the extra RTT latency introduced by RDMA atomic OPs
- [Hybrid-EP](https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep)
- A new backend implementation using TMA instructions for minimal SM usage and larger NVLink domain support
- Fine-grained communication-computation overlap for single-batch scenarios
- PCIe kernel support for non-NVLink environments
- NVFP4 data type support
- [AntGroup-Opt](https://github.com/deepseek-ai/DeepEP/tree/antgroup-opt)
- This optimization series is authored by **AntGroup Network Platform Department**
- [Normal-SMFree](https://github.com/deepseek-ai/DeepEP/pull/347) Eliminating SM from RDMA path by decoupling comm-kernel execution from NIC token transfer, freeing SMs for compute
- [LL-SBO](https://github.com/deepseek-ai/DeepEP/pull/483) Overlapping Down GEMM computation with Combine Send communication via signaling mechanism to reduce end-to-end latency
- [LL-Layered](https://github.com/deepseek-ai/DeepEP/pull/500) Optimizing cross-node LL operator communication using rail-optimized forwarding and data merging to reduce latency
- [Mori-EP](https://github.com/deepseek-ai/DeepEP/tree/mori-ep)
- ROCm/AMD GPU support powered by [MORI](https://github.com/ROCm/mori) backend (low-latency mode)
- [nvDev](https://github.com/deepseek-ai/DeepEP/tree/nvDev)
- V2-based branch with the latest CUDA features, such as Compute Fabric Transport (CFT) that brings better latency on small token sizes.
## Community forks
- [uccl/uccl-ep](https://github.com/uccl-project/uccl/tree/main/ep) - Enables running DeepEP on heterogeneous GPUs (e.g., Nvidia, AMD) and NICs (e.g., EFA, Broadcom, CX7)
- [Infrawaves/DeepEP_ibrc_dual-ports_multiQP](https://github.com/Infrawaves/DeepEP_ibrc_dual-ports_multiQP) - Adds multi-QP solution and dual-port NIC support in IBRC transport
- [antgroup/DeepXTrace](https://github.com/antgroup/DeepXTrace) - A diagnostic analyzer for efficient and precise localization of slow ranks
- [ROCm/mori](https://github.com/ROCm/mori) - AMD's next-generation communication library for performance-critical AI workloads (e.g., Wide EP, KVCache transfer, Collectives)
## Acknowledgement
DeepEP V2 is built on top of the [NCCL](https://github.com/nvidia/nccl) Gin backend. Thanks to @sjeaugey, @pakmarkthub, @sb17v, @xiaofanl-nvidia, and the NCCL team for their support!
## License
This code repository is released under [the MIT License](LICENSE).
## Citation
```bibtex
@misc{deepep2025,
title={DeepEP: an efficient expert-parallel communication library},
author={Chenggang Zhao and Shangyan Zhou and Liyue Zhang and Chengqi Deng and Zhean Xu and Yuxuan Liu and Kuai Yu and Jiashi Li and Liang Zhao},
year={2025},
publisher = {GitHub},
howpublished = {\url{https://github.com/deepseek-ai/DeepEP}},
}
```
README.wehub.md
+7
View File
@@ -0,0 +1,7 @@
# WeHub 来源说明
- 原始项目：`deepseek-ai/DeepEP`
- 原始仓库：https://github.com/deepseek-ai/DeepEP
- 导入方式：上游默认分支的最新快照
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
build.sh
Executable
+12
View File
@@ -0,0 +1,12 @@
# Change current directory into project root
original_dir=$(pwd)
script_dir=$(realpath "$(dirname "$0")")
cd "$script_dir"
# Remove old dist file, build files, and install
rm -rf build dist
rm -rf *.egg-info
python setup.py bdist_wheel
# Open users' original directory
cd "$original_dir"
csrc/CMakeLists.txt
+1
View File
@@ -0,0 +1 @@
add_subdirectory(kernels)
csrc/elastic/buffer.hpp
+1382
View File
File diff suppressed because it is too large Load Diff
csrc/elastic/utils.hpp
+28
View File
@@ -0,0 +1,28 @@
#pragma once
#include <torch/python.h>
#include <deep_ep/common/exception.cuh>
namespace deep_ep::elastic {
static at::cuda::CUDAStream get_global_comm_stream() {
static std::optional<at::cuda::CUDAStream> comm_stream = std::nullopt;
if (not comm_stream.has_value())
comm_stream = at::cuda::getStreamFromPool(true);
return comm_stream.value();
}
template <int kNumDims>
static auto get_shape(const torch::Tensor& t) {
EP_HOST_ASSERT(t.dim() == kNumDims);
return [&t] <size_t... Is> (std::index_sequence<Is...>) {
return std::make_tuple(static_cast<int>(t.sizes()[Is])...);
}(std::make_index_sequence<kNumDims>());
}
template <typename dtype_t = void>
static dtype_t* get_data_ptr(const std::optional<torch::Tensor>& t) {
return t.has_value() ? t->data_ptr<dtype_t>() : nullptr;
}
} // deep_ep::elastic
csrc/indexing/main.cu
+21
View File
@@ -0,0 +1,21 @@
// Utils
#include <deep_ep/impls/barrier.cuh>
// EP
#include <deep_ep/impls/dispatch.cuh>
#include <deep_ep/impls/combine.cuh>
#include <deep_ep/impls/dispatch_copy_epilogue.cuh>
#include <deep_ep/impls/combine_reduce_epilogue.cuh>
#include <deep_ep/impls/hybrid_dispatch.cuh>
#include <deep_ep/impls/hybrid_combine.cuh>
// Engram
#include <deep_ep/impls/engram_fetch.cuh>
#include <deep_ep/impls/engram_fetch_wait.cuh>
// PP
#include <deep_ep/impls/pp_send_recv.cuh>
int main() {
return 0;
}
csrc/jit/api.hpp
+20
View File
@@ -0,0 +1,20 @@
#pragma once
#include "compiler.hpp"
#include "include_parser.hpp"
#include "kernel_runtime.hpp"
namespace deep_ep::jit {
static void init(const std::string& library_root_path,
const std::string& cuda_home_path_by_python, const std::string& nccl_root_path_by_python) {
Compiler::prepare_init(library_root_path, cuda_home_path_by_python, nccl_root_path_by_python);
KernelRuntime::prepare_init(cuda_home_path_by_python);
IncludeParser::prepare_init(library_root_path);
}
static void register_apis(pybind11::module_& m) {
m.def("init_jit", &init);
}
} // namespace deep_ep::jit
csrc/jit/cache.hpp
+34
View File
@@ -0,0 +1,34 @@
#pragma once
#include <filesystem>
#include <memory>
#include <unordered_map>
#include "kernel_runtime.hpp"
namespace deep_ep::jit {
class KernelRuntimeCache {
std::unordered_map<std::string, std::shared_ptr<KernelRuntime>> cache;
public:
KernelRuntimeCache() = default;
void clear() {
cache.clear();
}
std::shared_ptr<KernelRuntime> get(const std::filesystem::path& dir_path) {
// Hit the runtime cache
if (const auto iterator = cache.find(dir_path); iterator != cache.end())
return iterator->second;
if (KernelRuntime::check_validity(dir_path))
return cache[dir_path] = std::make_shared<KernelRuntime>(dir_path);
return nullptr;
}
};
static auto kernel_runtime_cache = std::make_shared<KernelRuntimeCache>();
} // namespace deep_ep::jit
csrc/jit/compiler.hpp
+269
View File
@@ -0,0 +1,269 @@
#pragma once
#include <fcntl.h>
#include <filesystem>
#include <fstream>
#include <nvrtc.h>
#include <optional>
#include <regex>
#include <string>
#include <deep_ep/common/exception.cuh>
#include "../utils/format.hpp"
#include "../utils/hash.hpp"
#include "../utils/lazy_init.hpp"
#include "../utils/system.hpp"
#include "cache.hpp"
#include "device_runtime.hpp"
namespace deep_ep::jit {
class Compiler {
public:
static std::filesystem::path library_root_path;
static std::filesystem::path library_include_path;
static std::filesystem::path cuda_home;
static std::filesystem::path nccl_root;
static std::filesystem::path cuobjdump_path;
static void prepare_init(const std::string& library_root_path,
const std::string& cuda_home_path_by_python,
const std::string& nccl_root_path_by_python) {
// NOTES: if you are adding some third-party includes for kernels, please add its hash value
Compiler::library_root_path = library_root_path;
Compiler::library_include_path = Compiler::library_root_path / "include";
Compiler::cuda_home = cuda_home_path_by_python;
Compiler::nccl_root = nccl_root_path_by_python;
Compiler::cuobjdump_path = Compiler::cuda_home / "bin" / "cuobjdump";
}
std::string signature, flags;
std::filesystem::path cache_dir_path;
Compiler() {
EP_HOST_ASSERT(not library_root_path.empty());
EP_HOST_ASSERT(not library_include_path.empty());
EP_HOST_ASSERT(not cuda_home.empty());
EP_HOST_ASSERT(not nccl_root.empty());
EP_HOST_ASSERT(not cuobjdump_path.empty());
// Cache settings
cache_dir_path = std::filesystem::path(get_env<std::string>("HOME")) / ".deep_ep";
if (const auto env_cache_dir_path = get_env<std::string>("EP_JIT_CACHE_DIR"); not env_cache_dir_path.empty())
cache_dir_path = env_cache_dir_path;
// The compiler flags applied to all derived compilers
signature = "unknown-compiler";
flags = fmt::format("-std=c++{} --diag-suppress=39,161,174,177,186,940,3012 "
"--ptxas-options=--register-usage-level=10",
get_env<int>("EP_JIT_CPP_STANDARD", 20));
if (get_env("EP_JIT_DEBUG", 0) or get_env("EP_JIT_PTXAS_VERBOSE", 0))
flags += " --ptxas-options=--verbose";
if (get_env("EP_JIT_DEBUG", 0) or get_env("EP_JIT_WITH_LINEINFO", 0))
flags += " -Xcompiler -rdynamic -lineinfo";
if (get_env("EP_GIN_GDAKI_DEBUG", 0))
flags += " -DNCCL_DEVICE_GIN_GDAKI_ENABLE_DEBUG=1";
flags += fmt::format(" -I {}/include", nccl_root.c_str());
// Some special flags for EP
// TODO: make it more general, e.g. `EP_JIT_EXTRA_FLAGS`
if (int num_topk_idx_bits = get_env("EP_NUM_TOPK_IDX_BITS", 0); num_topk_idx_bits != 0)
flags += fmt::format(" -DEP_NUM_TOPK_IDX_BITS={}", num_topk_idx_bits);
}
virtual ~Compiler() = default;
std::filesystem::path make_tmp_dir() const {
return make_dirs(cache_dir_path / "tmp");
}
static void fsync_path(const std::filesystem::path& path) {
const auto fd = ::open(path.c_str(), O_RDONLY);
if (fd >= 0) {
::fsync(fd);
::close(fd);
}
}
// Recursively fsync a directory: files and subdirectories first (bottom-up), then the directory itself
// NOTES: ensures data and directory entries are visible on other nodes in distributed filesystems
static void fsync_dir(const std::filesystem::path& dir_path) { // NOLINT(*-no-recursion)
for (const auto& entry: std::filesystem::directory_iterator(dir_path)) {
if (entry.is_directory())
fsync_dir(entry.path());
else if (entry.is_regular_file())
fsync_path(entry.path());
}
fsync_path(dir_path);
}
static void put(const std::filesystem::path& path, const std::string& data) {
std::ofstream out(path, std::ios::binary);
EP_HOST_ASSERT(out.write(data.data(), data.size()));
out.close();
// NOTES: fsync to ensure the data is visible to other processes (e.g., NVCC)
// on distributed filesystems, where `close()` alone does not guarantee persistence
fsync_path(path);
}
std::shared_ptr<KernelRuntime> build(const std::string& name, const std::string& code) const {
const auto kernel_signature = fmt::format("{}$${}$${}$${}", name, signature, flags, code);
const auto dir_path = cache_dir_path / "cache" / fmt::format("kernel.{}.{}", name, get_hex_digest(kernel_signature));
// Hit the runtime cache
if (const auto runtime = kernel_runtime_cache->get(dir_path); runtime != nullptr)
return runtime;
// Compile into a temporary directory, then atomically rename the whole directory
// NOTES: renaming a directory is atomic on both local and distributed filesystems,
// avoiding the stale inode issue that occurs when renaming individual files
const auto tmp_dir_path = make_tmp_dir() / get_uuid();
make_dirs(tmp_dir_path);
// Compile into the temporary directory
const auto tmp_cubin_path = tmp_dir_path / "kernel.cubin";
if (get_env<int>("EP_JIT_DUMP_ASM") or get_env<int>("EP_JIT_DUMP_PTX")) {
const auto tmp_ptx_path = tmp_dir_path / "kernel.ptx";
compile(code, tmp_dir_path, tmp_cubin_path, tmp_ptx_path);
} else {
compile(code, tmp_dir_path, tmp_cubin_path);
}
// Disassemble if needed
if (get_env<int>("EP_JIT_DUMP_ASM") or get_env<int>("EP_JIT_DUMP_SASS")) {
const auto tmp_sass_path = tmp_dir_path / "kernel.sass";
disassemble(tmp_cubin_path, tmp_sass_path);
}
// Fsync before rename to ensure visibility on distributed filesystems
fsync_dir(tmp_dir_path);
// Atomically rename the temporary directory to the final cache path
// NOTES: if another rank already created dir_path, rename will fail — that's fine
make_dirs(dir_path.parent_path());
std::error_code error_code;
std::filesystem::rename(tmp_dir_path, dir_path, error_code);
if (error_code) {
// Another rank beat us, then clean up our dir and use the existing one
// NOTES: avoid `std::filesystem::remove_all` here — it can segfault on
// distributed filesystems, when concurrent processes operate
// on the same parent directory, causing stale directory entries
safe_remove_all(tmp_dir_path);
}
// Put into the runtime cache
const auto runtime = kernel_runtime_cache->get(dir_path);
EP_HOST_ASSERT(runtime != nullptr);
return runtime;
}
static void disassemble(const std::filesystem::path &cubin_path, const std::filesystem::path &sass_path) {
// Disassemble the CUBIN file to SASS
const auto command = fmt::format("{} --dump-sass {} > {}", cuobjdump_path.c_str(), cubin_path.c_str(), sass_path.c_str());
if (get_env("EP_JIT_DEBUG", 0) or get_env("EP_JIT_PRINT_COMPILER_COMMAND", 0))
printf("Running cuobjdump command: %s\n", command.c_str());
const auto [return_code, output] = call_external_command(command);
if (return_code != 0) {
printf("cuobjdump failed: %s\n", output.c_str());
EP_HOST_ASSERT(false and "cuobjdump failed");
}
}
virtual void compile(const std::string &code, const std::filesystem::path& dir_path, const std::filesystem::path &cubin_path, const std::optional<std::filesystem::path> &ptx_path = std::nullopt) const = 0;
};
EP_DECLARE_STATIC_VAR_IN_CLASS(Compiler, library_root_path);
EP_DECLARE_STATIC_VAR_IN_CLASS(Compiler, library_include_path);
EP_DECLARE_STATIC_VAR_IN_CLASS(Compiler, cuda_home);
EP_DECLARE_STATIC_VAR_IN_CLASS(Compiler, nccl_root);
EP_DECLARE_STATIC_VAR_IN_CLASS(Compiler, cuobjdump_path);
class NVCCCompiler final: public Compiler {
std::filesystem::path nvcc_path;
std::pair<int, int> get_nvcc_version() const {
EP_HOST_ASSERT(std::filesystem::exists(nvcc_path));
// Call the version command
const auto command = std::string(nvcc_path) + " --version";
const auto [return_code, output] = call_external_command(command);
EP_HOST_ASSERT(return_code == 0);
// The version should be at least 12.3
int major, minor;
std::smatch match;
EP_HOST_ASSERT(std::regex_search(output, match, std::regex(R"(release (\d+\.\d+))")));
std::sscanf(match[1].str().c_str(), "%d.%d", &major, &minor);
EP_HOST_ASSERT((major > 12 or (major == 12 and minor >= 3)) and "NVCC version should be >= 12.3");
return {major, minor};
}
public:
NVCCCompiler() {
// Override the compiler signature
nvcc_path = cuda_home / "bin" / "nvcc";
cuobjdump_path = cuda_home / "bin" / "cuobjdump";
if (const auto env_nvcc_path = get_env<std::string>("EP_JIT_NVCC_COMPILER"); not env_nvcc_path.empty())
nvcc_path = env_nvcc_path;
const auto [nvcc_major, nvcc_minor] = get_nvcc_version();
signature = fmt::format("NVCC{}.{}", nvcc_major, nvcc_minor);
// The override the compiler flags
// Only NVCC >= 12.9 supports arch-specific family suffix
const auto arch = device_runtime->get_arch(false, nvcc_major > 12 or nvcc_minor >= 9);
flags = fmt::format("{} -I{} --gpu-architecture=sm_{} "
"--compiler-options=-fPIC,-O3,-fconcepts,-Wno-deprecated-declarations,-Wno-abi "
"-O3 --expt-relaxed-constexpr --expt-extended-lambda",
flags, library_include_path.c_str(), arch);
}
void compile(const std::string &code, const std::filesystem::path& dir_path,
const std::filesystem::path &cubin_path,
const std::optional<std::filesystem::path> &ptx_path) const override {
// Write the code into the cache directory
const auto code_path = dir_path / "kernel.cu";
put(code_path, code);
// Compile to CUBIN
// Avoid cwd files shadowing C++ standard library headers
const auto compile_dir = make_tmp_dir();
const auto command = fmt::format("cd {} && {} {} -cubin -o {} {}",
compile_dir.c_str(), nvcc_path.c_str(), code_path.c_str(), cubin_path.c_str(), flags);
if (get_env("EP_JIT_DEBUG", 0) or get_env("EP_JIT_PRINT_COMPILER_COMMAND", 0))
printf("Running NVCC command: %s\n", command.c_str());
const auto [return_code, output] = call_external_command(command);
if (return_code != 0) {
printf("NVCC compilation failed: %s\n", output.c_str());
EP_HOST_ASSERT(false and "NVCC compilation failed");
}
// Compile to PTX if needed
if (ptx_path.has_value()) {
const auto ptx_command = fmt::format("cd {} && {} {} -ptx -o {} {}",
compile_dir.c_str(), nvcc_path.c_str(), code_path.c_str(), ptx_path->c_str(), flags);
if (get_env("EP_JIT_DEBUG", 0) or get_env("EP_JIT_PRINT_COMPILER_COMMAND", 0))
printf("Running NVCC PTX command: %s\n", ptx_command.c_str());
const auto [ptx_return_code, ptx_output] = call_external_command(ptx_command);
if (ptx_return_code != 0) {
printf("NVCC PTX compilation failed: %s\n", ptx_output.c_str());
EP_HOST_ASSERT(false and "NVCC PTX compilation failed");
}
}
// Check local memory usage
if (get_env("EP_JIT_PTXAS_CHECK", 0))
EP_HOST_ASSERT(not std::regex_search(output, std::regex(R"(Local memory used)")));
// Print PTXAS log
if (get_env("EP_JIT_DEBUG", 0) or get_env("EP_JIT_PTXAS_VERBOSE", 0))
printf("%s", output.c_str());
}
};
static auto compiler = LazyInit<Compiler>([]() -> std::shared_ptr<Compiler> {
return std::make_shared<NVCCCompiler>();
});
} // namespace deep_ep::jit
csrc/jit/device_runtime.hpp
+67
View File
@@ -0,0 +1,67 @@
#pragma once
#include <deep_ep/common/exception.cuh>
#include "../utils/lazy_init.hpp"
namespace deep_ep::jit {
class DeviceRuntime {
int64_t cached_clock_rate = 0;
std::shared_ptr<cudaDeviceProp> cached_prop;
std::shared_ptr<cudaDeviceProp> get_prop() {
if (cached_prop == nullptr) {
int device_idx;
cudaDeviceProp prop;
CUDA_RUNTIME_CHECK(cudaGetDevice(&device_idx));
CUDA_RUNTIME_CHECK(cudaGetDeviceProperties(&prop, device_idx));
cached_prop = std::make_shared<cudaDeviceProp>(prop);
}
return cached_prop;
}
public:
int64_t get_clock_rate() {
if (cached_clock_rate == 0) {
// NOTES: we should convert kHz into Hz
int device_idx, rate;
CUDA_RUNTIME_CHECK(cudaGetDevice(&device_idx));
CUDA_RUNTIME_CHECK(cudaDeviceGetAttribute(&rate, cudaDevAttrClockRate, device_idx));
cached_clock_rate = static_cast<int64_t>(rate) * 1000ll;
}
return cached_clock_rate;
}
int get_num_smem_bytes() {
return static_cast<int>(get_prop()->sharedMemPerBlockOptin);
}
int get_num_sms() {
return get_prop()->multiProcessorCount;
}
std::pair<int, int> get_arch_pair() {
const auto prop = get_prop();
return {prop->major, prop->minor};
}
std::string get_arch(const bool& number_only = false,
const bool& support_arch_family = false) {
const auto [major, minor] = get_arch_pair();
if (major == 10 and minor != 1) {
if (number_only)
return "100";
return support_arch_family ? "100f" : "100a";
}
return std::to_string(major * 10 + minor) + (number_only ? "" : "a");
}
int get_arch_major() {
return get_arch_pair().first;
}
};
static auto device_runtime = LazyInit<DeviceRuntime>([](){ return std::make_shared<DeviceRuntime>(); });
} // namespace deep_ep::jit
csrc/jit/handle.hpp
+155
View File
@@ -0,0 +1,155 @@
#pragma once
#include <cuda.h>
#include <filesystem>
#include <deep_ep/common/exception.cuh>
#include "../utils/lazy_driver.hpp"
namespace deep_ep::jit {
#if CUDART_VERSION >= 12080 and defined(EP_JIT_USE_RUNTIME_API)
// Use CUDA runtime API
using LibraryHandle = cudaLibrary_t;
using KernelHandle = cudaKernel_t;
using LaunchConfigHandle = cudaLaunchConfig_t;
using LaunchAttrHandle = cudaLaunchAttribute;
#define EP_CUDA_UNIFIED_CHECK CUDA_RUNTIME_CHECK
static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const std::string& func_name,
LibraryHandle *library_opt = nullptr) {
LibraryHandle library;
KernelHandle kernel{};
CUDA_RUNTIME_CHECK(cudaLibraryLoadFromFile(&library, cubin_path.c_str(), nullptr, nullptr, 0, nullptr, nullptr, 0));
nvshmemx_culibrary_init(library);
CUDA_RUNTIME_CHECK(cudaLibraryGetKernel(&kernel, library, func_name.c_str()));
if (library_opt != nullptr)
*library_opt = library;
return kernel;
}
static void unload_library(const LibraryHandle& library) {
nvshmemx_culibrary_finalize(library);
const auto error = cudaLibraryUnload(library);
EP_HOST_ASSERT(error == cudaSuccess or error == cudaErrorCudartUnloading);
}
static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
const cudaStream_t& stream, const int& smem_size,
const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim) {
if (smem_size > 0)
CUDA_RUNTIME_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
LaunchConfigHandle config;
config.gridDim = grid_dim;
config.blockDim = block_dim;
config.dynamicSmemBytes = smem_size;
config.stream = stream;
config.numAttrs = 0;
config.attrs = nullptr;
// TODO: support cooperative and dependent kernel launch
// NOTES: must use `static` or the `attr` will be deconstructed
static LaunchAttrHandle attr;
if (cluster_dim > 1) {
attr.id = cudaLaunchAttributeClusterDimension;
attr.val.clusterDim = {static_cast<unsigned>(cluster_dim), 1, 1};
config.attrs = &attr;
config.numAttrs = 1;
}
return config;
}
template<typename... ActTypes>
static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) {
void *ptr_args[] = { &args... };
return cudaLaunchKernelExC(&config, kernel, ptr_args);
}
#else
// Use CUDA driver API
using LibraryHandle = CUmodule;
using KernelHandle = CUfunction;
using LaunchConfigHandle = CUlaunchConfig;
using LaunchAttrHandle = CUlaunchAttribute;
#define EP_CUDA_UNIFIED_CHECK CUDA_DRIVER_CHECK
static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const std::string& func_name,
LibraryHandle *library_opt = nullptr) {
LibraryHandle library;
KernelHandle kernel;
CUDA_DRIVER_CHECK(lazy_cuModuleLoad(&library, cubin_path.c_str()));
CUDA_DRIVER_CHECK(lazy_cuModuleGetFunction(&kernel, library, func_name.c_str()));
if (library_opt != nullptr)
*library_opt = library;
return kernel;
}
static void unload_library(const LibraryHandle& library) {
const auto error = lazy_cuModuleUnload(library);
EP_HOST_ASSERT(error == CUDA_SUCCESS or error == CUDA_ERROR_DEINITIALIZED);
}
static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
const cudaStream_t& stream, const int& smem_size,
const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim,
const bool& cooperative, const bool& enable_pdl) {
if (smem_size > 0)
CUDA_DRIVER_CHECK(lazy_cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size));
LaunchConfigHandle config;
config.gridDimX = grid_dim.x;
config.gridDimY = grid_dim.y;
config.gridDimZ = grid_dim.z;
config.blockDimX = block_dim.x;
config.blockDimY = block_dim.y;
config.blockDimZ = block_dim.z;
config.sharedMemBytes = smem_size;
config.hStream = stream;
// Create attributes
static LaunchAttrHandle attrs[3];
config.attrs = attrs;
config.numAttrs = 0;
// Cooperative launch
if (cooperative) {
auto& attr = attrs[config.numAttrs ++];
attr.id = CU_LAUNCH_ATTRIBUTE_COOPERATIVE;
attr.value.cooperative = 1;
}
// Cluster size
// NOTES: must use `static` or the `attr` will be deconstructed
if (cluster_dim > 1) {
auto& attr = attrs[config.numAttrs ++];
attr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
attr.value.clusterDim.x = cluster_dim;
attr.value.clusterDim.y = 1;
attr.value.clusterDim.z = 1;
}
// Dependent kernel launch
if (enable_pdl) {
auto& attr = attrs[config.numAttrs ++];
attr.id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION;
attr.value.programmaticStreamSerializationAllowed = 1;
}
return config;
}
template<typename... ActTypes>
static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) {
void *ptr_args[] = { &args... };
return lazy_cuLaunchKernelEx(&config, kernel, ptr_args, nullptr);
}
#endif
} // namespace deep_ep::jit
csrc/jit/include_parser.hpp
+80
View File
@@ -0,0 +1,80 @@
#pragma once
#include <filesystem>
#include <regex>
#include <string>
#include <vector>
#include "../utils/format.hpp"
#include "../utils/system.hpp"
namespace deep_ep::jit {
class IncludeParser {
std::unordered_map<std::string, std::optional<std::string>> cache;
static std::vector<std::string> get_includes(const std::string& code, const std::filesystem::path& file_path = "") {
std::vector<std::string> includes;
const std::regex pattern(R"(#\s*include\s*[<"][^>"]+[>"])");
std::sregex_iterator iter(code.begin(), code.end(), pattern);
const std::sregex_iterator end;
// TODO: parse relative paths as well
for (; iter != end; ++ iter) {
const auto include_str = iter->str();
const int len = include_str.length();
if (include_str.substr(0, 10) == "#include <" and include_str[len - 1] == '>' and include_str[10] != ' ' and include_str[len - 2] != ' ') {
std::string filename = include_str.substr(10, len - 11);
if (filename.substr(0, 7) == "deep_ep") // We only parse `<deep_ep/*>`
includes.push_back(filename);
} else {
std::string error_info = fmt::format("Non-standard include: {}", include_str);
if (file_path != "")
error_info += fmt::format(" ({})", file_path.string());
EP_HOST_UNREACHABLE(error_info);
}
}
return includes;
}
public:
static std::filesystem::path library_include_path;
static void prepare_init(const std::string& library_root_path) {
library_include_path = std::filesystem::path(library_root_path) / "include";
}
std::string get_hash_value(const std::string& code, const bool& exclude_code = true) {
std::stringstream ss;
for (const auto& i: get_includes(code))
ss << get_hash_value_by_path(library_include_path / i) << "$";
if (not exclude_code)
ss << "#" << get_hex_digest(code);
return get_hex_digest(ss.str());
}
std::string get_hash_value_by_path(const std::filesystem::path& path) {
// Check whether hit in cache
// ReSharper disable once CppUseAssociativeContains
if (cache.count(path) > 0) {
const auto opt = cache[path];
if (not opt.has_value())
EP_HOST_UNREACHABLE(fmt::format("Circular include may occur: {}", path.string()));
return opt.value();
}
// Read file and calculate hash recursively
std::ifstream in(path);
if (not in.is_open())
EP_HOST_UNREACHABLE(fmt::format("Failed to open: {}", path.string()));
std::string code((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>());
cache[path] = std::nullopt;
return (cache[path] = get_hash_value(code, false)).value();
}
};
EP_DECLARE_STATIC_VAR_IN_CLASS(IncludeParser, library_include_path);
static auto include_parser = std::make_shared<IncludeParser>();
} // namespace deep_ep::jit
csrc/jit/kernel_runtime.hpp
+87
View File
@@ -0,0 +1,87 @@
#pragma once
#include <ATen/cuda/CUDAContext.h>
#include <deep_ep/common/exception.cuh>
#include "../utils/format.hpp"
#include "../utils/lazy_init.hpp"
#include "handle.hpp"
namespace deep_ep::jit {
class KernelRuntime final {
public:
static std::filesystem::path cuda_home;
LibraryHandle library;
KernelHandle kernel;
explicit KernelRuntime(const std::filesystem::path& dir_path) {
EP_HOST_ASSERT(not cuda_home.empty());
// NOLINT(*-pro-type-member-init)
const auto cuobjdump_path = cuda_home / "bin" / "cuobjdump";
const auto cubin_path = dir_path / "kernel.cubin";
if (get_env<int>("EP_JIT_DEBUG"))
printf("Loading CUBIN: %s\n", cubin_path.c_str());
// Find the only symbol
// TODO: use kernel enumeration for newer drivers
const std::vector<std::string> illegal_names = {"vprintf", "__instantiate_kernel", "__internal", "__assertfail"};
const auto [exit_code, symbols] = call_external_command(fmt::format("{} -symbols {}", cuobjdump_path.c_str(), cubin_path.c_str()));
EP_HOST_ASSERT(exit_code == 0);
std::istringstream iss(symbols);
std::vector<std::string> symbol_names;
for (std::string line; std::getline(iss, line); ) {
if (line.find("STT_FUNC") == 0 and line.find("STO_ENTRY") != std::string::npos and
std::none_of(illegal_names.begin(), illegal_names.end(),
[&](const auto name) { return line.find(name) != std::string::npos; })) {
const auto last_space = line.rfind(' ');
symbol_names.push_back(line.substr(last_space + 1));
}
}
// Print symbols
if (symbol_names.size() != 1) {
printf("Corrupted JIT cache directory (expected 1 kernel symbol, found %zu): %s, "
"please run `rm -rf %s` and restart your task.\n",
symbol_names.size(), dir_path.c_str(), dir_path.c_str());
printf("Symbol names: ");
for (const auto& symbol: symbol_names)
printf("%s, ", symbol.c_str());
printf("\n");
EP_HOST_ASSERT(false and "Corrupted JIT cache directory");
}
// Load from the library
kernel = load_kernel(cubin_path, symbol_names[0], &library);
}
static void prepare_init(const std::string& cuda_home_path_by_python) {
cuda_home = cuda_home_path_by_python;
}
static bool check_validity(const std::filesystem::path& dir_path) {
if (not std::filesystem::exists(dir_path))
return false;
// NOTES: if the directory exists, kernel.cu and kernel.cubin must both exist,
// because the directory is created atomically via rename
if (not std::filesystem::exists(dir_path / "kernel.cu") or
not std::filesystem::exists(dir_path / "kernel.cubin")) {
printf("Corrupted JIT cache directory (missing kernel.cu or kernel.cubin): %s, "
"please run `rm -rf %s` and restart your task.\n",
dir_path.c_str(), dir_path.c_str());
EP_HOST_ASSERT(false and "Corrupted JIT cache directory");
}
return true;
}
~KernelRuntime() noexcept(false) {
unload_library(library);
}
};
EP_DECLARE_STATIC_VAR_IN_CLASS(KernelRuntime, cuda_home);
} // namespace deep_ep::jit
csrc/jit/launch_runtime.hpp
+74
View File
@@ -0,0 +1,74 @@
#pragma once
#include <ATen/cuda/CUDAContext.h>
#include <deep_ep/common/exception.cuh>
#include "../utils/format.hpp"
#include "../utils/system.hpp"
#include "compiler.hpp"
#include "include_parser.hpp"
namespace deep_ep::jit {
struct LaunchArgs {
std::pair<int, int> grid_dim;
int num_threads;
int smem_size;
int cluster_dim;
bool cooperative;
bool pdl_enabled;
LaunchArgs(const int& grid_dim_x, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1, const bool& cooperative = false, const bool& pdl_enabled = false):
grid_dim({grid_dim_x, 1}), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim), cooperative(cooperative), pdl_enabled(pdl_enabled) {}
LaunchArgs(const std::pair<int, int>& grid_dim, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1, const bool& cooperative = false, const bool& pdl_enabled = false):
grid_dim(grid_dim), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim), cooperative(cooperative), pdl_enabled(pdl_enabled) {}
};
template <typename Derived>
class LaunchRuntime {
public:
template <typename Args>
static std::string generate(const Args& args) {
auto code = Derived::generate_impl(args);
// NOTES: we require that `generate_impl`'s includes never change
static std::string include_hash;
if (include_hash.empty())
include_hash = include_parser->get_hash_value(code);
// TODO: optimize string concat performance
code = fmt::format("// Includes' hash value: {}\n{}", include_hash, code);
if (get_env<int>("EP_JIT_DEBUG", 0))
printf("Generated kernel code:\n%s\n", code.c_str());
return code;
}
template <typename Args>
static void launch(const std::shared_ptr<KernelRuntime>& kernel_runtime, const Args& args,
const std::optional<at::cuda::CUDAStream>& stream_opt = std::nullopt) {
const auto kernel = kernel_runtime->kernel;
const auto stream = stream_opt.value_or(at::cuda::getCurrentCUDAStream());
const LaunchArgs& launch_args = args.launch_args;
const dim3& grid_dim = {static_cast<unsigned>(launch_args.grid_dim.first),
static_cast<unsigned>(launch_args.grid_dim.second),
1};
const dim3& block_dim = {static_cast<unsigned>(launch_args.num_threads), 1, 1};
auto config = construct_launch_config(kernel, stream, launch_args.smem_size,
grid_dim, block_dim, launch_args.cluster_dim,
launch_args.cooperative, launch_args.pdl_enabled);
// Launch in the derived class
if (get_env<int>("EP_JIT_DEBUG")) {
printf("Launch kernel with {%d, %d} x %d (cooperative: %d), shared memory: %d bytes, cluster: %d, stream: %ld\n",
launch_args.grid_dim.first, launch_args.grid_dim.second, launch_args.num_threads,
launch_args.cooperative,
launch_args.smem_size, launch_args.cluster_dim, stream.id());
}
Derived::launch_impl(kernel, config, args);
}
};
} // namespace deep_ep::jit
csrc/kernels/CMakeLists.txt
+15
View File
@@ -0,0 +1,15 @@
function(add_deep_ep_library target_name source_file)
add_library(${target_name} STATIC ${source_file})
set_target_properties(${target_name} PROPERTIES
POSITION_INDEPENDENT_CODE ON
CXX_STANDARD_REQUIRED ON
CUDA_STANDARD_REQUIRED ON
CXX_STANDARD 17
CUDA_STANDARD 17
CUDA_SEPARABLE_COMPILATION ON
)
target_link_libraries(${target_name} PUBLIC nvshmem_host nvshmem_device cudart cudadevrt mlx5)
endfunction()
add_subdirectory(legacy)
add_subdirectory(backend)
csrc/kernels/backend/CMakeLists.txt
+6
View File
@@ -0,0 +1,6 @@
add_deep_ep_library(runtime_nccl_cuda nccl.cu)
add_deep_ep_library(runtime_nvshmem_cuda nvshmem.cu)
add_deep_ep_library(runtime_cuda_driver cuda_driver.cu)
# Link these libraries later
set(RUNTIME_CUDA_LIBRARIES runtime_nccl_cuda runtime_nvshmem_cuda runtime_cuda_driver CACHE INTERNAL "Runtime kernels")
csrc/kernels/backend/api.cuh
+104
View File
@@ -0,0 +1,104 @@
#pragma once
#include <memory>
#include <vector>
#include <optional>
#include <nccl.h>
#include <nccl_device.h>
#include "symmetric.hpp"
// TODO: make a unified API
namespace deep_ep::nvshmem {
std::vector<uint8_t> get_unique_id();
int init(const std::vector<uint8_t>& root_unique_id_val,
const int& rank,
const int& num_ranks,
const int& team_split_stride);
void* alloc(const size_t& size, const size_t& alignment);
void free(void* ptr);
void barrier(const bool& with_cpu_sync, const std::optional<cudaStream_t>& stream_opt = std::nullopt);
void finalize();
} // deep_ep::nvshmem
namespace deep_ep::nccl {
pybind11::bytearray get_local_unique_id();
int64_t create_nccl_comm(const pybind11::bytearray& root_unique_id_bytes,
const int& num_ranks, const int& rank_idx);
void destroy_nccl_comm(const int64_t& nccl_comm);
std::tuple<int, int> get_physical_domain_size(const int64_t& nccl_comm);
std::tuple<int, int> get_logical_domain_size(const int64_t& nccl_comm, const bool& allow_hybrid_mode);
// TODO: make it header only?
struct NCCLSymmetricMemoryContext {
private:
// Can not use this unmapped pointer from outside
void* raw_window_ptr;
std::shared_ptr<symmetric::SymmetricMemory> symmetric_memory;
public:
// Global
int rank_idx;
int num_ranks;
// Logical
int num_scaleout_ranks, num_scaleup_ranks;
int scaleout_rank_idx, scaleup_rank_idx;
// Physical
int num_rdma_ranks, num_nvl_ranks;
int rdma_rank_idx, nvl_rank_idx;
bool is_scaleup_nvlink;
// NCCL handles
ncclComm_t comm;
ncclDevComm_t dev_comm;
ncclWindow_t window;
void* mapped_window_ptr;
std::vector<void*> nvl_window_ptrs;
// Configs
int num_allocated_qps;
// Buffer size
int64_t num_gpu_bytes;
int64_t num_cpu_bytes;
NCCLSymmetricMemoryContext(const int64_t& nccl_comm, const symmetric::cpu_comm_t& cpu_comm,
const int& num_ranks, const int& rank_idx,
const int64_t& num_bytes, const int64_t& num_cpu_bytes,
const bool& allow_hybrid_mode,
const int& sl_idx, const int& num_allocated_qps);
// TODO: finish this with `explicit_destroy`
// ~NCCLSymmetricMemoryContext();
void* get_sym_ptr(void* ptr, const int& dst_rank_idx) const;
void finalize();
};
} // deep_ep::nccl
namespace deep_ep::cuda_driver {
void batched_write(CUstream stream, const std::vector<void*>& ptrs, const int& value);
void batched_wait(CUstream stream, const std::vector<void*>& ptrs, const int& value);
void batched_write_and_wait(CUstream stream, const std::vector<void*>& write_ptrs, const std::vector<void*>& wait_ptrs, const int& value);
} // namespace deep_ep::cuda_driver
csrc/kernels/backend/cuda_driver.cu
+54
View File
@@ -0,0 +1,54 @@
#include <cuda.h>
#include <pybind11/pytypes.h>
#include <pybind11/stl.h>
#include <deep_ep/common/exception.cuh>
#include "api.cuh"
#include "../../utils/lazy_driver.hpp"
namespace deep_ep::cuda_driver {
static CUstreamBatchMemOpParams create_mem_op(
void *ptr, const int& value,
const CUstreamBatchMemOpType& type,
const CUstreamWaitValue_flags& wait_flag = CU_STREAM_WAIT_VALUE_EQ) {
CUstreamBatchMemOpParams params;
if (type == CU_STREAM_MEM_OP_WRITE_VALUE_32) {
params.operation = CU_STREAM_MEM_OP_WRITE_VALUE_32;
params.writeValue.address = reinterpret_cast<CUdeviceptr>(ptr);
params.writeValue.value = value;
params.writeValue.flags = 0;
} else {
params.operation = CU_STREAM_MEM_OP_WAIT_VALUE_32;
params.waitValue.address = reinterpret_cast<CUdeviceptr>(ptr);
params.waitValue.value = value;
params.waitValue.flags = wait_flag;
}
return params;
}
void batched_write(CUstream stream, const std::vector<void*>& ptrs, const int& value) {
std::vector<CUstreamBatchMemOpParams> ops(ptrs.size());
for (int i = 0; i < ptrs.size(); ++ i)
ops[i] = create_mem_op(ptrs[i], value, CU_STREAM_MEM_OP_WRITE_VALUE_32);
CUDA_DRIVER_CHECK(lazy_cuStreamBatchMemOp(stream, ops.size(), ops.data(), 0));
}
void batched_wait(CUstream stream, const std::vector<void*>& ptrs, const int& value) {
std::vector<CUstreamBatchMemOpParams> ops(ptrs.size());
for (int i = 0; i < ptrs.size(); ++ i)
ops[i] = create_mem_op(ptrs[i], value, CU_STREAM_MEM_OP_WAIT_VALUE_32, CU_STREAM_WAIT_VALUE_GEQ);
CUDA_DRIVER_CHECK(lazy_cuStreamBatchMemOp(stream, ops.size(), ops.data(), 0));
}
void batched_write_and_wait(CUstream stream, const std::vector<void*>& write_ptrs, const std::vector<void*>& wait_ptrs, const int& value) {
std::vector<CUstreamBatchMemOpParams> ops(write_ptrs.size() + wait_ptrs.size());
for (int i = 0; i < write_ptrs.size(); ++ i)
ops[i] = create_mem_op(write_ptrs[i], value, CU_STREAM_MEM_OP_WRITE_VALUE_32);
for (int i = 0; i < wait_ptrs.size(); ++ i)
ops[write_ptrs.size() + i] = create_mem_op(wait_ptrs[i], value, CU_STREAM_MEM_OP_WAIT_VALUE_32, CU_STREAM_WAIT_VALUE_GEQ);
CUDA_DRIVER_CHECK(lazy_cuStreamBatchMemOp(stream, ops.size(), ops.data(), 0));
}
} // namespace deep_ep::cuda_driver
csrc/kernels/backend/nccl.cu
+156
View File
@@ -0,0 +1,156 @@
#include <cstring>
#include <vector>
#include <string>
#include <pybind11/pytypes.h>
#include <pybind11/stl.h>
#include <sstream>
#include <nccl.h>
#include <nccl_device/core.h>
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/exception.cuh>
#include "api.cuh"
#include "../../utils/system.hpp"
namespace deep_ep::nccl {
pybind11::bytearray get_local_unique_id() {
ncclUniqueId unique_id;
NCCL_CHECK(ncclGetUniqueId(&unique_id));
std::vector<char> result(sizeof(ncclUniqueId));
std::memcpy(result.data(), &unique_id, sizeof(ncclUniqueId));
return {result.data(), result.size()};
}
int64_t create_nccl_comm(const pybind11::bytearray& root_unique_id_bytes,
const int& num_ranks, const int& rank_idx) {
// Copy unique ID
ncclUniqueId root_unique_id;
const auto root_unique_id_str = root_unique_id_bytes.cast<std::string>();
std::memcpy(&root_unique_id, root_unique_id_str.c_str(), sizeof(ncclUniqueId));
// Init
ncclComm_t comm;
NCCL_CHECK(ncclCommInitRank(&comm, num_ranks, root_unique_id, rank_idx));
if (get_env<int>("EP_BUFFER_DEBUG"))
printf("New NCCL host communicator created (%d/%d)\n", rank_idx, num_ranks);
return reinterpret_cast<int64_t>(comm);
}
void destroy_nccl_comm(const int64_t& nccl_comm) {
NCCL_CHECK(ncclCommAbort(reinterpret_cast<ncclComm_t>(nccl_comm)));
if (get_env<int>("EP_BUFFER_DEBUG"))
printf("NCCL host communicator aborted\n");
}
std::tuple<int, int> get_physical_domain_size(const int64_t& nccl_comm) {
const auto comm = reinterpret_cast<ncclComm_t>(nccl_comm);
const int num_ranks = ncclTeamWorld(comm).nRanks, num_nvl_ranks = ncclTeamLsa(comm).nRanks;
EP_HOST_ASSERT(num_ranks % num_nvl_ranks == 0);
return {num_ranks / num_nvl_ranks, num_nvl_ranks};
}
std::tuple<int, int> get_logical_domain_size(const int64_t& nccl_comm, const bool& allow_hybrid_mode) {
const auto [num_rdma_ranks, num_nvl_ranks] = get_physical_domain_size(nccl_comm);
return {allow_hybrid_mode ? num_rdma_ranks : 1,
allow_hybrid_mode ? num_nvl_ranks : num_rdma_ranks * num_nvl_ranks};
}
NCCLSymmetricMemoryContext::NCCLSymmetricMemoryContext(const int64_t& nccl_comm, const symmetric::cpu_comm_t& cpu_comm,
const int& num_ranks, const int& rank_idx,
const int64_t& num_bytes, const int64_t& num_cpu_bytes,
const bool& allow_hybrid_mode,
const int& sl_idx, const int& num_allocated_qps):
rank_idx(rank_idx), num_ranks(num_ranks), num_allocated_qps(num_allocated_qps) {
if (get_env("EP_BUFFER_DEBUG", 0)) {
int nccl_version;
NCCL_CHECK(ncclGetVersion(&nccl_version));
printf("DeepEP initialized with NCCL version: %d.%d.%d (loaded library)\n",
nccl_version / 10000, (nccl_version % 10000) / 100, nccl_version % 100);
}
// Reuse the NCCL communicator
comm = reinterpret_cast<ncclComm_t>(nccl_comm);
// Print number of allocated QPs
if (get_env<int>("EP_BUFFER_DEBUG"))
printf("EP NCCL device communicator has %d allocated QPs\n", num_allocated_qps);
// Initialize NCCL device communicator
ncclDevCommRequirements_t reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER;
if (num_ranks > 1 and get_env("EP_DISABLE_GIN", 0) == 0) {
// Query NCCL supported Gin Type
ncclCommProperties props = NCCL_COMM_PROPERTIES_INITIALIZER;
NCCL_CHECK(ncclCommQueryProperties(comm, &props));
EP_HOST_ASSERT(
(allow_hybrid_mode ? props.railedGinType : props.ginType) != NCCL_GIN_TYPE_NONE and
"NCCL GIN is unavailable. This is usually due to a network configuration issue, "
"such as `allow_hybrid_mode=0` (disable direct RDMA kernels) in multi-plane network.");
reqs.ginContextCount = num_allocated_qps;
reqs.ginExclusiveContexts = true;
reqs.ginQueueDepth = kGinQPDepth;
reqs.ginTrafficClass = sl_idx;
// Customized RDMA barrier needs extra signals
reqs.ginSignalCount = num_ranks + 2 * 2;
reqs.ginConnectionType = allow_hybrid_mode ? NCCL_GIN_CONNECTION_RAIL: NCCL_GIN_CONNECTION_FULL;
}
NCCL_CHECK(ncclDevCommCreate(comm, &reqs, &dev_comm));
// Now we know the NVLink domain size
num_nvl_ranks = dev_comm.lsaSize, nvl_rank_idx = dev_comm.lsaRank;
num_rdma_ranks = num_ranks / num_nvl_ranks, rdma_rank_idx = rank_idx / num_nvl_ranks;
EP_HOST_ASSERT(num_ranks % num_nvl_ranks == 0 and nvl_rank_idx == rank_idx % num_nvl_ranks);
EP_HOST_ASSERT(rank_idx == rdma_rank_idx * num_nvl_ranks + nvl_rank_idx);
// Calculate scaleout/up domain size
if (allow_hybrid_mode) {
num_scaleout_ranks = num_rdma_ranks, num_scaleup_ranks = num_nvl_ranks;
scaleout_rank_idx = rdma_rank_idx, scaleup_rank_idx = nvl_rank_idx;
} else {
num_scaleout_ranks = 1, num_scaleup_ranks = num_ranks;
scaleout_rank_idx = 0, scaleup_rank_idx = rank_idx;
}
is_scaleup_nvlink = num_scaleup_ranks == num_nvl_ranks;
// Create symmetric memory
// num_bytes = GPU + CPU, derive GPU portion
this->symmetric_memory = symmetric::alloc(
num_bytes - num_cpu_bytes, num_cpu_bytes,
allow_hybrid_mode, num_scaleup_ranks, scaleout_rank_idx,
cpu_comm);
// Create window
// NOTES: `ncclCommWindowRegister` is collective: it internally calls bootstrapBarrier
// across all ranks, so no explicit barrier is needed after this call.
raw_window_ptr = this->symmetric_memory->ptr;
this->num_gpu_bytes = this->symmetric_memory->num_gpu_bytes;
this->num_cpu_bytes = this->symmetric_memory->num_cpu_bytes;
NCCL_CHECK(ncclCommWindowRegister(comm, raw_window_ptr, this->symmetric_memory->num_bytes, &window, NCCL_WIN_STRICT_ORDERING));
NCCL_CHECK(ncclGetLsaDevicePointer(window, 0, nvl_rank_idx, &mapped_window_ptr));
// Get LSA pointers for all LSA peers
// TODO: check whether this is correct for network with RDMA
nvl_window_ptrs.resize(num_nvl_ranks);
for (int i = 0; i < num_nvl_ranks; ++ i)
NCCL_CHECK(ncclGetLsaDevicePointer(window, 0, i, &nvl_window_ptrs[i]));
}
void* NCCLSymmetricMemoryContext::get_sym_ptr(void* ptr, const int& dst_rank_idx) const {
const auto offset = static_cast<uint8_t*>(ptr) - static_cast<uint8_t*>(mapped_window_ptr);
return static_cast<uint8_t*>(nvl_window_ptrs[dst_rank_idx]) + offset;
}
void NCCLSymmetricMemoryContext::finalize() {
// Deregister window
NCCL_CHECK(ncclCommWindowDeregister(comm, window));
symmetric_memory.reset();
// Destroy device communicator
NCCL_CHECK(ncclDevCommDestroy(comm, &dev_comm));
}
} // namespace deep_ep::nccl
csrc/kernels/backend/nvshmem.cu
+87
View File
@@ -0,0 +1,87 @@
#include <cstring>
#include <optional>
#include <vector>
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/exception.cuh>
#include <nvshmem.h>
namespace deep_ep::nvshmem {
nvshmem_team_t cpu_rdma_team = NVSHMEM_TEAM_INVALID;
nvshmem_team_config_t cpu_rdma_team_config;
std::vector<uint8_t> get_unique_id() {
nvshmemx_uniqueid_t unique_id;
nvshmemx_get_uniqueid(&unique_id);
std::vector<uint8_t> result(sizeof(nvshmemx_uniqueid_t));
std::memcpy(result.data(), &unique_id, sizeof(nvshmemx_uniqueid_t));
return result;
}
void* alloc(const size_t& size, const size_t& alignment) {
return nvshmem_align(alignment, size);
}
void free(void* ptr) {
nvshmem_free(ptr);
}
void barrier(const bool& with_cpu_sync,
const std::optional<cudaStream_t>& stream_opt = std::nullopt) {
// Wait all streams to finish on this GPU
if (with_cpu_sync)
CUDA_RUNTIME_CHECK(cudaDeviceSynchronize());
// NOTES: this only launches kernels at GPU
if (stream_opt.has_value()) {
nvshmemx_barrier_all_on_stream(stream_opt.value());
} else {
nvshmem_barrier_all();
}
// Let CPU wait
if (with_cpu_sync)
CUDA_RUNTIME_CHECK(cudaDeviceSynchronize());
}
int init(const std::vector<uint8_t>& root_unique_id_val,
const int& rank,
const int& num_ranks,
const int& team_split_stride) {
nvshmemx_uniqueid_t root_unique_id;
nvshmemx_init_attr_t attr;
std::memcpy(&root_unique_id, root_unique_id_val.data(), sizeof(nvshmemx_uniqueid_t));
nvshmemx_set_attr_uniqueid_args(rank, num_ranks, &root_unique_id, &attr);
nvshmemx_init_attr(NVSHMEMX_INIT_WITH_UNIQUEID, &attr);
// Create sub-RDMA teams
if (team_split_stride > 0 and num_ranks > team_split_stride) {
EP_HOST_ASSERT(cpu_rdma_team == NVSHMEM_TEAM_INVALID);
EP_HOST_ASSERT(num_ranks % team_split_stride == 0);
EP_HOST_ASSERT(nvshmem_team_split_strided(NVSHMEM_TEAM_WORLD,
rank % team_split_stride,
team_split_stride,
num_ranks / team_split_stride,
&cpu_rdma_team_config,
0,
&cpu_rdma_team) == 0);
EP_HOST_ASSERT(cpu_rdma_team != NVSHMEM_TEAM_INVALID);
}
// Wait all GPUs to get ready
barrier(true);
return nvshmem_my_pe();
}
void finalize() {
barrier(true);
if (cpu_rdma_team != NVSHMEM_TEAM_INVALID) {
nvshmem_team_destroy(cpu_rdma_team);
cpu_rdma_team = NVSHMEM_TEAM_INVALID;
}
nvshmem_finalize();
}
} // namespace deep_ep::nvshmem
csrc/kernels/backend/symmetric.hpp
+319
View File
@@ -0,0 +1,319 @@
#pragma once
#include <cuda.h>
#include <nccl.h>
#include <vector>
#include <unistd.h>
#include <sys/syscall.h>
#include <deep_ep/common/exception.cuh>
#include <deep_ep/common/math.cuh>
#include "../../utils/lazy_driver.hpp"
namespace deep_ep::symmetric {
static constexpr int64_t kNumAlignmentBytes = 2097152;
// CPU communicator types: (pid, fd) handle and list of handles from all ranks
using cpu_handle_t = std::pair<int, int>;
using cpu_comm_t = std::vector<cpu_handle_t>;
struct DeviceContext {
int device_idx;
CUdevice device;
int numa_idx;
DeviceContext() {
CUDA_RUNTIME_CHECK(cudaGetDevice(&device_idx));
CUDA_DRIVER_CHECK(lazy_cuDeviceGet(&device, device_idx));
CUDA_DRIVER_CHECK(lazy_cuDeviceGetAttribute(&numa_idx, CU_DEVICE_ATTRIBUTE_HOST_NUMA_ID, device));
}
CUmemAllocationProp gpu_alloc_prop() const { return build_alloc_prop(CU_MEM_LOCATION_TYPE_DEVICE, device_idx); }
CUmemAllocationProp cpu_alloc_prop() const { return build_alloc_prop(CU_MEM_LOCATION_TYPE_HOST_NUMA, numa_idx); }
private:
CUmemAllocationProp build_alloc_prop(const CUmemLocationType& location_type, const int& location_idx) const {
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = location_type;
prop.location.id = location_idx;
int requested_handle_types = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
int flag = 0;
#if CUDART_VERSION >= 13000
CUDA_DRIVER_CHECK(lazy_cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device));
if (flag)
requested_handle_types |= CU_MEM_HANDLE_TYPE_FABRIC;
#endif
prop.requestedHandleTypes = static_cast<CUmemAllocationHandleType>(requested_handle_types);
if (location_type == CU_MEM_LOCATION_TYPE_DEVICE) {
flag = 0;
#if CUDART_VERSION >= 13000
CUDA_DRIVER_CHECK(lazy_cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED, device));
#else
CUDA_DRIVER_CHECK(lazy_cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_SUPPORTED, device));
#endif
EP_HOST_ASSERT(flag and "GPUDirect RDMA with CUDA VMM is not supported on this device");
prop.allocFlags.gpuDirectRDMACapable = 1;
}
// Check granularity
size_t num_granularity_bytes = 0;
CUDA_DRIVER_CHECK(lazy_cuMemGetAllocationGranularity(&num_granularity_bytes, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED));
EP_HOST_ASSERT((num_granularity_bytes == 0 or kNumAlignmentBytes % num_granularity_bytes == 0) and
"Alignment must be a multiple of CUDA allocation granularity");
return prop;
}
};
// Try cuMemCreate with FABRIC fallback (mirrors ncclMemAlloc logic)
static void cumem_create_with_fallback(CUmemGenericAllocationHandle* handle,
const int64_t& num_bytes, CUmemAllocationProp* prop) {
if (prop->requestedHandleTypes & CU_MEM_HANDLE_TYPE_FABRIC) {
CUresult err = lazy_cuMemCreate(handle, num_bytes, prop, 0);
if (err == CUDA_ERROR_NOT_PERMITTED or err == CUDA_ERROR_NOT_SUPPORTED) {
prop->requestedHandleTypes = static_cast<CUmemAllocationHandleType>(
prop->requestedHandleTypes & ~CU_MEM_HANDLE_TYPE_FABRIC);
CUDA_DRIVER_CHECK(lazy_cuMemCreate(handle, num_bytes, prop, 0));
} else {
CUDA_DRIVER_CHECK(err);
}
} else {
CUDA_DRIVER_CHECK(lazy_cuMemCreate(handle, num_bytes, prop, 0));
}
}
// Set read/write access for all peer GPUs, and optionally a NUMA node
static void set_access(const CUdeviceptr& addr, const int64_t& num_bytes,
const int& device_idx, const int& numa_idx = -1) {
const bool with_cpu = (numa_idx >= 0);
CUmemAccessDesc desc[2];
desc[0].location.type = CU_MEM_LOCATION_TYPE_DEVICE;
desc[0].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
if (with_cpu) {
desc[1].location.type = CU_MEM_LOCATION_TYPE_HOST_NUMA;
desc[1].location.id = numa_idx;
desc[1].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
}
int num_devices;
CUDA_RUNTIME_CHECK(cudaGetDeviceCount(&num_devices));
for (int i = 0; i < num_devices; ++ i) {
int can_access_peer = 0;
if (i == device_idx or (cudaDeviceCanAccessPeer(&can_access_peer, i, device_idx) == cudaSuccess and can_access_peer)) {
desc[0].location.id = i;
CUDA_DRIVER_CHECK(lazy_cuMemSetAccess(addr, num_bytes, desc, with_cpu ? 2 : 1));
}
}
}
class SymmetricMemory {
public:
void* ptr = nullptr;
int64_t num_bytes = 0;
int64_t num_gpu_bytes = 0;
int64_t num_cpu_bytes = 0;
virtual ~SymmetricMemory() noexcept(0) = default;
};
// Wraps ncclMemAlloc/ncclMemFree (pure GPU, current default)
class GPUSymmetricMemory final : public SymmetricMemory {
public:
explicit GPUSymmetricMemory(const int64_t& num_bytes) {
EP_HOST_ASSERT(num_bytes > 0 and num_bytes % kNumAlignmentBytes == 0);
NCCL_CHECK(ncclMemAlloc(&ptr, num_bytes));
EP_HOST_ASSERT(reinterpret_cast<uint64_t>(ptr) % kNumAlignmentBytes == 0);
this->num_bytes = num_bytes;
this->num_gpu_bytes = num_bytes;
}
~GPUSymmetricMemory() override {
if (ptr != nullptr) {
NCCL_CHECK(ncclMemFree(ptr));
ptr = nullptr;
}
}
};
// GPU + CPU mixed allocation via CUDA Driver API
// Memory layout: [GPU VRAM (front)] [CPU RAM / NUMA-local (back)]
// The entire contiguous VA range is compatible with `ncclCommWindowRegister`.
class ElasticSymmetricMemory : public SymmetricMemory {
CUmemGenericAllocationHandle gpu_handle = {};
CUmemGenericAllocationHandle cpu_handle = {};
public:
ElasticSymmetricMemory(const int64_t& num_gpu_bytes, const int64_t& num_cpu_bytes) {
EP_HOST_ASSERT(num_gpu_bytes > 0 and num_gpu_bytes % kNumAlignmentBytes == 0);
EP_HOST_ASSERT(num_cpu_bytes > 0 and num_cpu_bytes % kNumAlignmentBytes == 0);
DeviceContext ctx;
auto gpu_prop = ctx.gpu_alloc_prop();
auto cpu_prop = ctx.cpu_alloc_prop();
this->num_gpu_bytes = num_gpu_bytes;
this->num_cpu_bytes = num_cpu_bytes;
this->num_bytes = num_gpu_bytes + num_cpu_bytes;
// Reserve VA and map segments
CUdeviceptr addr;
CUDA_DRIVER_CHECK(lazy_cuMemAddressReserve(&addr, this->num_bytes, kNumAlignmentBytes, 0, 0));
this->ptr = reinterpret_cast<void*>(static_cast<uintptr_t>(addr));
cumem_create_with_fallback(&gpu_handle, num_gpu_bytes, &gpu_prop);
CUDA_DRIVER_CHECK(lazy_cuMemMap(addr, num_gpu_bytes, 0, gpu_handle, 0));
set_access(addr, num_gpu_bytes, ctx.device_idx);
cumem_create_with_fallback(&cpu_handle, num_cpu_bytes, &cpu_prop);
CUDA_DRIVER_CHECK(lazy_cuMemMap(addr + num_gpu_bytes, num_cpu_bytes, 0, cpu_handle, 0));
set_access(addr + num_gpu_bytes, num_cpu_bytes, ctx.device_idx, ctx.numa_idx);
EP_HOST_ASSERT(reinterpret_cast<uint64_t>(ptr) % kNumAlignmentBytes == 0);
}
~ElasticSymmetricMemory() override {
auto addr = static_cast<CUdeviceptr>(reinterpret_cast<uintptr_t>(ptr));
CUDA_DRIVER_CHECK(lazy_cuMemUnmap(addr, num_gpu_bytes));
CUDA_DRIVER_CHECK(lazy_cuMemRelease(gpu_handle));
CUDA_DRIVER_CHECK(lazy_cuMemUnmap(addr + num_gpu_bytes, num_cpu_bytes));
CUDA_DRIVER_CHECK(lazy_cuMemRelease(cpu_handle));
CUDA_DRIVER_CHECK(lazy_cuMemAddressFree(addr, num_bytes));
}
};
// Each rank creates its own NUMA-local CPU segment;
// CPU segments are imported and mapped contiguously into every rank's VA space.
// Layout: [GPU VRAM (front)] [CPU rank0 | CPU rank1 | ... | CPU rank(N-1) (back)]
class HybridElasticSymmetricMemory final : public SymmetricMemory {
int num_scaleup_ranks;
CUmemGenericAllocationHandle gpu_handle = {};
std::vector<CUmemGenericAllocationHandle> cpu_handles;
int local_export_fd = -1;
public:
HybridElasticSymmetricMemory(const cpu_comm_t& cpu_comm,
const int64_t& num_gpu_bytes, const int64_t& num_cpu_bytes,
const int& num_scaleup_ranks, const int& scaleout_rank_idx):
num_scaleup_ranks(num_scaleup_ranks),
cpu_handles(num_scaleup_ranks) {
EP_HOST_ASSERT(num_gpu_bytes > 0 and num_gpu_bytes % kNumAlignmentBytes == 0);
EP_HOST_ASSERT(num_cpu_bytes > 0 and num_cpu_bytes % kNumAlignmentBytes == 0);
DeviceContext ctx;
auto gpu_prop = ctx.gpu_alloc_prop();
this->num_gpu_bytes = num_gpu_bytes;
this->num_cpu_bytes = num_cpu_bytes;
this->num_bytes = num_gpu_bytes + num_cpu_bytes * num_scaleup_ranks;
// Reserve VA
CUdeviceptr addr;
CUDA_DRIVER_CHECK(lazy_cuMemAddressReserve(&addr, this->num_bytes, kNumAlignmentBytes, 0, 0));
this->ptr = reinterpret_cast<void*>(static_cast<uintptr_t>(addr));
// Map GPU segment
cumem_create_with_fallback(&gpu_handle, num_gpu_bytes, &gpu_prop);
CUDA_DRIVER_CHECK(lazy_cuMemMap(addr, num_gpu_bytes, 0, gpu_handle, 0));
set_access(addr, num_gpu_bytes, ctx.device_idx);
// Import and map all intra-node CPU segments
const auto local_pid = getpid();
for (int i = 0; i < num_scaleup_ranks; ++ i) {
auto [pid, fd] = cpu_comm[num_scaleup_ranks * scaleout_rank_idx + i];
int local_fd = fd;
if (pid != local_pid) {
int pidfd = syscall(SYS_pidfd_open, pid, 0);
EP_HOST_ASSERT(pidfd >= 0 and "`pidfd_open` failed");
local_fd = syscall(SYS_pidfd_getfd, pidfd, fd, 0);
EP_HOST_ASSERT(local_fd >= 0 and "`pidfd_getfd` failed");
close(pidfd);
}
const auto offset = num_gpu_bytes + i * num_cpu_bytes;
CUDA_DRIVER_CHECK(lazy_cuMemImportFromShareableHandle(
&cpu_handles[i],
reinterpret_cast<void*>(static_cast<uintptr_t>(local_fd)),
CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR));
CUDA_DRIVER_CHECK(lazy_cuMemMap(addr + offset, num_cpu_bytes, 0, cpu_handles[i], 0));
set_access(addr + offset, num_cpu_bytes, ctx.device_idx, ctx.numa_idx);
if (pid != local_pid) {
close(local_fd);
} else {
local_export_fd = local_fd;
}
}
EP_HOST_ASSERT(reinterpret_cast<uint64_t>(ptr) % kNumAlignmentBytes == 0);
}
~HybridElasticSymmetricMemory() override {
auto addr = static_cast<CUdeviceptr>(reinterpret_cast<uintptr_t>(ptr));
CUDA_DRIVER_CHECK(lazy_cuMemUnmap(addr, num_gpu_bytes));
CUDA_DRIVER_CHECK(lazy_cuMemRelease(gpu_handle));
if (local_export_fd >= 0)
close(local_export_fd);
CUDA_DRIVER_CHECK(lazy_cuMemUnmap(addr + num_gpu_bytes, num_cpu_bytes * num_scaleup_ranks));
for (int i = 0; i < num_scaleup_ranks; ++ i)
CUDA_DRIVER_CHECK(lazy_cuMemRelease(cpu_handles[i]));
CUDA_DRIVER_CHECK(lazy_cuMemAddressFree(addr, num_bytes));
}
// Create a NUMA-local CPU segment and export its POSIX FD handle.
// Returns (pid, fd) handle for cross-process sharing.
static cpu_handle_t create_cpu_handle(const int64_t& num_cpu_bytes) {
EP_HOST_ASSERT(num_cpu_bytes > 0 and num_cpu_bytes % kNumAlignmentBytes == 0);
DeviceContext ctx;
auto cpu_prop = ctx.cpu_alloc_prop();
CUmemGenericAllocationHandle handle;
cumem_create_with_fallback(&handle, num_cpu_bytes, &cpu_prop);
int fd = -1;
CUDA_DRIVER_CHECK(lazy_cuMemExportToShareableHandle(
&fd, handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0));
// Release the allocation handle -- the POSIX FD keeps the physical memory alive
CUDA_DRIVER_CHECK(lazy_cuMemRelease(handle));
return {getpid(), fd};
}
};
static std::shared_ptr<SymmetricMemory> alloc(const int64_t& num_gpu_bytes, const int64_t& num_cpu_bytes,
const bool& allow_hybrid_mode = false,
const int& num_scaleup_ranks = 0, const int& scaleout_rank_idx = 0,
const cpu_comm_t& cpu_comm = {}) {
EP_HOST_ASSERT(num_gpu_bytes > 0 and num_gpu_bytes % kNumAlignmentBytes == 0);
EP_HOST_ASSERT(num_cpu_bytes >= 0 and num_cpu_bytes % kNumAlignmentBytes == 0);
std::shared_ptr<SymmetricMemory> result;
if (num_cpu_bytes > 0) {
if (allow_hybrid_mode) {
result = std::make_shared<HybridElasticSymmetricMemory>(
cpu_comm, num_gpu_bytes, num_cpu_bytes,
num_scaleup_ranks, scaleout_rank_idx);
} else {
result = std::make_shared<ElasticSymmetricMemory>(
num_gpu_bytes, num_cpu_bytes);
}
} else {
result = std::make_shared<GPUSymmetricMemory>(num_gpu_bytes);
}
// TODO: move all variables into NCCL runtime
// Enable NCCL elastic buffer register if the allocator may produce CPU-backed segments
if (dynamic_cast<ElasticSymmetricMemory*>(result.get()) != nullptr)
setenv("NCCL_ELASTIC_BUFFER_REGISTER", "1", 0);
return result;
}
} // namespace deep_ep::symmetric
csrc/kernels/elastic/api.hpp
+9
View File
@@ -0,0 +1,9 @@
#pragma once
#include <deep_ep/common/compiled.cuh>
#include "barrier.hpp"
#include "dispatch.hpp"
#include "combine.hpp"
#include "engram.hpp"
#include "pp_send_recv.hpp"
csrc/kernels/elastic/barrier.hpp
+86
View File
@@ -0,0 +1,86 @@
#pragma once
#include <nccl.h>
#include <nccl_device.h>
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/exception.cuh>
#include "../../jit/compiler.hpp"
#include "../../jit/launch_runtime.hpp"
namespace deep_ep::elastic {
class BarrierRuntime final : public jit::LaunchRuntime<BarrierRuntime> {
public:
struct Args {
// Templated arguments
bool is_scaleup_nvlink;
int num_scaleout_ranks, num_scaleup_ranks;
int64_t num_timeout_cycles;
bool sequential;
// Parameters
ncclDevComm_t nccl_dev_comm;
ncclWindow_t nccl_window;
void* workspace;
int scaleout_rank_idx, scaleup_rank_idx;
jit::LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#include <deep_ep/impls/barrier.cuh>
using namespace deep_ep::elastic;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&barrier_impl<{}, {}, {}, {}, {}, {}, {}>);
}}
)", args.is_scaleup_nvlink,
args.launch_args.grid_dim.first, args.launch_args.num_threads,
args.num_scaleout_ranks, args.num_scaleup_ranks,
args.num_timeout_cycles, args.sequential);
}
static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) {
EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(
kernel, config,
args.nccl_dev_comm, args.nccl_window,
args.workspace, args.scaleout_rank_idx, args.scaleup_rank_idx
));
}
};
static void launch_barrier(const ncclDevComm_t& nccl_dev_comm, const ncclWindow_t& nccl_window,
void* workspace,
const int& scaleout_rank_idx, const int& scaleup_rank_idx,
const int& num_scaleout_ranks, const int& num_scaleup_ranks,
const int64_t& num_timeout_cycles,
const bool& is_scaleup_nvlink,
const bool& sequential,
const at::cuda::CUDAStream& stream) {
// Number of threads equals to the number of ranks
constexpr auto kNumThreads = 512;
// Generate, build and launch
// NOTES: only the parallel hybrid kernel needs 2 SMs; the sequential mode does scaleout and
// scaleup one after another, so a single SM is sufficient.
const auto num_sms = (not sequential and num_scaleout_ranks > 1) ? 2 : 1;
const BarrierRuntime::Args args = {
.is_scaleup_nvlink = is_scaleup_nvlink,
.num_scaleout_ranks = num_scaleout_ranks, .num_scaleup_ranks = num_scaleup_ranks,
.num_timeout_cycles = num_timeout_cycles,
.sequential = sequential,
.nccl_dev_comm = nccl_dev_comm,
.nccl_window = nccl_window,
.workspace = workspace,
.scaleout_rank_idx = scaleout_rank_idx, .scaleup_rank_idx = scaleup_rank_idx,
.launch_args = jit::LaunchArgs(num_sms, kNumThreads, 0, 1, true)};
const auto code = BarrierRuntime::generate(args);
const auto runtime = jit::compiler->build("barrier", code);
BarrierRuntime::launch(runtime, args, stream);
}
} // namespace deep_ep::elastic
csrc/kernels/elastic/combine.hpp
+289
View File
@@ -0,0 +1,289 @@
#pragma once
#include <nccl.h>
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/exception.cuh>
#include "../../jit/compiler.hpp"
#include "../../jit/launch_runtime.hpp"
namespace deep_ep::elastic {
class CombineRuntime final : public jit::LaunchRuntime<CombineRuntime> {
public:
struct Args {
// Templated arguments
bool is_scaleup_nvlink;
bool use_expanded_layout, allow_multiple_reduction;
int num_scaleup_warps, num_forward_warps;
int num_scaleout_ranks, num_scaleup_ranks;
int hidden;
int num_max_tokens_per_rank;
int num_experts;
int num_topk;
int num_qps;
int64_t num_timeout_cycles;
// Parameters
nv_bfloat16* x;
float* topk_weights;
int* src_metadata;
int* psum_num_recv_tokens_per_scaleup_rank;
int* token_metadata_at_forward;
int* channel_linked_list;
ncclDevComm_t nccl_dev_comm;
ncclWindow_t nccl_window;
void* buffer;
void* workspace;
int scaleout_rank_idx, scaleup_rank_idx;
int num_reduced_tokens;
jit::LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
std::string header_name, func_name;
if (args.num_scaleout_ranks == 1) {
header_name = "combine";
func_name = fmt::format("combine_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>",
args.is_scaleup_nvlink,
args.use_expanded_layout, args.allow_multiple_reduction,
args.launch_args.grid_dim.first,
args.launch_args.num_threads / 32,
args.num_scaleup_ranks * args.num_scaleout_ranks,
args.hidden,
args.num_max_tokens_per_rank,
args.num_experts,
args.num_topk,
args.num_qps, args.num_timeout_cycles);
} else {
header_name = "hybrid_combine";
func_name = fmt::format("hybrid_combine_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>",
args.use_expanded_layout, args.allow_multiple_reduction,
args.launch_args.grid_dim.first,
args.num_scaleup_warps, args.num_forward_warps,
args.num_scaleout_ranks, args.num_scaleup_ranks,
args.hidden,
args.num_max_tokens_per_rank,
args.num_experts,
args.num_topk,
args.num_qps,
args.num_timeout_cycles);
}
return fmt::format(R"(
#include <deep_ep/impls/{}.cuh>
using namespace deep_ep::elastic;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&{});
}}
)", header_name, func_name);
}
static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) {
if (args.num_scaleout_ranks == 1) {
EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(kernel, config,
args.x, args.topk_weights,
args.src_metadata, args.psum_num_recv_tokens_per_scaleup_rank,
args.nccl_dev_comm, args.nccl_window,
args.buffer, args.workspace,
args.scaleup_rank_idx,
args.num_reduced_tokens));
} else {
EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(kernel, config,
args.x, args.topk_weights,
args.src_metadata,
args.psum_num_recv_tokens_per_scaleup_rank,
args.token_metadata_at_forward,
args.channel_linked_list,
args.nccl_dev_comm, args.nccl_window,
args.buffer, args.workspace,
args.scaleout_rank_idx, args.scaleup_rank_idx,
args.num_reduced_tokens));
}
}
};
static layout::TokenLayout get_combine_token_layout(
const int& hidden, const int& elem_size, const int& num_topk) {
return layout::TokenLayout(hidden * elem_size, 0, num_topk, false);
}
static void* launch_combine(void* x,
void* topk_weights,
int* src_metadata,
int* psum_num_recv_tokens_per_scaleup_rank,
int* token_metadata_at_forward,
int* channel_linked_list,
const ncclDevComm_t& nccl_dev_comm, const ncclWindow_t& nccl_window,
void* buffer, void* workspace,
const int& num_reduced_tokens, const int& num_max_tokens_per_rank,
const int& hidden,
const int& num_experts, const int& num_topk,
const int& num_qps, const int64_t& num_timeout_cycles,
const int& num_scaleout_ranks, const int& num_scaleup_ranks,
const int& scaleout_rank_idx, const int& scaleup_rank_idx,
const bool& is_scaleup_nvlink,
const int& num_sms, const int& num_smem_bytes,
const int& num_channels,
const bool& use_expanded_layout, const bool& allow_multiple_reduction,
const at::cuda::CUDAStream& stream) {
// Maximize shared memory utilization
const auto token_layout = get_combine_token_layout(hidden, sizeof(nv_bfloat16), num_topk);
auto num_warps = std::min(num_smem_bytes / token_layout.get_num_bytes<true>(), 32);
// Decide warps
int num_scaleup_warps = 0, num_forward_warps = 0;
if (num_scaleout_ranks > 1) {
EP_HOST_ASSERT(num_channels % num_sms == 0 and
"Invalid number of channels or SMs, you may use a different SM count than dispatch");
EP_HOST_ASSERT(num_channels / num_sms <= 16);
num_scaleup_warps = num_forward_warps = num_channels / num_sms;
num_warps = num_scaleup_warps + num_forward_warps;
EP_HOST_ASSERT(num_warps * token_layout.get_num_bytes<true>() <= num_smem_bytes and
"Invalid combine SM count, please try to match your dispatch config");
}
// Generate, build and launch
const auto num_threads = num_warps * 32;
const CombineRuntime::Args args = {
.is_scaleup_nvlink = is_scaleup_nvlink,
.use_expanded_layout = use_expanded_layout,
.allow_multiple_reduction = allow_multiple_reduction,
.num_scaleup_warps = num_scaleup_warps, .num_forward_warps = num_forward_warps,
.num_scaleout_ranks = num_scaleout_ranks, .num_scaleup_ranks = num_scaleup_ranks,
.hidden = hidden,
.num_max_tokens_per_rank = num_max_tokens_per_rank,
.num_experts = num_experts,
.num_topk = num_topk,
.num_qps = num_qps, .num_timeout_cycles = num_timeout_cycles,
.x = static_cast<nv_bfloat16*>(x),
.topk_weights = static_cast<float*>(topk_weights),
.src_metadata = src_metadata,
.psum_num_recv_tokens_per_scaleup_rank = psum_num_recv_tokens_per_scaleup_rank,
.token_metadata_at_forward = token_metadata_at_forward,
.channel_linked_list = channel_linked_list,
.nccl_dev_comm = nccl_dev_comm, .nccl_window = nccl_window,
.buffer = buffer, .workspace = workspace,
.scaleout_rank_idx = scaleout_rank_idx, .scaleup_rank_idx = scaleup_rank_idx,
.num_reduced_tokens = num_reduced_tokens,
// NOTES: make cluster dim 2 to overlap with clustered computation kernels
.launch_args = jit::LaunchArgs(num_sms, num_threads, num_smem_bytes, 2 - (num_sms % 2), true)
};
const auto code = CombineRuntime::generate(args);
const auto runtime = jit::compiler->build("combine", code);
CombineRuntime::launch(runtime, args, stream);
// Return the buffer to be reduced
if (num_scaleout_ranks == 1)
return buffer;
// For hybrid mode, we have to skip the scale-up buffer
const bool is_scaleup_buffer_rank_layout =
allow_multiple_reduction ? (num_scaleup_ranks <= num_topk) : false;
const auto scaleup_buffer = layout::BufferLayout<false>(
token_layout,
is_scaleup_buffer_rank_layout ? num_scaleup_ranks : num_topk,
num_scaleout_ranks * num_max_tokens_per_rank,
buffer);
return scaleup_buffer.get_buffer_end_ptr();
}
class CombineReduceEpilogueRuntime final : public jit::LaunchRuntime<CombineReduceEpilogueRuntime> {
public:
struct Args {
// Templated arguments
bool use_expanded_layout, allow_multiple_reduction;
int num_scaleout_ranks, num_scaleup_ranks;
int hidden;
int num_max_tokens_per_rank;
int num_experts, num_topk;
// Parameters
nv_bfloat16* combined_x;
float* combined_topk_weights;
topk_idx_t* combined_topk_idx;
void* reduce_buffer;
void* bias_0;
void* bias_1;
int num_combined_tokens;
int scaleout_rank_idx, scaleup_rank_idx;
jit::LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#include <deep_ep/impls/combine_reduce_epilogue.cuh>
using namespace deep_ep::elastic;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&combine_reduce_epilogue_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}>);
}}
)", args.use_expanded_layout, args.allow_multiple_reduction,
args.launch_args.grid_dim.first,
args.launch_args.num_threads / 32,
args.num_scaleout_ranks, args.num_scaleup_ranks,
args.hidden,
args.num_max_tokens_per_rank,
args.num_experts, args.num_topk);
}
static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) {
EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(kernel, config,
args.combined_x,
args.combined_topk_weights,
args.combined_topk_idx,
args.reduce_buffer,
args.bias_0, args.bias_1,
args.num_combined_tokens,
args.scaleout_rank_idx, args.scaleup_rank_idx));
}
};
static void launch_combine_reduce_epilogue(void* combined_x,
float* combined_topk_weights,
topk_idx_t* combined_topk_idx,
const int& num_combined_tokens, const int& num_max_tokens_per_rank,
const int& hidden,
const int& num_experts, const int& num_topk,
void* reduce_buffer,
void* bias_0, void* bias_1,
const int& num_scaleout_ranks, const int& num_scaleup_ranks,
const int& scaleout_rank_idx, const int& scaleup_rank_idx,
const int& num_sms, const int& num_smem_bytes,
const bool& use_expanded_layout, const bool& allow_multiple_reduction,
const at::cuda::CUDAStream& stream) {
// Maximize shared memory utilization
// Too many warps may cause performance degrade, so we limit into 1024
const auto token_layout = layout::TokenLayout(hidden * sizeof(nv_bfloat16), 0, 0, false);
const auto num_warps = std::min<int>(num_smem_bytes / token_layout.get_num_bytes<false>(), 32);
const auto num_threads = num_warps * 32;
// Generate, build and launch
const CombineReduceEpilogueRuntime::Args args = {
.use_expanded_layout = use_expanded_layout,
.allow_multiple_reduction = allow_multiple_reduction,
.num_scaleout_ranks = num_scaleout_ranks, .num_scaleup_ranks = num_scaleup_ranks,
.hidden = hidden,
.num_max_tokens_per_rank = num_max_tokens_per_rank,
.num_experts = num_experts, .num_topk = num_topk,
.combined_x = static_cast<nv_bfloat16*>(combined_x),
.combined_topk_weights = combined_topk_weights,
.combined_topk_idx = combined_topk_idx,
.reduce_buffer = reduce_buffer,
.bias_0 = bias_0, .bias_1 = bias_1,
.num_combined_tokens = num_combined_tokens,
.scaleout_rank_idx = scaleout_rank_idx, .scaleup_rank_idx = scaleup_rank_idx,
.launch_args = jit::LaunchArgs(num_sms, num_threads, num_smem_bytes, 1, false, true)
};
const auto code = CombineReduceEpilogueRuntime::generate(args);
const auto runtime = jit::compiler->build("combine_reduce_epilogue", code);
CombineReduceEpilogueRuntime::launch(runtime, args, stream);
}
} // namespace deep_ep::elastic
csrc/kernels/elastic/dispatch.hpp
+342
View File
@@ -0,0 +1,342 @@
#pragma once
#include <nccl.h>
#include <nccl_device.h>
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/exception.cuh>
#include "../../jit/compiler.hpp"
#include "../../jit/launch_runtime.hpp"
namespace deep_ep::elastic {
class DispatchRuntime final : public jit::LaunchRuntime<DispatchRuntime> {
public:
struct Args {
// Templated arguments
bool is_scaleup_nvlink;
bool do_cpu_sync;
bool reuse_slot_indices;
int num_notify_warps;
int num_dispatch_warps; // For hybrid dispatch
int num_scaleout_warps, num_forward_warps; // For direct dispatch
int num_scaleout_ranks, num_scaleup_ranks;
int num_hidden_bytes, num_sf_packs;
int num_max_tokens_per_rank;
int num_experts, num_topk, expert_alignment;
int num_qps;
int64_t num_timeout_cycles;
// Parameters
void* x; sf_pack_t* sf; topk_idx_t* topk_idx; float* topk_weights;
topk_idx_t* copied_topk_idx;
int* cumulative_local_expert_recv_stats;
int* psum_num_recv_tokens_per_scaleup_rank;
int* psum_num_recv_tokens_per_expert;
int* num_unaligned_recv_tokens_per_expert;
int* dst_buffer_slot_idx;
int* token_metadata_at_forward;
int num_tokens;
int sf_token_stride, sf_hidden_stride;
ncclDevComm_t nccl_dev_comm;
ncclWindow_t nccl_window;
void* buffer;
void* workspace; void* mapped_host_workspace;
int scaleout_rank_idx, scaleup_rank_idx;
jit::LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
std::string header_name, func_name;
if (args.num_scaleout_ranks == 1) {
header_name = "dispatch";
func_name = fmt::format("dispatch_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>",
args.is_scaleup_nvlink,
args.do_cpu_sync,
args.reuse_slot_indices,
args.launch_args.grid_dim.first,
args.num_notify_warps, args.num_dispatch_warps,
args.num_scaleup_ranks,
args.num_hidden_bytes, args.num_sf_packs,
args.num_max_tokens_per_rank,
args.num_experts, args.num_topk, args.expert_alignment,
args.num_qps, args.num_timeout_cycles);
} else {
header_name = "hybrid_dispatch";
func_name = fmt::format("hybrid_dispatch_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>",
args.do_cpu_sync,
args.reuse_slot_indices,
args.launch_args.grid_dim.first,
args.num_notify_warps, args.num_scaleout_warps, args.num_forward_warps,
args.num_scaleout_ranks, args.num_scaleup_ranks,
args.num_hidden_bytes, args.num_sf_packs,
args.num_max_tokens_per_rank,
args.num_experts, args.num_topk, args.expert_alignment,
args.num_qps, args.num_timeout_cycles);
}
return fmt::format(R"(
#include <deep_ep/impls/{}.cuh>
using namespace deep_ep::elastic;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&{});
}}
)", header_name, func_name);
}
static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) {
if (args.num_scaleout_ranks == 1) {
EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(
kernel, config,
args.x, args.sf, args.topk_idx, args.topk_weights,
args.copied_topk_idx,
args.cumulative_local_expert_recv_stats,
args.psum_num_recv_tokens_per_scaleup_rank,
args.psum_num_recv_tokens_per_expert,
args.num_unaligned_recv_tokens_per_expert,
args.dst_buffer_slot_idx,
args.num_tokens,
args.sf_token_stride, args.sf_hidden_stride,
args.nccl_dev_comm, args.nccl_window,
args.buffer,
args.workspace, args.mapped_host_workspace,
args.scaleup_rank_idx));
} else {
EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(
kernel, config,
args.x, args.sf, args.topk_idx, args.topk_weights,
args.copied_topk_idx,
args.cumulative_local_expert_recv_stats,
args.psum_num_recv_tokens_per_scaleup_rank,
args.psum_num_recv_tokens_per_expert,
args.num_unaligned_recv_tokens_per_expert,
args.dst_buffer_slot_idx,
args.token_metadata_at_forward,
args.num_tokens,
args.sf_token_stride, args.sf_hidden_stride,
args.nccl_dev_comm, args.nccl_window,
args.buffer,
args.workspace, args.mapped_host_workspace,
args.scaleout_rank_idx, args.scaleup_rank_idx
));
}
}
};
constexpr int kNumNotifyWarps = 4;
static int get_num_notify_smem_bytes(const int& num_ranks, const int& num_experts) {
return math::align(num_ranks + num_experts, kNumNotifyWarps * 32) * sizeof(int);
}
static layout::TokenLayout get_dispatch_token_layout(
const int& hidden, const int& elem_size, const int& num_sf_packs, const int& num_topk) {
return layout::TokenLayout(hidden * elem_size, num_sf_packs * sizeof(sf_pack_t), num_topk, true);
}
static void launch_dispatch(void* x, void* sf,
topk_idx_t* topk_idx, float* topk_weights,
topk_idx_t* copied_topk_idx,
int* cumulative_local_expert_recv_stats,
int* psum_num_recv_tokens_per_scaleup_rank,
int* psum_num_recv_tokens_per_expert,
int* num_unaligned_recv_tokens_per_expert,
int* dst_buffer_slot_idx,
int* token_metadata_at_forward,
const int& num_tokens, const int& num_max_tokens_per_rank,
const int& hidden, const int& elem_size,
const int& num_sf_packs, const int& sf_token_stride, const int& sf_hidden_stride,
const int& num_experts, const int& num_topk, const int& expert_alignment,
const ncclDevComm_t& nccl_dev_comm, const ncclWindow_t& nccl_window,
void* buffer,
void* workspace, void* mapped_host_workspace,
const int& scaleout_rank_idx, const int& scaleup_rank_idx,
const int& num_scaleout_ranks, const int& num_scaleup_ranks,
const bool& is_scaleup_nvlink,
const int& num_sms, const int& num_channels_per_sm,
const int& num_smem_bytes,
const int& num_qps, const int64_t& num_timeout_cycles,
const bool& cached_mode,
const bool& do_cpu_sync,
const at::cuda::CUDAStream& stream) {
// Cached mode does not support expert token counting
if (cached_mode)
EP_HOST_ASSERT(cumulative_local_expert_recv_stats == nullptr);
// Utils
const auto num_ranks = num_scaleout_ranks * num_scaleup_ranks;
// Notify warps
// TODO: why don't we use 4 notify warps?
const int num_notify_warps = cached_mode ? 0 : kNumNotifyWarps;
const bool reuse_slot_indices = cached_mode;
const int num_notify_smem_bytes = cached_mode ? 0 : get_num_notify_smem_bytes(num_ranks, num_experts);
EP_HOST_ASSERT(num_notify_warps % 4 == 0);
// Other warps
int num_dispatch_warps = 0;
int num_scaleout_warps = 0, num_forward_warps = 0;
int num_threads = 0;
// Maximize shared memory utilization
if (num_scaleout_ranks == 1) {
const auto token_layout = get_dispatch_token_layout(hidden, elem_size, num_sf_packs, num_topk);
num_dispatch_warps = std::min<int>(
(num_smem_bytes - num_notify_smem_bytes) / token_layout.get_num_bytes<true>(), 32 - num_notify_warps);
num_threads = (num_notify_warps + num_dispatch_warps) * 32;
} else {
// Hybrid kernels
num_scaleout_warps = num_channels_per_sm;
num_forward_warps = num_channels_per_sm;
num_threads = (num_notify_warps + num_scaleout_warps + num_forward_warps) * 32;
}
// Generate, build and launch
const DispatchRuntime::Args args = {
.is_scaleup_nvlink = is_scaleup_nvlink,
.do_cpu_sync = do_cpu_sync,
.reuse_slot_indices = reuse_slot_indices,
.num_notify_warps = num_notify_warps,
.num_dispatch_warps = num_dispatch_warps,
.num_scaleout_warps = num_scaleout_warps, .num_forward_warps = num_forward_warps,
.num_scaleout_ranks = num_scaleout_ranks, .num_scaleup_ranks = num_scaleup_ranks,
.num_hidden_bytes = hidden * elem_size, .num_sf_packs = num_sf_packs,
.num_max_tokens_per_rank = num_max_tokens_per_rank,
.num_experts = num_experts, .num_topk = num_topk, .expert_alignment = expert_alignment,
.num_qps = num_qps, .num_timeout_cycles = num_timeout_cycles,
.x = x, .sf = static_cast<sf_pack_t*>(sf), .topk_idx = topk_idx, .topk_weights = topk_weights,
.copied_topk_idx = copied_topk_idx,
.cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats,
.psum_num_recv_tokens_per_scaleup_rank = psum_num_recv_tokens_per_scaleup_rank,
.psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert,
.num_unaligned_recv_tokens_per_expert = num_unaligned_recv_tokens_per_expert,
.dst_buffer_slot_idx = dst_buffer_slot_idx,
.token_metadata_at_forward = token_metadata_at_forward,
.num_tokens = num_tokens,
.sf_token_stride = sf_token_stride, .sf_hidden_stride = sf_hidden_stride,
.nccl_dev_comm = nccl_dev_comm, .nccl_window = nccl_window,
.buffer = buffer,
.workspace = workspace, .mapped_host_workspace = mapped_host_workspace,
.scaleout_rank_idx = scaleout_rank_idx, .scaleup_rank_idx = scaleup_rank_idx,
// NOTES: make cluster dim 2 to overlap with clustered computation kernels
.launch_args = jit::LaunchArgs(num_sms, num_threads, num_smem_bytes, 2 - (num_sms % 2), true)};
const auto code = DispatchRuntime::generate(args);
const auto runtime = jit::compiler->build("dispatch", code);
DispatchRuntime::launch(runtime, args, stream);
}
class DispatchCopyEpilogueRuntime final : public jit::LaunchRuntime<DispatchCopyEpilogueRuntime> {
public:
struct Args {
// Templated arguments
bool do_expand, cached_mode, do_zero_padding;
int num_channels;
int num_warps;
int num_scaleout_ranks, num_scaleup_ranks;
int num_hidden_bytes, num_sf_packs;
int num_max_tokens_per_rank;
int num_experts, num_topk, expert_alignment;
// Parameters
void *buffer, *workspace;
int* psum_num_recv_tokens_per_scaleup_rank;
int* psum_num_recv_tokens_per_expert;
void* recv_x; void* recv_sf;
topk_idx_t* recv_topk_idx; float* recv_topk_weights;
int* recv_src_metadata;
int* channel_linked_list;
int* num_unaligned_recv_tokens_per_expert;
int num_recv_tokens;
int recv_sf_token_stride, recv_sf_hidden_stride;
int scaleout_rank_idx, scaleup_rank_idx;
jit::LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#include <deep_ep/impls/dispatch_copy_epilogue.cuh>
using namespace deep_ep::elastic;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&dispatch_copy_epilogue_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>);
}}
)",
args.do_expand, args.cached_mode, args.do_zero_padding,
args.launch_args.grid_dim.first, args.num_channels, args.num_warps,
args.num_scaleout_ranks, args.num_scaleup_ranks,
args.num_hidden_bytes, args.num_sf_packs,
args.num_max_tokens_per_rank,
args.num_experts, args.num_topk, args.expert_alignment);
}
static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) {
EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(kernel, config,
args.buffer, args.workspace,
args.psum_num_recv_tokens_per_scaleup_rank,
args.psum_num_recv_tokens_per_expert,
args.recv_x, args.recv_sf, args.recv_topk_idx, args.recv_topk_weights,
args.recv_src_metadata,
args.channel_linked_list,
args.num_unaligned_recv_tokens_per_expert,
args.num_recv_tokens,
args.recv_sf_token_stride, args.recv_sf_hidden_stride,
args.scaleout_rank_idx, args.scaleup_rank_idx));
}
};
static void launch_dispatch_copy_epilogue(void* buffer, void* workspace,
int* psum_num_recv_tokens_per_scaleup_rank,
int* psum_num_recv_tokens_per_expert,
void* recv_x, void* recv_sf,
topk_idx_t* recv_topk_idx, float* recv_topk_weights,
int* recv_src_metadata,
int* channel_linked_list,
int* num_unaligned_recv_tokens_per_expert,
const int& num_recv_tokens, const int& num_max_tokens_per_rank,
const int& num_hidden_bytes,
const int& num_sf_packs, const int& recv_sf_token_stride, const int& recv_sf_hidden_stride,
const int& num_experts, const int& num_topk, const int& expert_alignment,
const int& scaleout_rank_idx, const int& scaleup_rank_idx,
const int& num_scaleout_ranks, const int& num_scaleup_ranks,
const int& num_sms, const int& num_smem_bytes,
const int& num_channels,
const bool& do_expand, const bool& cached_mode,
const bool& do_zero_padding,
const at::cuda::CUDAStream& stream) {
// Maximize shared memory utilization
const auto token_layout = layout::TokenLayout(num_hidden_bytes, num_sf_packs * sizeof(sf_pack_t), num_topk, true);
const auto num_warps = std::min(num_smem_bytes / token_layout.get_num_bytes<true>(), 32);
const auto num_threads = num_warps * 32;
// Generate, build and launch
const DispatchCopyEpilogueRuntime::Args args = {
.do_expand = do_expand, .cached_mode = cached_mode, .do_zero_padding = do_zero_padding,
.num_channels = num_channels, .num_warps = num_warps,
.num_scaleout_ranks = num_scaleout_ranks, .num_scaleup_ranks = num_scaleup_ranks,
.num_hidden_bytes = num_hidden_bytes, .num_sf_packs = num_sf_packs,
.num_max_tokens_per_rank = num_max_tokens_per_rank,
.num_experts = num_experts, .num_topk = num_topk, .expert_alignment = expert_alignment,
.buffer = buffer, .workspace = workspace,
.psum_num_recv_tokens_per_scaleup_rank = psum_num_recv_tokens_per_scaleup_rank,
.psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert,
.recv_x = recv_x, .recv_sf = recv_sf,
.recv_topk_idx = recv_topk_idx, .recv_topk_weights = recv_topk_weights,
.recv_src_metadata = recv_src_metadata,
.channel_linked_list = channel_linked_list,
.num_unaligned_recv_tokens_per_expert = num_unaligned_recv_tokens_per_expert,
.num_recv_tokens = num_recv_tokens,
.recv_sf_token_stride = recv_sf_token_stride, .recv_sf_hidden_stride = recv_sf_hidden_stride,
.scaleout_rank_idx = scaleout_rank_idx, .scaleup_rank_idx = scaleup_rank_idx,
.launch_args = jit::LaunchArgs(num_sms, num_threads, num_smem_bytes, 1, false, true)};
const auto code = DispatchCopyEpilogueRuntime::generate(args);
const auto runtime = jit::compiler->build("dispatch_copy_epilogue", code);
DispatchCopyEpilogueRuntime::launch(runtime, args, stream);
}
} // namespace deep_ep::elastic
csrc/kernels/elastic/engram.hpp
+191
View File
@@ -0,0 +1,191 @@
#pragma once
#include <nccl.h>
#include <nccl_device.h>
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/exception.cuh>
#include <deep_ep/common/layout.cuh>
#include "../../jit/compiler.hpp"
#include "../../jit/launch_runtime.hpp"
namespace deep_ep::elastic {
class EngramFetchRuntime final : public jit::LaunchRuntime<EngramFetchRuntime> {
public:
struct Args {
// Templated arguments
int num_entries_per_rank;
int num_hidden_bytes, num_sf_packs;
int num_entries_per_token;
int num_scaleout_ranks, num_scaleup_ranks;
int64_t num_cpu_bytes_per_rank;
int num_qps;
bool allow_hybrid_mode;
// Parameters
ncclDevComm_t nccl_dev_comm;
ncclWindow_t nccl_window;
void* storage;
void* fetched;
int* indices;
ncclGinRequest_t* last_gin_requests;
sf_pack_t* sf_table; sf_pack_t* fetched_sf;
int sf_token_stride; int sf_hidden_stride;
int num_tokens;
jit::LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
int num_rdma_peers, num_ranks_per_rdma_peer;
std::string team_tag;
if (args.allow_hybrid_mode) {
num_rdma_peers = args.num_scaleout_ranks;
num_ranks_per_rdma_peer = args.num_scaleup_ranks;
team_tag = "ncclTeamTagRail";
} else {
num_rdma_peers = args.num_scaleout_ranks * args.num_scaleup_ranks;
num_ranks_per_rdma_peer = 1;
team_tag = "ncclTeamTagWorld";
}
auto func_name = fmt::format("engram_fetch_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}>",
args.num_qps, args.num_entries_per_rank, args.num_hidden_bytes, args.num_sf_packs,
args.num_entries_per_token, num_rdma_peers, num_ranks_per_rdma_peer,
args.num_cpu_bytes_per_rank, args.launch_args.num_threads, team_tag);
return fmt::format(R"(
#include <deep_ep/impls/engram_fetch.cuh>
using namespace deep_ep::elastic;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&{});
}}
)", func_name);
}
static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) {
EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(
kernel, config,
args.nccl_dev_comm, args.nccl_window,
args.storage, args.fetched,
args.indices,
args.last_gin_requests,
args.sf_table, args.fetched_sf,
args.sf_token_stride, args.sf_hidden_stride,
args.num_tokens
));
}
};
static void launch_engram_fetch(const ncclDevComm_t& nccl_dev_comm, const ncclWindow_t& nccl_window,
void* storage, void* fetched,
int* indices,
ncclGinRequest_t* last_gin_requests,
void* sf_table, void* fetched_sf,
const int& sf_token_stride, const int& sf_hidden_stride,
const int& num_entries_per_rank,
const int& hidden, const int& elem_size, const int& num_sf_packs,
const int& num_entries_per_token,
const int& num_tokens,
const int& num_scaleout_ranks, const int& num_scaleup_ranks,
const int64_t& num_cpu_bytes_per_rank,
const int& num_qps,
const bool& allow_hybrid_mode,
const at::cuda::CUDAStream& stream) {
constexpr int kNumEngramFetchThreads = 1024;
// Generate, build and launch
const EngramFetchRuntime::Args args = {
.num_entries_per_rank = num_entries_per_rank,
.num_hidden_bytes = hidden * elem_size,
.num_sf_packs = num_sf_packs,
.num_entries_per_token = num_entries_per_token,
.num_scaleout_ranks = num_scaleout_ranks,
.num_scaleup_ranks = num_scaleup_ranks,
.num_cpu_bytes_per_rank = num_cpu_bytes_per_rank,
.num_qps = num_qps,
.allow_hybrid_mode = allow_hybrid_mode,
.nccl_dev_comm = nccl_dev_comm,
.nccl_window = nccl_window,
.storage = storage,
.fetched = fetched,
.indices = indices,
.last_gin_requests = last_gin_requests,
.sf_table = static_cast<sf_pack_t*>(sf_table),
.fetched_sf = static_cast<sf_pack_t*>(fetched_sf),
.sf_token_stride = sf_token_stride,
.sf_hidden_stride = sf_hidden_stride,
.num_tokens = num_tokens,
.launch_args = jit::LaunchArgs(num_qps, kNumEngramFetchThreads)};
const auto code = EngramFetchRuntime::generate(args);
const auto runtime = jit::compiler->build("engram_fetch", code);
EngramFetchRuntime::launch(runtime, args, stream);
}
class EngramFetchWaitRuntime final : public jit::LaunchRuntime<EngramFetchWaitRuntime> {
public:
struct Args {
// Templated arguments
int num_scaleout_ranks, num_scaleup_ranks;
bool allow_hybrid_mode;
ncclDevComm_t nccl_dev_comm;
ncclWindow_t nccl_window;
ncclGinRequest_t* last_gin_requests;
jit::LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
const int num_rdma_peers = args.allow_hybrid_mode
? args.num_scaleout_ranks
: args.num_scaleout_ranks * args.num_scaleup_ranks;
auto func_name = fmt::format("engram_fetch_wait_impl<{}, {}>",
num_rdma_peers, args.launch_args.num_threads);
return fmt::format(R"(
#include <deep_ep/impls/engram_fetch_wait.cuh>
using namespace deep_ep::elastic;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&{});
}}
)", func_name);
}
static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) {
EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(
kernel, config,
args.nccl_dev_comm, args.nccl_window,
args.last_gin_requests
));
}
};
static void launch_engram_fetch_wait(ncclGinRequest_t* last_gin_requests,
const ncclDevComm_t& nccl_dev_comm, const ncclWindow_t& nccl_window,
const int& num_scaleout_ranks, const int& num_scaleup_ranks,
const int& num_qps,
const bool& allow_hybrid_mode,
const at::cuda::CUDAStream& stream) {
constexpr int kNumEngramFetchWaitThreads = 1024;
// Generate, build and launch
const EngramFetchWaitRuntime::Args args = {
.num_scaleout_ranks = num_scaleout_ranks,
.num_scaleup_ranks = num_scaleup_ranks,
.allow_hybrid_mode = allow_hybrid_mode,
.nccl_dev_comm = nccl_dev_comm,
.nccl_window = nccl_window,
.last_gin_requests = last_gin_requests,
.launch_args = jit::LaunchArgs(num_qps, kNumEngramFetchWaitThreads)};
const auto code = EngramFetchWaitRuntime::generate(args);
const auto runtime = jit::compiler->build("engram_fetch_wait", code);
EngramFetchWaitRuntime::launch(runtime, args, stream);
}
} // namespace deep_ep::elastic
csrc/kernels/elastic/pp_send_recv.hpp
+182
View File
@@ -0,0 +1,182 @@
#pragma once
#include <nccl.h>
#include <nccl_device.h>
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/exception.cuh>
#include <deep_ep/common/layout.cuh>
#include "../../jit/compiler.hpp"
#include "../../jit/launch_runtime.hpp"
namespace deep_ep::elastic {
class PPSendRuntime final : public jit::LaunchRuntime<PPSendRuntime> {
public:
struct Args {
// Templated arguments
int num_ranks;
int num_smem_bytes;
int64_t num_timeout_cycles;
// Parameters
ncclDevComm_t nccl_dev_comm;
ncclWindow_t nccl_window;
void* x;
int64_t num_x_bytes;
void* buffer;
void* workspace;
int rank_idx, dst_rank_idx;
int64_t num_max_tensor_bytes;
int num_max_inflight_tensors;
jit::LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#include <deep_ep/impls/pp_send_recv.cuh>
using namespace deep_ep::elastic;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&pp_send_impl<{}, {}, {}, {}>);
}}
)", args.launch_args.grid_dim.first,
args.num_ranks,
args.num_smem_bytes,
args.num_timeout_cycles);
}
static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) {
EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(
kernel, config,
args.nccl_dev_comm, args.nccl_window,
args.x, args.num_x_bytes,
args.buffer, args.workspace,
args.rank_idx, args.dst_rank_idx,
args.num_max_tensor_bytes, args.num_max_inflight_tensors
));
}
};
static void launch_pp_send(const ncclDevComm_t& nccl_dev_comm,
const ncclWindow_t& nccl_window,
void* x, const int64_t& num_x_bytes,
void* buffer, void* workspace,
const int& rank_idx, const int& dst_rank_idx, const int& num_ranks,
const int64_t& num_max_tensor_bytes,
const int num_max_inflight_tensors,
const int& num_sms,
const int64_t& num_timeout_cycles,
const int& num_smem_bytes,
const at::cuda::CUDAStream& stream) {
// Generate, build and launch
const PPSendRuntime::Args args = {
.num_ranks = num_ranks,
.num_smem_bytes = num_smem_bytes,
.num_timeout_cycles = num_timeout_cycles,
.nccl_dev_comm = nccl_dev_comm,
.nccl_window = nccl_window,
.x = x,
.num_x_bytes = num_x_bytes,
.buffer = buffer,
.workspace = workspace,
.rank_idx = rank_idx,
.dst_rank_idx = dst_rank_idx,
.num_max_tensor_bytes = num_max_tensor_bytes,
.num_max_inflight_tensors = num_max_inflight_tensors,
.launch_args = jit::LaunchArgs(num_sms, 32, num_smem_bytes, 1, true)
};
const auto code = PPSendRuntime::generate(args);
const auto runtime = jit::compiler->build("pp_send", code);
PPSendRuntime::launch(runtime, args, stream);
}
class PPRecvRuntime final : public jit::LaunchRuntime<PPRecvRuntime> {
public:
struct Args {
// Templated arguments
int num_ranks;
int num_smem_bytes;
int64_t num_timeout_cycles;
// Parameters
ncclDevComm_t nccl_dev_comm;
ncclWindow_t nccl_window;
void* x;
int64_t num_x_bytes;
void* buffer;
void* workspace;
int rank_idx;
int src_rank_idx;
int64_t num_max_tensor_bytes;
int num_max_inflight_tensors;
jit::LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#include <deep_ep/impls/pp_send_recv.cuh>
using namespace deep_ep::elastic;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&pp_recv_impl<{}, {}, {}, {}>);
}}
)", args.launch_args.grid_dim.first,
args.num_ranks,
args.num_smem_bytes,
args.num_timeout_cycles);
}
static void launch_impl(const jit::KernelHandle& kernel, const jit::LaunchConfigHandle& config, Args args) {
EP_CUDA_UNIFIED_CHECK(jit::launch_kernel(
kernel, config,
args.nccl_dev_comm, args.nccl_window,
args.x, args.num_x_bytes,
args.buffer, args.workspace,
args.rank_idx, args.src_rank_idx,
args.num_max_tensor_bytes,
args.num_max_inflight_tensors
));
}
};
static void launch_pp_recv(const ncclDevComm_t& nccl_dev_comm,
const ncclWindow_t& nccl_window,
void* x,
const int64_t& num_x_bytes,
void* buffer, void* workspace,
const int& rank_idx, const int& src_rank_idx, const int& num_ranks,
const int64_t& num_max_tensor_bytes,
const int& num_max_inflight_tensors,
const int& num_sms,
const int64_t& num_timeout_cycles,
const int& num_smem_bytes,
const at::cuda::CUDAStream& stream) {
// Generate, build and launch
const PPRecvRuntime::Args args = {
.num_ranks = num_ranks,
.num_smem_bytes = num_smem_bytes,
.num_timeout_cycles = num_timeout_cycles,
.nccl_dev_comm = nccl_dev_comm,
.nccl_window = nccl_window,
.x = x,
.num_x_bytes = num_x_bytes,
.buffer = buffer,
.workspace = workspace,
.rank_idx = rank_idx,
.src_rank_idx = src_rank_idx,
.num_max_tensor_bytes = num_max_tensor_bytes,
.num_max_inflight_tensors = num_max_inflight_tensors,
.launch_args = jit::LaunchArgs(num_sms, 32, num_smem_bytes, 1, true)
};
const auto code = PPRecvRuntime::generate(args);
const auto runtime = jit::compiler->build("pp_recv", code);
PPRecvRuntime::launch(runtime, args, stream);
}
} // namespace deep_ep::elastic
csrc/kernels/legacy/CMakeLists.txt
+7
View File
@@ -0,0 +1,7 @@
add_deep_ep_library(legacy_layout_cuda layout.cu)
add_deep_ep_library(legacy_intranode_cuda intranode.cu)
add_deep_ep_library(legacy_internode_cuda internode.cu)
add_deep_ep_library(legacy_internode_ll_cuda internode_ll.cu)
# Link these libraries later
set(LEGACY_CUDA_LIBRARIES legacy_layout_cuda legacy_intranode_cuda legacy_internode_cuda legacy_internode_ll_cuda CACHE INTERNAL "Legacy kernels")
csrc/kernels/legacy/api.cuh
+326
View File
@@ -0,0 +1,326 @@
#pragma once
#include "compiled.cuh"
namespace deep_ep::legacy {
// Layout kernels
namespace layout {
void get_dispatch_layout(const topk_idx_t* topk_idx,
int* num_tokens_per_rank,
int* num_tokens_per_rdma_rank,
int* num_tokens_per_expert,
bool* is_token_in_rank,
int num_tokens,
int num_topk,
int num_ranks,
int num_experts,
cudaStream_t stream);
} // namespace layout
// Intranode kernels
namespace intranode {
void barrier(int** barrier_signal_ptrs, int rank, int num_ranks, cudaStream_t stream);
void notify_dispatch(const int* num_tokens_per_rank,
int* moe_recv_counter_mapped,
int num_ranks,
const int* num_tokens_per_expert,
int* moe_recv_expert_counter_mapped,
int num_experts,
int num_tokens,
const bool* is_token_in_rank,
int* channel_prefix_matrix,
int* rank_prefix_matrix_copy,
int num_memset_int,
int expert_alignment,
void** buffer_ptrs,
int** barrier_signal_ptrs,
int rank,
cudaStream_t stream,
int num_sms);
void cached_notify_dispatch(const int* rank_prefix_matrix,
int num_memset_int,
void** buffer_ptrs,
int** barrier_signal_ptrs,
int rank,
int num_ranks,
cudaStream_t stream);
void dispatch(void* recv_x,
float* recv_x_scales,
int* recv_src_idx,
topk_idx_t* recv_topk_idx,
float* recv_topk_weights,
int* recv_channel_offset,
int* send_head,
const void* x,
const float* x_scales,
const topk_idx_t* topk_idx,
const float* topk_weights,
const bool* is_token_in_rank,
const int* channel_prefix_matrix,
int num_tokens,
int num_worst_tokens,
int hidden_int4,
int num_topk,
int num_experts,
int num_scales,
int scale_token_stride,
int scale_hidden_stride,
void** buffer_ptrs,
int rank,
int num_ranks,
cudaStream_t stream,
int num_sms,
int num_max_send_tokens,
int num_recv_buffer_tokens);
void cached_notify_combine(void** buffer_ptrs,
int* send_head,
int num_channels,
int num_recv_tokens,
int num_memset_int,
int** barrier_signal_ptrs,
int rank,
int num_ranks,
cudaStream_t stream);
void combine(cudaDataType_t type,
void* recv_x,
float* recv_topk_weights,
const void* x,
const float* topk_weights,
const void* bias_0,
const void* bias_1,
const int* src_idx,
const int* rank_prefix_matrix,
const int* channel_prefix_matrix,
int* send_head,
int num_tokens,
int num_recv_tokens,
int hidden,
int num_topk,
void** buffer_ptrs,
int rank,
int num_ranks,
cudaStream_t stream,
int num_sms,
int num_max_send_tokens,
int num_recv_buffer_tokens);
} // namespace intranode
// Internode kernels
namespace internode {
int get_source_meta_bytes();
void notify_dispatch(const int* num_tokens_per_rank,
int* moe_recv_counter_mapped,
int num_ranks,
const int* num_tokens_per_rdma_rank,
int* moe_recv_rdma_counter_mapped,
const int* num_tokens_per_expert,
int* moe_recv_expert_counter_mapped,
int num_experts,
const bool* is_token_in_rank,
int num_tokens,
int num_worst_tokens,
int num_channels,
int hidden_int4,
int num_scales,
int num_topk,
int expert_alignment,
int* rdma_channel_prefix_matrix,
int* recv_rdma_rank_prefix_sum,
int* gbl_channel_prefix_matrix,
int* recv_gbl_rank_prefix_sum,
void* rdma_buffer_ptr,
int num_max_rdma_chunked_recv_tokens,
void** buffer_ptrs,
int num_max_nvl_chunked_recv_tokens,
int** barrier_signal_ptrs,
int rank,
cudaStream_t stream,
int64_t num_rdma_bytes,
int64_t num_nvl_bytes,
bool low_latency_mode);
void dispatch(void* recv_x,
float* recv_x_scales,
topk_idx_t* recv_topk_idx,
float* recv_topk_weights,
void* recv_src_meta,
const void* x,
const float* x_scales,
const topk_idx_t* topk_idx,
const float* topk_weights,
int* send_rdma_head,
int* send_nvl_head,
int* recv_rdma_channel_prefix_matrix,
int* recv_gbl_channel_prefix_matrix,
const int* rdma_channel_prefix_matrix,
const int* recv_rdma_rank_prefix_sum,
const int* gbl_channel_prefix_matrix,
const int* recv_gbl_rank_prefix_sum,
const bool* is_token_in_rank,
int num_tokens,
int num_worst_tokens,
int hidden_int4,
int num_scales,
int num_topk,
int num_experts,
int scale_token_stride,
int scale_hidden_stride,
void* rdma_buffer_ptr,
int num_max_rdma_chunked_send_tokens,
int num_max_rdma_chunked_recv_tokens,
void** buffer_ptrs,
int num_max_nvl_chunked_send_tokens,
int num_max_nvl_chunked_recv_tokens,
int rank,
int num_ranks,
bool is_cached_dispatch,
cudaStream_t stream,
int num_channels,
bool low_latency_mode);
void cached_notify(int hidden_int4,
int num_scales,
int num_topk_idx,
int num_topk_weights,
int num_ranks,
int num_channels,
int num_combined_tokens,
int* combined_rdma_head,
const int* rdma_channel_prefix_matrix,
const int* rdma_rank_prefix_sum,
int* combined_nvl_head,
void* rdma_buffer_ptr,
int num_max_rdma_chunked_recv_tokens,
void** buffer_ptrs,
int num_max_nvl_chunked_recv_tokens,
int** barrier_signal_ptrs,
int rank,
cudaStream_t stream,
int64_t num_rdma_bytes,
int64_t num_nvl_bytes,
bool is_cached_dispatch,
bool low_latency_mode);
void combine(cudaDataType_t type,
void* combined_x,
float* combined_topk_weights,
const bool* is_combined_token_in_rank,
const void* x,
const float* topk_weights,
const void* bias_0,
const void* bias_1,
const int* combined_rdma_head,
const int* combined_nvl_head,
const void* src_meta,
const int* rdma_channel_prefix_matrix,
const int* rdma_rank_prefix_sum,
const int* gbl_channel_prefix_matrix,
int num_tokens,
int num_combined_tokens,
int hidden,
int num_topk,
void* rdma_buffer_ptr,
int num_max_rdma_chunked_send_tokens,
int num_max_rdma_chunked_recv_tokens,
void** buffer_ptrs,
int num_max_nvl_chunked_send_tokens,
int num_max_nvl_chunked_recv_tokens,
int rank,
int num_ranks,
cudaStream_t stream,
int num_channels,
bool low_latency_mode);
} // namespace internode
// Internode low-latency kernels
namespace internode_ll {
void clean_low_latency_buffer(int* clean_0,
int num_clean_int_0,
int* clean_1,
int num_clean_int_1,
int rank,
int num_ranks,
int* mask_buffer,
int* sync_buffer,
cudaStream_t stream);
void dispatch(void* packed_recv_x,
void* packed_recv_x_scales,
int* packed_recv_src_info,
int64_t* packed_recv_layout_range,
int* packed_recv_count,
int* mask_buffer,
int* cumulative_local_expert_recv_stats,
int64_t* dispatch_wait_recv_cost_stats,
void* rdma_recv_x,
int* rdma_recv_count,
void* rdma_x,
const void* x,
const topk_idx_t* topk_idx,
int* next_clean,
int num_next_clean_int,
int num_tokens,
int hidden,
int num_max_dispatch_tokens_per_rank,
int num_topk,
int num_experts,
int rank,
int num_ranks,
bool use_fp8,
bool round_scale,
bool use_ue8m0,
void* workspace,
int num_device_sms,
cudaStream_t stream,
int phases);
void combine(void* combined_x,
void* rdma_recv_x,
int* rdma_recv_flag,
void* rdma_send_x,
const void* x,
const topk_idx_t* topk_idx,
const float* topk_weights,
const int* src_info,
const int64_t* layout_range,
int* mask_buffer,
int64_t* combine_wait_recv_cost_stats,
int* next_clean,
int num_next_clean_int,
int num_combined_tokens,
int hidden,
int num_max_dispatch_tokens_per_rank,
int num_topk,
int num_experts,
int rank,
int num_ranks,
bool use_logfmt,
void* workspace,
int num_device_sms,
cudaStream_t stream,
int phases,
bool zero_copy);
void query_mask_buffer(int* mask_buffer_ptr, int num_ranks, int* output_mask_tensor, cudaStream_t stream);
void update_mask_buffer(int* mask_buffer_ptr, int rank_to_mask, bool mask, cudaStream_t stream);
void clean_mask_buffer(int* mask_buffer_ptr, int num_ranks, cudaStream_t stream);
} // namespace internode_ll
} // namespace deep_ep::legacy
csrc/kernels/legacy/buffer.cuh
+132
View File
@@ -0,0 +1,132 @@
#pragma once
#include <deep_ep/common/exception.cuh>
#include "compiled.cuh"
namespace deep_ep::legacy {
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::legacy
csrc/kernels/legacy/compiled.cuh
+18
View File
@@ -0,0 +1,18 @@
#pragma once
#include <deep_ep/common/compiled.cuh>
#define LEGACY_NUM_MAX_NVL_PEERS 8
#define LEGACY_NUM_MAX_RDMA_PEERS 20
#define LEGACY_NUM_WORKSPACE_BYTES (32 * 1024 * 1024)
#define LEGACY_NUM_MAX_LOCAL_EXPERTS 1024
#define LEGACY_NUM_BUFFER_ALIGNMENT_BYTES 128
#define LEGACY_LOW_LATENCY_SEND_PHASE 1
#define LEGACY_LOW_LATENCY_RECV_PHASE 2
#define LEGACY_FINISHED_SUM_TAG 1024
#define LEGACY_NUM_WAIT_NANOSECONDS 500
#define LEGACY_NUM_CPU_TIMEOUT_SECS 100
#define LEGACY_NUM_TIMEOUT_CYCLES 200000000000ull // 200G cycles ~= 100s
csrc/kernels/legacy/ibgda_device.cuh
+496
View File
@@ -0,0 +1,496 @@
// Portions derived from NVSHMEM (https://developer.nvidia.com/nvshmem)
// Copyright (c) NVIDIA Corporation.
// Licensed under the NVSHMEM Software License Agreement (version: September 3, 2019).
// See full license at: https://docs.nvidia.com/nvshmem/api/sla.html
//
// Modified from original source:
// - nvshmem/src/include/non_abi/device/pt-to-pt/ibgda_device.cuh
#pragma once
#include <nvshmem.h>
#include <device_host_transport/nvshmem_common_ibgda.h>
#include <non_abi/device/threadgroup/nvshmemi_common_device_defines.cuh>
#include <deep_ep/common/exception.cuh>
#include "utils.cuh"
namespace deep_ep::legacy {
EP_STATIC_ASSERT(NVSHMEMI_IBGDA_MIN_QP_DEPTH >= 64, "Invalid QP minimum depth");
__device__ static __forceinline__ uint64_t HtoBE64(uint64_t x) {
uint64_t ret;
asm("{\n\t"
".reg .b32 ign;\n\t"
".reg .b32 lo;\n\t"
".reg .b32 hi;\n\t"
".reg .b32 new_lo;\n\t"
".reg .b32 new_hi;\n\t"
"mov.b64 {lo,hi}, %1;\n\t"
"prmt.b32 new_hi, lo, ign, 0x0123;\n\t"
"prmt.b32 new_lo, hi, ign, 0x0123;\n\t"
"mov.b64 %0, {new_lo,new_hi};\n\t"
"}"
: "=l"(ret)
: "l"(x));
return ret;
}
__device__ static __forceinline__ uint32_t HtoBE32(uint32_t x) {
uint32_t ret;
asm("{\n\t"
".reg .b32 ign;\n\t"
"prmt.b32 %0, %1, ign, 0x0123;\n\t"
"}"
: "=r"(ret)
: "r"(x));
return ret;
}
__device__ static __forceinline__ uint16_t HtoBE16(uint16_t x) {
// TODO: simplify PTX using 16-bit instructions
auto a = static_cast<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;
}
__device__ static __forceinline__ nvshmemi_ibgda_device_qp_t* ibgda_get_rc(int pe, int id) {
auto state = ibgda_get_state();
const auto num_rc_per_pe = ibgda_get_state()->num_rc_per_pe;
return &state->globalmem
.rcs[pe * num_rc_per_pe * state->num_devices_initialized + id % (num_rc_per_pe * state->num_devices_initialized)];
}
__device__ static __forceinline__ void ibgda_lock_acquire(int* lock) {
while (atomicCAS(lock, 0, 1) == 1)
;
// Prevent reordering before the lock is acquired
memory_fence_cta();
}
__device__ static __forceinline__ void ibgda_lock_release(int* lock) {
memory_fence_cta();
// Prevent reordering before lock is released
st_na_relaxed(lock, 0);
}
__device__ static __forceinline__ void ibgda_update_dbr(nvshmemi_ibgda_device_qp_t* qp, uint32_t dbrec_head) {
// `DBREC` contains the index of the next empty `WQEBB`
__be32 dbrec_val;
__be32* dbrec_ptr = qp->tx_wq.dbrec;
// This is equivalent to `WRITE_ONCE(dbrec_ptr, HtoBE32(dbrec_head & 0xffff))`
asm("{\n\t"
".reg .b32 dbrec_head_16b;\n\t"
".reg .b32 ign;\n\t"
"and.b32 dbrec_head_16b, %1, 0xffff;\n\t"
"prmt.b32 %0, dbrec_head_16b, ign, 0x123;\n\t"
"}"
: "=r"(dbrec_val)
: "r"(dbrec_head));
st_na_release(dbrec_ptr, dbrec_val);
}
__device__ static __forceinline__ void ibgda_ring_db(nvshmemi_ibgda_device_qp_t* qp, uint16_t prod_idx) {
auto bf_ptr = reinterpret_cast<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::legacy
csrc/kernels/legacy/internode.cu
+2384
View File
File diff suppressed because it is too large Load Diff
csrc/kernels/legacy/internode_ll.cu
+1289
View File
File diff suppressed because it is too large Load Diff
csrc/kernels/legacy/intranode.cu
+1117
View File
File diff suppressed because it is too large Load Diff
csrc/kernels/legacy/launch.cuh
+134
View File
@@ -0,0 +1,134 @@
#pragma once
#include "compiled.cuh"
#ifndef SETUP_LAUNCH_CONFIG
#ifndef DISABLE_SM90_FEATURES
#define SETUP_LAUNCH_CONFIG(num_sms, num_threads, stream) \
cudaLaunchConfig_t cfg = {(num_sms), (num_threads), 0, stream, nullptr, 0}; \
cudaLaunchAttribute attr[2]; \
attr[0].id = cudaLaunchAttributeCooperative; \
attr[0].val.cooperative = 1; \
attr[1].id = cudaLaunchAttributeClusterDimension; \
attr[1].val.clusterDim.x = (num_sms % 2 == 0 ? 2 : 1); \
attr[1].val.clusterDim.y = 1; \
attr[1].val.clusterDim.z = 1; \
cfg.attrs = attr; \
cfg.numAttrs = 2
#else
#define SETUP_LAUNCH_CONFIG(sms, threads, stream) \
int __num_sms = (sms); \
int __num_threads = (threads); \
auto __stream = (stream)
#endif
#endif
#ifndef LAUNCH_KERNEL
#ifndef DISABLE_SM90_FEATURES
#define LAUNCH_KERNEL(config, kernel, ...) CUDA_RUNTIME_CHECK(cudaLaunchKernelEx(config, kernel, ##__VA_ARGS__))
#else
#define LAUNCH_KERNEL(config, kernel, ...) \
do { \
kernel<<<__num_sms, __num_threads, 0, __stream>>>(__VA_ARGS__); \
cudaError_t e = cudaGetLastError(); \
if (e != cudaSuccess) { \
EPException cuda_exception("CUDA", __FILE__, __LINE__, cudaGetErrorString(e)); \
fprintf(stderr, "%s\n", cuda_exception.what()); \
throw cuda_exception; \
} \
} while (0)
#endif
#endif
#ifndef SET_SHARED_MEMORY_FOR_TMA
#ifndef DISABLE_SM90_FEATURES
#define SET_SHARED_MEMORY_FOR_TMA(kernel) \
EP_HOST_ASSERT(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size) == cudaSuccess); \
cfg.dynamicSmemBytes = smem_size;
#else
#define SET_SHARED_MEMORY_FOR_TMA(kernel) void()
#endif
#endif
#define SWITCH_RANKS(case_macro) \
switch (num_ranks) { \
case 2: \
case_macro(2); \
case 4: \
case_macro(4); \
case 8: \
case_macro(8); \
default: \
EP_HOST_ASSERT(false and "Unsupported ranks"); \
} \
while (false)
#define SWITCH_RDMA_RANKS(case_macro) \
switch (num_ranks / LEGACY_NUM_MAX_NVL_PEERS) { \
case 2: \
case_macro(2); \
case 3: \
case_macro(3); \
case 4: \
case_macro(4); \
case 6: \
case_macro(6); \
case 8: \
case_macro(8); \
case 12: \
case_macro(12); \
case 16: \
case_macro(16); \
case 18: \
case_macro(18); \
case 20: \
case_macro(20); \
default: \
EP_HOST_ASSERT(false and "Unsupported RDMA ranks"); \
} \
while (false)
#define SWITCH_RANKS_WITH_DTYPE(dtype, case_macro) \
switch (num_ranks) { \
case 2: \
case_macro(dtype, 2); \
case 4: \
case_macro(dtype, 4); \
case 8: \
case_macro(dtype, 8); \
default: \
EP_HOST_ASSERT(false and "Unsupported ranks"); \
} \
while (false)
#define SWITCH_TYPES(case_macro) \
switch (type) { \
case CUDA_R_16BF: \
case_macro(nv_bfloat16); \
default: \
EP_HOST_ASSERT(false and "Unsupported type"); \
} \
while (false)
#define SWITCH_HIDDEN(case_macro) \
switch (hidden) { \
case 2048: \
case_macro(2048); \
case 2560: \
case_macro(2560); \
case 3072: \
case_macro(3072); /* for gpt-oss */ \
case 4096: \
case_macro(4096); \
case 5120: \
case_macro(5120); \
case 6144: \
case_macro(6144); /* For qwen3 coder */ \
case 7168: \
case_macro(7168); \
case 8192: \
case_macro(8192); \
default: \
EP_HOST_ASSERT(false and "Unsupported hidden"); \
} \
while (false)
csrc/kernels/legacy/layout.cu
+154
View File
@@ -0,0 +1,154 @@
#include <deep_ep/common/exception.cuh>
#include "compiled.cuh"
#include "launch.cuh"
namespace deep_ep::legacy {
namespace layout {
template <int kNumThreads, int kNumExpertsPerSM, int kNumRanksPerSM>
__global__ void get_dispatch_layout(const topk_idx_t* topk_idx,
int* num_tokens_per_rank,
int* num_tokens_per_rdma_rank,
int* num_tokens_per_expert,
bool* is_token_in_rank,
int num_tokens,
int num_topk,
int num_ranks,
int num_experts) {
auto sm_id = static_cast<int>(blockIdx.x);
auto thread_id = static_cast<int>(threadIdx.x);
// Count expert statistics
__shared__ int num_tokens_per_expert_per_thread[kNumThreads][kNumExpertsPerSM];
int expert_begin_idx = sm_id * kNumExpertsPerSM, expert_end_idx = min(expert_begin_idx + kNumExpertsPerSM, num_experts);
if (expert_begin_idx < expert_end_idx) {
// Per-thread count
#pragma unroll
for (int i = 0; i < kNumExpertsPerSM; ++i)
num_tokens_per_expert_per_thread[thread_id][i] = 0;
#pragma unroll
for (int i = thread_id; i < num_tokens; i += kNumThreads) {
auto shifted_topk_idx = topk_idx + i * num_topk;
#pragma unroll
for (int j = 0, expert_idx; j < num_topk; ++j) {
expert_idx = static_cast<int>(shifted_topk_idx[j]);
if (expert_begin_idx <= expert_idx and expert_idx < expert_end_idx)
++num_tokens_per_expert_per_thread[thread_id][expert_idx - expert_begin_idx];
}
}
__syncthreads();
// Sum up
EP_STATIC_ASSERT(kNumExpertsPerSM <= kNumThreads, "Too many experts per SM");
if (expert_begin_idx + thread_id < expert_end_idx) {
int sum = 0;
#pragma unroll
for (int i = 0; i < kNumThreads; ++i)
sum += num_tokens_per_expert_per_thread[i][thread_id];
num_tokens_per_expert[expert_begin_idx + thread_id] = sum;
}
return;
}
if (num_tokens_per_rdma_rank != nullptr)
EP_DEVICE_ASSERT(num_ranks % LEGACY_NUM_MAX_NVL_PEERS == 0 and num_ranks > LEGACY_NUM_MAX_NVL_PEERS);
// Count rank statistics
constexpr int kNumRDMARanksPerSM = kNumRanksPerSM / LEGACY_NUM_MAX_NVL_PEERS;
__shared__ int num_tokens_per_rank_per_thread[kNumThreads][kNumRanksPerSM];
__shared__ int num_tokens_per_rdma_rank_per_thread[kNumThreads][kNumRDMARanksPerSM];
auto sm_begin = (num_experts + kNumExpertsPerSM - 1) / kNumExpertsPerSM;
int rank_begin_idx = (sm_id - sm_begin) * kNumRanksPerSM, rank_end_idx = min(rank_begin_idx + kNumRanksPerSM, num_ranks);
int rdma_rank_begin_idx = rank_begin_idx / LEGACY_NUM_MAX_NVL_PEERS, rdma_rank_end_idx = rank_end_idx / LEGACY_NUM_MAX_NVL_PEERS;
if (rank_begin_idx < rank_end_idx) {
const auto num_expert_per_rank = num_experts / num_ranks;
auto expert_begin = rank_begin_idx * num_expert_per_rank;
auto expert_end = rank_end_idx * num_expert_per_rank;
// Per-thread count
#pragma unroll
for (int i = 0; i < kNumRanksPerSM; ++i)
num_tokens_per_rank_per_thread[thread_id][i] = 0;
#pragma unroll
for (int i = 0; i < kNumRDMARanksPerSM; ++i)
num_tokens_per_rdma_rank_per_thread[thread_id][i] = 0;
#pragma unroll
for (int i = thread_id; i < num_tokens; i += kNumThreads) {
auto shifted_topk_idx = topk_idx + i * num_topk;
int is_in_rank[kNumRanksPerSM] = {0}, is_in_rdma_rank[kNumRDMARanksPerSM] = {0};
#pragma unroll
for (int j = 0, expert_idx, rank_idx; j < num_topk; ++j) {
expert_idx = static_cast<int>(shifted_topk_idx[j]);
if (expert_begin <= expert_idx and expert_idx < expert_end) {
// Count single rank
rank_idx = expert_idx / num_expert_per_rank - rank_begin_idx;
is_in_rank[rank_idx]++, is_in_rdma_rank[rank_idx / LEGACY_NUM_MAX_NVL_PEERS]++;
}
}
auto shifted_is_token_in_rank = is_token_in_rank + i * num_ranks;
#pragma unroll
for (int j = 0; j + rank_begin_idx < rank_end_idx; ++j) {
shifted_is_token_in_rank[j + rank_begin_idx] = (is_in_rank[j] > 0);
num_tokens_per_rank_per_thread[thread_id][j] += (is_in_rank[j] > 0);
}
#pragma unroll
for (int j = 0; j + rdma_rank_begin_idx < rdma_rank_end_idx; ++j)
num_tokens_per_rdma_rank_per_thread[thread_id][j] += (is_in_rdma_rank[j] > 0);
}
__syncthreads();
// Sum up
EP_STATIC_ASSERT(kNumRanksPerSM <= kNumThreads, "Too many ranks per SM");
if (rank_begin_idx + thread_id < rank_end_idx) {
int sum = 0;
#pragma unroll
for (int i = 0; i < kNumThreads; ++i)
sum += num_tokens_per_rank_per_thread[i][thread_id];
num_tokens_per_rank[rank_begin_idx + thread_id] = sum;
}
if (num_tokens_per_rdma_rank != nullptr and rdma_rank_begin_idx + thread_id < rdma_rank_end_idx) {
int sum = 0;
#pragma unroll
for (int i = 0; i < kNumThreads; ++i)
sum += num_tokens_per_rdma_rank_per_thread[i][thread_id];
num_tokens_per_rdma_rank[rdma_rank_begin_idx + thread_id] = sum;
}
}
}
void get_dispatch_layout(const topk_idx_t* topk_idx,
int* num_tokens_per_rank,
int* num_tokens_per_rdma_rank,
int* num_tokens_per_expert,
bool* is_token_in_rank,
int num_tokens,
int num_topk,
int num_ranks,
int num_experts,
cudaStream_t stream) {
constexpr int kNumThreads = 256, kNumExpertsPerSM = 4, kNumRanksPerSM = 8;
int num_sms = ((num_experts + kNumExpertsPerSM - 1) / kNumExpertsPerSM) + (num_ranks + kNumRanksPerSM - 1) / kNumRanksPerSM;
EP_STATIC_ASSERT(kNumRanksPerSM % LEGACY_NUM_MAX_NVL_PEERS == 0, "Invalid number of ranks per SM");
SETUP_LAUNCH_CONFIG(num_sms, kNumThreads, stream);
LAUNCH_KERNEL(&cfg,
(get_dispatch_layout<kNumThreads, kNumExpertsPerSM, kNumRanksPerSM>),
topk_idx,
num_tokens_per_rank,
num_tokens_per_rdma_rank,
num_tokens_per_expert,
is_token_in_rank,
num_tokens,
num_topk,
num_ranks,
num_experts);
}
} // namespace layout
} // namespace deep_ep::legacy
csrc/kernels/legacy/utils.cuh
+650
View File
@@ -0,0 +1,650 @@
#pragma once
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/exception.cuh>
#define UNROLLED_WARP_COPY(UNROLL_FACTOR, LANE_ID, N, DST, SRC, LD_FUNC, ST_FUNC) \
{ \
constexpr int kLoopStride = 32 * (UNROLL_FACTOR); \
typename std::remove_reference<decltype(LD_FUNC((SRC) + 0))>::type unrolled_values[(UNROLL_FACTOR)]; \
auto __src = (SRC); \
auto __dst = (DST); \
for (int __i = (LANE_ID); __i < ((N) / kLoopStride) * kLoopStride; __i += kLoopStride) { \
_Pragma("unroll") for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) unrolled_values[__j] = LD_FUNC(__src + __i + __j * 32); \
_Pragma("unroll") for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) ST_FUNC(__dst + __i + __j * 32, unrolled_values[__j]); \
} \
{ \
int __i = ((N) / kLoopStride) * kLoopStride + (LANE_ID); \
_Pragma("unroll") for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) { \
if (__i + __j * 32 < (N)) { \
unrolled_values[__j] = LD_FUNC(__src + __i + __j * 32); \
} \
} \
_Pragma("unroll") for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) { \
if (__i + __j * 32 < (N)) { \
ST_FUNC(__dst + __i + __j * 32, unrolled_values[__j]); \
} \
} \
} \
}
namespace deep_ep::legacy {
template <int kBytes>
struct VecInt {};
template <>
struct VecInt<1> {
using vec_t = int8_t;
};
template <>
struct VecInt<2> {
using vec_t = int16_t;
};
template <>
struct VecInt<4> {
using vec_t = int;
};
template <>
struct VecInt<8> {
using vec_t = int64_t;
};
template <>
struct VecInt<16> {
using vec_t = int4;
};
template <typename FuncT>
struct PatternVisitor {
FuncT func;
__device__ __host__ explicit PatternVisitor(FuncT&& func) : func(std::forward<FuncT>(func)) {}
__device__ __host__ auto operator[](const uint32_t& i) { return func(i); }
};
__device__ __forceinline__ void trap() {
asm("trap;");
}
__device__ __forceinline__ void memory_fence() {
asm volatile("fence.acq_rel.sys;" ::: "memory");
}
__device__ __forceinline__ void memory_fence_gpu() {
asm volatile("fence.acq_rel.gpu;" ::: "memory");
}
__device__ __forceinline__ void memory_fence_cta() {
asm volatile("fence.acq_rel.cta;" ::: "memory");
}
__device__ __forceinline__ void st_relaxed_sys_global(const int* ptr, int val) {
asm volatile("st.relaxed.sys.global.s32 [%0], %1;" ::"l"(ptr), "r"(val) : "memory");
}
__device__ __forceinline__ void st_release_sys_global(const int* ptr, int val) {
asm volatile("st.release.sys.global.s32 [%0], %1;" ::"l"(ptr), "r"(val) : "memory");
}
__device__ __forceinline__ void st_release_cta(const int* ptr, int val) {
asm volatile("st.release.cta.s32 [%0], %1;" ::"l"(ptr), "r"(val) : "memory");
}
__device__ __forceinline__ int ld_acquire_sys_global(const int* ptr) {
int ret;
asm volatile("ld.acquire.sys.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
return ret;
}
__device__ __forceinline__ uint64_t ld_acquire_sys_global(const uint64_t* ptr) {
uint64_t ret;
asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(ret) : "l"(ptr));
return ret;
}
__device__ __forceinline__ int ld_acquire_global(const int* ptr) {
int ret;
asm volatile("ld.acquire.gpu.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
return ret;
}
__device__ __forceinline__ int atomic_add_release_sys_global(const int* ptr, int value) {
int ret;
asm volatile("atom.add.release.sys.global.s32 %0, [%1], %2;" : "=r"(ret) : "l"(ptr), "r"(value));
return ret;
}
__device__ __forceinline__ int atomic_add_release_global(const int* ptr, int value) {
int ret;
asm volatile("atom.add.release.gpu.global.s32 %0, [%1], %2;" : "=r"(ret) : "l"(ptr), "r"(value));
return ret;
}
__device__ __forceinline__ int ld_acquire_cta(const int* ptr) {
int ret;
asm volatile("ld.acquire.cta.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
return ret;
}
__device__ __forceinline__ uint8_t ld_na_relaxed(const uint8_t* ptr) {
uint16_t ret;
asm volatile("ld.relaxed.gpu.global.L1::no_allocate.b8 %0, [%1];" : "=h"(ret) : "l"(ptr));
return static_cast<uint8_t>(ret);
}
__device__ __forceinline__ uint16_t ld_na_relaxed(const uint16_t* ptr) {
uint16_t ret;
asm volatile("ld.relaxed.gpu.global.L1::no_allocate.b16 %0, [%1];" : "=h"(ret) : "l"(ptr));
return ret;
}
__device__ __forceinline__ uint32_t ld_na_relaxed(const uint32_t* ptr) {
uint32_t ret;
asm volatile("ld.relaxed.gpu.global.L1::no_allocate.b32 %0, [%1];" : "=r"(ret) : "l"(ptr));
return ret;
}
__device__ __forceinline__ uint64_t ld_na_relaxed(const uint64_t* ptr) {
uint64_t ret;
asm volatile("ld.relaxed.gpu.global.L1::no_allocate.b64 %0, [%1];" : "=l"(ret) : "l"(ptr));
return ret;
}
__device__ __forceinline__ int ld_volatile_global(const int* ptr) {
int ret;
asm volatile("ld.volatile.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
return ret;
}
__device__ __forceinline__ float ld_volatile_global(const float* ptr) {
float ret;
asm volatile("ld.volatile.global.f32 %0, [%1];" : "=f"(ret) : "l"(ptr));
return ret;
}
__device__ __forceinline__ int64_t ld_volatile_global(const int64_t* ptr) {
int64_t ret;
asm volatile("ld.volatile.global.s64 %0, [%1];" : "=l"(ret) : "l"(ptr));
return ret;
}
__device__ __forceinline__ int64_t ld_volatile_global(const uint64_t* ptr) {
int64_t ret;
asm volatile("ld.volatile.global.u64 %0, [%1];" : "=l"(ret) : "l"(ptr));
return ret;
}
#ifndef DISABLE_AGGRESSIVE_PTX_INSTRS
#define LD_NC_FUNC "ld.global.nc.L1::no_allocate.L2::256B"
#else
#define LD_NC_FUNC "ld.volatile.global"
#endif
// `ld.global.nc.L1::no_allocate` will be translated into `LDG.E.NA.[width].CONSTANT` in SASS
template <typename dtype_t>
__device__ __forceinline__ dtype_t ld_nc_global(const dtype_t* ptr) {
auto ret = ld_nc_global(reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t*>(ptr));
return *reinterpret_cast<dtype_t*>(&ret);
}
template <>
__device__ __forceinline__ uint8_t ld_nc_global(const uint8_t* ptr) {
uint16_t ret;
// NOTES: we must use `uint16_t` as inline ASM does not support 8-bit constraint letter (`h` below means unsigned 16-bit)
asm volatile(LD_NC_FUNC ".u8 %0, [%1];" : "=h"(ret) : "l"(ptr));
return static_cast<uint8_t>(ret);
}
template <>
__device__ __forceinline__ int ld_nc_global(const int* ptr) {
int ret;
asm volatile(LD_NC_FUNC ".s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
return ret;
}
template <>
__device__ __forceinline__ int64_t ld_nc_global(const int64_t* ptr) {
int64_t ret;
asm volatile(LD_NC_FUNC ".s64 %0, [%1];" : "=l"(ret) : "l"(ptr));
return ret;
}
template <>
__device__ __forceinline__ float ld_nc_global(const float* ptr) {
float ret;
asm volatile(LD_NC_FUNC ".f32 %0, [%1];" : "=f"(ret) : "l"(ptr));
return ret;
}
template <>
__device__ __forceinline__ int2 ld_nc_global(const int2* ptr) {
int2 ret;
asm volatile(LD_NC_FUNC ".v2.s32 {%0, %1}, [%2];" : "=r"(ret.x), "=r"(ret.y) : "l"(ptr));
return ret;
}
template <>
__device__ __forceinline__ int4 ld_nc_global(const int4* ptr) {
int4 ret;
asm volatile(LD_NC_FUNC ".v4.s32 {%0, %1, %2, %3}, [%4];" : "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w) : "l"(ptr));
return ret;
}
__device__ __forceinline__ void st_na_relaxed(const uint8_t* ptr, uint8_t val) {
asm volatile("st.relaxed.gpu.global.L1::no_allocate.b8 [%0], %1;" : : "l"(ptr), "h"(static_cast<uint16_t>(val)));
}
__device__ __forceinline__ void st_na_relaxed(const uint16_t* ptr, uint16_t val) {
asm volatile("st.relaxed.gpu.global.L1::no_allocate.b16 [%0], %1;" : : "l"(ptr), "h"(val));
}
__device__ __forceinline__ void st_na_relaxed(const uint32_t* ptr, uint32_t val) {
asm volatile("st.relaxed.gpu.global.L1::no_allocate.b32 [%0], %1;" : : "l"(ptr), "r"(val));
}
__device__ __forceinline__ void st_na_relaxed(const int* ptr, int val) {
asm volatile("st.relaxed.gpu.global.L1::no_allocate.b32 [%0], %1;" : : "l"(ptr), "r"(val));
}
__device__ __forceinline__ void st_na_relaxed(const int4* ptr, int4 val) {
asm volatile("st.relaxed.gpu.global.L1::no_allocate.v4.s32 [%0], {%1, %2, %3, %4};"
:
: "l"(ptr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w));
}
__device__ __forceinline__ void st_na_release(const int* ptr, int val) {
asm volatile("st.release.gpu.global.L1::no_allocate.b32 [%0], %1;" : : "l"(ptr), "r"(val));
}
__device__ __forceinline__ void st_na_release(const uint32_t* ptr, uint32_t val) {
asm volatile("st.release.gpu.global.L1::no_allocate.b32 [%0], %1;" : : "l"(ptr), "r"(val));
}
__device__ __forceinline__ void st_na_release(const uint64_t* ptr, uint64_t val) {
asm volatile("st.release.gpu.global.L1::no_allocate.b64 [%0], %1;" : : "l"(ptr), "l"(val));
}
// `st.global.L1::no_allocate` will be translated into `ST.E.NA.[width]` in SASS
#ifndef DISABLE_AGGRESSIVE_PTX_INSTRS
#define ST_NA_FUNC "st.global.L1::no_allocate"
#else
#define ST_NA_FUNC "st.global"
#endif
template <typename dtype_t>
__device__ __forceinline__ void st_na_global(const dtype_t* ptr, const dtype_t& value) {
st_na_global(reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t*>(ptr),
*reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t*>(&value));
}
template <>
__device__ __forceinline__ void st_na_global(const int* ptr, const int& value) {
asm volatile(ST_NA_FUNC ".s32 [%0], %1;" ::"l"(ptr), "r"(value));
}
template <>
__device__ __forceinline__ void st_na_global(const int64_t* ptr, const int64_t& value) {
asm volatile(ST_NA_FUNC ".s64 [%0], %1;" ::"l"(ptr), "l"(value));
}
template <>
__device__ __forceinline__ void st_na_global(const float* ptr, const float& value) {
asm volatile(ST_NA_FUNC ".f32 [%0], %1;" ::"l"(ptr), "f"(value));
}
template <>
__device__ __forceinline__ void st_na_global(const int4* ptr, const int4& value) {
asm volatile(ST_NA_FUNC ".v4.s32 [%0], {%1, %2, %3, %4};" ::"l"(ptr), "r"(value.x), "r"(value.y), "r"(value.z), "r"(value.w));
}
__device__ __forceinline__ float log2f_approx(const float& x) {
float ret;
asm volatile("lg2.approx.f32 %0, %1;" : "=f"(ret) : "f"(x));
return ret;
}
__device__ __forceinline__ float exp2f_approx(const float& x) {
float ret;
asm volatile("ex2.approx.f32 %0, %1;" : "=f"(ret) : "f"(x));
return ret;
}
__forceinline__ __device__ int get_lane_id() {
int lane_id;
asm("mov.s32 %0, %laneid;" : "=r"(lane_id));
return lane_id;
}
__device__ __forceinline__ uint32_t elect_one_sync() {
#ifndef DISABLE_SM90_FEATURES
uint32_t pred = 0;
asm volatile(
"{\n"
".reg .b32 %%rx;\n"
".reg .pred %%px;\n"
" elect.sync %%rx|%%px, %1;\n"
"@%%px mov.s32 %0, 1;\n"
"}\n"
: "+r"(pred)
: "r"(0xffffffff));
return pred;
#else
return get_lane_id() == 0;
#endif
}
// TMA PTX instructions
#ifndef DISABLE_SM90_FEATURES
__device__ __forceinline__ void fence_barrier_init() {
asm volatile("fence.mbarrier_init.release.cluster; \n" ::);
}
__device__ __forceinline__ void mbarrier_init(uint64_t* mbar_ptr, uint32_t arrive_count) {
auto mbar_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_ptr));
asm volatile("mbarrier.init.shared::cta.b64 [%1], %0;" ::"r"(arrive_count), "r"(mbar_int_ptr));
}
__device__ __forceinline__ void mbarrier_inval(uint64_t* mbar_ptr) {
auto mbar_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_ptr));
asm volatile("mbarrier.inval.shared::cta.b64 [%0];" ::"r"(mbar_int_ptr));
}
template <bool kWithMultiStages = false>
__device__ __forceinline__ void mbarrier_wait(uint64_t* mbar_ptr, uint32_t& phase, int stage_idx = 0) {
auto mbar_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_ptr));
const auto wait = kWithMultiStages ? (phase >> stage_idx) & 1 : phase;
asm volatile(
"{\n\t"
".reg .pred P1; \n\t"
"LAB_WAIT: \n\t"
"mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2; \n\t"
"@P1 bra DONE; \n\t"
"bra LAB_WAIT; \n\t"
"DONE: \n\t"
"}" ::"r"(mbar_int_ptr),
"r"(wait),
"r"(0x989680));
phase ^= kWithMultiStages ? (1 << stage_idx) : 1;
}
__device__ __forceinline__ void mbarrier_arrive_and_expect_tx(uint64_t* mbar_ptr, int num_bytes) {
auto mbar_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_ptr));
asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%1], %0; \n\t" ::"r"(num_bytes), "r"(mbar_int_ptr));
}
__device__ __forceinline__ void mbarrier_arrive(uint64_t* mbar_ptr) {
auto mbar_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_ptr));
asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0]; \n\t" ::"r"(mbar_int_ptr));
}
__device__ __forceinline__ void tma_store_fence() {
asm volatile("fence.proxy.async.shared::cta;");
}
__device__ __forceinline__ void fence_view_async_shared() {
asm volatile (
"{\n\t"
"fence.proxy.async.shared::cta; \n"
"}"
::
: "memory");
}
constexpr uint64_t kEvictFirst = 0x12f0000000000000;
constexpr uint64_t kEvictNormal = 0x1000000000000000;
__device__ __forceinline__ void tma_load_1d(
const void* smem_ptr, const void* gmem_ptr, uint64_t* mbar_ptr, int num_bytes, bool evict_first = true) {
auto mbar_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_ptr));
auto smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
const auto cache_hint = evict_first ? kEvictFirst : kEvictNormal;
asm volatile(
"cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;\n" ::"r"(smem_int_ptr),
"l"(gmem_ptr),
"r"(num_bytes),
"r"(mbar_int_ptr),
"l"(cache_hint)
: "memory");
}
__device__ __forceinline__ void tma_store_1d(const void* smem_ptr, const void* gmem_ptr, int num_bytes, bool evict_first = true) {
auto smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
const auto cache_hint = evict_first ? kEvictFirst : kEvictNormal;
asm volatile("cp.async.bulk.global.shared::cta.bulk_group.L2::cache_hint [%0], [%1], %2, %3;\n" ::"l"(gmem_ptr),
"r"(smem_int_ptr),
"r"(num_bytes),
"l"(cache_hint)
: "memory");
asm volatile("cp.async.bulk.commit_group;");
}
template <int N>
__device__ __forceinline__ void tma_store_wait() {
asm volatile("cp.async.bulk.wait_group %0;" ::"n"(N) : "memory");
}
#endif
template <typename dtype_t>
__host__ __device__ constexpr dtype_t ceil_div(dtype_t a, dtype_t b) {
return (a + b - 1) / b;
}
template <typename dtype_t>
__host__ __device__ constexpr dtype_t align_up(dtype_t a, dtype_t b) {
return ceil_div<dtype_t>(a, b) * b;
}
template <typename dtype_t>
__host__ __device__ constexpr dtype_t align_down(dtype_t a, dtype_t b) {
return a / b * b;
}
__forceinline__ __device__ void get_channel_task_range(int num_tokens, int num_sms, int sm_id, int& token_start_idx, int& token_end_idx) {
int num_tokens_per_sm = ceil_div(num_tokens, num_sms);
token_start_idx = min(num_tokens_per_sm * sm_id, num_tokens);
token_end_idx = min(token_start_idx + num_tokens_per_sm, num_tokens);
}
template <typename dtype_a_t, typename dtype_b_t>
__device__ __forceinline__ dtype_b_t pack2(const dtype_a_t& x, const dtype_a_t& y) {
EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), "Invalid dtypes");
dtype_b_t packed;
auto unpacked_ptr = reinterpret_cast<dtype_a_t*>(&packed);
unpacked_ptr[0] = x, unpacked_ptr[1] = y;
return packed;
}
template <typename dtype_a_t, typename dtype_b_t>
__device__ __forceinline__ void unpack2(const dtype_b_t& packed, dtype_a_t& x, dtype_a_t& y) {
EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), "Invalid dtypes");
auto unpacked_ptr = reinterpret_cast<const dtype_a_t*>(&packed);
x = unpacked_ptr[0], y = unpacked_ptr[1];
}
template <typename dtype_t>
__device__ __forceinline__ dtype_t broadcast(dtype_t& ptr, int src_lane_idx) {
EP_STATIC_ASSERT(sizeof(dtype_t) % sizeof(int) == 0, "");
auto send_int_values = reinterpret_cast<int*>(&ptr);
int recv_int_values[sizeof(dtype_t) / sizeof(int)];
#pragma unroll
for (int i = 0; i < sizeof(dtype_t) / sizeof(int); ++i)
recv_int_values[i] = __shfl_sync(0xffffffff, send_int_values[i], src_lane_idx);
return *reinterpret_cast<dtype_t*>(recv_int_values);
}
constexpr float kFP8Margin = 1e-4;
constexpr float kFinfoAmaxE4M3 = 448.0f;
constexpr float kFinfoAmaxInvE4M3 = 1 / 448.0f;
__forceinline__ __device__ float fast_pow2(int x) {
// We can ensure `-126 <= x and x <= 127`
uint32_t bits_x = (x + 127) << 23;
return *reinterpret_cast<float*>(&bits_x);
}
__forceinline__ __device__ int fast_log2_ceil(float x) {
auto bits_x = *reinterpret_cast<uint32_t*>(&x);
auto exp_x = (bits_x >> 23) & 0xff;
auto man_bits = bits_x & ((1 << 23) - 1);
return exp_x - 127 + (man_bits != 0);
}
__forceinline__ __device__ void calculate_fp8_scales(float amax, float& scale, float& scale_inv, bool round_scale) {
if (round_scale) {
auto exp_scale_inv = fast_log2_ceil(amax * kFinfoAmaxInvE4M3);
scale = fast_pow2(-exp_scale_inv);
scale_inv = fast_pow2(exp_scale_inv);
} else {
scale_inv = amax * kFinfoAmaxInvE4M3;
scale = kFinfoAmaxE4M3 / amax;
}
}
template <bool kIsUE8M0, typename out_dtype_t = std::conditional_t<kIsUE8M0, uint8_t, float>>
__forceinline__ __device__ out_dtype_t extract_required_scale_format(float value) {
if constexpr (kIsUE8M0) {
return static_cast<uint8_t>((*reinterpret_cast<uint32_t*>(&value)) >> 23);
} else {
return value;
}
}
template <int kNumRanks, bool kSyncOnly = false>
__forceinline__ __device__ void barrier_block(int** barrier_signal_ptrs, int rank) {
auto thread_id = static_cast<int>(threadIdx.x);
// For non-sync-only cases, the memory operations by other threads in the block must be visible to the `sys` scope
if constexpr (not kSyncOnly) {
memory_fence();
__syncthreads();
}
// Add self-ranks, sub other ranks
if (thread_id < kNumRanks) {
atomicAdd_system(barrier_signal_ptrs[rank] + thread_id, LEGACY_FINISHED_SUM_TAG);
atomicSub_system(barrier_signal_ptrs[thread_id] + rank, LEGACY_FINISHED_SUM_TAG);
}
EP_DEVICE_ASSERT(kNumRanks <= blockDim.x);
// Check timeout
auto start_time = clock64();
while (true) {
auto value = thread_id < kNumRanks ? ld_volatile_global(barrier_signal_ptrs[rank] + thread_id) : 0;
if (__all_sync(0xffffffff, value <= 0))
break;
if (clock64() - start_time > LEGACY_NUM_TIMEOUT_CYCLES and thread_id < kNumRanks) {
printf("DeepEP timeout check failed: rank = %d, thread = %d, value = %d)\n", rank, thread_id, value);
trap();
}
}
__syncthreads();
}
__forceinline__ __device__ int atomic_cas_cta_acquire(int* addr, int x, int y) {
int ret;
asm volatile("atom.acquire.cta.shared::cta.cas.b32 %0, [%1], %2, %3;" : "=r"(ret) : "l"(addr), "r"(x), "r"(y) : "memory");
return ret;
}
__forceinline__ __device__ int atomic_exch_cta_release(int* addr, int x) {
int ret;
asm volatile("atom.release.cta.shared::cta.exch.b32 %0, [%1], %2;" : "=r"(ret) : "l"(addr), "r"(x) : "memory");
return ret;
}
__forceinline__ __device__ void acquire_lock(int* mutex) {
// To make later memory operations valid, we must use `acquire` for memory semantics
while (atomic_cas_cta_acquire(mutex, 0, 1) != 0)
;
}
__forceinline__ __device__ void release_lock(int* mutex) {
// To make previous memory operations visible to other threads, we must use `release` for memory semantics
atomic_exch_cta_release(mutex, 0);
}
// Operation functors
template <typename T>
struct ReduceSum {
__device__ T operator()(T a, T b) const { return a + b; }
};
template <typename T>
struct ReduceMax {
__device__ T operator()(T a, T b) const { return a > b ? a : b; }
};
template <typename T>
struct ReduceMin {
__device__ T operator()(T a, T b) const { return a < b ? a : b; }
};
template <typename T>
struct ReduceAnd {
__device__ T operator()(T a, T b) const { return a & b; }
};
template <typename T>
struct ReduceOr {
__device__ T operator()(T a, T b) const { return a | b; }
};
// Unified reduction function
template <int kNumLanesPerGroup, bool kIntergroupReduce, typename T, typename Op>
__forceinline__ __device__ T warp_reduce(T value, Op op) {
EP_STATIC_ASSERT(kNumLanesPerGroup == 32 or kNumLanesPerGroup == 16 or kNumLanesPerGroup == 8 or kNumLanesPerGroup == 4 or
kNumLanesPerGroup == 2 or kNumLanesPerGroup == 1,
"Invalid number of lanes");
constexpr uint32_t mask = 0xffffffff;
if constexpr (kIntergroupReduce) {
if constexpr (kNumLanesPerGroup <= 1)
value = op(value, __shfl_xor_sync(mask, value, 1));
if constexpr (kNumLanesPerGroup <= 2)
value = op(value, __shfl_xor_sync(mask, value, 2));
if constexpr (kNumLanesPerGroup <= 4)
value = op(value, __shfl_xor_sync(mask, value, 4));
if constexpr (kNumLanesPerGroup <= 8)
value = op(value, __shfl_xor_sync(mask, value, 8));
if constexpr (kNumLanesPerGroup <= 16)
value = op(value, __shfl_xor_sync(mask, value, 16));
} else {
if constexpr (kNumLanesPerGroup >= 32)
value = op(value, __shfl_xor_sync(mask, value, 16));
if constexpr (kNumLanesPerGroup >= 16)
value = op(value, __shfl_xor_sync(mask, value, 8));
if constexpr (kNumLanesPerGroup >= 8)
value = op(value, __shfl_xor_sync(mask, value, 4));
if constexpr (kNumLanesPerGroup >= 4)
value = op(value, __shfl_xor_sync(mask, value, 2));
if constexpr (kNumLanesPerGroup >= 2)
value = op(value, __shfl_xor_sync(mask, value, 1));
}
return value;
}
// Convenience aliases
template <int kNumLanesPerGroup = 32, bool kIntergroupReduce = false, typename T>
__forceinline__ __device__ T warp_reduce_sum(T value) {
return warp_reduce<kNumLanesPerGroup, kIntergroupReduce, T>(value, ReduceSum<T>{});
}
template <int kNumLanesPerGroup = 32, bool kIntergroupReduce = false, typename T>
__forceinline__ __device__ T warp_reduce_max(T value) {
return warp_reduce<kNumLanesPerGroup, kIntergroupReduce, T>(value, ReduceMax<T>{});
}
template <int kNumLanesPerGroup = 32, bool kIntergroupReduce = false, typename T>
__forceinline__ __device__ T warp_reduce_min(T value) {
return warp_reduce<kNumLanesPerGroup, kIntergroupReduce, T>(value, ReduceMin<T>{});
}
template <int kNumLanesPerGroup = 32, bool kIntergroupReduce = false, typename T>
__forceinline__ __device__ T warp_reduce_and(T value) {
return warp_reduce<kNumLanesPerGroup, kIntergroupReduce, T>(value, ReduceAnd<T>{});
}
template <int kNumLanesPerGroup = 32, bool kIntergroupReduce = false, typename T>
__forceinline__ __device__ T warp_reduce_or(T value) {
return warp_reduce<kNumLanesPerGroup, kIntergroupReduce, T>(value, ReduceOr<T>{});
}
} // namespace deep_ep::legacy
csrc/legacy/buffer.hpp
+1794
View File
File diff suppressed because it is too large Load Diff
csrc/legacy/config.hpp
+190
View File
@@ -0,0 +1,190 @@
#pragma once
#include <deep_ep/common/exception.cuh>
#include "../kernels/legacy/api.cuh"
namespace deep_ep::legacy {
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 < LEGACY_NUM_MAX_NVL_PEERS or num_ranks % LEGACY_NUM_MAX_NVL_PEERS == 0);
EP_HOST_ASSERT(num_ranks <= LEGACY_NUM_MAX_NVL_PEERS or num_sms % 2 == 0);
const auto num_rdma_ranks = std::max(num_ranks / LEGACY_NUM_MAX_NVL_PEERS, 1);
const auto num_nvl_ranks = std::min(num_ranks, LEGACY_NUM_MAX_NVL_PEERS);
const int num_channels = num_sms / 2;
size_t num_bytes = 0;
num_bytes += num_channels * num_nvl_ranks * (2 * num_rdma_ranks + 3) * sizeof(int);
num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * hidden_bytes;
num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * internode::get_source_meta_bytes();
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 {
// Legacy mode
if (num_ranks <= LEGACY_NUM_MAX_NVL_PEERS)
return 0;
// Below are some assumptions
// TODO: add assertions
constexpr int kNumMaxTopK = 128;
constexpr int kNumMaxScales = 128;
EP_HOST_ASSERT(num_ranks % LEGACY_NUM_MAX_NVL_PEERS == 0);
EP_HOST_ASSERT(num_sms % 2 == 0);
const int num_rdma_ranks = num_ranks / LEGACY_NUM_MAX_NVL_PEERS;
const int num_channels = num_sms / 2;
size_t num_bytes = 0;
num_bytes += num_channels * num_rdma_ranks * (LEGACY_NUM_MAX_NVL_PEERS * 2 + 2) * 2 * sizeof(int);
num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * hidden_bytes * 2;
num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * internode::get_source_meta_bytes() * 2;
num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * kNumMaxTopK * sizeof(topk_idx_t) * 2;
num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * kNumMaxTopK * sizeof(float) * 2;
num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * kNumMaxScales * sizeof(float) * 2;
num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * sizeof(int4) * 2;
num_bytes = ((num_bytes + 127) / 128) * 128;
return num_bytes;
}
};
struct 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 + LEGACY_NUM_BUFFER_ALIGNMENT_BYTES) / LEGACY_NUM_BUFFER_ALIGNMENT_BYTES) * LEGACY_NUM_BUFFER_ALIGNMENT_BYTES;
}
} // namespace deep_ep
csrc/python_api.cpp
+39
View File
@@ -0,0 +1,39 @@
#include <pybind11/pybind11.h>
#include <torch/python.h>
#include <deep_ep/common/compiled.cuh>
#include "jit/api.hpp"
#include "elastic/buffer.hpp"
#include "legacy/buffer.hpp"
#ifndef TORCH_EXTENSION_NAME
#define TORCH_EXTENSION_NAME _C
#endif
bool is_sm90_compiled() {
#ifndef DISABLE_SM90_FEATURES
return true;
#else
return false;
#endif
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.doc() = "DeepEP: an efficient expert-parallel communication library";
// Whether support FP8 and TMA features
m.def("is_sm90_compiled", []() { return deep_ep::kEnableSM90Features; });
// The integer type of top-k indices
m.attr("topk_idx_t") = py::cast(c10::CppTypeToScalarType<deep_ep::topk_idx_t>::value);
// JIT API
deep_ep::jit::register_apis(m);
// Register legacy buffer APIs
deep_ep::legacy::register_apis(m);
// Register elastic buffer (DeepEP V2) APIs
deep_ep::elastic::register_apis(m);
}
csrc/utils/event.hpp
+44
View File
@@ -0,0 +1,44 @@
#pragma once
#include <ATen/cuda/CUDAContext.h>
#include <memory>
#include <deep_ep/common/exception.cuh>
namespace deep_ep {
struct EventHandle {
std::shared_ptr<torch::Event> event;
std::vector<std::optional<torch::Tensor>> tensors_to_record;
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); }
};
static torch::Event create_event(const at::cuda::CUDAStream& s) {
auto event = torch::Event(torch::kCUDA);
event.record(s);
return event;
}
static void stream_wait(const at::cuda::CUDAStream& s_0, const at::cuda::CUDAStream& s_1) {
EP_HOST_ASSERT(s_0.id() != s_1.id());
s_0.unwrap().wait(create_event(s_1));
}
static void stream_wait(const at::cuda::CUDAStream& s, const EventHandle& event) {
s.unwrap().wait(*event.event);
}
} // namespace deep_ep
csrc/utils/format.hpp
+6
View File
@@ -0,0 +1,6 @@
#pragma once
// Just a wrapper for the `fmt` headers
#define FMT_HEADER_ONLY
#include <fmt/base.h>
#include <fmt/format.h>
csrc/utils/hash.hpp
+40
View File
@@ -0,0 +1,40 @@
#pragma once
#include <string>
#include <cstdint>
namespace deep_ep {
static uint64_t fnv1a(const std::vector<char>& data, const uint64_t& seed) {
uint64_t h = seed;
constexpr uint64_t prime = 0x100000001b3ull;
for (const char& c: data) {
h ^= static_cast<uint8_t>(c);
h *= prime;
}
return h;
}
static std::string get_hex_digest(const std::vector<char>& data) {
const auto state_0 = fnv1a(data, 0xc6a4a7935bd1e995ull);
const auto state_1 = fnv1a(data, 0x9e3779b97f4a7c15ull);
// Split-mix 64
const auto split_mix = [](uint64_t z) {
z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ull;
z = (z ^ (z >> 27)) * 0x94d049bb133111ebull;
return z ^ (z >> 31);
};
std::ostringstream oss;
oss << std::hex << std::setfill('0')
<< std::setw(16) << split_mix(state_0)
<< std::setw(16) << split_mix(state_1);
return oss.str();
}
static std::string get_hex_digest(const std::string& data) {
return get_hex_digest(std::vector<char>{data.begin(), data.end()});
}
} // namespace deep_ep
csrc/utils/lazy_driver.hpp
+57
View File
@@ -0,0 +1,57 @@
#pragma once
#include <cuda.h>
#include <dlfcn.h>
#include <deep_ep/common/exception.cuh>
namespace deep_ep {
// Lazy loading all driver symbols
static void* get_driver_handle() {
static void* handle = nullptr;
if (handle == nullptr) {
handle = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL);
EP_HOST_ASSERT(handle != nullptr and "Failed to load CUDA driver `libcuda.so.1`");
}
return handle;
}
// Macro to define wrapper functions named `lazy_cu{API name}`
#define DECL_LAZY_CUDA_DRIVER_FUNCTION(name) \
template <typename... Args> \
static auto lazy_##name(Args&&... args) -> decltype(name(args...)) { \
using FuncType = decltype(&name); \
static FuncType func = nullptr; \
if (func == nullptr) { \
func = reinterpret_cast<FuncType>(dlsym(get_driver_handle(), #name)); \
EP_HOST_ASSERT(func != nullptr and "Failed to load CUDA driver API"); \
} \
return func(std::forward<decltype(args)>(args)...); \
}
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorName);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorString);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuFuncSetAttribute);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleLoad);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleUnload);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleGetFunction);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLaunchKernelEx);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemSetAccess);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemRetainAllocationHandle);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemGetAddressRange_v2);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemAddressReserve);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemAddressFree);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemMap);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemUnmap);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemCreate);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemRelease);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuCtxGetDevice);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemImportFromShareableHandle);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemExportToShareableHandle);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuMemGetAllocationGranularity);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuDeviceGet);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuDeviceGetAttribute);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuStreamBatchMemOp);
} // namespace deep_ep
csrc/utils/lazy_init.hpp
+27
View File
@@ -0,0 +1,27 @@
#pragma once
#include <functional>
#include <memory>
#define EP_DECLARE_STATIC_VAR_IN_CLASS(cls, name) decltype(cls::name) cls::name
namespace deep_ep {
template <typename T>
class LazyInit {
public:
explicit LazyInit(std::function<std::shared_ptr<T>()> factory)
: factory(std::move(factory)) {}
T* operator -> () {
if (ptr == nullptr)
ptr = factory();
return ptr.get();
}
private:
std::shared_ptr<T> ptr;
std::function<std::shared_ptr<T>()> factory;
};
} // namespace deep_ep
csrc/utils/shared_memory.hpp
+127
View File
@@ -0,0 +1,127 @@
#pragma once
#include <cuda.h>
#include <deep_ep/common/exception.cuh>
#include "lazy_driver.hpp"
namespace deep_ep::shared_memory {
union MemHandleInner {
cudaIpcMemHandle_t cuda_ipc_mem_handle;
CUmemFabricHandle cu_mem_fabric_handle;
};
struct MemHandle {
MemHandleInner inner;
size_t size;
};
static void cu_mem_set_access_all(void* ptr, size_t size) {
int device_count;
CUDA_RUNTIME_CHECK(cudaGetDeviceCount(&device_count));
constexpr int kMaxDeviceCount = 8;
EP_HOST_ASSERT(0 < device_count and device_count <= kMaxDeviceCount);
CUmemAccessDesc access_desc[kMaxDeviceCount];
for (int i = 0; i < device_count; ++ i) {
access_desc[i].location.type = CU_MEM_LOCATION_TYPE_DEVICE;
access_desc[i].location.id = i;
access_desc[i].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
}
CUDA_DRIVER_CHECK(lazy_cuMemSetAccess(reinterpret_cast<CUdeviceptr>(ptr), size, access_desc, device_count));
}
static void cu_mem_free(void* ptr) {
CUmemGenericAllocationHandle handle;
CUDA_DRIVER_CHECK(lazy_cuMemRetainAllocationHandle(&handle, ptr));
size_t size = 0;
CUDA_DRIVER_CHECK(lazy_cuMemGetAddressRange_v2(nullptr, &size, reinterpret_cast<CUdeviceptr>(ptr)));
CUDA_DRIVER_CHECK(lazy_cuMemUnmap(reinterpret_cast<CUdeviceptr>(ptr), size));
CUDA_DRIVER_CHECK(lazy_cuMemAddressFree(reinterpret_cast<CUdeviceptr>(ptr), size));
CUDA_DRIVER_CHECK(lazy_cuMemRelease(handle));
}
class SharedMemoryAllocator {
public:
explicit SharedMemoryAllocator(const bool& use_fabric) : use_fabric(use_fabric) {}
void malloc(void** ptr, size_t size) const {
if (use_fabric) {
CUdevice device;
CUDA_DRIVER_CHECK(lazy_cuCtxGetDevice(&device));
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
prop.location.id = device;
size_t alignment = 0;
EP_HOST_ASSERT(size > 0);
CUDA_DRIVER_CHECK(lazy_cuMemGetAllocationGranularity(&alignment, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM));
size = ((size + alignment - 1) / alignment) * alignment;
CUmemGenericAllocationHandle handle;
CUDA_DRIVER_CHECK(lazy_cuMemCreate(&handle, size, &prop, 0));
CUDA_DRIVER_CHECK(lazy_cuMemAddressReserve(reinterpret_cast<CUdeviceptr*>(ptr), size, alignment, 0, 0));
CUDA_DRIVER_CHECK(lazy_cuMemMap(reinterpret_cast<CUdeviceptr>(*ptr), size, 0, handle, 0));
cu_mem_set_access_all(*ptr, size);
} else {
CUDA_RUNTIME_CHECK(cudaMalloc(ptr, size));
}
}
void free(void* ptr) const {
if (use_fabric) {
cu_mem_free(ptr);
} else {
CUDA_RUNTIME_CHECK(cudaFree(ptr));
}
}
void get_mem_handle(MemHandle* mem_handle, void* ptr) const {
size_t size = 0;
CUDA_DRIVER_CHECK(lazy_cuMemGetAddressRange_v2(nullptr, &size, reinterpret_cast<CUdeviceptr>(ptr)));
mem_handle->size = size;
if (use_fabric) {
CUmemGenericAllocationHandle handle;
CUDA_DRIVER_CHECK(lazy_cuMemRetainAllocationHandle(&handle, ptr));
CUDA_DRIVER_CHECK(lazy_cuMemExportToShareableHandle(&mem_handle->inner.cu_mem_fabric_handle, handle, CU_MEM_HANDLE_TYPE_FABRIC, 0));
} else {
CUDA_RUNTIME_CHECK(cudaIpcGetMemHandle(&mem_handle->inner.cuda_ipc_mem_handle, ptr));
}
}
void open_mem_handle(void** ptr, MemHandle* mem_handle) const {
if (use_fabric) {
size_t size = mem_handle->size;
CUmemGenericAllocationHandle handle;
CUDA_DRIVER_CHECK(lazy_cuMemImportFromShareableHandle(&handle, &mem_handle->inner.cu_mem_fabric_handle, CU_MEM_HANDLE_TYPE_FABRIC));
CUDA_DRIVER_CHECK(lazy_cuMemAddressReserve(reinterpret_cast<CUdeviceptr*>(ptr), size, 0, 0, 0));
CUDA_DRIVER_CHECK(lazy_cuMemMap(reinterpret_cast<CUdeviceptr>(*ptr), size, 0, handle, 0));
cu_mem_set_access_all(*ptr, size);
} else {
CUDA_RUNTIME_CHECK(cudaIpcOpenMemHandle(ptr, mem_handle->inner.cuda_ipc_mem_handle, cudaIpcMemLazyEnablePeerAccess));
}
}
void close_mem_handle(void* ptr) const {
if (use_fabric) {
cu_mem_free(ptr);
} else {
CUDA_RUNTIME_CHECK(cudaIpcCloseMemHandle(ptr));
}
}
private:
bool use_fabric;
};
} // namespace deep_ep::shared_memory
csrc/utils/system.hpp
+111
View File
@@ -0,0 +1,111 @@
#pragma once
#include <algorithm>
#include <array>
#include <filesystem>
#include <functional>
#include <random>
#include <string>
#include <memory>
#include <unistd.h>
#include <deep_ep/common/exception.cuh>
#include "format.hpp"
namespace deep_ep {
// ReSharper disable once CppNotAllPathsReturnValue
template <typename dtype_t>
static dtype_t get_env(const std::string& name, const dtype_t& default_value = dtype_t()) {
const auto c_str = std::getenv(name.c_str());
if (c_str == nullptr)
return default_value;
// Read the env and convert to the desired type
if constexpr (std::is_same_v<dtype_t, std::string>) {
return std::string(c_str);
} else if constexpr (std::is_same_v<dtype_t, int>) {
int value;
std::sscanf(c_str, "%d", &value);
return value;
} else {
EP_HOST_ASSERT(false and "Unexpected type");
}
}
static std::tuple<int, std::string> call_external_command(std::string command) {
command = command + " 2>&1";
const auto deleter = [](FILE* f) { if (f) pclose(f); };
std::unique_ptr<FILE, decltype(deleter)> pipe(popen(command.c_str(), "r"), deleter);
EP_HOST_ASSERT(pipe != nullptr);
std::array<char, 512> buffer;
std::string output;
while (fgets(buffer.data(), buffer.size(), pipe.get()))
output += buffer.data();
const auto status = pclose(pipe.release());
// NOTES: if the child was killed by a signal (e.g., SIGINT from Ctrl+C),
// WEXITSTATUS would incorrectly return 0. Treat signal death as failure.
const auto exit_code = WIFEXITED(status) ? WEXITSTATUS(status) : 128 + WTERMSIG(status);
return {exit_code, output};
}
static std::filesystem::path make_dirs(const std::filesystem::path& path) {
// OK if existed
std::error_code capture;
const bool created = std::filesystem::create_directories(path, capture);
if (not (created or capture.value() == 0)) {
EP_HOST_UNREACHABLE(fmt::format("Failed to make directory: {}, created: {}, value: {}",
path.c_str(), created, capture.value()));
}
if (created and get_env<int>("EP_JIT_DEBUG"))
printf("Create directory: %s\n", path.c_str());
return path;
}
static std::string get_uuid() {
static std::random_device rd;
static std::mt19937 gen([]() {
return rd() ^ std::chrono::steady_clock::now().time_since_epoch().count();
}());
static std::uniform_int_distribution<uint32_t> dist;
std::stringstream ss;
ss << getpid() << "-"
<< std::hex << std::setfill('0')
<< std::setw(8) << dist(gen) << "-"
<< std::setw(8) << dist(gen) << "-"
<< std::setw(8) << dist(gen);
return ss.str();
}
static void safe_remove_all(const std::filesystem::path& path) {
std::error_code ec;
if (not std::filesystem::exists(path, ec) or ec)
return;
// A single file
if (not std::filesystem::is_directory(path, ec) or ec) {
std::filesystem::remove(path, ec);
return;
}
// Remove directory
auto it = std::filesystem::directory_iterator(path,
std::filesystem::directory_options::skip_permission_denied, ec);
for (auto end = std::filesystem::directory_iterator(); it != end and not ec;) {
const auto entry_path = it->path();
// Increase firstly to avoid failures
it.increment(ec);
if (ec)
break;
// Recursively clean
safe_remove_all(entry_path);
}
std::filesystem::remove(path, ec);
}
} // deep_ep
deep_ep/__init__.py
+97
View File
@@ -0,0 +1,97 @@
import filecmp
import functools
import glob
import subprocess
import torch
import os
from .utils.find_pkgs import find_nccl_root
# Set some default environment provided at setup
try:
# noinspection PyUnresolvedReferences
from .envs import persistent_envs
for key, value in persistent_envs.items():
if key not in os.environ:
os.environ[key] = value
except ImportError:
pass
# Initialize
@functools.lru_cache()
def find_cuda_home() -> str:
"""
Find the CUDA installation directory, cached.
Returns:
cuda_home: the CUDA installation path.
"""
# TODO: reuse PyTorch API later
# For some PyTorch versions, the original `_find_cuda_home` will initialize CUDA, which is incompatible with process forks
cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
if cuda_home is None:
# noinspection PyBroadException
try:
with open(os.devnull, 'w') as devnull:
nvcc = subprocess.check_output(['which', 'nvcc'], stderr=devnull).decode().rstrip('\r\n')
cuda_home = os.path.dirname(os.path.dirname(nvcc))
except Exception:
cuda_home = '/usr/local/cuda'
if not os.path.exists(cuda_home):
cuda_home = None
assert cuda_home is not None
return cuda_home
def check_nccl_so():
"""
Verify that the NCCL library loaded at runtime matches the linked version.
Aborts if duplicate NCCL libraries are found or if versions mismatch.
"""
if int(os.environ.get('EP_SUPPRESS_NCCL_CHECK', 0)):
return
# PyTorch may load another NCCL library, which is different to the linked one
with open('/proc/self/maps', 'r') as f:
loaded_nccl_so = None
for so in [line.strip().split(' ')[-1] for line in f if 'libnccl' in line]:
loaded_nccl_so = so if loaded_nccl_so is None else loaded_nccl_so
assert so == loaded_nccl_so, f'Duplicate NCCL runtime found in the current system: {so} and {loaded_nccl_so}'
linked_nccl_so_candidates = sorted(glob.glob(f'{find_nccl_root()}/lib/libnccl.so*'))
assert linked_nccl_so_candidates, f'No libnccl.so found in {find_nccl_root()}/lib/'
linked_nccl_so = linked_nccl_so_candidates[0]
# So checking binary-level equalness is necessary
# noinspection PyTypeChecker
assert filecmp.cmp(loaded_nccl_so, linked_nccl_so, shallow=False), \
(f'Invalid NCCL versions: {loaded_nccl_so} (loaded) v.s. {linked_nccl_so} (expected), '
f'please contact Chenggang or Shangyan to upgrade PyTorch NCCL version')
def init_jit():
"""
Initialize the JIT compilation runtime. Sets up CUDA and NCCL root paths for the JIT compiler.
"""
# noinspection PyUnresolvedReferences
import deep_ep._C as _C
library_root_path = os.path.dirname(os.path.abspath(__file__))
_C.init_jit(library_root_path, # Library root directory path
find_cuda_home(), # CUDA home
find_nccl_root()) # NCCL root
# Run initialization
check_nccl_so()
init_jit()
# Import APIs after initialization
from .buffers.legacy import Buffer
from .buffers.elastic import ElasticBuffer, EPHandle
# noinspection PyUnresolvedReferences
from .utils.event import EventOverlap, EventHandle
from .utils.envs import get_physical_domain_size, get_logical_domain_size
# noinspection PyUnresolvedReferences
from deep_ep._C import Config, topk_idx_t
__version__ = '2.1.0'
deep_ep/buffers/__init__.py
View File
deep_ep/buffers/elastic.py
+1107
View File
File diff suppressed because it is too large Load Diff
deep_ep/buffers/legacy.py
+713
View File
@@ -0,0 +1,713 @@
import os
import torch
import torch.distributed as dist
from typing import Callable, List, Tuple, Optional, Union
# noinspection PyUnresolvedReferences
import deep_ep._C as _C
# noinspection PyUnresolvedReferences
from deep_ep._C import Config, EventHandle
from ..utils.event import EventOverlap
from ..utils.envs import check_nvlink_connections, check_torch_deterministic
class Buffer:
"""
The core expert-parallel (EP) communication buffers for Mixture of Experts (MoE) model, which supports:
- high-throughput intranode all-to-all (dispatch and combine, using NVLink)
- high-throughput internode all-to-all (dispatch and combine, using RDMA and NVLink)
- low-latency all-to-all (dispatch and combine, using RDMA)
Attributes:
num_sms: the SMs used in high-throughput kernels.
rank: the local rank number.
group_size: the number of ranks in the group.
group: the communication group.
num_nvl_bytes: the buffer size for intranode NVLink communication.
num_rdma_bytes: the buffer size for internode (also for intranode with low-latency mode) RDMA communication.
runtime: the C++ runtime.
"""
num_sms: int = 20
def __init__(self,
group: Optional[dist.ProcessGroup],
num_nvl_bytes: int = 0,
num_rdma_bytes: int = 0,
low_latency_mode: bool = False,
num_qps_per_rank: int = 24,
allow_nvlink_for_low_latency_mode: bool = True,
allow_mnnvl: bool = False,
explicitly_destroy: bool = False,
enable_shrink: bool = False,
comm: Optional["mpi4py.MPI.Comm"] = None) -> None: # noqa: F821
"""
Initialize the communication buffer.
Arguments:
group: the communication group.
num_nvl_bytes: the buffer size for intranode NVLink communication.
num_rdma_bytes: the buffer size for internode (also for intranode with low-latency mode) RDMA communication.
low_latency_mode: whether to enable low-latency mode.
num_qps_per_rank: the number of QPs for RDMA, the low-latency mode requires that this number equals
to the number of local experts.
allow_nvlink_for_low_latency_mode: whether allow NVLink traffic for low-latency mode, you should notice
this is somehow incompatible with the hook-based overlapping.
Warning: PCIe connections may lead to errors due to memory ordering issues,
please make sure all connections are via NVLink.
allow_mnnvl: whether to allow MNNVL
use_fabric: whether to use fabric API for memory buffers.
enable_shrink: whether to enable shrink mode. The enable mode allocates a mask buffer to support masking ranks dynamically.
explicitly_destroy: If this flag is set to True, you need to explicitly call `destroy()` to release resources;
otherwise, the resources will be released by the destructor.
Note: Releasing resources in the destructor may cause Python's exception handling process to hang.
comm: the `mpi4py.MPI.Comm` communicator to use in case the group parameter is absent.
"""
check_nvlink_connections(group)
# Initialize the CPP runtime
if group is not None:
self.rank = group.rank()
self.group = group
self.group_size = group.size()
def all_gather_object(obj):
object_list = [None] * self.group_size
dist.all_gather_object(object_list, obj, group)
return object_list
elif comm is not None:
self.rank = comm.Get_rank()
self.group = comm
self.group_size = comm.Get_size()
def all_gather_object(obj):
return comm.allgather(obj)
else:
raise ValueError("Either 'group' or 'comm' must be provided.")
self.num_nvl_bytes = num_nvl_bytes
self.num_rdma_bytes = num_rdma_bytes
self.low_latency_mode = low_latency_mode
self.explicitly_destroy = explicitly_destroy
self.enable_shrink = enable_shrink
self.runtime = _C.Buffer(self.rank, self.group_size, num_nvl_bytes, num_rdma_bytes, low_latency_mode,
explicitly_destroy, enable_shrink, allow_mnnvl)
# Synchronize device IDs
local_device_id = self.runtime.get_local_device_id()
device_ids = all_gather_object(local_device_id)
# Synchronize IPC handles
local_ipc_handle = self.runtime.get_local_ipc_handle()
ipc_handles = all_gather_object(local_ipc_handle)
# Synchronize NVSHMEM unique IDs
root_unique_id = None
if self.runtime.get_num_rdma_ranks() > 1 or low_latency_mode:
# Enable IBGDA
assert num_qps_per_rank > 0
os.environ['NVSHMEM_DISABLE_P2P'] = '0' if allow_nvlink_for_low_latency_mode else '1'
os.environ['NVSHMEM_IB_ENABLE_IBGDA'] = '1'
os.environ['NVSHMEM_IBGDA_NUM_RC_PER_PE'] = f'{num_qps_per_rank}'
# Make sure QP depth is always larger than the number of on-flight WRs, so that we can skip WQ slot check
self.nvshmem_qp_depth = int(os.environ.get('NVSHMEM_QP_DEPTH', '1024'))
os.environ['NVSHMEM_QP_DEPTH'] = str(self.nvshmem_qp_depth)
# Reduce gpu memory usage
# 6 default teams + 1 extra team
os.environ['NVSHMEM_MAX_TEAMS'] = '7'
# Disable NVLink SHARP
os.environ['NVSHMEM_DISABLE_NVLS'] = '1'
# NOTES: NVSHMEM initialization requires at least 256 MiB
os.environ['NVSHMEM_CUMEM_GRANULARITY'] = f'{2 ** 29}'
if not allow_mnnvl:
# Disable multi-node NVLink detection
os.environ['NVSHMEM_DISABLE_MNNVL'] = '1'
# Synchronize using the root ID
if (low_latency_mode and self.rank == 0) or (not low_latency_mode and self.runtime.get_rdma_rank() == 0):
root_unique_id = self.runtime.get_local_nvshmem_unique_id()
nvshmem_unique_ids = all_gather_object(root_unique_id)
root_unique_id = nvshmem_unique_ids[0 if low_latency_mode else self.runtime.get_root_rdma_rank(True)]
# Make CPP runtime available
self.runtime.sync(device_ids, ipc_handles, root_unique_id)
assert self.runtime.is_available()
def destroy(self):
"""
Destroy the cpp runtime and release resources.
"""
assert self.explicitly_destroy, '`explicitly_destroy` flag must be set'
self.runtime.destroy()
self.runtime = None
@staticmethod
def is_sm90_compiled():
return _C.is_sm90_compiled()
@staticmethod
def set_num_sms(new_num_sms: int) -> None:
"""
Set the number of SMs to use in high-throughput kernels.
Arguments:
new_num_sms: the new number to be set.
"""
assert new_num_sms % 2 == 0, 'The SM count must be even'
Buffer.num_sms = new_num_sms
@staticmethod
def capture() -> EventOverlap:
"""
Capture a CUDA event on the current stream, i.e. `torch.cuda.current_stream()`.
Returns:
event: the captured event.
"""
return EventOverlap(EventHandle())
@staticmethod
def get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int) -> int:
"""
Get a minimum size requirement for the RDMA buffer. The size calculation will be done with BF16.
Arguments:
num_max_dispatch_tokens_per_rank: the maximum number of tokens to dispatch, all the ranks must hold the same value.
hidden: the hidden dimension of each token.
num_ranks: the number of EP group ranks.
num_experts: the number of all experts.
Returns:
size: the RDMA buffer size recommended.
"""
return _C.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts)
def get_comm_stream(self) -> torch.Stream:
"""
Get the communication stream.
Returns:
stream: the communication stream.
"""
ts: torch.Stream = self.runtime.get_comm_stream()
return torch.cuda.Stream(stream_id=ts.stream_id, device_index=ts.device_index, device_type=ts.device_type)
def get_local_buffer_tensor(self,
dtype: torch.dtype,
size: Optional[torch.Size] = None,
offset: int = 0,
use_rdma_buffer: bool = False) -> torch.Tensor:
"""
Get the raw buffer (slice supported) as a PyTorch tensor.
Argument:
dtype: the data type (PyTorch `dtype`) for the tensor.
size: the slice size (by elements) to get from the buffer.
offset: the offset of the beginning element.
use_rdma_buffer: whether to return the RDMA buffer.
"""
tensor = self.runtime.get_local_buffer_tensor(dtype, offset, use_rdma_buffer)
if size is None:
return tensor
assert tensor.numel() >= size.numel()
return tensor[:size.numel()].view(size)
@staticmethod
def _unpack_bias(bias: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]):
bias_0, bias_1 = None, None
if isinstance(bias, torch.Tensor):
bias_0 = bias
elif isinstance(bias, tuple):
assert len(bias) == 2
bias_0, bias_1 = bias
return bias_0, bias_1
@staticmethod
def get_dispatch_config(num_ranks: int) -> Config:
"""
Get a recommended dispatch config.
Argument:
num_ranks: the number of ranks.
Returns:
config: the recommended config.
"""
# TODO: automatically tune
config_map = {
2: Config(Buffer.num_sms, 24, 256, 6, 128),
4: Config(Buffer.num_sms, 6, 256, 6, 128),
8: Config(Buffer.num_sms, 6, 256, 6, 128),
16: Config(Buffer.num_sms, 36, 288, 20, 128),
24: Config(Buffer.num_sms, 32, 288, 8, 128),
32: Config(Buffer.num_sms, 32, 288, 8, 128),
48: Config(Buffer.num_sms, 32, 288, 8, 128),
64: Config(Buffer.num_sms, 32, 288, 8, 128),
96: Config(Buffer.num_sms, 20, 480, 12, 128),
128: Config(Buffer.num_sms, 20, 560, 12, 128),
144: Config(Buffer.num_sms, 32, 720, 12, 128),
160: Config(Buffer.num_sms, 28, 720, 12, 128),
}
assert num_ranks in config_map, f'Unsupported number of EP ranks: {num_ranks}'
return config_map[num_ranks]
@staticmethod
def get_combine_config(num_ranks: int) -> Config:
"""
Get a recommended combine config.
Argument:
num_ranks: the number of ranks.
Returns:
config: the recommended config.
"""
# TODO: automatically tune
config_map = {
2: Config(Buffer.num_sms, 10, 256, 6, 128),
4: Config(Buffer.num_sms, 9, 256, 6, 128),
8: Config(Buffer.num_sms, 4, 256, 6, 128),
16: Config(Buffer.num_sms, 4, 288, 12, 128),
24: Config(Buffer.num_sms, 1, 288, 8, 128),
32: Config(Buffer.num_sms, 1, 288, 8, 128),
48: Config(Buffer.num_sms, 1, 288, 8, 128),
64: Config(Buffer.num_sms, 1, 288, 8, 128),
96: Config(Buffer.num_sms, 1, 480, 8, 128),
128: Config(Buffer.num_sms, 1, 560, 8, 128),
144: Config(Buffer.num_sms, 2, 720, 8, 128),
160: Config(Buffer.num_sms, 2, 720, 8, 128),
}
assert num_ranks in config_map, f'Unsupported number of EP ranks: {num_ranks}'
return config_map[num_ranks]
# noinspection PyTypeChecker
def get_dispatch_layout(self, topk_idx: torch.Tensor, num_experts: int,
previous_event: Optional[EventOverlap] = None, async_finish: bool = False,
allocate_on_comm_stream: bool = False) -> \
Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor, EventOverlap]:
"""
Calculate the layout required for later communication.
Arguments:
topk_idx: `[num_tokens, num_topk]`, dtype must be `deep_ep.topk_idx_t` (typically `torch.int64`), the expert
indices selected by each token, `-1` means no selections.
num_experts: the number of experts.
previous_event: the event to wait before actually executing the kernel.
async_finish: the current stream will not wait for the communication kernels to be finished if set.
allocate_on_comm_stream: control whether all the allocated tensors' ownership to be on the communication stream.
Returns:
num_tokens_per_rank: `[num_ranks]` with `torch.int`, the number of tokens to be sent to each rank.
num_tokens_per_rdma_rank: `[num_rdma_ranks]` with `torch.int`, the number of tokens to be sent to each RDMA
rank (with the same GPU index), return `None` for intranode settings.
num_tokens_per_expert: `[num_experts]` with `torch.int`, the number of tokens to be sent to each expert.
is_token_in_rank: `[num_tokens, num_ranks]` with `torch.bool`, whether a token be sent to a rank.
event: the event after executing the kernel (valid only if `async_finish` is set).
"""
num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, event = \
self.runtime.get_dispatch_layout(topk_idx, num_experts, getattr(previous_event, 'event', None),
async_finish, allocate_on_comm_stream)
return num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, EventOverlap(event)
# noinspection PyTypeChecker
def dispatch(self, x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
handle: Optional[Tuple] = None,
num_tokens_per_rank: Optional[torch.Tensor] = None, num_tokens_per_rdma_rank: Optional[torch.Tensor] = None,
is_token_in_rank: Optional[torch.Tensor] = None, num_tokens_per_expert: Optional[torch.Tensor] = None,
topk_idx: Optional[torch.Tensor] = None, topk_weights: Optional[torch.Tensor] = None,
expert_alignment: int = 1, num_worst_tokens: int = 0,
config: Optional[Config] = None,
previous_event: Optional[EventOverlap] = None, async_finish: bool = False,
allocate_on_comm_stream: bool = False) -> \
Tuple[Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor], Optional[torch.Tensor],
Optional[torch.Tensor], List[int], Tuple, EventOverlap]:
"""
Dispatch tokens to different ranks, both intranode and internode settings are supported.
Intranode kernels require all the ranks should be visible via NVLink.
Internode kernels require the ranks in a node should be visible via NVLink, while the ranks with the same GPU
index should be visible via RDMA.
Arguments:
x: `torch.Tensor` or tuple of `torch.Tensor`, for the first type, the shape must be `[num_tokens, hidden]`,
and type must be `torch.bfloat16`; for the second type, the first element of the tuple must be shaped as
`[num_tokens, hidden]` with type `torch.float8_e4m3fn`, the second must be `[num_tokens, hidden // 128]`
(requiring divisible) with type `torch.float`.
handle: an optional communication handle, if set, the CPU will reuse the layout information to save some time.
num_tokens_per_rank: `[num_ranks]` with `torch.int`, the number of tokens to be sent to each rank.
num_tokens_per_rdma_rank: `[num_rdma_ranks]` with `torch.int`, the number of tokens to be sent to each RDMA
rank (with the same GPU index), return `None` for intranode settings.
is_token_in_rank: `[num_tokens, num_ranks]` with `torch.bool`, whether a token be sent to a rank.
num_tokens_per_expert: `[num_experts]` with `torch.int`, the number of tokens to be sent to each expert.
topk_idx: `[num_tokens, num_topk]` with `deep_ep.topk_idx_t` (typically `torch.int64`), the expert indices
selected by each token, `-1` means no selections.
topk_weights: `[num_tokens, num_topk]` with `torch.float`, the expert weights of each token to dispatch.
expert_alignment: align the number of tokens received by each local expert to this variable.
num_worst_tokens: the worst number of tokens to receive, if specified, there will be no CPU sync, and it
will be CUDA-graph compatible. Please also notice that this flag is for intranode only.
config: the performance tuning config.
previous_event: the event to wait before actually executing the kernel.
async_finish: the current stream will not wait for the communication kernels to be finished if set.
allocate_on_comm_stream: control whether all the allocated tensors' ownership to be on the communication stream.
Returns:
recv_x: received tokens, the same type and tuple as the input `x`, but the number of tokens equals to the
received token count.
recv_topk_idx: received expert indices.
recv_topk_weights: received expert weights.
num_recv_tokens_per_expert_list: Python list shaped `[num_local_experts]`, the received token count by
each local expert, aligned to the input `expert_alignment`. If `num_worst_tokens` is specified, the list
will be empty.
handle: the returned communication handle.
event: the event after executing the kernel (valid only if `async_finish` is set).
"""
check_torch_deterministic()
# Default config
config = self.get_dispatch_config(self.group_size) if config is None else config
# Internode
if self.runtime.get_num_rdma_ranks() > 1:
assert num_worst_tokens == 0, 'Internode dispatch does not support `num_worst_tokens > 0`'
return self.internode_dispatch(x, handle, num_tokens_per_rank, num_tokens_per_rdma_rank, is_token_in_rank,
num_tokens_per_expert, topk_idx, topk_weights, expert_alignment, config, previous_event,
async_finish, allocate_on_comm_stream)
# Launch the kernel with cached or non-cached mode
x, x_scales = x if isinstance(x, tuple) else (x, None)
if handle is not None:
assert topk_idx is None and topk_weights is None
rank_prefix_matrix, channel_prefix_matrix, recv_channel_prefix_matrix, recv_src_idx, is_token_in_rank, send_head = handle
num_recv_tokens = recv_src_idx.size(0)
recv_x, recv_x_scales, _, _, _, _, _, _, _, _, event = self.runtime.intranode_dispatch(
x, x_scales, None, None, None, is_token_in_rank, None, num_recv_tokens, rank_prefix_matrix, channel_prefix_matrix,
expert_alignment, num_worst_tokens, config, getattr(previous_event, 'event', None), async_finish, allocate_on_comm_stream)
return (recv_x, recv_x_scales) if x_scales is not None else recv_x, None, None, None, None, EventOverlap(event)
else:
assert num_tokens_per_rank is not None and is_token_in_rank is not None and num_tokens_per_expert is not None
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 = \
self.runtime.intranode_dispatch(x, x_scales, topk_idx, topk_weights,
num_tokens_per_rank, is_token_in_rank, num_tokens_per_expert, 0, None, None,
expert_alignment, num_worst_tokens, config,
getattr(previous_event, 'event', None), async_finish, allocate_on_comm_stream)
handle = (rank_prefix_matrix, channel_prefix_matrix, recv_channel_prefix_matrix, recv_src_idx, is_token_in_rank, send_head)
return (
recv_x, recv_x_scales
) if x_scales is not None else recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, EventOverlap(
event)
# noinspection PyTypeChecker
def combine(self, x: torch.Tensor, handle: Tuple,
topk_weights: Optional[torch.Tensor] = None,
bias: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] = None,
config: Optional[Config] = None,
previous_event: Optional[EventOverlap] = None, async_finish: bool = False,
allocate_on_comm_stream: bool = False) -> \
Tuple[torch.Tensor, Optional[torch.Tensor], EventOverlap]:
"""
Combine (reduce) tokens (addition **without** weights) from different ranks, both intranode and internode
settings are supported.
Intranode kernels require all the ranks should be visible via NVLink.
Internode kernels require the ranks in a node should be visible via NVLink, while the ranks with the same GPU
index should be visible via RDMA.
Arguments:
x: `[num_tokens, hidden]` with `torch.bfloat16`, the tokens to send for reducing to its original ranks.
handle: a must-set communication handle, you can obtain this from the dispatch function.
topk_weights: `[num_tokens, num_topk]` with `torch.float`, the tokens' top-k weights for reducing to its original ranks.
bias: 0, 1 or 2 `[num_tokens, hidden]` with `torch.bfloat16` final bias to the output.
config: the performance tuning config.
previous_event: the event to wait before actually executing the kernel.
async_finish: the current stream will not wait for the communication kernels to be finished if set.
allocate_on_comm_stream: control whether all the allocated tensors' ownership to be on the communication stream.
Returns:
recv_x: the reduced token from its dispatched ranks.
recv_topk_weights: the reduced top-k weights from its dispatch ranks.
event: the event after executing the kernel (valid only if `async_finish` is set).
"""
check_torch_deterministic()
# Default config
config = self.get_combine_config(self.group_size) if config is None else config
# Internode
if self.runtime.get_num_rdma_ranks() > 1:
return self.internode_combine(x, handle, topk_weights, bias, config, previous_event, async_finish, allocate_on_comm_stream)
# NOTES: the second `_` is for the sending side, so we should use the third one
rank_prefix_matrix, _, channel_prefix_matrix, src_idx, is_recv_token_in_rank, send_head = handle
bias_0, bias_1 = Buffer._unpack_bias(bias)
# Launch the kernel
recv_x, recv_topk_weights, event = self.runtime.intranode_combine(x, topk_weights, bias_0, bias_1, src_idx, rank_prefix_matrix,
channel_prefix_matrix, send_head, config,
getattr(previous_event, 'event',
None), async_finish, allocate_on_comm_stream)
return recv_x, recv_topk_weights, EventOverlap(event)
# noinspection PyTypeChecker
def internode_dispatch(self, x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
handle: Optional[Tuple] = None,
num_tokens_per_rank: Optional[torch.Tensor] = None, num_tokens_per_rdma_rank: Optional[torch.Tensor] = None,
is_token_in_rank: Optional[torch.Tensor] = None, num_tokens_per_expert: Optional[torch.Tensor] = None,
topk_idx: Optional[torch.Tensor] = None, topk_weights: Optional[torch.Tensor] = None, expert_alignment: int = 1,
config: Optional[Config] = None,
previous_event: Optional[EventOverlap] = None, async_finish: bool = False,
allocate_on_comm_stream: bool = False) -> \
Tuple[Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor], Optional[torch.Tensor],
Optional[torch.Tensor], List[int], Tuple, EventOverlap]:
"""
Internode dispatch implementation, for more details, please refer to the `dispatch` docs.
Normally, you should not directly call this function.
"""
assert config is not None
# Launch the kernel with cached or non-cached mode
x, x_scales = x if isinstance(x, tuple) else (x, None)
if handle is not None:
assert topk_idx is None and topk_weights is None
is_token_in_rank, \
rdma_channel_prefix_matrix, gbl_channel_prefix_matrix, \
recv_rdma_channel_prefix_matrix, recv_rdma_rank_prefix_sum, recv_gbl_channel_prefix_matrix, recv_gbl_rank_prefix_sum, \
recv_src_meta, send_rdma_head, send_nvl_head = handle
num_recv_tokens = recv_src_meta.size(0)
num_rdma_recv_tokens = send_nvl_head.size(0)
recv_x, recv_x_scales, _, _, _, _, _, _, _, _, _, _, _, _, event = self.runtime.internode_dispatch(
x, x_scales, topk_idx, topk_weights, None, None, is_token_in_rank, None, num_recv_tokens, num_rdma_recv_tokens,
rdma_channel_prefix_matrix, recv_rdma_rank_prefix_sum, gbl_channel_prefix_matrix, recv_gbl_rank_prefix_sum,
expert_alignment, 0, config, getattr(previous_event, 'event', None), async_finish, allocate_on_comm_stream)
return (recv_x, recv_x_scales) if x_scales is not None else recv_x, None, None, None, None, EventOverlap(event)
else:
assert num_tokens_per_rank is not None and is_token_in_rank is not None and num_tokens_per_expert is not None
recv_x, recv_x_scales, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, \
rdma_channel_prefix_matrix, gbl_channel_prefix_matrix, \
recv_rdma_channel_prefix_matrix, recv_rdma_rank_prefix_sum, \
recv_gbl_channel_prefix_matrix, recv_gbl_rank_prefix_sum, \
recv_src_meta, send_rdma_head, send_nvl_head, event = self.runtime.internode_dispatch(
x, x_scales, topk_idx, topk_weights,
num_tokens_per_rank, num_tokens_per_rdma_rank, is_token_in_rank, num_tokens_per_expert,
0, 0, None, None, None, None,
expert_alignment, 0, config, getattr(previous_event, 'event', None), async_finish, allocate_on_comm_stream)
handle = (is_token_in_rank, rdma_channel_prefix_matrix, gbl_channel_prefix_matrix, recv_rdma_channel_prefix_matrix,
recv_rdma_rank_prefix_sum, recv_gbl_channel_prefix_matrix, recv_gbl_rank_prefix_sum, recv_src_meta, send_rdma_head,
send_nvl_head)
return (
recv_x, recv_x_scales
) if x_scales is not None else recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, EventOverlap(
event)
# noinspection PyTypeChecker
def internode_combine(self, x: torch.Tensor, handle: Union[tuple, list],
topk_weights: Optional[torch.Tensor] = None,
bias: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] = None,
config: Optional[Config] = None,
previous_event: Optional[EventOverlap] = None, async_finish: bool = False,
allocate_on_comm_stream: bool = False) -> \
Tuple[torch.Tensor, Optional[torch.Tensor], EventOverlap]:
"""
Internode combine implementation, for more details, please refer to the `combine` docs.
Normally, you should not directly call this function.
"""
assert config is not None
# Unpack handle and bias
is_combined_token_in_rank, \
_, _, \
rdma_channel_prefix_matrix, rdma_rank_prefix_sum, gbl_channel_prefix_matrix, gbl_rank_prefix_sum, \
src_meta, send_rdma_head, send_nvl_head = handle
bias_0, bias_1 = Buffer._unpack_bias(bias)
# Launch the kernel
combined_x, combined_topk_weights, event = self.runtime.internode_combine(x, topk_weights, bias_0, bias_1, src_meta,
is_combined_token_in_rank, rdma_channel_prefix_matrix,
rdma_rank_prefix_sum, gbl_channel_prefix_matrix,
send_rdma_head, send_nvl_head, config,
getattr(previous_event, 'event',
None), async_finish, allocate_on_comm_stream)
return combined_x, combined_topk_weights, EventOverlap(event)
def clean_low_latency_buffer(self, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int) -> None:
"""
As low-latency kernels require part of the buffer to be zero-initialized, so it is vital to clean the buffer
if the buffer is dirty at some time.
For example, after running the normal dispatch/combine, you must run this function before executing any
low-latency kernel.
Arguments:
num_max_dispatch_tokens_per_rank: the maximum number of tokens to dispatch, all the ranks must hold the same value.
hidden: the hidden dimension of each token.
num_experts: the number of all experts.
"""
self.runtime.clean_low_latency_buffer(num_max_dispatch_tokens_per_rank, hidden, num_experts)
# noinspection PyTypeChecker
def low_latency_dispatch(self, x: torch.Tensor, topk_idx: torch.Tensor,
num_max_dispatch_tokens_per_rank: int, num_experts: int,
cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None,
dispatch_wait_recv_cost_stats: Optional[torch.Tensor] = None,
use_fp8: bool = True, round_scale: bool = False, use_ue8m0: bool = False,
async_finish: bool = False, return_recv_hook: bool = False) -> \
Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor, Tuple, EventOverlap, Callable]:
"""
A low-latency implementation for dispatching with IBGDA.
This kernel requires all the ranks (no matter intranode or internode) should be visible via RDMA
(specifically, IBGDA must be enabled).
Warning: as there are only two buffers, and the returned tensors reuse the buffer, you cannot hold more than 2
low-latency kernels' result tensors at a single moment.
Arguments:
x: `torch.Tensor` with `torch.bfloat16`, shaped as `[num_tokens, hidden]`, only several hidden shapes are
supported. The number of tokens to be dispatched must be less than `num_max_dispatch_tokens_per_rank`.
topk_idx: `torch.Tensor` with `deep_ep.topk_idx_t` (typically `torch.int64`), shaped as `[num_tokens, num_topk]`,
only several top-k shapes are supported. `-1` indices (not selecting any expert) are supported.
num_max_dispatch_tokens_per_rank: the maximum number of tokens to dispatch, all the ranks must hold the same value.
num_experts: the number of all experts.
cumulative_local_expert_recv_stats: a cumulative expert count tensor for statistics, which should have shape
`[num_local_experts]` and be typed as `torch.int`. This is useful for online service EP load balance
monitoring.
dispatch_wait_recv_cost_stats: a cumulative time spent waiting to receive each token tensor for statistics,
which should have shape `[num_ranks, num_ranks]` and be typed as `torch.int64`.
This is useful for detecting and precisely localizing slow anomalies.
use_fp8: whether to enable FP8 casting, with this, the received data will be a tuple of FP8 tensor and scaling factors.
round_scale: whether round the scaling factors into power of 2.
use_ue8m0: whether use UE8M0 as scaling factor format (available only with `round_scale=True`).
async_finish: the current stream will not wait for the communication kernels to be finished if set.
return_recv_hook: return a receiving hook if set. If set, the kernel will just do the RDMA request issues,
but **without actually receiving the data**. You must call the received hook to make sure the data's arrival.
If you do not set this flag, the kernel will ensure the data's arrival.
Returns:
recv_x: a tensor or tuple with received tokens for each expert.
With `use_fp8=True`: the first element is a `torch.Tensor` shaped as
`[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden]` with `torch.float8_e4m3fn`.
The second tensor is the corresponding scales for the first element with shape
`[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden // 128]` with `torch.float`,
if `use_ue8m0=False`. With `use_ue8m0=True`, the second one is packed and shaped as
`[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden // 512]` with type `torch.int`.
Notice that, the last-two-dimension of the scaling tensors are in column-major for TMA compatibility.
With `use_fp8=False`, the result would be a tensor shaped as
`[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden]` with `torch.bfloat16`.
Moreover, not all tokens are valid, only some of the `num_max_dispatch_tokens_per_rank * num_ranks` are,
as we do not synchronize CPU received count with GPU (also not incompatible with CUDA graph if synced).
recv_count: a tensor shaped `[num_local_experts]` with type `torch.int`, indicating how many tokens each
expert receives. As mentioned before, not all tokens are valid in `recv_x`.
handle: the communication handle to be used in the `low_latency_combine` function.
event: the event after executing the kernel (valid only if `async_finish` is set).
hook: the receiving hook function (valid only if `return_recv_hook` is set).
"""
check_torch_deterministic()
assert self.nvshmem_qp_depth >= (num_max_dispatch_tokens_per_rank + 1) * 2
packed_recv_x, packed_recv_x_scales, packed_recv_count, packed_recv_src_info, packed_recv_layout_range, event, hook = \
self.runtime.low_latency_dispatch(x, topk_idx,
cumulative_local_expert_recv_stats,
dispatch_wait_recv_cost_stats,
num_max_dispatch_tokens_per_rank, num_experts,
use_fp8, round_scale, use_ue8m0,
async_finish, return_recv_hook)
handle = (packed_recv_src_info, packed_recv_layout_range, num_max_dispatch_tokens_per_rank, x.size(1), num_experts)
tensors_to_record = (x, topk_idx, packed_recv_x, packed_recv_x_scales, packed_recv_count, packed_recv_src_info,
packed_recv_layout_range, cumulative_local_expert_recv_stats)
return (packed_recv_x, packed_recv_x_scales) if use_fp8 else packed_recv_x, packed_recv_count, handle, \
EventOverlap(event, tensors_to_record if async_finish else None), hook
# noinspection PyTypeChecker
def low_latency_combine(self, x: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor,
handle: tuple, use_logfmt: bool = False, zero_copy: bool = False, async_finish: bool = False,
return_recv_hook: bool = False, out: Optional[torch.Tensor] = None,
combine_wait_recv_cost_stats: Optional[torch.Tensor] = None) -> \
Tuple[torch.Tensor, EventOverlap, Callable]:
"""
A low-latency implementation for combining tokens (reduce **with weights**) with IBGDA.
This kernel requires all the ranks (no matter intranode or internode) should be visible via RDMA
(specifically, IBGDA must be enabled).
Warning: as there are only two buffers, and the returned tensors reuse the buffer, you cannot hold more than 2
low-latency kernels' result tensors at a single moment.
Arguments:
x: `[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden]` with `torch.bfloat16`,
the local calculated tokens to be sent to this original rank and reduced.
topk_idx: `[num_combined_tokens, num_topk]` with `deep_ep.topk_idx_t` (typically `torch.int64`), the expert
indices selected by the dispatched tokens. `-1` indices (not selecting any expert) are supported. Note that,
`num_combined_tokens` equals to the number of dispatched tokens.
topk_weights: `[num_combined_tokens, num_topk]` with `torch.float`, the expert weights selected by the dispatched
tokens. The received tokens will be reduced with the weights in this tensor.
handle: the communication handle given by the `dispatch` function.
use_logfmt: whether to use an internal "LogFMT with dynamic per-64-channel cast" format (10 bits).
zero_copy: whether the tensor is already copied into the RDMA buffer, should be cooperative
with `get_next_low_latency_combine_buffer`.
async_finish: the current stream will not wait for the communication kernels to be finished if set.
return_recv_hook: return a receiving hook if set. If set, the kernel will just do the RDMA request issues,
but **without actually receiving the data**. You must call the received hook to make sure the data's arrival.
If you do not set this flag, the kernel will ensure the data's arrival.
out: the in-place output tensor, if set, the kernel will write the result to this tensor and return it directly.
combine_wait_recv_cost_stats: a cumulative time spent waiting to receive each token tensor for statistics,
which should have shape `[num_ranks, num_ranks]` and be typed as `torch.int64`.
This is useful for detecting and precisely localizing slow anomalies.
Returns:
combined_x: the reduced token tensor, with shape `[num_combined_tokens, hidden]` and type `torch.bfloat16`.
event: the event after executing the kernel (valid only if `async_finish` is set).
hook: the receiving hook function (valid only if `return_recv_hook` is set).
"""
check_torch_deterministic()
src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
assert self.nvshmem_qp_depth >= (num_max_dispatch_tokens_per_rank + 1) * 2
combined_x, event, hook = self.runtime.low_latency_combine(x, topk_idx, topk_weights, src_info, layout_range,
combine_wait_recv_cost_stats, num_max_dispatch_tokens_per_rank,
num_experts, use_logfmt, zero_copy, async_finish, return_recv_hook, out)
tensors_to_record = (x, topk_idx, topk_weights, src_info, layout_range, combined_x)
return combined_x, EventOverlap(event, tensors_to_record if async_finish else None), hook
def low_latency_update_mask_buffer(self, rank_to_mask: int, mask: bool = False):
"""
Mask (unmask) a rank during communication (dispatch, combine, and clean)
Arguments:
rank_to_mask: the rank to mask (unmask).
mask: if True, will mask the rank (do not recv from/send to the rank), otherwise will unmask the rank.
"""
self.runtime.low_latency_update_mask_buffer(rank_to_mask, mask)
def low_latency_query_mask_buffer(self, mask_status: torch.Tensor):
"""
Query the mask status of all ranks
Arguments:
mask_status: `[num_ranks]` with `torch.int`, the mask status of each rank. `1` means mask and `0` means unmasked.
"""
self.runtime.low_latency_query_mask_buffer(mask_status)
def low_latency_clean_mask_buffer(self):
"""
Clean the mask buffer
"""
self.runtime.low_latency_clean_mask_buffer()
def get_next_low_latency_combine_buffer(self, handle: object):
"""
Get the raw registered RDMA buffer tensor for next low-latency combine, so that the next combine kernel can skip the copying.
Arguments:
handle: the communication handle given by the `dispatch` function.
Returns:
buffer: the raw RDMA low-latency buffer as a BF16 PyTorch tensor with shape
`[num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank, hidden]`, you should fill this buffer
by yourself.
"""
src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
return self.runtime.get_next_low_latency_combine_buffer(num_max_dispatch_tokens_per_rank, hidden, num_experts)
deep_ep/include/deep_ep/common/comm.cuh
+266
View File
@@ -0,0 +1,266 @@
#pragma once
#include <cooperative_groups.h>
#include <nccl.h>
#include <nccl_device.h>
#include <deep_ep/common/handle.cuh>
#include <deep_ep/common/ptx.cuh>
#include <deep_ep/common/layout.cuh>
namespace deep_ep::elastic::comm {
static constexpr int64_t kNumOneSecCycles = 2000000000; // An approximation of the GPU clock at 2000 MHz
// Some reserved tags
static constexpr int kDeviceBarrierTag = 0;
static constexpr int kKernelBarrierTag = 1;
static constexpr int kDispatchTag0 = 2;
static constexpr int kDispatchTag1 = 3;
static constexpr int kCombineTag0 = 4;
static constexpr int kCombineTag1 = 5;
static constexpr int kHybridDispatchTag0 = 6;
static constexpr int kHybridDispatchTag1 = 7;
static constexpr int kHybridCombineTag0 = 8;
static constexpr int kHybridCombineTag1 = 9;
// Some reserved count
static constexpr int kFlushAllAllocatedQPs = -1;
template <int64_t kNumTimeoutCycles, typename func_t>
__device__ __forceinline__ void timeout_while(const bool& condition, const func_t& func,
int64_t start_clock = 0) {
// User may share a start clock for multiple waits
if (start_clock == 0)
start_clock = clock64();
while (condition) {
const bool timeout = clock64() - start_clock >= kNumTimeoutCycles;
if (func(timeout))
break;
if (timeout) {
// Wait another 1 second to let all threads print information and trap
start_clock = clock64();
while (clock64() - start_clock < kNumOneSecCycles) {}
ptx::trap();
}
}
}
template <int64_t kNumTimeoutCycles, typename func_t>
__device__ __forceinline__ void timeout_while(const func_t& func, const int64_t& start_clock = 0) {
timeout_while<kNumTimeoutCycles, func_t>(true, func, start_clock);
}
template <int kNumSMs, int kNumQPs, int kNumChannelsPerSM, bool kWithNotifyWarps = false>
__device__ __forceinline__ std::pair<int, ncclGinResourceSharingMode> get_qp_mode(
const int& sm_idx, const int& channel_in_sm_idx, const bool& is_notify_warp = false) {
constexpr auto kSharingCTA = NCCL_GIN_RESOURCE_SHARING_CTA;
constexpr auto kSharingGrid = kNumSMs == 1 ? NCCL_GIN_RESOURCE_SHARING_CTA : NCCL_GIN_RESOURCE_SHARING_GPU;
// Only one QP
if constexpr (kNumQPs == 1)
return {0, kSharingGrid};
// The notify warp always use 1 SM and 1 QP
if (is_notify_warp)
return {0, kSharingCTA};
// Data channels
constexpr int kQPStartIdx = static_cast<int>(kWithNotifyWarps);
constexpr int kNumAvailableQPs = kNumQPs - kQPStartIdx;
if constexpr (kNumSMs <= kNumAvailableQPs) {
// A single SM uses an entire QP
// e.g., 3 SMs and 10 QPs
// SM 0: 0 3 6 9
// SM 1: 1 4 7
// SM 2: 2 5 8
const int num_qps_in_sm = (kNumAvailableQPs / kNumSMs) + (sm_idx < (kNumAvailableQPs % kNumSMs));
return {kQPStartIdx + sm_idx + (channel_in_sm_idx % num_qps_in_sm) * kNumSMs, kSharingCTA};
} else {
// All SMs share all QPs
const auto global_channel_idx = sm_idx * kNumChannelsPerSM + channel_in_sm_idx;
return {kQPStartIdx + (global_channel_idx % kNumAvailableQPs), kSharingGrid};
}
}
template <int kNumRanks, int kNumSMs, int kNumThreads, int64_t kNumTimeoutCycles, int kTag = kDeviceBarrierTag>
__forceinline__ __device__ void nvlink_barrier_wo_local_sync(
const handle::NCCLGin& gin,
const layout::WorkspaceLayout& workspace,
const int& rank_idx, const int& sm_idx, const int& thread_idx) {
// This barrier only uses 1 SM
if (kNumSMs > 1 and sm_idx > 0)
return;
// Read the current barrier phase first
const int status = static_cast<int>((*workspace.get_nvl_barrier_counter_ptr()) & 3);
const int phase = status & 1, sign = status >> 1;
EP_STATIC_ASSERT(kNumRanks <= kNumThreads, "Insufficient threads");
if (thread_idx < kNumRanks) {
const auto dst_ptr =
gin.get_sym_ptr<ncclTeamTagLsa>(workspace.get_nvl_barrier_signal_ptr(phase), thread_idx);
ptx::red_add_rel_sys(dst_ptr, sign ? -1 : 1);
}
__syncthreads();
// NOTES: we need `2^64 / 1e6 / 3600 / 24 / 365 = 571000` years to make the counter overflow (1 barrier per us)
// Add the phase counter
if (thread_idx == 0)
atomicAdd(workspace.get_nvl_barrier_counter_ptr(), 1);
// Check timeout
const auto target = sign ? 0 : kNumRanks;
timeout_while<kNumTimeoutCycles>(thread_idx == 0, [=](const bool& is_last_check) {
const auto signal = ptx::ld_acquire_sys<int>(workspace.get_nvl_barrier_signal_ptr(phase));
if (signal == target)
return true;
if (is_last_check) {
printf("DeepEP NVLink barrier timeout, tag: %d, nvl: %d, thread: %d, "
"status: %d, signal: %d, phase: %d, target: %d, counter: %llu\n",
kTag, rank_idx, thread_idx, status, signal, phase, target,
*workspace.get_nvl_barrier_counter_ptr());
}
return false;
});
}
template <int kNumRanks, int kNumSMs, int kNumThreads, int kNumQPs, int64_t kNumTimeoutCycles,
typename team_t, int kTag = kDeviceBarrierTag,
bool kFlushStores = true,
int kNumWarps = kNumThreads / 32>
__forceinline__ __device__ void gin_barrier_wo_local_sync(
const ncclDevComm_t& nccl_dev_comm,
const int& scaleout_rank_idx, const int& scaleup_rank_idx,
const int& sm_idx, const int& thread_idx) {
const auto global_warp_idx = sm_idx * kNumWarps + (thread_idx / 32);
const int& rank_idx = (std::is_same_v<team_t, ncclTeamTagWorld>) ? scaleup_rank_idx : scaleout_rank_idx;
const int num_qps = kNumQPs == kFlushAllAllocatedQPs ? nccl_dev_comm.ginContextCount : kNumQPs;
// Flush all QPs by all SMs (only needed for release semantics)
if constexpr (kFlushStores) {
for (int i = global_warp_idx; i < num_qps; i += kNumSMs * kNumWarps) {
ncclGin(nccl_dev_comm, i, NCCL_GIN_RESOURCE_SHARING_CTA).flush(ncclCoopWarp());
}
// NOTES: we can not use `kNumSMs` to judge, as maybe only part of the SMs will call this function
(gridDim.x > 1) ? cooperative_groups::this_grid().sync() : __syncthreads();
}
if (sm_idx == 0) {
// Use QP 0 to do barrier
const auto team = (std::is_same_v<team_t, ncclTeamTagWorld>) ?
ncclTeamWorld(nccl_dev_comm) : ncclTeamRail(nccl_dev_comm);
const ncclGin gin(nccl_dev_comm, 0, NCCL_GIN_RESOURCE_SHARING_CTA);
for (int i = thread_idx; i < kNumRanks; i += kNumThreads)
gin.signal(team, i, ncclGin_SignalInc{static_cast<ncclGinSignal_t>(rank_idx)});
// TODO(NCCL): Using the official NCCL wait signal API, after they added timeout check.
for (int i = thread_idx; i < kNumRanks; i += kNumThreads) {
const auto signal_idx = static_cast<ncclGinSignal_t>(i);
const auto shadow_ptr = gin.getSignalShadowPtr(signal_idx);
const auto target = ++(*shadow_ptr);
const auto gdaki = static_cast<struct ncclGinGdakiGPUContext*>(gin._ginHandle) + gin.contextId;
const auto signal_ptr = reinterpret_cast<uint64_t*>(__ldg(reinterpret_cast<uint64_t*>(&gdaki->signals_table.buffer))) + signal_idx;
timeout_while<kNumTimeoutCycles>([=](const bool& is_last_check) {
const auto signal = ptx::ld_acquire_sys<uint64_t>(signal_ptr);
if (signal >= target)
return true;
if (is_last_check) {
printf("DeepEP Gin barrier timeout, tag: %d, scaleout: %d, scaleup: %d, thread: %d, "
"signal: %lu, target: %lu\n", kTag, scaleout_rank_idx, scaleup_rank_idx, thread_idx, signal, target);
}
return false;
});
}
}
}
template <bool kIsScaleupNVLink, int kNumRanks, int kNumSMs, int kNumThreads, int kNumQPs,
int64_t kNumTimeoutCycles, int kTag = kDeviceBarrierTag, bool kFlushStores = true>
__forceinline__ __device__ void scaleup_barrier_wo_local_sync(
const handle::NCCLGin& gin,
const layout::WorkspaceLayout& workspace,
const int& rank_idx, const int& sm_idx, const int& thread_idx) {
if constexpr (kIsScaleupNVLink) {
nvlink_barrier_wo_local_sync<kNumRanks, kNumSMs, kNumThreads, kNumTimeoutCycles, kTag>(
gin, workspace, rank_idx, sm_idx, thread_idx);
} else {
gin_barrier_wo_local_sync<kNumRanks, kNumSMs, kNumThreads, kNumQPs, kNumTimeoutCycles, ncclTeamTagWorld, kTag, kFlushStores>(
gin.nccl_dev_comm, 1, rank_idx, sm_idx, thread_idx);
}
}
template <int kNumRanks, int kNumSMs, int kNumThreads, int kNumQPs, int64_t kNumTimeoutCycles, int kTag = kDeviceBarrierTag,
bool kFlushStores = true>
__forceinline__ __device__ void scaleout_barrier_wo_local_sync(
const handle::NCCLGin& gin,
const int& scaleout_rank_idx, const int& scaleup_rank_idx,
const int& sm_idx, const int& thread_idx) {
gin_barrier_wo_local_sync<kNumRanks, kNumSMs, kNumThreads, kNumQPs, kNumTimeoutCycles, ncclTeamTagRail, kTag, kFlushStores>(
gin.nccl_dev_comm, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx);
}
template <bool kIsScaleupNVLink,
int kNumScaleoutRanks, int kNumScaleupRanks,
int kNumSMs, int kNumThreads, int kNumQPs,
int64_t kNumTimeoutCycles, int kTag = kDeviceBarrierTag,
bool kFlushStores = true, bool kSyncAtStart = true, bool kSyncAtEnd = true>
__forceinline__ __device__ void gpu_barrier(const handle::NCCLGin& gin,
const layout::WorkspaceLayout& workspace,
const int& scaleout_rank_idx, const int& scaleup_rank_idx,
const int& sm_idx, const int& thread_idx,
bool do_scaleout = true, bool do_scaleup = true) {
// A general TMA store wait to prevent proxy memory issues
if constexpr (kFlushStores) {
ptx::tma_store_commit();
ptx::tma_store_wait();
__syncwarp();
}
// All the SMs should wait
if constexpr (kSyncAtStart) {
cooperative_groups::this_grid().sync();
} else {
EP_STATIC_ASSERT(not kFlushStores, "No data to be flushed");
}
do_scaleout &= kNumScaleoutRanks > 1;
do_scaleup &= kNumScaleupRanks > 1;
if (do_scaleup and do_scaleout) {
// Do scaleup and scaleout barrier in parallel
EP_DEVICE_ASSERT(kNumSMs >= 2 and "At least 2 SMs for a hybrid barrier");
if (sm_idx == 0) {
// First SM do the scaleup barrier
scaleup_barrier_wo_local_sync<kIsScaleupNVLink, kNumScaleupRanks, kNumSMs, kNumThreads, kNumQPs, kNumTimeoutCycles, kTag, kFlushStores>(
gin, workspace, scaleup_rank_idx, sm_idx, thread_idx);
// We need an extra grid sync, as the scaleout barrier will do a sync after flush, before the barrier
// NOTES: this is kind of hacky
if constexpr (kFlushStores)
cooperative_groups::this_grid().sync();
} else {
// The remaining SMs do the scaleout barrier
scaleout_barrier_wo_local_sync<kNumScaleoutRanks, kNumSMs - 1, kNumThreads, kNumQPs, kNumTimeoutCycles, kTag, kFlushStores>(
gin, scaleout_rank_idx, scaleup_rank_idx, sm_idx - 1, thread_idx);
}
} else if (do_scaleup) {
// Scaleup only
scaleup_barrier_wo_local_sync<kIsScaleupNVLink, kNumScaleupRanks, kNumSMs, kNumThreads, kNumQPs, kNumTimeoutCycles, kTag, kFlushStores>(
gin, workspace, scaleup_rank_idx, sm_idx, thread_idx);
} else if (do_scaleout) {
// Scaleout only
scaleout_barrier_wo_local_sync<kNumScaleoutRanks, kNumSMs, kNumThreads, kNumQPs, kNumTimeoutCycles, kTag, kFlushStores>(
gin, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx);
}
// All the SMs should wait
if constexpr (kSyncAtEnd)
cooperative_groups::this_grid().sync();
}
} // namespace deep_ep::elastic::comm
deep_ep/include/deep_ep/common/compiled.cuh
+87
View File
@@ -0,0 +1,87 @@
#pragma once
// Make CLion CUDA indexing work
#ifdef __CLION_IDE__
#define __CUDA_ARCH__ 900
#define __CUDACC_RDC__
#define __CUDACC__
#endif
// Remove Torch restrictions
#ifdef __CUDA_NO_HALF_CONVERSIONS__
#undef __CUDA_NO_HALF_CONVERSIONS__
#endif
#ifdef __CUDA_NO_HALF_OPERATORS__
#undef __CUDA_NO_HALF_OPERATORS__
#endif
#ifdef __CUDA_NO_HALF2_OPERATORS__
#undef __CUDA_NO_HALF2_OPERATORS__
#endif
#ifdef __CUDA_NO_BFLOAT16_CONVERSIONS__
#undef __CUDA_NO_BFLOAT16_CONVERSIONS__
#endif
#ifdef __CUDA_NO_BFLOAT162_OPERATORS__
#undef __CUDA_NO_BFLOAT162_OPERATORS__
#endif
#include <cstdint>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#ifndef DISABLE_SM90_FEATURES
#include <cuda_fp8.h>
#else
// Ampere does not support FP8 features
#define __NV_E4M3 0
#define __NV_E5M2 1
typedef int __nv_fp8_interpretation_t;
typedef int __nv_fp8x4_e4m3;
typedef uint8_t __nv_fp8_storage_t;
#endif
// Compatibility: 256 bits LD/ST instructions
#if defined(CUDART_VERSION) and CUDART_VERSION >= 13000
using longlong4_t = longlong4_32a;
#define make_longlong4_t make_longlong4_32a
#else
struct alignas(32) longlong4_t { long long x, y, z, w; };
__device__ __forceinline__ longlong4_t make_longlong4_t(
const long long& x, const long long& y, const long long& z, const long long& w) {
return {x, y, z, w};
}
#endif
#ifndef EP_NUM_TOPK_IDX_BITS
#define EP_NUM_TOPK_IDX_BITS 64
#endif
namespace deep_ep {
#ifndef DISABLE_SM90_FEATURES
constexpr bool kEnableSM90Features = true;
#else
constexpr bool kEnableSM90Features = false;
#endif
template <int kNumBits> struct int_with_bits;
template <> struct int_with_bits<8> { using type = int8_t; };
template <> struct int_with_bits<16> { using type = int16_t; };
template <> struct int_with_bits<32> { using type = int32_t; };
template <> struct int_with_bits<64> { using type = int64_t; };
using topk_idx_t = int_with_bits<EP_NUM_TOPK_IDX_BITS>::type;
union sf_pack_t {
float fp32;
int ue8m0x4;
};
constexpr int kNumTMAAlignedBytes = 16;
constexpr int kNumAlignedSFPacks = 16 / sizeof(sf_pack_t);
// Some communication channel settings
constexpr int kNumMaxChannels = 1024;
constexpr int kGinQPDepth = 1024;
constexpr int kGinQPFlushDepth = 768;
} // namespace deep_ep
deep_ep/include/deep_ep/common/exception.cuh
+94
View File
@@ -0,0 +1,94 @@
#pragma once
#include <exception>
#include <string>
#include <sstream>
#ifndef EP_STATIC_ASSERT
#define EP_STATIC_ASSERT(cond, reason) static_assert(cond, reason)
#endif
class EPException : public std::exception {
private:
std::string message = {};
public:
explicit EPException(const char* name, const char* file, const int line, const std::string& error) {
std::stringstream ss;
ss << name << " exception (" << file << ":" << line << "): " << error;
message = ss.str();
}
const char* what() const noexcept override { return message.c_str(); }
};
#define EPExceptionWithLineInfo(name, message) EPException(name, __FILE__, __LINE__, message)
#ifndef EP_HOST_ASSERT
#define EP_HOST_ASSERT(cond) \
do { \
if (not(cond)) { \
throw EPException("Assertion", __FILE__, __LINE__, #cond); \
} \
} while (0)
#endif
#ifndef EP_HOST_UNREACHABLE
#define EP_HOST_UNREACHABLE(reason) (throw EPException("Assertion", __FILE__, __LINE__, reason))
#endif
#ifndef EP_DEVICE_ASSERT
#define EP_DEVICE_ASSERT(cond) \
do { \
if (not(cond)) { \
printf("Assertion failed: %s:%d, condition: %s\n", __FILE__, __LINE__, #cond); \
asm("trap;"); \
} \
} while (0)
#endif
#ifndef EP_UNIFIED_ASSERT
#ifdef __CUDA_ARCH__
#define EP_UNIFIED_ASSERT(cond) EP_DEVICE_ASSERT(cond)
#else
#define EP_UNIFIED_ASSERT(cond) EP_HOST_ASSERT(cond)
#endif
#endif
#ifndef CUDA_RUNTIME_CHECK
#define CUDA_RUNTIME_CHECK(cmd) \
do { \
const auto e = (cmd); \
if (e != cudaSuccess) { \
std::stringstream ss; \
ss << static_cast<int>(e) << " (" << cudaGetErrorName(e) << ", " << cudaGetErrorString(e) << ")"; \
throw EPException("CUDA runtime", __FILE__, __LINE__, ss.str()); \
} \
} while (0)
#endif
#ifndef CUDA_DRIVER_CHECK
#define CUDA_DRIVER_CHECK(cmd) \
do { \
const auto e = (cmd); \
if (e != CUDA_SUCCESS) { \
std::stringstream ss; \
const char *name, *info; \
lazy_cuGetErrorName(e, &name), lazy_cuGetErrorString(e, &info); \
ss << static_cast<int>(e) << " (" << name << ", " << info << ")"; \
throw EPException("CUDA driver", __FILE__, __LINE__, ss.str()); \
} \
} while (0)
#endif
#ifndef NCCL_CHECK
#define NCCL_CHECK(cmd) \
do { \
const auto e = (cmd); \
if (e != ncclSuccess) { \
std::stringstream ss; \
ss << static_cast<int>(e) << " (" << ncclGetLastError(nullptr) << ")"; \
throw EPException("NCCL", __FILE__, __LINE__, ss.str()); \
} \
} while (0)
#endif
deep_ep/include/deep_ep/common/handle.cuh
+230
View File
@@ -0,0 +1,230 @@
#pragma once
#include <nccl.h>
#include <nccl_device.h>
#include <deep_ep/common/exception.cuh>
#include <deep_ep/common/ptx.cuh>
namespace deep_ep::elastic::handle {
struct NCCLGin {
#define IS_TEAM_WORLD(code) if constexpr (std::is_same_v<team_t, ncclTeamTagWorld>) { code }
#define IS_TEAM_LSA(code) if constexpr (std::is_same_v<team_t, ncclTeamTagLsa>) { code }
#define IS_TEAM_RAIL(code) if constexpr (std::is_same_v<team_t, ncclTeamTagRail>) { code }
#define IS_TEAM_WORLD_RAIL(code) if constexpr (std::is_same_v<team_t, ncclTeamTagWorld> or std::is_same_v<team_t, ncclTeamTagRail>) { code }
#define IS_TEAM_WORLD_LSA(code) if constexpr (std::is_same_v<team_t, ncclTeamTagWorld> or std::is_same_v<team_t, ncclTeamTagLsa>) { code }
#define TEAM_WORLD_RAIL() ((std::is_same_v<team_t, ncclTeamTagWorld>) ? team_world : team_rail)
const ncclDevComm_t& nccl_dev_comm;
const ncclWindow_t& nccl_window;
ncclGin gin;
ncclTeam team_world, team_lsa, team_rail;
uint64_t lsa_base_ptr;
// TODO(NCCL): QP index should just be a hint or the users maintain the mapping?
__device__ __forceinline__
NCCLGin(const ncclDevComm_t& nccl_dev_comm, const ncclWindow_t& nccl_window,
const int& qp_idx = 0,
const ncclGinResourceSharingMode& resource_sharing_mode = NCCL_GIN_RESOURCE_SHARING_GPU):
nccl_dev_comm(nccl_dev_comm), nccl_window(nccl_window),
gin(ncclGin(nccl_dev_comm, qp_idx, resource_sharing_mode)),
team_world(ncclTeamWorld(nccl_dev_comm)), team_lsa(ncclTeamLsa(nccl_dev_comm)), team_rail(ncclTeamRail(nccl_dev_comm)) {
// TODO: what if we only have 1 NVLink rank
lsa_base_ptr = reinterpret_cast<uint64_t>(ncclGetLsaPointer(nccl_window, 0, team_lsa.rank));
}
template <typename team_t>
__device__ __forceinline__ bool is_nvlink_accessible(const int& dst_rank_idx) const {
IS_TEAM_LSA({
return true;
})
IS_TEAM_WORLD({
// TODO(NCCL): optimize this function's cycles
// return ncclTeamRankIsMember(team_lsa, team_world, dst_rank_idx);
return team_rail.rank * team_lsa.nRanks <= dst_rank_idx and
dst_rank_idx < (team_rail.rank + 1) * team_lsa.nRanks;
})
IS_TEAM_RAIL({
// TODO(NCCL): some ranks may be connected by NVLink, e.g., "2 + 2 + 4"
return team_rail.rank == dst_rank_idx;
})
}
// ReSharper disable once CppNotAllPathsReturnValue
template <typename dtype_t = void*>
__device__ __forceinline__
uint64_t get_sym_offset(dtype_t* ptr) const {
return reinterpret_cast<uint64_t>(ptr) - lsa_base_ptr;
}
// ReSharper disable once CppNotAllPathsReturnValue
template <typename team_t, typename dtype_t = void*>
__device__ __forceinline__
dtype_t* get_sym_ptr(dtype_t* ptr, const int& dst_rank_idx) const {
IS_TEAM_RAIL({
return team_rail.rank == dst_rank_idx ? ptr : nullptr;
})
IS_TEAM_WORLD_LSA({
constexpr bool kIsTeamLSA = (std::is_same_v<team_t, ncclTeamTagLsa>);
// Team world and not accessible by symmetric pointers
if (not is_nvlink_accessible<team_t>(dst_rank_idx))
return nullptr;
// Translate into NVLink rank index
const auto dst_nvl_rank_idx = kIsTeamLSA ?
dst_rank_idx : (dst_rank_idx - team_rail.rank * team_lsa.nRanks);
// Local rank bypass
// TODO(NCCL): support this
if (dst_nvl_rank_idx == team_lsa.rank)
return ptr;
// Get base ptr
const auto dst_ptr = ncclGetLsaPointer(
nccl_window, get_sym_offset(ptr), dst_nvl_rank_idx);
return static_cast<dtype_t*>(dst_ptr);
});
}
// NOTES: take care of this function when `team_t` is not LSA
// Do not mix atomic add with gin signal into a single position
template <typename team_t, typename dtype_t>
__device__ __forceinline__
void red_add_rel(dtype_t* sym_ptr, const dtype_t& value, const int& dst_rank_idx,
const int& extra_options = 0) const {
const auto dst_ptr = get_sym_ptr<team_t>(sym_ptr, dst_rank_idx);
// Use symmetric pointers as much as possible, RDMA otherwise
if (dst_ptr != nullptr) {
if (std::is_same_v<team_t, ncclTeamTagRail> or dst_ptr == sym_ptr) {
ptx::red_add_rel_gpu(dst_ptr, value);
} else {
ptx::red_add_rel_sys(dst_ptr, value);
}
} else {
EP_DEVICE_ASSERT((not std::is_same_v<team_t, ncclTeamTagLsa>));
EP_DEVICE_ASSERT((std::is_same_v<dtype_t, int64_t>) or (std::is_same_v<dtype_t, uint64_t>));
// TODO(NCCL): support all dtypes
gin.signal(TEAM_WORLD_RAIL(), dst_rank_idx,
ncclGin_VASignalAdd(nccl_window, reinterpret_cast<int64_t>(sym_ptr) - lsa_base_ptr, static_cast<uint64_t>(value)),
ncclCoopThread(),
ncclGin_None(),
cuda::thread_scope_thread,
cuda::thread_scope_device,
ncclGinOptFlagsDefault | extra_options);
}
}
__device__ __forceinline__
void wait(ncclGinRequest_t& request) const {
gin.wait(request);
}
template <typename team_t, typename coop_t = ncclCoopThread, typename segment_t = ncclGin_SegmentDevice>
__device__ __forceinline__
void get(void* src_ptr, void* dst_ptr, const int& num_bytes, const int& src_rank_idx,
const int& extra_options = 0) const {
IS_TEAM_WORLD_RAIL({
gin.get(
TEAM_WORLD_RAIL(),
src_rank_idx,
nccl_window, reinterpret_cast<int64_t>(src_ptr) - lsa_base_ptr,
nccl_window, reinterpret_cast<int64_t>(dst_ptr) - lsa_base_ptr,
num_bytes,
coop_t(),
ncclGin_None(),
ncclGinOptFlagsDefault | extra_options,
segment_t()
);
});
}
template <typename coop_t = ncclCoopThread>
__device__ __forceinline__
void flush() const {
gin.flush(coop_t());
}
template <typename team_t, typename coop_t = ncclCoopThread>
__device__ __forceinline__
void flush_async(const int& src_rank_idx, ncclGinRequest_t* request,
const int& extra_options = 0) const {
IS_TEAM_WORLD_RAIL({
gin.flushAsync(
TEAM_WORLD_RAIL(),
src_rank_idx,
request,
coop_t(),
ncclGinOptFlagsDefault | extra_options
);
});
}
template <typename team_t, typename remote_action_t>
__device__ __forceinline__
void signal(const int& dst_rank_idx, const remote_action_t& remote_action) const {
IS_TEAM_WORLD_RAIL({
gin.signal(TEAM_WORLD_RAIL(), dst_rank_idx, remote_action);
});
}
template <typename team_t,
typename remote_action_t = ncclGin_None>
__device__ __forceinline__
void put(void* recv_sym_ptr, void* send_sym_ptr, const int& num_bytes, const int& dst_rank_idx,
const int& extra_options = 0,
const remote_action_t& remote_action = remote_action_t()) const {
// NOTES: local or NVLink put will also go through NIC via this API
IS_TEAM_WORLD_RAIL({
gin.put(TEAM_WORLD_RAIL(),
dst_rank_idx,
// TODO: can we don't repeat the window?
// TODO: can we pass raw pointers?
nccl_window, reinterpret_cast<int64_t>(recv_sym_ptr) - lsa_base_ptr,
nccl_window, reinterpret_cast<int64_t>(send_sym_ptr) - lsa_base_ptr,
num_bytes,
remote_action,
ncclGin_None(),
ncclCoopThread(),
ncclGin_None(),
cuda::thread_scope_thread,
cuda::thread_scope_device,
ncclGinOptFlagsDefault | extra_options);
});
}
template <typename team_t, typename dtype_t>
__device__ __forceinline__
void put_value(dtype_t* sym_ptr, const dtype_t& value, const int& dst_rank_idx,
const int& extra_options = 0) const {
const auto dst_ptr = get_sym_ptr<team_t>(sym_ptr, dst_rank_idx);
if (dst_ptr != nullptr) {
ptx::st_relaxed_sys(dst_ptr, value);
} else {
EP_DEVICE_ASSERT((not std::is_same_v<team_t, ncclTeamTagLsa>));
gin.putValue(TEAM_WORLD_RAIL(),
dst_rank_idx,
nccl_window, reinterpret_cast<int64_t>(sym_ptr) - lsa_base_ptr,
value,
ncclGin_None(),
ncclCoopThread(),
ncclGin_None(),
cuda::thread_scope_thread,
cuda::thread_scope_device,
ncclGinOptFlagsDefault | extra_options);
}
}
#undef IS_TEAM_WORLD
#undef IS_TEAM_LSA
#undef IS_TEAM_RAIL
#undef IS_TEAM_WORLD_RAIL
#undef IS_TEAM_WORLD_LSA
#undef TEAM_WORLD_RAIL
};
} // namespace deep_ep::elastic::handle
deep_ep/include/deep_ep/common/layout.cuh
+313
View File
@@ -0,0 +1,313 @@
#pragma once
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/exception.cuh>
#include <deep_ep/common/math.cuh>
#include <deep_ep/common/ptx.cuh>
namespace deep_ep::elastic::layout {
struct WorkspaceLayout {
void* workspace;
int num_ranks;
int num_scaleout_ranks, num_scaleup_ranks;
int num_experts, num_experts_per_rank;
// We want to fix the layout position for all settings,
// so that one buffer can be reused for all cases
static constexpr int kNumMaxRanks = 1024;
static constexpr int kNumMaxExperts = 2048;
static constexpr int kNumMaxExpertsPerRank = 256;
static constexpr int kNumMaxInflightAGRS = 32;
static constexpr int64_t kNumBarrierSignalBytes = 16;
__forceinline__ __device__ __host__
WorkspaceLayout(void* workspace,
const int& num_scaleout_ranks,
const int& num_scaleup_ranks,
const int& num_experts):
workspace(workspace),
num_ranks(num_scaleout_ranks * num_scaleup_ranks),
num_scaleout_ranks(num_scaleout_ranks),
num_scaleup_ranks(num_scaleup_ranks),
num_experts(num_experts) {
num_experts_per_rank = num_experts / num_ranks;
EP_UNIFIED_ASSERT(num_experts % num_ranks == 0);
EP_UNIFIED_ASSERT(num_ranks <= kNumMaxRanks);
EP_UNIFIED_ASSERT(num_experts <= kNumMaxExperts);
EP_UNIFIED_ASSERT(num_experts_per_rank <= kNumMaxExpertsPerRank);
}
static int64_t get_num_bytes() {
// Pure NVLink scaleup barrier signals
int64_t num_bytes = 0;
num_bytes += kNumBarrierSignalBytes;
// Notify reduction workspace
num_bytes += (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t);
// Scaleup notify threads
// Rank send/recv count
num_bytes += kNumMaxRanks * sizeof(int64_t) * 2;
// Expert send/recv count
num_bytes += kNumMaxExperts * sizeof(int64_t) * 2;
// Scaleup atomic sender count
num_bytes += kNumMaxRanks * sizeof(int);
// Scaleout notify threads
// Rank send/recv count
num_bytes += kNumMaxRanks * sizeof(int) * 2;
// Expert send/recv count
num_bytes += kNumMaxExperts * sizeof(int) * 2;
// Scaleout channel metadata (finish flag and tails)
num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int64_t);
// Channel aggregated into the scaleup domains
// Also reused for channel scaleup tail
num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int);
// Rank send/recv count, for PP prev/next ranks
num_bytes += 2 * 2 * sizeof(int64_t);
// AGRS signals
num_bytes += (kNumMaxInflightAGRS + 1) * kNumMaxRanks * sizeof(int);
return num_bytes;
}
__forceinline__ __device__ __host__ unsigned long long* get_nvl_barrier_counter_ptr() const {
return static_cast<unsigned long long*>(workspace);
}
__forceinline__ __device__ __host__ int* get_nvl_barrier_signal_ptr(const int& phase) const {
return math::advance_ptr<int>(workspace, (2 + phase) * sizeof(int));
}
__forceinline__ __device__ __host__ int64_t* get_notify_reduction_workspace_ptr() const {
return math::advance_ptr<int64_t>(workspace, kNumBarrierSignalBytes);
}
template <bool kIsSendBuffer>
__forceinline__ __device__ __host__ int64_t* get_scaleup_rank_expert_count_ptr() const {
const auto base_ptr =
math::advance_ptr<int64_t>(get_notify_reduction_workspace_ptr(), (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t));
return base_ptr + (kIsSendBuffer ? 0 : kNumMaxRanks + kNumMaxExperts);
}
template <bool kIsSendBuffer>
__forceinline__ __device__ __host__ int64_t* get_scaleup_rank_count_ptr() const {
return get_scaleup_rank_expert_count_ptr<kIsSendBuffer>();
}
template <bool kIsSendBuffer>
__forceinline__ __device__ __host__ int64_t* get_scaleup_expert_count_ptr() const {
return get_scaleup_rank_expert_count_ptr<kIsSendBuffer>() + num_scaleup_ranks;
}
__forceinline__ __device__ __host__ int* get_scaleup_atomic_sender_counter() const {
return math::advance_ptr<int>(
get_scaleup_rank_expert_count_ptr<true>(), 2 * (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t));
}
template <bool kIsSendBuffer>
__forceinline__ __device__ __host__ int* get_scaleout_rank_expert_count_ptr() const {
const auto base_ptr =
math::advance_ptr<int>(get_scaleup_atomic_sender_counter(), kNumMaxRanks * sizeof(int));
return base_ptr + (kIsSendBuffer ? 0 : kNumMaxRanks + kNumMaxExperts);
}
template <bool kIsSendBuffer>
__forceinline__ __device__ __host__ int* get_scaleout_rank_count_ptr(
const int& scaleout_rank_idx = 0, const int& scaleup_rank_idx = 0) const {
const auto base_ptr = get_scaleout_rank_expert_count_ptr<kIsSendBuffer>();
return base_ptr + scaleout_rank_idx * num_scaleup_ranks + scaleup_rank_idx;
}
template <bool kIsSendBuffer>
__forceinline__ __device__ __host__ int* get_scaleout_expert_count_ptr(
const int& scaleout_rank_idx = 0, const int& expert_idx = 0) const {
const auto base_ptr = get_scaleout_rank_expert_count_ptr<kIsSendBuffer>() + num_ranks;
return base_ptr + scaleout_rank_idx * (num_scaleup_ranks * num_experts_per_rank) + expert_idx;
}
__forceinline__ __device__ __host__ int64_t* get_scaleout_channel_signaled_tail_ptr(
const int& channel_idx, const int& scaleout_rank_idx) const {
const auto base_ptr = math::advance_ptr<int64_t>(
get_scaleout_rank_expert_count_ptr<true>(),
(kNumMaxRanks + kNumMaxExperts) * sizeof(int) * 2);
return base_ptr + (channel_idx * num_scaleout_ranks + scaleout_rank_idx);
}
__forceinline__ __device__ __host__ int* get_channel_scaleup_tail_ptr(
const int& channel_idx, const int& scaleup_rank_idx) const {
const auto base_ptr = math::advance_ptr<int>(
get_scaleout_channel_signaled_tail_ptr(0, 0),
kNumMaxRanks * kNumMaxChannels * sizeof(int64_t));
return base_ptr + (channel_idx * num_scaleup_ranks + scaleup_rank_idx);
}
__forceinline__ __device__ __host__ int64_t* get_pp_send_count_ptr(const int& offset) const {
const auto base_ptr = math::advance_ptr<int64_t>(
get_channel_scaleup_tail_ptr(0, 0),
kNumMaxRanks * kNumMaxChannels * sizeof(int));
return base_ptr + offset;
}
__forceinline__ __device__ __host__ int64_t* get_pp_recv_count_ptr(const int& offset) const {
const auto base_ptr = math::advance_ptr<int64_t>(
get_pp_send_count_ptr(0), 2 * sizeof(int64_t));
return base_ptr + offset;
}
__forceinline__ __device__ __host__ int* get_agrs_recv_signal_ptr(const int& slot, const int& rank_idx) const {
const auto base_ptr = math::advance_ptr<int>(
get_pp_recv_count_ptr(0), 2 * sizeof(int64_t));
return base_ptr + slot * kNumMaxRanks + rank_idx;
}
__forceinline__ __device__ __host__ int* get_agrs_session_signal_ptr(const int& rank_idx) const {
const auto base_ptr = math::advance_ptr<int>(
get_agrs_recv_signal_ptr(0, 0), kNumMaxInflightAGRS * kNumMaxRanks * sizeof(int));
return base_ptr + rank_idx;
}
};
struct TokenLayout {
int num_hidden_bytes, num_sf_bytes;
// NOTES: the top-k index is always 32-bit
bool with_metadata;
int num_topk, num_metadata_bytes;
void* base;
__forceinline__ __device__ __host__
TokenLayout(const int& num_hidden_bytes, const int& num_sf_bytes,
const int& num_topk, const bool& with_metadata, void* base = nullptr) :
num_hidden_bytes(num_hidden_bytes),
num_sf_bytes(num_sf_bytes),
// Metadata includes: top-k indices, weight and source rank/token index
with_metadata(with_metadata),
num_topk(num_topk),
num_metadata_bytes(num_topk * (sizeof(int) + sizeof(float)) +
(with_metadata ? (1 + num_topk) * sizeof(int) : 0)),
base(base) {
EP_STATIC_ASSERT(sizeof(int) == sizeof(float), "Invalid size assumption");
EP_UNIFIED_ASSERT(num_hidden_bytes % ptx::kNumTMAAlignBytes == 0);
}
template <bool kWithMBarrier, typename dtype_t = int>
__forceinline__ __device__ __host__ dtype_t get_num_bytes() const {
const auto num_bytes = math::align(num_hidden_bytes, ptx::kNumTMAAlignBytes) +
math::align(num_sf_bytes, ptx::kNumTMAAlignBytes) +
math::align(num_metadata_bytes, ptx::kNumTMAAlignBytes) +
math::align<int>(kWithMBarrier ? sizeof(ptx::mbarrier) : 0, ptx::kNumTMAAlignBytes);
return static_cast<dtype_t>(num_bytes);
}
__forceinline__ __device__ __host__ void* get_base_ptr() const {
return base;
}
__forceinline__ __device__ __host__ void set_base_ptr(void* ptr) {
base = ptr;
}
__forceinline__ __device__ __host__ void* get_hidden_ptr() const {
return get_base_ptr();
}
__forceinline__ __device__ __host__ sf_pack_t* get_sf_ptr() const {
return math::advance_ptr<sf_pack_t>(base, math::align(num_hidden_bytes, ptx::kNumTMAAlignBytes));
}
__forceinline__ __device__ __host__ int* get_metadata_ptr() const {
return math::advance_ptr<int>(get_sf_ptr(), math::align(num_sf_bytes, ptx::kNumTMAAlignBytes));
}
__forceinline__ __device__ __host__ int* get_topk_idx_ptr() const {
return get_metadata_ptr();
}
__forceinline__ __device__ __host__ float* get_topk_weights_ptr() const {
return math::advance_ptr<float>(get_metadata_ptr(), num_topk * sizeof(int));
}
__forceinline__ __device__ __host__ int* get_src_token_global_idx_ptr() const {
return math::advance_ptr<int>(get_topk_weights_ptr(), num_topk * sizeof(float));
}
__forceinline__ __device__ __host__ int* get_linked_list_idx_ptr() const {
return get_src_token_global_idx_ptr() + 1;
}
__forceinline__ __device__ ptx::mbarrier* get_mbarrier_ptr() const {
return math::advance_ptr<ptx::mbarrier>(get_metadata_ptr(), math::align(num_metadata_bytes, ptx::kNumTMAAlignBytes));
}
};
template <bool kWithMBarrier>
struct BufferLayout {
TokenLayout token_layout;
int num_ranks;
int num_max_tokens_per_rank;
void* base;
__forceinline__ __device__ __host__
BufferLayout(const TokenLayout& token_layout,
const int& num_ranks,
const int& max_num_tokens_per_rank,
void* base = nullptr) :
token_layout(token_layout),
num_ranks(num_ranks), num_max_tokens_per_rank(max_num_tokens_per_rank),
base(base) {}
__forceinline__ __device__ __host__
int64_t get_num_bytes_per_token() const {
return token_layout.get_num_bytes<kWithMBarrier, int64_t>();
}
__forceinline__ __device__ __host__
int64_t get_num_bytes_per_rank() const {
return num_max_tokens_per_rank * get_num_bytes_per_token();
}
__forceinline__ __device__ __host__
int64_t get_num_bytes() const {
return get_num_bytes_per_rank() * num_ranks;
}
__forceinline__ __device__ __host__
void* get_buffer_end_ptr() const {
return math::advance_ptr(base, get_num_bytes());
}
__forceinline__ __device__ __host__
BufferLayout get_rank_buffer(const int& rank_idx) const {
return BufferLayout(token_layout,
1, num_max_tokens_per_rank,
static_cast<int8_t*>(base) + get_num_bytes_per_rank() * rank_idx);
}
template <int kNumTokensPerChannel>
__forceinline__ __device__ __host__
BufferLayout get_channel_buffer(const int& channel_idx) const {
EP_UNIFIED_ASSERT(num_max_tokens_per_rank % kNumTokensPerChannel == 0);
return BufferLayout(token_layout,
// Do not use `num_max_tokens_per_rank / kNumTokensPerChannel` as the false stride
num_ranks, num_max_tokens_per_rank,
static_cast<int8_t*>(base) + get_num_bytes_per_token() * kNumTokensPerChannel * channel_idx);
}
__forceinline__ __device__ __host__
TokenLayout get_token_buffer(const int& token_idx, const bool& global = false) const {
EP_UNIFIED_ASSERT(num_ranks == 1 or global);
return TokenLayout(token_layout.num_hidden_bytes, token_layout.num_sf_bytes, token_layout.num_topk, token_layout.with_metadata,
static_cast<int8_t*>(base) + token_layout.get_num_bytes<kWithMBarrier, int64_t>() * token_idx);
}
};
} // namespace deep_ep::elastic
deep_ep/include/deep_ep/common/math.cuh
+68
View File
@@ -0,0 +1,68 @@
#pragma once
namespace deep_ep::elastic::math {
template <typename T>
__forceinline__ __device__ __host__ T ceil_div(T a, T b) {
return (a + b - 1) / b;
}
template <typename T>
__forceinline__ __device__ __host__ constexpr T constexpr_ceil_div(T a, T b) {
return (a + b - 1) / b;
}
template <typename T, bool kDoCeilAlignment = true>
__forceinline__ __device__ __host__ T align(T a, T b) {
return (kDoCeilAlignment ? ceil_div(a, b) : (a / b)) * b;
}
template <typename T, bool kDoCeilAlignment = true>
__forceinline__ __device__ __host__ constexpr T constexpr_align(T a, T b) {
return (kDoCeilAlignment ? constexpr_ceil_div(a, b) : (a / b)) * b;
}
template <typename dtype_t>
__forceinline__ __device__ __host__ bool is_decoded_positive_ready(const dtype_t& value) {
return value >= 0;
}
template <typename dtype_t>
__forceinline__ __device__ __host__ dtype_t encode_decode_positive(const dtype_t& value) {
return -value - static_cast<dtype_t>(1);
}
template <typename dtype_t = void>
__forceinline__ __device__ __host__ dtype_t* advance_ptr(void* ptr, const int64_t num_bytes) {
return reinterpret_cast<dtype_t*>(static_cast<int8_t*>(ptr) + num_bytes);
}
__forceinline__ __device__ __host__ ptrdiff_t ptr_diff(const void* ptr, const void* base) {
return static_cast<const int8_t*>(ptr) - static_cast<const int8_t*>(base);
}
template <typename dtype_a_t, typename dtype_b_t>
__device__ __forceinline__ dtype_b_t pack2(const dtype_a_t& x, const dtype_a_t& y) {
EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), "Invalid dtypes");
dtype_b_t packed;
auto unpacked_ptr = reinterpret_cast<dtype_a_t*>(&packed);
unpacked_ptr[0] = x, unpacked_ptr[1] = y;
return packed;
}
template <typename dtype_a_t, typename dtype_b_t>
__device__ __forceinline__ std::tuple<dtype_a_t, dtype_a_t> unpack2(const dtype_b_t& packed) {
EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), "Invalid dtypes");
auto unpacked_ptr = reinterpret_cast<const dtype_a_t*>(&packed);
dtype_a_t x = unpacked_ptr[0], y = unpacked_ptr[1];
return {x, y};
}
template <typename dtype_a_t, typename dtype_b_t>
__device__ __forceinline__ void unpack2(const dtype_b_t& packed, dtype_a_t& x, dtype_a_t& y) {
EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), "Invalid dtypes");
auto unpacked_ptr = reinterpret_cast<const dtype_a_t*>(&packed);
x = unpacked_ptr[0], y = unpacked_ptr[1];
}
} // namespace deep_ep::elastic::math
deep_ep/include/deep_ep/common/ptx.cuh
+458
View File
@@ -0,0 +1,458 @@
#pragma once
#include <cuda_bf16.h>
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/exception.cuh>
namespace deep_ep::elastic::ptx {
// Host-side placeholder with the same size/alignment as cuda::barrier<thread_scope_block>
// (a single uint64_t atomic), so that sizeof(mbarrier) is consistent across host and device.
struct alignas(8) mbarrier { uint64_t __placeholder; };
using arrival_phase = uint32_t;
// More than TMA, `longlong4` requires 32 bytes aligned
static constexpr int kNumTMAAlignBytes = 32;
#ifdef __CUDACC__
/// Exceptions
__forceinline__ __device__ void trap() {
asm volatile("trap;");
}
/// Thread layout
__forceinline__ __device__ int get_warp_idx() {
return __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
}
__forceinline__ __device__ int get_lane_idx() {
int lane_idx;
asm volatile("mov.s32 %0, %laneid;" : "=r"(lane_idx));
return lane_idx;
}
/// Election
__forceinline__ __device__ int elect_one_sync() {
#ifndef DISABLE_SM90_FEATURES
int pred = 0;
asm volatile(
"{\n"
".reg .b32 %%rx;\n"
".reg .pred %%px;\n"
" elect.sync %%rx|%%px, %1;\n"
"@%%px mov.s32 %0, 1;\n"
"}\n"
: "+r"(pred)
: "r"(0xffffffff));
return pred;
#else
return get_lane_idx() == 0;
#endif
}
/// TMA and `cp.async`
__forceinline__ __device__ void mbarrier_init_with_fence(mbarrier* ptr, const int& arrive_count = 1) {
asm volatile("mbarrier.init.shared::cta.b64 [%1], %0;" ::
"r"(arrive_count), "r"(static_cast<uint32_t>(__cvta_generic_to_shared(ptr))));
asm volatile("fence.mbarrier_init.release.cluster;" ::);
}
__forceinline__ __device__ void mbarrier_invalidate(mbarrier* ptr) {
asm volatile("mbarrier.inval.shared::cta.b64 [%0];" ::
"r"(static_cast<uint32_t>(__cvta_generic_to_shared(ptr))));
}
__forceinline__ __device__ void mbarrier_arrive(mbarrier* ptr) {
asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0]; \n\t" ::
"r"(static_cast<uint32_t>(__cvta_generic_to_shared(ptr))));
}
__forceinline__ __device__ void mbarrier_arrive_and_set_tx(mbarrier* ptr, const int& num_bytes) {
asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%1], %0; \n\t" ::
"r"(num_bytes), "r"(static_cast<uint32_t>(__cvta_generic_to_shared(ptr))));
}
__forceinline__ __device__ void mbarrier_wait_and_flip_phase(mbarrier* ptr, arrival_phase& phase) {
asm volatile(
"{\n\t"
".reg .pred P1; \n\t"
"LAB_WAIT: \n\t"
"mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2; \n\t"
"@P1 bra DONE; \n\t"
"bra LAB_WAIT; \n\t"
"DONE: \n\t"
"}" ::
"r"(static_cast<uint32_t>(__cvta_generic_to_shared(ptr))),
"r"(phase), "r"(0x989680));
phase ^= 1;
}
template <int kNumBytes>
__forceinline__ __device__ void st_bulk(void* smem_ptr) {
EP_STATIC_ASSERT(kNumBytes % 8 == 0, "`st.bulk` requires size to be a multiple of 8");
#if defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)
if (elect_one_sync()) {
asm volatile("st.bulk.weak.shared::cta [%0], %1, 0;\n" ::
"r"(static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr))),
"r"(kNumBytes)
: "memory");
}
#else
#pragma unroll
for (int i = get_lane_idx(); i < kNumBytes / 8; i += 32)
static_cast<uint64_t*>(smem_ptr)[i] = 0;
#endif
__syncwarp();
}
__forceinline__ __device__ void tma_store_fence() {
asm volatile("fence.proxy.async.shared::cta;");
}
template <int kNumRemainingWaits = 0>
__forceinline__ __device__ void tma_store_wait() {
asm volatile("cp.async.bulk.wait_group %0;" ::"n"(kNumRemainingWaits) : "memory");
}
enum TMACacheHint: int64_t {
kEvictFirst = 0x12f0000000000000ll,
kEvictNormal = 0x1000000000000000ll
};
__forceinline__ __device__ void tma_load_1d(
const void* dst_ptr, const void* src_ptr, mbarrier* ptr, const int& num_bytes,
const TMACacheHint& hint = TMACacheHint::kEvictFirst) {
// NOTES: normally, the loaded part will be evicted soon
asm volatile(
"cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;\n" ::
"r"(static_cast<uint32_t>(__cvta_generic_to_shared(dst_ptr))),
"l"(src_ptr),
"r"(num_bytes),
"r"(static_cast<uint32_t>(__cvta_generic_to_shared(ptr))),
"l"(hint)
: "memory");
}
__forceinline__ __device__ void tma_store_1d(
const void* dst_ptr, const void* src_ptr, const int& num_bytes,
const TMACacheHint& hint = TMACacheHint::kEvictNormal) {
// NOTES: normally, the stored part will be used soon
asm volatile("cp.async.bulk.global.shared::cta.bulk_group.L2::cache_hint [%0], [%1], %2, %3;\n" ::
"l"(dst_ptr),
"r"(static_cast<uint32_t>(__cvta_generic_to_shared(src_ptr))),
"r"(num_bytes),
"l"(hint)
: "memory");
}
__forceinline__ __device__ void tma_store_commit() {
asm volatile("cp.async.bulk.commit_group;");
}
template <class dtype_t>
__forceinline__ __device__ void cp_async_ca(const dtype_t* gmem_src, const dtype_t* smem_dst) {
EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8 or sizeof(dtype_t) == 16, "Invalid dtype bytes");
asm volatile("cp.async.ca.shared::cta.global.L2::128B [%0], [%1], %2;\n" ::
"r"(static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst))),
"l"(gmem_src),
"n"(sizeof(dtype_t)));
}
__forceinline__ __device__ void cp_async_mbarrier_arrive(mbarrier* ptr) {
asm volatile("cp.async.mbarrier.arrive.shared::cta.b64 [%0];\n" ::
"r"(static_cast<uint32_t>(__cvta_generic_to_shared(ptr))));
}
/// Barriers
template <int kNumThreads>
__forceinline__ __device__ void named_barrier(const int& idx) {
// Equivalent to `barrier.sync.aligned`, which requires all threads run the same location of code
asm volatile("bar.sync %0, %1;" ::"r"(idx), "r"(kNumThreads));
}
/// LD/ST instructions
__forceinline__ __device__ int4 ldg_with_gez_pred(const int4* ptr, const int& value, const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) {
int4 ret = make_int4(0, 0, 0, 0);
asm volatile(
"{\n\t"
" .reg .pred p;\n\t"
" setp.ge.s32 p, %5, 0;\n\t"
" @p ld.L1::no_allocate.L2::cache_hint.global.nc.v4.s32 {%0, %1, %2, %3}, [%4], %6;\n\t"
"}"
: "+r"(ret.x), "+r"(ret.y), "+r"(ret.z), "+r"(ret.w)
: "l"(ptr), "r"(value), "l"(cache_hint)
: "memory");
return ret;
}
__forceinline__ __device__ int4 ldg_with_gtz_pred(const int4* ptr, const int& value, const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) {
int4 ret = make_int4(0, 0, 0, 0);
asm volatile(
"{\n\t"
" .reg .pred p;\n\t"
" setp.gt.s32 p, %5, 0;\n\t"
" @p ld.L1::no_allocate.L2::cache_hint.global.nc.v4.s32 {%0, %1, %2, %3}, [%4], %6;\n\t"
"}"
: "+r"(ret.x), "+r"(ret.y), "+r"(ret.z), "+r"(ret.w)
: "l"(ptr), "r"(value), "l"(cache_hint)
: "memory");
return ret;
}
__forceinline__ __device__ int4 ld_with_gez_pred(const int4* ptr, const int& value, const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) {
int4 ret = make_int4(0, 0, 0, 0);
asm volatile(
"{\n\t"
" .reg .pred p;\n\t"
" setp.ge.s32 p, %5, 0;\n\t"
" @p ld.L1::no_allocate.L2::cache_hint.global.v4.s32 {%0, %1, %2, %3}, [%4], %6;\n\t"
"}"
: "+r"(ret.x), "+r"(ret.y), "+r"(ret.z), "+r"(ret.w)
: "l"(ptr), "r"(value), "l"(cache_hint)
: "memory");
return ret;
}
#if defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)
__forceinline__ __device__ longlong4_t ldg_with_gez_pred(const longlong4_t* ptr, const int& value, const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) {
longlong4_t ret = make_longlong4_t(0, 0, 0, 0);
asm volatile(
"{\n\t"
" .reg .pred p;\n\t"
" setp.ge.s32 p, %5, 0;\n\t"
" @p ld.L1::no_allocate.L2::cache_hint.global.nc.v4.s64 {%0, %1, %2, %3}, [%4], %6;\n\t"
"}"
: "+l"(ret.x), "+l"(ret.y), "+l"(ret.z), "+l"(ret.w)
: "l"(ptr), "r"(value), "l"(cache_hint)
: "memory");
return ret;
}
__forceinline__ __device__ longlong4_t ldg(const longlong4_t* ptr) {
longlong4_t ret;
asm volatile(
"ld.L1::no_allocate.global.nc.v4.s64 {%0, %1, %2, %3}, [%4];\n\t"
: "=l"(ret.x), "=l"(ret.y), "=l"(ret.z), "=l"(ret.w)
: "l"(ptr)
: "memory");
return ret;
}
#endif
__forceinline__ __device__ int4 ldg(const int4* ptr) {
return __ldg(ptr);
}
template <typename dtype_t>
__forceinline__ __device__ void st_with_gez_pred(dtype_t* ptr, dtype_t value, const int& condition) {
EP_STATIC_ASSERT(sizeof(dtype_t) == 4, "Invalid data type");
auto view = *reinterpret_cast<int*>(&value);
asm volatile(
"{\n\t"
" .reg .pred p;\n\t"
" setp.ge.s32 p, %2, 0;\n\t"
" @p st.global.s32 [%0], %1;\n\t"
"}"
:: "l"(ptr), "r"(view), "r"(condition)
: "memory");
}
template <typename dtype_t>
__forceinline__ __device__ dtype_t ld_volatile(const void* ptr) {
if constexpr (sizeof(dtype_t) == 4) {
uint32_t value;
asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(value) : "l"(ptr));
return reinterpret_cast<const dtype_t&>(value);
} else if constexpr (sizeof(dtype_t) == 8) {
uint64_t value;
asm volatile("ld.volatile.global.u64 %0, [%1];" : "=l"(value) : "l"(ptr));
return reinterpret_cast<const dtype_t&>(value);
} else {
EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, "Invalid data type length");
}
}
__forceinline__ __device__ void red_add(const int64_t* ptr, const int64_t& value) {
// TODO(NVCC): why don't NVCC support `s64`?
asm volatile("red.gpu.global.add.u64 [%0], %1;" :: "l"(ptr), "l"(value));
}
__forceinline__ __device__ void red_add_rel_sys(const int* ptr, const int& value) {
asm volatile("red.release.sys.global.add.s32 [%0], %1;" :: "l"(ptr), "r"(value));
}
__forceinline__ __device__ void red_add_rel_sys(const int64_t* ptr, const int64_t& value) {
asm volatile("red.release.sys.global.add.u64 [%0], %1;" :: "l"(ptr), "l"(value));
}
__forceinline__ __device__ void red_add_rel_gpu(const int* ptr, const int& value) {
asm volatile("red.release.gpu.global.add.s32 [%0], %1;" :: "l"(ptr), "r"(value));
}
__forceinline__ __device__ void red_add_rel_gpu(const int64_t* ptr, const int64_t& value) {
asm volatile("red.release.gpu.global.add.u64 [%0], %1;" :: "l"(ptr), "l"(value));
}
template <typename dtype_t>
__forceinline__ __device__ dtype_t ld_acquire_sys(const dtype_t* ptr) {
if constexpr (sizeof(dtype_t) == 4) {
uint32_t value;
asm volatile("ld.acquire.sys.L1::no_allocate.global.u32 %0, [%1];" : "=r"(value) : "l"(ptr));
return reinterpret_cast<const dtype_t&>(value);
} else if constexpr (sizeof(dtype_t) == 8) {
uint64_t value;
asm volatile("ld.acquire.sys.L1::no_allocate.global.u64 %0, [%1];" : "=l"(value) : "l"(ptr));
return reinterpret_cast<const dtype_t&>(value);
} else {
EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, "Invalid data type length");
}
}
template <typename dtype_t>
__forceinline__ __device__ void st_relaxed_sys(void* ptr, dtype_t value) {
if constexpr (sizeof(dtype_t) == 4) {
uint32_t int_value = reinterpret_cast<const uint32_t&>(value);
asm volatile("st.relaxed.sys.global.u32 [%0], %1;" :: "l"(ptr), "r"(int_value));
} else if constexpr (sizeof(dtype_t) == 8) {
uint64_t int_value = reinterpret_cast<const uint64_t&>(value);
asm volatile("st.relaxed.sys.global.u64 [%0], %1;" :: "l"(ptr), "l"(int_value));
} else {
EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, "Invalid data type length");
}
}
template <typename dtype_t>
__forceinline__ __device__ void st_release_sys(void* ptr, dtype_t value) {
if constexpr (sizeof(dtype_t) == 4) {
uint32_t int_value = reinterpret_cast<const uint32_t&>(value);
asm volatile("st.release.sys.global.u32 [%0], %1;" :: "l"(ptr), "r"(int_value));
} else if constexpr (sizeof(dtype_t) == 8) {
uint64_t int_value = reinterpret_cast<const uint64_t&>(value);
asm volatile("st.release.sys.global.u64 [%0], %1;" :: "l"(ptr), "l"(int_value));
} else {
EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, "Invalid data type length");
}
}
// Adjust registers
template <int kNumRegs>
__device__ __forceinline__ void warpgroup_reg_alloc(){
asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(kNumRegs));
}
template <int kNumRegs>
__device__ __forceinline__ void warpgroup_reg_dealloc(){
asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(kNumRegs));
}
/// General fences
__device__ __forceinline__ void fence_acq_rel_sys() {
asm volatile("fence.acq_rel.sys;" ::: "memory");
}
/// Intrinsics
template <typename dtype_t>
__device__ __forceinline__ dtype_t exchange(dtype_t ptr, const int& src_lane_idx) {
EP_STATIC_ASSERT(sizeof(dtype_t) % sizeof(int) == 0, "");
const auto send_int_values = reinterpret_cast<int*>(&ptr);
dtype_t recv_dtype;
auto recv_int_values = reinterpret_cast<int*>(&recv_dtype);
#pragma unroll
for (int i = 0; i < sizeof(dtype_t) / sizeof(int); ++i)
recv_int_values[i] = __shfl_sync(0xffffffff, send_int_values[i], src_lane_idx);
return recv_dtype;
}
__device__ __forceinline__ unsigned gather(const bool& value) {
return __ballot_sync(0xffffffff, value);
}
__device__ __forceinline__ bool all(const bool& value) {
return __all_sync(0xffffffff, value);
}
__device__ __forceinline__ bool any(const bool& value) {
return __any_sync(0xffffffff, value);
}
__device__ __forceinline__ unsigned reduce_or(const unsigned& value) {
return __reduce_or_sync(0xffffffff, value);
}
__device__ __forceinline__ unsigned long long reduce_or(const unsigned long long& value) {
const auto low = __reduce_or_sync(0xffffffff, static_cast<unsigned>(value));
const auto high = __reduce_or_sync(0xffffffff, static_cast<unsigned>(value >> 32));
return (static_cast<unsigned long long>(high) << 32) | low;
}
__device__ __forceinline__ int reduce_add(const int& value) {
return __reduce_add_sync(0xffffffff, value);
}
__device__ __forceinline__ unsigned match(const int& value) {
return __match_any_sync(0xffffffff, value);
}
__device__ __forceinline__ int fns(const unsigned& value, const int& offset) {
return __fns(value, 0, offset);
}
template <typename dtype_t>
__device__ __forceinline__ auto ffs(const dtype_t& value) {
if constexpr (sizeof(dtype_t) == 4) {
return __ffs(static_cast<int>(value)) - 1;
} else {
EP_STATIC_ASSERT(sizeof(dtype_t) == 8, "Invalid data type");
return __ffsll(static_cast<long long>(value)) - 1;
}
}
__device__ __forceinline__ int get_master_lane_idx(const unsigned& mask) {
// Equivalent to `31 - __clz(mask)`
int highest_idx;
asm volatile("bfind.u32 %0, %1;" : "=r"(highest_idx) : "r"(mask));
return highest_idx;
}
__device__ __forceinline__ bool deduplicate(const int& value, const int& lane_idx) {
return get_master_lane_idx(match(value)) == lane_idx;
}
__device__ __forceinline__ int warp_inclusive_sum(int value, const int& lane_idx) {
#pragma unroll
for (int offset = 1; offset < 32; offset <<= 1) {
const auto synced = __shfl_up_sync(0xffffffff, value, offset);
if (lane_idx >= offset)
value += synced;
}
return value;
}
__device__ __forceinline__ int warp_exclusive_sum(const int& value, const int& lane_idx) {
return warp_inclusive_sum(value, lane_idx) - value;
}
__device__ __forceinline__ float2 fadd2(const float2& a, const float2& b) {
#if defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)
return __fadd2_rn(a, b);
#else
return {a.x + b.x, a.y + b.y};
#endif
}
__device__ __forceinline__ void accumulate(float2& a, nv_bfloat162 b) {
#if defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)
// Use `add.rn.f32.bf16` instruction to perform fused (cast + add) operation on SM100
asm("add.rn.f32.bf16 %0, %1, %0;\n" : "+f"(a.x) : "h"(*reinterpret_cast<uint16_t*>(&b.x)));
asm("add.rn.f32.bf16 %0, %1, %0;\n" : "+f"(a.y) : "h"(*reinterpret_cast<uint16_t*>(&b.y)));
#else
const auto [x, y] = __bfloat1622float2(b);
a.x += x, a.y += y;
#endif
}
#endif
} // namespace deep_ep::elastic::ptx
deep_ep/include/deep_ep/impls/barrier.cuh
+42
View File
@@ -0,0 +1,42 @@
#pragma once
#include <deep_ep/common/comm.cuh>
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/layout.cuh>
#include <deep_ep/common/ptx.cuh>
namespace deep_ep::elastic {
template <bool kIsScaleupNVLink,
int kNumSMs, int kNumThreads,
int kNumScaleoutRanks, int kNumScaleupRanks,
int64_t kNumTimeoutCycles,
bool kSequential>
__global__ void __launch_bounds__(kNumThreads, 1)
barrier_impl(const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window, void* workspace,
const int scaleout_rank_idx, const int scaleup_rank_idx) {
const auto sm_idx = static_cast<int>(blockIdx.x), thread_idx = static_cast<int>(threadIdx.x);
// Barrier only uses the first part of workspace, so making `num_experts` as 0 is fine
const auto workspace_layout = layout::WorkspaceLayout(workspace, kNumScaleoutRanks, kNumScaleupRanks, 0);
const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, 0);
if constexpr (kSequential) {
// Scaleout barrier
if constexpr (kNumScaleoutRanks > 1)
comm::gpu_barrier<kIsScaleupNVLink, kNumScaleoutRanks, kNumScaleupRanks,
kNumSMs, kNumThreads, comm::kFlushAllAllocatedQPs, kNumTimeoutCycles, comm::kKernelBarrierTag, false, false, false>(
gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx, true, false);
// Scaleup barrier, and it needs to flush the RDMA requests issued by scaleout barrier
comm::gpu_barrier<kIsScaleupNVLink, kNumScaleoutRanks, kNumScaleupRanks,
kNumSMs, kNumThreads, comm::kFlushAllAllocatedQPs, kNumTimeoutCycles, comm::kKernelBarrierTag, true, true, false>(
gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx, false, true);
} else {
comm::gpu_barrier<kIsScaleupNVLink, kNumScaleoutRanks, kNumScaleupRanks,
kNumSMs, kNumThreads, comm::kFlushAllAllocatedQPs, kNumTimeoutCycles, comm::kKernelBarrierTag, false, false, false>(
gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx);
}
}
} // namespace deep_ep::elastic
deep_ep/include/deep_ep/impls/combine.cuh
+245
View File
@@ -0,0 +1,245 @@
#pragma once
#include <nccl_device.h>
#include <deep_ep/common/comm.cuh>
#include <deep_ep/common/layout.cuh>
#include <deep_ep/common/math.cuh>
#include <deep_ep/common/ptx.cuh>
#include <deep_ep/impls/combine_utils.cuh>
namespace deep_ep::elastic {
template <bool kIsScaleupNVLink,
bool kUseExpandedLayout, bool kAllowMultipleReduction,
int kNumSMs, int kNumWarps,
int kNumRanks,
int kHidden,
int kNumMaxTokensPerRank,
int kNumExperts, int kNumTopk,
int kNumQPs, int64_t kNumTimeoutCycles,
int kNumThreads = kNumWarps * 32,
int kNumHiddenBytes = kHidden * sizeof(nv_bfloat16),
bool kUseRankLayout = use_rank_layout<kAllowMultipleReduction, kNumRanks, kNumTopk>(),
int kNumTokensInLayout = get_num_tokens_in_layout<kAllowMultipleReduction, kNumRanks, kNumTopk>(),
typename team_t = std::conditional_t<kIsScaleupNVLink, ncclTeamTagLsa, ncclTeamTagWorld>>
__global__ void __launch_bounds__(kNumThreads, 1)
combine_impl(nv_bfloat16* x,
float* topk_weights,
int* src_metadata, int* psum_num_recv_tokens_per_scaleup_rank,
const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window,
void* buffer, void* workspace,
const int rank_idx,
int num_reduced_tokens) {
// Utils
const auto sm_idx = static_cast<int>(blockIdx.x);
const auto thread_idx = static_cast<int>(threadIdx.x);
const auto warp_idx = (ptx::get_warp_idx() + rank_idx) % kNumWarps;
const auto lane_idx = ptx::get_lane_idx();
const auto global_warp_idx = warp_idx * kNumSMs + sm_idx;
constexpr bool kDoExpandedSend = not kAllowMultipleReduction and kUseExpandedLayout;
// We should assign the real number of received tokens if without CPU sync
if (num_reduced_tokens == kNumMaxTokensPerRank * kNumRanks)
num_reduced_tokens = __ldg(psum_num_recv_tokens_per_scaleup_rank + kNumRanks - 1);
// Buffer layouts
extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[];
const auto token_layout = layout::TokenLayout(kNumHiddenBytes, 0, kNumTopk, false);
const auto tma_buffer = layout::BufferLayout<true>(token_layout, kNumWarps, 1, smem)
.get_rank_buffer(warp_idx).get_token_buffer(0);
const auto recv_buffer = layout::BufferLayout<false>(
token_layout, kNumTokensInLayout, kNumMaxTokensPerRank, buffer);
const auto send_buffer = layout::BufferLayout<false>(
token_layout, kNumRanks,
kNumMaxTokensPerRank * (kDoExpandedSend ? kNumTopk : 1),
recv_buffer.get_buffer_end_ptr());
// Init TMA
ptx::arrival_phase phase = 0;
const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr();
if (ptx::elect_one_sync())
ptx::mbarrier_init_with_fence(mbarrier_ptr, 1);
__syncwarp();
// Expanding send mode must not be backward
if constexpr (kDoExpandedSend)
EP_DEVICE_ASSERT(topk_weights == nullptr);
// Gin handle
// We treat each warp as a "channel"
const auto [qp_idx, sharing_mode] = comm::get_qp_mode<kNumSMs, kNumQPs, kNumWarps>(sm_idx, warp_idx);
const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, sharing_mode);
// Full barrier to ensure the remote buffer is available
const auto workspace_layout = layout::WorkspaceLayout(workspace, 1, kNumRanks, kNumExperts);
comm::gpu_barrier<kIsScaleupNVLink, 1, kNumRanks,
kNumSMs, kNumThreads, kNumQPs, kNumTimeoutCycles, comm::kCombineTag0, false, false, true>(
gin, workspace_layout, 0, rank_idx, sm_idx, thread_idx);
// Do TMA writes into the remote buffers
int num_tokens_per_warp = math::ceil_div(num_reduced_tokens, kNumSMs * kNumWarps);
const int token_start_idx = num_tokens_per_warp * global_warp_idx;
const int token_end_idx = min(token_start_idx + num_tokens_per_warp, num_reduced_tokens);
for (int i = token_start_idx; i < token_end_idx; ++ i) {
// The master slot index during dispatch
constexpr int kMetadataStride = 2 + kNumTopk;
const int src_token_idx = __ldg(src_metadata + i * kMetadataStride) % kNumMaxTokensPerRank;
const int src_rank_topk_idx = __ldg(src_metadata + i * kMetadataStride + 1);
const int src_rank_idx = src_rank_topk_idx / kNumTopk;
const int src_topk_idx = src_rank_topk_idx % kNumTopk;
// Directly to the remote or via RDMA
const bool nvlink_bypass = gin.is_nvlink_accessible<team_t>(src_rank_idx);
layout::TokenLayout master_token_buffer = [=]() {
// NVLink bypass
if (nvlink_bypass) {
auto token_buffer = recv_buffer.get_rank_buffer(kUseRankLayout ? rank_idx : src_topk_idx).get_token_buffer(src_token_idx);
token_buffer.set_base_ptr(gin.get_sym_ptr<team_t>(token_buffer.get_base_ptr(), src_rank_idx));
return token_buffer;
}
// Use RDMA
return send_buffer.get_rank_buffer(src_rank_idx).get_token_buffer(src_token_idx);
}();
// Hidden requirements
EP_STATIC_ASSERT(kHidden % (32 * sizeof(int4) / sizeof(nv_bfloat16)) == 0, "Invalid hidden");
using combine_vec_t = typename CombineVecTraits<kHidden * sizeof(nv_bfloat16)>::vec_t;
constexpr int kHiddenVec = kHidden * sizeof(nv_bfloat16) / sizeof(combine_vec_t);
// Read source indices for expand mode
int stored_topk_slot_idx = -1;
if constexpr (kUseExpandedLayout) {
if (lane_idx < kNumTopk)
stored_topk_slot_idx = __ldg(src_metadata + i * kMetadataStride + (2 + lane_idx));
__syncwarp();
}
// 3 cases:
// - no expand + no reduce, or expand + no reduce
// - expand + reduce
// - expand + send all
auto reduce_valid_mask = ptx::gather(stored_topk_slot_idx >= 0);
auto no_local_reduce = not kUseExpandedLayout or (kAllowMultipleReduction and __popc(reduce_valid_mask) == 1);
if (no_local_reduce) {
int token_idx_in_tensor = i;
if constexpr (kUseExpandedLayout)
token_idx_in_tensor = ptx::exchange(stored_topk_slot_idx, ptx::get_master_lane_idx(reduce_valid_mask));
// No reduce
if (ptx::elect_one_sync()) {
const auto load_ptr =
math::advance_ptr(x, static_cast<int64_t>(token_idx_in_tensor) * kNumHiddenBytes);
ptx::tma_store_wait();
ptx::tma_load_1d(tma_buffer.get_base_ptr(), load_ptr, mbarrier_ptr, kNumHiddenBytes);
ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes);
ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase);
ptx::tma_store_1d(master_token_buffer.get_base_ptr(), tma_buffer.get_base_ptr(), kNumHiddenBytes);
ptx::tma_store_commit();
}
__syncwarp();
} else if constexpr (kAllowMultipleReduction) {
// Do local reduction
// Sort valid top-k indices to front
int topk_slot_idx[kNumTopk];
compute_topk_slots(
topk_slot_idx, reduce_valid_mask,
[=](const int& idx) {
return ptx::exchange(stored_topk_slot_idx, idx);
}
);
// Reduce into shared memory
constexpr int kUnrollFactor = get_max_unroll_factor<kHiddenVec, 4>();
combine_reduce<kHiddenVec, kUnrollFactor, math::constexpr_ceil_div(kNumTopk, kNumRanks)>(
lane_idx, topk_slot_idx, static_cast<combine_vec_t*>(tma_buffer.get_base_ptr()),
/* Get source base */ [=](const int& slot_idx) {
return math::advance_ptr<combine_vec_t>(
x, slot_idx * static_cast<int64_t>(kNumHiddenBytes));
},
/* Wait buffer release */ [=]() {
ptx::tma_store_wait();
__syncwarp();
}
);
ptx::tma_store_fence();
__syncwarp();
// Issue TMA stores
if (ptx::elect_one_sync()) {
ptx::tma_store_1d(master_token_buffer.get_base_ptr(), tma_buffer.get_base_ptr(), kNumHiddenBytes);
ptx::tma_store_commit();
}
__syncwarp();
} else {
// No local reduction, send all data (expanded send)
#pragma unroll
for (int k = 0; k < kNumTopk; ++ k) {
const auto slot_idx = ptx::exchange(stored_topk_slot_idx, k);
if (slot_idx >= 0) {
const auto src_token_ptr = math::advance_ptr<int4>(x, slot_idx * static_cast<int64_t>(kNumHiddenBytes));
const auto token_buffer = recv_buffer.get_rank_buffer(k).get_token_buffer(src_token_idx);
if (ptx::elect_one_sync()) {
// Load
ptx::tma_store_wait();
ptx::tma_load_1d(tma_buffer.get_base_ptr(), src_token_ptr, mbarrier_ptr, kNumHiddenBytes);
ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes);
ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase);
if (nvlink_bypass) {
// Write into the same position
ptx::tma_store_1d(gin.get_sym_ptr<team_t>(token_buffer.get_base_ptr(), src_rank_idx),
tma_buffer.get_base_ptr(), kNumHiddenBytes);
ptx::tma_store_commit();
} else {
// Write to the RDMA send buffer
const auto send_token_buffer =
send_buffer.get_rank_buffer(src_rank_idx).get_token_buffer(src_token_idx * kNumTopk + k);
ptx::tma_store_1d(send_token_buffer.get_base_ptr(), tma_buffer.get_base_ptr(), kNumHiddenBytes);
ptx::tma_store_commit();
ptx::tma_store_wait();
// Issue RDMA
gin.put<team_t>(token_buffer.get_base_ptr(), send_token_buffer.get_base_ptr(),
kNumHiddenBytes, src_rank_idx);
}
}
__syncwarp();
}
}
}
// Write topk weights
if (not kDoExpandedSend and topk_weights != nullptr and lane_idx < kNumTopk) {
float value = 0;
if constexpr (kUseExpandedLayout) {
if (stored_topk_slot_idx >= 0)
value = __ldg(topk_weights + stored_topk_slot_idx);
} else {
value = __ldg(topk_weights + (i * kNumTopk + lane_idx));
}
master_token_buffer.get_topk_weights_ptr()[lane_idx] = value;
}
__syncwarp();
// Wait send buffer's TMA store and issue RDMA send
// NOTES: `kDoExpandedSend` mode has already issued
if (not kDoExpandedSend and not nvlink_bypass and ptx::elect_one_sync()) {
ptx::tma_store_wait();
const auto dst_ptr = recv_buffer.get_rank_buffer(kUseRankLayout ? rank_idx : src_topk_idx)
.get_token_buffer(src_token_idx).get_base_ptr();
gin.put<team_t>(dst_ptr, master_token_buffer.get_base_ptr(),
master_token_buffer.get_num_bytes<false>(), src_rank_idx);
}
}
// Final barrier to ensure data arrival
comm::gpu_barrier<kIsScaleupNVLink, 1, kNumRanks,
kNumSMs, kNumThreads, kNumQPs, kNumTimeoutCycles, comm::kCombineTag1, true, true, false>(
gin, workspace_layout, 0, rank_idx, sm_idx, thread_idx);
}
} // deep_ep::elastic
deep_ep/include/deep_ep/impls/combine_reduce_epilogue.cuh
+145
View File
@@ -0,0 +1,145 @@
#pragma once
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/ptx.cuh>
#include <deep_ep/common/layout.cuh>
#include <deep_ep/impls/combine_utils.cuh>
namespace deep_ep::elastic {
template <bool kUseExpandedLayout, bool kAllowMultipleReduction,
int kNumSMs, int kNumWarps,
// TODO: merge these two variables into one (ensure the whole file does not contain "scaleup")
int kNumScaleoutRanks, int kNumScaleupRanks,
int kHidden,
int kNumMaxTokensPerRank,
int kNumExperts, int kNumTopk,
int kNumThreads = kNumWarps * 32,
int kNumHiddenBytes = kHidden * sizeof(nv_bfloat16),
int kNumRanks = kNumScaleoutRanks == 1 ? kNumScaleupRanks : kNumScaleoutRanks,
bool kUseRankLayout = use_rank_layout<kAllowMultipleReduction, kNumRanks, kNumTopk>(),
int kNumTokensInLayout = get_num_tokens_in_layout<kAllowMultipleReduction, kNumRanks, kNumTopk>()>
__global__ void __launch_bounds__(kNumThreads, 1)
combine_reduce_epilogue_impl(nv_bfloat16* combined_x,
float* combined_topk_weights,
topk_idx_t* combined_topk_idx,
void* recv_buffer,
void* bias_0, void* bias_1,
const int num_combined_tokens,
const int scaleout_rank_idx, const int scaleup_rank_idx) {
constexpr int kNumExpertsPerScaleout = kNumExperts / kNumScaleoutRanks;
constexpr int kNumExpertsPerRank = kNumExperts / (kNumScaleupRanks * kNumScaleoutRanks);
EP_STATIC_ASSERT(kNumExperts % (kNumScaleupRanks * kNumScaleoutRanks) == 0, "Invalid number of experts or ranks");
// Utils
const auto sm_idx = static_cast<int>(blockIdx.x);
const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx();
const auto global_warp_idx = warp_idx * kNumSMs + sm_idx; // NOTES: Here we prioritize distributing tasks to different SMs to ensure that the last wave is evenly concentrated on each SM.
// Load buffers from scale-out or scale-up ranks
extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[];
const auto comm_token_layout = layout::TokenLayout(kNumHiddenBytes, 0, kNumTopk, false);
const auto comm_buffer = layout::BufferLayout<false>(
comm_token_layout, kNumTokensInLayout, kNumMaxTokensPerRank, recv_buffer);
// Store buffers
const auto output_token_layout = layout::TokenLayout(kNumHiddenBytes, 0, 0, false);
const auto output_buffer = layout::BufferLayout<false>(output_token_layout, 1, num_combined_tokens, combined_x);
const auto tma_buffer = layout::BufferLayout<false>(output_token_layout, kNumWarps, 1, smem)
.get_rank_buffer(warp_idx).get_token_buffer(0);
// Bias layout
const auto bias_0_buffer = layout::BufferLayout<false>(output_token_layout, 1, num_combined_tokens, bias_0);
const auto bias_1_buffer = layout::BufferLayout<false>(output_token_layout, 1, num_combined_tokens, bias_1);
// Will block until the main combine kernel has finished and all data are visible
// NOTES: PDL is used, please do not use `__ldg`
cudaGridDependencySynchronize();
// Read from buffers and do reduction
for (int token_idx = global_warp_idx; token_idx < num_combined_tokens; token_idx += kNumWarps * kNumSMs) {
// Preprocess all indices
int stored_dst_rank_idx = -1, stored_dst_expert_idx = -1;
EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k selections");
if (lane_idx < kNumTopk) {
stored_dst_expert_idx = static_cast<int>(combined_topk_idx[token_idx * kNumTopk + lane_idx]);
stored_dst_rank_idx = stored_dst_expert_idx >= 0 ?
stored_dst_expert_idx / (kNumScaleoutRanks == 1 ? kNumExpertsPerRank : kNumExpertsPerScaleout) : -1;
}
__syncwarp();
// Sort valid top-k indices to front
const auto [should_deduplicate, deduplicate_key] = [&]() -> std::pair<bool, int> {
if constexpr (kUseExpandedLayout and not kAllowMultipleReduction) {
// Activations are never reduced before
return {false, 0};
} else if constexpr (kNumScaleoutRanks != 1 and not kUseExpandedLayout and not kAllowMultipleReduction) {
// Hybrid mode without expanded layout and multiple reduction. Should deduplicate on a per-rank basis
return {true, stored_dst_expert_idx >= 0 ? stored_dst_expert_idx / kNumExpertsPerRank : -1};
} else {
// Should deduplicate on a per-rank (for non-hybrid mode) or a per-scale-rank (for hybrid mode) basis
return {true, stored_dst_rank_idx};
}
}();
auto reduce_valid_mask = should_deduplicate ?
ptx::gather(ptx::deduplicate(deduplicate_key, lane_idx) and stored_dst_rank_idx >= 0) :
ptx::gather(stored_dst_rank_idx >= 0);
int topk_slot_idx[kNumTokensInLayout];
compute_topk_slots(
topk_slot_idx, reduce_valid_mask,
[=](const int& idx) {
return kUseRankLayout ? ptx::exchange(stored_dst_rank_idx, idx) : idx;
}
);
// Iterate over per-hidden-chunk stage
using combine_vec_t = typename CombineVecTraits<kHidden * sizeof(nv_bfloat16)>::vec_t;
constexpr int kHiddenVec = kHidden * sizeof(nv_bfloat16) / sizeof(combine_vec_t);
constexpr int kUnrollFactor = get_max_unroll_factor<kHiddenVec, 4>();
combine_reduce<kHiddenVec, kUnrollFactor, kNumTokensInLayout>(
lane_idx, topk_slot_idx, static_cast<combine_vec_t*>(tma_buffer.get_base_ptr()),
/* Get source base */ [=](const int& slot_idx) {
return static_cast<combine_vec_t*>(
comm_buffer.get_rank_buffer(slot_idx).get_token_buffer(token_idx).get_base_ptr());
},
/* Wait buffer release */ [=]() {
ptx::tma_store_wait();
__syncwarp();
},
/* Bias 0 */ bias_0 == nullptr ?
nullptr : static_cast<combine_vec_t*>(bias_0_buffer.get_token_buffer(token_idx).get_base_ptr()),
/* Bias 1 */ bias_1 == nullptr ?
nullptr : static_cast<combine_vec_t*>(bias_1_buffer.get_token_buffer(token_idx).get_base_ptr())
);
ptx::tma_store_fence();
__syncwarp();
// Issue TMA copy
if (ptx::elect_one_sync()) {
ptx::tma_store_1d(output_buffer.get_token_buffer(token_idx).get_base_ptr(),
tma_buffer.get_base_ptr(), kNumHiddenBytes);
ptx::tma_store_commit();
}
__syncwarp();
// Write top-k weights
if (combined_topk_weights != nullptr) {
const auto master_lane_idx = ptx::get_master_lane_idx(ptx::match(stored_dst_rank_idx));
if (lane_idx < kNumTopk) {
float value = 0;
if (stored_dst_rank_idx >= 0) {
const auto dst_ptr = comm_buffer
.get_rank_buffer(kUseRankLayout ? stored_dst_rank_idx : master_lane_idx)
.get_token_buffer(token_idx).get_topk_weights_ptr() + lane_idx;
value = *dst_ptr;
}
combined_topk_weights[token_idx * kNumTopk + lane_idx] = value;
}
__syncwarp();
}
}
}
} // deep_ep::elastic
deep_ep/include/deep_ep/impls/combine_utils.cuh
+172
View File
@@ -0,0 +1,172 @@
#pragma once
#include <stdexcept>
namespace deep_ep::elastic {
template <bool kAllowMultipleReduction, int kNumRanks, int kNumTopk>
constexpr bool use_rank_layout() {
if constexpr (not kAllowMultipleReduction)
return false;
return kNumRanks <= kNumTopk;
}
template <bool kAllowMultipleReduction, int kNumRanks, int kNumTopk>
constexpr int get_num_tokens_in_layout() {
return use_rank_layout<kAllowMultipleReduction, kNumRanks, kNumTopk>() ? kNumRanks : kNumTopk;
}
template <int kLength, int kMaxUnrollFactor, int kWarpSize = 32>
constexpr int get_max_unroll_factor() {
for (int i = kMaxUnrollFactor; i >= 1; -- i)
if (kLength % (kWarpSize * i) == 0)
return i;
throw std::logic_error("Invalid length, cannot find unrolling factor");
}
// Determine the vector type for combine loads/stores based on arch and hidden size alignment
template <int kHiddenBytes>
struct CombineVecTraits {
#if defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)
// On SM100+, use `longlong4_t` (32 bytes) if hidden is aligned, otherwise fall back to `int4` (16 bytes)
// NOTES: we observe some performance degrade with `longlong4_t`, temporarily disable it
static constexpr bool kUseLonglong4 = false;
using vec_t = std::conditional_t<kUseLonglong4, longlong4_t, int4>;
#else
using vec_t = int4;
#endif
};
template <int kNumValidTopk, typename fetch_func_t>
__device__ __forceinline__
void compute_topk_slots(int (&topk_slot_idx)[kNumValidTopk], uint32_t mask,
const fetch_func_t& fetch_func) {
#pragma unroll
for (int k = 0; k < kNumValidTopk; ++ k) {
const int lowest_idx = __ffs(mask) - 1;
// Here we perform the exchange unconditionally to avoid `BRA.DIV`
const auto fetched = fetch_func(lowest_idx);
mask &= mask - 1;
topk_slot_idx[k] = lowest_idx >= 0 ? fetched : -1;
}
}
template <int kHiddenVec, int kUnrollFactor, int kNumExpectedTopk, int kNumValidTopk,
typename vec_t, typename get_src_buffer_ptr_func_t, typename wait_buffer_func_t>
__device__ __forceinline__
void combine_reduce(const int& lane_idx, int (&topk_slot_idx)[kNumValidTopk],
vec_t* dst_buffer_ptr,
const get_src_buffer_ptr_func_t& get_src_buffer_ptr_func,
const wait_buffer_func_t& wait_buffer_func,
vec_t* bias_0 = nullptr, vec_t* bias_1 = nullptr) {
constexpr int kNumElemsPerVec = sizeof(vec_t) / sizeof(nv_bfloat16);
EP_STATIC_ASSERT(kNumElemsPerVec % 2 == 0, "Invalid number of elements");
EP_STATIC_ASSERT(kHiddenVec % (kUnrollFactor * 32) == 0, "Invalid unrolling");
// We use BF16 add as much as possible, as casting is slow
const bool enable_hadd_bypass =
(bias_0 == nullptr and bias_1 == nullptr) and
(kNumValidTopk <= 2 or topk_slot_idx[2] < 0);
EP_STATIC_ASSERT(kNumValidTopk > 0, "Invalid top-k");
if (enable_hadd_bypass) {
#pragma unroll 1
for (int i = 0; i < kHiddenVec / (kUnrollFactor * 32); ++ i) {
// Read values 0
const auto slot_0 = topk_slot_idx[0];
const auto src_base_ptr_0 = get_src_buffer_ptr_func(slot_0);
vec_t values_0[kUnrollFactor] = {};
#pragma unroll
for (int j = 0; j < kUnrollFactor; ++ j) {
values_0[j] = ptx::ldg_with_gez_pred(
src_base_ptr_0 + (i * (kUnrollFactor * 32) + j * 32 + lane_idx), slot_0);
}
// Read values 1
vec_t values_1[kUnrollFactor] = {};
const auto slot_1 = kNumValidTopk == 1 ? -1 : topk_slot_idx[1];
const auto src_base_ptr_1 = get_src_buffer_ptr_func(slot_1);
#pragma unroll
for (int j = 0; j < kUnrollFactor; ++ j) {
values_1[j] = ptx::ldg_with_gez_pred(
src_base_ptr_1 + (i * (kUnrollFactor * 32) + j * 32 + lane_idx), slot_1);
}
// Wait buffer releases for the first write
if (i == 0)
wait_buffer_func();
// Reduce into shared memory
const auto bf162_view_0 = reinterpret_cast<nv_bfloat162*>(values_0);
const auto bf162_view_1 = reinterpret_cast<nv_bfloat162*>(values_1);
#pragma unroll
for (int j = 0; j < kUnrollFactor; ++ j) {
#pragma unroll
for (int l = 0; l < kNumElemsPerVec / 2; ++ l)
bf162_view_0[j * (kNumElemsPerVec / 2) + l] += bf162_view_1[j * (kNumElemsPerVec / 2) + l];
dst_buffer_ptr[i * (kUnrollFactor * 32) + j * 32 + lane_idx] = values_0[j];
}
}
} else {
#pragma unroll 1
for (int i = 0; i < kHiddenVec / (kUnrollFactor * 32); ++ i) {
// Add bias
float2 reduced[kUnrollFactor * kNumElemsPerVec / 2] = {};
const auto add_bias = [&](const vec_t* base_ptr) {
// Read
vec_t values[kUnrollFactor];
#pragma unroll
for (int j = 0; j < kUnrollFactor; ++ j)
values[j] = ptx::ldg(base_ptr + i * (kUnrollFactor * 32) + j * 32 + lane_idx);
// Reduce
const auto bf162_view = reinterpret_cast<nv_bfloat162*>(values);
#pragma unroll
for (int j = 0; j < kUnrollFactor * kNumElemsPerVec / 2; ++ j)
ptx::accumulate(reduced[j], bf162_view[j]);
};
bias_0 != nullptr ? add_bias(bias_0) : void();
bias_1 != nullptr ? add_bias(bias_1) : void();
#pragma unroll
for (int k = 0; k < kNumValidTopk; ++ k) {
// We have a limitation on `k` to reduce the branch instruction count
if (k >= kNumExpectedTopk and topk_slot_idx[k] < 0)
break;
// Read values
const auto src_base_ptr = get_src_buffer_ptr_func(topk_slot_idx[k]);
vec_t values[kUnrollFactor] = {};
#pragma unroll
for (int j = 0; j < kUnrollFactor; ++ j) {
values[j] = ptx::ldg_with_gez_pred(
src_base_ptr + (i * (kUnrollFactor * 32) + j * 32 + lane_idx), topk_slot_idx[k]);
}
// Reduce
const auto bf162_view = reinterpret_cast<nv_bfloat162*>(values);
#pragma unroll
for (int j = 0; j < kUnrollFactor * kNumElemsPerVec / 2; ++ j)
ptx::accumulate(reduced[j], bf162_view[j]);
}
// Wait buffer releases for the first write
if (i == 0)
wait_buffer_func();
// Cast into shared memory
#pragma unroll
for (int j = 0; j < kUnrollFactor; ++ j) {
vec_t casted_value;
auto bf162_view = reinterpret_cast<nv_bfloat162*>(&casted_value);
#pragma unroll
for (int l = 0; l < kNumElemsPerVec / 2; ++ l)
bf162_view[l] = __float22bfloat162_rn(reduced[j * (kNumElemsPerVec / 2) + l]);
dst_buffer_ptr[i * (kUnrollFactor * 32) + j * 32 + lane_idx] = casted_value;
}
}
}
}
} // namespace deep_ep::elastic
deep_ep/include/deep_ep/impls/dispatch.cuh
+411
View File
@@ -0,0 +1,411 @@
#pragma once
#include <nccl.h>
#include <nccl_device.h>
#include <deep_ep/common/comm.cuh>
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/exception.cuh>
#include <deep_ep/common/handle.cuh>
#include <deep_ep/common/layout.cuh>
#include <deep_ep/common/math.cuh>
#include <deep_ep/common/ptx.cuh>
namespace deep_ep::elastic {
template <bool kIsScaleupNVLink,
bool kDoCPUSync,
bool kReuseSlotIndices,
int kNumSMs,
int kNumNotifyWarps, int kNumDispatchWarps,
int kNumRanks,
int kNumHiddenBytes, int kNumSFPacks,
int kNumMaxTokensPerRank,
int kNumExperts, int kNumTopk, int kExpertAlignment,
int kNumQPs, int64_t kNumTimeoutCycles,
int kNumNotifyThreads = kNumNotifyWarps * 32,
int kNumDispatchThreads = kNumDispatchWarps * 32,
int kNumThreads = kNumNotifyThreads + kNumDispatchThreads,
typename team_t = std::conditional_t<kIsScaleupNVLink, ncclTeamTagLsa, ncclTeamTagWorld>>
__global__ void __launch_bounds__(kNumThreads, 1)
dispatch_impl(
void* x, sf_pack_t* sf, topk_idx_t* topk_idx, float* topk_weights,
topk_idx_t* copied_topk_idx,
int* cumulative_local_expert_recv_stats,
int* psum_num_recv_tokens_per_scaleup_rank,
int* psum_num_recv_tokens_per_expert,
int* num_unaligned_recv_tokens_per_expert,
int* dst_buffer_slot_idx,
const int num_tokens,
const int sf_token_stride, const int sf_hidden_stride,
const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window, void* buffer,
void* workspace, void* mapped_host_workspace,
const int rank_idx
) {
constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks;
EP_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks");
EP_STATIC_ASSERT(kNumNotifyWarps % 4 == 0, "Invalid warpgroup size");
// Utils
const auto sm_idx = static_cast<int>(blockIdx.x), thread_idx = static_cast<int>(threadIdx.x);
const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx();
// Workspaces
const auto workspace_layout = layout::WorkspaceLayout(workspace, 1, kNumRanks, kNumExperts);
const auto host_workspace_layout = layout::WorkspaceLayout(mapped_host_workspace, 1, kNumRanks, kNumExperts);
// The kernel uses a fixed space of dynamic shared memory (no static shared memory)
extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[];
constexpr int kNumSmemBytesForNotify = kNumNotifyThreads > 0 ?
math::constexpr_align(kNumRanks + kNumExperts, kNumNotifyThreads) * sizeof(int) : 0;
EP_STATIC_ASSERT(kNumSmemBytesForNotify % ptx::kNumTMAAlignBytes == 0, "Invalid TMA alignment");
// Named barrier indices
constexpr int kNotifyBarrierIndex = 1;
// Gin handle
// We treat each warp as a "channel"
const auto [qp_idx, sharing_mode] = comm::get_qp_mode<kNumSMs, kNumQPs, kNumDispatchWarps, (kNumNotifyWarps > 0)>(
sm_idx, warp_idx - kNumNotifyWarps, warp_idx < kNumNotifyWarps);
const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, sharing_mode);
// Barrier without TMA store flush, without prologue grid sync
comm::gpu_barrier<kIsScaleupNVLink, 1, kNumRanks,
kNumSMs, kNumThreads, kNumQPs, kNumTimeoutCycles, comm::kDispatchTag0, false, false, true>(
gin, workspace_layout, 0, rank_idx, sm_idx, thread_idx);
// Different warp roles
if (warp_idx < kNumNotifyWarps) {
// Assign shared memory
constexpr int kNumAlignedElems = kNumSmemBytesForNotify / sizeof(int);
const auto rank_expert_count = math::advance_ptr<int>(smem, 0);
// Clean initial counts
// NOTES: if you want to change the order of different warp roles, please take care of the `thread_idx`
int *rank_count = rank_expert_count, *expert_count = rank_expert_count + kNumRanks;
#pragma unroll
for (int i = 0; i < kNumAlignedElems / kNumNotifyThreads; ++ i)
rank_expert_count[i * kNumNotifyThreads + thread_idx] = 0;
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);
// Atomic add on shared memory
EP_STATIC_ASSERT(kNumTopk <= 32, "Insufficient lanes");
const auto global_warp_idx = warp_idx * kNumSMs + sm_idx;
for (int i = global_warp_idx; i < num_tokens; i += kNumNotifyWarps * kNumSMs) {
// Expert choice can not be redundant
// NOTES: no assertions here as they are expensive
const auto dst_expert_idx = lane_idx < kNumTopk ?
static_cast<int>(__ldg(topk_idx + i * kNumTopk + lane_idx)) : -1;
if (dst_expert_idx >= 0)
atomicAdd_block(expert_count + dst_expert_idx, 1);
// Rank choice should do deduplication here
const auto dst_rank_idx = dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerRank : -1;
if (ptx::deduplicate(dst_rank_idx, lane_idx) and dst_rank_idx >= 0)
atomicAdd_block(rank_count + dst_rank_idx, 1);
}
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);
// Do full-grid reduction
#pragma unroll
for (int i = thread_idx; i < kNumRanks + kNumExperts; i += kNumNotifyThreads) {
const int64_t counter = (1ll << 32ll) | rank_expert_count[i];
ptx::red_add(workspace_layout.get_notify_reduction_workspace_ptr() + i, counter);
}
// Do the remaining work by SM 0
if (sm_idx == 0) {
// Reduce all SM's count
// Wait all SMs' arrival
#pragma unroll
for (int i = thread_idx; i < kNumRanks + kNumExperts; i += kNumNotifyThreads) {
comm::timeout_while<kNumTimeoutCycles>(true, [=](const bool& is_last_check) {
const auto status = ptx::ld_volatile<int64_t>(workspace_layout.get_notify_reduction_workspace_ptr() + i);
if ((status >> 32) == kNumSMs) {
// Write into shared memory
// Write into send buffer if with RDMA
const auto encoded =
math::encode_decode_positive(static_cast<int>(status & 0xffffffffll));
rank_expert_count[i] = encoded;
if constexpr (not kIsScaleupNVLink)
workspace_layout.get_scaleup_rank_expert_count_ptr<true>()[i] = encoded;
// Clean for the next usage
workspace_layout.get_notify_reduction_workspace_ptr()[i] = 0;
return true;
}
if (is_last_check) {
printf("DeepEP notify (GPU reduction) timeout, rank: %d/%d, "
"thread: %d, status: %d | %d, expected: %d\n",
rank_idx, kNumRanks, thread_idx,
static_cast<int>(status >> 32), static_cast<int>(status & 0xffffffff), kNumSMs);
}
return false;
});
}
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);
// TODO: for further optimization, we can fuse rank and expert counters
// Issue scaleup rank count writes to peers
for (int i = thread_idx; i < kNumRanks; i += kNumNotifyThreads) {
// Rank counters
const auto dst_rank_counter =
workspace_layout.get_scaleup_rank_count_ptr<false>() + rank_idx;
gin.put_value<team_t>(dst_rank_counter, static_cast<int64_t>(rank_count[i]), i,
ncclGinOptFlagsAggregateRequests);
}
__syncwarp();
// Issue scaleup expert count writes to peers
if constexpr (kIsScaleupNVLink) {
// NVLink per-element copy
// We don't use TMA as the dtype of shared memory and global is different
for (int i = thread_idx; i < kNumExperts; i += kNumNotifyThreads) {
const auto idx = kNumExpertsPerRank * rank_idx + (i % kNumExpertsPerRank);
gin.put_value<team_t>(
workspace_layout.get_scaleup_expert_count_ptr<false>() + idx,
static_cast<int64_t>(expert_count[i]), i / kNumExpertsPerRank);
}
} else {
// RDMA bulk copy
for (int i = thread_idx; i < kNumRanks; i += kNumNotifyThreads) {
const auto src_ptr = workspace_layout.get_scaleup_expert_count_ptr<true>() + kNumExpertsPerRank * i;
const auto dst_ptr = workspace_layout.get_scaleup_expert_count_ptr<false>() + kNumExpertsPerRank * rank_idx;
gin.put<team_t>(dst_ptr, src_ptr, kNumExpertsPerRank * sizeof(int64_t), i);
}
}
// This is necessary, as the waited results will rewrite the shared memory
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);
// Wait for rank and expert count
const auto start_clock = clock64();
for (int i = thread_idx; i < kNumRanks + kNumExperts; i += kNumNotifyThreads) {
comm::timeout_while<kNumTimeoutCycles>([=](const bool& is_last_check) {
// NOTES: the global memory type has 64 bits
const auto count = static_cast<int>(
ptx::ld_volatile<int64_t>(workspace_layout.get_scaleup_rank_expert_count_ptr<false>() + i));
const auto decoded = math::encode_decode_positive(count);
if (math::is_decoded_positive_ready(decoded)) {
workspace_layout.get_scaleup_rank_expert_count_ptr<false>()[i] = 0;
rank_expert_count[i] = decoded;
return true;
}
if (is_last_check)
printf("DeepEP notify timeout, rank: %d, thread: %d, count: %d\n", rank_idx, i, decoded);
return false;
}, start_clock);
}
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);
// Reduce expert count and add stats
for (int i = thread_idx; i < kNumExpertsPerRank; i += kNumNotifyThreads) {
int sum = 0;
#pragma unroll
for (int j = 0; j < kNumRanks; ++ j)
sum += expert_count[j * kNumExpertsPerRank + i];
// Write unaligned count before aligning
if (num_unaligned_recv_tokens_per_expert != nullptr)
num_unaligned_recv_tokens_per_expert[i] = sum;
expert_count[i] = math::align(sum, kExpertAlignment);
// Update statistics counters
if (cumulative_local_expert_recv_stats != nullptr)
atomicAdd(cumulative_local_expert_recv_stats + i, sum);
}
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);
// Write host workspace
if constexpr (kDoCPUSync) {
for (int i = thread_idx; i < kNumRanks + kNumExpertsPerRank; i += kNumNotifyThreads) {
host_workspace_layout.get_scaleup_rank_expert_count_ptr<false>()[i] =
math::encode_decode_positive(rank_expert_count[i]);
}
__syncwarp();
}
// Do prefix sum by the warps
// NOTES: we may have fast implementation with `cub::BlockScan`, but it is too heavy to use
const auto do_psum = [=](const int* count, int* out, const int n, const int is_exclusive) {
int psum = 0;
#pragma unroll
for (int i = 0; i < math::ceil_div(n + is_exclusive, 32); ++ i) {
const auto idx = i * 32 + lane_idx;
const auto mem_idx = idx - is_exclusive;
const auto value = (0 <= mem_idx and mem_idx < n) ? count[mem_idx] : 0;
const auto sum = psum + ptx::warp_inclusive_sum(value, lane_idx);
// Store into global memory
if (idx < n + is_exclusive)
out[idx] = sum;
// Update `psum` by using the last lane's value
psum = ptx::exchange(sum, 31);
}
};
if (warp_idx == 0) {
// Inclusive prefix sum
do_psum(rank_count, psum_num_recv_tokens_per_scaleup_rank, kNumRanks, 0);
} else if (warp_idx == 1) {
// Exclusive prefix sum for later expanding
do_psum(expert_count, psum_num_recv_tokens_per_expert, kNumExpertsPerRank, 1);
}
}
} else {
const int dispatch_warp_idx = warp_idx - kNumNotifyWarps;
// Buffer layouts
const auto token_layout = layout::TokenLayout(kNumHiddenBytes, kNumSFPacks * sizeof(sf_pack_t), kNumTopk, true);
const auto tma_buffer = layout::BufferLayout<true>(token_layout, kNumDispatchWarps, 1,
math::advance_ptr<int>(smem, kNumSmemBytesForNotify)).get_rank_buffer(dispatch_warp_idx).get_token_buffer(0);
auto recv_buffer = layout::BufferLayout<false>(token_layout, kNumRanks, kNumMaxTokensPerRank, buffer);
auto send_buffer = layout::BufferLayout<false>(token_layout, 1, kNumMaxTokensPerRank, recv_buffer.get_buffer_end_ptr());
recv_buffer = recv_buffer.get_rank_buffer(rank_idx);
// Init TMA
ptx::arrival_phase phase = 0;
const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr();
if (ptx::elect_one_sync())
ptx::mbarrier_init_with_fence(mbarrier_ptr, 1);
__syncwarp();
// Iterate all tokens
const auto token_start = dispatch_warp_idx * kNumSMs + sm_idx;
const auto token_stride = kNumDispatchWarps * kNumSMs;
for (int token_idx = token_start; token_idx < num_tokens; token_idx += token_stride) {
const auto token_i64_idx = static_cast<int64_t>(token_idx);
// Wait TMA store arrivals
ptx::tma_store_wait();
__syncwarp();
// Issue data TMA
if (ptx::elect_one_sync()) {
ptx::tma_load_1d(tma_buffer.get_hidden_ptr(), math::advance_ptr(x, token_i64_idx * kNumHiddenBytes),
mbarrier_ptr, kNumHiddenBytes);
}
__syncwarp();
// Issue SF TMA or cp.async
if constexpr (kNumSFPacks > 0) {
EP_STATIC_ASSERT(sizeof(sf_pack_t) % 4 == 0, "Unaligned SF element type");
const auto gmem_src_ptr = math::advance_ptr<sf_pack_t>(sf, token_i64_idx * sf_token_stride * sizeof(sf_pack_t));
const auto smem_dst_ptr = tma_buffer.get_sf_ptr();
constexpr auto kNumFullIters = kNumSFPacks / 32;
#pragma unroll
for (int k = 0; k < kNumFullIters; ++ k) {
ptx::cp_async_ca(gmem_src_ptr + (k * 32 + lane_idx) * sf_hidden_stride,
smem_dst_ptr + k * 32 + lane_idx);
}
if (kNumFullIters * 32 + lane_idx < kNumSFPacks) {
ptx::cp_async_ca(gmem_src_ptr + (kNumFullIters * 32 + lane_idx) * sf_hidden_stride,
smem_dst_ptr + kNumFullIters * 32 + lane_idx);
}
ptx::cp_async_mbarrier_arrive(mbarrier_ptr);
__syncwarp();
}
// Load top-k indices and weights
EP_STATIC_ASSERT(kNumTopk <= 32, "Insufficient lanes for loading top-k indices");
int stored_dst_rank_idx = -1;
if (lane_idx < kNumTopk) {
const auto uncasted_dst_expert_idx = __ldg(topk_idx + token_idx * kNumTopk + lane_idx);
const auto dst_expert_idx = static_cast<int>(uncasted_dst_expert_idx);
stored_dst_rank_idx = dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerRank : -1;
tma_buffer.get_topk_idx_ptr()[lane_idx] = dst_expert_idx;
if (topk_weights != nullptr)
tma_buffer.get_topk_weights_ptr()[lane_idx] = __ldg(topk_weights + token_idx * kNumTopk + lane_idx);
if (copied_topk_idx != nullptr)
copied_topk_idx[token_idx * kNumTopk + lane_idx] = uncasted_dst_expert_idx;
}
__syncwarp();
// Add source metadata (rank index and token index)
// Please ensure no TMA buffer shared memory writes after this part
if (ptx::elect_one_sync())
*tma_buffer.get_src_token_global_idx_ptr() = rank_idx * kNumMaxTokensPerRank + token_idx;
ptx::tma_store_fence();
__syncwarp();
// Deduplicate ranks and assign slots
int stored_dst_slot_idx = -1;
if constexpr (kReuseSlotIndices) {
if (lane_idx < kNumTopk)
stored_dst_slot_idx = __ldg(dst_buffer_slot_idx + token_idx * kNumTopk + lane_idx);
stored_dst_slot_idx = stored_dst_slot_idx >= 0 ?
(stored_dst_slot_idx - rank_idx * kNumMaxTokensPerRank) : -1;
} else {
if (ptx::deduplicate(stored_dst_rank_idx, lane_idx) and stored_dst_rank_idx >= 0)
stored_dst_slot_idx = atomicAdd(workspace_layout.get_scaleup_atomic_sender_counter() + stored_dst_rank_idx, 1);
if (lane_idx < kNumTopk) {
const auto value = stored_dst_slot_idx >= 0 ?
rank_idx * kNumMaxTokensPerRank + stored_dst_slot_idx : -1;
dst_buffer_slot_idx[token_idx * kNumTopk + lane_idx] = value;
}
}
__syncwarp();
// Wait TMA load arrival
// NOTES: this arrive must be after the `ptx::cp_async_mbarrier_arrive`
if (ptx::elect_one_sync()) {
ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes);
ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase);
}
__syncwarp();
// TMA store to send buffer
auto send_buffer_ptr = send_buffer.get_token_buffer(token_idx).get_base_ptr();
if constexpr (not kIsScaleupNVLink) {
if (ptx::elect_one_sync())
ptx::tma_store_1d(send_buffer_ptr, tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes<false>());
ptx::tma_store_commit();
__syncwarp();
}
// Issue TMA NVLink stores
EP_STATIC_ASSERT(kNumTopk <= 32, "Invalid top-k selection");
const auto dst_ptr = stored_dst_slot_idx >= 0 ?
gin.get_sym_ptr<team_t>(recv_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(), stored_dst_rank_idx) :
nullptr;
if (dst_ptr != nullptr)
ptx::tma_store_1d(dst_ptr, tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes<false>());
ptx::tma_store_commit();
__syncwarp();
// Issue RDMA put
if constexpr (not kIsScaleupNVLink) {
// Wait the send buffer store to arrive
ptx::tma_store_wait<1>();
__syncwarp();
// NOTES: we should skip the NVLink accessible ranks
if (stored_dst_slot_idx >= 0 and dst_ptr == nullptr) {
gin.put<team_t>(recv_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(),
send_buffer_ptr, tma_buffer.get_num_bytes<false>(), stored_dst_rank_idx);
}
__syncwarp();
}
}
}
// Barrier to ensure data arrival
comm::gpu_barrier<kIsScaleupNVLink, 1, kNumRanks,
kNumSMs, kNumThreads, kNumQPs, kNumTimeoutCycles, comm::kDispatchTag1, true, true, false>(
gin, workspace_layout, 0, rank_idx, sm_idx, thread_idx);
// Trigger the copy epilogue kernel
cudaTriggerProgrammaticLaunchCompletion();
// Clean atomic counters
EP_STATIC_ASSERT(kNumRanks <= kNumThreads, "Insufficient threads");
if (not kReuseSlotIndices and sm_idx == 0 and thread_idx < kNumRanks)
workspace_layout.get_scaleup_atomic_sender_counter()[thread_idx] = 0;
}
} // namespace deep_ep::elastic
deep_ep/include/deep_ep/impls/dispatch_copy_epilogue.cuh
+325
View File
@@ -0,0 +1,325 @@
#pragma once
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/layout.cuh>
#include <deep_ep/common/math.cuh>
#include <deep_ep/common/ptx.cuh>
namespace deep_ep::elastic {
template <bool kDoExpand, bool kCachedMode, bool kDoZeroPadding,
// NOTES: this channel concept only applies for scale-out ranks
int kNumSMs, int kNumChannels, int kNumWarps,
int kNumScaleoutRanks, int kNumScaleupRanks,
int kNumHiddenBytes, int kNumSFPacks,
int kNumMaxTokensPerRank,
int kNumExperts, int kNumTopk, int kExpertAlignment,
int kNumRanks = kNumScaleoutRanks * kNumScaleupRanks,
int kNumThreads = kNumWarps * 32,
int kNumMaxTokensPerChannel = math::constexpr_ceil_div(kNumMaxTokensPerRank, kNumChannels),
bool kDoCreateLinkedList = (kNumScaleoutRanks > 1 and not kCachedMode)>
__global__ void __launch_bounds__(kNumThreads, 1)
dispatch_copy_epilogue_impl(void* buffer, void* workspace,
int* psum_num_recv_tokens_per_scaleup_rank,
int* psum_num_recv_tokens_per_expert,
void* recv_x, sf_pack_t* recv_sf,
topk_idx_t* recv_topk_idx, float* recv_topk_weights,
int* recv_src_metadata,
int* channel_linked_list,
int* num_unaligned_recv_tokens_per_expert,
int num_recv_tokens,
const int recv_sf_token_stride, const int recv_sf_hidden_stride,
const int scaleout_rank_idx, const int scaleup_rank_idx) {
// Utils
const auto sm_idx = static_cast<int>(blockIdx.x), thread_idx = static_cast<int>(threadIdx.x);
const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx();
const auto global_warp_idx = warp_idx * kNumSMs + sm_idx;
// For top-k index transformations
constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks;
const auto rank_idx = scaleout_rank_idx * kNumScaleupRanks + scaleup_rank_idx;
const auto expert_start_idx = kNumExpertsPerRank * rank_idx, expert_end_idx = kNumExpertsPerRank * (rank_idx + 1);
// Buffer layouts
extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[];
const auto token_layout = layout::TokenLayout(kNumHiddenBytes, kNumSFPacks * sizeof(sf_pack_t), kNumTopk, true);
const auto tma_buffer = layout::BufferLayout<true>(token_layout, kNumWarps, 1, smem)
.get_rank_buffer(warp_idx).get_token_buffer(0);
const auto scaleup_buffer = layout::BufferLayout<false>(token_layout, kNumScaleupRanks, kNumScaleoutRanks * kNumMaxTokensPerRank, buffer);
// Init TMA
ptx::arrival_phase phase = 0;
const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr();
if (ptx::elect_one_sync())
ptx::mbarrier_init_with_fence(mbarrier_ptr, 1);
__syncwarp();
// Will block until the main dispatch kernel has finished and all data are visible
// NOTES: PDL is used, please do not use `__ldg`
cudaGridDependencySynchronize();
// For no CPU sync case, the number of received tokens should be read from the GPU tensor
if (num_recv_tokens == kNumMaxTokensPerRank * kNumRanks)
num_recv_tokens = psum_num_recv_tokens_per_scaleup_rank[kNumScaleupRanks - 1];
// Current rank indices should be maintained
int current_rank_idx = -1, stored_psum_num_recv_tokens;
int current_rank_start = 0, current_rank_end = 0;
#pragma unroll
for (int i = global_warp_idx; i < num_recv_tokens; i += kNumWarps * kNumSMs) {
// Calculate token index in the buffer
while (i >= current_rank_end) {
current_rank_idx += 1;
EP_DEVICE_ASSERT(current_rank_idx < kNumScaleupRanks);
const auto stored_lane_idx = current_rank_idx % 32;
if (stored_lane_idx == 0 and current_rank_idx + lane_idx < kNumScaleupRanks)
stored_psum_num_recv_tokens = psum_num_recv_tokens_per_scaleup_rank[current_rank_idx + lane_idx];
current_rank_start = current_rank_end;
current_rank_end = ptx::exchange(stored_psum_num_recv_tokens, stored_lane_idx);
}
const auto buffer_token = scaleup_buffer.get_rank_buffer(current_rank_idx).get_token_buffer(i - current_rank_start);
// Wait buffer releases
ptx::tma_store_wait();
__syncwarp();
// Issue TMA loads
// Including all stuffs: data, SF, top-k metadata
if (ptx::elect_one_sync()) {
ptx::tma_load_1d(tma_buffer.get_base_ptr(), buffer_token.get_base_ptr(),
mbarrier_ptr, tma_buffer.get_num_bytes<false>());
ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, tma_buffer.get_num_bytes<false>());
}
__syncwarp();
// Load target expert indices separately to tolerate TMA load latency
EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k selections");
int dst_expert_idx = -1;
if (lane_idx < kNumTopk)
dst_expert_idx = buffer_token.get_topk_idx_ptr()[lane_idx];
__syncwarp();
// Validate target expert indices and store for non-expand mode
const auto in_range = expert_start_idx <= dst_expert_idx and dst_expert_idx < expert_end_idx;
const auto master_src_topk_idx = ptx::get_master_lane_idx(ptx::gather(in_range));
dst_expert_idx = in_range ? dst_expert_idx - expert_start_idx : -1;
EP_DEVICE_ASSERT(ptx::deduplicate(dst_expert_idx, lane_idx) or dst_expert_idx == -1);
if (not kDoExpand and lane_idx < kNumTopk)
recv_topk_idx[i * kNumTopk + lane_idx] = static_cast<topk_idx_t>(dst_expert_idx);
__syncwarp();
// Calculate target indices in the tensor
constexpr int kMetadataStride = 2 + kNumTopk;
int dst_tensor_idx = -1;
if (not kDoExpand and ptx::elect_one_sync()) {
dst_tensor_idx = i;
} else if (kDoExpand and kCachedMode and lane_idx < kNumTopk) {
// Cached expand mode: read pre-computed dst_tensor_idx from metadata
dst_tensor_idx = recv_src_metadata[i * kMetadataStride + 2 + lane_idx];
} else if (kDoExpand and not kCachedMode and dst_expert_idx >= 0) {
dst_tensor_idx = atomicAdd(psum_num_recv_tokens_per_expert + dst_expert_idx, 1);
}
__syncwarp();
// Wait for TMA arrival
if (ptx::elect_one_sync())
ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase);
__syncwarp();
// Maintain linked list
if constexpr (kDoCreateLinkedList) {
if (ptx::elect_one_sync())
channel_linked_list[tma_buffer.get_linked_list_idx_ptr()[master_src_topk_idx]] = i;
__syncwarp();
}
// Issue TMA stores for data
if (kDoExpand ? (dst_tensor_idx >= 0) : ptx::elect_one_sync()) {
ptx::tma_store_1d(math::advance_ptr(recv_x, static_cast<int64_t>(dst_tensor_idx) * kNumHiddenBytes),
tma_buffer.get_hidden_ptr(), kNumHiddenBytes);
ptx::tma_store_commit();
}
__syncwarp();
// Store SF
if constexpr (kNumSFPacks > 0) {
constexpr auto kNumFullIters = kNumSFPacks / 32;
const bool do_last_iter = (kNumSFPacks % 32 != 0) and (kNumFullIters * 32 + lane_idx < kNumSFPacks);
EP_STATIC_ASSERT(sizeof(sf_pack_t) % 4 == 0, "Unaligned SF element type");
// Load into registers
const auto smem_src_ptr = tma_buffer.get_sf_ptr();
sf_pack_t reg_src[kNumFullIters + 1];
#pragma unroll
for (int k = 0; k < kNumFullIters; ++ k)
reg_src[k] = smem_src_ptr[k * 32 + lane_idx];
if (do_last_iter)
reg_src[kNumFullIters] = smem_src_ptr[kNumFullIters * 32 + lane_idx];
// Prepare strides
const auto recv_sf_token_stride_i64 = static_cast<int64_t>(recv_sf_token_stride);
const auto recv_sf_hidden_stride_i64 = static_cast<int64_t>(recv_sf_hidden_stride);
// Iterate through all valid indices and store into output buffer
auto mask = kDoExpand ? ptx::gather(dst_tensor_idx >= 0) : 1;
while (mask) {
const int valid_lane_idx = __ffs(mask) - 1;
const auto gmem_dst = math::advance_ptr<sf_pack_t>(recv_sf,
ptx::exchange(dst_tensor_idx, valid_lane_idx) * (recv_sf_token_stride_i64 * sizeof(sf_pack_t)));
#pragma unroll
for (int k = 0; k < kNumFullIters; ++ k)
gmem_dst[(k * 32 + lane_idx) * recv_sf_hidden_stride_i64] = reg_src[k];
if (do_last_iter)
gmem_dst[(kNumFullIters * 32 + lane_idx) * recv_sf_hidden_stride_i64] = reg_src[kNumFullIters];
mask ^= 1 << valid_lane_idx;
}
}
// Store the top-k weights
if (kDoExpand and recv_topk_weights != nullptr and dst_tensor_idx >= 0) {
recv_topk_weights[dst_tensor_idx] = tma_buffer.get_topk_weights_ptr()[lane_idx];
} else if (not kDoExpand and recv_topk_weights != nullptr and lane_idx < kNumTopk) {
// For backward, weights are optional
recv_topk_weights[i * kNumTopk + lane_idx] = tma_buffer.get_topk_weights_ptr()[lane_idx];
}
__syncwarp();
// Write source token index (skip in cached mode as metadata is reused)
// And:
// - Non-hybrid mode: the source scaleup peer rank index and master top-k lane index
// - Hybrid mode: the slot index and master top-k lane index
if constexpr (not kCachedMode) {
if (ptx::elect_one_sync()) {
recv_src_metadata[i * kMetadataStride + 0] = *tma_buffer.get_src_token_global_idx_ptr();
if constexpr (kNumScaleoutRanks == 1) {
recv_src_metadata[i * kMetadataStride + 1] = current_rank_idx * kNumTopk + master_src_topk_idx;
} else {
recv_src_metadata[i * kMetadataStride + 1] = (i - current_rank_start) * kNumTopk + master_src_topk_idx;
}
}
__syncwarp();
// Write reduction source indices
if (kDoExpand and lane_idx < kNumTopk)
recv_src_metadata[i * kMetadataStride + 2 + lane_idx] = dst_tensor_idx;
__syncwarp();
}
}
// Maintain linked list's ending
// Or you can understand it as writing the tail at once
if constexpr (kDoCreateLinkedList) {
constexpr int kNumScaleupRanksPerLane = math::constexpr_ceil_div(kNumScaleupRanks, 32);
const auto workspace_layout = layout::WorkspaceLayout(workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts);
for (int i = global_warp_idx; i < kNumChannels; i += kNumSMs * kNumWarps) {
#pragma unroll
for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) {
if (const auto k = j * 32 + lane_idx; j < (kNumScaleupRanksPerLane - 1) or k < kNumScaleupRanks) {
channel_linked_list[
*workspace_layout.get_channel_scaleup_tail_ptr(i, k)
] = -1;
// Clean for combine usages
*workspace_layout.get_channel_scaleup_tail_ptr(i, k) = 0;
}
}
__syncwarp();
}
}
// Zero padding: clear the alignment gaps between experts in the expanded output
if constexpr (kDoZeroPadding and kDoExpand) {
// Wait for last TMA store from main loop to complete
ptx::tma_store_wait();
__syncwarp();
// Zero out the TMA buffer's hidden region in smem
ptx::st_bulk<kNumHiddenBytes>(tma_buffer.get_hidden_ptr());
// Read all expert unaligned counts in parallel
constexpr int kNumExpertsPerLane = math::constexpr_ceil_div(kNumExpertsPerRank, 32);
int num_experts_per_lane[kNumExpertsPerLane];
#pragma unroll
for (int i = 0; i < kNumExpertsPerLane; ++ i) {
const int expert_idx = i * 32 + lane_idx;
num_experts_per_lane[i] = expert_idx < kNumExpertsPerRank ?
num_unaligned_recv_tokens_per_expert[expert_idx] : 0;
}
// Single while loop over all padding tokens
int wave_idx = 0, wave_pad_psum = 0, wave_tensor_psum = 0, pad_idx = global_warp_idx;
while (true) {
// Jump into the right position
int dst_tensor_idx;
while (wave_idx < kNumExpertsPerLane) {
// Assign current wave expert token count
int wave_num_experts_per_lane;
#pragma unroll
for (int i = 0; i < kNumExpertsPerLane; ++ i)
wave_num_experts_per_lane = i == wave_idx ? num_experts_per_lane[i] : wave_num_experts_per_lane;
int wave_num_pads_per_lane = 0;
if (wave_idx * 32 + lane_idx < kNumExpertsPerRank) {
wave_num_pads_per_lane = kExpertAlignment - (wave_num_experts_per_lane % kExpertAlignment);
wave_num_pads_per_lane = wave_num_pads_per_lane == kExpertAlignment ? 0 : wave_num_pads_per_lane;
}
int wave_num_pads = ptx::reduce_add(wave_num_pads_per_lane);
// Check whether hit the current wave
if (pad_idx < wave_pad_psum + wave_num_pads) {
const int local_pad_idx = pad_idx - wave_pad_psum;
const int pad_psum = ptx::warp_inclusive_sum(wave_num_pads_per_lane, lane_idx);
const int owner_lane_idx = __ffs(__ballot_sync(0xffffffff, pad_psum > local_pad_idx)) - 1;
dst_tensor_idx = wave_tensor_psum + ptx::exchange(
ptx::warp_exclusive_sum(wave_num_experts_per_lane + wave_num_pads_per_lane, lane_idx) +
wave_num_experts_per_lane + local_pad_idx - (pad_psum - wave_num_pads_per_lane),
owner_lane_idx);
break;
}
// Move to the next wave
wave_idx ++;
wave_pad_psum += wave_num_pads;
wave_tensor_psum += ptx::reduce_add(wave_num_experts_per_lane + wave_num_pads_per_lane);
}
if (wave_idx >= kNumExpertsPerLane)
break;
// Zero data via TMA store
if (ptx::elect_one_sync()) {
ptx::tma_store_1d(math::advance_ptr(recv_x, static_cast<int64_t>(dst_tensor_idx) * kNumHiddenBytes),
tma_buffer.get_hidden_ptr(), kNumHiddenBytes);
ptx::tma_store_commit();
}
__syncwarp();
// Zero weights
if (recv_topk_weights != nullptr and ptx::elect_one_sync())
recv_topk_weights[dst_tensor_idx] = 0.0f;
__syncwarp();
// Zero SF
if constexpr (kNumSFPacks > 0) {
const auto recv_sf_token_stride_i64 = static_cast<int64_t>(recv_sf_token_stride);
const auto recv_sf_hidden_stride_i64 = static_cast<int64_t>(recv_sf_hidden_stride);
constexpr sf_pack_t zero_sf_pack = {0};
const auto gmem_dst = math::advance_ptr<sf_pack_t>(recv_sf,
dst_tensor_idx * (recv_sf_token_stride_i64 * sizeof(sf_pack_t)));
constexpr auto kNumFullIters = kNumSFPacks / 32;
#pragma unroll
for (int k = 0; k < kNumFullIters; ++ k)
gmem_dst[(k * 32 + lane_idx) * recv_sf_hidden_stride_i64] = zero_sf_pack;
if constexpr (kNumSFPacks % 32 != 0) {
if (kNumFullIters * 32 + lane_idx < kNumSFPacks)
gmem_dst[(kNumFullIters * 32 + lane_idx) * recv_sf_hidden_stride_i64] = zero_sf_pack;
}
}
__syncwarp();
// Try to fill the next one
pad_idx += kNumSMs * kNumWarps;
}
}
}
} // namespace deep_ep::elastic
deep_ep/include/deep_ep/impls/engram_fetch.cuh
+108
View File
@@ -0,0 +1,108 @@
#pragma once
#include <deep_ep/common/comm.cuh>
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/layout.cuh>
#include <deep_ep/common/ptx.cuh>
namespace deep_ep::elastic {
template <int kNumQPs,
int kNumEntriesPerRank,
int kNumHiddenBytes,
int kNumSFPacks,
int kNumEntriesPerToken,
int kNumRDMAPeers,
int kNumRanksPerRDMAPeer,
size_t kIntraRankStorageStride,
int kNumThreads,
typename team_t,
int kNumWarps = kNumThreads / 32>
__global__ void __launch_bounds__(kNumThreads, 1)
engram_fetch_impl(const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window,
void* storage, void* fetched, int* indices,
ncclGinRequest_t* last_gin_requests,
const sf_pack_t* sf_table, sf_pack_t* fetched_sf,
const int sf_token_stride, const int sf_hidden_stride,
const int num_tokens) {
const auto qp_idx = static_cast<int>(blockIdx.x);
const auto warp_idx = ptx::get_warp_idx();
const auto global_warp_idx = qp_idx * kNumWarps + warp_idx;
const auto thread_idx = static_cast<int>(threadIdx.x);
// Gin handle
const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, NCCL_GIN_RESOURCE_SHARING_CTA);
__shared__ int num_requests_per_peer[kNumRDMAPeers];
EP_STATIC_ASSERT(kNumRDMAPeers <= kNumThreads, "Too many RDMA peers");
if (thread_idx < kNumRDMAPeers)
num_requests_per_peer[thread_idx] = 0;
__syncthreads();
// Issue RDMA
const auto issue_rdma_get = [=](const int& token_idx, const int& peer_idx,
const int64_t& src_byte_offset, const int& extra_options = 0) {
gin.get<team_t, ncclCoopThread, ncclGin_SegmentMixed>(math::advance_ptr(storage, src_byte_offset),
math::advance_ptr(fetched, static_cast<int64_t>(token_idx) * kNumHiddenBytes),
kNumHiddenBytes, peer_idx, extra_options);
};
// Each warp fetches one token cooperatively via RDMA gin.get
// TODO: deal with padded tokens
if (ptx::elect_one_sync()) {
#pragma unroll 4
for (int i = global_warp_idx; i < num_tokens * kNumEntriesPerToken; i += kNumQPs * kNumWarps) {
const auto global_idx = __ldg(indices + i);
const auto owner_rank_idx = global_idx / kNumEntriesPerRank;
const auto local_entry_idx = global_idx % kNumEntriesPerRank;
// Route owner rank to RDMA peer and intra-peer rank
const auto peer_idx = owner_rank_idx / kNumRanksPerRDMAPeer;
const auto intra_peer_rank_idx = owner_rank_idx % kNumRanksPerRDMAPeer;
// Byte offset into storage layout
const auto src_byte_offset = static_cast<int64_t>(intra_peer_rank_idx) * kIntraRankStorageStride +
static_cast<int64_t>(local_entry_idx) * kNumHiddenBytes;
// Issue RDMA get
const auto request_idx = atomicAdd_block(num_requests_per_peer + peer_idx, 1);
issue_rdma_get(
i, peer_idx, src_byte_offset,
// NOTES: requests may exceed the queue depth, flush if needed
(request_idx % kGinQPFlushDepth == (kGinQPFlushDepth - 1)) ? 0 : ncclGinOptFlagsAggregateRequests
);
// TODO: once NCCL supports ncclCoopWarp gin.get, drop the elect_one_sync and let the whole warp
// gather SF packs in parallel.
if constexpr (kNumSFPacks > 0) {
EP_STATIC_ASSERT(sizeof(sf_pack_t) == sizeof(int), "SF pack must be 4 bytes");
const auto token_idx = i / kNumEntriesPerToken;
const auto entry_idx = i % kNumEntriesPerToken;
const auto* src = reinterpret_cast<const int*>(sf_table) + global_idx * kNumSFPacks;
auto* dst = reinterpret_cast<int*>(fetched_sf)
+ token_idx * sf_token_stride + entry_idx * kNumSFPacks * sf_hidden_stride;
#pragma unroll
for (int p = 0; p < kNumSFPacks; ++p)
dst[p * sf_hidden_stride] = __ldg(src + p);
}
}
}
__syncthreads();
// Issue flush per peer we sent to; its unconditional DB ring flushes all
// prior aggregated gets on the same QP.
if (ptx::elect_one_sync()) {
for (int i = warp_idx; i < kNumRDMAPeers; i += kNumWarps) {
const auto request_ptr = last_gin_requests + qp_idx * kNumRDMAPeers + i;
if (num_requests_per_peer[i] > 0) {
gin.flush_async<team_t, ncclCoopThread>(i, request_ptr);
} else {
EP_STATIC_ASSERT(sizeof(ncclGinRequest_t) == sizeof(int4), "Invalid request size");
*reinterpret_cast<int4*>(request_ptr) = make_int4(0, 0, 0, 0);
}
}
}
}
} // namespace deep_ep::elastic
deep_ep/include/deep_ep/impls/engram_fetch_wait.cuh
+31
View File
@@ -0,0 +1,31 @@
#pragma once
#include <deep_ep/common/comm.cuh>
#include <deep_ep/common/compiled.cuh>
namespace deep_ep::elastic {
template <int kNumRDMAPeers, int kNumThreads>
__global__ void __launch_bounds__(kNumThreads, 1)
engram_fetch_wait_impl(const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window,
ncclGinRequest_t* last_gin_requests) {
const auto qp_idx = static_cast<int>(blockIdx.x);
const auto thread_idx = static_cast<int>(threadIdx.x);
// Gin handle
const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, NCCL_GIN_RESOURCE_SHARING_CTA);
// Wait for all RDMA gets to complete
for (int i = thread_idx; i < kNumRDMAPeers; i += kNumThreads) {
EP_STATIC_ASSERT(sizeof(ncclGinRequest_t) == sizeof(int4), "Invalid request size");
auto last_gin_req_int4 = __ldg(reinterpret_cast<int4*>(last_gin_requests + qp_idx * kNumRDMAPeers + i));
if (last_gin_req_int4.x != 0 or last_gin_req_int4.y != 0 or
last_gin_req_int4.z != 0 or last_gin_req_int4.w != 0) {
auto last_gin_req = *reinterpret_cast<ncclGinRequest_t*>(&last_gin_req_int4);
gin.wait(last_gin_req);
}
}
}
} // namespace deep_ep::elastic
deep_ep/include/deep_ep/impls/hybrid_combine.cuh
+626
View File
@@ -0,0 +1,626 @@
#pragma once
#include <deep_ep/common/comm.cuh>
#include <deep_ep/common/layout.cuh>
#include <deep_ep/common/math.cuh>
#include <deep_ep/common/ptx.cuh>
#include <deep_ep/impls/combine_utils.cuh>
namespace deep_ep::elastic {
template <bool kUseExpandedLayout, bool kAllowMultipleReduction,
int kNumSMs,
int kNumScaleupWarps, int kNumForwardWarps,
int kNumScaleoutRanks, int kNumScaleupRanks,
int kHidden,
int kNumMaxTokensPerRank,
int kNumExperts, int kNumTopk,
int kNumQPs, int64_t kNumTimeoutCycles,
int kNumScaleupRanksPerLane = math::constexpr_ceil_div(kNumScaleupRanks, 32),
int kNumScaleupUpdateInterval = 3,
int kNumChannelsPerSM = kNumForwardWarps,
int kNumChannels = kNumChannelsPerSM * kNumSMs,
int kNumMaxTokensPerChannel = math::constexpr_ceil_div(kNumMaxTokensPerRank, kNumChannels),
int kNumRanks = kNumScaleoutRanks * kNumScaleupRanks,
int kNumWarps = kNumScaleupWarps + kNumForwardWarps,
int kNumThreads = kNumWarps * 32,
int kNumHiddenBytes = kHidden * sizeof(nv_bfloat16),
bool kUseScaleoutRankLayout = use_rank_layout<kAllowMultipleReduction, kNumScaleoutRanks, kNumTopk>(),
bool kUseScaleupRankLayout = use_rank_layout<kAllowMultipleReduction, kNumScaleupRanks, kNumTopk>(),
int kNumTokensInScaleoutLayout = get_num_tokens_in_layout<kAllowMultipleReduction, kNumScaleoutRanks, kNumTopk>(),
int kNumTokensInScaleupLayout = get_num_tokens_in_layout<kAllowMultipleReduction, kNumScaleupRanks, kNumTopk>()>
__global__ void __launch_bounds__(kNumThreads, 1)
hybrid_combine_impl(nv_bfloat16* x,
float* topk_weights,
int* src_metadata,
int* psum_num_recv_tokens_per_scaleup_rank,
int* token_metadata_at_forward,
int* channel_linked_list,
const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window,
void* buffer, void* workspace,
const int scaleout_rank_idx, const int scaleup_rank_idx,
int num_reduced_tokens) {
// Utils
const auto sm_idx = static_cast<int>(blockIdx.x);
const auto thread_idx = static_cast<int>(threadIdx.x);
const auto warp_idx = ptx::get_warp_idx();
const auto lane_idx = ptx::get_lane_idx();
constexpr bool kDoExpandedSend = not kAllowMultipleReduction and kUseExpandedLayout;
// Combine vector type selection
using combine_vec_t = typename CombineVecTraits<kNumHiddenBytes>::vec_t;
constexpr int kHiddenVec = kNumHiddenBytes / sizeof(combine_vec_t);
// Workspaces
const auto workspace_layout = layout::WorkspaceLayout(workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts);
// We should assign the real number of received tokens if without CPU sync
if (num_reduced_tokens == kNumMaxTokensPerRank * kNumRanks)
num_reduced_tokens = __ldg(psum_num_recv_tokens_per_scaleup_rank + kNumScaleupRanks - 1);
// Token layouts
const auto token_layout = layout::TokenLayout(kNumHiddenBytes, 0, kNumTopk, false);
// TMA buffers
extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[];
const auto tma_buffer = layout::BufferLayout<true>(
token_layout, kNumWarps, 1, smem).get_rank_buffer(warp_idx).get_token_buffer(0);
// All the buffer layouts
auto scaleup_buffer = layout::BufferLayout<false>(
token_layout, kNumTokensInScaleupLayout, kNumScaleoutRanks * kNumMaxTokensPerRank,
buffer);
auto scaleout_recv_buffer = layout::BufferLayout<false>(
token_layout, kNumTokensInScaleoutLayout, kNumMaxTokensPerRank,
scaleup_buffer.get_buffer_end_ptr());
auto scaleout_send_buffer = layout::BufferLayout<false>(
token_layout, kAllowMultipleReduction ? 1 : kNumTopk, kNumChannels * (kNumScaleoutRanks * kNumMaxTokensPerChannel),
scaleout_recv_buffer.get_buffer_end_ptr());
// Init TMA for scale-up and forward warps
ptx::arrival_phase phase = 0;
const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr();
if (ptx::elect_one_sync())
ptx::mbarrier_init_with_fence(mbarrier_ptr, 1);
__syncwarp();
// NCCL Gin handle
// Each warp is a channel
const auto [qp_idx, sharing_mode] =
comm::get_qp_mode<kNumSMs, kNumQPs, kNumChannelsPerSM>(sm_idx, warp_idx % kNumChannelsPerSM);
const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, sharing_mode);
// Global parallel barriers for scale-out subteam and scale-up subteam
// NOTES: this barrier needs a grid sync, as there are channel scale-up tail cleaning before
comm::gpu_barrier<true, kNumScaleoutRanks, kNumScaleupRanks,
kNumSMs, kNumThreads, kNumQPs, kNumTimeoutCycles, comm::kHybridCombineTag0, false, true, true>(
gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx);
// Adjust register count at certain cases
// TODO: support more cases, or try to make channel count more aligned
const bool kAdjustRegisters = (kNumChannelsPerSM == 4 or kNumChannelsPerSM == 8) and not kUseExpandedLayout;
constexpr int kNumRegistersForScaleupWarps = 40;
constexpr int kNumRegistersForForwardWarps = 256 - kNumRegistersForScaleupWarps;
// Different warp roles
if (warp_idx < kNumScaleupWarps) {
const auto channel_idx = sm_idx * kNumChannelsPerSM + warp_idx;
// Adjust registers
if constexpr (kAdjustRegisters)
ptx::warpgroup_reg_dealloc<kNumRegistersForScaleupWarps>();
// Shift into the right buffer if using rank layout
if constexpr (kUseScaleupRankLayout)
scaleup_buffer = scaleup_buffer.get_rank_buffer(scaleup_rank_idx);
// Expanding send mode must not be backward
if constexpr (kDoExpandedSend)
EP_DEVICE_ASSERT(topk_weights == nullptr);
// Tail issuer
// `st.release.sys` is pretty slow, so do it by an interval
int update_counter = 0;
int stored_num_tokens_sent[kNumScaleupRanksPerLane] = {};
int stored_old_num_tokens_sent[kNumScaleupRanksPerLane] = {};
const auto tail_ptr = workspace_layout.get_channel_scaleup_tail_ptr(channel_idx, scaleup_rank_idx);
const auto update_tails = [&](const bool& finish = false) {
++ update_counter;
if (finish or update_counter == kNumScaleupUpdateInterval) {
// Wait all TMA stores to finish
ptx::tma_store_wait();
__syncwarp();
// Issue
#pragma unroll
for (int i = 0; i < kNumScaleupRanksPerLane; ++ i) {
if (const auto j = i * 32 + lane_idx; i < (kNumScaleupRanksPerLane - 1) or j < kNumScaleupRanks) {
// NOTES: save some traffic with `stored_old_num_tokens_sent`
// Also, we cannot rewrite a finished slot, if the peer is going to clean it
if (stored_num_tokens_sent[i] != stored_old_num_tokens_sent[i])
ptx::st_release_sys(gin.get_sym_ptr<ncclTeamTagLsa>(tail_ptr, j), stored_num_tokens_sent[i]);
stored_old_num_tokens_sent[i] = stored_num_tokens_sent[i];
}
}
update_counter = 0;
}
__syncwarp();
};
// Shape of `channel_linked_list`: `[kNumChannels, kNumMaxTokensPerChannel + 1, kNumScaleupRanks]`
// Iterate until all scale-up peers finish
int dst_scaleup_rank_idx = channel_idx;
int stored_ll_idx[kNumScaleupRanksPerLane] = {}, stored_token_idx[kNumScaleupRanksPerLane] = {};
#pragma unroll
for (int i = 0; i < kNumScaleupRanksPerLane; ++ i)
stored_token_idx[i] = -1;
while (true) {
// Load token indices in the list
#pragma unroll
for (int i = 0; i < kNumScaleupRanksPerLane; ++ i) {
const auto j = i * 32 + lane_idx;
stored_token_idx[i] = i < (kNumScaleupRanksPerLane - 1) or j < kNumScaleupRanks ?
__ldg(channel_linked_list +
channel_idx * (kNumScaleoutRanks * kNumMaxTokensPerChannel + 1) * kNumScaleupRanks +
stored_ll_idx[i] * kNumScaleupRanks + j) : -1;
}
__syncwarp();
// Check whether all ranks are finished
bool exited = true;
#pragma unroll
for (int i = 0; i < kNumScaleupRanksPerLane; ++ i)
exited &= ptx::all(stored_token_idx[i] < 0);
if (exited)
break;
// Process tokens for all ranks together using bitmask to skip inactive ranks
EP_STATIC_ASSERT(kNumScaleupRanks <= 64, "Too many scale-up ranks for 64-bit mask");
using mask_t = std::conditional_t<(kNumScaleupRanks <= 32), uint32_t, uint64_t>;
mask_t wip_mask = 0;
#pragma unroll
for (int j = 0; j < kNumScaleupRanksPerLane; ++ j)
wip_mask |= static_cast<mask_t>(ptx::gather(stored_token_idx[j] >= 0)) << (j * 32);
while (wip_mask) {
// Find next active rank after `dst_scaleup_rank_idx` (round-robin)
const auto start = (dst_scaleup_rank_idx + 1) % kNumScaleupRanks;
const auto hi_mask = (wip_mask >> start) << start;
dst_scaleup_rank_idx = hi_mask ? ptx::ffs(hi_mask) : ptx::ffs(wip_mask);
wip_mask ^= static_cast<mask_t>(1) << dst_scaleup_rank_idx;
// Exchange token index from the owning lane using static partition iteration
int token_idx = -1;
#pragma unroll
for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) {
const auto src_lane_idx = dst_scaleup_rank_idx - j * 32;
token_idx = src_lane_idx == lane_idx ? stored_token_idx[j] : token_idx;
}
token_idx = ptx::exchange(token_idx, dst_scaleup_rank_idx % 32);
// Get source metadata and decide the destination buffer
constexpr int kMetadataStride = 2 + kNumTopk;
const auto src_global_token_idx = __ldg(src_metadata + token_idx * kMetadataStride + 0);
const auto src_token_idx = src_global_token_idx % kNumMaxTokensPerRank;
const auto src_scaleout_rank_idx = src_global_token_idx / (kNumMaxTokensPerRank * kNumScaleupRanks);
auto token_buffer = [&]() {
if constexpr (kUseScaleupRankLayout) {
const auto src_slot_idx = __ldg(src_metadata + token_idx * kMetadataStride + 1) / kNumTopk;
return scaleup_buffer.get_token_buffer(src_slot_idx);
} else {
const auto master_topk_idx = __ldg(src_metadata + token_idx * kMetadataStride + 1) % kNumTopk;
return scaleup_buffer
.get_rank_buffer(master_topk_idx)
.get_token_buffer(src_scaleout_rank_idx * kNumMaxTokensPerRank + src_token_idx);
}
}();
token_buffer.set_base_ptr(gin.get_sym_ptr<ncclTeamTagLsa>(token_buffer.get_base_ptr(), dst_scaleup_rank_idx));
// Some checks
EP_STATIC_ASSERT(kHidden % (32 * sizeof(int4) / sizeof(nv_bfloat16)) == 0, "Invalid hidden");
// Read source indices for expand mode
int stored_topk_slot_idx = -1;
if constexpr (kUseExpandedLayout) {
if (lane_idx < kNumTopk)
stored_topk_slot_idx = __ldg(src_metadata + token_idx * kMetadataStride + (2 + lane_idx));
__syncwarp();
}
// 3 cases:
// - no-expand, expand + no-reduce
// - expand + reduce
// - expand + send all
auto reduce_valid_mask = ptx::gather(stored_topk_slot_idx >= 0);
auto no_local_reduce = not kUseExpandedLayout or (kAllowMultipleReduction and __popc(reduce_valid_mask) == 1);
if (no_local_reduce) {
int token_idx_in_tensor = token_idx;
if constexpr (kUseExpandedLayout)
token_idx_in_tensor = ptx::exchange(stored_topk_slot_idx, ptx::get_master_lane_idx(reduce_valid_mask));
// Directly load
if (ptx::elect_one_sync()) {
const auto load_ptr =
math::advance_ptr(x, static_cast<int64_t>(token_idx_in_tensor) * kNumHiddenBytes);
ptx::tma_store_wait();
ptx::tma_load_1d(tma_buffer.get_base_ptr(), load_ptr, mbarrier_ptr, kNumHiddenBytes);
}
__syncwarp();
} else if constexpr (kAllowMultipleReduction) {
// Do local reduction
// Sort valid top-k indices to front
int topk_slot_idx[kNumTopk];
compute_topk_slots(
topk_slot_idx, reduce_valid_mask,
[=](const int& idx) {
return ptx::exchange(stored_topk_slot_idx, idx);
}
);
// Reduce into shared memory
constexpr int kUnrollFactor = get_max_unroll_factor<kHiddenVec, 4>();
combine_reduce<kHiddenVec, kUnrollFactor, math::constexpr_ceil_div(kNumTopk, kNumRanks)>(
lane_idx, topk_slot_idx, static_cast<combine_vec_t*>(tma_buffer.get_base_ptr()),
/* Get source base */ [=](const int& slot_idx) {
return math::advance_ptr<combine_vec_t>(
x, slot_idx * static_cast<int64_t>(kNumHiddenBytes));
},
/* Wait buffer release */ [=]() {
ptx::tma_store_wait();
__syncwarp();
}
);
ptx::tma_store_fence();
__syncwarp();
} else {
// No local reduction, send all data (expanded send)
#pragma unroll
for (int k = 0; k < kNumTopk; ++ k) {
int topk_slot_idx = ptx::exchange(stored_topk_slot_idx, k);
if (topk_slot_idx < 0)
continue;
if (ptx::elect_one_sync()) {
// Load
const auto load_ptr = math::advance_ptr(x, static_cast<int64_t>(kDoExpandedSend ? topk_slot_idx : token_idx) * kNumHiddenBytes);
ptx::tma_store_wait();
ptx::tma_load_1d(tma_buffer.get_base_ptr(), load_ptr, mbarrier_ptr, kNumHiddenBytes);
ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes);
ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase);
// NOTES: We don't need to care about `topk_weights` since we are in expand mode
// Store
const auto dst_token_buffer = scaleup_buffer
.get_rank_buffer(k)
.get_token_buffer(src_scaleout_rank_idx * kNumMaxTokensPerRank + src_token_idx);
ptx::tma_store_1d(
gin.get_sym_ptr<ncclTeamTagLsa>(dst_token_buffer.get_base_ptr(), dst_scaleup_rank_idx),
tma_buffer.get_base_ptr(), token_layout.get_num_bytes<false>());
ptx::tma_store_commit();
}
__syncwarp();
}
}
// Write top-k weights (expanded send handled inside the loop above)
if (not kDoExpandedSend and topk_weights != nullptr and lane_idx < kNumTopk) {
float value = 0;
if constexpr (kUseExpandedLayout) {
if (stored_topk_slot_idx >= 0)
value = __ldg(topk_weights + stored_topk_slot_idx);
} else {
value = __ldg(topk_weights + (token_idx * kNumTopk + lane_idx));
}
tma_buffer.get_topk_weights_ptr()[lane_idx] = value;
ptx::tma_store_fence();
}
__syncwarp();
// Issue TMA stores into remote scale-up buffer
// NOTES: `kDoExpandedSend` mode has already issued
if (not kDoExpandedSend and ptx::elect_one_sync()) {
// Wait TMA arrival (only for non-reduced cases)
if (no_local_reduce) {
ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes);
ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase);
}
// Issue stores
ptx::tma_store_1d(
token_buffer.get_base_ptr(), tma_buffer.get_base_ptr(),
token_layout.get_num_bytes<false>());
ptx::tma_store_commit();
}
#pragma unroll
for (int j = 0; j < kNumScaleupRanksPerLane; ++ j)
stored_num_tokens_sent[j] += (j * 32 + lane_idx) == dst_scaleup_rank_idx;
__syncwarp();
}
// Update the tails together
// NOTES: TMA wait is inside
update_tails();
// Move linked list
#pragma unroll
for (int i = 0; i < kNumScaleupRanksPerLane; ++ i)
stored_ll_idx[i] += (stored_token_idx[i] >= 0);
}
// Update for the unissued ones
update_tails(true);
} else {
const auto forward_warp_idx = warp_idx - kNumScaleupWarps;
const auto channel_idx = sm_idx * kNumChannelsPerSM + forward_warp_idx;
// Adjust registers
if constexpr (kAdjustRegisters)
ptx::warpgroup_reg_alloc<kNumRegistersForForwardWarps>();
// Shift into the right buffer
scaleout_send_buffer = scaleout_send_buffer.get_channel_buffer<kNumScaleoutRanks * kNumMaxTokensPerChannel>(channel_idx);
// Shape of `token_metadata_at_forward`: `[kNumChannels, kNumScaleoutRanks * kNumMaxTokensPerChannel + 1, kNumForwardMetadataDims]`
constexpr int kNumForwardMetadataDims = 2 + kNumTopk * 2;
token_metadata_at_forward += channel_idx * ((kNumScaleoutRanks * kNumMaxTokensPerChannel + 1) * kNumForwardMetadataDims);
// Overlap TMA stores and reduction
int last_src_scaleout_rank_idx = -1;
int last_is_token_last_in_chunk = 0;
void* last_recv_token_buffer_ptr = nullptr;
void* last_send_token_buffer_ptr = nullptr;
const auto flush_last_tma_and_issue_rdma = [&]() {
if (last_src_scaleout_rank_idx >= 0 and ptx::elect_one_sync()) {
ptx::tma_store_wait();
// Issue only if not local rank
if (last_src_scaleout_rank_idx != scaleout_rank_idx) {
gin.put<ncclTeamTagRail>(
last_recv_token_buffer_ptr,
last_send_token_buffer_ptr,
token_layout.get_num_bytes<false>(),
last_src_scaleout_rank_idx,
last_is_token_last_in_chunk ? 0 : ncclGinOptFlagsAggregateRequests
);
}
}
__syncwarp();
};
// Replay the dispatch
int stored_num_tokens_recv[kNumScaleupRanksPerLane] = {}, stored_cached_scaleup_tail[kNumScaleupRanksPerLane] = {};
for (int i = 0; ; ++ i) {
const auto src_token_global_idx = __ldg(token_metadata_at_forward + i * kNumForwardMetadataDims);
const auto is_token_last_in_chunk = __ldg(token_metadata_at_forward + i * kNumForwardMetadataDims + 1);
const auto src_rank_idx = src_token_global_idx / kNumMaxTokensPerRank;
const auto src_scaleout_rank_idx = src_rank_idx / kNumScaleupRanks;
const auto src_token_idx = src_token_global_idx % kNumMaxTokensPerRank;
auto stored_src_scaleup_rank_idx = lane_idx < kNumTopk ?
__ldg(token_metadata_at_forward + i * kNumForwardMetadataDims + 2 + lane_idx) : -1;
auto stored_src_slot_idx = lane_idx < kNumTopk ?
__ldg(token_metadata_at_forward + i * kNumForwardMetadataDims + 2 + kNumTopk + lane_idx) : -1;
if (src_token_global_idx < 0)
break;
// Scaleup rank mask
EP_STATIC_ASSERT(kNumScaleupRanks <= 64, "Too many scale-up peers");
using mask_t = std::conditional_t<kNumScaleupRanks <= 32, unsigned, unsigned long long>;
const auto scaleup_mask = ptx::reduce_or(
stored_src_scaleup_rank_idx >= 0 ?
(mask_t(1) << stored_src_scaleup_rank_idx) : mask_t(0));
bool stored_is_scaleup_rank_needed[kNumScaleupRanksPerLane];
#pragma unroll
for (int j = 0; j < kNumScaleupRanksPerLane; ++ j)
stored_is_scaleup_rank_needed[j] = (scaleup_mask >> (j * 32 + lane_idx)) & 1;
// Wait all tails to arrive
comm::timeout_while<kNumTimeoutCycles>([&](const bool& is_last_check) {
bool arrived = true;
#pragma unroll
for (int j = 0; j < kNumScaleupRanksPerLane; ++ j)
arrived &= not stored_is_scaleup_rank_needed[j] or stored_num_tokens_recv[j] < stored_cached_scaleup_tail[j];
if (ptx::all(arrived))
return true;
// Reload cached
#pragma unroll
for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) {
const auto k = j * 32 + lane_idx;
stored_cached_scaleup_tail[j] = j < (kNumScaleupRanksPerLane - 1) or k < kNumScaleupRanks ?
ptx::ld_acquire_sys(workspace_layout.get_channel_scaleup_tail_ptr(channel_idx, k)) : -1;
}
// Timeout
if (is_last_check) {
#pragma unroll
for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) {
printf("DeepEP combine (scale-up wait) timeout, scale-out: %d/%d, scale-up: %d/%d, "
"channel: %d, lane: %d, recv: %d, tail: %d (wait=%d)\n",
scaleout_rank_idx, kNumScaleoutRanks, scaleup_rank_idx, kNumScaleupRanks,
channel_idx, j * 32 + lane_idx,
stored_num_tokens_recv[j],
stored_cached_scaleup_tail[j],
stored_is_scaleup_rank_needed[j]);
}
}
return false;
});
// Increase received count
#pragma unroll
for (int j = 0; j < kNumScaleupRanksPerLane; ++ j)
stored_num_tokens_recv[j] += static_cast<int>(stored_is_scaleup_rank_needed[j]);
if constexpr (not kAllowMultipleReduction) {
// Cases where multiple reduction is disabled. We need to forward all data from scaleup peers to scaleout peers
// TODO: Let scale-up warps directly put data into `send_buffer`?
const auto src_slot_idx = src_scaleout_rank_idx * kNumMaxTokensPerRank + src_token_idx;
auto topk_valid_mask = kUseExpandedLayout ?
ptx::gather(stored_src_scaleup_rank_idx >= 0) :
ptx::gather(ptx::deduplicate(stored_src_scaleup_rank_idx, lane_idx) and stored_src_scaleup_rank_idx >= 0); // Deduplicate w.r.t. scaleup rank index if expanded mode is disabled
if (ptx::elect_one_sync()) {
#pragma unroll
for (int k = 0; k < kNumTopk; ++ k) {
if ((topk_valid_mask & (1u << k)) == 0u)
continue;
// Issue TMA load, and wait
ptx::tma_load_1d(
tma_buffer.get_base_ptr(), scaleup_buffer.get_rank_buffer(k).get_token_buffer(src_slot_idx).get_base_ptr(),
mbarrier_ptr, token_layout.get_num_bytes<false>());
ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, token_layout.get_num_bytes<false>());
ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase);
// Issue TMA store, and wait
const auto recv_buffer_ptr = scaleout_recv_buffer.get_rank_buffer(k).get_token_buffer(src_token_idx).get_base_ptr();
const auto send_buffer_ptr = src_scaleout_rank_idx == scaleout_rank_idx ?
recv_buffer_ptr : scaleout_send_buffer.get_rank_buffer(k).get_token_buffer(i).get_base_ptr();
ptx::tma_store_1d(send_buffer_ptr, tma_buffer.get_base_ptr(), token_layout.get_num_bytes<false>());
ptx::tma_store_commit();
ptx::tma_store_wait();
// Issue IBGDA
topk_valid_mask ^= 1u << k;
if (src_scaleout_rank_idx != scaleout_rank_idx) {
gin.put<ncclTeamTagRail>(
recv_buffer_ptr,
send_buffer_ptr,
token_layout.get_num_bytes<false>(),
src_scaleout_rank_idx,
topk_valid_mask == 0 and is_token_last_in_chunk ? 0 : ncclGinOptFlagsAggregateRequests
);
}
}
}
__syncwarp();
} else {
// NOTES: we must do deduplicate and only add once from one rank
auto reduce_valid_mask = ptx::gather(
ptx::deduplicate(stored_src_scaleup_rank_idx, lane_idx) and stored_src_scaleup_rank_idx >= 0);
// Calculate the source buffer index
int stored_src_buffer_idx = 0;
if constexpr (kUseScaleupRankLayout) {
stored_src_buffer_idx =
stored_src_scaleup_rank_idx * scaleup_buffer.num_max_tokens_per_rank + stored_src_slot_idx;
} else {
const auto src_slot_idx = src_scaleout_rank_idx * kNumMaxTokensPerRank + src_token_idx;
stored_src_buffer_idx = stored_src_slot_idx == -1 ? -1 :
lane_idx * scaleup_buffer.num_max_tokens_per_rank + src_slot_idx;
}
// Preprocess top-k indices
int topk_slot_idx[kNumTokensInScaleupLayout];
compute_topk_slots(
topk_slot_idx, reduce_valid_mask,
[=](const int& idx) {
return ptx::exchange(stored_src_buffer_idx, idx);
}
);
// Do reduce
constexpr int kUnrollFactor = get_max_unroll_factor<kHiddenVec, kAdjustRegisters ? 8 : 4>();
combine_reduce<kHiddenVec, kUnrollFactor, math::constexpr_ceil_div(kNumTopk, kNumScaleoutRanks)>(
lane_idx, topk_slot_idx, static_cast<combine_vec_t*>(tma_buffer.get_base_ptr()),
/* Get source base */ [=](const int& slot_idx) {
return static_cast<combine_vec_t*>(scaleup_buffer.get_token_buffer(slot_idx, true).get_base_ptr());
},
/* Wait buffer release */ [=]() {
flush_last_tma_and_issue_rdma();
}
);
// Merge topk weights
// NOTES: the slot indices must follow the master lane
stored_src_buffer_idx = ptx::exchange(
stored_src_buffer_idx, ptx::get_master_lane_idx(ptx::match(stored_src_scaleup_rank_idx)));
if (stored_src_scaleup_rank_idx >= 0) {
tma_buffer.get_topk_weights_ptr()[lane_idx] =
scaleup_buffer.get_token_buffer(stored_src_buffer_idx, true)
.get_topk_weights_ptr()[lane_idx];
}
ptx::tma_store_fence();
__syncwarp(); // Necessary to let the leader lane see the writes
// Assign send and receive buffers
// NOTES: as we only have 1 destination, we will use "send" as "recv" for local transfer
int scaleout_recv_buffer_rank_idx;
if constexpr (kUseScaleoutRankLayout) {
scaleout_recv_buffer_rank_idx = scaleout_rank_idx;
} else {
const int src_topk_idx = ptx::get_master_lane_idx(ptx::gather(stored_src_scaleup_rank_idx >= 0));
scaleout_recv_buffer_rank_idx = src_topk_idx;
}
const auto recv_token_buffer = scaleout_recv_buffer.get_rank_buffer(scaleout_recv_buffer_rank_idx).get_token_buffer(src_token_idx);
const auto send_token_buffer = src_scaleout_rank_idx == scaleout_rank_idx ?
recv_token_buffer :
scaleout_send_buffer.get_token_buffer(i);
// Write into scale-out send buffer or local rank recv buffer bypass
if (ptx::elect_one_sync()) {
ptx::tma_store_1d(send_token_buffer.get_base_ptr(), tma_buffer.get_base_ptr(),
token_layout.get_num_bytes<false>());
ptx::tma_store_commit();
}
__syncwarp();
// Record RDMA info to issue later
last_src_scaleout_rank_idx = src_scaleout_rank_idx;
last_is_token_last_in_chunk = is_token_last_in_chunk;
last_recv_token_buffer_ptr = recv_token_buffer.get_base_ptr();
last_send_token_buffer_ptr = send_token_buffer.get_base_ptr();
}
}
// Issue the last RDMA
if constexpr (kAllowMultipleReduction)
flush_last_tma_and_issue_rdma();
// Clean scaleup tails
#pragma unroll
for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) {
const auto k = j * 32 + lane_idx;
if (j < (kNumScaleupRanksPerLane - 1) or k < kNumScaleupRanks)
*workspace_layout.get_channel_scaleup_tail_ptr(channel_idx, k) = 0;
}
__syncwarp();
// Update, wait and clean
EP_STATIC_ASSERT(kNumScaleoutRanks <= 32, "Invalid ranks");
const auto expected_signal = math::pack2<int, int64_t>(1, 0);
gin.flush<ncclCoopWarp>();
if (lane_idx < kNumScaleoutRanks) {
// Update remote tails
gin.red_add_rel<ncclTeamTagRail>(
workspace_layout.get_scaleout_channel_signaled_tail_ptr(channel_idx, scaleout_rank_idx),
expected_signal, lane_idx);
}
__syncwarp();
// Wait tail arrival
if (lane_idx < kNumScaleoutRanks) {
const auto wait_ptr = workspace_layout.get_scaleout_channel_signaled_tail_ptr(channel_idx, lane_idx);
comm::timeout_while<kNumTimeoutCycles>([=](const bool& is_last_check) {
const auto signal = ptx::ld_acquire_sys<int64_t>(wait_ptr);
if (signal == expected_signal) {
// Clean for next usages
*wait_ptr = 0;
return true;
}
if (is_last_check) {
printf("DeepEP combine (scale-out wait all) timeout, scale-out: %d/%d, scale-up: %d/%d, "
"channel: %d, lane: %d, signal: %lld, expected: %lld\n",
scaleout_rank_idx, kNumScaleoutRanks, scaleup_rank_idx, kNumScaleupRanks,
channel_idx, lane_idx,
signal, expected_signal);
}
return false;
});
}
__syncwarp();
}
// No barrier at epilogue
}
} // namespace deep_ep::elastic
deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh
+677
View File
@@ -0,0 +1,677 @@
#pragma once
#include <deep_ep/common/comm.cuh>
#include <deep_ep/common/compiled.cuh>
#include <deep_ep/common/exception.cuh>
#include <deep_ep/common/layout.cuh>
#include <deep_ep/common/math.cuh>
#include <deep_ep/common/ptx.cuh>
namespace deep_ep::elastic {
template <bool kDoCPUSync,
bool kReuseSlotIndices,
int kNumSMs,
int kNumNotifyWarps, int kNumScaleoutWarps, int kNumForwardWarps,
int kNumScaleoutRanks, int kNumScaleupRanks,
int kNumHiddenBytes, int kNumSFPacks,
int kNumMaxTokensPerRank,
int kNumExperts, int kNumTopk, int kExpertAlignment,
int kNumQPs, int64_t kNumTimeoutCycles,
int kNumScaleupRanksPerLane = math::constexpr_ceil_div(kNumScaleupRanks, 32),
int kNumChannelsPerSM = kNumScaleoutWarps,
int kNumChannels = kNumScaleoutWarps * kNumSMs,
int kNumMaxTokensPerChannel = math::constexpr_ceil_div(kNumMaxTokensPerRank, kNumChannels),
int kScaleoutUpdateInterval = 6,
int kNumSlotsPerForwardChunk = kScaleoutUpdateInterval,
int kNumRanks = kNumScaleoutRanks * kNumScaleupRanks,
int kNumNotifyThreads = kNumNotifyWarps * 32,
int kNumScaleoutSendThreads = kNumScaleoutWarps * 32,
int kNumForwardThreads = kNumForwardWarps * 32,
int kNumThreads = kNumNotifyThreads + kNumScaleoutSendThreads + kNumForwardThreads>
__global__ void __launch_bounds__(kNumThreads, 1)
hybrid_dispatch_impl(
void* x, sf_pack_t* sf, topk_idx_t* topk_idx, float* topk_weights,
topk_idx_t* copied_topk_idx,
int* cumulative_local_expert_recv_stats,
int* psum_num_recv_tokens_per_scaleup_rank,
int* psum_num_recv_tokens_per_expert,
int* num_unaligned_recv_tokens_per_expert,
int* dst_buffer_slot_idx,
int* token_metadata_at_forward,
const int num_tokens,
const int sf_token_stride, const int sf_hidden_stride,
// TODO(NCCL): so many params, plans to optimize?
const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window,
void* buffer,
void* workspace, void* mapped_host_workspace,
const int scaleout_rank_idx, const int scaleup_rank_idx) {
constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks;
constexpr int kNumExpertsPerScaleout = kNumExperts / kNumScaleoutRanks;
EP_STATIC_ASSERT(kNumExperts % kNumScaleupRanks == 0, "Invalid number of experts or ranks");
EP_STATIC_ASSERT(kNumNotifyWarps % 4 == 0, "Invalid warpgroup size");
EP_STATIC_ASSERT(kNumScaleoutWarps == kNumForwardWarps, "Invalid warp size");
// Utils
// NOTES: a warp is a channel (different channels may share QPs)
const auto sm_idx = static_cast<int>(blockIdx.x), thread_idx = static_cast<int>(threadIdx.x);
const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx();
const auto rank_idx = scaleout_rank_idx * kNumScaleupRanks + scaleup_rank_idx;
// Workspaces
const auto workspace_layout = layout::WorkspaceLayout(workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts);
const auto host_workspace_layout = layout::WorkspaceLayout(mapped_host_workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts);
// The kernel uses a fixed space of dynamic shared memory (no static shared memory)
extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[];
constexpr int kNumSmemBytesForNotify = kNumNotifyThreads > 0 ?
math::constexpr_align(kNumRanks + kNumExperts, kNumNotifyThreads) * sizeof(int) : 0;
EP_STATIC_ASSERT(kNumSmemBytesForNotify % ptx::kNumTMAAlignBytes == 0, "Invalid TMA alignment");
// Named barrier indices
constexpr int kNotifyBarrierIndex = 1;
// NCCL Gin handle
// Each warp is a channel
const auto [qp_idx, sharing_mode] = comm::get_qp_mode<kNumSMs, kNumQPs, kNumChannelsPerSM, (kNumNotifyWarps > 0)>(
sm_idx, (warp_idx - kNumNotifyWarps) % kNumChannelsPerSM, warp_idx < kNumNotifyWarps);
const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, sharing_mode);
// Global parallel barriers for scale-out subteam and scale-up subteam
comm::gpu_barrier<true, kNumScaleoutRanks, kNumScaleupRanks,
kNumSMs, kNumThreads, kNumQPs, kNumTimeoutCycles, comm::kHybridDispatchTag0, false, false, true>(
gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx);
// The golden layout during the whole process for both scale-out and forward warps
const auto token_layout = layout::TokenLayout(kNumHiddenBytes, kNumSFPacks * sizeof(sf_pack_t), kNumTopk, true);
const auto tma_buffer = layout::BufferLayout<true>(token_layout, kNumScaleoutWarps + kNumForwardWarps, 1,
math::advance_ptr<int>(smem, kNumSmemBytesForNotify)).get_rank_buffer(warp_idx - kNumNotifyWarps).get_token_buffer(0);
// All the buffers
auto scaleup_buffer = layout::BufferLayout<false>(
token_layout, kNumScaleupRanks, kNumScaleoutRanks * kNumMaxTokensPerRank, buffer);
auto scaleout_send_buffer = layout::BufferLayout<false>(
token_layout, 1, kNumMaxTokensPerRank, scaleup_buffer.get_buffer_end_ptr());
auto scaleout_recv_buffer = layout::BufferLayout<false>(
token_layout, kNumScaleoutRanks, kNumChannels * kNumMaxTokensPerChannel, scaleout_send_buffer.get_buffer_end_ptr());
// Init TMA for scale-out and forward warps
ptx::arrival_phase phase = 0;
const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr();
if (warp_idx >= kNumNotifyWarps and ptx::elect_one_sync())
ptx::mbarrier_init_with_fence(mbarrier_ptr, 1);
__syncwarp();
// Different warp roles
if (warp_idx < kNumNotifyWarps) {
// Assign shared memory
constexpr int kNumAlignedElems = kNumSmemBytesForNotify / sizeof(int);
const auto rank_expert_count = math::advance_ptr<int>(smem, 0);
// Clean initial counts
// NOTES: if you want to change the order of different warp roles, please take care of the `thread_idx`
int *rank_count = rank_expert_count, *expert_count = rank_expert_count + kNumRanks;
#pragma unroll
for (int i = 0; i < kNumAlignedElems / kNumNotifyThreads; ++ i)
rank_expert_count[i * kNumNotifyThreads + thread_idx] = 0;
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);
// Atomic add on shared memory
EP_STATIC_ASSERT(kNumTopk <= 32, "Insufficient lanes");
const auto global_warp_idx = sm_idx * kNumNotifyWarps + warp_idx;
for (int i = global_warp_idx; i < num_tokens; i += kNumNotifyWarps * kNumSMs) {
// Expert choice can not be redundant
// NOTES: no assertions here as they are expensive
const auto dst_expert_idx = lane_idx < kNumTopk ?
static_cast<int>(__ldg(topk_idx + i * kNumTopk + lane_idx)) : -1;
if (dst_expert_idx >= 0)
atomicAdd_block(expert_count + dst_expert_idx, 1);
// Rank choice should do deduplication here
const auto dst_rank_idx = dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerRank : -1;
if (ptx::deduplicate(dst_rank_idx, lane_idx) and dst_rank_idx >= 0)
atomicAdd_block(rank_count + dst_rank_idx, 1);
}
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);
// Do full-grid reduction
#pragma unroll
for (int i = thread_idx; i < kNumRanks + kNumExperts; i += kNumNotifyThreads) {
const int64_t counter = (1ll << 32ll) | rank_expert_count[i];
ptx::red_add(workspace_layout.get_notify_reduction_workspace_ptr() + i, counter);
}
// Do the remaining work by SM 0
if (sm_idx == 0) {
// Reduce all SM's count
// Wait all SMs' arrival
#pragma unroll
for (int i = thread_idx; i < kNumRanks + kNumExperts; i += kNumNotifyThreads) {
comm::timeout_while<kNumTimeoutCycles>([=](const bool& is_last_check) {
const auto status = ptx::ld_volatile<int64_t>(workspace_layout.get_notify_reduction_workspace_ptr() + i);
if ((status >> 32) == kNumSMs) {
// Encode and write into the send buffer
workspace_layout.get_scaleout_rank_expert_count_ptr<true>()[i] =
math::encode_decode_positive<int>(status & 0xffffffffll);
// Clean for the next usage
workspace_layout.get_notify_reduction_workspace_ptr()[i] = 0;
return true;
}
if (is_last_check) {
printf("DeepEP hybrid notify (GPU reduction) timeout, scale-out: %d/%d, scale-up: %d/%d, "
"thread: %d, status: %d | %d, expected: %d\n",
scaleout_rank_idx, kNumScaleoutRanks, scaleup_rank_idx, kNumScaleupRanks, thread_idx,
static_cast<int>(status >> 32), static_cast<int>(status & 0xffffffff), kNumSMs);
}
return false;
});
}
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);
// Issue scaleout writes to peers
EP_STATIC_ASSERT(kReuseSlotIndices or kNumScaleoutRanks <= kNumNotifyThreads,
"kNumScaleoutRanks must be less than kNumNotifyThreads");
if (thread_idx < kNumScaleoutRanks) {
const auto dst_scaleout_rank_idx = thread_idx;
gin.put<ncclTeamTagRail>(
workspace_layout.get_scaleout_rank_count_ptr<false>(scaleout_rank_idx),
workspace_layout.get_scaleout_rank_count_ptr<true>(dst_scaleout_rank_idx),
kNumScaleupRanks * sizeof(int), dst_scaleout_rank_idx,
ncclGinOptFlagsAggregateRequests);
gin.put<ncclTeamTagRail>(
workspace_layout.get_scaleout_expert_count_ptr<false>(scaleout_rank_idx),
workspace_layout.get_scaleout_expert_count_ptr<true>(dst_scaleout_rank_idx),
kNumExpertsPerScaleout * sizeof(int), dst_scaleout_rank_idx);
}
__syncwarp();
// Util functions to get metadata from scale-out peers
// NOTES: this is correct as RDMA operations has a minimum write granularity of 1024 bytes (a whole integer write is atomic)
const auto recv_and_reduce = [=](const auto& get_ptr_func, const bool& is_expert_reduction = false) -> int {
int count = 0;
#pragma unroll
for (int j = 0; j < kNumScaleoutRanks; ++ j) {
const auto ptr = get_ptr_func(j);
int decoded;
comm::timeout_while<kNumTimeoutCycles>([&](const bool& is_last_check){
decoded = math::encode_decode_positive(ptx::ld_acquire_sys<int>(ptr));
if (math::is_decoded_positive_ready(decoded))
return true;
if (is_last_check) {
printf("DeepEP hybrid notify (scale-out %s reduction) timeout, "
"scale-out: %d, scale-up: %d, "
"thread: %d, wait scale-out: %d, decoded: %d\n",
is_expert_reduction ? "expert" : "rank",
scaleout_rank_idx, scaleup_rank_idx, thread_idx, j,
decoded);
}
return false;
});
// Add and clean for next usages
count += decoded, *ptr = 0;
}
return count;
};
// Write into all scale-up peers' rank-level counters
#pragma unroll
for (int i = thread_idx; i < kNumScaleupRanks; i += kNumNotifyThreads) {
// Wait scale-out arrival and reduce
const auto count = recv_and_reduce([=](const int& scaleout_peer_idx) {
return workspace_layout.get_scaleout_rank_count_ptr<false>(scaleout_peer_idx, i);
});
// Write into the remote scale-up peer
const int64_t counter = (static_cast<int64_t>(kNumScaleupRanks) << 32ll) | count;
gin.put_value<ncclTeamTagLsa>(
workspace_layout.get_scaleup_rank_count_ptr<false>() + scaleup_rank_idx,
counter, i);
}
__syncwarp();
// Atomic add into all scale-up peers' expert-level counters
#pragma unroll
for (int i = thread_idx; i < kNumExpertsPerScaleout; i += kNumNotifyThreads) {
// Wait scale-out arrival and reduce
const auto count = recv_and_reduce([=](const int& scaleout_peer_idx) {
return workspace_layout.get_scaleout_expert_count_ptr<false>(scaleout_peer_idx, i);
}, true);
// Write into the remote scale-up peer
const int64_t counter = (1ll << 32ll) | count;
const auto dst_scaleup_rank_idx = i / kNumExpertsPerRank;
const auto expert_idx_in_dst_rank = i % kNumExpertsPerRank;
gin.red_add_rel<ncclTeamTagLsa>(
workspace_layout.get_scaleup_expert_count_ptr<false>() + expert_idx_in_dst_rank,
counter, dst_scaleup_rank_idx);
}
// There are shared memory reads above, a barrier is necessary
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);
// NOTES: from now on, the `rank` and `expert`s size change into the local size
expert_count = rank_expert_count + kNumScaleupRanks;
// Wait local counters to be ready
// NOTES: here we only care the prefix sum by scale-up peers (used for later epilogue), not all ranks
EP_STATIC_ASSERT(kNumNotifyWarps == 0 or kNumScaleupRanks + kNumExpertsPerRank <= kNumNotifyWarps * 32,
"Insufficient notify threads");
comm::timeout_while<kNumTimeoutCycles>(thread_idx < kNumScaleupRanks + kNumExpertsPerRank,
[&](const bool& is_last_check) {
const auto status = ptx::ld_volatile<int64_t>(workspace_layout.get_scaleup_rank_expert_count_ptr<false>() + thread_idx);
if ((status >> 32ull) == kNumScaleupRanks) {
// Clean GPU workspace and write into host workspace
const auto count = static_cast<int>(status & 0xffffffffll);
const auto aligned_count = math::align<int>(
count, thread_idx < kNumScaleupRanks ? 1 : kExpertAlignment);
workspace_layout.get_scaleup_rank_expert_count_ptr<false>()[thread_idx] = 0;
if constexpr (kDoCPUSync) {
host_workspace_layout.get_scaleup_rank_expert_count_ptr<false>()[thread_idx] =
math::encode_decode_positive(aligned_count);
}
// Update statistics counters
if (cumulative_local_expert_recv_stats != nullptr and thread_idx >= kNumScaleupRanks)
atomicAdd(cumulative_local_expert_recv_stats + (thread_idx - kNumScaleupRanks), count);
// Write unaligned count before aligning
if (num_unaligned_recv_tokens_per_expert != nullptr and thread_idx >= kNumScaleupRanks)
num_unaligned_recv_tokens_per_expert[thread_idx - kNumScaleupRanks] = count;
// Save for later prefix sum calculation
rank_expert_count[thread_idx] = aligned_count;
return true;
}
if (is_last_check) {
printf("DeepEP hybrid notify (scale-up reduction) timeout,"
"scale-out: %d/%d, scale-up: %d/%d, "
"thread: %d, status: %d | %d, expected: %d\n",
scaleout_rank_idx, kNumScaleoutRanks, scaleup_rank_idx, kNumScaleupRanks, thread_idx,
static_cast<int>(status >> 32), static_cast<int>(status & 0xffffffff), kNumScaleupRanks);
}
return false;
});
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);
// Do prefix sum by the warps of the first SM
// NOTES: we may have fast implementation with `cub::BlockScan`, but it is too heavy to use
const auto do_psum = [=](const int* count, int* out, const int n, const int is_exclusive) {
int psum = 0;
#pragma unroll
for (int i = 0; i < math::ceil_div(n + is_exclusive, 32); ++ i) {
const auto idx = i * 32 + lane_idx;
const auto mem_idx = idx - is_exclusive;
const auto value = (0 <= mem_idx and mem_idx < n) ? count[mem_idx] : 0;
const auto sum = psum + ptx::warp_inclusive_sum(value, lane_idx);
// Store into global memory
if (idx < n + is_exclusive)
out[idx] = sum;
// Update `psum` by using the last lane's value
psum = ptx::exchange(sum, 31);
}
};
if (warp_idx == 0) {
// Inclusive prefix sum
do_psum(rank_count, psum_num_recv_tokens_per_scaleup_rank, kNumScaleupRanks, 0);
} else if (warp_idx == 1) {
// Exclusive prefix sum for later expanding
do_psum(expert_count, psum_num_recv_tokens_per_expert, kNumExpertsPerRank, 1);
}
}
} else if (warp_idx < kNumNotifyWarps + kNumScaleoutWarps) {
const int scaleout_warp_idx = warp_idx - kNumNotifyWarps;
const int channel_idx = sm_idx * kNumChannelsPerSM + scaleout_warp_idx;
scaleout_recv_buffer = scaleout_recv_buffer.get_rank_buffer(scaleout_rank_idx);
scaleout_recv_buffer = scaleout_recv_buffer.get_channel_buffer<kNumMaxTokensPerChannel>(channel_idx);
// Channel metadata maintenance
EP_STATIC_ASSERT(kNumScaleoutRanks <= 32, "Invalid number of scale-out ranks");
int stored_scaleout_tail = 0, stored_old_scaleout_tail = 0;
const auto update_scaleout_tail = [&](const bool& finish_flag = false) {
if (lane_idx < kNumScaleoutRanks and
(stored_scaleout_tail >= stored_old_scaleout_tail + kScaleoutUpdateInterval or finish_flag)) {
const auto signaled_tail = math::pack2<int, int64_t>(finish_flag, stored_scaleout_tail);
const auto ptr = workspace_layout.get_scaleout_channel_signaled_tail_ptr(channel_idx, scaleout_rank_idx);
const auto old_signaled_tail = math::pack2<int, int64_t>(0, stored_old_scaleout_tail);
// NOTES: the "release" scope will be `sys` for the local rank (we may involve NVLink so not `gpu`)
// For RDMA requests, "release" is ensured by "atomic"
gin.red_add_rel<ncclTeamTagRail>(ptr, signaled_tail - old_signaled_tail, lane_idx);
stored_old_scaleout_tail = stored_scaleout_tail;
}
__syncwarp();
};
// Preload next token
const auto preload_next_token = [&](const int& token_idx) {
if (token_idx >= num_tokens)
return;
// Issue TMA load
const auto token_i64_idx = static_cast<int64_t>(token_idx);
if (ptx::elect_one_sync()) {
ptx::tma_load_1d(tma_buffer.get_hidden_ptr(), math::advance_ptr(x, token_i64_idx * kNumHiddenBytes),
mbarrier_ptr, kNumHiddenBytes);
}
__syncwarp();
// Issue SF `cp.async`
if constexpr (kNumSFPacks > 0) {
EP_STATIC_ASSERT(sizeof(sf_pack_t) % 4 == 0, "Unaligned SF element type");
const auto gmem_src_ptr = math::advance_ptr<sf_pack_t>(sf, token_i64_idx * sf_token_stride * sizeof(sf_pack_t));
const auto smem_dst_ptr = tma_buffer.get_sf_ptr();
constexpr auto kNumFullIters = kNumSFPacks / 32;
#pragma unroll
for (int k = 0; k < kNumFullIters; ++ k) {
ptx::cp_async_ca(gmem_src_ptr + (k * 32 + lane_idx) * sf_hidden_stride,
smem_dst_ptr + k * 32 + lane_idx);
}
if (kNumFullIters * 32 + lane_idx < kNumSFPacks) {
ptx::cp_async_ca(gmem_src_ptr + (kNumFullIters * 32 + lane_idx) * sf_hidden_stride,
smem_dst_ptr + kNumFullIters * 32 + lane_idx);
}
ptx::cp_async_mbarrier_arrive(mbarrier_ptr);
__syncwarp();
}
};
// Iterate all tokens
preload_next_token(channel_idx);
for (int token_idx = channel_idx; token_idx < num_tokens; token_idx += kNumChannels) {
// Load top-k indices and weights
EP_STATIC_ASSERT(kNumTopk <= 32, "Insufficient lanes for loading top-k indices");
int stored_dst_scaleout_rank_idx = -1;
if (lane_idx < kNumTopk) {
const auto uncasted_dst_expert_idx = __ldg(topk_idx + token_idx * kNumTopk + lane_idx);
const auto dst_expert_idx = static_cast<int>(uncasted_dst_expert_idx);
stored_dst_scaleout_rank_idx = dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerScaleout : -1;
tma_buffer.get_topk_idx_ptr()[lane_idx] = dst_expert_idx;
if (topk_weights != nullptr)
tma_buffer.get_topk_weights_ptr()[lane_idx] = __ldg(topk_weights + token_idx * kNumTopk + lane_idx);
if (copied_topk_idx != nullptr)
copied_topk_idx[token_idx * kNumTopk + lane_idx] = uncasted_dst_expert_idx;
}
__syncwarp();
// Add source metadata (rank index and token index)
if (ptx::elect_one_sync())
*tma_buffer.get_src_token_global_idx_ptr() = rank_idx * kNumMaxTokensPerRank + token_idx;
ptx::tma_store_fence();
__syncwarp();
// Deduplicate ranks and assign slots
int stored_dst_slot_idx = -1;
const auto stored_old_slot_idx = ptx::exchange(
stored_scaleout_tail, stored_dst_scaleout_rank_idx >= 0 ? stored_dst_scaleout_rank_idx : 0);
if (ptx::deduplicate(stored_dst_scaleout_rank_idx, lane_idx) and stored_dst_scaleout_rank_idx >= 0)
stored_dst_slot_idx = stored_old_slot_idx;
// Update scale-out tail
const auto scaleout_rank_mask = ptx::reduce_or(stored_dst_scaleout_rank_idx >= 0 ? (1u << stored_dst_scaleout_rank_idx) : 0u);
stored_scaleout_tail += (scaleout_rank_mask >> lane_idx) & 1;
// Wait TMA arrival and issue the TMA store into send buffer
if (ptx::elect_one_sync()) {
ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes);
ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase);
// So if no ranks will go by RDMA, we skip the send buffer stores
if (scaleout_rank_mask ^ (1 << scaleout_rank_idx)) {
ptx::tma_store_1d(scaleout_send_buffer.get_token_buffer(token_idx).get_base_ptr(),
tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes<false>());
}
}
__syncwarp();
// Local rank can be bypassed
if (stored_dst_slot_idx >= 0 and stored_dst_scaleout_rank_idx == scaleout_rank_idx) {
ptx::tma_store_1d(scaleout_recv_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(),
tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes<false>());
}
ptx::tma_store_commit();
ptx::tma_store_wait();
__syncwarp();
// Preload the next token (overlapping with the IBGDA issues)
preload_next_token(token_idx + kNumChannels);
// Issue IBGDA requests
if (stored_dst_slot_idx >= 0 and stored_dst_scaleout_rank_idx != scaleout_rank_idx) {
gin.put<ncclTeamTagRail>(
scaleout_recv_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(),
scaleout_send_buffer.get_token_buffer(token_idx).get_base_ptr(),
tma_buffer.get_num_bytes<false>(),
stored_dst_scaleout_rank_idx,
ncclGinOptFlagsAggregateRequests);
}
__syncwarp();
// Issue scale-out tail update
update_scaleout_tail();
}
// Flush unflushed tails
update_scaleout_tail(true);
} else {
const int forward_warp_idx = warp_idx - (kNumNotifyWarps + kNumScaleoutWarps);
const int channel_idx = sm_idx * kNumChannelsPerSM + forward_warp_idx;
scaleout_recv_buffer = scaleout_recv_buffer.get_channel_buffer<kNumMaxTokensPerChannel>(channel_idx);
scaleup_buffer = scaleup_buffer.get_rank_buffer(scaleup_rank_idx);
// Shape of `token_metadata_at_forward`: `[kNumChannels, kNumScaleoutRanks * kNumMaxTokensPerChannel + 1, kNumForwardMetadataDims]`
constexpr int kNumForwardMetadataDims = 2 + kNumTopk * 2;
token_metadata_at_forward += channel_idx * ((kNumScaleoutRanks * kNumMaxTokensPerChannel + 1) * kNumForwardMetadataDims);
// Shape of `dst_buffer_slot_idx`: `[kNumChannels, kNumScaleoutRanks, kNumMaxTokensPerChannel, kNumTopk]`
dst_buffer_slot_idx += channel_idx * (kNumScaleoutRanks * kNumMaxTokensPerChannel * kNumTopk);
// Transform linked list index
const auto transform_linked_list_idx = [=](const int& idx) {
constexpr int kNumTokensInLinkedList = kNumMaxTokensPerChannel * kNumScaleoutRanks + 1;
return channel_idx * (kNumTokensInLinkedList * kNumScaleupRanks) +
idx * kNumScaleupRanks + scaleup_rank_idx;
};
// Forward tokens from scale-out ranks
EP_STATIC_ASSERT(kNumScaleoutRanks <= 32, "Too many scale-out ranks");
int num_tokens_processed = 0;
int stored_scaleout_old_tail_idx = 0;
int stored_scaleup_send_counters[kNumScaleupRanksPerLane] = {};
int stored_finish_flag = lane_idx >= kNumScaleoutRanks;
int stored_scaleout_tail_idx = 0;
int recv_scaleout_rank_idx = channel_idx % kNumScaleoutRanks;
uint32_t wip_mask;
while ((wip_mask = ptx::gather(stored_scaleout_tail_idx > stored_scaleout_old_tail_idx or stored_finish_flag == 0))) {
// Pick next rank in round-robin
const auto offset = (recv_scaleout_rank_idx + 1) % kNumScaleoutRanks;
const auto hi_mask = (wip_mask >> offset) << offset;
recv_scaleout_rank_idx = hi_mask ? ptx::ffs(hi_mask) : ptx::ffs(wip_mask);
// Wait for this rank to have data (or finish)
comm::timeout_while<kNumTimeoutCycles>([&](const bool& is_last_check) {
const uint32_t arrived_or_finished =
stored_scaleout_tail_idx > stored_scaleout_old_tail_idx or stored_finish_flag > 0;
if (ptx::exchange(arrived_or_finished, recv_scaleout_rank_idx))
return true;
// Timeout
if (is_last_check) {
if (lane_idx < kNumScaleoutRanks) {
printf("DeepEP hybrid dispatch (forwarding) timeout, scale-out: %d, scale-up: %d, "
"channel: %d, lane: %d, old scale-out tail: %d, scale-out tail: (%d, %d)\n",
scaleout_rank_idx, scaleup_rank_idx,
channel_idx, lane_idx, stored_scaleout_old_tail_idx,
stored_finish_flag, stored_scaleout_tail_idx);
}
return false;
}
// Read new signaled tails
if (lane_idx < kNumScaleoutRanks) {
const auto signaled_tail = ptx::ld_acquire_sys<int64_t>(
workspace_layout.get_scaleout_channel_signaled_tail_ptr(channel_idx, lane_idx));
math::unpack2<int, int64_t>(signaled_tail, stored_finish_flag, stored_scaleout_tail_idx);
}
__syncwarp();
return false;
});
// Process one chunk from the current rank
const auto start_slot_idx = ptx::exchange(stored_scaleout_old_tail_idx, recv_scaleout_rank_idx);
const auto end_slot_idx = std::min(
ptx::exchange(stored_scaleout_tail_idx, recv_scaleout_rank_idx),
start_slot_idx + kNumSlotsPerForwardChunk
);
if (lane_idx == recv_scaleout_rank_idx)
stored_scaleout_old_tail_idx = end_slot_idx;
const auto recv_buffer = scaleout_recv_buffer.get_rank_buffer(recv_scaleout_rank_idx);
for (int slot_idx = start_slot_idx; slot_idx < end_slot_idx; ++ slot_idx) {
const auto token_buffer = recv_buffer.get_token_buffer(slot_idx);
// Wait TMA arrival
ptx::tma_store_wait();
__syncwarp();
// TMA load into shared memory
if (ptx::elect_one_sync()) {
ptx::tma_load_1d(tma_buffer.get_base_ptr(), token_buffer.get_base_ptr(),
mbarrier_ptr, token_layout.get_num_bytes<false>());
ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, token_layout.get_num_bytes<false>());
ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase);
}
__syncwarp();
// Read top-k indices
EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k selections");
int stored_dst_scaleup_rank_idx = -1;
auto dst_expert_idx = lane_idx < kNumTopk ? tma_buffer.get_topk_idx_ptr()[lane_idx] : -1;
dst_expert_idx -= scaleout_rank_idx * kNumExpertsPerScaleout;
stored_dst_scaleup_rank_idx = 0 <= dst_expert_idx and dst_expert_idx < kNumExpertsPerScaleout ?
dst_expert_idx / kNumExpertsPerRank : -1;
// Write the per-scaleup channel index for this token
int linked_list_idx = -1;
#pragma unroll
for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) {
const auto src_lane_idx = stored_dst_scaleup_rank_idx - j * 32;
const bool valid = 0 <= src_lane_idx and src_lane_idx < 32;
const auto exchanged = ptx::exchange(
stored_scaleup_send_counters[j], valid ? src_lane_idx : 0);
linked_list_idx = valid ? exchanged : linked_list_idx;
}
if (not kReuseSlotIndices and lane_idx < kNumTopk) {
tma_buffer.get_linked_list_idx_ptr()[lane_idx] = transform_linked_list_idx(linked_list_idx);
ptx::tma_store_fence();
}
__syncwarp();
// Deduplicate for scale-up ranks
int stored_dst_slot_idx = -1;
const auto dst_slot_idx_ptr = dst_buffer_slot_idx +
recv_scaleout_rank_idx * (kNumMaxTokensPerChannel * kNumTopk) + slot_idx * kNumTopk;
if constexpr (kReuseSlotIndices) {
if (lane_idx < kNumTopk)
stored_dst_slot_idx = __ldg(dst_slot_idx_ptr + lane_idx);
} else {
// Deduplicate for NVLink ranks
if (ptx::deduplicate(stored_dst_scaleup_rank_idx, lane_idx) and stored_dst_scaleup_rank_idx >= 0)
stored_dst_slot_idx = atomicAdd(workspace_layout.get_scaleup_atomic_sender_counter() + stored_dst_scaleup_rank_idx, 1);
}
__syncwarp();
// Issue TMAs
if (stored_dst_slot_idx >= 0) {
const auto dst_ptr = gin.get_sym_ptr<ncclTeamTagLsa>(
scaleup_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(),
stored_dst_scaleup_rank_idx);
ptx::tma_store_1d(dst_ptr, tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes<false>());
ptx::tma_store_commit();
}
__syncwarp();
// Add per-scale-up counter
EP_STATIC_ASSERT(kNumScaleupRanks <= 64, "Invalid number of scale-up peers");
using mask_t = std::conditional_t<kNumScaleupRanks <= 32, unsigned, unsigned long long>;
const auto scaleup_send_mask = ptx::reduce_or(
stored_dst_scaleup_rank_idx >= 0 ?
(mask_t(1) << stored_dst_scaleup_rank_idx) : mask_t(0));
#pragma unroll
for (int j = 0; j < kNumScaleupRanksPerLane; ++ j)
stored_scaleup_send_counters[j] += (scaleup_send_mask >> (j * 32 + lane_idx)) & 1;
// Record metadata at forward
if constexpr (not kReuseSlotIndices) {
EP_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of selections");
const auto metadata_ptr = token_metadata_at_forward +
num_tokens_processed * kNumForwardMetadataDims;
// Source token index and last token index flag
if (ptx::elect_one_sync()) {
metadata_ptr[0] = tma_buffer.get_src_token_global_idx_ptr()[0];
metadata_ptr[1] = slot_idx == (end_slot_idx - 1);
}
// Second, original top-k indices and destination slots
if (lane_idx < kNumTopk) {
metadata_ptr[2 + lane_idx] = stored_dst_scaleup_rank_idx;
metadata_ptr[2 + kNumTopk + lane_idx] = stored_dst_slot_idx;
dst_slot_idx_ptr[lane_idx] = stored_dst_slot_idx;
}
}
num_tokens_processed += 1;
__syncwarp();
}
}
// Assign the source token index part of the metadata into `-1` as an ending mark
if (not kReuseSlotIndices and ptx::elect_one_sync())
token_metadata_at_forward[num_tokens_processed * kNumForwardMetadataDims] = -1;
__syncwarp();
// Update linked list's ending position
if constexpr (not kReuseSlotIndices) {
const auto tail_ptr = workspace_layout.get_channel_scaleup_tail_ptr(channel_idx, scaleup_rank_idx);
#pragma unroll
for (int i = 0; i < kNumScaleupRanksPerLane; ++ i) {
if (const auto j = i * 32 + lane_idx; i < (kNumScaleupRanksPerLane - 1) or j < kNumScaleupRanks) {
ptx::st_relaxed_sys(
gin.get_sym_ptr<ncclTeamTagLsa>(tail_ptr, j),
transform_linked_list_idx(stored_scaleup_send_counters[i]));
}
}
}
__syncwarp();
// Clean tails for next usages
if (lane_idx < kNumScaleoutRanks)
*workspace_layout.get_scaleout_channel_signaled_tail_ptr(channel_idx, lane_idx) = 0;
__syncwarp();
}
// Scale-up barrier to ensure data arrival
// As scale-out tokens have already been consumed by forwarders, no need to do scale-out barrier again
comm::gpu_barrier<true, kNumScaleoutRanks, kNumScaleupRanks,
kNumSMs, kNumThreads, kNumQPs, kNumTimeoutCycles, comm::kHybridDispatchTag1, true, true, false>(
gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, thread_idx, /* do not scale-out */ false, true);
// Trigger the copy epilogue kernel
cudaTriggerProgrammaticLaunchCompletion();
// Clean scale-up counters
// All scale-out counters should be cleaned before
EP_STATIC_ASSERT(kNumScaleupRanks <= kNumThreads, "Insufficient threads");
if (not kReuseSlotIndices and sm_idx == 0 and thread_idx < kNumScaleupRanks)
workspace_layout.get_scaleup_atomic_sender_counter()[thread_idx] = 0;
}
} // namespace deep_ep::elastic
deep_ep/include/deep_ep/impls/pp_send_recv.cuh
+214
View File
@@ -0,0 +1,214 @@
#pragma once
#include <cooperative_groups.h>
#include <deep_ep/common/comm.cuh>
#include <deep_ep/common/layout.cuh>
#include <deep_ep/common/ptx.cuh>
namespace deep_ep::elastic {
template <int kNumRanks>
__device__ __forceinline__ std::pair<int, int> get_buffer_offset(
const int& src_rank_idx, const int& dst_rank_idx) {
const auto next_rank_idx = (src_rank_idx + 1) % kNumRanks;
return dst_rank_idx == next_rank_idx ? std::make_pair(0, 1) : std::make_pair(1, 0);
}
template <int64_t kNumTimeoutCycles, typename timeout_print_t>
__device__ __forceinline__ void check_signal(
const handle::NCCLGin& gin,
const ncclGinSignal_t& signal_idx,
const int64_t& target,
const timeout_print_t& timeout_print) {
const auto gdaki = static_cast<struct ncclGinGdakiGPUContext*>(gin.gin._ginHandle) + gin.gin.contextId;
const auto signal_ptr = reinterpret_cast<int64_t*>(
__ldg(reinterpret_cast<int64_t*>(&gdaki->signals_table.buffer))) + signal_idx;
comm::timeout_while<kNumTimeoutCycles>([=](const bool& is_last_check) {
const auto signal = ptx::ld_acquire_sys<int64_t>(signal_ptr);
if (signal >= target)
return true;
if (is_last_check)
timeout_print();
return false;
});
}
template <int kNumSMs,
int kNumSmemBytes,
int kNumStages = 2,
int kNumTMABytesPerStage = math::constexpr_align<int, false>(
(kNumSmemBytes - kNumStages * sizeof(ptx::mbarrier)) / kNumStages, ptx::kNumTMAAlignBytes),
int kNumTMABlocksPerStage = kNumTMABytesPerStage / ptx::kNumTMAAlignBytes>
__device__ __forceinline__ void tma_copy(
void* src_ptr, void* dst_ptr,
const int64_t& num_bytes, const int& sm_idx) {
extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[];
const auto tma_buffers = smem;
const auto mbarriers = reinterpret_cast<ptx::mbarrier*>(smem + kNumStages * kNumTMABytesPerStage);
EP_STATIC_ASSERT(kNumTMABytesPerStage > 0, "Invalid shared memory bytes");
EP_STATIC_ASSERT(kNumStages >= 2, "Need at least 2 stages for pipelining");
// Init mbarriers
ptx::arrival_phase phases[kNumStages];
#pragma unroll
for (int s = 0; s < kNumStages; ++ s)
phases[s] = 0, ptx::mbarrier_init_with_fence(mbarriers + s, 1);
// Work partitioning across SMs
EP_DEVICE_ASSERT(num_bytes % ptx::kNumTMAAlignBytes == 0);
const auto num_tma_blocks = num_bytes / ptx::kNumTMAAlignBytes;
const auto num_tma_blocks_per_sm = math::ceil_div<int64_t>(num_tma_blocks, kNumSMs);
const auto start_block_idx = sm_idx * num_tma_blocks_per_sm;
const auto end_block_idx = std::min(start_block_idx + num_tma_blocks_per_sm, num_tma_blocks);
const auto num_iterations = math::ceil_div<int64_t>(end_block_idx - start_block_idx, kNumTMABlocksPerStage);
auto get_iter_info = [&](const int64_t& iter_idx) {
const auto i = start_block_idx + iter_idx * kNumTMABlocksPerStage;
const auto offset = i * ptx::kNumTMAAlignBytes;
const auto num_transaction_bytes =
std::min<int>(kNumTMABlocksPerStage, end_block_idx - i) * ptx::kNumTMAAlignBytes;
return std::make_pair(offset, num_transaction_bytes);
};
// Fill pipeline: issue loads for the first kNumStages iterations
for (int64_t iter_idx = 0; iter_idx < kNumStages and iter_idx < num_iterations; ++ iter_idx) {
const auto [load_offset, num_load_bytes] = get_iter_info(iter_idx);
ptx::tma_load_1d(
tma_buffers + iter_idx * kNumTMABytesPerStage,
math::advance_ptr(src_ptr, load_offset),
mbarriers + iter_idx, num_load_bytes);
ptx::mbarrier_arrive_and_set_tx(mbarriers + iter_idx, num_load_bytes);
}
for (int64_t iter_idx = 0; iter_idx < num_iterations; ++ iter_idx) {
const auto stage_idx = static_cast<int>(iter_idx % kNumStages);
const auto [store_offset, num_store_bytes] = get_iter_info(iter_idx);
// Wait this stage's load and issue store
ptx::mbarrier_wait_and_flip_phase(mbarriers + stage_idx, phases[stage_idx]);
ptx::tma_store_1d(
math::advance_ptr(dst_ptr, store_offset),
tma_buffers + stage_idx * kNumTMABytesPerStage,
num_store_bytes);
ptx::tma_store_commit();
// Prefetch: wait until this stage's store is completed, then issue next load
const auto next_iter_idx = iter_idx + kNumStages;
if (next_iter_idx < num_iterations) {
ptx::tma_store_wait();
const auto [load_offset, num_load_bytes] = get_iter_info(next_iter_idx);
ptx::tma_load_1d(
tma_buffers + stage_idx * kNumTMABytesPerStage,
math::advance_ptr(src_ptr, load_offset),
mbarriers + stage_idx, num_load_bytes);
ptx::mbarrier_arrive_and_set_tx(mbarriers + stage_idx, num_load_bytes);
}
}
// Drain all outstanding stores
ptx::tma_store_wait();
}
template <int kNumSMs,
int kNumRanks,
int kNumSmemBytes,
int64_t kNumTimeoutCycles>
__global__ void __launch_bounds__(32, 1)
pp_send_impl(const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window,
void* x, const int64_t num_x_bytes,
void* buffer, void* workspace,
const int rank_idx, const int dst_rank_idx,
const int64_t num_max_tensor_bytes,
const int num_max_inflight_tensors) {
const auto sm_idx = static_cast<int>(blockIdx.x);
const auto workspace_layout = layout::WorkspaceLayout(workspace, 1, kNumRanks, 0);
const auto [local_idx_in_dst, dst_idx_in_local] = get_buffer_offset<kNumRanks>(rank_idx, dst_rank_idx);
// Gin handle
const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, 0, NCCL_GIN_RESOURCE_SHARING_CTA);
// Buffer offsets
const auto send_count_ptr = workspace_layout.get_pp_send_count_ptr(dst_idx_in_local);
const auto send_count = __ldg(send_count_ptr);
const auto slot_idx = send_count % num_max_inflight_tensors;
auto send_buffer_ptr = math::advance_ptr(
buffer, ((dst_idx_in_local + 2) * num_max_inflight_tensors + slot_idx) * num_max_tensor_bytes);
auto recv_buffer_ptr = math::advance_ptr(
buffer, ((local_idx_in_dst + 0) * num_max_inflight_tensors + slot_idx) * num_max_tensor_bytes);
// Wait buffer slot release and do TMA
if (ptx::elect_one_sync()) {
check_signal<kNumTimeoutCycles>(
gin,
static_cast<ncclGinSignal_t>(kNumRanks + dst_idx_in_local + 2),
send_count - num_max_inflight_tensors + 1,
// TODO: print more info, and control the SM who prints it
[]() { printf("DeepEP PP send timeout, recv buffer is full"); }
);
tma_copy<kNumSMs, kNumSmemBytes>(x, send_buffer_ptr, num_x_bytes, sm_idx);
}
cooperative_groups::this_grid().sync();
// Issue RDMA put
if (sm_idx == 0 and ptx::elect_one_sync()) {
gin.put<ncclTeamTagWorld>(
recv_buffer_ptr,
send_buffer_ptr,
num_x_bytes, dst_rank_idx,
0,
// TODO: is this signal highly optimized?
ncclGin_SignalInc(static_cast<ncclGinSignal_t>(local_idx_in_dst + kNumRanks)));
*send_count_ptr += 1;
}
}
template <int kNumSMs,
int kNumRanks,
int kNumSmemBytes,
int64_t kNumTimeoutCycles>
__global__ void __launch_bounds__(32, 1)
pp_recv_impl(const ncclDevComm_t nccl_dev_comm, const ncclWindow_t nccl_window,
void* x, int64_t num_x_bytes,
void* buffer, void* workspace,
const int rank_idx, const int src_rank_idx,
const int64_t num_max_tensor_bytes,
const int num_max_inflight_tensors) {
const auto sm_idx = static_cast<int>(blockIdx.x);
const auto workspace_layout = layout::WorkspaceLayout(workspace, 1, kNumRanks, 0);
const auto [src_idx_in_local, local_idx_in_src] = get_buffer_offset<kNumRanks>(src_rank_idx, rank_idx);
// Gin handle
const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, 0, NCCL_GIN_RESOURCE_SHARING_CTA);
// Buffer offsets
const auto recv_count_ptr = workspace_layout.get_pp_recv_count_ptr(src_idx_in_local);
const auto recv_count = __ldg(recv_count_ptr);
const auto slot_idx = recv_count % num_max_inflight_tensors;
const auto recv_buffer_ptr = math::advance_ptr(
buffer, ((src_idx_in_local + 0) * num_max_inflight_tensors + slot_idx) * num_max_tensor_bytes);
// Copy from the buffer into a new tensor
if (ptx::elect_one_sync()) {
check_signal<kNumTimeoutCycles>(
gin,
static_cast<ncclGinSignal_t>(src_idx_in_local + kNumRanks),
recv_count + 1,
// TODO: print more info, and control the SM who prints it
[]() { printf("DeepEP PP recv timeout, recv buffer is empty\n"); }
);
tma_copy<kNumSMs, kNumSmemBytes>(recv_buffer_ptr, x, num_x_bytes, sm_idx);
}
cooperative_groups::this_grid().sync();
// TODO: add a comment
if (sm_idx == 0 and ptx::elect_one_sync()) {
gin.signal<ncclTeamTagWorld>(
src_rank_idx, ncclGin_SignalInc(static_cast<ncclGinSignal_t>(kNumRanks + local_idx_in_src + 2))
);
*recv_count_ptr += 1;
}
}
} // namespace deep_ep::elastic
deep_ep/utils/__init__.py
+3
View File
@@ -0,0 +1,3 @@
# For forward compatibility
# noinspection PyUnresolvedReferences
from .event import EventHandle
deep_ep/utils/comm.py
+83
View File
@@ -0,0 +1,83 @@
import os
import time
import torch
import torch.distributed as dist
# noinspection PyUnresolvedReferences
import deep_ep._C as _C
class NCCLCommHandle:
"""
A wrapper around a raw NCCL communicator. Manages the lifecycle of the communicator if created by DeepEP,
or simply wraps an existing one if obtained from PyTorch.
Attributes:
nccl_comm: the raw NCCL communicator.
managed: whether the communicator was created by DeepEP and should be destroyed when this handle is dropped.
"""
def __init__(self, nccl_comm: int, managed: bool):
self.nccl_comm = nccl_comm
self.managed = managed
self.destroy = _C.destroy_nccl_comm
def __del__(self):
if self.managed:
self.destroy(self.nccl_comm)
def get(self) -> int:
"""
Get the raw NCCL communicator.
Returns:
nccl_comm: the raw NCCL communicator.
"""
return self.nccl_comm
_storage = dict()
def get_nccl_comm_handle(group: dist.ProcessGroup, force_new_comm: bool = False) -> NCCLCommHandle:
"""
Get or create an NCCL communicator handle for the given process group.
Results are cached, so subsequent calls with the same group return the same handle.
Arguments:
group: the communication group.
force_new_comm: if set, never reuse PyTorch's communicator and never hit the cache; always
create a fresh DeepEP-managed comm.
Returns:
handle: the NCCL communicator handle.
"""
# Check cache hit
global _storage
if not force_new_comm and group in _storage:
return _storage[group]
# New PyTorch has such API
backend = group._get_backend(torch.device('cuda'))
if not force_new_comm and hasattr(backend, '_comm_ptr') and int(os.getenv('EP_REUSE_NCCL_COMM', '1')):
_storage[group] = NCCLCommHandle(backend._comm_ptr(), False)
return _storage[group]
# For old PyTorch, we have to recreate a NCCL comm
nccl_unique_ids = [None, ] * group.size()
dist.all_gather_object(nccl_unique_ids, _C.get_local_nccl_unique_id(), group)
root_unique_id = nccl_unique_ids[0]
# Create a new communicator
key = time.time_ns() if force_new_comm else group
_storage[key] = NCCLCommHandle(
_C.create_nccl_comm(root_unique_id, group.size(), group.rank()), True)
return _storage[key]
def destroy_all_managed_nccl_comm() -> None:
"""
Destroy all cached NCCL communicator handles and clear the cache.
"""
_storage.clear()
deep_ep/utils/envs.py
+268
View File
@@ -0,0 +1,268 @@
import functools
import inspect
import os
import random
import re
import subprocess
import torch
import torch.distributed as dist
from typing import Tuple
# noinspection PyUnresolvedReferences
import deep_ep._C as _C
from .comm import get_nccl_comm_handle
_local_rank = None
_local_seed = 0
_global_seed = 0
# Default NIC name for RDMA operations, configurable via environment variable
_DEFAULT_NIC_NAME = os.getenv('EP_NIC_NAME', 'mlx5_0')
def init_seed(global_seed: int) -> None:
"""
Initialize the random seed for reproducibility. The local seed is derived from the global seed plus rank.
Arguments:
global_seed: the global random seed.
"""
global _local_seed, _global_seed
_local_seed = global_seed + dist.get_rank()
_global_seed = global_seed
torch.manual_seed(_local_seed)
random.seed(_local_seed)
def get_local_seed() -> int:
"""
Get the local random seed.
Returns:
seed: the local random seed.
"""
return _local_seed
def get_global_seed() -> int:
"""
Get the global random seed.
Returns:
seed: the global random seed.
"""
return _global_seed
def dist_print(s: str = '', once_in_node: bool = False) -> None:
"""
Print a message from all ranks, or only from rank 0 of each node, followed by a barrier.
Arguments:
s: the message to print.
once_in_node: if `True`, only the first local rank in each node prints.
"""
global _local_rank
assert _local_rank is not None
if not once_in_node or _local_rank == 0:
print(s, flush=True)
dist.barrier()
def init_dist(local_rank: int, num_local_ranks: int, seed: int = 0) -> Tuple[int, int, dist.ProcessGroup]:
"""
Initialize the distributed environment with NCCL backend.
Arguments:
local_rank: the local rank index.
num_local_ranks: the number of local ranks.
seed: the global random seed.
Returns:
rank: the global rank index.
world_size: the total number of ranks.
group: the communication group.
"""
# NOTES: you may rewrite this function with your own cluster settings
ip = os.getenv('MASTER_ADDR', '127.0.0.1')
port = int(os.getenv('MASTER_PORT', '8361'))
num_nodes = int(os.getenv('WORLD_SIZE', 1))
node_rank = int(os.getenv('RANK', 0))
# Set local rank
global _local_rank
_local_rank = local_rank
sig = inspect.signature(dist.init_process_group)
params = {
'backend': 'nccl',
'init_method': f'tcp://{ip}:{port}',
'world_size': num_nodes * num_local_ranks,
'rank': node_rank * num_local_ranks + local_rank,
}
if 'device_id' in sig.parameters:
# noinspection PyTypeChecker
params['device_id'] = torch.device(f'cuda:{local_rank}')
dist.init_process_group(**params)
torch.set_default_dtype(torch.bfloat16)
torch.set_default_device('cuda')
torch.cuda.set_device(local_rank)
init_seed(seed)
return dist.get_rank(), dist.get_world_size(), dist.new_group(list(range(num_local_ranks * num_nodes)))
def get_physical_domain_size(group: dist.ProcessGroup) -> Tuple[int, int]:
"""
Get the physical domain sizes (RDMA ranks and NVLink ranks).
Arguments:
group: the communication group.
Returns:
num_rdma_ranks: the number of physical RDMA ranks.
num_nvlink_ranks: the number of physical NVLink ranks.
"""
return _C.get_physical_domain_size(get_nccl_comm_handle(group).get())
def get_logical_domain_size(group: dist.ProcessGroup, allow_hybrid_mode: bool = True) -> Tuple[int, int]:
"""
Get the logical domain sizes (scaleout ranks and scaleup ranks).
Arguments:
group: the communication group.
allow_hybrid_mode: whether to enable hybrid mode.
Returns:
num_scaleout_ranks: the number of logical scaleout ranks.
num_scaleup_ranks: the number of logical scaleup ranks.
"""
return _C.get_logical_domain_size(get_nccl_comm_handle(group).get(), allow_hybrid_mode)
def check_nvlink_connections(group: dist.ProcessGroup) -> None:
"""
Check NVLink connection between every pair of GPUs.
Arguments:
group: the communication group.
"""
# Check NVLink connection
# NOTES: some A100 PCIE GPUs only have pairwise NVLink connection, so that we can only use EP2
# TODO: check all cases, all local-node GPUs in the group should be connected via NVLink
if 'PCIE' in torch.cuda.get_device_name():
assert group.size() <= 2, 'PCIe GPUs only have pairwise NVLink connections'
# noinspection PyUnresolvedReferences
import pynvml
pynvml.nvmlInit()
# noinspection PyTypeChecker
devices = os.environ.get('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7').strip(',').split(',')
physical_device_idx = int(devices[torch.cuda.current_device()])
physical_device_indices = [0, ] * group.size()
dist.all_gather_object(physical_device_indices, physical_device_idx, group)
# Check whether they are all connected via NVLink
# Reference: https://github.com/vllm-project/vllm/blob/b8e809a057765c574726a6077fd124db5077ce1f/vllm/platforms/cuda.py#L438
handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_indices]
for i, handle in enumerate(handles):
for j, peer_handle in enumerate(handles):
if i >= j:
continue
status = pynvml.nvmlDeviceGetP2PStatus(handle, peer_handle, pynvml.NVML_P2P_CAPS_INDEX_NVLINK)
assert status == pynvml.NVML_P2P_STATUS_OK, \
f'GPU {physical_device_indices[i]} and GPU {physical_device_indices[j]} are not connected via NVLink'
# Close NVML
pynvml.nvmlShutdown()
def check_torch_deterministic() -> None:
"""
Ensure PyTorch deterministic algorithms and fill_uninitialized_memory are not both enabled.
When both are on, `torch.empty()` calls an initialization kernel that may overlap with communication streams,
causing errors.
"""
assert not (torch.are_deterministic_algorithms_enabled() and torch.utils.deterministic.fill_uninitialized_memory)
@functools.lru_cache()
def get_nvlink_gbs(factor: float = 0.9) -> float:
"""
Get the total NVLink bandwidth in GB/s, cached.
Arguments:
factor: the bandwidth efficiency factor.
Returns:
gbs: the total NVLink bandwidth in GB/s (0 if detection fails).
"""
# noinspection PyBroadException
try:
result = subprocess.run(['nvidia-smi', 'nvlink', '-s'],
capture_output=True, text=True, check=True)
output = result.stdout
pattern = r'GPU \d+:.*?(?=^GPU \d+:|^$)'
match = re.search(pattern, output, re.MULTILINE | re.DOTALL)
assert match
gpu_block = match.group(0)
link_pattern = r'Link \d+:\s*([\d\.]+) GB/s'
link_matches = re.findall(link_pattern, gpu_block)
assert link_matches
return sum(float(bw) for bw in link_matches) * factor
except Exception as e:
print(f'Failed to get NVLink connection speed: {e}')
return 0
@functools.lru_cache()
def check_fast_rdma_atomic_support(nic_name: str = _DEFAULT_NIC_NAME) -> bool:
"""
Check whether the NIC supports fast RDMA atomic operations (MT4131 or newer).
Arguments:
nic_name: the NIC device name.
Returns:
supported: `True` if fast RDMA atomics are supported.
"""
# noinspection PyBroadException
try:
result = subprocess.run(['ibstat'], capture_output=True, text=True, check=True)
output = result.stdout
pattern = rf"CA '{nic_name}'.*?CA type:\s*(\S+)"
match = re.search(pattern, output, re.DOTALL)
assert match
return match.group(1) == 'MT4131'
except Exception:
return False
@functools.lru_cache()
def get_rdma_gbs(nic_name: str = _DEFAULT_NIC_NAME) -> float:
"""
Get the RDMA bandwidth in GB/s, cached.
Arguments:
nic_name: the NIC device name.
Returns:
gbs: the RDMA bandwidth in GB/s (0 if detection fails).
"""
# noinspection PyBroadException
try:
result = subprocess.run(['ibstat'], capture_output=True, text=True, check=True)
output = result.stdout
pattern = rf"CA '{nic_name}'.*?Port \d+:\s*.*?Rate:\s*(\d+)"
match = re.search(pattern, output, re.DOTALL)
assert match
rate = int(match.group(1))
return rate / 8
except Exception as e:
print(f'Failed to get RDMA connection speed: {e}')
return 0
deep_ep/utils/event.py
+96
View File
@@ -0,0 +1,96 @@
import torch
from typing import Any, Optional, Tuple, Callable
# noinspection PyUnresolvedReferences
from deep_ep._C import EventHandle
class EventOverlap:
"""
A wrapper class to manage CUDA events, also for better overlapping convenience.
Attributes:
event: the CUDA event captured.
extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph.
"""
def __init__(self, event: Optional[EventHandle] = None, extra_tensors: Optional[Tuple[torch.Tensor]] = None) -> None:
"""
Initialize the class.
Arguments:
event: the CUDA event captured.
extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph.
"""
self.event = event
# NOTES: we use extra tensors to achieve stream recording, otherwise,
# stream recording will be incompatible with CUDA graph.
# TODO: `extra_tensors` is not longer useful for EPv2, as objects are stored in `self.event`
self.extra_tensors = extra_tensors
# A wrapper for `with event_overlap(release_handle=True)`
self._release_handle_by_call = False
# A hook that will be triggered after `current_stream_wait()`
# Useful for deterministic dispatch, which requires a sort (on the current stream) after `self.current_stream_wait()` is invoked
self.hook_after_wait: Optional[Callable] = None
def current_stream_wait(self, release_handle: bool = False) -> None:
"""
The current stream `torch.cuda.current_stream()` waits for the event to be finished.
"""
assert self.event is not None
self.event.current_stream_wait()
if self.hook_after_wait is not None:
self.hook_after_wait()
self.hook_after_wait = None
# In `self.event`, we also have some V2 APIs storing tensors to record in it,
# So, after waiting the current stream, those tensors can be released by deleting `self.event`
# However, you better do it by yourself (to be compatible with multi-stream waits)
if release_handle:
self.event = None
def register_hook_after_wait(self, hook_after_wait: Callable) -> None:
"""
Register a hook, which will be invoked after `self.current_stream_wait()`
"""
assert self.hook_after_wait is None, "A hook is already registered on this `EventOverlap`"
self.hook_after_wait = hook_after_wait
def __call__(self, release_handle: bool = False) -> "EventOverlap":
"""
Configures the 'release_handle' behavior for the upcoming context manager usage.
Usage:
with event_overlap(release_handle=True):
...
Returns `self` to ensure no new wrapper object is created, keeping the reference count of the underlying event unchanged/managed solely by this instance.
"""
self._release_handle_by_call = release_handle
return self
def __enter__(self) -> Any:
"""
Utility for overlapping and Python `with` syntax.
You can overlap the kernels on the current stream with the following example:
```python
event_overlap = event_after_all_to_all_kernels()
with event_overlap:
do_something_on_current_stream()
# After exiting the `with` scope, the current stream with wait the event to be finished.
```
"""
return self
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
"""
Utility for overlapping and Python `with` syntax.
Please follow the example in the `__enter__` function.
"""
if self.event is not None:
self.current_stream_wait(release_handle=self._release_handle_by_call)
self._release_handle_by_call = False
deep_ep/utils/find_pkgs.py
+82
View File
@@ -0,0 +1,82 @@
import functools
import os
import sys
from importlib.metadata import distributions
from typing import Optional
def find_pkg_root(name: str, lib_name: Optional[str] = None, optional: bool = False):
"""
Find the root directory of an installed NVIDIA package by inspecting Python package metadata.
Checks environment variables `EP_{NAME}_ROOT_DIR` and `{NAME}_DIR` first.
Arguments:
name: the package name (e.g., `'nccl'`, `'nvshmem'`).
lib_name: the library filename to search for within the package files.
optional: if ``False``, raises an assertion error when the package is not found.
Returns:
root: the package root directory, or `None` if not found and optional.
"""
upper = name.upper()
for env_name in (f'EP_{upper}_ROOT_DIR', f'{upper}_DIR'):
if env_name in os.environ:
return os.environ[env_name]
path_priority = {p: i for i, p in enumerate(sys.path)}
best, best_priority = None, len(sys.path)
for dist in distributions():
dist_name = (dist.metadata['Name'] or '').lower()
if f'nvidia-{name}' not in dist_name and f'nvidia_{name}' not in dist_name:
continue
dist_site = str(dist._path.parent)
priority = path_priority.get(dist_site, len(sys.path))
if priority > best_priority:
continue
if lib_name is not None:
for f in (dist.files or []):
if lib_name in str(f):
lib_dir = os.path.dirname(str(f.locate()))
root = os.path.dirname(lib_dir) if os.path.basename(lib_dir) == 'lib' else lib_dir
best, best_priority = root, priority
break
else:
pkg_dir = os.path.join(dist_site, 'nvidia', name)
if os.path.isdir(pkg_dir):
best, best_priority = pkg_dir, priority
# Raise error if not optional
if not optional:
assert best is not None, f'Cannot find package: {name}'
return best
@functools.lru_cache()
def find_nccl_root(optional: bool = False):
"""
Find the NCCL installation root directory, cached.
Arguments:
optional: if `False`, raises an assertion error when NCCL is not found.
Returns:
root: the NCCL root directory.
"""
return find_pkg_root('nccl', lib_name='libnccl.so', optional=optional)
@functools.lru_cache()
def find_nvshmem_root(optional: bool = False):
"""
Find the NVSHMEM installation root directory, cached.
Arguments:
optional: if `False`, raises an assertion error when NVSHMEM is not found.
Returns:
root: the NVSHMEM root directory.
"""
return find_pkg_root('nvshmem', optional=optional)
deep_ep/utils/gate.py
+180
View File
@@ -0,0 +1,180 @@
import torch
def generate_topk_idx(rank_count: torch.Tensor, num_tokens: int, num_experts: int, num_ranks: int, num_topk: int) -> torch.Tensor:
"""
Map rank count to expert indices
"""
assert torch.equal(torch.sum(rank_count, dim=1), torch.ones(num_tokens, dtype=torch.int, device='cuda') * num_topk)
assert (num_tokens, num_ranks) == rank_count.shape
num_experts_per_rank = num_experts // num_ranks
# Generate base value
base_vals = torch.arange(num_experts, device='cuda').view(1, num_ranks, num_experts_per_rank).expand(num_tokens, num_ranks, num_experts_per_rank)
# Randomize the ordering within each row
rand_vals = torch.rand(num_tokens, num_ranks, num_experts_per_rank, device='cuda')
perm_indices = torch.argsort(rand_vals, dim=-1)
permuted = torch.gather(base_vals, 2, perm_indices)
# Create the mask
k_idx = torch.arange(num_experts_per_rank, device='cuda').view(1, 1, num_experts_per_rank).expand(num_tokens, num_ranks, num_experts_per_rank)
rank_count_expanded = rank_count.unsqueeze(2).expand(num_tokens, num_ranks, num_experts_per_rank)
mask = k_idx < rank_count_expanded
# Get the final indices by masking and reshaping
selected = permuted[mask] # (num_tokens * num_topk,)
topk_idx = selected.view(num_tokens, num_topk)
return topk_idx
def generate_rank_count(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float) -> torch.Tensor:
"""
Generate rank count tensor for a given number of tokens, experts, ranks, and top-k.
This function generates a tensor of shape `(num_tokens, num_ranks)` where each element `[i, j]` represents
the number of topk experts that token `i` have on rank `j`. The distribution is such that
one special rank gets `ratio` times more traffic than the others.
"""
num_experts_per_rank = num_experts // num_ranks
num_normal_ranks = num_ranks - 1
assert ratio >= 1.0, 'ratio must be no less than 1.0'
# Generate rank count of each token from random distribution
random_scores = torch.rand(num_tokens, num_experts, device='cuda')
topk_weights_, topk_indices = torch.topk(random_scores, num_topk, dim=1, largest=True, sorted=False)
topk_indices //= num_experts_per_rank
sorted_topk_indices = torch.sort(topk_indices, dim=1)[0]
topk_indices_diff_mask = sorted_topk_indices[:, 1:] != sorted_topk_indices[:, :-1]
a = topk_indices_diff_mask.sum(dim=1) + 1
# Upper bound for this generating algorithm
upper_bound_per_token = int(num_normal_ranks / ratio) + 1
# Clamp the value in range [1, upper_bound_per_token] for each token
a = torch.clamp(a, None, upper_bound_per_token)
# Consider the special rank
sum_a = torch.sum(a).item()
normal_token_count = int(sum_a / (num_normal_ranks + ratio))
special_token_count = sum_a - normal_token_count * num_normal_ranks
special_token_count = min(special_token_count, int(normal_token_count * ratio) + 1)
# Tokens that the special rank must be in topk
must_mask = (a == num_ranks)
must_count = int(must_mask.sum().item())
special_token_count = max(must_count, special_token_count)
assert must_count <= special_token_count, 'Too many tokens with full rank assignment'
# Tokens that the special rank can optionally be in topk
optional_token_indices = torch.where(must_mask == 0)[0]
optional_token_indices = optional_token_indices[torch.randperm(num_tokens - must_count, device='cuda')][:special_token_count - must_count]
must_token_indices = torch.where(must_mask != 0)[0]
special_token_row_index = torch.cat(([must_token_indices, optional_token_indices]))
# Generate permutations for normal ranks
rank_perm = (torch.randperm(num_normal_ranks, device='cuda') + 1).repeat(num_tokens * num_topk // num_normal_ranks + 1)
# Compute cumulative sum of a to get starting indices in b for each row
a_cumsum = torch.cumsum(torch.cat((torch.tensor([0], device='cuda'), a)), dim=0)
row_starts = a_cumsum[:-1] # Starting indices for each row in b, shape (n,)
# Insert special rank index into the permutation for special tokens
rank_perm_with_special_rank = torch.zeros(num_tokens * num_topk, dtype=torch.long, device='cuda') # (n * k,)
special_token_mask = torch.zeros(num_tokens * num_topk, dtype=torch.bool, device='cuda')
special_token_flattened_row_index = row_starts[special_token_row_index]
special_token_mask[special_token_flattened_row_index] = 1
all_indices = torch.arange(num_tokens * num_topk, device='cuda')
non_special_indices = all_indices[special_token_mask != True]
rank_perm_with_special_rank[non_special_indices] = rank_perm[:len(non_special_indices)]
# Create column index grids
col_idx = torch.arange(num_topk, device='cuda').view(1, num_topk) # (1, num_topk)
# Compute modulo indices: col_idx % a[i] for each row
# torch.max is used to avoid zeros in case a[i] = 0 (which happens when the only topk rank is the special rank)
mod_idx = col_idx % a.view(num_tokens, 1) # (n, num_topk)
# Compute indices in b: row_start + (col % a[i])
b_indices = row_starts.view(num_tokens, 1) + mod_idx # (n, k)
# Gather values from b using computed indices
result = rank_perm_with_special_rank[b_indices]
# Shuffle rows randomly to avoid any pattern
shuffle_indices = torch.randperm(num_tokens, device='cuda')
result = result[shuffle_indices] # Shuffle rows
# Create rank count tensor
rank_count = torch.zeros((num_tokens, num_ranks), dtype=torch.int32, device='cuda')
rank_count.scatter_add_(dim=1, index=result, src=torch.ones_like(result, dtype=torch.int32))
return rank_count
def get_precise_unbalanced_scores(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float):
"""
Generate precise unbalanced scores for testing.
Note that this function generates a distribution with precise unbalanced distribution,
which **differs from real distribution**.
"""
# Generate num topk experts for each rank
rank_count = generate_rank_count(num_tokens, num_experts, num_ranks, num_topk, ratio)
# Generate scores in a low distribution
threshold = 0.9
scores = torch.empty((num_tokens, num_experts), dtype=torch.float32, device='cuda')
scores.uniform_(to=threshold)
# Generate topk indices and change their scores to a high distribution
topk_idx = generate_topk_idx(rank_count, num_tokens, num_experts, num_ranks, num_topk)
topk_scores = torch.empty((num_tokens, num_topk), dtype=torch.float32, device='cuda')
topk_scores.uniform_(threshold + 1e-6, 1.0)
row_idx = torch.arange(num_tokens).unsqueeze(1).expand(num_tokens, num_topk)
scores[row_idx, topk_idx] = topk_scores
return scores
def get_scores_by_factor(num_tokens: int, num_experts: int, num_ranks: int, factor: float) -> torch.Tensor:
num_experts_per_rank = num_experts // num_ranks
scores = torch.empty((num_tokens, num_experts), dtype=torch.float32, device='cuda')
scores[:, :num_experts_per_rank].uniform_(to=factor)
scores[:, num_experts_per_rank:].uniform_(to=1)
return scores
def map_unbalanced_ratio_to_factor(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float) -> float:
num_iterations = 20
factor_l, factor_r = 1.0, 100.0
num_experts_per_rank = num_experts // num_ranks
for _i in range(num_iterations):
factor_mid = (factor_l + factor_r) / 2
scores = get_scores_by_factor(num_tokens, num_experts, num_ranks, factor_mid)
_, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)
rank_idx = topk_idx // num_experts_per_rank
one_hot = torch.nn.functional.one_hot(rank_idx, num_ranks)
counts = one_hot.any(dim=1).to(torch.float).sum(dim=0)
if counts[0].item() > counts[1:].mean().item() * ratio:
factor_r = factor_mid
else:
factor_l = factor_mid
return factor_l
def get_random_unbalanced_scores(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float):
"""Generate unbalanced scores with a given ratio.
"""
factor = 1.0
if ratio != 1.0:
factor = map_unbalanced_ratio_to_factor(num_tokens, num_experts, num_ranks, num_topk, ratio)
return get_scores_by_factor(num_tokens, num_experts, num_ranks, factor)
def get_unbalanced_scores(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float, precise: bool):
if precise:
return get_precise_unbalanced_scores(num_tokens, num_experts, num_ranks, num_topk, ratio)
else:
return get_random_unbalanced_scores(num_tokens, num_experts, num_ranks, num_topk, ratio)
deep_ep/utils/math.py
+103
View File
@@ -0,0 +1,103 @@
import torch
from typing import Tuple
def calc_diff(x: torch.Tensor, y: torch.Tensor) -> float:
x, y = x.double() + 1, y.double() + 1
denominator = (x * x + y * y).sum()
sim = 2 * (x * y).sum() / denominator
return (1 - sim).item()
def safe_div(a, b) -> float:
try:
return a / b
except ZeroDivisionError as e:
if a == 0:
return 0
else:
raise
def ceil_div(x: int, y: int) -> int:
return (x + y - 1) // y
def align(x: int, y: int) -> int:
return ceil_div(x, y) * y
@torch.compile(dynamic=True)
def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2
m, n = x.shape
aligned_n = align(n, 128)
x_padded = torch.nn.functional.pad(x, (0, aligned_n - n), mode='constant', value=0)
x_padded_view = x_padded.view(m, -1, 128)
x_amax = x_padded_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
return (x_padded_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view(
m, aligned_n)[:, :n].contiguous(), (x_amax / 448.0).view(m, -1)
@torch.compile(dynamic=True)
def per_token_cast_back(x_fp8: torch.Tensor, x_scales: torch.Tensor) -> torch.Tensor:
if x_fp8.numel() == 0:
return x_fp8.to(torch.bfloat16)
assert x_fp8.dim() == 2
m, n = x_fp8.shape
aligned_n = align(n, 128)
x_fp8_padded = torch.nn.functional.pad(x_fp8, (0, aligned_n - n), mode='constant', value=0)
if x_scales.dtype == torch.int:
x_scales = x_scales.view(dtype=torch.uint8).to(torch.int) << 23
x_scales = x_scales.view(dtype=torch.float)
x_fp32_padded = x_fp8_padded.to(torch.float32).view(x_fp8.shape[0], -1, 128)
x_scales = x_scales.view(x_fp8.shape[0], -1, 1)
return (x_fp32_padded * x_scales).view(x_fp8_padded.shape).to(torch.bfloat16)[:, :n].contiguous()
def inplace_unique(x: torch.Tensor, num_slots: int) -> None:
assert x.dim() == 2
mask = x < 0
x_padded = x.masked_fill(mask, num_slots)
bin_count = torch.zeros((x.size(0), num_slots + 1), dtype=x.dtype, device=x.device)
bin_count.scatter_add_(1, x_padded, torch.ones_like(x_padded))
bin_count = bin_count[:, :num_slots]
sorted_bin_count, sorted_bin_idx = torch.sort(bin_count, dim=-1, descending=True)
sorted_bin_idx.masked_fill_(sorted_bin_count == 0, -1)
sorted_bin_idx = torch.sort(sorted_bin_idx, descending=True, dim=-1).values
x[:, :].fill_(-1)
valid_len = min(num_slots, x.size(1))
x[:, :valid_len] = sorted_bin_idx[:, :valid_len]
def create_grouped_scores(scores: torch.Tensor, group_idx: torch.Tensor, num_groups: int) -> torch.Tensor:
num_tokens, num_experts = scores.shape
scores = scores.view(num_tokens, num_groups, -1)
mask = torch.zeros((num_tokens, num_groups), dtype=torch.bool, device=scores.device)
mask = mask.scatter_(1, group_idx, True).unsqueeze(-1).expand_as(scores)
return (scores * mask).view(num_tokens, num_experts)
def hash_tensor(t: torch.Tensor) -> int:
return t.view(torch.int).sum().item()
def hash_tensors(*tensors) -> int:
value = 0
for t in tensors:
if isinstance(t, (tuple, list)):
value ^= hash_tensors(*t)
elif t is not None and isinstance(t, torch.Tensor):
value ^= hash_tensor(t)
return value
def count_bytes(*tensors) -> int:
total = 0
for t in tensors:
if isinstance(t, (tuple, list)):
total += count_bytes(*t)
elif t is not None:
total += t.numel() * t.element_size()
return total
deep_ep/utils/refs.py
+243
View File
@@ -0,0 +1,243 @@
import math
import torch
import torch.distributed as dist
from typing import Optional, Tuple, Union
from .envs import get_global_seed
from .math import ceil_div
def dispatch(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
topk_idx: torch.Tensor, topk_weights: Optional[torch.Tensor],
num_max_tokens_per_rank: int, num_experts: int):
"""
The reference implementation of dispatching tokens to experts across multiple ranks.
Not expanded. Sorted by rank and then by token within each rank (i.e. sorted by `src_token_global_idx`).
Arguments:
- `x`: Input tokens, `[num_tokens, hidden]` or (`[num_tokens, hidden], [num_tokens, hidden_sf]`)
- `topk_idx`: Top-k expert indices for each token, `[num_tokens, num_topk]`
- `topk_weights`: Top-k weights for each token, can be None, `[num_tokens, num_topk]`
- `num_max_tokens_per_rank`: Maximum number of tokens per rank, must >= actual number of tokens per rank and aligned with DeepEP's `num_max_tokens_per_rank` since we're going to calculate `src_token_global_idx = src_rank_idx*num_max_tokens_per_rank + src_token_local_idx`
- `num_experts`: Total number of experts across all ranks
Returns:
- `recv_x`, `recv_topk_idx`, and `recv_topk_weights`: Received tokens, top-k indices, and top-k weights after dispatching. Out of range `recv_topk_idx` (i.e. that expert is not on the current rank) are set to -1.
- `recv_src_token_idx`: Received `src_token_global_idx` for each received token
- `num_recv_tokens_per_rank`: Number of received tokens from each rank, `[num_ranks]`
"""
# TODO: support forwarding
# TODO: make top-k weight fully optional
rank_idx = dist.get_rank()
num_ranks = dist.get_world_size()
assert num_experts % num_ranks == 0
num_experts_per_rank = num_experts // num_ranks
# Unpack SF
use_fp8 = isinstance(x, tuple)
x, sf = x if use_fp8 else (x, None)
# TODO: use SF bytes instead of hardcoded recipe
num_tokens, hidden = x.size()
num_tokens_, num_topk = topk_idx.size()
assert num_tokens == num_tokens_
if sf is not None:
num_tokens__, hidden_sf = sf.size()
assert num_tokens == num_tokens__
assert hidden_sf == ceil_div(hidden, 128)
if topk_weights is not None:
num_tokens__, num_topk_ = topk_weights.size()
assert num_tokens == num_tokens__
assert num_topk == num_topk_
# Prepare per-peer send buffers
send_x_list = []
send_sf_list = []
send_topk_idx_list = []
send_topk_weights_list = []
send_src_token_idx_list = []
num_send_tokens_per_rank = torch.zeros((num_ranks, ), dtype=torch.int, device=x.device)
for dst_rank_idx in range(num_ranks):
expert_start_idx = dst_rank_idx * num_experts_per_rank
expert_end_idx = expert_start_idx + num_experts_per_rank
# Get the indices of tokens
mask_to_send = ((expert_start_idx <= topk_idx) & (topk_idx < expert_end_idx)).any(dim=1)
indices_to_send = mask_to_send.nonzero(as_tuple=True)[0]
num_send_tokens_per_rank[dst_rank_idx] = indices_to_send.numel()
# Select the data for tokens
x_to_send = x[indices_to_send]
sf_to_send = sf[indices_to_send] if use_fp8 else None
topk_idx_to_send = topk_idx[indices_to_send]
topk_weights_to_send = topk_weights[indices_to_send]
masked_topk_idx = torch.where((expert_start_idx <= topk_idx_to_send) & (topk_idx_to_send < expert_end_idx),
topk_idx_to_send, torch.full_like(topk_idx_to_send, -1))
send_x_list.append(x_to_send)
send_sf_list.append(sf_to_send)
send_topk_idx_list.append(masked_topk_idx)
send_topk_weights_list.append(topk_weights_to_send)
send_src_token_idx_list.append(indices_to_send)
send_x = torch.cat(send_x_list, dim=0)
send_sf = torch.cat(send_sf_list, dim=0) if use_fp8 else None
send_topk_idx = torch.cat(send_topk_idx_list, dim=0)
send_topk_weights = torch.cat(send_topk_weights_list, dim=0)
send_src_token_idx = torch.cat(send_src_token_idx_list, dim=0).to(torch.int)
send_src_token_idx += rank_idx * num_max_tokens_per_rank
# Exchange size
num_recv_tokens_per_rank = torch.empty((num_ranks, ), dtype=torch.int, device=x.device)
dist.all_to_all_single(num_recv_tokens_per_rank, num_send_tokens_per_rank)
num_recv_tokens = int(num_recv_tokens_per_rank.sum().item())
# Exchange main data
num_send_tokens_per_rank = num_send_tokens_per_rank.tolist()
num_recv_tokens_per_rank = num_recv_tokens_per_rank.tolist()
recv_x = torch.empty((num_recv_tokens, hidden), dtype=x.dtype, device=x.device)
recv_sf = torch.empty((num_recv_tokens, hidden_sf), dtype=sf.dtype, device=x.device) if use_fp8 else None
recv_topk_idx = torch.empty((num_recv_tokens, num_topk), dtype=topk_idx.dtype, device=x.device)
recv_topk_weights = torch.empty((num_recv_tokens, num_topk), dtype=topk_weights.dtype, device=x.device)
recv_src_token_idx = torch.empty((num_recv_tokens, ), dtype=torch.int, device=x.device)
dist.all_to_all_single(recv_x, send_x, num_recv_tokens_per_rank, num_send_tokens_per_rank)
if use_fp8:
dist.all_to_all_single(recv_sf, send_sf, num_recv_tokens_per_rank, num_send_tokens_per_rank)
dist.all_to_all_single(recv_topk_idx, send_topk_idx, num_recv_tokens_per_rank, num_send_tokens_per_rank)
dist.all_to_all_single(recv_topk_weights, send_topk_weights, num_recv_tokens_per_rank, num_send_tokens_per_rank)
dist.all_to_all_single(recv_src_token_idx, send_src_token_idx, num_recv_tokens_per_rank, num_send_tokens_per_rank)
# Mask top-k indices
expert_start_idx = rank_idx * num_experts_per_rank
expert_end_idx = expert_start_idx + num_experts_per_rank
mask = (expert_start_idx <= recv_topk_idx) & (recv_topk_idx < expert_end_idx)
recv_topk_idx = recv_topk_idx - expert_start_idx
recv_topk_idx.masked_fill_(~mask, -1)
# Pack SF
recv_x = (recv_x, recv_sf) if use_fp8 else recv_x
return (recv_x, recv_topk_idx, recv_topk_weights,
recv_src_token_idx, torch.tensor(num_recv_tokens_per_rank, dtype=torch.int))
def generate_pre_combine_data(src_token_global_idx: torch.Tensor,
num_max_tokens_per_rank: int, num_topk: int, hidden: int) -> torch.Tensor:
"""
Generate data needed for combine from `src_token_global_idx`.
Recall that `src_token_global_idx = src_rank_idx * num_max_tokens_per_rank + src_token_local_idx`.
The generated data (denoted as `y`) of the i-th token has a shape of [num_topk, hidden], with
`y[j, k] = sin((token_seeds * P % max_seed + 1) / max_seed * (k + 1) + sin(seed))`
where `P=131071` is a large prime, `token_seeds` is calculated via `token_seeds = src_token_global_idx[i] * num_topk + j`,
and `max_seed = num_ranks * num_max_tokens_per_rank * num_topk`.
Arguments:
- `src_token_global_idx`: Source token global indices, `[num_tokens]`
Returns:
- Generated data, `[num_tokens, num_topk, hidden]`
"""
num_ranks = dist.get_world_size()
token_seeds = (src_token_global_idx.unsqueeze(1) * num_topk +
torch.arange(num_topk, device=src_token_global_idx.device).unsqueeze(0)) # [num_tokens, num_topk]
max_seed = num_ranks * num_max_tokens_per_rank * num_topk
result = torch.sin(
(((token_seeds * 131071 % max_seed).float() + 1) / max_seed).unsqueeze(-1) *
torch.arange(1, hidden + 1, device=src_token_global_idx.device, dtype=torch.float32).broadcast_to(1, 1, hidden) +
math.sin(float(get_global_seed()))
)
return result.to(torch.bfloat16)
def ordered_accumulate(data: torch.Tensor, initial_value: Optional[torch.Tensor] = None) -> torch.Tensor:
"""
Accumulate `data` in order along the num_topk dimension.
Arguments:
- `data`: Data to be accumulated, `[num_tokens, num_topk, hidden]`
- `initial_value`: Initial value for accumulation, `[num_tokens, hidden]`
Returns:
- Result, `[num_tokens, hidden]`
"""
num_topk = data.shape[1]
if initial_value is None:
result = torch.zeros((data.shape[0], data.shape[2]), dtype=torch.float32, device=data.device)
else:
result = initial_value.clone()
for i in range(num_topk):
result += data[:, i, :].float()
return result.to(data.dtype)
def combine(y: torch.Tensor, topk_idx: torch.Tensor,
num_scaleout_ranks: int, num_scaleup_ranks: int, num_experts: int,
bias: Optional[Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]],
reduce_in_local: bool, reduce_in_scaleup: bool) -> torch.Tensor:
"""
The reference implementation of (possibly multi-level reduction) combining tokens.
Arguments:
- `y`: Input tokens to be combined, `[num_tokens, num_topk, hidden]`
- `topk_idx`: `[num_tokens, num_topk]`
- `reduce_in_local` and `reduce_in_scaleup`: Whether to do reduction within rank and within scale-up group.
- `(True, True)` -> Hybrid combine
- `(True, False)` -> Non-hybrid combine
- `(False, False)` -> Equivalent to `allow_multiple_reduction` is False
Pay attention that `reduce_in_scaleup` = `True` or `False` is NOT equivalent even if `num_scaleout_ranks == 1` due to `bias` handling.
Returns:
- Combined result, `[num_tokens, hidden]`
"""
num_ranks = num_scaleout_ranks * num_scaleup_ranks
num_tokens, hidden = y.shape[0], y.shape[2]
num_topk = y.shape[1]
assert not (not reduce_in_local and reduce_in_scaleup), 'Invalid reduction configuration'
def grouped_reduce(data_to_reduce: torch.Tensor, group_id: torch.Tensor) -> torch.Tensor:
"""
Perform in-place grouped reduction on `data_to_reduce` according to `group_id`.
The summation within each group are performed in strict order along the `num_topk` dimension.
The result for each group is stored at the rightmost token of that group, and other tokens are set to zero.
Arguments:
- `data_to_reduce`: Data to be reduced, `[num_tokens, num_topk, hidden]`
- `group_id`: group IDs for each token, `[num_tokens, num_topk]`
"""
# Shuffle to make tokens with the same group_id contiguous
group_id, src_indices = torch.sort(group_id, dim=-1, stable=True)
# transformed_src_indices[i, j] = i * num_topk + src_indices[i, j]
transformed_src_indices = (
(src_indices + torch.arange(0, num_tokens, device=y.device).unsqueeze(-1) * num_topk).flatten())
data_to_reduce = data_to_reduce.view(-1, hidden)[transformed_src_indices].view(num_tokens, num_topk, hidden)
# Perform segmented reduce within each group
cur_accum_buf = torch.zeros((num_tokens, hidden), dtype=torch.float32, device=y.device)
for i in range(num_topk):
is_segment_break = torch.full((num_tokens, ), True, dtype=torch.bool, device=y.device) \
if i == num_topk - 1 else group_id[:, i] != group_id[:, i + 1]
cur_accum_buf += data_to_reduce[:, i, :].float()
# For one token, if `is_segment_break` is True,
# save the accumulated value and clear the buffer, otherwise, clear `data_to_reduce[:, i, :]`
segment_break_token_indices = torch.where(is_segment_break)[0]
data_to_reduce[segment_break_token_indices, i] = cur_accum_buf[segment_break_token_indices].to(data_to_reduce.dtype)
cur_accum_buf[segment_break_token_indices] = 0.0
non_segment_break_token_indices = torch.where(~is_segment_break)[0]
data_to_reduce[non_segment_break_token_indices, i] = 0.0
# Unshuffle
# noinspection PyShadowingNames
result = torch.empty_like(data_to_reduce)
result.view(-1, hidden)[transformed_src_indices] = data_to_reduce.view(-1, hidden)
return result.view(num_tokens, num_topk, hidden)
num_experts_per_rank = num_experts // num_ranks
if reduce_in_local:
y = grouped_reduce(y, topk_idx // num_experts_per_rank)
if reduce_in_scaleup:
y = grouped_reduce(y, topk_idx // (num_experts_per_rank * num_scaleup_ranks))
bias_sum = bias[0].float() + bias[1].float() if isinstance(bias, tuple) else bias.float() if bias is not None else None
result = ordered_accumulate(y, bias_sum)
return result
deep_ep/utils/semantic.py
+27
View File
@@ -0,0 +1,27 @@
from typing import Any, Optional
import weakref
import functools
def value_or(value: Optional[Any], default: Any) -> Any:
return default if value is None else value
def weak_lru(maxsize: Optional[int] = 128, typed: bool = False):
"""
LRU Cache decorator that keeps a weak reference to `self`
Useful for caching methods in classes that may cause memory leaks if `functools.lru_cache` is used directly.
From https://stackoverflow.com/a/68052994/16569836
"""
def wrapper(func):
@functools.lru_cache(maxsize, typed)
def _func(_self, *args, **kwargs):
return func(_self(), *args, **kwargs)
@functools.wraps(func)
def inner(self, *args, **kwargs):
return _func(weakref.ref(self), *args, **kwargs)
return inner
return wrapper
deep_ep/utils/testing.py
+219
View File
@@ -0,0 +1,219 @@
import json
import os
import sys
import torch
import numpy as np
import tempfile
import torch.distributed as dist
from pathlib import Path
from typing import Callable, Optional, Union
def flush_l2_cache(enabled: bool = True):
"""
Flush the GPU L2 cache by writing a large zero-initialized tensor.
Arguments:
enabled: if `False`, does nothing.
"""
l2_flush_cache_size = 256e6
if enabled:
torch.empty(int(l2_flush_cache_size // 4), dtype=torch.int, device='cuda').zero_()
def bench(fn, num_warmups: int = 50, num_tests: int = 50,
post_fn: Optional[Callable] = None, flush_l2: bool = True):
"""
Benchmark a function using CUDA events.
Arguments:
fn: the function to benchmark.
num_warmups: the number of warmup iterations.
num_tests: the number of measurement iterations.
post_fn: an optional function to call after each test iteration.
flush_l2: whether to flush the L2 cache before each iteration.
Returns:
avg: the average execution time in seconds.
min: the minimum execution time in seconds.
max: the maximum execution time in seconds.
"""
torch.cuda.synchronize()
# Warmup
for _ in range(num_warmups):
fn()
# Testing
start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)]
end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)]
for i in range(num_tests):
flush_l2_cache(flush_l2)
start_events[i].record()
fn()
end_events[i].record()
if post_fn is not None:
post_fn()
torch.cuda.synchronize()
times = np.array([s.elapsed_time(e) / 1e3 for s, e in zip(start_events, end_events)])[1:]
return np.average(times), np.min(times), np.max(times)
class empty_suppress:
def __enter__(self):
return self
def __exit__(self, *_):
pass
class suppress_stdout_stderr:
"""
Context manager to suppress stdout and stderr output.
"""
def __enter__(self):
self.outnull_file = open(os.devnull, 'w')
self.errnull_file = open(os.devnull, 'w')
self.old_stdout_fileno_undup = sys.stdout.fileno()
self.old_stderr_fileno_undup = sys.stderr.fileno()
self.old_stdout_fileno = os.dup(sys.stdout.fileno())
self.old_stderr_fileno = os.dup(sys.stderr.fileno())
self.old_stdout = sys.stdout
self.old_stderr = sys.stderr
os.dup2(self.outnull_file.fileno(), self.old_stdout_fileno_undup)
os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup)
sys.stdout = self.outnull_file
sys.stderr = self.errnull_file
return self
def __exit__(self, *_):
sys.stdout = self.old_stdout
sys.stderr = self.old_stderr
os.dup2(self.old_stdout_fileno, self.old_stdout_fileno_undup)
os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup)
os.close(self.old_stdout_fileno)
os.close(self.old_stderr_fileno)
self.outnull_file.close()
self.errnull_file.close()
def bench_kineto(fn,
kernel_names: Union[str, tuple],
num_tests: int = 30,
suppress_kineto_output: bool = False,
trace_path: Optional[str] = None,
flush_l2: bool = True,
barrier_comm_profiling: bool = False,
num_kernels_per_period: int = 1,
barrier: Optional[Callable] = None):
"""
Benchmark a function using the PyTorch profiler (kineto) to get per-kernel timing.
Arguments:
fn: the function to benchmark.
kernel_names: the CUDA kernel name(s) to profile.
num_tests: the number of test iterations.
suppress_kineto_output: whether to suppress profiler output.
trace_path: the path to save the Chrome trace (`None` to skip).
flush_l2: whether to flush the L2 cache before each iteration.
barrier_comm_profiling: whether to insert a barrier before each iteration to reduce
unbalanced CPU launch overhead.
num_kernels_per_period: the number of kernels launched per test period.
barrier: a custom barrier function to use instead of `dist.all_reduce`.
Returns:
durations: the average kernel duration(s) in seconds.
"""
assert isinstance(kernel_names, (str, tuple))
is_tuple = isinstance(kernel_names, tuple)
# Skip profiling
# Conflict with Nsight Systems, Nsight Compute and Compute Sanitizer
if int(os.environ.get('EP_USE_NVIDIA_TOOLS', 0)):
return (1, ) * len(kernel_names) if is_tuple else 1
# For some auto-tuning kernels with prints
fn()
torch.cuda.synchronize()
# Profile
suppress = suppress_stdout_stderr if suppress_kineto_output else empty_suppress
barrier_comm_profiling &= int(os.environ.get('EP_DISABLE_BARRIER_PROFILING', 0)) == 0
with suppress():
schedule = torch.profiler.schedule(wait=0, warmup=1, active=1, repeat=1)
profiler = torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA], schedule=schedule, acc_events=True)
dummy = torch.ones(1, dtype=torch.float, device='cuda')
with profiler:
for i in range(2):
for _ in range(num_tests):
# Flush L2 cache
flush_l2_cache(flush_l2)
# NOTES: use a large kernel and a barrier to eliminate the unbalanced CPU launch overhead
if barrier_comm_profiling:
torch.cuda._sleep(int(2e7)) # ~10ms
# Some network may have ring-based implement, so be careful to use `all_reduce`
if barrier is None:
dist.all_reduce(dummy)
else:
barrier()
fn()
torch.cuda.synchronize()
profiler.step()
# Parse the profiling table
prof_lines = profiler.key_averages().table(sort_by='cuda_time_total', max_name_column_width=100).split('\n')
kernel_names = (kernel_names, ) if isinstance(kernel_names, str) else kernel_names
assert all([isinstance(name, str) for name in kernel_names])
for name in kernel_names:
assert sum([name in line for line in prof_lines]) <= 1, f'Errors of the kernel {name} in the profiling table: {prof_lines}'
# Save chrome traces
if trace_path is not None:
profiler.export_chrome_trace(trace_path)
# Return average kernel durations
units = {'ms': 1e3, 'us': 1e6}
kernel_durations = []
for name in kernel_names:
total_time = 0
total_num = 0
for line in prof_lines:
if name in line:
time_str = line.split()[-2]
num_str = line.split()[-1]
for unit, scale in units.items():
if unit in time_str:
total_time += float(time_str.replace(unit, '')) / scale * int(num_str)
total_num += int(num_str)
break
kernel_durations.append(total_time / total_num if total_num > 0 else 0)
# Expand the kernels by periods
if num_kernels_per_period > 1:
with tempfile.NamedTemporaryFile(suffix='.json') as tmp:
profiler.export_chrome_trace(tmp.name)
profile_data = json.loads(Path(tmp.name).read_text())
for i, kernel_name in enumerate(kernel_names):
events = [event for event in profile_data['traceEvents'] if f'::{kernel_name}' in event['name']]
events = sorted(events, key=lambda event: event['ts'])
durations = [event['dur'] / 1e6 for event in events]
assert len(durations) % num_kernels_per_period == 0
num_kernel_patterns = len(durations) // num_kernels_per_period
kernel_durations[i] = [sum(durations[j::num_kernels_per_period]) / num_kernel_patterns for j in range(num_kernels_per_period)]
# Return execution durations
return kernel_durations if is_tuple else kernel_durations[0]
develop.sh
Executable
+21
View File
@@ -0,0 +1,21 @@
# Change current directory into project root
original_dir=$(pwd)
script_dir=$(realpath "$(dirname "$0")")
cd "$script_dir"
# Remove old dist file, build files, and build
rm -rf build dist
rm -rf *.egg-info
python setup.py build
# Find the .so file in build directory and create symlink in current directory
so_file=$(find build -name "*.so" -type f | head -n 1)
if [ -n "$so_file" ]; then
ln -sf "../$so_file" deep_ep/
else
echo "Error: No SO file found in build directory" >&2
exit 1
fi
# Open users' original directory
cd "$original_dir"
docs/legacy.md
+320
View File
@@ -0,0 +1,320 @@
# DeepEP V1 (Legacy)
> **Note:** This is the archived documentation for DeepEP V1 (NVSHMEM-based). For the latest V2 documentation, see the [main README](../README.md).
---
DeepEP (DeepEveryParallel) V1 is the original high-performance communication library for modern machine learning, focused on expert parallelism (EP). It provides high-throughput and low-latency all-to-all GPU kernels, which are also known as MoE dispatch and combine. The library also supports low-precision operations, including FP8.
To align with the group-limited gating algorithm proposed in the [DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3) paper, DeepEP V1 offers a set of kernels optimized for asymmetric-domain bandwidth forwarding, such as forwarding data from NVLink domain to RDMA domain. These kernels deliver high throughput, making them suitable for both training and inference prefilling tasks. Additionally, they support SM (Streaming Multiprocessors) number control.
For latency-sensitive inference decoding, DeepEP V1 includes a set of low-latency kernels with pure RDMA to minimize delays. The library also introduces a hook-based communication-computation overlapping method that does not occupy any SM resource.
Notice: the implementation in this library may have some slight differences from the [DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3) paper.
## Performance
### Normal kernels with NVLink and RDMA forwarding
We test normal kernels on H800 (~160 GB/s NVLink maximum bandwidth), with each connected to a CX7 InfiniBand 400 Gb/s RDMA network card (~50 GB/s maximum bandwidth). And we follow the DeepSeek-V3/R1 pretraining setting (4096 tokens per batch, 7168 hidden, top-4 groups, top-8 experts, FP8 dispatching and BF16 combining).
| Type | Dispatch #EP | Bottleneck bandwidth | Combine #EP | Bottleneck bandwidth |
|:---------:|:------------:|:--------------------:|:-----------:|:--------------------:|
| Intranode | 8 | 153 GB/s (NVLink) | 8 | 158 GB/s (NVLink) |
| Internode | 16 | 43 GB/s (RDMA) | 16 | 43 GB/s (RDMA) |
| Internode | 32 | 58 GB/s (RDMA) | 32 | 57 GB/s (RDMA) |
| Internode | 64 | 51 GB/s (RDMA) | 64 | 50 GB/s (RDMA) |
### Low-latency kernels with pure RDMA
We test low-latency kernels on H800 with each connected to a CX7 InfiniBand 400 Gb/s RDMA network card (~50 GB/s maximum bandwidth). And we follow a typical DeepSeek-V3/R1 production setting (128 tokens per batch, 7168 hidden, top-8 experts, FP8 dispatching and BF16 combining).
| Dispatch #EP | Latency | RDMA bandwidth | Combine #EP | Latency | RDMA bandwidth |
|:------------:|:-------:|:--------------:|:-----------:|:-------:|:--------------:|
| 8 | 77 us | 98 GB/s | 8 | 114 us | 127 GB/s |
| 16 | 118 us | 63 GB/s | 16 | 195 us | 74 GB/s |
| 32 | 155 us | 48 GB/s | 32 | 273 us | 53 GB/s |
| 64 | 173 us | 43 GB/s | 64 | 314 us | 46 GB/s |
| 128 | 192 us | 39 GB/s | 128 | 369 us | 39 GB/s |
| 256 | 194 us | 39 GB/s | 256 | 360 us | 40 GB/s |
## Quick Start
### Requirements
- Ampere (SM80), Hopper (SM90) GPUs, or other architectures with SM90 PTX ISA support
- Python 3.8 and above
- CUDA version
- CUDA 11.0 and above for SM80 GPUs
- CUDA 12.3 and above for SM90 GPUs
- PyTorch 2.1 and above
- NVLink for intranode communication
- RDMA network for internode communication
### Download and install NVSHMEM dependency
DeepEP V1 depends on NVSHMEM. Please refer to the NVSHMEM Installation Guide for instructions.
### Development
```bash
# Build and make symbolic links for SO files
NVSHMEM_DIR=/path/to/installed/nvshmem python setup.py build
# You may modify the specific SO names according to your own platform
ln -s build/lib.linux-x86_64-cpython-38/deep_ep_cpp.cpython-38-x86_64-linux-gnu.so
# Run test cases
# NOTES: you may modify the `init_dist` function in `tests/utils.py`
# according to your own cluster settings, and launch into multiple nodes
python tests/test_intranode.py
python tests/test_internode.py
python tests/test_low_latency.py
```
### Installation
```bash
NVSHMEM_DIR=/path/to/installed/nvshmem python setup.py install
```
#### Installation environment variables
- `NVSHMEM_DIR`: the path to the NVSHMEM directory, disable all internode and low-latency features if not specified
- `DISABLE_SM90_FEATURES`: 0 or 1, whether to disable SM90 features, it is required for SM90 devices or CUDA 11
- `TORCH_CUDA_ARCH_LIST`: the list of target architectures, e.g. `TORCH_CUDA_ARCH_LIST="9.0"`
- `DISABLE_AGGRESSIVE_PTX_INSTRS`: 0 or 1, whether to disable aggressive load/store instructions, see [Undefined-behavior PTX usage](#undefined-behavior-ptx-usage) for more details
## Network Configurations
DeepEP is fully tested with InfiniBand networks. However, it is theoretically compatible with RDMA over Converged Ethernet (RoCE) as well.
### Traffic isolation
Traffic isolation is supported by InfiniBand through Virtual Lanes (VL).
To prevent interference between different types of traffic, we recommend segregating workloads across different virtual lanes as follows:
- workloads using normal kernels
- workloads using low-latency kernels
- other workloads
For DeepEP V1, you can control the virtual lane assignment by setting the `NVSHMEM_IB_SL` environment variable.
### Adaptive routing
Adaptive routing is an advanced routing feature provided by InfiniBand switches that can evenly distribute traffic across multiple paths. Enabling adaptive routing can completely eliminate network congestion caused by routing conflicts, but it also introduces additional latency. We recommend the following configuration for optimal performance:
- enable adaptive routing in environments with heavy network loads
- use static routing in environments with light network loads
### Congestion control
Congestion control is disabled as we have not observed significant congestion in our production environment.
## Interfaces and Examples
### Example use in model training or inference prefilling
The normal kernels can be used in model training or the inference prefilling phase (without the backward part) as the below example code shows.
```python
import torch
import torch.distributed as dist
from typing import List, Tuple, Optional, Union
from deep_ep import Buffer, EventOverlap
# Communication buffer (will allocate at runtime)
_buffer: Optional[Buffer] = None
# Set the number of SMs to use
# NOTES: this is a static variable
Buffer.set_num_sms(24)
# You may call this function at the framework initialization
def get_buffer(group: dist.ProcessGroup, hidden_bytes: int) -> Buffer:
global _buffer
# NOTES: you may also replace `get_*_config` with your auto-tuned results via all the tests
num_nvl_bytes, num_rdma_bytes = 0, 0
for config in (Buffer.get_dispatch_config(group.size()), Buffer.get_combine_config(group.size())):
num_nvl_bytes = max(config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes)
num_rdma_bytes = max(config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes)
# Allocate a buffer if not existed or not enough buffer size
if _buffer is None or _buffer.group != group or _buffer.num_nvl_bytes < num_nvl_bytes or _buffer.num_rdma_bytes < num_rdma_bytes:
_buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes)
return _buffer
def get_hidden_bytes(x: torch.Tensor) -> int:
t = x[0] if isinstance(x, tuple) else x
return t.size(1) * max(t.element_size(), 2)
def dispatch_forward(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
topk_idx: torch.Tensor, topk_weights: torch.Tensor,
num_experts: int, previous_event: Optional[EventOverlap] = None) -> \
Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], torch.Tensor, torch.Tensor, List, Tuple, EventOverlap]:
# NOTES: an optional `previous_event` means a CUDA event captured that you want to make it as a dependency
# of the dispatch kernel, it may be useful with communication-computation overlap. For more information, please
# refer to the docs of `Buffer.dispatch`
global _buffer
# Calculate layout before actual dispatch
num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, previous_event = \
_buffer.get_dispatch_layout(topk_idx, num_experts,
previous_event=previous_event, async_finish=True,
allocate_on_comm_stream=previous_event is not None)
# Do MoE dispatch
# NOTES: the CPU will wait for GPU's signal to arrive, so this is not compatible with CUDA graph
# Unless you specify `num_worst_tokens`, but this flag is for intranode only
# For more advanced usages, please refer to the docs of the `dispatch` function
recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event = \
_buffer.dispatch(x, topk_idx=topk_idx, topk_weights=topk_weights,
num_tokens_per_rank=num_tokens_per_rank, num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
is_token_in_rank=is_token_in_rank, num_tokens_per_expert=num_tokens_per_expert,
previous_event=previous_event, async_finish=True,
allocate_on_comm_stream=True)
# For event management, please refer to the docs of the `EventOverlap` class
return recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event
def dispatch_backward(grad_recv_x: torch.Tensor, grad_recv_topk_weights: torch.Tensor, handle: Tuple) -> \
Tuple[torch.Tensor, torch.Tensor, EventOverlap]:
global _buffer
# The backward process of MoE dispatch is actually a combine
# For more advanced usages, please refer to the docs of the `combine` function
combined_grad_x, combined_grad_recv_topk_weights, event = \
_buffer.combine(grad_recv_x, handle, topk_weights=grad_recv_topk_weights, async_finish=True)
# For event management, please refer to the docs of the `EventOverlap` class
return combined_grad_x, combined_grad_recv_topk_weights, event
def combine_forward(x: torch.Tensor, handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \
Tuple[torch.Tensor, EventOverlap]:
global _buffer
# Do MoE combine
# For more advanced usages, please refer to the docs of the `combine` function
combined_x, _, event = _buffer.combine(x, handle, async_finish=True, previous_event=previous_event,
allocate_on_comm_stream=previous_event is not None)
# For event management, please refer to the docs of the `EventOverlap` class
return combined_x, event
def combine_backward(grad_combined_x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \
Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], EventOverlap]:
global _buffer
# The backward process of MoE combine is actually a dispatch
# For more advanced usages, please refer to the docs of the `dispatch` function
grad_x, _, _, _, _, event = _buffer.dispatch(grad_combined_x, handle=handle, async_finish=True,
previous_event=previous_event,
allocate_on_comm_stream=previous_event is not None)
# For event management, please refer to the docs of the `EventOverlap` class
return grad_x, event
```
Moreover, inside the dispatch function, we may not know how many tokens to receive for the current rank. So an implicit CPU wait for GPU received count signal will be involved, as the following figure shows.
![normal](../figures/normal.png)
### Example use in inference decoding
The low latency kernels can be used in the inference decoding phase as the below example code shows.
```python
import torch
import torch.distributed as dist
from typing import Tuple, Optional
from deep_ep import Buffer
# Communication buffer (will allocate at runtime)
# NOTES: there is no SM control API for the low-latency kernels
_buffer: Optional[Buffer] = None
# You may call this function at the framework initialization
def get_buffer(group: dist.ProcessGroup, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int) -> Buffer:
# NOTES: the low-latency mode will consume much more space than the normal mode
# So we recommend that `num_max_dispatch_tokens_per_rank` (the actual batch size in the decoding engine) should be less than 256
global _buffer
num_rdma_bytes = Buffer.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, group.size(), num_experts)
# Allocate a buffer if not existed or not enough buffer size
if _buffer is None or _buffer.group != group or not _buffer.low_latency_mode or _buffer.num_rdma_bytes < num_rdma_bytes:
# NOTES: for the best performance, the QP number **must** be equal to the number of the local experts
assert num_experts % group.size() == 0
_buffer = Buffer(group, 0, num_rdma_bytes, low_latency_mode=True, num_qps_per_rank=num_experts // group.size())
return _buffer
def low_latency_dispatch(hidden_states: torch.Tensor, topk_idx: torch.Tensor, num_max_dispatch_tokens_per_rank: int, num_experts: int):
global _buffer
# Do MoE dispatch, compatible with CUDA graph (but you may restore some buffer status once you replay)
recv_hidden_states, recv_expert_count, handle, event, hook = \
_buffer.low_latency_dispatch(hidden_states, topk_idx, num_max_dispatch_tokens_per_rank, num_experts,
async_finish=False, return_recv_hook=True)
# NOTES: the actual tensor will not be received only if you call `hook()`,
# it is useful for double-batch overlapping, but **without any SM occupation**
# If you don't want to overlap, please set `return_recv_hook=False`
# Later, you can use our GEMM library to do the computation with this specific format
return recv_hidden_states, recv_expert_count, handle, event, hook
def low_latency_combine(hidden_states: torch.Tensor,
topk_idx: torch.Tensor, topk_weights: torch.Tensor, handle: Tuple):
global _buffer
# Do MoE combine, compatible with CUDA graph (but you may restore some buffer status once you replay)
combined_hidden_states, event_overlap, hook = \
_buffer.low_latency_combine(hidden_states, topk_idx, topk_weights, handle,
async_finish=False, return_recv_hook=True)
# NOTES: the same behavior as described in the dispatch kernel
return combined_hidden_states, event_overlap, hook
```
For two-micro-batch overlapping, you can refer to the following figure. With our receiving hook interface, the RDMA network traffic is happening in the background, without costing any GPU SMs from the computation part. But notice, the overlapped parts can be adjusted, i.e., the 4 parts of attention/dispatch/MoE/combine may not have the exact same execution time. You may adjust the stage settings according to your workload.
![low-latency](../figures/low-latency.png)
## Roadmap (V1)
- [x] AR support
- [x] Refactor low-latency mode AR code
- [x] A100 support (intranode only)
- [x] Support BF16 for the low-latency dispatch kernel
- [x] Support NVLink protocol for intranode low-latency kernels
- [ ] TMA copy instead of LD/ST
- [x] Intranode kernels
- [ ] Internode kernels
- [ ] Low-latency kernels
- [ ] SM-free kernels and refactors
- [ ] Fully remove undefined-behavior PTX instructions
## Notices
#### Easier potential overall design
The V1 implementation uses queues for communication buffers which save memory but introduce complexity and potential deadlocks. If you're implementing your own version based on DeepEP V1, consider using fixed-size buffers allocated to maximum capacity for simplicity and better performance. For a detailed discussion of this alternative approach, see https://github.com/deepseek-ai/DeepEP/issues/39.
#### Undefined-behavior PTX usage
- For extreme performance, we discover and use an undefined-behavior PTX usage: using read-only PTX `ld.global.nc.L1::no_allocate.L2::256B` to **read volatile data**. The PTX modifier `.nc` indicates that a non-coherent cache is used. But the correctness is tested to be guaranteed with `.L1::no_allocate` on Hopper architectures, and performance will be much better. The reason we guess may be: the non-coherent cache is unified with L1, and the L1 modifier is not just a hint but a strong option, so that the correctness can be guaranteed by no dirty data in L1.
- Initially, because NVCC could not automatically unroll volatile read PTX, we tried using `__ldg` (i.e., `ld.nc`). Even compared to manually unrolled volatile reads, it was significantly faster (likely due to additional compiler optimizations). However, the results could be incorrect or dirty. After consulting the PTX documentation, we discovered that L1 and non-coherent cache are unified on Hopper architectures. We speculated that `.L1::no_allocate` might resolve the issue, leading to this discovery.
- If you find kernels not working on some other platforms, you may add `DISABLE_AGGRESSIVE_PTX_INSTRS=1` to `setup.py` and disable this, or file an issue.
#### Auto-tuning on your cluster
For better performance on your cluster, we recommend to run all the tests and use the best auto-tuned configuration. The default configurations are optimized on the DeepSeek's internal cluster.
docs/nvshmem.md
+77
View File
@@ -0,0 +1,77 @@
# Install NVSHMEM
## Important notices
**This project is neither sponsored nor supported by NVIDIA.**
**Use of NVIDIA NVSHMEM is governed by the terms at [NVSHMEM Software License Agreement](https://docs.nvidia.com/nvshmem/api/sla.html).**
## Prerequisites
Hardware requirements:
- GPUs inside one node needs to be connected by NVLink
- GPUs across different nodes needs to be connected by RDMA devices, see [GPUDirect RDMA Documentation](https://docs.nvidia.com/cuda/gpudirect-rdma/)
- InfiniBand GPUDirect Async (IBGDA) support, see [IBGDA Overview](https://developer.nvidia.com/blog/improving-network-performance-of-hpc-systems-using-nvidia-magnum-io-nvshmem-and-gpudirect-async/)
- For more detailed requirements, see [NVSHMEM Hardware Specifications](https://docs.nvidia.com/nvshmem/release-notes-install-guide/install-guide/abstract.html#hardware-requirements)
Software requirements:
- NVSHMEM v3.3.9 or later
## Installation procedure
### 1. Install NVSHMEM binaries
NVSHMEM 3.3.9 binaries are available in several formats:
- Tarballs for [x86_64](https://developer.download.nvidia.com/compute/nvshmem/redist/libnvshmem/linux-x86_64/libnvshmem-linux-x86_64-3.3.9_cuda12-archive.tar.xz) and [aarch64](https://developer.download.nvidia.com/compute/nvshmem/redist/libnvshmem/linux-sbsa/libnvshmem-linux-sbsa-3.3.9_cuda12-archive.tar.xz)
- RPM and deb packages: instructions can be found on the [NVSHMEM installer page](https://developer.nvidia.com/nvshmem-downloads?target_os=Linux)
- Conda packages through conda-forge
- pip wheels through PyPI: `pip install nvidia-nvshmem-cu12`
DeepEP is compatible with upstream NVSHMEM 3.3.9 and later.
### 2. Enable NVSHMEM IBGDA support
NVSHMEM Supports two modes with different requirements. Either of the following methods can be used to enable IBGDA support.
#### 2.1 Configure NVIDIA driver
This configuration enables traditional IBGDA support.
Modify `/etc/modprobe.d/nvidia.conf`:
```bash
options nvidia NVreg_EnableStreamMemOPs=1 NVreg_RegistryDwords="PeerMappingOverride=1;"
```
Update kernel configuration:
```bash
sudo update-initramfs -u
sudo reboot
```
#### 2.2 Install GDRCopy and load the gdrdrv kernel module
This configuration enables IBGDA through asynchronous post-send operations assisted by the CPU. More information about CPU-assisted IBGDA can be found in [this blog](https://developer.nvidia.com/blog/enhancing-application-portability-and-compatibility-across-new-platforms-using-nvidia-magnum-io-nvshmem-3-0/#cpu-assisted_infiniband_gpu_direct_async%C2%A0).
It comes with a small performance penalty, but can be used when modifying the driver regkeys is not an option.
Download GDRCopy
GDRCopy is available as prebuilt deb and rpm packages [here](https://developer.download.nvidia.com/compute/redist/gdrcopy/). or as source code on the [GDRCopy github repository](https://github.com/NVIDIA/gdrcopy).
Install GDRCopy following the instructions on the [GDRCopy github repository](https://github.com/NVIDIA/gdrcopy?tab=readme-ov-file#build-and-installation).
## Post-installation configuration
When not installing NVSHMEM from RPM or deb packages, set the following environment variables in your shell configuration:
```bash
export NVSHMEM_DIR=/path/to/your/dir/to/install # Use for DeepEP installation
export LD_LIBRARY_PATH="${NVSHMEM_DIR}/lib:$LD_LIBRARY_PATH"
export PATH="${NVSHMEM_DIR}/bin:$PATH"
```
## Verification
```bash
nvshmem-info -a # Should display details of nvshmem
```
figures/low-latency.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 672 KiB

figures/normal.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 506 KiB

format.sh
Executable
+194
View File
@@ -0,0 +1,194 @@
#!/usr/bin/env bash
# Usage:
# # Do work and commit your work.
# # Format files that differ from origin/main.
# bash format.sh
# # Commit changed files with message 'Run yapf and ruff'
#
#
# YAPF + Clang formatter (if installed). This script formats all changed files from the last mergebase.
# You are encouraged to run this locally before pushing changes for review.
# Cause the script to exit if a single command fails
set -eo pipefail
# If yapf/ruff is not installed, install according to the requirements
if ! (yapf --version &>/dev/null && ruff --version &>/dev/null); then
pip install -r requirements-lint.txt
fi
YAPF_VERSION=$(yapf --version | awk '{print $2}')
RUFF_VERSION=$(ruff --version | awk '{print $2}')
echo 'yapf: Check Start'
YAPF_FLAGS=(
'--recursive'
'--parallel'
)
YAPF_EXCLUDES=(
'--exclude' 'build/**'
)
# Format specified files
format() {
yapf --in-place "${YAPF_FLAGS[@]}" "$@"
}
# Format all files
format_all() {
yapf --in-place "${YAPF_FLAGS[@]}" "${YAPF_EXCLUDES[@]}" .
}
# Format files that differ from main branch
format_changed() {
# The `if` guard ensures that the list of filenames is not empty, which
# could cause ruff to receive 0 positional arguments, making it hang
# waiting for STDIN.
#
# `diff-filter=ACM` and $MERGEBASE is to ensure we only lint files that
# exist on both branches.
if git show-ref --verify --quiet refs/remotes/origin/main; then
BASE_BRANCH="origin/main"
else
BASE_BRANCH="main"
fi
MERGEBASE="$(git merge-base $BASE_BRANCH HEAD)"
if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then
git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs -P 5 \
yapf --in-place "${YAPF_EXCLUDES[@]}" "${YAPF_FLAGS[@]}"
fi
}
# If `--all` is passed, then any further arguments are ignored and the
# entire python directory is formatted.
if [[ "$1" == '--all' ]]; then
format_all
else
# Format only the files that changed in last commit.
format_changed
fi
echo 'yapf: Done'
echo 'ruff: Check Start'
# Lint specified files
lint() {
ruff check "$@"
}
# Lint files that differ from main branch. Ignores dirs that are not slated
# for autolint yet.
lint_changed() {
# The `if` guard ensures that the list of filenames is not empty, which
# could cause ruff to receive 0 positional arguments, making it hang
# waiting for STDIN.
#
# `diff-filter=ACM` and $MERGEBASE is to ensure we only lint files that
# exist on both branches.
if git show-ref --verify --quiet refs/remotes/origin/main; then
BASE_BRANCH="origin/main"
else
BASE_BRANCH="main"
fi
MERGEBASE="$(git merge-base $BASE_BRANCH HEAD)"
if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then
git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs \
ruff check
fi
}
# Run Ruff
# If `--all` is passed, then any further arguments are ignored and the
# entire python directory is linted.
if [[ "$1" == '--all' ]]; then
lint .
else
# Check only the files that changed in last commit.
lint_changed
fi
echo 'ruff: Done'
# # params: tool name, tool version, required version
tool_version_check() {
if [[ $2 != $3 ]]; then
echo "Wrong $1 version installed: $3 is required, not $2."
pip install -r requirements-lint.txt
fi
}
echo 'clang-format: Check Start'
# If clang-format is available, run it; otherwise, skip
if command -v clang-format &>/dev/null; then
CLANG_FORMAT_VERSION=$(clang-format --version | awk '{print $3}')
tool_version_check "clang-format" "$CLANG_FORMAT_VERSION" "$(grep clang-format requirements-lint.txt | cut -d'=' -f3)"
CLANG_FORMAT_FLAGS=("-i")
# Format all C/C++ files in the repo, excluding specified directories
clang_format_all() {
# Replace "#pragma unroll" by "// #pragma unroll"
find . -type f \( -name '*.c' -o -name '*.cc' -o -name '*.cpp' -o -name '*.h' -o -name '*.hpp' -o -name '*.cu' -o -name '*.cuh' \) \
-not -path "./build/*" \
-exec perl -pi -e 's/#pragma unroll/\/\/#pragma unroll/g' {} +
find . -type f \( -name '*.c' -o -name '*.cc' -o -name '*.cpp' -o -name '*.h' -o -name '*.hpp' -o -name '*.cu' -o -name '*.cuh' \) \
-not -path "./build/*" \
-exec clang-format -i {} +
# Replace "// #pragma unroll" by "#pragma unroll"
find . -type f \( -name '*.c' -o -name '*.cc' -o -name '*.cpp' -o -name '*.h' -o -name '*.hpp' -o -name '*.cu' -o -name '*.cuh' \) \
-not -path "./build/*" \
-exec perl -pi -e 's/\/\/ *#pragma unroll/#pragma unroll/g' {} +
}
# Format changed C/C++ files relative to main
clang_format_changed() {
if git show-ref --verify --quiet refs/remotes/origin/main; then
BASE_BRANCH="origin/main"
else
BASE_BRANCH="main"
fi
MERGEBASE="$(git merge-base $BASE_BRANCH HEAD)"
if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.c' '*.cc' '*.cpp' '*.h' '*.hpp' '*.cu' '*.cuh' &>/dev/null; then
git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.c' '*.cc' '*.cpp' '*.h' '*.hpp' '*.cu' '*.cuh' | xargs perl -pi -e 's/#pragma unroll/\/\/#pragma unroll/g'
git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.c' '*.cc' '*.cpp' '*.h' '*.hpp' '*.cu' '*.cuh' | xargs clang-format -i
git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.c' '*.cc' '*.cpp' '*.h' '*.hpp' '*.cu' '*.cuh' | xargs perl -pi -e 's/\/\/ *#pragma unroll/#pragma unroll/g'
fi
}
if [[ "$1" == '--all' ]]; then
# If --all is given, format all eligible C/C++ files
clang_format_all
else
# Otherwise, format only changed C/C++ files
clang_format_changed
fi
else
echo "clang-format not found. Skipping C/C++ formatting."
fi
echo 'clang-format: Done'
# Check if there are any uncommitted changes after all formatting steps.
# If there are, ask the user to review and stage them.
if ! git diff --quiet &>/dev/null; then
echo 'Reformatted files. Please review and stage the changes.'
echo 'Changes not staged for commit:'
echo
git --no-pager diff --name-only
echo 'You can also copy-paste the diff below to fix the lint:'
echo
git --no-pager diff
exit 1
fi
echo 'All checks passed'
install.sh
Executable
+12
View File
@@ -0,0 +1,12 @@
# Change current directory into project root
original_dir=$(pwd)
script_dir=$(dirname "$0")
cd "$script_dir"
# Remove old dist file, build, and install
rm -rf dist
python setup.py bdist_wheel
pip install dist/*.whl
# Open users' original directory
cd "$original_dir"
pyproject.toml
+39
View File
@@ -0,0 +1,39 @@
[tool.yapf]
based_on_style = "pep8"
column_limit = 140
indent_width = 4
[tool.ruff.lint]
select = [
# pycodestyle
"E", "W",
# Pyflakes
"F",
# pyupgrade
# "UP",
# flake8-bugbear
"B",
# flake8-simplify
"SIM",
# isort
# "I",
]
ignore = [
# Module level import not at top of file
"E402",
# star imports
"F405", "F403",
# ambiguous name
"E741",
# line too long
"E501",
# key in dict.keys()
"SIM118",
# memory leaks
"B019",
# No such file or directory
"E902",
]
exclude = [
"deep_ep/__init__.py"
]
requirements-lint.txt
+3
View File
@@ -0,0 +1,3 @@
clang-format==15.0.7
yapf==0.40.2
ruff==0.6.5
setup.py
+218
View File
@@ -0,0 +1,218 @@
import ast
import re
import os
import subprocess
import setuptools
import importlib
from pathlib import Path
from setuptools.command.build_py import build_py
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
current_dir = os.path.dirname(os.path.realpath(__file__))
persistent_env_names = ('EP_JIT_CACHE_DIR', 'EP_JIT_PRINT_COMPILER_COMMAND', 'EP_NUM_TOPK_IDX_BITS', 'EP_NCCL_ROOT_DIR')
# Load discover module without triggering `deep_ep.__init__`
find_pkgs_spec = importlib.util.spec_from_file_location('find_pkgs', os.path.join(current_dir, 'deep_ep', 'utils', 'find_pkgs.py'))
find_pkgs = importlib.util.module_from_spec(find_pkgs_spec)
find_pkgs_spec.loader.exec_module(find_pkgs)
# Wheel specific: NVIDIA pip wheels (nvidia-nvshmem-cu12, nvidia-nccl-cu12)
# only ship the SO name of the host library, e.g. `libnvshmem_host.so.3`,
# without the unversioned `libnvshmem_host.so` symlink. So `-l:libnvshmem_host.so`
# (exact-name link) cannot resolve. Resolve the real file name at build time
# and pass it through to the linker instead.
def _find_versioned_so(base_dir, prefix):
"""Return the real filename of the first ``{prefix}.so*`` under ``base_dir/lib``.
Prefers an unversioned ``{prefix}.so`` symlink when present so we keep
behaving identically to the Tarball install. Falls back to the SONAME
file (``{prefix}.so.X``) shipped by pip wheels.
"""
lib_dir = Path(base_dir).joinpath('lib')
unversioned = lib_dir / f'{prefix}.so'
if unversioned.exists():
return unversioned.name
for file in sorted(lib_dir.rglob(f'{prefix}.so.*')):
return file.name
raise ModuleNotFoundError(f'{prefix}.so not found under {lib_dir}')
def get_nvshmem_host_lib_name(base_dir):
return _find_versioned_so(base_dir, 'libnvshmem_host')
def get_nccl_lib_name(base_dir):
return _find_versioned_so(base_dir, 'libnccl')
def get_package_version():
with open(Path(current_dir) / 'deep_ep' / '__init__.py', 'r') as f:
version_match = re.search(r'^__version__\s*=\s*(.*)$', f.read(), re.MULTILINE)
public_version = ast.literal_eval(version_match.group(1))
# noinspection PyBroadException
try:
status_cmd = ['git', 'status', '--porcelain']
status_output = subprocess.check_output(status_cmd).decode('ascii').strip()
if status_output:
print(f'Warning: Git working directory is not clean. Uncommitted changes:\n{status_output}')
assert False, 'Git working directory is not clean'
cmd = ['git', 'rev-parse', '--short', 'HEAD']
revision = '+' + subprocess.check_output(cmd).decode('ascii').rstrip()
except:
revision = '+local'
return f'{public_version}{revision}'
class CustomBuildPy(build_py):
def run(self):
# Make clusters' cache setting default into `envs.py`
self.generate_default_envs()
# Finally, run the regular build
build_py.run(self)
def generate_default_envs(self):
code = '# Pre-installed environment variables\n'
code += 'persistent_envs = dict()\n'
# noinspection PyShadowingNames
for name in persistent_env_names:
code += f"persistent_envs['{name}'] = '{os.environ[name]}'\n" if name in os.environ else ''
# Create temporary build directory
build_include_dir = os.path.join(self.build_lib, 'deep_ep')
os.makedirs(build_include_dir, exist_ok=True)
with open(os.path.join(self.build_lib, 'deep_ep', 'envs.py'), 'w') as f:
f.write(code)
if __name__ == '__main__':
# TODO: make NVSHMEM and legacy optional
nvshmem_root_dir = find_pkgs.find_nvshmem_root()
nccl_root_dir = find_pkgs.find_nccl_root()
# `128,2417` is used to suppress warnings of `fmt`
cxx_flags = ['-O3', '-Wno-deprecated-declarations', '-Wno-unused-variable', '-Wno-sign-compare', '-Wno-reorder', '-Wno-attributes']
nvcc_flags = ['-O3', '-Xcompiler', '-O3', '--extended-lambda', '--diag-suppress=128,2417']
sources = ['csrc/python_api.cpp', 'csrc/kernels/legacy/layout.cu', 'csrc/kernels/legacy/intranode.cu']
include_dirs = [f'{current_dir}/deep_ep/include',
f'{current_dir}/third-party/fmt/include',
'/usr/local/cuda/include/cccl']
library_dirs = []
nvcc_dlink = []
extra_link_args = ['-lcuda']
# NVSHMEM flags. Use the real on-disk file name (which may be SONAME-only
# like ``libnvshmem_host.so.3`` when NVSHMEM came from a pip wheel) so
# that ``-l:NAME`` can resolve. The static device library always ships
# under its canonical name, so it stays hard-coded.
sources.extend(['csrc/kernels/legacy/internode.cu', 'csrc/kernels/legacy/internode_ll.cu', 'csrc/kernels/backend/nvshmem.cu'])
include_dirs.extend([f'{nvshmem_root_dir}/include'])
library_dirs.extend([f'{nvshmem_root_dir}/lib'])
nvcc_dlink.extend(['-dlink', f'-L{nvshmem_root_dir}/lib', '-lnvshmem_device'])
nvshmem_host_lib = get_nvshmem_host_lib_name(nvshmem_root_dir)
extra_link_args.extend([f'-l:{nvshmem_host_lib}', '-l:libnvshmem_device.a', f'-Wl,-rpath,{nvshmem_root_dir}/lib'])
# NCCL flags. Same story as NVSHMEM above — pip wheels ship
# ``libnccl.so.2`` only, so resolve the real name dynamically.
sources.extend(['csrc/kernels/backend/nccl.cu'])
include_dirs.extend([f'{nccl_root_dir}/include'])
nccl_lib = get_nccl_lib_name(nccl_root_dir)
extra_link_args.extend([f'-l:{nccl_lib}', f'-Wl,-rpath,{nccl_root_dir}/lib'])
# CUDA driver sources
sources.extend(['csrc/kernels/backend/cuda_driver.cu'])
# TODO: remove these
if int(os.getenv('DISABLE_SM90_FEATURES', 0)):
# Prefer A100
os.environ['TORCH_CUDA_ARCH_LIST'] = os.getenv('TORCH_CUDA_ARCH_LIST', '8.0')
# Disable some SM90 features: FP8, launch methods, and TMA
cxx_flags.append('-DDISABLE_SM90_FEATURES')
nvcc_flags.append('-DDISABLE_SM90_FEATURES')
# Disable internode and low-latency kernels
assert False, 'Not implemented'
else:
# Prefer H800 series
os.environ['TORCH_CUDA_ARCH_LIST'] = os.getenv('TORCH_CUDA_ARCH_LIST', '9.0')
# CUDA 12 flags
nvcc_flags.extend(['-rdc=true', '--ptxas-options=--register-usage-level=10'])
# Disable LD/ST tricks, as some CUDA version does not support `.L1::no_allocate`
if os.environ['TORCH_CUDA_ARCH_LIST'].strip() != '9.0':
assert int(os.getenv('DISABLE_AGGRESSIVE_PTX_INSTRS', 1)) == 1
os.environ['DISABLE_AGGRESSIVE_PTX_INSTRS'] = '1'
# Disable aggressive PTX instructions
if int(os.getenv('DISABLE_AGGRESSIVE_PTX_INSTRS', '1')):
cxx_flags.append('-DDISABLE_AGGRESSIVE_PTX_INSTRS')
nvcc_flags.append('-DDISABLE_AGGRESSIVE_PTX_INSTRS')
# Legacy environment name
if 'TOPK_IDX_BITS' in os.environ:
assert 'EP_NUM_TOPK_IDX_BITS' not in os.environ
os.environ['EP_NUM_TOPK_IDX_BITS'] = os.environ['TOPK_IDX_BITS']
# Bits of `topk_idx.dtype`, choices are 32 and 64
if 'EP_NUM_TOPK_IDX_BITS' in os.environ:
num_topk_idx_bits = int(os.environ['EP_NUM_TOPK_IDX_BITS'])
cxx_flags.append(f'-DEP_NUM_TOPK_IDX_BITS={num_topk_idx_bits}')
nvcc_flags.append(f'-DEP_NUM_TOPK_IDX_BITS={num_topk_idx_bits}')
# Put them together
extra_compile_args = {
'cxx': cxx_flags,
'nvcc': nvcc_flags,
}
if len(nvcc_dlink) > 0:
extra_compile_args['nvcc_dlink'] = nvcc_dlink
# Summary
print('Build summary:')
print(f' > Sources: {sources}')
print(f' > Includes: {include_dirs}')
print(f' > Libraries: {library_dirs}')
print(f' > Compilation flags: {extra_compile_args}')
print(f' > Link flags: {extra_link_args}')
print(f' > Arch list: {os.environ["TORCH_CUDA_ARCH_LIST"]}')
print(f' > NVSHMEM path: {nvshmem_root_dir}')
print(f' > NCCL path: {nccl_root_dir}')
# Print persistent env variables
persistent_envs = []
for name in persistent_env_names:
if name in os.environ:
persistent_envs.append((name, os.environ[name]))
if len(persistent_envs) > 0:
print(f' > Persistent envs:')
for k, v in persistent_envs:
print(f' > {k}: {v}')
print()
setuptools.setup(
name='deep_ep',
version=get_package_version(),
packages=setuptools.find_packages(include=['deep_ep', 'deep_ep.*']),
package_data={
'deep_ep': [
'include/deep_ep/**/*',
]
},
ext_modules=[
CUDAExtension(name='deep_ep._C',
include_dirs=include_dirs,
library_dirs=library_dirs,
sources=sources,
extra_compile_args=extra_compile_args,
extra_link_args=extra_link_args)
],
cmdclass={
'build_ext': BuildExtension,
'build_py': CustomBuildPy
}
)
tests/elastic/test_agrs.py
+202
View File
@@ -0,0 +1,202 @@
import argparse
import math
import random
import torch
import torch.distributed as dist
import numpy as np
import deep_ep
from deep_ep.utils.envs import init_dist, dist_print
def all_gather_ref(shape: tuple, rank_idx: int, num_ranks: int, round_idx: int = 0):
ref_list = []
for i in range(num_ranks):
torch.manual_seed(42 + round_idx * 43 + i)
ref_list.append(torch.randn(shape, dtype=torch.bfloat16, device='cuda'))
return ref_list[rank_idx], torch.stack(ref_list, dim=0)
def generate_stress_ops(
num_ops: int,
num_max_inflight_agrs: int,
shape: tuple,
rank_idx: int,
num_ranks: int,
) -> tuple[list[tuple], tuple[torch.Tensor], tuple[torch.Tensor]]:
tensors, refs = zip(*(all_gather_ref(shape, rank_idx, num_ranks, round_idx=i) for i in range(num_ops)), strict=True)
unprocessed = random.sample(range(num_ops), num_ops)
inflight, ops = [], [('create_session', (-1,))]
limit = num_max_inflight_agrs
while unprocessed or inflight:
max_g = min(len(unprocessed), limit)
choices = []
if max_g > 0:
choices.append('ag')
if inflight:
choices.append('fetch')
else:
choices.append('destroy')
op = random.choice(choices)
if op == 'ag':
b = tuple(unprocessed[-random.randint(1, max_g):])
limit -= len(b)
del unprocessed[-len(b):]
inflight.append(b)
ops.append(('ag', b))
elif op == 'fetch':
ops.append(('fetch', inflight.pop(random.randrange(len(inflight)))))
else:
ops.extend([('destroy_session', (-1,)), ('create_session', (-1,))])
limit = num_max_inflight_agrs
ops.append(('destroy_session', (-1,)))
return ops, tensors, refs
def do_all_gather(buffer: deep_ep.ElasticBuffer,
is_inplace: bool, is_batched: bool,
tensors: tuple[torch.Tensor, ...],
start_event: torch.cuda.Event | None = None):
# Copy into buffer if inplace
if is_inplace:
ag_tensors = buffer.agrs_get_inplace_tensor(tuple(t.shape for t in tensors), torch.bfloat16)
for x, y in zip(ag_tensors, tensors, strict=True):
x.copy_(y)
else:
ag_tensors = tensors
# Record event
if start_event is not None:
torch.zeros(int(256e6 // 4), dtype=torch.int, device='cuda') # flush L2 cache
start_event.record()
# Do all-gather
if is_batched:
*out_tensors, handle = buffer.all_gather(ag_tensors)
return out_tensors, [handle]
else:
out_tensors, handles = [], []
for t in ag_tensors:
out_tensor, handle = buffer.all_gather(t)
out_tensors.append(out_tensor)
handles.append(handle)
return out_tensors, handles
# noinspection PyTypeChecker,PyCallingNonCallable,PyShadowingNames
@torch.inference_mode()
def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks)
# Print configs
shape = (32, 64, 2048)
num_max_inflight_agrs = args.num_max_inflight_agrs
num_max_session_bytes = deep_ep.ElasticBuffer.get_agrs_num_max_session_bytes(
group,
[shape for _ in range(num_max_inflight_agrs)],
torch.bfloat16
)
num_max_session_bytes = deep_ep.ElasticBuffer.get_agrs_buffer_size_hint(
group, num_max_session_bytes)
dist_print(f'Config:\n'
f' > Ranks: {num_ranks}\n'
f' > Shape: {shape}\n'
f' > Max inflight AGRS: {num_max_inflight_agrs}\n',
once_in_node=True)
# Create buffer
buffer = deep_ep.ElasticBuffer(group, explicitly_destroy=True, num_bytes=num_max_session_bytes)
buffer.agrs_set_config(num_max_session_bytes, num_max_inflight_agrs)
# Run stress tests
dist_print('Running stress tests:', once_in_node=True)
for seed in range(args.num_stress_iterations):
random.seed(42 + seed)
num_ops = 128
ops, tensors, refs = generate_stress_ops(num_ops, num_max_inflight_agrs, shape, rank_idx, num_ranks)
results = [None] * num_ops
handles = dict()
torch.cuda.synchronize()
for op, indices in ops:
if op == 'create_session':
buffer.create_agrs_session()
elif op == 'destroy_session':
buffer.destroy_agrs_session()
elif op == 'ag':
is_inplace, is_batched = random.random() < 0.5, random.random() < 0.8
handles[indices] = do_all_gather(buffer, is_inplace, is_batched, tuple(tensors[i] for i in indices))
elif op == 'fetch':
out_tensors, wait_handles = handles[indices]
for h in wait_handles:
h()
for out, idx in zip(out_tensors, indices, strict=True):
results[idx] = out.clone()
for i in range(num_ops):
assert results[i] is not None and torch.equal(results[i], refs[i]), \
f'Rank {rank_idx}: stress mismatch at seed={seed}, op={i}'
dist_print(f' > Seed {seed} passed ({num_ops} ops)', once_in_node=True)
dist_print(once_in_node=True)
# Destroy the buffer
dist_print(f'Profiling all-gather:', once_in_node=True)
buffer.destroy()
# Profiling
num_max_session_bytes = deep_ep.ElasticBuffer.get_agrs_num_max_session_bytes(
group,
[(2 ** 26,) for _ in range(num_max_inflight_agrs)],
torch.bfloat16
)
num_max_session_bytes = deep_ep.ElasticBuffer.get_agrs_buffer_size_hint(
group, num_max_session_bytes)
buffer = deep_ep.ElasticBuffer(group, explicitly_destroy=True, num_bytes=num_max_session_bytes)
buffer.agrs_set_config(num_max_session_bytes, num_max_inflight_agrs)
for num_bytes in (2 ** p for p in range(20, 27)):
# Create tensors
shape = (num_bytes // 2, )
tensors = tuple(torch.randn(shape, dtype=torch.bfloat16, device='cuda') for _ in range(num_max_inflight_agrs))
# Tests
for is_inplace in (False, True):
for is_batched in (False, True):
num_tests = 50
start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)]
end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)]
torch.cuda.synchronize()
for i in range(num_tests):
with buffer.agrs_new_session():
_, wait_handles = do_all_gather(buffer, is_inplace, is_batched, tensors, start_event=start_events[i])
for h in wait_handles:
h()
end_events[i].record()
torch.cuda.synchronize()
times = np.array([s.elapsed_time(e) / 1e3 for s, e in zip(start_events, end_events, strict=True)])[1:]
avg_t = np.average(times)
unit = ('MB', 1e6) if num_bytes >= 1e6 else ('KB', 1e3)
bandwidth_info = f', {num_bytes * num_ranks * num_max_inflight_agrs / avg_t / 1e9:.3f} GB/s' if num_ranks > 1 else ''
dist_print(
f' > Rank: {rank_idx:3}/{num_ranks:3} | '
f'{num_ranks} x {(num_bytes / unit[1]):.0f} {unit[0]} | '
f'avg: {avg_t / num_max_inflight_agrs * 1e6:.3f} us'
f'{bandwidth_info}'
f' (inplace={int(is_inplace)}, batched={int(is_batched)})')
dist_print(once_in_node=True)
# Destroy the runtime and communication group
buffer.destroy()
dist.destroy_process_group()
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Test all_gather kernels')
parser.add_argument('--num-processes', type=int, default=8)
parser.add_argument('--num-max-inflight-agrs', type=int, default=4)
parser.add_argument('--num-stress-iterations', type=int, default=4)
args = parser.parse_args()
torch.multiprocessing.spawn(test, args=(args.num_processes, args), nprocs=args.num_processes)
tests/elastic/test_barrier.py
+62
View File
@@ -0,0 +1,62 @@
import argparse
import torch
import torch.distributed as dist
import deep_ep
from deep_ep.utils.envs import init_dist, dist_print
from deep_ep.utils.testing import bench_kineto
def test_barrier(buffer: deep_ep.ElasticBuffer, args: argparse.Namespace):
dist_print('Profiling barrier:', once_in_node=True)
num_scaleout_ranks, num_scaleup_ranks = buffer.get_logical_domain_size()
dist_print(f'Config:\n'
f' > Ranks: {num_scaleout_ranks} x {num_scaleup_ranks}\n'
f' > #QPs: {buffer.num_allocated_qps}\n',
once_in_node=True)
# Test barrier time
def loop_barrier(num_tests=1000):
for i in range(num_tests):
buffer.barrier()
t = bench_kineto(lambda: loop_barrier(), 'barrier', barrier_comm_profiling=True, barrier=buffer.barrier)
dist_print(f' > EP: {buffer.rank_idx:3}/{buffer.num_ranks:3}, '
f'barrier time: {t * 1e6:.3f} us')
dist_print(once_in_node=True)
# noinspection PyShadowingNames
@torch.inference_mode()
def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks)
do_pressure_test = args.do_pressure_test
for i in range(int(1e9) if do_pressure_test else 1):
buffer = deep_ep.ElasticBuffer(
group, num_bytes=2 ** 30,
allow_hybrid_mode=args.allow_hybrid_mode,
num_allocated_qps=args.num_allocated_qps,
explicitly_destroy=True
)
# Test barrier
test_barrier(buffer, args)
# Destroy the runtime and communication group
buffer.destroy()
dist.destroy_process_group()
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Test elastic EP barrier performance')
parser.add_argument('--num-processes', type=int, default=8, help='Number of processes to spawn (default: 8)')
parser.add_argument('--allow-hybrid-mode', type=int, default=1, help='Whether to allow hybrid mode')
parser.add_argument('--num-allocated-qps', type=int, default=8, help='Number of QPs to use (0 means auto)')
parser.add_argument('--do-pressure-test', action='store_true', help='Whether to do pressure test')
args = parser.parse_args()
# Launch test processes
num_processes = args.num_processes
torch.multiprocessing.spawn(test_loop, args=(num_processes, args), nprocs=num_processes)

Some files were not shown because too many files have changed in this diff Show More