forked from mooncake-track/Mooncake
[Misc] Mooncake EP & Mooncake Backend (#805)
* Initialize a mooncake backend * Add pybind * Fix incorrect backend registration * Fix wheel building of mooncake_ep * Add a fake allreduce implementation * Introduce transfer_engine to mooncake_backend * Add a basic CPU proxy execution framework * Implement a seemingly working allgather * Remove mooncake_ep's dependency on etcd * Implement `_allgather_base` * Implement `allreduce` * Implement `alltoall` * Use an even-odd pattern for data transfer * Add a `set_host_ip` method * Switch to an extended-API implementation of the Mooncake backend * Implement `broadcast` * Implement `barrier` * Extend Mooncake backend to CPU * Support more operations for reduction * Fix the backend-worker coordination logic * Optimize CPU worker with a callback pattern * Add a timeout-based broken-ranks detection * Merge EP module into Mooncake's build system * Share transfer buffer across all worker instances * Switch to a more robust approach to detect broken ranks * Specify CUDA device for test_mooncake_backend.py * Explicitly stop mooncake worker * Use transfer engine's notifications to implement collective signals * Remove the unused `all_reduce_without` API * Switch to mooncake backend for test_mooncake_ep.py * Support both IB and RoCE * Fix EP unit test * Pass the auto-detected nic_id to EP Buffer * Fix CMake conditional branches when `PYTORCH_CMAKE_PATH` is not set * Fix ibgda syncing for RoCE * Revert "Share transfer buffer across all worker instances" This reverts commit964e0a96* Implement `_reduce_scatter_base` * Make CPU backends aware of broken ranks * Fix .typos.toml * Add a perf test for mooncake backend * Support more dtypes for reduction * Revert "Use transfer engine's notifications to implement collective signals" This reverts commitf20ffb21* Share worker thread among all process groups * Share transfer engine among all process groups * Fix unit tests * Add a warmup phase for transfer engine * Fix transfer engine buffer locations * Fix incorrect calculation of mooncake ep buffer * Do not use timeout detection in mooncake_ep tests * Update mooncake backend perf test * Demangle per-group buffer offset from the shared taskId * Stop allocating the useless `cuda_counter_buffer` and `cuda_data_buffer` * Split the task list into a CPU region and a CUDA region * Add a warmup for test_mooncake_backend_perf.py * Switch from raw cudaEvent to `torch::Event` * Fix MooncakeWorkCuda::wait() to make it compatible with cuda graphs * Add doc * Fix perf test * Implement all-gather for perf test * Move impl of `MooncakeEpBuffer`'s member functions to .cpp * Change `gathered_experts` to `broken_nodes` to make the API more consistent * `broken_nodes` should be `broken_ranks` * API rename * Fix format * Enable WITH_EP option in CI * Try installing torch in advance in CI * Set `TORCH_CUDA_ARCH_LIST` in CMakeLists.txt * Install required dependencies in the CI CUDA environment * [CI] Add the matching PyTorch * [CI] Add a workaround for missing `CUDA::nvToolsExt` * Remove unused pybind base class declaration of `MooncakeBackendOptions` * Support `set_device_filter` * Remove unused headers for ep_py.cpp * Build the EP-wheel with setuptools on CI * [CI] Add the build-with-ep process to release.yaml * Minor format fix * Update build guide * Fix docs * Only build EP wheel with torch==2.8.0 * Add a torch version assertion for Mooncake Backend * Fix some python typing * Use the correct group for EP's initial data sharing * API: invert `broken_ranks` and change into `active_ranks` * Followup fix for inverting the API * Fix format * Bug-fix in mooncake_ep_kernel.cu * Mooncake EP has to be built with USE_CUDA on * Fixed some issues according to the review * Fix bug
This commit is contained in:
Xun Sun
GitHub
parent
9028227711
commit
c5829aad1b
|
|
@ -294,6 +294,98 @@ jobs:
|
|||
name: mooncake-wheel-ubuntu-py${{ steps.generate_tag_flags.outputs.python_version_tag }}
|
||||
path: mooncake-wheel/dist-py${{ steps.generate_tag_flags.outputs.python_version_tag }}/*.whl
|
||||
|
||||
build-with-ep:
|
||||
runs-on: ubuntu-22.04
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ['3.10', '3.12']
|
||||
env:
|
||||
BUILD_WITH_EP: "1"
|
||||
SCCACHE_GHA_ENABLED: "true"
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
|
||||
- name: Free up disk space
|
||||
run: |
|
||||
sudo rm -rf /usr/share/dotnet
|
||||
sudo rm -rf /opt/ghc
|
||||
sudo rm -rf /opt/hostedtoolcache/CodeQL
|
||||
|
||||
- name: Install CUDA Toolkit
|
||||
uses: Jimver/cuda-toolkit@v0.2.24
|
||||
with:
|
||||
cuda: '12.8.1'
|
||||
linux-local-args: '["--toolkit"]'
|
||||
method: 'network'
|
||||
sub-packages: '["nvcc", "nvrtc-dev"]'
|
||||
non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]'
|
||||
|
||||
- name: Run sccache-cache
|
||||
uses: mozilla-actions/sccache-action@v0.0.9
|
||||
|
||||
- name: Configure sccache
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
script: |
|
||||
core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
|
||||
core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
|
||||
|
||||
- name: Run sccache stat for check
|
||||
shell: bash
|
||||
run: ${SCCACHE_PATH} --show-stats
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
sudo apt update -y
|
||||
sudo bash -x dependencies.sh -y
|
||||
pip install toml-cli # for updating the version
|
||||
pip install torch==2.8.0
|
||||
shell: bash
|
||||
|
||||
- name: Build transfer engine with EP
|
||||
run: |
|
||||
mkdir build
|
||||
cd build
|
||||
export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
|
||||
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
|
||||
cmake .. -DUSE_ETCD=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DUSE_CUDA=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_EP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
|
||||
make -j
|
||||
sudo make install
|
||||
shell: bash
|
||||
|
||||
- name: Build nvlink_allocator.so
|
||||
run: |
|
||||
mkdir -p build/mooncake-transfer-engine/nvlink-allocator
|
||||
cd mooncake-transfer-engine/nvlink-allocator
|
||||
bash build.sh --ci-build ../../build/mooncake-transfer-engine/nvlink-allocator/
|
||||
shell: bash
|
||||
|
||||
- name: Generate Python version tag
|
||||
id: generate_tag_flags
|
||||
run: |
|
||||
echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
|
||||
shell: bash
|
||||
|
||||
- name: Build Python wheel
|
||||
run: |
|
||||
BASE_VERSION=$(toml get --toml-path mooncake-wheel/pyproject.toml project.version | tr -d '"')
|
||||
toml set --toml-path mooncake-wheel/pyproject.toml project.version "${BASE_VERSION}+ep"
|
||||
# Build wheel with specific Python version
|
||||
PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_flags.outputs.python_version_tag }} ./scripts/build_wheel.sh
|
||||
shell: bash
|
||||
|
||||
- name: Upload Python wheel artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: mooncake-wheel-ubuntu-py${{ steps.generate_tag_flags.outputs.python_version_tag }}+ep
|
||||
path: mooncake-wheel/dist-py${{ steps.generate_tag_flags.outputs.python_version_tag }}/*.whl
|
||||
|
||||
build-docker:
|
||||
name: Build Docker Image
|
||||
runs-on: ubuntu-22.04
|
||||
|
|
|
|||
|
|
@ -95,8 +95,104 @@ jobs:
|
|||
name: mooncake-wheel-py${{ steps.generate_tag_release.outputs.python_version_tag }}
|
||||
path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
|
||||
|
||||
build-with-ep:
|
||||
runs-on: ubuntu-22.04
|
||||
permissions:
|
||||
contents: write
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ['3.10', '3.12']
|
||||
env:
|
||||
BUILD_WITH_EP: "1"
|
||||
steps:
|
||||
- name: Checkout source
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
|
||||
- name: Free up disk space
|
||||
run: |
|
||||
sudo rm -rf /usr/share/dotnet
|
||||
sudo rm -rf /opt/ghc
|
||||
sudo rm -rf /opt/hostedtoolcache/CodeQL
|
||||
|
||||
- name: Install CUDA Toolkit
|
||||
uses: Jimver/cuda-toolkit@v0.2.24
|
||||
with:
|
||||
cuda: '12.8.1'
|
||||
linux-local-args: '["--toolkit"]'
|
||||
method: 'network'
|
||||
sub-packages: '["nvcc", "nvrtc-dev"]'
|
||||
non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]'
|
||||
|
||||
- name: Run sccache-cache
|
||||
uses: mozilla-actions/sccache-action@v0.0.9
|
||||
|
||||
- name: Configure sccache
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
script: |
|
||||
core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
|
||||
core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
|
||||
|
||||
- name: Run sccache stat for check
|
||||
shell: bash
|
||||
run: ${SCCACHE_PATH} --show-stats
|
||||
|
||||
- name: Configure project
|
||||
run: |
|
||||
sudo apt update -y
|
||||
sudo bash -x dependencies.sh -y
|
||||
pip install toml-cli # for updating the version
|
||||
pip install torch==2.8.0
|
||||
mkdir build
|
||||
cd build
|
||||
cmake .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
|
||||
shell: bash
|
||||
|
||||
- name: Build project
|
||||
run: |
|
||||
cd build
|
||||
make -j
|
||||
sudo make install
|
||||
shell: bash
|
||||
|
||||
- name: Build nvlink_allocator.so
|
||||
run: |
|
||||
mkdir -p build/mooncake-transfer-engine/nvlink-allocator
|
||||
cd mooncake-transfer-engine/nvlink-allocator
|
||||
bash build.sh --ci-build ../../build/mooncake-transfer-engine/nvlink-allocator/
|
||||
shell: bash
|
||||
|
||||
- name: Generate Python version tag
|
||||
id: generate_tag_release
|
||||
run: |
|
||||
echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
|
||||
shell: bash
|
||||
|
||||
- name: Build Python wheel
|
||||
run: |
|
||||
BASE_VERSION=$(toml get --toml-path mooncake-wheel/pyproject.toml project.version | tr -d '"')
|
||||
toml set --toml-path mooncake-wheel/pyproject.toml project.version "${BASE_VERSION}+ep"
|
||||
# Set LD_LIBRARY_PATH for wheel building
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
|
||||
PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
|
||||
env:
|
||||
VERSION: ${{ env.VERSION }}
|
||||
|
||||
- name: Upload Python wheel artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: mooncake-wheel-py${{ steps.generate_tag_release.outputs.python_version_tag }}+ep
|
||||
path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
|
||||
|
||||
publish-release:
|
||||
needs: build
|
||||
needs:
|
||||
- build
|
||||
- build-with-ep
|
||||
runs-on: ubuntu-22.04
|
||||
permissions:
|
||||
contents: write
|
||||
|
|
|
|||
|
|
@ -1,8 +1,7 @@
|
|||
[default]
|
||||
extend-ignore-words = ["CANN"]
|
||||
|
||||
[files]
|
||||
extend-exclude = ["mooncake-ep/csrc/*.h"]
|
||||
extend-ignore-words = ["CANN", "ASO", "fre"]
|
||||
|
||||
[default.extend-words]
|
||||
CANN = "CANN"
|
||||
ASO = "ASO"
|
||||
fre = "fre"
|
||||
|
|
@ -15,6 +15,7 @@ endif()
|
|||
option(WITH_STORE "build mooncake store library and sample code" ON)
|
||||
option(WITH_P2P_STORE "build p2p store library and sample code" OFF)
|
||||
option(WITH_RUST_EXAMPLE "build the Rust interface and sample code for the transfer engine" OFF)
|
||||
option(WITH_EP "build mooncake with expert parallelism support" OFF)
|
||||
|
||||
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extern/pybind11)
|
||||
set(PYTHON_EXECUTABLE "python3")
|
||||
|
|
@ -51,6 +52,12 @@ if (WITH_STORE)
|
|||
include_directories(mooncake-store/include)
|
||||
endif()
|
||||
|
||||
if (WITH_EP)
|
||||
message(STATUS "Mooncake EP will be built")
|
||||
add_subdirectory(mooncake-ep)
|
||||
include_directories(mooncake-ep/include)
|
||||
endif()
|
||||
|
||||
add_subdirectory(mooncake-integration)
|
||||
|
||||
if (WITH_P2P_STORE)
|
||||
|
|
|
|||
|
|
@ -6,6 +6,11 @@ This document describes how to build Mooncake from source.
|
|||
```bash
|
||||
pip3 install mooncake-transfer-engine --upgrade
|
||||
```
|
||||
- To install with the Mooncake Backend and Mooncake EP support, use the following command:
|
||||
```bash
|
||||
# replace torch2.8.0 with the corresponding version
|
||||
pip3 install mooncake-transfer-engine==0.3.7+ep --upgrade
|
||||
```
|
||||
|
||||
## Automatic
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,70 @@
|
|||
# Mooncake EP & Mooncake Backend
|
||||
|
||||
## Overview
|
||||
|
||||
Mooncake EP is an adaption of [DeepEP](https://github.com/deepseek-ai/DeepEP) that supports **fault tolerance** and fast data transfer with **IBGDA**, designed as a critical component for large-scale, latency-sensitive MoE (Mixture of Experts) inference. Mooncake EP aims to retain full compatibility with the DeepEP API, with the addition of an `active_ranks` tensor passed to both the `dispatch` and `combine` functions to capture information about rank activeness. By integrating with the EPLB module, Mooncake EP ensures fault tolerance during MoE inference, enabling robust performance even in large-scale, fault-prone environments.
|
||||
|
||||
Mooncake Backend is a PyTorch distributed backend (a replacement for NCCL and Gloo) that provides **fault-tolerant collective communication primitives** and can be seamlessly integrated into machine learning systems. Built with the [Transfer Engine](transfer-engine.md), Mooncake Backend ensures that collective communications can continue even in the event of rank failures. Furthermore, it reports these failures to the upper layers of the system, allowing for graceful error handling without disrupting ongoing operations.
|
||||
|
||||
## Usage
|
||||
|
||||
### Mooncake EP
|
||||
|
||||
> **Note:** Mooncake EP currently supports only the low-latency transfer mode.
|
||||
|
||||
The API is largely consistent with DeepEP's, with only minor differences in a few parameters. Mooncake EP exposes a `Buffer` that can be imported from `mooncake.mooncake_ep_buffer`. For example, refer to `mooncake-wheel/tests/test_mooncake_ep.py`.
|
||||
|
||||
#### Buffer.get_buffer_size_hint()
|
||||
|
||||
**Signature:**
|
||||
|
||||
```python
|
||||
@staticmethod
|
||||
def get_ep_buffer_size_hint(num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int) -> int
|
||||
```
|
||||
|
||||
Calculates the number of bytes to pre-allocate for data transfer.
|
||||
|
||||
#### Buffer.\_\_init\_\_()
|
||||
|
||||
**Signature:**
|
||||
|
||||
```python
|
||||
def __init__(self, group: dist.ProcessGroup, num_ep_buffer_bytes: int = 0)
|
||||
```
|
||||
|
||||
The constructor. Ensure that only one instance is created.
|
||||
|
||||
- **group**: Must be a Mooncake Backend process group.
|
||||
- **num_ep_buffer_bytes**: The number of bytes acquired with `Buffer.get_buffer_size_hint()`
|
||||
|
||||
#### Buffer.dispatch/Buffer.combine
|
||||
|
||||
**Signature:** Similar to DeepEP's `low_latency_dispatch`/`low_latency_combine`, with two additional parameters:
|
||||
|
||||
- **active_ranks**: A tensor of shape `(num_ranks,)` containing values of 0 or 1. The indices of the broken ranks will be set to 0.
|
||||
- **timeout_us**: The timeout in microseconds for a rank to be considered broken. Set to -1 for infinite timeout.
|
||||
|
||||
### Mooncake Backend
|
||||
|
||||
Basic usage:
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from mooncake import ep
|
||||
|
||||
active_ranks = torch.ones((world_size,), dtype=torch.int32, device="cuda")
|
||||
dist.init_process_group(
|
||||
backend="mooncake",
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
pg_options=ep.MooncakeBackendOptions(active_ranks),
|
||||
)
|
||||
|
||||
dist.all_gather(...) # Standard API usage
|
||||
assert active_ranks.all() # Verify that no ranks are broken
|
||||
```
|
||||
|
||||
For a full example, see `mooncake-wheel/tests/test_mooncake_backend.py`.
|
||||
|
||||
|
|
@ -0,0 +1,42 @@
|
|||
cmake_minimum_required(VERSION 3.16)
|
||||
project(mooncake-ep)
|
||||
|
||||
# Find PyTorch's CMake prefix path
|
||||
execute_process(
|
||||
COMMAND ${PYTHON_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)"
|
||||
OUTPUT_VARIABLE PYTORCH_CMAKE_PATH
|
||||
OUTPUT_STRIP_TRAILING_WHITESPACE
|
||||
)
|
||||
if(NOT PYTORCH_CMAKE_PATH)
|
||||
message(WARNING "Could not find PyTorch CMake path! Please set Torch_DIR.")
|
||||
else ()
|
||||
message(STATUS "Found PyTorch CMake path: ${PYTORCH_CMAKE_PATH}")
|
||||
list(APPEND CMAKE_PREFIX_PATH "${PYTORCH_CMAKE_PATH}/Torch")
|
||||
endif()
|
||||
|
||||
set(TORCH_CUDA_ARCH_LIST "8.0;9.0")
|
||||
|
||||
find_package(CUDAToolkit REQUIRED)
|
||||
# https://discuss.pytorch.org/t/failed-to-find-nvtoolsext/179635/13
|
||||
if(NOT TARGET CUDA::nvToolsExt AND TARGET CUDA::nvtx3)
|
||||
add_library(CUDA::nvToolsExt INTERFACE IMPORTED)
|
||||
target_compile_definitions(
|
||||
CUDA::nvToolsExt INTERFACE
|
||||
TORCH_CUDA_USE_NVTX3
|
||||
)
|
||||
target_link_libraries(CUDA::nvToolsExt INTERFACE CUDA::nvtx3)
|
||||
endif()
|
||||
find_package(Torch REQUIRED)
|
||||
include_directories(${TORCH_INCLUDE_DIRS})
|
||||
|
||||
include_directories(include)
|
||||
add_subdirectory(include)
|
||||
add_subdirectory(src)
|
||||
|
||||
if (BUILD_UNIT_TESTS)
|
||||
add_subdirectory(tests)
|
||||
endif()
|
||||
|
||||
if (BUILD_EXAMPLES)
|
||||
add_subdirectory(example)
|
||||
endif()
|
||||
|
|
@ -1,25 +0,0 @@
|
|||
# NOTES: this CMake is only for debugging; for setup, please use Torch extension
|
||||
cmake_minimum_required(VERSION 3.10)
|
||||
project(mxa_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 "-O3")
|
||||
list(APPEND CUDA_NVCC_FLAGS "--ptxas-options=--verbose,--register-usage-level=10,--warn-on-local-memory-usage")
|
||||
|
||||
find_package(CUDAToolkit REQUIRED)
|
||||
find_package(pybind11 REQUIRED)
|
||||
find_package(Torch REQUIRED)
|
||||
|
||||
# Seems bugs with CMake, NVCC 12 and C++ 17
|
||||
set(CMAKE_CXX_STANDARD 17)
|
||||
set(CMAKE_CUDA_STANDARD 14)
|
||||
|
||||
include_directories(${CUDA_TOOLKIT_ROOT_DIR}/include ${TORCH_INCLUDE_DIRS} ${PYTHON_INCLUDE_DIRS} ${TRANSFER_ENGINE_INCLUDE_DIR})
|
||||
link_directories(${TORCH_INSTALL_PREFIX}/lib ${CUDA_TOOLKIT_ROOT_DIR}/lib)
|
||||
|
||||
# Link CPP and CUDA together
|
||||
pybind11_add_module(mxa_ep_cpp mxa_ep.cpp mlx5gda.cpp mxa_kernel.cu)
|
||||
target_link_libraries(mxa_ep_cpp PRIVATE ${EP_CUDA_LIBRARIES} ${TORCH_LIBRARIES} torch_python)
|
||||
|
|
@ -1,594 +0,0 @@
|
|||
#include <fstream>
|
||||
#include <pybind11/functional.h>
|
||||
#include <torch/python.h>
|
||||
|
||||
#include "api.cuh"
|
||||
#include "event.hpp"
|
||||
#include "exception.cuh"
|
||||
#include "memheap.h"
|
||||
#include "mlx5gda.h"
|
||||
|
||||
#ifndef TORCH_EXTENSION_NAME
|
||||
#define TORCH_EXTENSION_NAME mxa_ep_cpp
|
||||
#endif
|
||||
|
||||
namespace mxa_ep {
|
||||
|
||||
struct BufferLayout {
|
||||
int* rdma_send_signal_buffer;
|
||||
int* rdma_recv_signal_buffer;
|
||||
void* rdma_send_data_buffer;
|
||||
void* rdma_recv_data_buffer;
|
||||
|
||||
void* cuda_counter_buffer;
|
||||
void* cuda_data_buffer;
|
||||
};
|
||||
|
||||
struct BufferPair {
|
||||
size_t total_bytes = 0;
|
||||
BufferLayout buffers[2];
|
||||
|
||||
template <typename out_ptr_t = void*, typename count_ptr_t = uint8_t*,
|
||||
typename in_ptr_t = void*>
|
||||
static 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);
|
||||
}
|
||||
|
||||
BufferPair(void* rdma_buffer, int num_max_dispatch_tokens_per_rank,
|
||||
int hidden, int num_ranks, int num_experts,
|
||||
size_t bytes_reserved) {
|
||||
total_bytes = bytes_reserved;
|
||||
size_t signaling_buffer_bytes = num_experts * sizeof(int);
|
||||
size_t send_recv_buffer_bytes =
|
||||
num_experts * num_max_dispatch_tokens_per_rank *
|
||||
(2 * sizeof(int4) + hidden * sizeof(nv_bfloat16));
|
||||
for (int i = 0; i < 2; ++i) {
|
||||
size_t rdma_base_offset = total_bytes +
|
||||
2 * i * signaling_buffer_bytes +
|
||||
2 * i * send_recv_buffer_bytes;
|
||||
buffers[i] = {
|
||||
advance<int*>(rdma_buffer, rdma_base_offset),
|
||||
advance<int*>(rdma_buffer,
|
||||
rdma_base_offset + signaling_buffer_bytes),
|
||||
advance<int*>(rdma_buffer,
|
||||
rdma_base_offset + 2 * signaling_buffer_bytes),
|
||||
advance<int*>(rdma_buffer, rdma_base_offset +
|
||||
2 * signaling_buffer_bytes +
|
||||
send_recv_buffer_bytes),
|
||||
advance<void*>(rdma_buffer, rdma_base_offset +
|
||||
2 * signaling_buffer_bytes +
|
||||
2 * send_recv_buffer_bytes),
|
||||
advance<void*>(rdma_buffer, rdma_base_offset +
|
||||
3 * signaling_buffer_bytes +
|
||||
2 * send_recv_buffer_bytes),
|
||||
};
|
||||
}
|
||||
total_bytes += 4 * signaling_buffer_bytes + 4 * send_recv_buffer_bytes;
|
||||
}
|
||||
};
|
||||
|
||||
struct Buffer {
|
||||
private:
|
||||
// Device info and communication
|
||||
int device_id;
|
||||
int rank, num_ranks;
|
||||
int clock_rate_khz;
|
||||
|
||||
// MXA Buffer
|
||||
int buffer_idx{};
|
||||
int64_t num_mxa_bytes;
|
||||
void* gdr_buffer = nullptr;
|
||||
size_t bytes_reserved; // For all-reduce
|
||||
|
||||
// IBGDA
|
||||
const size_t ctrl_buf_size = 256 * 1024 * 1024; // 256 MiB
|
||||
void* ctrl_buf = nullptr;
|
||||
ibv_mr* mr;
|
||||
std::vector<mlx5gda_qp*> qps;
|
||||
ibv_gid gid;
|
||||
void* raddrs = nullptr;
|
||||
void* rkeys = nullptr;
|
||||
void* qp_devctxs = nullptr;
|
||||
|
||||
// Stream for communication
|
||||
at::cuda::CUDAStream comm_stream;
|
||||
|
||||
// Workspace
|
||||
void* workspace = nullptr;
|
||||
|
||||
public:
|
||||
Buffer(int rank, int num_ranks, int64_t num_mxa_bytes,
|
||||
size_t bytes_reserved)
|
||||
: rank(rank),
|
||||
num_ranks(num_ranks),
|
||||
num_mxa_bytes(num_mxa_bytes),
|
||||
bytes_reserved(bytes_reserved),
|
||||
comm_stream(at::cuda::getStreamFromPool(true)) {
|
||||
// Get ranks
|
||||
CUDA_CHECK(cudaGetDevice(&device_id));
|
||||
CUDA_CHECK(cudaDeviceGetAttribute(&clock_rate_khz, cudaDevAttrClockRate,
|
||||
device_id));
|
||||
CUDA_CHECK(cudaMalloc(&gdr_buffer, num_mxa_bytes));
|
||||
CUDA_CHECK(cudaMalloc(&raddrs, num_ranks * sizeof(uint64_t)));
|
||||
CUDA_CHECK(cudaMalloc(&rkeys, num_ranks * sizeof(uint32_t)));
|
||||
CUDA_CHECK(
|
||||
cudaMalloc(&qp_devctxs, num_ranks * sizeof(mlx5gda_qp_devctx)));
|
||||
init_ibgda();
|
||||
|
||||
// Create 32 MiB workspace
|
||||
CUDA_CHECK(cudaMalloc(&workspace, NUM_WORKSPACE_BYTES));
|
||||
CUDA_CHECK(
|
||||
cudaMemsetAsync(workspace, 0, NUM_WORKSPACE_BYTES, comm_stream));
|
||||
}
|
||||
|
||||
~Buffer() noexcept(false) {
|
||||
cudaFree(gdr_buffer);
|
||||
cudaFree(raddrs);
|
||||
cudaFree(rkeys);
|
||||
cudaFree(qp_devctxs);
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>, torch::Tensor,
|
||||
torch::Tensor, torch::Tensor, std::optional<EventHandle>,
|
||||
std::optional<std::function<void()>>>
|
||||
dispatch(const torch::Tensor& x, const torch::Tensor& topk_idx,
|
||||
torch::Tensor& broken_nodes, int num_max_dispatch_tokens_per_rank,
|
||||
int num_experts, int timeout_us, bool use_fp8, bool async,
|
||||
bool return_recv_hook) {
|
||||
// Tensor checks
|
||||
// By default using `ptp128c` FP8 cast
|
||||
EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous() and
|
||||
x.scalar_type() == torch::kBFloat16);
|
||||
EP_HOST_ASSERT(x.size(1) % sizeof(int4) == 0 and x.size(1) % 128 == 0);
|
||||
EP_HOST_ASSERT(topk_idx.dim() == 2 and topk_idx.is_contiguous());
|
||||
EP_HOST_ASSERT(x.size(0) == topk_idx.size(0) and
|
||||
x.size(0) <= num_max_dispatch_tokens_per_rank);
|
||||
EP_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt64);
|
||||
EP_HOST_ASSERT(num_experts % num_ranks == 0);
|
||||
|
||||
auto num_tokens = static_cast<int>(x.size(0)),
|
||||
hidden = static_cast<int>(x.size(1));
|
||||
auto num_scales = hidden / 128,
|
||||
num_topk = static_cast<int>(topk_idx.size(1));
|
||||
int num_local_experts = num_experts / num_ranks;
|
||||
|
||||
// Buffer control
|
||||
BufferPair layout(gdr_buffer, num_max_dispatch_tokens_per_rank, hidden,
|
||||
num_ranks, num_experts, bytes_reserved);
|
||||
EP_HOST_ASSERT(layout.total_bytes <= num_mxa_bytes + bytes_reserved);
|
||||
auto buffer = layout.buffers[buffer_idx];
|
||||
auto next_buffer = layout.buffers[buffer_idx ^= 1];
|
||||
|
||||
// Wait previous tasks to be finished
|
||||
// NOTES: the hook mode will always use the default stream
|
||||
auto compute_stream = at::cuda::getCurrentCUDAStream();
|
||||
auto launch_stream = return_recv_hook ? compute_stream : comm_stream;
|
||||
EP_HOST_ASSERT(not(async and return_recv_hook));
|
||||
if (not return_recv_hook) stream_wait(launch_stream, compute_stream);
|
||||
|
||||
// Allocate packed tensors
|
||||
auto packed_recv_x =
|
||||
torch::empty({num_local_experts,
|
||||
num_ranks * num_max_dispatch_tokens_per_rank, hidden},
|
||||
x.options().dtype(use_fp8 ? torch::kFloat8_e4m3fn
|
||||
: torch::kBFloat16));
|
||||
auto packed_recv_src_info = torch::empty(
|
||||
{num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank},
|
||||
torch::dtype(torch::kInt32).device(torch::kCUDA));
|
||||
auto packed_recv_layout_range =
|
||||
torch::empty({num_local_experts, num_ranks},
|
||||
torch::dtype(torch::kInt64).device(torch::kCUDA));
|
||||
auto packed_recv_count =
|
||||
torch::zeros({num_local_experts},
|
||||
torch::dtype(torch::kInt32).device(torch::kCUDA));
|
||||
|
||||
// Allocate column-majored scales
|
||||
auto packed_recv_x_scales = std::optional<torch::Tensor>();
|
||||
float* packed_recv_x_scales_ptr = nullptr;
|
||||
if (use_fp8) {
|
||||
EP_HOST_ASSERT(
|
||||
(num_ranks * num_max_dispatch_tokens_per_rank) % 4 == 0 and
|
||||
"TMA requires the number of tokens to be multiple of 4");
|
||||
packed_recv_x_scales = torch::empty(
|
||||
{num_local_experts, num_scales,
|
||||
num_ranks * num_max_dispatch_tokens_per_rank},
|
||||
torch::dtype(torch::kFloat32).device(torch::kCUDA));
|
||||
packed_recv_x_scales =
|
||||
torch::transpose(packed_recv_x_scales.value(), 1, 2);
|
||||
packed_recv_x_scales_ptr = packed_recv_x_scales->data_ptr<float>();
|
||||
}
|
||||
|
||||
int64_t timeout_ticks =
|
||||
timeout_us == -1
|
||||
? -1
|
||||
: (int64_t)clock_rate_khz * (int64_t)timeout_us / 1000;
|
||||
|
||||
auto launcher = [=](int phases) {
|
||||
cudaMemsetAsync(buffer.cuda_counter_buffer, 0,
|
||||
num_experts * sizeof(int), launch_stream);
|
||||
mxa_ep::dispatch(
|
||||
packed_recv_x.data_ptr(), packed_recv_x_scales_ptr,
|
||||
packed_recv_src_info.data_ptr<int>(),
|
||||
packed_recv_layout_range.data_ptr<int64_t>(),
|
||||
packed_recv_count.data_ptr<int>(),
|
||||
broken_nodes.data_ptr<int32_t>(), gdr_buffer,
|
||||
buffer.rdma_send_signal_buffer, buffer.rdma_recv_signal_buffer,
|
||||
buffer.rdma_send_data_buffer, buffer.rdma_recv_data_buffer,
|
||||
buffer.cuda_counter_buffer, buffer.cuda_data_buffer, raddrs,
|
||||
rkeys, qp_devctxs, x.data_ptr(), topk_idx.data_ptr<int64_t>(),
|
||||
next_buffer.rdma_recv_signal_buffer, num_tokens, hidden,
|
||||
num_max_dispatch_tokens_per_rank, num_topk, num_experts, rank,
|
||||
num_ranks, use_fp8, workspace, launch_stream, timeout_ticks,
|
||||
phases);
|
||||
};
|
||||
launcher(return_recv_hook
|
||||
? LOW_LATENCY_SEND_PHASE
|
||||
: (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE));
|
||||
|
||||
// Wait streams
|
||||
std::optional<EventHandle> event;
|
||||
if (async) {
|
||||
// NOTES: we must ensure the all tensors will not be deallocated
|
||||
// before the stream-wait happens, so in Python API, we must wrap
|
||||
// all tensors into the event handle.
|
||||
event = EventHandle(launch_stream);
|
||||
} else if (not return_recv_hook) {
|
||||
stream_wait(compute_stream, launch_stream);
|
||||
}
|
||||
|
||||
// Receiver callback
|
||||
std::optional<std::function<void()>> recv_hook = std::nullopt;
|
||||
if (return_recv_hook)
|
||||
recv_hook = [=]() { launcher(LOW_LATENCY_RECV_PHASE); };
|
||||
|
||||
// Return values
|
||||
return {packed_recv_x,
|
||||
packed_recv_x_scales,
|
||||
packed_recv_count,
|
||||
packed_recv_src_info,
|
||||
packed_recv_layout_range,
|
||||
event,
|
||||
recv_hook};
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<EventHandle>,
|
||||
std::optional<std::function<void()>>>
|
||||
combine(const torch::Tensor& x, const torch::Tensor& topk_idx,
|
||||
const torch::Tensor& topk_weights, const torch::Tensor& src_info,
|
||||
const torch::Tensor& layout_range, torch::Tensor& gathered_experts,
|
||||
int num_max_dispatch_tokens_per_rank, int num_experts,
|
||||
int timeout_us, bool zero_copy, bool async, bool return_recv_hook,
|
||||
const std::optional<torch::Tensor>& out) {
|
||||
// Tensor checks
|
||||
EP_HOST_ASSERT(x.dim() == 3 and x.is_contiguous() and
|
||||
x.scalar_type() == torch::kBFloat16);
|
||||
EP_HOST_ASSERT(x.size(0) == num_experts / num_ranks);
|
||||
EP_HOST_ASSERT(x.size(1) ==
|
||||
num_ranks * num_max_dispatch_tokens_per_rank);
|
||||
EP_HOST_ASSERT(x.size(2) % sizeof(int4) == 0 and x.size(2) % 128 == 0);
|
||||
EP_HOST_ASSERT(topk_idx.dim() == 2 and topk_idx.is_contiguous());
|
||||
EP_HOST_ASSERT(topk_idx.size(0) == topk_weights.size(0) and
|
||||
topk_idx.size(1) == topk_weights.size(1));
|
||||
EP_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt64);
|
||||
EP_HOST_ASSERT(topk_weights.dim() == 2 and
|
||||
topk_weights.is_contiguous());
|
||||
EP_HOST_ASSERT(topk_weights.size(0) <=
|
||||
num_max_dispatch_tokens_per_rank);
|
||||
EP_HOST_ASSERT(topk_weights.scalar_type() == torch::kFloat32);
|
||||
EP_HOST_ASSERT(src_info.dim() == 2 and src_info.is_contiguous());
|
||||
EP_HOST_ASSERT(src_info.scalar_type() == torch::kInt32 and
|
||||
x.size(0) == src_info.size(0));
|
||||
EP_HOST_ASSERT(layout_range.dim() == 2 and
|
||||
layout_range.is_contiguous());
|
||||
EP_HOST_ASSERT(layout_range.scalar_type() == torch::kInt64);
|
||||
EP_HOST_ASSERT(layout_range.size(0) == num_experts / num_ranks and
|
||||
layout_range.size(1) == num_ranks);
|
||||
auto hidden = static_cast<int>(x.size(2));
|
||||
auto num_local_experts = num_experts / num_ranks,
|
||||
num_topk = static_cast<int>(topk_weights.size(1));
|
||||
auto num_combined_tokens = static_cast<int>(topk_weights.size(0));
|
||||
|
||||
// Buffer control
|
||||
BufferPair layout(gdr_buffer, num_max_dispatch_tokens_per_rank, hidden,
|
||||
num_ranks, num_experts, bytes_reserved);
|
||||
EP_HOST_ASSERT(layout.total_bytes <= num_mxa_bytes + bytes_reserved);
|
||||
auto buffer = layout.buffers[buffer_idx];
|
||||
auto next_buffer = layout.buffers[buffer_idx ^= 1];
|
||||
|
||||
// Wait previous tasks to be finished
|
||||
// NOTES: the hook mode will always use the default stream
|
||||
auto compute_stream = at::cuda::getCurrentCUDAStream();
|
||||
auto launch_stream = return_recv_hook ? compute_stream : comm_stream;
|
||||
EP_HOST_ASSERT(not(async and return_recv_hook));
|
||||
if (not return_recv_hook) stream_wait(launch_stream, compute_stream);
|
||||
|
||||
// Allocate output tensor
|
||||
torch::Tensor combined_x;
|
||||
if (out.has_value()) {
|
||||
EP_HOST_ASSERT(out->dim() == 2 and out->is_contiguous());
|
||||
EP_HOST_ASSERT(out->size(0) == num_combined_tokens and
|
||||
out->size(1) == hidden);
|
||||
EP_HOST_ASSERT(out->scalar_type() == x.scalar_type());
|
||||
combined_x = out.value();
|
||||
} else {
|
||||
combined_x =
|
||||
torch::empty({num_combined_tokens, hidden}, x.options());
|
||||
}
|
||||
|
||||
int64_t timeout_ticks =
|
||||
timeout_us == -1
|
||||
? -1
|
||||
: (int64_t)clock_rate_khz * (int64_t)timeout_us / 1000;
|
||||
|
||||
// Kernel launch
|
||||
auto launcher = [=](int phases) {
|
||||
mxa_ep::combine(
|
||||
combined_x.data_ptr(), gathered_experts.data_ptr<int32_t>(),
|
||||
gdr_buffer, buffer.rdma_send_signal_buffer,
|
||||
buffer.rdma_recv_signal_buffer, buffer.rdma_send_data_buffer,
|
||||
buffer.rdma_recv_data_buffer, buffer.cuda_counter_buffer,
|
||||
buffer.cuda_data_buffer, raddrs, rkeys, qp_devctxs,
|
||||
x.data_ptr(), topk_idx.data_ptr<int64_t>(),
|
||||
topk_weights.data_ptr<float>(), src_info.data_ptr<int>(),
|
||||
layout_range.data_ptr<int64_t>(),
|
||||
next_buffer.rdma_recv_signal_buffer, num_combined_tokens,
|
||||
hidden, num_max_dispatch_tokens_per_rank, num_topk, num_experts,
|
||||
rank, num_ranks, workspace, launch_stream, timeout_ticks,
|
||||
phases, zero_copy);
|
||||
};
|
||||
launcher(return_recv_hook
|
||||
? LOW_LATENCY_SEND_PHASE
|
||||
: (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE));
|
||||
|
||||
// Wait streams
|
||||
std::optional<EventHandle> event;
|
||||
if (async) {
|
||||
// NOTES: we must ensure the all tensors will not be deallocated
|
||||
// before the stream-wait happens, so in Python API, we must wrap
|
||||
// all tensors into the event handle.
|
||||
event = EventHandle(launch_stream);
|
||||
} else if (not return_recv_hook) {
|
||||
stream_wait(compute_stream, launch_stream);
|
||||
}
|
||||
|
||||
// Receiver callback
|
||||
std::optional<std::function<void()>> recv_hook = std::nullopt;
|
||||
if (return_recv_hook)
|
||||
recv_hook = [=]() { launcher(LOW_LATENCY_RECV_PHASE); };
|
||||
|
||||
// Return values
|
||||
return {combined_x, event, recv_hook};
|
||||
}
|
||||
|
||||
torch::Tensor get_next_combine_buffer(int num_max_dispatch_tokens_per_rank,
|
||||
int hidden, int num_experts) {
|
||||
BufferPair layout(gdr_buffer, num_max_dispatch_tokens_per_rank, hidden,
|
||||
num_ranks, num_experts, bytes_reserved);
|
||||
|
||||
auto buffer = layout.buffers[buffer_idx];
|
||||
auto dtype = torch::kBFloat16;
|
||||
size_t num_bytes_per_combine_msg = hidden * sizeof(nv_bfloat16);
|
||||
auto num_msg_elems = static_cast<int>(num_bytes_per_combine_msg /
|
||||
elementSize(torch::kBFloat16));
|
||||
|
||||
EP_HOST_ASSERT(
|
||||
num_bytes_per_combine_msg % elementSize(torch::kBFloat16) == 0);
|
||||
return torch::from_blob(
|
||||
buffer.rdma_send_data_buffer,
|
||||
{num_experts / num_ranks,
|
||||
num_ranks * num_max_dispatch_tokens_per_rank, hidden},
|
||||
{num_ranks * num_max_dispatch_tokens_per_rank * num_msg_elems,
|
||||
num_msg_elems, 1},
|
||||
torch::TensorOptions().dtype(dtype).device(torch::kCUDA));
|
||||
}
|
||||
|
||||
void all_reduce_without(const torch::Tensor& broken_nodes,
|
||||
torch::Tensor& x) {
|
||||
auto mxa_buffer = reinterpret_cast<int*>(gdr_buffer);
|
||||
int size = x.numel();
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
cudaMemsetAsync(mxa_buffer, 0, (1 + size) * num_ranks * sizeof(int),
|
||||
stream);
|
||||
mxa_ep::all_reduce_without(broken_nodes.data_ptr<int32_t>(),
|
||||
x.data_ptr<int>(), mxa_buffer, raddrs, rkeys,
|
||||
qp_devctxs, size, rank, num_ranks, stream);
|
||||
}
|
||||
|
||||
void init_ibgda() {
|
||||
std::ifstream config("gpu_to_nic.txt");
|
||||
if (!config)
|
||||
throw std::runtime_error("Cannot open config file gpu_to_nic.txt");
|
||||
std::unordered_map<int, int> gpu_to_nic;
|
||||
std::string line;
|
||||
for (size_t lineno = 1; std::getline(config, line); ++lineno) {
|
||||
// Strip everything after ‘#’
|
||||
if (auto pos = line.find('#'); pos != std::string::npos)
|
||||
line.erase(pos);
|
||||
std::istringstream iss(line);
|
||||
int gpu, nic;
|
||||
if (!(iss >> gpu >> nic)) {
|
||||
if (iss.rdbuf()->in_avail() == 0)
|
||||
continue; // blank or comment line
|
||||
throw std::runtime_error("Parse error in gpu_to_nic.txt line " +
|
||||
std::to_string(lineno));
|
||||
}
|
||||
if (!gpu_to_nic.emplace(gpu, nic).second)
|
||||
throw std::runtime_error("Duplicate GPU id on line " +
|
||||
std::to_string(lineno));
|
||||
}
|
||||
|
||||
auto nic = gpu_to_nic.find(device_id);
|
||||
if (nic == gpu_to_nic.end())
|
||||
throw std::out_of_range("GPU id not found in config");
|
||||
int num_devices;
|
||||
ibv_device** dev_list = ibv_get_device_list(&num_devices);
|
||||
printf("GPU %d uses NIC %d out of %d NIC(s)\n", device_id, nic->second,
|
||||
num_devices);
|
||||
ibv_context* ctx = ibv_open_device(dev_list[nic->second]);
|
||||
if (!ctx) {
|
||||
perror("Failed to open device");
|
||||
exit(1);
|
||||
}
|
||||
if (ibv_query_gid(ctx, 1, 3, &gid)) {
|
||||
perror("Failed to query gid");
|
||||
}
|
||||
ibv_free_device_list(dev_list);
|
||||
|
||||
ibv_pd* pd = ibv_alloc_pd(ctx);
|
||||
if (!pd) {
|
||||
perror("Failed to allocate protection domain");
|
||||
exit(1);
|
||||
}
|
||||
mlx5dv_pd mpd;
|
||||
mlx5dv_obj dv_obj = {};
|
||||
dv_obj.pd.in = pd;
|
||||
dv_obj.pd.out = &mpd;
|
||||
if (mlx5dv_init_obj(&dv_obj, MLX5DV_OBJ_PD)) {
|
||||
perror("Failed to initialize mlx5dv object");
|
||||
}
|
||||
mr = ibv_reg_mr(pd, gdr_buffer, num_mxa_bytes,
|
||||
IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ |
|
||||
IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_ATOMIC);
|
||||
if (!mr) {
|
||||
perror("Failed to reg mr");
|
||||
}
|
||||
|
||||
CUDA_CHECK(cudaMalloc(&ctrl_buf, ctrl_buf_size));
|
||||
CUDA_CHECK(cudaMemset(ctrl_buf, 0, ctrl_buf_size));
|
||||
mlx5dv_devx_umem* ctrl_buf_umem = mlx5dv_devx_umem_reg(
|
||||
ctx, ctrl_buf, ctrl_buf_size, IBV_ACCESS_LOCAL_WRITE);
|
||||
if (!ctrl_buf_umem) {
|
||||
perror("Failed to register control buffer as umem");
|
||||
fprintf(stderr,
|
||||
"If the error is `Bad address`, probably because your GPU "
|
||||
"does not support GPUDirect RDMA.\n");
|
||||
exit(1);
|
||||
}
|
||||
memheap* ctrl_buf_heap = memheap_create(ctrl_buf_size);
|
||||
if (!ctrl_buf_heap) {
|
||||
perror("Failed to create memory heap");
|
||||
exit(1);
|
||||
}
|
||||
for (int i = 0; i < num_ranks; ++i) {
|
||||
mlx5gda_qp* qp = mlx5gda_create_rc_qp(mpd, ctrl_buf, ctrl_buf_umem,
|
||||
ctrl_buf_heap, pd, 16384, 1);
|
||||
if (!qp) {
|
||||
perror("Failed to create QP");
|
||||
exit(1);
|
||||
}
|
||||
if (mlx5gda_modify_rc_qp_rst2init(qp, 0)) {
|
||||
perror("Failed to mlx5gda_modify_rc_qp_rst2init");
|
||||
exit(1);
|
||||
}
|
||||
mlx5gda_qp_devctx qp_devctx = {
|
||||
.qpn = qp->qpn,
|
||||
.wqeid_mask = qp->num_wqebb - 1,
|
||||
.wq = (mlx5gda_wqebb*)(ctrl_buf + qp->wq_offset),
|
||||
.cq = (mlx5_cqe64*)(ctrl_buf + qp->send_cq->cq_offset),
|
||||
.dbr = (mlx5gda_wq_dbr*)(ctrl_buf + qp->dbr_offset),
|
||||
.bf = (char*)qp->uar->reg_addr,
|
||||
};
|
||||
cudaMemcpy(qp_devctxs + i * sizeof(mlx5gda_qp_devctx), &qp_devctx,
|
||||
sizeof(mlx5gda_qp_devctx), cudaMemcpyHostToDevice);
|
||||
qps.push_back(qp);
|
||||
}
|
||||
}
|
||||
|
||||
void sync(const std::vector<int64_t>& remote_addrs,
|
||||
const std::vector<int32_t>& remote_keys,
|
||||
const std::vector<int32_t>& remote_qpns,
|
||||
const std::vector<int64_t>& subnet_prefixes,
|
||||
const std::vector<int64_t>& interface_ids) {
|
||||
for (int i = 0; i < num_ranks; ++i) {
|
||||
ibv_gid remote_gid{};
|
||||
remote_gid.global.subnet_prefix = subnet_prefixes[i];
|
||||
remote_gid.global.interface_id = interface_ids[i];
|
||||
ibv_ah_attr ah_attr = {};
|
||||
ah_attr.is_global = 1;
|
||||
ah_attr.grh.dgid = remote_gid;
|
||||
ah_attr.grh.sgid_index = 3;
|
||||
ah_attr.grh.hop_limit = 1;
|
||||
ah_attr.port_num = 1;
|
||||
ah_attr.dlid = qps[i]->port_attr.lid | 0xC000;
|
||||
if (mlx5gda_modify_rc_qp_init2rtr(
|
||||
qps[i], ah_attr, (uint32_t)remote_qpns[i], IBV_MTU_4096)) {
|
||||
perror("Failed to mlx5gda_modify_rc_qp_init2rtr");
|
||||
exit(1);
|
||||
}
|
||||
if (mlx5gda_modify_rc_qp_rtr2rts(qps[i])) {
|
||||
perror("Failed to mlx5gda_modify_rc_qp_rtr2rts");
|
||||
exit(1);
|
||||
}
|
||||
uint64_t raddr =
|
||||
i == rank ? (uint64_t)mr->addr : (uint64_t)remote_addrs[i];
|
||||
cudaMemcpy(raddrs + i * sizeof(uint64_t), &raddr, sizeof(uint64_t),
|
||||
cudaMemcpyHostToDevice);
|
||||
uint32_t rkey = i == rank ? mr->lkey : (uint32_t)remote_keys[i];
|
||||
cudaMemcpy(rkeys + i * sizeof(uint32_t), &rkey, sizeof(uint32_t),
|
||||
cudaMemcpyHostToDevice);
|
||||
}
|
||||
}
|
||||
|
||||
std::tuple<int64_t, int32_t> get_mr_info() {
|
||||
return {(int64_t)mr->addr, (int32_t)mr->rkey};
|
||||
}
|
||||
|
||||
std::tuple<int64_t, int64_t> get_gid() {
|
||||
return {(int64_t)gid.global.subnet_prefix,
|
||||
(int64_t)gid.global.interface_id};
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> get_local_qpns() {
|
||||
std::vector<torch::Tensor> local_qpns;
|
||||
for (int i = 0; i < num_ranks; ++i) {
|
||||
local_qpns.push_back(
|
||||
torch::full({1}, qps[i]->qpn,
|
||||
torch::dtype(torch::kInt32).device(torch::kCUDA)));
|
||||
}
|
||||
return local_qpns;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> get_local_lids() {
|
||||
std::vector<torch::Tensor> local_lids;
|
||||
for (int i = 0; i < num_ranks; ++i) {
|
||||
local_lids.push_back(
|
||||
torch::full({1}, qps[i]->port_attr.lid,
|
||||
torch::dtype(torch::kInt32).device(torch::kCUDA)));
|
||||
}
|
||||
return local_lids;
|
||||
}
|
||||
};
|
||||
|
||||
size_t get_mxa_size_hint(int num_max_dispatch_tokens_per_rank, int hidden,
|
||||
int num_ranks, int num_experts,
|
||||
size_t bytes_reserved) {
|
||||
return BufferPair(nullptr, num_max_dispatch_tokens_per_rank, hidden,
|
||||
num_ranks, num_experts, bytes_reserved)
|
||||
.total_bytes;
|
||||
}
|
||||
|
||||
} // namespace mxa_ep
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.doc() = "MXA-EP: Expert parallelism with MXA";
|
||||
|
||||
m.def("get_mxa_size_hint", &mxa_ep::get_mxa_size_hint);
|
||||
|
||||
pybind11::class_<mxa_ep::EventHandle>(m, "EventHandle")
|
||||
.def(pybind11::init<>())
|
||||
.def("current_stream_wait", &mxa_ep::EventHandle::current_stream_wait);
|
||||
|
||||
pybind11::class_<mxa_ep::Buffer>(m, "Buffer")
|
||||
.def(pybind11::init<int, int, int64_t, size_t>())
|
||||
.def("sync", &mxa_ep::Buffer::sync)
|
||||
.def("get_mr_info", &mxa_ep::Buffer::get_mr_info)
|
||||
.def("get_gid", &mxa_ep::Buffer::get_gid)
|
||||
.def("get_local_qpns", &mxa_ep::Buffer::get_local_qpns)
|
||||
.def("get_local_lids", &mxa_ep::Buffer::get_local_lids)
|
||||
.def("all_reduce_without", &mxa_ep::Buffer::all_reduce_without)
|
||||
.def("dispatch", &mxa_ep::Buffer::dispatch)
|
||||
.def("combine", &mxa_ep::Buffer::combine)
|
||||
.def("get_next_combine_buffer",
|
||||
&mxa_ep::Buffer::get_next_combine_buffer);
|
||||
}
|
||||
|
|
@ -0,0 +1,86 @@
|
|||
#ifndef MOONCAKE_BACKEND_H
|
||||
#define MOONCAKE_BACKEND_H
|
||||
|
||||
#include <mooncake_worker.cuh>
|
||||
#include <torch/torch.h>
|
||||
#include <torch/csrc/distributed/c10d/Backend.hpp>
|
||||
#include <transfer_engine.h>
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
class MooncakeBackend final : public ::c10d::Backend {
|
||||
public:
|
||||
struct MooncakeBackendOptions final : ::c10d::Backend::Options {
|
||||
explicit MooncakeBackendOptions(at::Tensor activeRanks)
|
||||
: Options{"mooncake"}, activeRanks_{activeRanks} {}
|
||||
|
||||
~MooncakeBackendOptions() override = default;
|
||||
|
||||
at::Tensor activeRanks_;
|
||||
};
|
||||
|
||||
MooncakeBackend(c10::intrusive_ptr<::c10d::Store> store, int rank, int size,
|
||||
c10::intrusive_ptr<MooncakeBackendOptions> options,
|
||||
bool isCpu = false);
|
||||
|
||||
~MooncakeBackend() override;
|
||||
|
||||
const std::string getBackendName() const override;
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> broadcast(
|
||||
std::vector<at::Tensor>& tensors,
|
||||
const c10d::BroadcastOptions& opts) override;
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> allreduce(
|
||||
std::vector<at::Tensor>& tensors,
|
||||
const c10d::AllreduceOptions& opts) override;
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> allgather(
|
||||
std::vector<std::vector<at::Tensor>>& outputTensors,
|
||||
std::vector<at::Tensor>& inputTensors,
|
||||
const c10d::AllgatherOptions& opts) override;
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> _allgather_base(
|
||||
at::Tensor& outputBuffer, at::Tensor& inputBuffer,
|
||||
const c10d::AllgatherOptions& opts) override;
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> _reduce_scatter_base(
|
||||
at::Tensor& outputBuffer, at::Tensor& inputBuffer,
|
||||
const c10d::ReduceScatterOptions& opts) override;
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> alltoall(
|
||||
std::vector<at::Tensor>& outputTensors,
|
||||
std::vector<at::Tensor>& inputTensors,
|
||||
const c10d::AllToAllOptions& opts) override;
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> barrier(
|
||||
const c10d::BarrierOptions& opts) override;
|
||||
|
||||
static void setHostIp(const std::string& hostIp) { hostIp_ = hostIp; }
|
||||
|
||||
static void setDeviceFilter(std::vector<std::string> filters) {
|
||||
engine_.setWhitelistFilters(std::move(filters));
|
||||
}
|
||||
|
||||
std::string getPreferredHca(std::string location) {
|
||||
auto matrix = engine_.getLocalTopology()->getMatrix();
|
||||
return matrix[location].preferred_hca[0];
|
||||
}
|
||||
|
||||
private:
|
||||
static TransferEngine engine_;
|
||||
static Transport* transport_;
|
||||
static int backendIndex_;
|
||||
bool isCpu_{false};
|
||||
static std::string hostIp_;
|
||||
void* send_buffer_[2];
|
||||
void* recv_buffer_[2];
|
||||
int32_t* cpu_sync_send_region_[2];
|
||||
int32_t* cpu_sync_recv_region_[2];
|
||||
static MooncakeWorker worker_;
|
||||
TransferGroupMeta meta_;
|
||||
};
|
||||
|
||||
} // namespace mooncake
|
||||
|
||||
#endif // MOONCAKE_BACKEND_H
|
||||
|
|
@ -2,11 +2,11 @@
|
|||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
namespace mxa_ep {
|
||||
namespace mooncake {
|
||||
|
||||
void dispatch(void* packed_recv_x, float* packed_recv_x_scales,
|
||||
int* packed_recv_src_info, int64_t* packed_recv_layout_range,
|
||||
int* packed_recv_count, int32_t* broken_nodes, void* mxa_buffer,
|
||||
int* packed_recv_count, int32_t* active_ranks, void* mxa_buffer,
|
||||
int* rdma_send_signal_buffer, int* rdma_recv_signal_buffer,
|
||||
void* rdma_send_data_buffer, void* rdma_recv_data_buffer,
|
||||
void* cuda_counter_buffer, void* cuda_data_buffer, void* raddrs,
|
||||
|
|
@ -17,7 +17,7 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales,
|
|||
void* workspace, cudaStream_t stream, int64_t timeout_ticks,
|
||||
int phases);
|
||||
|
||||
void combine(void* combined_x, int32_t* gathered_experts, void* mxa_buffer,
|
||||
void combine(void* combined_x, int32_t* active_ranks, void* mxa_buffer,
|
||||
int* rdma_send_signal_buffer, int* rdma_recv_signal_buffer,
|
||||
void* rdma_send_data_buffer, void* rdma_recv_data_buffer,
|
||||
void* cuda_counter_buffer, void* cuda_data_buffer, void* raddrs,
|
||||
|
|
@ -30,9 +30,4 @@ void combine(void* combined_x, int32_t* gathered_experts, void* mxa_buffer,
|
|||
cudaStream_t stream, int64_t timeout_ticks, int phases,
|
||||
bool zero_copy);
|
||||
|
||||
void all_reduce_without(const int32_t* broken_nodes, int* x, int* mxa_buffer,
|
||||
void* raddrs, void* rkeys, void* qp_devctxs,
|
||||
int num_experts, int rank, int num_ranks,
|
||||
cudaStream_t stream);
|
||||
|
||||
} // namespace mxa_ep
|
||||
} // namespace mooncake
|
||||
|
|
@ -0,0 +1,168 @@
|
|||
#ifndef MOONCAKE_EP_BUFFER_H
|
||||
#define MOONCAKE_EP_BUFFER_H
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <fstream>
|
||||
#include <mooncake_ibgda/memheap.h>
|
||||
#include <mooncake_ibgda/mlx5gda.h>
|
||||
#include <mooncake_ep_api.cuh>
|
||||
#include <mooncake_ep_configs.cuh>
|
||||
#include <mooncake_ep_event.h>
|
||||
#include <mooncake_ep_exception.cuh>
|
||||
#include <torch/torch.h>
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
struct BufferLayout {
|
||||
int* rdma_send_signal_buffer;
|
||||
int* rdma_recv_signal_buffer;
|
||||
void* rdma_send_data_buffer;
|
||||
void* rdma_recv_data_buffer;
|
||||
};
|
||||
|
||||
struct BufferPair {
|
||||
size_t total_bytes = 0;
|
||||
BufferLayout buffers[2];
|
||||
|
||||
template <typename out_ptr_t = void*, typename count_ptr_t = uint8_t*,
|
||||
typename in_ptr_t = void*>
|
||||
static 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);
|
||||
}
|
||||
|
||||
BufferPair(void* rdma_buffer, int num_max_dispatch_tokens_per_rank,
|
||||
int hidden, int num_ranks, int num_experts) {
|
||||
size_t signaling_buffer_bytes = num_experts * sizeof(int);
|
||||
size_t send_recv_buffer_bytes =
|
||||
num_experts * num_max_dispatch_tokens_per_rank *
|
||||
(2 * sizeof(int4) + hidden * sizeof(nv_bfloat16));
|
||||
for (int i = 0; i < 2; ++i) {
|
||||
size_t rdma_base_offset = total_bytes +
|
||||
2 * i * signaling_buffer_bytes +
|
||||
2 * i * send_recv_buffer_bytes;
|
||||
buffers[i] = {
|
||||
advance<int*>(rdma_buffer, rdma_base_offset),
|
||||
advance<int*>(rdma_buffer,
|
||||
rdma_base_offset + signaling_buffer_bytes),
|
||||
advance<int*>(rdma_buffer,
|
||||
rdma_base_offset + 2 * signaling_buffer_bytes),
|
||||
advance<int*>(rdma_buffer, rdma_base_offset +
|
||||
2 * signaling_buffer_bytes +
|
||||
send_recv_buffer_bytes),
|
||||
};
|
||||
}
|
||||
total_bytes += 4 * signaling_buffer_bytes + 4 * send_recv_buffer_bytes;
|
||||
}
|
||||
};
|
||||
|
||||
struct MooncakeEpBuffer {
|
||||
private:
|
||||
// Device info and communication
|
||||
int device_id;
|
||||
int rank, num_ranks;
|
||||
int clock_rate_khz;
|
||||
|
||||
// MXA Buffer
|
||||
int buffer_idx{};
|
||||
int64_t num_ep_buffer_bytes;
|
||||
void* gdr_buffer = nullptr;
|
||||
|
||||
// IBGDA
|
||||
static constexpr size_t CTRL_BUF_SIZE = 1024 * 1024 * 1024; // 1024 MiB
|
||||
void* ctrl_buf = nullptr;
|
||||
ibv_mr* mr;
|
||||
std::vector<mlx5gda_qp*> qps;
|
||||
ibv_gid gid;
|
||||
void* raddrs = nullptr;
|
||||
void* rkeys = nullptr;
|
||||
void* qp_devctxs = nullptr;
|
||||
int nic_id;
|
||||
bool is_roce_ = false;
|
||||
|
||||
// Stream for communication
|
||||
at::cuda::CUDAStream comm_stream;
|
||||
|
||||
// Workspace
|
||||
void* workspace = nullptr;
|
||||
|
||||
public:
|
||||
MooncakeEpBuffer(int rank, int num_ranks, int64_t num_ep_buffer_bytes,
|
||||
int nic_id);
|
||||
|
||||
~MooncakeEpBuffer() noexcept(false);
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>, torch::Tensor,
|
||||
torch::Tensor, torch::Tensor, std::optional<EventHandle>,
|
||||
std::optional<std::function<void()>>>
|
||||
dispatch(const torch::Tensor& x, const torch::Tensor& topk_idx,
|
||||
torch::Tensor& active_ranks, int num_max_dispatch_tokens_per_rank,
|
||||
int num_experts, int timeout_us, bool use_fp8, bool async,
|
||||
bool return_recv_hook);
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<EventHandle>,
|
||||
std::optional<std::function<void()>>>
|
||||
combine(const torch::Tensor& x, const torch::Tensor& topk_idx,
|
||||
const torch::Tensor& topk_weights, const torch::Tensor& src_info,
|
||||
const torch::Tensor& layout_range, torch::Tensor& active_ranks,
|
||||
int num_max_dispatch_tokens_per_rank, int num_experts,
|
||||
int timeout_us, bool zero_copy, bool async, bool return_recv_hook,
|
||||
const std::optional<torch::Tensor>& out);
|
||||
|
||||
torch::Tensor get_next_combine_buffer(int num_max_dispatch_tokens_per_rank,
|
||||
int hidden, int num_experts);
|
||||
|
||||
void init_ibgda();
|
||||
|
||||
bool is_roce() { return is_roce_; }
|
||||
|
||||
void sync_ib(const std::vector<int64_t>& remote_addrs,
|
||||
const std::vector<int32_t>& remote_keys,
|
||||
const std::vector<int32_t>& remote_qpns,
|
||||
const std::vector<int32_t>& remote_lids);
|
||||
|
||||
void sync_roce(const std::vector<int64_t>& remote_addrs,
|
||||
const std::vector<int32_t>& remote_keys,
|
||||
const std::vector<int32_t>& remote_qpns,
|
||||
const std::vector<int64_t>& subnet_prefixes,
|
||||
const std::vector<int64_t>& interface_ids);
|
||||
|
||||
std::tuple<int64_t, int32_t> get_mr_info() {
|
||||
return {(int64_t)mr->addr, (int32_t)mr->rkey};
|
||||
}
|
||||
|
||||
std::tuple<int64_t, int64_t> get_gid() {
|
||||
return {(int64_t)gid.global.subnet_prefix,
|
||||
(int64_t)gid.global.interface_id};
|
||||
}
|
||||
|
||||
std::vector<int32_t> get_local_qpns() {
|
||||
std::vector<int32_t> local_qpns;
|
||||
for (int i = 0; i < MAX_QP_COUNT; ++i) {
|
||||
local_qpns.push_back((int32_t)qps[i]->qpn);
|
||||
}
|
||||
return local_qpns;
|
||||
}
|
||||
|
||||
std::vector<int32_t> get_local_lids() {
|
||||
std::vector<int32_t> local_lids;
|
||||
for (int i = 0; i < MAX_QP_COUNT; ++i) {
|
||||
local_lids.push_back((int32_t)qps[i]->port_attr.lid);
|
||||
}
|
||||
return local_lids;
|
||||
}
|
||||
};
|
||||
|
||||
inline size_t get_ep_buffer_size_hint(int num_max_dispatch_tokens_per_rank,
|
||||
int hidden, int num_ranks,
|
||||
int num_experts) {
|
||||
return BufferPair(nullptr, num_max_dispatch_tokens_per_rank, hidden,
|
||||
num_ranks, num_experts)
|
||||
.total_bytes;
|
||||
}
|
||||
|
||||
} // namespace mooncake
|
||||
|
||||
#endif // MOONCAKE_EP_BUFFER_H
|
||||
|
|
@ -1,9 +1,11 @@
|
|||
#pragma once
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <memory>
|
||||
#include <mooncake_ep_exception.cuh>
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include "exception.cuh"
|
||||
|
||||
namespace mxa_ep {
|
||||
namespace mooncake {
|
||||
|
||||
struct EventHandle {
|
||||
std::shared_ptr<torch::Event> event;
|
||||
|
|
@ -25,20 +27,21 @@ struct EventHandle {
|
|||
}
|
||||
};
|
||||
|
||||
torch::Event create_event(const at::cuda::CUDAStream& s) {
|
||||
inline torch::Event create_event(const at::cuda::CUDAStream& s) {
|
||||
auto event = torch::Event(torch::kCUDA);
|
||||
event.record(s);
|
||||
return event;
|
||||
}
|
||||
|
||||
void stream_wait(const at::cuda::CUDAStream& s_0,
|
||||
const at::cuda::CUDAStream& s_1) {
|
||||
inline void stream_wait(const at::cuda::CUDAStream& s_0,
|
||||
const at::cuda::CUDAStream& s_1) {
|
||||
EP_HOST_ASSERT(s_0.id() != s_1.id());
|
||||
s_0.unwrap().wait(create_event(s_1));
|
||||
}
|
||||
|
||||
void stream_wait(const at::cuda::CUDAStream& s, const EventHandle& event) {
|
||||
inline void stream_wait(const at::cuda::CUDAStream& s,
|
||||
const EventHandle& event) {
|
||||
s.unwrap().wait(*event.event);
|
||||
}
|
||||
|
||||
} // namespace mxa_ep
|
||||
} // namespace mooncake
|
||||
|
|
@ -3,8 +3,6 @@
|
|||
#include <string>
|
||||
#include <exception>
|
||||
|
||||
#include "configs.cuh"
|
||||
|
||||
#ifndef EP_STATIC_ASSERT
|
||||
#define EP_STATIC_ASSERT(cond, reason) static_assert(cond, reason)
|
||||
#endif
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
#pragma once
|
||||
|
||||
#include "configs.cuh"
|
||||
#include <mooncake_ep_configs.cuh>
|
||||
|
||||
#ifndef SETUP_LAUNCH_CONFIG
|
||||
#define SETUP_LAUNCH_CONFIG(num_sms, num_threads, stream) \
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
#pragma once
|
||||
|
||||
#include "exception.cuh"
|
||||
#include <mooncake_ep_exception.cuh>
|
||||
|
||||
#define UNROLLED_WARP_COPY(UNROLL_FACTOR, LANE_ID, N, DST, SRC, LD_FUNC, \
|
||||
ST_FUNC) \
|
||||
|
|
@ -22,7 +22,7 @@
|
|||
ST_FUNC(__dst + __i, LD_FUNC(__src + __i)); \
|
||||
}
|
||||
|
||||
namespace mxa_ep {
|
||||
namespace mooncake {
|
||||
|
||||
template <int kBytes>
|
||||
struct VecInt {};
|
||||
|
|
@ -535,4 +535,4 @@ __forceinline__ __device__ void barrier_device(int **task_fifo_ptrs, int head,
|
|||
timeout_check<kNumRanks>(task_fifo_ptrs, head, rank, 0, tag);
|
||||
}
|
||||
|
||||
} // namespace mxa_ep
|
||||
} // namespace mooncake
|
||||
|
|
@ -6,6 +6,8 @@
|
|||
#include <infiniband/verbs.h>
|
||||
#include <infiniband/mlx5dv.h>
|
||||
|
||||
#define MAX_QP_COUNT 256
|
||||
|
||||
struct mlx5gda_cq_dbr {
|
||||
uint64_t unused;
|
||||
};
|
||||
|
|
@ -0,0 +1,82 @@
|
|||
#ifndef MOONCAKE_WORKER_CUH
|
||||
#define MOONCAKE_WORKER_CUH
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <torch/torch.h>
|
||||
#include <torch/csrc/distributed/c10d/Types.hpp>
|
||||
#include <torch/csrc/distributed/c10d/Work.hpp>
|
||||
#include <transfer_engine.h>
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
struct TransferGroupMeta {
|
||||
int rank;
|
||||
int size;
|
||||
int taskCount;
|
||||
bool* activeRanks;
|
||||
bool* activeRanksDevice;
|
||||
at::Tensor activeRanksTensor;
|
||||
TransferEngine* engine;
|
||||
int bufferBaseIndex;
|
||||
std::vector<TransferMetadata::SegmentID> segmentIDs;
|
||||
std::vector<std::shared_ptr<TransferMetadata::SegmentDesc>> segmentDescs;
|
||||
};
|
||||
|
||||
__global__ struct Task {
|
||||
volatile bool active = false;
|
||||
c10d::OpType opType = c10d::OpType::UNKNOWN;
|
||||
size_t tensorSize; // In bytes
|
||||
int64_t broadcastRoot;
|
||||
int bufferOffset;
|
||||
BatchID batchID;
|
||||
void* transferGroupMeta;
|
||||
};
|
||||
|
||||
static constexpr size_t kBufferSize = 1u << 29;
|
||||
static constexpr size_t kMaxNumRanks = 64;
|
||||
|
||||
void launchReduceKernel(at::Tensor dst, void* src, size_t numRanks,
|
||||
c10d::ReduceOp op, bool* activeRanks,
|
||||
cudaStream_t stream);
|
||||
|
||||
void launchReduceCpu(at::Tensor dst, void* src, size_t numRanks,
|
||||
c10d::ReduceOp op);
|
||||
|
||||
class MooncakeWorker {
|
||||
public:
|
||||
explicit MooncakeWorker();
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> putTaskCpu(
|
||||
c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
|
||||
TransferGroupMeta* meta,
|
||||
const std::function<void(void* dst)>& tensorToBuffer,
|
||||
const std::function<void(void* src)>& bufferToTensor);
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> putTaskCuda(
|
||||
c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
|
||||
TransferGroupMeta* meta, const at::cuda::CUDAStream& stream,
|
||||
const std::function<void(void* dst)>& tensorToBuffer,
|
||||
const std::function<void(void* src)>& bufferToTensor);
|
||||
|
||||
void startWorker();
|
||||
|
||||
void stopWorker() { running_ = false; }
|
||||
|
||||
private:
|
||||
static constexpr size_t kNumTasks_ = 4;
|
||||
|
||||
bool running_ = false;
|
||||
|
||||
Task *tasks_, *tasks_device_;
|
||||
bool hasCallback_[kNumTasks_]{};
|
||||
std::function<void()> callbacks_[kNumTasks_]{};
|
||||
|
||||
int cpuTaskCount = 0;
|
||||
int cudaTaskCount = 0;
|
||||
};
|
||||
|
||||
} // namespace mooncake
|
||||
|
||||
#endif // MOONCAKE_WORKER_CUH
|
||||
|
|
@ -1,2 +0,0 @@
|
|||
from .buffer import Buffer
|
||||
from .utils import EventOverlap
|
||||
|
|
@ -1,89 +0,0 @@
|
|||
import numpy as np
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from typing import Callable, List, Tuple, Optional, Union
|
||||
|
||||
# noinspection PyUnresolvedReferences
|
||||
import mxa_ep_cpp
|
||||
from .utils import EventOverlap
|
||||
|
||||
class Buffer:
|
||||
def __init__(self, group: dist.ProcessGroup, num_mxa_bytes: int = 0, bytes_reserved: int = 0):
|
||||
# Initialize the CPP runtime
|
||||
self.rank = group.rank()
|
||||
self.group_size = group.size()
|
||||
self.group = group
|
||||
self.num_mxa_bytes = num_mxa_bytes
|
||||
self.runtime = mxa_ep_cpp.Buffer(self.rank, self.group_size, num_mxa_bytes, bytes_reserved)
|
||||
|
||||
(raddr, rkey) = self.runtime.get_mr_info()
|
||||
|
||||
raddr = torch.tensor([raddr], dtype=torch.int64)
|
||||
raddrs = [torch.empty(1, dtype=torch.int64) for _ in range(self.group_size)]
|
||||
dist.all_gather(raddrs, raddr)
|
||||
raddrs = torch.cat(raddrs).tolist()
|
||||
|
||||
rkey = torch.tensor([rkey], dtype=torch.int32)
|
||||
rkeys = [torch.empty(1, dtype=torch.int32) for _ in range(self.group_size)]
|
||||
dist.all_gather(rkeys, rkey)
|
||||
rkeys = torch.cat(rkeys).tolist()
|
||||
|
||||
(subnet_prefix, interface_id) = self.runtime.get_gid()
|
||||
|
||||
subnet_prefix = torch.tensor([subnet_prefix], dtype=torch.int64)
|
||||
subnet_prefixes = [torch.empty(1, dtype=torch.int64) for _ in range(self.group_size)]
|
||||
dist.all_gather(subnet_prefixes, subnet_prefix)
|
||||
subnet_prefixes = torch.cat(subnet_prefixes).tolist()
|
||||
|
||||
interface_id = torch.tensor([interface_id], dtype=torch.int64)
|
||||
interface_ids = [torch.empty(1, dtype=torch.int64) for _ in range(self.group_size)]
|
||||
dist.all_gather(interface_ids, interface_id)
|
||||
interface_ids = torch.cat(interface_ids).tolist()
|
||||
|
||||
local_qpns = self.runtime.get_local_qpns()
|
||||
remote_qpns = [torch.empty(1, dtype=torch.int32) for _ in range(self.group_size)]
|
||||
dist.all_to_all(remote_qpns, local_qpns)
|
||||
remote_qpns = torch.cat(remote_qpns).tolist()
|
||||
|
||||
self.runtime.sync(raddrs, rkeys, remote_qpns, subnet_prefixes, interface_ids)
|
||||
|
||||
@staticmethod
|
||||
def get_mxa_size_hint(num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int, bytes_reserved: int = 0) -> int:
|
||||
return mxa_ep_cpp.get_mxa_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts, bytes_reserved)
|
||||
|
||||
# noinspection PyTypeChecker
|
||||
def dispatch(self, x: torch.Tensor, topk_idx: torch.Tensor, broken_nodes: torch.Tensor,
|
||||
num_max_dispatch_tokens_per_rank: int, num_experts: int, timeout_us: int,
|
||||
use_fp8: bool = True, async_finish: bool = False, return_recv_hook: bool = False) -> \
|
||||
Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor, Tuple, EventOverlap, Callable]:
|
||||
packed_recv_x, packed_recv_x_scales, packed_recv_count, packed_recv_src_info, packed_recv_layout_range, event, hook = \
|
||||
self.runtime.dispatch(x, topk_idx, broken_nodes,
|
||||
num_max_dispatch_tokens_per_rank, num_experts, timeout_us,
|
||||
use_fp8, 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)
|
||||
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 combine(self, x: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor,
|
||||
gathered_experts: torch.Tensor, timeout_us: int,
|
||||
handle: tuple, zero_copy: bool = False, async_finish: bool = False,
|
||||
return_recv_hook: bool = False, out: Optional[torch.Tensor] = None) -> \
|
||||
Tuple[torch.Tensor, torch.Tensor, EventOverlap, Callable]:
|
||||
src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
|
||||
combined_x, event, hook = self.runtime.combine(x, topk_idx, topk_weights, src_info, layout_range,
|
||||
gathered_experts,
|
||||
num_max_dispatch_tokens_per_rank, num_experts, timeout_us,
|
||||
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 get_next_combine_buffer(self, handle: object):
|
||||
src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
|
||||
return self.runtime.get_next_combine_buffer(num_max_dispatch_tokens_per_rank, hidden, num_experts)
|
||||
|
||||
def all_reduce_without(self, broken_nodes: torch.Tensor, x: torch.Tensor):
|
||||
self.runtime.all_reduce_without(broken_nodes, x)
|
||||
|
|
@ -1,60 +0,0 @@
|
|||
import torch
|
||||
from typing import Any, Optional, Tuple
|
||||
|
||||
# noinspection PyUnresolvedReferences
|
||||
from mxa_ep_cpp import EventHandle
|
||||
|
||||
|
||||
class EventOverlap:
|
||||
"""
|
||||
A wrapper class to manage CUDA events, also for better overlapping convenience.
|
||||
|
||||
Attributes:
|
||||
event: the CUDA event captured.
|
||||
extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph.
|
||||
"""
|
||||
|
||||
def __init__(self, event: Optional[EventHandle] = None,
|
||||
extra_tensors: Optional[Tuple[torch.Tensor]] = None) -> None:
|
||||
"""
|
||||
Initialize the class.
|
||||
|
||||
Arguments:
|
||||
event: the CUDA event captured.
|
||||
extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph.
|
||||
"""
|
||||
self.event = event
|
||||
|
||||
# NOTES: we use extra tensors to achieve stream recording, otherwise,
|
||||
# stream recording will be incompatible with CUDA graph.
|
||||
self.extra_tensors = extra_tensors
|
||||
|
||||
def current_stream_wait(self) -> None:
|
||||
"""
|
||||
The current stream `torch.cuda.current_stream()` waits for the event to be finished.
|
||||
"""
|
||||
assert self.event is not None
|
||||
self.event.current_stream_wait()
|
||||
|
||||
def __enter__(self) -> Any:
|
||||
"""
|
||||
Utility for overlapping and Python `with` syntax.
|
||||
|
||||
You can overlap the kernels on the current stream with the following example:
|
||||
```python
|
||||
event_overlap = event_after_all_to_all_kernels()
|
||||
with event_overlap():
|
||||
do_something_on_current_stream()
|
||||
# After exiting the `with` scope, the current stream with wait the event to be finished.
|
||||
```
|
||||
"""
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""
|
||||
Utility for overlapping and Python `with` syntax.
|
||||
|
||||
Please follow the example in the `__enter__` function.
|
||||
"""
|
||||
if self.event is not None:
|
||||
self.event.current_stream_wait()
|
||||
|
|
@ -1,42 +0,0 @@
|
|||
import os
|
||||
import subprocess
|
||||
import setuptools
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
include_dirs = ['csrc/']
|
||||
sources = ['csrc/mlx5gda.cpp', 'csrc/mxa_ep.cpp', 'csrc/mxa_kernel.cu']
|
||||
libraries = ['ibverbs', 'mlx5']
|
||||
extra_link_args = ['-libverbs', '-Wl,--no-as-needed', '-L/usr/lib/aarch64-linux-gnu', '-lcuda']
|
||||
|
||||
# noinspection PyBroadException
|
||||
try:
|
||||
cmd = ['git', 'rev-parse', '--short', 'HEAD']
|
||||
revision = '+' + subprocess.check_output(cmd).decode('ascii').rstrip()
|
||||
except Exception as _:
|
||||
revision = ''
|
||||
|
||||
setuptools.setup(
|
||||
name='mxa_ep',
|
||||
version='1.0.0' + revision,
|
||||
packages=setuptools.find_packages(
|
||||
include=['mxa_ep']
|
||||
),
|
||||
ext_modules=[
|
||||
CUDAExtension(
|
||||
name='mxa_ep_cpp',
|
||||
include_dirs=include_dirs,
|
||||
sources=sources,
|
||||
extra_compile_args={
|
||||
'cxx': ['-DTORCH_USE_CUDA_DSA'],
|
||||
'nvcc': ['-DTORCH_USE_CUDA_DSA'],
|
||||
},
|
||||
extra_link_args=extra_link_args,
|
||||
libraries=libraries
|
||||
)
|
||||
],
|
||||
cmdclass={
|
||||
'build_ext': BuildExtension
|
||||
}
|
||||
)
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
add_library(mooncake_ep mooncake_backend.cpp mooncake_ep_buffer.cpp mooncake_ep_kernel.cu mooncake_worker.cu mooncake_worker_thread.cpp mooncake_ibgda/mlx5gda.cpp)
|
||||
|
||||
set_target_properties(mooncake_ep PROPERTIES POSITION_INDEPENDENT_CODE ON)
|
||||
target_link_libraries(mooncake_ep PUBLIC ${TORCH_LIBRARIES} transfer_engine ibverbs mlx5)
|
||||
|
|
@ -0,0 +1,414 @@
|
|||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <torch/torch.h>
|
||||
#include <torch/csrc/distributed/c10d/Backend.hpp>
|
||||
#include <mooncake_backend.h>
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
constexpr const char* REGISTER_BUFFER_ERROR_MSG =
|
||||
"Failed to register local memory.";
|
||||
constexpr const char* MULTI_DEVICE_ERROR_MSG =
|
||||
"Expecting one tensor only but got multiple.";
|
||||
constexpr const char* SYNC_OP_ERROR_MSG = "Expecting async op but got sync op.";
|
||||
constexpr const char* REDUCE_OP_ERROR_MSG = "Only support SUM.";
|
||||
constexpr const char* SPARSE_ERROR_MSG = "Sparse op not supported.";
|
||||
constexpr const char* REDUCE_DTYPE_ERROR_MSG = "Unsupported reduce dtype: ";
|
||||
|
||||
std::string MooncakeBackend::hostIp_ = "127.0.0.1";
|
||||
TransferEngine MooncakeBackend::engine_ = TransferEngine(true);
|
||||
Transport* MooncakeBackend::transport_ = nullptr;
|
||||
int MooncakeBackend::backendIndex_ = 0;
|
||||
MooncakeWorker MooncakeBackend::worker_;
|
||||
|
||||
MooncakeBackend::MooncakeBackend(
|
||||
c10::intrusive_ptr<::c10d::Store> store, int rank, int size,
|
||||
c10::intrusive_ptr<MooncakeBackendOptions> options, bool isCpu)
|
||||
: Backend(rank, size), isCpu_(isCpu) {
|
||||
// Get device data
|
||||
int deviceId_;
|
||||
cudaError err = cudaGetDevice(&deviceId_);
|
||||
TORCH_CHECK(!err, c10::str("Failed to get device id"));
|
||||
|
||||
// Initialize transfer engine
|
||||
if (!transport_) {
|
||||
engine_.init(P2PHANDSHAKE, hostIp_);
|
||||
transport_ = engine_.installTransport("rdma", nullptr);
|
||||
TORCH_CHECK(transport_ != nullptr,
|
||||
c10::str("Failed to install transport"));
|
||||
}
|
||||
auto localRpcMeta = transport_->meta()->localRpcMeta();
|
||||
std::string localServerName = localRpcMeta.ip_or_host_name + ":" +
|
||||
std::to_string(localRpcMeta.rpc_port);
|
||||
|
||||
// Register buffers
|
||||
std::string location = "cuda:" + std::to_string(deviceId_);
|
||||
if (isCpu) {
|
||||
for (size_t i = 0; i < 2; i++) {
|
||||
send_buffer_[i] = malloc(kBufferSize);
|
||||
TORCH_CHECK(send_buffer_[i],
|
||||
c10::str("Failed to allocate CPU send buffer"));
|
||||
|
||||
int rc = engine_.registerLocalMemory(send_buffer_[i], kBufferSize,
|
||||
location);
|
||||
TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < 2; i++) {
|
||||
recv_buffer_[i] = malloc(kBufferSize);
|
||||
TORCH_CHECK(recv_buffer_[i],
|
||||
c10::str("Failed to allocate CPU recv buffer"));
|
||||
|
||||
int rc = engine_.registerLocalMemory(recv_buffer_[i], kBufferSize,
|
||||
location);
|
||||
TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG);
|
||||
}
|
||||
} else {
|
||||
for (size_t i = 0; i < 2; i++) {
|
||||
err = cudaMalloc(&send_buffer_[i], kBufferSize);
|
||||
TORCH_CHECK(!err, c10::str("Failed to allocate CUDA send buffer"));
|
||||
|
||||
int rc = engine_.registerLocalMemory(send_buffer_[i], kBufferSize,
|
||||
location);
|
||||
TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < 2; i++) {
|
||||
err = cudaMalloc(&recv_buffer_[i], kBufferSize);
|
||||
TORCH_CHECK(!err, c10::str("Failed to allocate CUDA recv buffer"));
|
||||
|
||||
int rc = engine_.registerLocalMemory(recv_buffer_[i], kBufferSize,
|
||||
location);
|
||||
TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG);
|
||||
}
|
||||
}
|
||||
|
||||
// Register CPU sync regions
|
||||
TORCH_CHECK(size <= kMaxNumRanks, "The number of ranks exceeds the limit.");
|
||||
for (size_t i = 0; i < 2; i++) {
|
||||
cpu_sync_send_region_[i] = new int32_t[kMaxNumRanks];
|
||||
int rc = engine_.registerLocalMemory(
|
||||
cpu_sync_send_region_[i], kMaxNumRanks * sizeof(int32_t), location);
|
||||
TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < 2; i++) {
|
||||
cpu_sync_recv_region_[i] = new int32_t[kMaxNumRanks];
|
||||
int rc = engine_.registerLocalMemory(
|
||||
cpu_sync_recv_region_[i], kMaxNumRanks * sizeof(int32_t), location);
|
||||
TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG);
|
||||
}
|
||||
|
||||
// Sync metadata
|
||||
store->set("server_name_" + std::to_string(rank), localServerName);
|
||||
|
||||
std::vector<std::string> server_names;
|
||||
for (int i = 0; i < size; i++) {
|
||||
server_names.push_back(
|
||||
store->get_to_str({"server_name_" + std::to_string(i)}));
|
||||
}
|
||||
|
||||
meta_.rank = rank;
|
||||
meta_.size = size;
|
||||
meta_.taskCount = 0;
|
||||
cudaHostAlloc(&meta_.activeRanks, kMaxNumRanks * sizeof(bool),
|
||||
cudaHostAllocMapped);
|
||||
cudaHostGetDevicePointer(&meta_.activeRanksDevice, meta_.activeRanks, 0);
|
||||
for (size_t i = 0; i < kMaxNumRanks; ++i) {
|
||||
meta_.activeRanks[i] = true;
|
||||
}
|
||||
if (options) {
|
||||
TORCH_CHECK(options->activeRanks_.dtype() == at::kInt,
|
||||
"activeRanks must be int.");
|
||||
if (isCpu) {
|
||||
TORCH_CHECK(options->activeRanks_.device().is_cpu(),
|
||||
"activeRanks must be on CPU.");
|
||||
} else {
|
||||
TORCH_CHECK(options->activeRanks_.device().is_cuda(),
|
||||
"activeRanks must be on CUDA.");
|
||||
}
|
||||
meta_.activeRanksTensor = options->activeRanks_;
|
||||
} else {
|
||||
meta_.activeRanksTensor =
|
||||
at::ones({size}, torch::dtype(torch::kInt32)
|
||||
.device(isCpu ? torch::kCPU : torch::kCUDA));
|
||||
}
|
||||
meta_.engine = &engine_;
|
||||
meta_.bufferBaseIndex = backendIndex_ * 8;
|
||||
for (int i = 0; i < size_; ++i) {
|
||||
auto segment_id = engine_.openSegment(server_names[i]);
|
||||
meta_.segmentIDs.emplace_back(segment_id);
|
||||
auto segment_desc =
|
||||
engine_.getMetadata()->getSegmentDescByID(segment_id, true);
|
||||
meta_.segmentDescs.emplace_back(segment_desc);
|
||||
}
|
||||
|
||||
// Let the default process group warm up the transfer engine
|
||||
if (backendIndex_ == 0) {
|
||||
std::vector<TransferRequest> entries;
|
||||
for (int i = rank_; i < size_; ++i) {
|
||||
entries.push_back(TransferRequest{
|
||||
.opcode = TransferRequest::READ,
|
||||
.source =
|
||||
(int32_t*)meta_.segmentDescs[rank_]->buffers[4].addr + 1,
|
||||
.target_id = meta_.segmentIDs[i],
|
||||
.target_offset = meta_.segmentDescs[i]->buffers[6].addr,
|
||||
.length = sizeof(int32_t),
|
||||
});
|
||||
}
|
||||
auto batchID = engine_.allocateBatchID(entries.size());
|
||||
engine_.submitTransfer(batchID, entries);
|
||||
|
||||
while (true) {
|
||||
bool batch_done = true;
|
||||
TransferStatus status;
|
||||
for (int i = 0; i < size_ - rank_; ++i) {
|
||||
engine_.getTransferStatus(batchID, i, status);
|
||||
if (status.s != TransferStatusEnum::COMPLETED &&
|
||||
status.s != TransferStatusEnum::FAILED) {
|
||||
batch_done = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (batch_done) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
store->set("warmup_done_" + std::to_string(rank_), "1");
|
||||
for (int i = 0; i < size_; i++) {
|
||||
store->get_to_str("warmup_done_" + std::to_string(i));
|
||||
}
|
||||
}
|
||||
|
||||
// Increment backend index
|
||||
++backendIndex_;
|
||||
}
|
||||
|
||||
MooncakeBackend::~MooncakeBackend() {
|
||||
for (size_t i = 0; i < 2; i++) {
|
||||
delete[] cpu_sync_send_region_[i];
|
||||
delete[] cpu_sync_recv_region_[i];
|
||||
if (isCpu_) {
|
||||
free(send_buffer_[i]);
|
||||
free(recv_buffer_[i]);
|
||||
} else {
|
||||
cudaFree(send_buffer_[i]);
|
||||
cudaFree(recv_buffer_[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const std::string MooncakeBackend::getBackendName() const { return "mooncake"; }
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> MooncakeBackend::broadcast(
|
||||
std::vector<at::Tensor>& tensors, const c10d::BroadcastOptions& opts) {
|
||||
TORCH_CHECK(tensors.size() == 1, MULTI_DEVICE_ERROR_MSG);
|
||||
auto tensor = tensors.back();
|
||||
size_t tensorSize = tensor.numel() * tensor.element_size();
|
||||
int64_t root = opts.rootRank + opts.rootTensor;
|
||||
bool isRoot = (root == rank_);
|
||||
if (isCpu_) {
|
||||
return worker_.putTaskCpu(
|
||||
c10d::OpType::BROADCAST, tensorSize, root, &meta_,
|
||||
[=](void* dst) {
|
||||
if (isRoot) {
|
||||
memcpy(dst, tensor.data_ptr(), tensorSize);
|
||||
}
|
||||
},
|
||||
[=](void* src) { memcpy(tensor.data_ptr(), src, tensorSize); });
|
||||
} else {
|
||||
at::cuda::CUDAStream stream =
|
||||
at::cuda::getCurrentCUDAStream(tensor.device().index());
|
||||
return worker_.putTaskCuda(
|
||||
c10d::OpType::BROADCAST, tensorSize, root, &meta_, stream,
|
||||
[&](void* dst) {
|
||||
if (isRoot) {
|
||||
cudaMemcpyAsync(dst, tensor.data_ptr(), tensorSize,
|
||||
cudaMemcpyHostToDevice, stream);
|
||||
}
|
||||
},
|
||||
[&](void* src) {
|
||||
cudaMemcpyAsync(tensor.data_ptr(), src, tensorSize,
|
||||
cudaMemcpyDeviceToHost, stream);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> MooncakeBackend::allreduce(
|
||||
std::vector<at::Tensor>& tensors, const c10d::AllreduceOptions& opts) {
|
||||
TORCH_CHECK(tensors.size() == 1, MULTI_DEVICE_ERROR_MSG);
|
||||
TORCH_CHECK(opts.sparseIndices == std::nullopt, SPARSE_ERROR_MSG);
|
||||
auto tensor = tensors.back();
|
||||
size_t tensorSize = tensor.numel() * tensor.element_size();
|
||||
if (isCpu_) {
|
||||
auto numRanks = size_;
|
||||
return worker_.putTaskCpu(
|
||||
c10d::OpType::ALLREDUCE, tensorSize, 0, &meta_,
|
||||
[=](void* dst) { memcpy(dst, tensor.data_ptr(), tensorSize); },
|
||||
[=](void* src) {
|
||||
memset(tensor.data_ptr(), 0, tensorSize);
|
||||
launchReduceCpu(tensor, src, numRanks, opts.reduceOp);
|
||||
});
|
||||
} else {
|
||||
auto stream = at::cuda::getCurrentCUDAStream(tensor.device().index());
|
||||
return worker_.putTaskCuda(
|
||||
c10d::OpType::ALLREDUCE, tensorSize, 0, &meta_, stream,
|
||||
[&](void* dst) {
|
||||
cudaMemcpyAsync(dst, tensor.data_ptr(), tensorSize,
|
||||
cudaMemcpyHostToDevice, stream);
|
||||
},
|
||||
[&](void* src) {
|
||||
cudaMemsetAsync(tensor.data_ptr(), 0, tensorSize, stream);
|
||||
launchReduceKernel(tensor, src, size_, opts.reduceOp,
|
||||
meta_.activeRanksDevice, stream);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> MooncakeBackend::allgather(
|
||||
std::vector<std::vector<at::Tensor>>& outputTensors,
|
||||
std::vector<at::Tensor>& inputTensors, const c10d::AllgatherOptions& opts) {
|
||||
TORCH_CHECK(inputTensors.size() == 1, MULTI_DEVICE_ERROR_MSG);
|
||||
TORCH_CHECK(outputTensors.size() == 1, MULTI_DEVICE_ERROR_MSG);
|
||||
auto inputTensor = inputTensors.back();
|
||||
auto outputTensors_ = outputTensors.back();
|
||||
size_t tensorSize = inputTensor.numel() * inputTensor.element_size();
|
||||
if (isCpu_) {
|
||||
return worker_.putTaskCpu(
|
||||
c10d::OpType::ALLGATHER, tensorSize, 0, &meta_,
|
||||
[=](void* dst) { memcpy(dst, inputTensor.data_ptr(), tensorSize); },
|
||||
[=](void* src) {
|
||||
for (const auto j : c10::irange(outputTensors_.size())) {
|
||||
memcpy(outputTensors_[j].data_ptr(), src + j * tensorSize,
|
||||
tensorSize);
|
||||
}
|
||||
});
|
||||
} else {
|
||||
auto stream =
|
||||
at::cuda::getCurrentCUDAStream(inputTensor.device().index());
|
||||
return worker_.putTaskCuda(
|
||||
c10d::OpType::ALLGATHER, tensorSize, 0, &meta_, stream,
|
||||
[&](void* dst) {
|
||||
cudaMemcpyAsync(dst, inputTensor.data_ptr(), tensorSize,
|
||||
cudaMemcpyHostToDevice, stream);
|
||||
},
|
||||
[&](void* src) {
|
||||
for (const auto j : c10::irange(outputTensors_.size())) {
|
||||
cudaMemcpyAsync(outputTensors_[j].data_ptr(),
|
||||
src + j * tensorSize, tensorSize,
|
||||
cudaMemcpyDeviceToHost, stream);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> MooncakeBackend::_allgather_base(
|
||||
at::Tensor& outputBuffer, at::Tensor& inputBuffer,
|
||||
const c10d::AllgatherOptions& opts) {
|
||||
size_t tensorSize = inputBuffer.numel() * inputBuffer.element_size();
|
||||
if (isCpu_) {
|
||||
auto numRanks = size_;
|
||||
return worker_.putTaskCpu(
|
||||
c10d::OpType::_ALLGATHER_BASE, tensorSize, 0, &meta_,
|
||||
[=](void* dst) { memcpy(dst, inputBuffer.data_ptr(), tensorSize); },
|
||||
[=](void* src) {
|
||||
memcpy(outputBuffer.data_ptr(), src, tensorSize * numRanks);
|
||||
});
|
||||
} else {
|
||||
auto stream =
|
||||
at::cuda::getCurrentCUDAStream(inputBuffer.device().index());
|
||||
return worker_.putTaskCuda(
|
||||
c10d::OpType::_ALLGATHER_BASE, tensorSize, 0, &meta_, stream,
|
||||
[&](void* dst) {
|
||||
cudaMemcpyAsync(dst, inputBuffer.data_ptr(), tensorSize,
|
||||
cudaMemcpyHostToDevice, stream);
|
||||
},
|
||||
[&](void* src) {
|
||||
cudaMemcpyAsync(outputBuffer.data_ptr(), src,
|
||||
tensorSize * size_, cudaMemcpyDeviceToHost,
|
||||
stream);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> MooncakeBackend::_reduce_scatter_base(
|
||||
at::Tensor& outputBuffer, at::Tensor& inputBuffer,
|
||||
const c10d::ReduceScatterOptions& opts) {
|
||||
size_t tensorSize = outputBuffer.numel() * outputBuffer.element_size();
|
||||
if (isCpu_) {
|
||||
auto numRanks = size_;
|
||||
return worker_.putTaskCpu(
|
||||
c10d::OpType::REDUCE_SCATTER, tensorSize, 0, &meta_,
|
||||
[=](void* dst) {
|
||||
memcpy(dst, inputBuffer.data_ptr(), tensorSize * numRanks);
|
||||
},
|
||||
[=](void* src) {
|
||||
memset(outputBuffer.data_ptr(), 0, tensorSize);
|
||||
launchReduceCpu(outputBuffer, src, numRanks, opts.reduceOp);
|
||||
});
|
||||
} else {
|
||||
auto stream =
|
||||
at::cuda::getCurrentCUDAStream(inputBuffer.device().index());
|
||||
return worker_.putTaskCuda(
|
||||
c10d::OpType::REDUCE_SCATTER, tensorSize, 0, &meta_, stream,
|
||||
[&](void* dst) {
|
||||
cudaMemcpyAsync(dst, inputBuffer.data_ptr(), tensorSize * size_,
|
||||
cudaMemcpyHostToDevice, stream);
|
||||
},
|
||||
[&](void* src) {
|
||||
cudaMemsetAsync(outputBuffer.data_ptr(), 0, tensorSize, stream);
|
||||
launchReduceKernel(outputBuffer, src, size_, opts.reduceOp,
|
||||
meta_.activeRanksDevice, stream);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> MooncakeBackend::alltoall(
|
||||
std::vector<at::Tensor>& outputTensors,
|
||||
std::vector<at::Tensor>& inputTensors, const c10d::AllToAllOptions& opts) {
|
||||
size_t tensorSize =
|
||||
inputTensors[0].numel() * inputTensors[0].element_size();
|
||||
if (isCpu_) {
|
||||
return worker_.putTaskCpu(
|
||||
c10d::OpType::ALLTOALL, tensorSize, 0, &meta_,
|
||||
[=](void* dst) {
|
||||
for (const auto j : c10::irange(inputTensors.size())) {
|
||||
memcpy(dst + j * tensorSize, inputTensors[j].data_ptr(),
|
||||
tensorSize);
|
||||
}
|
||||
},
|
||||
[=](void* src) {
|
||||
for (const auto j : c10::irange(outputTensors.size())) {
|
||||
memcpy(outputTensors[j].data_ptr(), src + j * tensorSize,
|
||||
tensorSize);
|
||||
}
|
||||
});
|
||||
} else {
|
||||
auto stream =
|
||||
at::cuda::getCurrentCUDAStream(inputTensors[0].device().index());
|
||||
return worker_.putTaskCuda(
|
||||
c10d::OpType::ALLTOALL, tensorSize, 0, &meta_, stream,
|
||||
[&](void* dst) {
|
||||
for (const auto j : c10::irange(inputTensors.size())) {
|
||||
cudaMemcpyAsync(dst + j * tensorSize,
|
||||
inputTensors[j].data_ptr(), tensorSize,
|
||||
cudaMemcpyHostToDevice, stream);
|
||||
}
|
||||
},
|
||||
[&](void* src) {
|
||||
for (const auto j : c10::irange(outputTensors.size())) {
|
||||
cudaMemcpyAsync(outputTensors[j].data_ptr(),
|
||||
src + j * tensorSize, tensorSize,
|
||||
cudaMemcpyDeviceToHost, stream);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
c10::intrusive_ptr<c10d::Work> MooncakeBackend::barrier(
|
||||
const c10d::BarrierOptions& opts) {
|
||||
TORCH_CHECK(isCpu_, "Barrier is available only for CPU.")
|
||||
return worker_.putTaskCpu(
|
||||
c10d::OpType::BARRIER, 0, 0, &meta_, [=](void*) {}, [=](void*) {});
|
||||
}
|
||||
} // namespace mooncake
|
||||
|
|
@ -0,0 +1,426 @@
|
|||
#include <mooncake_ep_buffer.h>
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
MooncakeEpBuffer::MooncakeEpBuffer(int rank, int num_ranks,
|
||||
int64_t num_ep_buffer_bytes, int nic_id)
|
||||
: rank(rank),
|
||||
num_ranks(num_ranks),
|
||||
num_ep_buffer_bytes(num_ep_buffer_bytes),
|
||||
nic_id(nic_id),
|
||||
comm_stream(at::cuda::getStreamFromPool(true)) {
|
||||
// Get ranks
|
||||
CUDA_CHECK(cudaGetDevice(&device_id));
|
||||
CUDA_CHECK(cudaDeviceGetAttribute(&clock_rate_khz, cudaDevAttrClockRate,
|
||||
device_id));
|
||||
CUDA_CHECK(cudaMalloc(&gdr_buffer, num_ep_buffer_bytes));
|
||||
CUDA_CHECK(cudaMalloc(&raddrs, num_ranks * sizeof(uint64_t)));
|
||||
CUDA_CHECK(cudaMalloc(&rkeys, num_ranks * sizeof(uint32_t)));
|
||||
CUDA_CHECK(
|
||||
cudaMalloc(&qp_devctxs, MAX_QP_COUNT * sizeof(mlx5gda_qp_devctx)));
|
||||
init_ibgda();
|
||||
|
||||
// Create 32 MiB workspace
|
||||
CUDA_CHECK(cudaMalloc(&workspace, NUM_WORKSPACE_BYTES));
|
||||
CUDA_CHECK(cudaMemsetAsync(workspace, 0, NUM_WORKSPACE_BYTES, comm_stream));
|
||||
}
|
||||
|
||||
MooncakeEpBuffer::~MooncakeEpBuffer() noexcept(false) {
|
||||
cudaFree(gdr_buffer);
|
||||
cudaFree(raddrs);
|
||||
cudaFree(rkeys);
|
||||
cudaFree(qp_devctxs);
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>, torch::Tensor,
|
||||
torch::Tensor, torch::Tensor, std::optional<EventHandle>,
|
||||
std::optional<std::function<void()>>>
|
||||
MooncakeEpBuffer::dispatch(const torch::Tensor& x,
|
||||
const torch::Tensor& topk_idx,
|
||||
torch::Tensor& active_ranks,
|
||||
int num_max_dispatch_tokens_per_rank,
|
||||
int num_experts, int timeout_us, bool use_fp8,
|
||||
bool async, bool return_recv_hook) {
|
||||
// Tensor checks
|
||||
// By default using `ptp128c` FP8 cast
|
||||
EP_HOST_ASSERT(x.dim() == 2 and x.is_contiguous() and
|
||||
x.scalar_type() == torch::kBFloat16);
|
||||
EP_HOST_ASSERT(x.size(1) % sizeof(int4) == 0 and x.size(1) % 128 == 0);
|
||||
EP_HOST_ASSERT(topk_idx.dim() == 2 and topk_idx.is_contiguous());
|
||||
EP_HOST_ASSERT(x.size(0) == topk_idx.size(0) and
|
||||
x.size(0) <= num_max_dispatch_tokens_per_rank);
|
||||
EP_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt64);
|
||||
EP_HOST_ASSERT(num_experts % num_ranks == 0);
|
||||
EP_HOST_ASSERT(MAX_QP_COUNT % num_ranks == 0);
|
||||
|
||||
auto num_tokens = static_cast<int>(x.size(0)),
|
||||
hidden = static_cast<int>(x.size(1));
|
||||
auto num_scales = hidden / 128,
|
||||
num_topk = static_cast<int>(topk_idx.size(1));
|
||||
int num_local_experts = num_experts / num_ranks;
|
||||
|
||||
// Buffer control
|
||||
BufferPair layout(gdr_buffer, num_max_dispatch_tokens_per_rank, hidden,
|
||||
num_ranks, num_experts);
|
||||
EP_HOST_ASSERT(layout.total_bytes <= num_ep_buffer_bytes);
|
||||
auto buffer = layout.buffers[buffer_idx];
|
||||
auto next_buffer = layout.buffers[buffer_idx ^= 1];
|
||||
|
||||
// Wait previous tasks to be finished
|
||||
// NOTES: the hook mode will always use the default stream
|
||||
auto compute_stream = at::cuda::getCurrentCUDAStream();
|
||||
auto launch_stream = return_recv_hook ? compute_stream : comm_stream;
|
||||
EP_HOST_ASSERT(not(async and return_recv_hook));
|
||||
if (not return_recv_hook) stream_wait(launch_stream, compute_stream);
|
||||
|
||||
// Allocate packed tensors
|
||||
auto packed_recv_x = torch::empty(
|
||||
{num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank,
|
||||
hidden},
|
||||
x.options().dtype(use_fp8 ? torch::kFloat8_e4m3fn : torch::kBFloat16));
|
||||
auto packed_recv_src_info = torch::empty(
|
||||
{num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank},
|
||||
torch::dtype(torch::kInt32).device(torch::kCUDA));
|
||||
auto packed_recv_layout_range =
|
||||
torch::empty({num_local_experts, num_ranks},
|
||||
torch::dtype(torch::kInt64).device(torch::kCUDA));
|
||||
auto packed_recv_count = torch::zeros(
|
||||
{num_local_experts}, torch::dtype(torch::kInt32).device(torch::kCUDA));
|
||||
|
||||
// Allocate column-majored scales
|
||||
auto packed_recv_x_scales = std::optional<torch::Tensor>();
|
||||
float* packed_recv_x_scales_ptr = nullptr;
|
||||
if (use_fp8) {
|
||||
EP_HOST_ASSERT((num_ranks * num_max_dispatch_tokens_per_rank) % 4 ==
|
||||
0 and
|
||||
"TMA requires the number of tokens to be multiple of 4");
|
||||
packed_recv_x_scales =
|
||||
torch::empty({num_local_experts, num_scales,
|
||||
num_ranks * num_max_dispatch_tokens_per_rank},
|
||||
torch::dtype(torch::kFloat32).device(torch::kCUDA));
|
||||
packed_recv_x_scales =
|
||||
torch::transpose(packed_recv_x_scales.value(), 1, 2);
|
||||
packed_recv_x_scales_ptr = packed_recv_x_scales->data_ptr<float>();
|
||||
}
|
||||
|
||||
int64_t timeout_ticks =
|
||||
timeout_us == -1 ? -1
|
||||
: (int64_t)clock_rate_khz * (int64_t)timeout_us / 1000;
|
||||
|
||||
auto launcher = [=](int phases) {
|
||||
mooncake::dispatch(
|
||||
packed_recv_x.data_ptr(), packed_recv_x_scales_ptr,
|
||||
packed_recv_src_info.data_ptr<int>(),
|
||||
packed_recv_layout_range.data_ptr<int64_t>(),
|
||||
packed_recv_count.data_ptr<int>(), active_ranks.data_ptr<int32_t>(),
|
||||
gdr_buffer, buffer.rdma_send_signal_buffer,
|
||||
buffer.rdma_recv_signal_buffer, buffer.rdma_send_data_buffer,
|
||||
buffer.rdma_recv_data_buffer, nullptr, nullptr, raddrs, rkeys,
|
||||
qp_devctxs, x.data_ptr(), topk_idx.data_ptr<int64_t>(),
|
||||
next_buffer.rdma_recv_signal_buffer, num_tokens, hidden,
|
||||
num_max_dispatch_tokens_per_rank, num_topk, num_experts, rank,
|
||||
num_ranks, use_fp8, workspace, launch_stream, timeout_ticks,
|
||||
phases);
|
||||
};
|
||||
launcher(return_recv_hook
|
||||
? LOW_LATENCY_SEND_PHASE
|
||||
: (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE));
|
||||
|
||||
// Wait streams
|
||||
std::optional<EventHandle> event;
|
||||
if (async) {
|
||||
// NOTES: we must ensure the all tensors will not be deallocated
|
||||
// before the stream-wait happens, so in Python API, we must wrap
|
||||
// all tensors into the event handle.
|
||||
event = EventHandle(launch_stream);
|
||||
} else if (not return_recv_hook) {
|
||||
stream_wait(compute_stream, launch_stream);
|
||||
}
|
||||
|
||||
// Receiver callback
|
||||
std::optional<std::function<void()>> recv_hook = std::nullopt;
|
||||
if (return_recv_hook)
|
||||
recv_hook = [=]() { launcher(LOW_LATENCY_RECV_PHASE); };
|
||||
|
||||
// Return values
|
||||
return {packed_recv_x,
|
||||
packed_recv_x_scales,
|
||||
packed_recv_count,
|
||||
packed_recv_src_info,
|
||||
packed_recv_layout_range,
|
||||
event,
|
||||
recv_hook};
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<EventHandle>,
|
||||
std::optional<std::function<void()>>>
|
||||
MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx,
|
||||
const torch::Tensor& topk_weights,
|
||||
const torch::Tensor& src_info,
|
||||
const torch::Tensor& layout_range,
|
||||
torch::Tensor& active_ranks,
|
||||
int num_max_dispatch_tokens_per_rank, int num_experts,
|
||||
int timeout_us, bool zero_copy, bool async,
|
||||
bool return_recv_hook,
|
||||
const std::optional<torch::Tensor>& out) {
|
||||
// Tensor checks
|
||||
EP_HOST_ASSERT(x.dim() == 3 and x.is_contiguous() and
|
||||
x.scalar_type() == torch::kBFloat16);
|
||||
EP_HOST_ASSERT(x.size(0) == num_experts / num_ranks);
|
||||
EP_HOST_ASSERT(x.size(1) == num_ranks * num_max_dispatch_tokens_per_rank);
|
||||
EP_HOST_ASSERT(x.size(2) % sizeof(int4) == 0 and x.size(2) % 128 == 0);
|
||||
EP_HOST_ASSERT(topk_idx.dim() == 2 and topk_idx.is_contiguous());
|
||||
EP_HOST_ASSERT(topk_idx.size(0) == topk_weights.size(0) and
|
||||
topk_idx.size(1) == topk_weights.size(1));
|
||||
EP_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt64);
|
||||
EP_HOST_ASSERT(topk_weights.dim() == 2 and topk_weights.is_contiguous());
|
||||
EP_HOST_ASSERT(topk_weights.size(0) <= num_max_dispatch_tokens_per_rank);
|
||||
EP_HOST_ASSERT(topk_weights.scalar_type() == torch::kFloat32);
|
||||
EP_HOST_ASSERT(src_info.dim() == 2 and src_info.is_contiguous());
|
||||
EP_HOST_ASSERT(src_info.scalar_type() == torch::kInt32 and
|
||||
x.size(0) == src_info.size(0));
|
||||
EP_HOST_ASSERT(layout_range.dim() == 2 and layout_range.is_contiguous());
|
||||
EP_HOST_ASSERT(layout_range.scalar_type() == torch::kInt64);
|
||||
EP_HOST_ASSERT(layout_range.size(0) == num_experts / num_ranks and
|
||||
layout_range.size(1) == num_ranks);
|
||||
auto hidden = static_cast<int>(x.size(2));
|
||||
auto num_local_experts = num_experts / num_ranks,
|
||||
num_topk = static_cast<int>(topk_weights.size(1));
|
||||
auto num_combined_tokens = static_cast<int>(topk_weights.size(0));
|
||||
|
||||
// Buffer control
|
||||
BufferPair layout(gdr_buffer, num_max_dispatch_tokens_per_rank, hidden,
|
||||
num_ranks, num_experts);
|
||||
EP_HOST_ASSERT(layout.total_bytes <= num_ep_buffer_bytes);
|
||||
auto buffer = layout.buffers[buffer_idx];
|
||||
auto next_buffer = layout.buffers[buffer_idx ^= 1];
|
||||
|
||||
// Wait previous tasks to be finished
|
||||
// NOTES: the hook mode will always use the default stream
|
||||
auto compute_stream = at::cuda::getCurrentCUDAStream();
|
||||
auto launch_stream = return_recv_hook ? compute_stream : comm_stream;
|
||||
EP_HOST_ASSERT(not(async and return_recv_hook));
|
||||
if (not return_recv_hook) stream_wait(launch_stream, compute_stream);
|
||||
|
||||
// Allocate output tensor
|
||||
torch::Tensor combined_x;
|
||||
if (out.has_value()) {
|
||||
EP_HOST_ASSERT(out->dim() == 2 and out->is_contiguous());
|
||||
EP_HOST_ASSERT(out->size(0) == num_combined_tokens and
|
||||
out->size(1) == hidden);
|
||||
EP_HOST_ASSERT(out->scalar_type() == x.scalar_type());
|
||||
combined_x = out.value();
|
||||
} else {
|
||||
combined_x = torch::empty({num_combined_tokens, hidden}, x.options());
|
||||
}
|
||||
|
||||
int64_t timeout_ticks =
|
||||
timeout_us == -1 ? -1
|
||||
: (int64_t)clock_rate_khz * (int64_t)timeout_us / 1000;
|
||||
|
||||
// Kernel launch
|
||||
auto launcher = [=](int phases) {
|
||||
mooncake::combine(
|
||||
combined_x.data_ptr(), active_ranks.data_ptr<int32_t>(), gdr_buffer,
|
||||
buffer.rdma_send_signal_buffer, buffer.rdma_recv_signal_buffer,
|
||||
buffer.rdma_send_data_buffer, buffer.rdma_recv_data_buffer, nullptr,
|
||||
nullptr, raddrs, rkeys, qp_devctxs, x.data_ptr(),
|
||||
topk_idx.data_ptr<int64_t>(), topk_weights.data_ptr<float>(),
|
||||
src_info.data_ptr<int>(), layout_range.data_ptr<int64_t>(),
|
||||
next_buffer.rdma_recv_signal_buffer, num_combined_tokens, hidden,
|
||||
num_max_dispatch_tokens_per_rank, num_topk, num_experts, rank,
|
||||
num_ranks, workspace, launch_stream, timeout_ticks, phases,
|
||||
zero_copy);
|
||||
};
|
||||
launcher(return_recv_hook
|
||||
? LOW_LATENCY_SEND_PHASE
|
||||
: (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE));
|
||||
|
||||
// Wait streams
|
||||
std::optional<EventHandle> event;
|
||||
if (async) {
|
||||
// NOTES: we must ensure the all tensors will not be deallocated
|
||||
// before the stream-wait happens, so in Python API, we must wrap
|
||||
// all tensors into the event handle.
|
||||
event = EventHandle(launch_stream);
|
||||
} else if (not return_recv_hook) {
|
||||
stream_wait(compute_stream, launch_stream);
|
||||
}
|
||||
|
||||
// Receiver callback
|
||||
std::optional<std::function<void()>> recv_hook = std::nullopt;
|
||||
if (return_recv_hook)
|
||||
recv_hook = [=]() { launcher(LOW_LATENCY_RECV_PHASE); };
|
||||
|
||||
// Return values
|
||||
return {combined_x, event, recv_hook};
|
||||
}
|
||||
|
||||
torch::Tensor MooncakeEpBuffer::get_next_combine_buffer(
|
||||
int num_max_dispatch_tokens_per_rank, int hidden, int num_experts) {
|
||||
BufferPair layout(gdr_buffer, num_max_dispatch_tokens_per_rank, hidden,
|
||||
num_ranks, num_experts);
|
||||
|
||||
auto buffer = layout.buffers[buffer_idx];
|
||||
auto dtype = torch::kBFloat16;
|
||||
size_t num_bytes_per_combine_msg = hidden * sizeof(nv_bfloat16);
|
||||
auto num_msg_elems = static_cast<int>(num_bytes_per_combine_msg /
|
||||
elementSize(torch::kBFloat16));
|
||||
|
||||
EP_HOST_ASSERT(num_bytes_per_combine_msg % elementSize(torch::kBFloat16) ==
|
||||
0);
|
||||
return torch::from_blob(
|
||||
buffer.rdma_send_data_buffer,
|
||||
{num_experts / num_ranks, num_ranks * num_max_dispatch_tokens_per_rank,
|
||||
hidden},
|
||||
{num_ranks * num_max_dispatch_tokens_per_rank * num_msg_elems,
|
||||
num_msg_elems, 1},
|
||||
torch::TensorOptions().dtype(dtype).device(torch::kCUDA));
|
||||
}
|
||||
|
||||
void MooncakeEpBuffer::init_ibgda() {
|
||||
int num_devices;
|
||||
ibv_device** dev_list = ibv_get_device_list(&num_devices);
|
||||
LOG(INFO) << "[EP] GPU " << device_id << " uses NIC " << nic_id
|
||||
<< " out of " << num_devices << " NIC(s)";
|
||||
ibv_context* ctx = ibv_open_device(dev_list[nic_id]);
|
||||
if (!ctx) {
|
||||
perror("Failed to open device");
|
||||
exit(1);
|
||||
}
|
||||
if (ibv_query_gid(ctx, 1, 3, &gid)) {
|
||||
perror("Failed to query gid");
|
||||
}
|
||||
ibv_free_device_list(dev_list);
|
||||
|
||||
ibv_pd* pd = ibv_alloc_pd(ctx);
|
||||
if (!pd) {
|
||||
perror("Failed to allocate protection domain");
|
||||
exit(1);
|
||||
}
|
||||
mlx5dv_pd mpd;
|
||||
mlx5dv_obj dv_obj = {};
|
||||
dv_obj.pd.in = pd;
|
||||
dv_obj.pd.out = &mpd;
|
||||
if (mlx5dv_init_obj(&dv_obj, MLX5DV_OBJ_PD)) {
|
||||
perror("Failed to initialize mlx5dv object");
|
||||
}
|
||||
mr = ibv_reg_mr(pd, gdr_buffer, num_ep_buffer_bytes,
|
||||
IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ |
|
||||
IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_ATOMIC);
|
||||
if (!mr) {
|
||||
perror("Failed to reg mr");
|
||||
}
|
||||
|
||||
CUDA_CHECK(cudaMalloc(&ctrl_buf, CTRL_BUF_SIZE));
|
||||
CUDA_CHECK(cudaMemset(ctrl_buf, 0, CTRL_BUF_SIZE));
|
||||
mlx5dv_devx_umem* ctrl_buf_umem = mlx5dv_devx_umem_reg(
|
||||
ctx, ctrl_buf, CTRL_BUF_SIZE, IBV_ACCESS_LOCAL_WRITE);
|
||||
if (!ctrl_buf_umem) {
|
||||
perror("Failed to register control buffer as umem");
|
||||
fprintf(stderr,
|
||||
"If the error is `Bad address`, probably because your GPU "
|
||||
"does not support GPUDirect RDMA.\n");
|
||||
exit(1);
|
||||
}
|
||||
memheap* ctrl_buf_heap = memheap_create(CTRL_BUF_SIZE);
|
||||
if (!ctrl_buf_heap) {
|
||||
perror("Failed to create memory heap");
|
||||
exit(1);
|
||||
}
|
||||
for (int i = 0; i < MAX_QP_COUNT; ++i) {
|
||||
mlx5gda_qp* qp = mlx5gda_create_rc_qp(mpd, ctrl_buf, ctrl_buf_umem,
|
||||
ctrl_buf_heap, pd, 16384, 1);
|
||||
if (!qp) {
|
||||
perror("Failed to create QP");
|
||||
exit(1);
|
||||
}
|
||||
is_roce_ = qp->port_attr.link_layer == IBV_LINK_LAYER_ETHERNET;
|
||||
if (mlx5gda_modify_rc_qp_rst2init(qp, 0)) {
|
||||
perror("Failed to mlx5gda_modify_rc_qp_rst2init");
|
||||
exit(1);
|
||||
}
|
||||
mlx5gda_qp_devctx qp_devctx = {
|
||||
.qpn = qp->qpn,
|
||||
.wqeid_mask = qp->num_wqebb - 1,
|
||||
.wq = (mlx5gda_wqebb*)(ctrl_buf + qp->wq_offset),
|
||||
.cq = (mlx5_cqe64*)(ctrl_buf + qp->send_cq->cq_offset),
|
||||
.dbr = (mlx5gda_wq_dbr*)(ctrl_buf + qp->dbr_offset),
|
||||
.bf = (char*)qp->uar->reg_addr,
|
||||
};
|
||||
cudaMemcpy(qp_devctxs + i * sizeof(mlx5gda_qp_devctx), &qp_devctx,
|
||||
sizeof(mlx5gda_qp_devctx), cudaMemcpyHostToDevice);
|
||||
qps.push_back(qp);
|
||||
}
|
||||
}
|
||||
|
||||
void MooncakeEpBuffer::sync_ib(const std::vector<int64_t>& remote_addrs,
|
||||
const std::vector<int32_t>& remote_keys,
|
||||
const std::vector<int32_t>& remote_qpns,
|
||||
const std::vector<int32_t>& remote_lids) {
|
||||
for (int i = 0; i < MAX_QP_COUNT; ++i) {
|
||||
ibv_ah_attr ah_attr = {
|
||||
.dlid = (uint16_t)remote_lids[i],
|
||||
.port_num = 0,
|
||||
};
|
||||
if (mlx5gda_modify_rc_qp_init2rtr(
|
||||
qps[i], ah_attr, (uint32_t)remote_qpns[i], IBV_MTU_4096)) {
|
||||
perror("Failed to mlx5gda_modify_rc_qp_init2rtr");
|
||||
exit(1);
|
||||
}
|
||||
if (mlx5gda_modify_rc_qp_rtr2rts(qps[i])) {
|
||||
perror("Failed to mlx5gda_modify_rc_qp_rtr2rts");
|
||||
exit(1);
|
||||
}
|
||||
}
|
||||
for (int i = 0; i < num_ranks; ++i) {
|
||||
uint64_t raddr =
|
||||
i == rank ? (uint64_t)mr->addr : (uint64_t)remote_addrs[i];
|
||||
cudaMemcpy(raddrs + i * sizeof(uint64_t), &raddr, sizeof(uint64_t),
|
||||
cudaMemcpyHostToDevice);
|
||||
uint32_t rkey = i == rank ? mr->lkey : (uint32_t)remote_keys[i];
|
||||
cudaMemcpy(rkeys + i * sizeof(uint32_t), &rkey, sizeof(uint32_t),
|
||||
cudaMemcpyHostToDevice);
|
||||
}
|
||||
}
|
||||
|
||||
void MooncakeEpBuffer::sync_roce(const std::vector<int64_t>& remote_addrs,
|
||||
const std::vector<int32_t>& remote_keys,
|
||||
const std::vector<int32_t>& remote_qpns,
|
||||
const std::vector<int64_t>& subnet_prefixes,
|
||||
const std::vector<int64_t>& interface_ids) {
|
||||
for (int i = 0; i < MAX_QP_COUNT; ++i) {
|
||||
ibv_gid remote_gid{};
|
||||
remote_gid.global.subnet_prefix =
|
||||
subnet_prefixes[i * num_ranks / MAX_QP_COUNT];
|
||||
remote_gid.global.interface_id =
|
||||
interface_ids[i * num_ranks / MAX_QP_COUNT];
|
||||
ibv_ah_attr ah_attr = {};
|
||||
ah_attr.is_global = 1;
|
||||
ah_attr.grh.dgid = remote_gid;
|
||||
ah_attr.grh.sgid_index = 3;
|
||||
ah_attr.grh.hop_limit = 1;
|
||||
ah_attr.port_num = 1;
|
||||
ah_attr.dlid = qps[i]->port_attr.lid | 0xC000;
|
||||
if (mlx5gda_modify_rc_qp_init2rtr(
|
||||
qps[i], ah_attr, (uint32_t)remote_qpns[i], IBV_MTU_4096)) {
|
||||
perror("Failed to mlx5gda_modify_rc_qp_init2rtr");
|
||||
exit(1);
|
||||
}
|
||||
if (mlx5gda_modify_rc_qp_rtr2rts(qps[i])) {
|
||||
perror("Failed to mlx5gda_modify_rc_qp_rtr2rts");
|
||||
exit(1);
|
||||
}
|
||||
}
|
||||
for (int i = 0; i < num_ranks; ++i) {
|
||||
uint64_t raddr =
|
||||
i == rank ? (uint64_t)mr->addr : (uint64_t)remote_addrs[i];
|
||||
cudaMemcpy(raddrs + i * sizeof(uint64_t), &raddr, sizeof(uint64_t),
|
||||
cudaMemcpyHostToDevice);
|
||||
uint32_t rkey = i == rank ? mr->lkey : (uint32_t)remote_keys[i];
|
||||
cudaMemcpy(rkeys + i * sizeof(uint32_t), &rkey, sizeof(uint32_t),
|
||||
cudaMemcpyHostToDevice);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace mooncake
|
||||
|
|
@ -2,15 +2,13 @@
|
|||
#include <cstdio>
|
||||
#include <cuda/atomic>
|
||||
|
||||
#include "configs.cuh"
|
||||
#include "exception.cuh"
|
||||
#include "launch.cuh"
|
||||
#include "mlx5gda.h"
|
||||
#include "utils.cuh"
|
||||
#include <mooncake_ep_configs.cuh>
|
||||
#include <mooncake_ep_exception.cuh>
|
||||
#include <mooncake_ep_launch.cuh>
|
||||
#include <mooncake_ibgda/mlx5gda.h>
|
||||
#include <mooncake_ep_utils.cuh>
|
||||
|
||||
#define TIMEOUT_TICKS 100000000000l
|
||||
|
||||
namespace mxa_ep {
|
||||
namespace mooncake {
|
||||
|
||||
static __device__ void device_mutex_lock_system(uint32_t *mutex) {
|
||||
cuda::atomic_ref<uint32_t, cuda::thread_scope_system> lock(*mutex);
|
||||
|
|
@ -154,7 +152,7 @@ template <bool kUseFP8, int kNumWarpGroups, int kNumWarpsPerGroup, int kHidden>
|
|||
__global__ __launch_bounds__(kNumWarpGroups * kNumWarpsPerGroup * 32, 1) void
|
||||
dispatch(void* packed_recv_x, float* packed_recv_x_scales,
|
||||
int* packed_recv_src_info, int64_t* packed_recv_layout_range,
|
||||
int* packed_recv_count, int32_t* broken_nodes,
|
||||
int* packed_recv_count, int32_t* active_ranks,
|
||||
void* mxa_buffer,
|
||||
int* rdma_send_signal_buffer, int* rdma_recv_signal_buffer,
|
||||
void* rdma_send_data_buffer, void* rdma_recv_data_buffer,
|
||||
|
|
@ -195,6 +193,7 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
|
|||
auto raddr_array = reinterpret_cast<uint64_t*>(raddrs);
|
||||
auto rkey_array = reinterpret_cast<uint32_t*>(rkeys);
|
||||
auto ctx_array = reinterpret_cast<mlx5gda_qp_devctx*>(qp_devctxs);
|
||||
const size_t num_qp_per_rank = MAX_QP_COUNT / num_ranks;
|
||||
|
||||
// Sending phase
|
||||
if ((phases & LOW_LATENCY_SEND_PHASE) == 0)
|
||||
|
|
@ -276,7 +275,7 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
|
|||
if (dst_rank != rank) {
|
||||
if (lane_id == 0) {
|
||||
uint64_t req_rptr_actual = raddr_array[dst_rank] + ((char *)dst_ptr - (char *)(mxa_buffer));
|
||||
auto ctx = ctx_array + dst_rank;
|
||||
auto ctx = ctx_array + dst_rank * num_qp_per_rank + dst_expert_local_idx % num_qp_per_rank;
|
||||
device_mutex_lock_system(&ctx->mutex);
|
||||
__mlx5gda_device_write_rdma_write_wqe(ctx, src_ptr, device_byteswap(rkey_array[rank]), req_rptr_actual, device_byteswap(rkey_array[dst_rank]), num_bytes_per_msg);
|
||||
__mlx5gda_device_post_send_db(ctx);
|
||||
|
|
@ -345,7 +344,7 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
|
|||
if (dst_rank != rank) {
|
||||
uint64_t laddr = (uint64_t)((char *)(raddr_array[rank]) + ((char *)(rdma_send_signal_buffer + dst_expert_local_idx * num_ranks + rank) - (char *)(mxa_buffer)));
|
||||
uint64_t rptr_actual = (uint64_t)((char *)(raddr_array[dst_rank]) + ((char *)(rdma_recv_signal_buffer + dst_expert_local_idx * num_ranks + rank) - (char *)(mxa_buffer)));
|
||||
auto ctx = ctx_array + dst_rank;
|
||||
auto ctx = ctx_array + dst_rank * num_qp_per_rank + dst_expert_local_idx % num_qp_per_rank;
|
||||
device_mutex_lock_system(&ctx->mutex);
|
||||
__mlx5gda_device_write_rdma_atomic_add_wqe(ctx, -num_tokens_sent - 1, laddr, device_byteswap(rkey_array[rank]), rptr_actual, device_byteswap(rkey_array[dst_rank]));
|
||||
__mlx5gda_device_post_send_db(ctx);
|
||||
|
|
@ -397,9 +396,11 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
|
|||
unsigned long long start_time = clock64();
|
||||
while ((num_recv_tokens = ld_acquire_sys_global(rdma_recv_signal_buffer + local_expert_idx * num_ranks + src_rank)) == 0) {
|
||||
unsigned long long end_time = clock64();
|
||||
if ((timeout_ticks != -1 && end_time - start_time > timeout_ticks) || broken_nodes[src_rank]) {
|
||||
if (timeout_ticks != -1 && end_time - start_time > timeout_ticks) {
|
||||
active_ranks[src_rank] = 0;
|
||||
}
|
||||
if (!active_ranks[src_rank]) {
|
||||
num_recv_tokens = -1;
|
||||
broken_nodes[src_rank] = 1;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
|
@ -444,7 +445,7 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
|
|||
|
||||
void dispatch(void* packed_recv_x, float* packed_recv_x_scales,
|
||||
int* packed_recv_src_info, int64_t* packed_recv_layout_range,
|
||||
int* packed_recv_count, int32_t* broken_nodes,
|
||||
int* packed_recv_count, int32_t* active_ranks,
|
||||
void* mxa_buffer,
|
||||
int* rdma_send_signal_buffer, int* rdma_recv_signal_buffer,
|
||||
void* rdma_send_data_buffer, void* rdma_recv_data_buffer,
|
||||
|
|
@ -475,7 +476,7 @@ auto dispatch_func = use_fp8 ? dispatch<true, kNumWarpGroups, kNumWarpsPerGroup,
|
|||
LAUNCH_KERNEL(&cfg, dispatch_func, \
|
||||
packed_recv_x, packed_recv_x_scales, \
|
||||
packed_recv_src_info, packed_recv_layout_range, \
|
||||
packed_recv_count, broken_nodes, \
|
||||
packed_recv_count, active_ranks, \
|
||||
mxa_buffer, \
|
||||
rdma_send_signal_buffer, rdma_recv_signal_buffer, \
|
||||
rdma_send_data_buffer, rdma_recv_data_buffer, \
|
||||
|
|
@ -494,7 +495,7 @@ LAUNCH_KERNEL(&cfg, dispatch_func, \
|
|||
|
||||
template <int kNumWarpGroups, int kNumWarpsPerGroup, int kHidden, int kNumMaxTopk>
|
||||
__global__ __launch_bounds__(kNumWarpGroups * kNumWarpsPerGroup * 32, 1) void
|
||||
combine(void* combined_x, int32_t* gathered_experts,
|
||||
combine(void* combined_x, int32_t* active_ranks,
|
||||
void* mxa_buffer,
|
||||
int* rdma_send_signal_buffer, int* rdma_recv_signal_buffer,
|
||||
void* rdma_send_data_buffer, void* rdma_recv_data_buffer,
|
||||
|
|
@ -531,6 +532,7 @@ combine(void* combined_x, int32_t* gathered_experts,
|
|||
auto raddr_array = reinterpret_cast<uint64_t*>(raddrs);
|
||||
auto rkey_array = reinterpret_cast<uint32_t*>(rkeys);
|
||||
auto ctx_array = reinterpret_cast<mlx5gda_qp_devctx*>(qp_devctxs);
|
||||
const size_t num_qp_per_rank = MAX_QP_COUNT / num_ranks;
|
||||
|
||||
// Sending phase
|
||||
if ((phases & LOW_LATENCY_SEND_PHASE) == 0)
|
||||
|
|
@ -585,7 +587,7 @@ combine(void* combined_x, int32_t* gathered_experts,
|
|||
|
||||
if (lane_id == 0) {
|
||||
uint64_t req_rptr_actual = raddr_array[dst_rank] + ((char *)dst_ptr - (char *)(mxa_buffer));
|
||||
auto ctx = ctx_array + dst_rank;
|
||||
auto ctx = ctx_array + dst_rank * num_qp_per_rank + local_expert_idx % num_qp_per_rank;
|
||||
device_mutex_lock_system(&ctx->mutex);
|
||||
__mlx5gda_device_write_rdma_write_wqe(ctx, (uint64_t) buf_ptr, device_byteswap(rkey_array[rank]), req_rptr_actual, device_byteswap(rkey_array[dst_rank]), num_bytes_per_slot);
|
||||
__mlx5gda_device_post_send_db(ctx);
|
||||
|
|
@ -602,7 +604,7 @@ combine(void* combined_x, int32_t* gathered_experts,
|
|||
if (dst_rank != rank) {
|
||||
uint64_t laddr = (uint64_t)((char *)(raddr_array[rank]) + ((char *)(rdma_send_signal_buffer + global_expert_idx) - (char *)(mxa_buffer)));
|
||||
uint64_t req_rptr_actual = (uint64_t)((char *)(raddr_array[dst_rank]) + ((char *)(rdma_recv_signal_buffer + global_expert_idx) - (char *)(mxa_buffer)));
|
||||
auto ctx = ctx_array + dst_rank;
|
||||
auto ctx = ctx_array + dst_rank * num_qp_per_rank + local_expert_idx % num_qp_per_rank;
|
||||
device_mutex_lock_system(&ctx->mutex);
|
||||
__mlx5gda_device_write_rdma_atomic_add_wqe(ctx, 1, laddr, device_byteswap(rkey_array[rank]), req_rptr_actual, device_byteswap(rkey_array[dst_rank]));
|
||||
__mlx5gda_device_post_send_db(ctx);
|
||||
|
|
@ -622,19 +624,19 @@ combine(void* combined_x, int32_t* gathered_experts,
|
|||
|
||||
// Wait all ranks to arrive
|
||||
if (responsible_expert_idx < num_experts) {
|
||||
const auto src_rank = responsible_expert_idx / num_local_experts;
|
||||
EP_STATIC_ASSERT(kNumWarpsPerGroup > 1, "Invalid number of warps per group");
|
||||
if (sub_warp_id == 0 and lane_id == 0) {
|
||||
unsigned long long start_time = clock64();
|
||||
bool timeout = false;
|
||||
while (ld_acquire_sys_global(rdma_recv_signal_buffer + responsible_expert_idx) == 0) {
|
||||
unsigned long long end_time = clock64();
|
||||
if ((timeout_ticks != -1 && end_time - start_time > timeout_ticks) || gathered_experts[responsible_expert_idx]) {
|
||||
timeout = true;
|
||||
if (timeout_ticks != -1 && end_time - start_time > timeout_ticks) {
|
||||
active_ranks[src_rank] = 0;
|
||||
}
|
||||
if (!active_ranks[src_rank]) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!timeout)
|
||||
gathered_experts[responsible_expert_idx] = 1;
|
||||
}
|
||||
}
|
||||
cooperative_groups::this_grid().sync();
|
||||
|
|
@ -655,7 +657,7 @@ combine(void* combined_x, int32_t* gathered_experts,
|
|||
|
||||
float combined_values[kNumElemsPerInt4] = {0.0f};
|
||||
#pragma unroll
|
||||
for (int i = 0; i < num_topk; ++ i) if (reg_topk_idx[i] >= 0 && ld_acquire_global(gathered_experts + reg_topk_idx[i])) {
|
||||
for (int i = 0; i < num_topk; ++ i) if (reg_topk_idx[i] >= 0) {
|
||||
// Read from sources
|
||||
auto rdma_buffer_type = reinterpret_cast<const int*>(reinterpret_cast<uint8_t*>(rdma_recv_data_buffer) + (reg_topk_idx[i] * num_max_dispatch_tokens_per_rank + token_idx) * num_bytes_per_slot);
|
||||
auto rdma_buffer_row = reinterpret_cast<const uint8_t*>(rdma_buffer_type);
|
||||
|
|
@ -679,7 +681,7 @@ combine(void* combined_x, int32_t* gathered_experts,
|
|||
}
|
||||
}
|
||||
|
||||
void combine(void* combined_x, int32_t* gathered_experts,
|
||||
void combine(void* combined_x, int32_t* active_ranks,
|
||||
void* mxa_buffer,
|
||||
int* rdma_send_signal_buffer, int* rdma_recv_signal_buffer,
|
||||
void* rdma_send_data_buffer, void* rdma_recv_data_buffer,
|
||||
|
|
@ -707,7 +709,7 @@ void combine(void* combined_x, int32_t* gathered_experts,
|
|||
#define COMBINE_LAUNCH_CASE(hidden) { \
|
||||
auto combine_func = combine<kNumWarpGroups, kNumWarpsPerGroup, hidden, kNumMaxTopk>; \
|
||||
LAUNCH_KERNEL(&cfg, combine_func, \
|
||||
combined_x, gathered_experts, \
|
||||
combined_x, active_ranks, \
|
||||
mxa_buffer, \
|
||||
rdma_send_signal_buffer, rdma_recv_signal_buffer, \
|
||||
rdma_send_data_buffer, rdma_recv_data_buffer, \
|
||||
|
|
@ -726,94 +728,4 @@ LAUNCH_KERNEL(&cfg, combine_func, \
|
|||
#undef COMBINE_LAUNCH_CASE
|
||||
}
|
||||
|
||||
__global__ void all_reduce_kernel_without(const int32_t* broken_nodes, int* x, int* mxa_buffer,
|
||||
void* raddrs, void* rkeys, void* qp_devctxs,
|
||||
int size, int rank, int num_ranks) {
|
||||
// IBGDA
|
||||
auto raddr_array = reinterpret_cast<uint64_t*>(raddrs);
|
||||
auto rkey_array = reinterpret_cast<uint32_t*>(rkeys);
|
||||
auto ctx_array = reinterpret_cast<mlx5gda_qp_devctx*>(qp_devctxs);
|
||||
int root = -1;
|
||||
for (int i = 0; i < num_ranks; ++i) {
|
||||
if (!broken_nodes[i]) {
|
||||
root = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
int tid = threadIdx.x;
|
||||
int stride = blockDim.x;
|
||||
|
||||
if (rank == root) {
|
||||
// Receive data from all valid ranks and reduce into x
|
||||
for (int i = 0; i < num_ranks; ++i) {
|
||||
if (i != root && !broken_nodes[i]) {
|
||||
while (ld_acquire_sys_global(mxa_buffer + (1 + size) * i) == 0) ;
|
||||
int* src = mxa_buffer + (1 + size) * i + 1;
|
||||
for (int j = tid; j < size; j += stride) {
|
||||
atomicAdd(&x[j], src[j]);
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Broadcast result to other ranks
|
||||
for (int i = tid; i < size; i += stride) {
|
||||
mxa_buffer[i + 1] = x[i];
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (tid == 0) {
|
||||
for (int i = 0; i < num_ranks; ++i) {
|
||||
if (i != root && !broken_nodes[i]) {
|
||||
auto ctx = ctx_array + i;
|
||||
device_mutex_lock_system(&ctx->mutex);
|
||||
__mlx5gda_device_write_rdma_write_wqe(ctx, (uint64_t) mxa_buffer + sizeof(int), device_byteswap(rkey_array[rank]), raddr_array[i] + sizeof(int), device_byteswap(rkey_array[i]), size * sizeof(int));
|
||||
__mlx5gda_device_post_send_db(ctx);
|
||||
__mlx5gda_device_write_rdma_atomic_add_wqe(ctx, 1, (uint64_t) mxa_buffer, device_byteswap(rkey_array[rank]), raddr_array[i], device_byteswap(rkey_array[i]));
|
||||
__mlx5gda_device_post_send_db(ctx);
|
||||
device_mutex_unlock_system(&ctx->mutex);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// prepare buffer
|
||||
for (int i = tid; i < size; i += stride) {
|
||||
mxa_buffer[i + 1] = x[i];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
uint64_t raddr_base = raddr_array[root] + (1 + size) * rank * sizeof(int);
|
||||
auto ctx = ctx_array + root;
|
||||
|
||||
if (tid == 0) {
|
||||
device_mutex_lock_system(&ctx->mutex);
|
||||
__mlx5gda_device_write_rdma_write_wqe(ctx, (uint64_t) mxa_buffer + sizeof(int), device_byteswap(rkey_array[rank]), raddr_base + sizeof(int), device_byteswap(rkey_array[root]), size * sizeof(int));
|
||||
__mlx5gda_device_post_send_db(ctx);
|
||||
__mlx5gda_device_write_rdma_atomic_add_wqe(ctx, 1, (uint64_t) mxa_buffer, device_byteswap(rkey_array[rank]), raddr_base, device_byteswap(rkey_array[root]));
|
||||
__mlx5gda_device_post_send_db(ctx);
|
||||
device_mutex_unlock_system(&ctx->mutex);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Wait for broadcast result from root
|
||||
if (tid == 0) {
|
||||
while (ld_acquire_sys_global(mxa_buffer) == 0) ;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int i = tid; i < size; i += stride) {
|
||||
x[i] = mxa_buffer[i + 1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void all_reduce_without(const int32_t* broken_nodes, int* x, int* mxa_buffer,
|
||||
void* raddrs, void* rkeys, void* qp_devctxs,
|
||||
int size, int rank, int num_ranks, cudaStream_t stream) {
|
||||
all_reduce_kernel_without<<<1, 256, 0, stream>>>(broken_nodes, x, mxa_buffer, raddrs, rkeys, qp_devctxs, size, rank, num_ranks);
|
||||
}
|
||||
|
||||
}
|
||||
} // namespace mooncake
|
||||
|
|
@ -5,10 +5,10 @@
|
|||
#include <infiniband/verbs.h>
|
||||
#include <infiniband/mlx5dv.h>
|
||||
|
||||
#include "memheap.h"
|
||||
#include "mlx5gda.h"
|
||||
#include "mlx5_ifc.h"
|
||||
#include "mlx5_prm.h"
|
||||
#include <mooncake_ibgda/memheap.h>
|
||||
#include <mooncake_ibgda/mlx5gda.h>
|
||||
#include <mooncake_ibgda/mlx5_ifc.h>
|
||||
#include <mooncake_ibgda/mlx5_prm.h>
|
||||
|
||||
template <typename T>
|
||||
inline T IBGDA_ILOG2(T _n) {
|
||||
|
|
@ -0,0 +1,284 @@
|
|||
#include <mooncake_backend.h>
|
||||
#include <mooncake_worker.cuh>
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
class MooncakeWorkCpu : public ::c10d::Work {
|
||||
public:
|
||||
MooncakeWorkCpu(c10d::OpType opType,
|
||||
c10::intrusive_ptr<c10::ivalue::Future> future)
|
||||
: Work(-1, opType), future_(future) {}
|
||||
|
||||
bool isCompleted() override { return future_->completed(); }
|
||||
|
||||
bool wait(std::chrono::milliseconds timeout) override {
|
||||
future_->wait();
|
||||
return future_->completed() && !future_->hasError();
|
||||
}
|
||||
|
||||
private:
|
||||
c10::intrusive_ptr<c10::ivalue::Future> future_;
|
||||
};
|
||||
|
||||
class MooncakeWorkCuda : public ::c10d::Work {
|
||||
public:
|
||||
MooncakeWorkCuda(c10d::OpType opType, std::shared_ptr<torch::Event> event)
|
||||
: Work(-1, opType), event_(event) {}
|
||||
|
||||
bool isCompleted() override { return event_->query(); }
|
||||
|
||||
bool wait(std::chrono::milliseconds timeout) override {
|
||||
return true; // This should be a no-op
|
||||
}
|
||||
|
||||
private:
|
||||
std::shared_ptr<torch::Event> event_;
|
||||
};
|
||||
|
||||
__global__ void enqueueTaskKernel(c10d::OpType opType, size_t tensorSize,
|
||||
int64_t broadcastRoot, int bufferOffset,
|
||||
void* meta, Task* tasks, int numRanks,
|
||||
const bool* activeRanks,
|
||||
int* activeRanksTensor, size_t taskId) {
|
||||
// Copy task into slot
|
||||
tasks[taskId].opType = opType;
|
||||
tasks[taskId].tensorSize = tensorSize;
|
||||
tasks[taskId].broadcastRoot = broadcastRoot;
|
||||
tasks[taskId].bufferOffset = bufferOffset;
|
||||
tasks[taskId].transferGroupMeta = meta;
|
||||
|
||||
// Mark active
|
||||
__threadfence(); // Ensure writes visible to host
|
||||
tasks[taskId].active = true;
|
||||
|
||||
// Spin-wait until CPU proxy sets DONE
|
||||
while (tasks[taskId].active) {
|
||||
__threadfence();
|
||||
}
|
||||
for (int i = 0; i < numRanks; ++i) {
|
||||
activeRanksTensor[i] = activeRanks[i] ? 1 : 0;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void reduceKernel(scalar_t* dst, const scalar_t* src,
|
||||
size_t numElements, size_t numRanks,
|
||||
bool* activeRanks) {
|
||||
size_t thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
size_t stride = blockDim.x * gridDim.x;
|
||||
for (size_t elem_idx = thread_idx; elem_idx < numElements;
|
||||
elem_idx += stride) {
|
||||
scalar_t sum = 0;
|
||||
for (size_t rank = 0; rank < numRanks; ++rank) {
|
||||
if (activeRanks[rank]) {
|
||||
sum += src[rank * numElements + elem_idx];
|
||||
}
|
||||
}
|
||||
dst[elem_idx] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
void launchReduceKernel(at::Tensor dst, void* src, size_t numRanks,
|
||||
c10d::ReduceOp op, bool* activeRanks,
|
||||
cudaStream_t stream) {
|
||||
TORCH_CHECK(op == c10d::ReduceOp::SUM, "Only support SUM for reduction.");
|
||||
switch (dst.scalar_type()) {
|
||||
case c10::kByte:
|
||||
reduceKernel<<<64, 256, 0, stream>>>(dst.data_ptr<uint8_t>(),
|
||||
(uint8_t*)src, dst.numel(),
|
||||
numRanks, activeRanks);
|
||||
break;
|
||||
case c10::kChar:
|
||||
reduceKernel<<<64, 256, 0, stream>>>(dst.data_ptr<int8_t>(),
|
||||
(int8_t*)src, dst.numel(),
|
||||
numRanks, activeRanks);
|
||||
break;
|
||||
case c10::kShort:
|
||||
reduceKernel<<<64, 256, 0, stream>>>(dst.data_ptr<int16_t>(),
|
||||
(int16_t*)src, dst.numel(),
|
||||
numRanks, activeRanks);
|
||||
break;
|
||||
case c10::kInt:
|
||||
reduceKernel<<<64, 256, 0, stream>>>(dst.data_ptr<int>(), (int*)src,
|
||||
dst.numel(), numRanks,
|
||||
activeRanks);
|
||||
break;
|
||||
case c10::kLong:
|
||||
reduceKernel<<<64, 256, 0, stream>>>(dst.data_ptr<int64_t>(),
|
||||
(int64_t*)src, dst.numel(),
|
||||
numRanks, activeRanks);
|
||||
break;
|
||||
case c10::kFloat:
|
||||
reduceKernel<<<64, 256, 0, stream>>>(dst.data_ptr<float>(),
|
||||
(float*)src, dst.numel(),
|
||||
numRanks, activeRanks);
|
||||
break;
|
||||
case c10::kDouble:
|
||||
reduceKernel<<<64, 256, 0, stream>>>(dst.data_ptr<double>(),
|
||||
(double*)src, dst.numel(),
|
||||
numRanks, activeRanks);
|
||||
break;
|
||||
case c10::kBool:
|
||||
reduceKernel<<<64, 256, 0, stream>>>(dst.data_ptr<bool>(),
|
||||
(bool*)src, dst.numel(),
|
||||
numRanks, activeRanks);
|
||||
break;
|
||||
case c10::kBFloat16:
|
||||
reduceKernel<<<64, 256, 0, stream>>>(
|
||||
dst.data_ptr<at::BFloat16>(), (at::BFloat16*)src, dst.numel(),
|
||||
numRanks, activeRanks);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, c10::str("Unsupported reduce dtype: ",
|
||||
dst.scalar_type()));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T applyReduceOp(const T& a, const T& b, c10d::ReduceOp op) {
|
||||
switch (op) {
|
||||
case c10d::ReduceOp::SUM:
|
||||
return a + b;
|
||||
case c10d::ReduceOp::PRODUCT:
|
||||
return a * b;
|
||||
case c10d::ReduceOp::MIN:
|
||||
return std::min(a, b);
|
||||
case c10d::ReduceOp::MAX:
|
||||
return std::max(a, b);
|
||||
default:
|
||||
TORCH_CHECK(false, c10::str("Unsupported reduce op: ", op));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void reduceCpu(T* dst, const T* src, size_t numElements, size_t numRanks,
|
||||
c10d::ReduceOp op) {
|
||||
at::parallel_for(0, numElements, 1024, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t i = begin; i < end; ++i) {
|
||||
T acc = src[i];
|
||||
for (int64_t rank = 1; rank < numRanks; ++rank) {
|
||||
acc = applyReduceOp(acc, src[i + rank * numElements], op);
|
||||
}
|
||||
dst[i] = acc;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
void launchReduceCpu(at::Tensor dst, void* src, size_t numRanks,
|
||||
c10d::ReduceOp op) {
|
||||
switch (dst.scalar_type()) {
|
||||
case c10::kByte:
|
||||
reduceCpu(dst.data_ptr<uint8_t>(), (uint8_t*)src, dst.numel(),
|
||||
numRanks, op);
|
||||
break;
|
||||
case c10::kChar:
|
||||
reduceCpu(dst.data_ptr<int8_t>(), (int8_t*)src, dst.numel(),
|
||||
numRanks, op);
|
||||
break;
|
||||
case c10::kShort:
|
||||
reduceCpu(dst.data_ptr<int16_t>(), (int16_t*)src, dst.numel(),
|
||||
numRanks, op);
|
||||
break;
|
||||
case c10::kInt:
|
||||
reduceCpu(dst.data_ptr<int>(), (int*)src, dst.numel(), numRanks,
|
||||
op);
|
||||
break;
|
||||
case c10::kLong:
|
||||
reduceCpu(dst.data_ptr<int64_t>(), (int64_t*)src, dst.numel(),
|
||||
numRanks, op);
|
||||
break;
|
||||
case c10::kFloat:
|
||||
reduceCpu(dst.data_ptr<float>(), (float*)src, dst.numel(), numRanks,
|
||||
op);
|
||||
break;
|
||||
case c10::kDouble:
|
||||
reduceCpu(dst.data_ptr<double>(), (double*)src, dst.numel(),
|
||||
numRanks, op);
|
||||
break;
|
||||
case c10::kBool:
|
||||
reduceCpu(dst.data_ptr<bool>(), (bool*)src, dst.numel(), numRanks,
|
||||
op);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, c10::str("Unsupported reduce dtype: ",
|
||||
dst.scalar_type()));
|
||||
}
|
||||
}
|
||||
|
||||
MooncakeWorker::MooncakeWorker() {
|
||||
// Pin memory for task array
|
||||
cudaHostAlloc(&tasks_, kNumTasks_ * sizeof(Task), cudaHostAllocMapped);
|
||||
cudaHostGetDevicePointer(&tasks_device_, tasks_, 0);
|
||||
for (size_t i = 0; i < kNumTasks_; ++i) {
|
||||
tasks_[i].active = false;
|
||||
}
|
||||
|
||||
// Start worker
|
||||
startWorker();
|
||||
}
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> MooncakeWorker::putTaskCpu(
|
||||
c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
|
||||
TransferGroupMeta* meta,
|
||||
const std::function<void(void* dst)>& tensorToBuffer,
|
||||
const std::function<void(void* src)>& bufferToTensor) {
|
||||
TORCH_CHECK(tensorSize * meta->size < kBufferSize, "Too large!");
|
||||
auto future = c10::make_intrusive<c10::ivalue::Future>(
|
||||
c10::ListType::create(c10::TensorType::get()));
|
||||
// Alternately use even-odd items to maintain tasks
|
||||
int taskId = cpuTaskCount % 2;
|
||||
TORCH_CHECK(!tasks_[taskId].active);
|
||||
int bufferOffset = meta->bufferBaseIndex + meta->taskCount % 2;
|
||||
|
||||
tasks_[taskId].opType = opType;
|
||||
tasks_[taskId].tensorSize = tensorSize;
|
||||
tasks_[taskId].broadcastRoot = broadcastRoot;
|
||||
tasks_[taskId].bufferOffset = bufferOffset;
|
||||
tasks_[taskId].transferGroupMeta = meta;
|
||||
tensorToBuffer(
|
||||
(void*)meta->segmentDescs[meta->rank]->buffers[bufferOffset].addr);
|
||||
|
||||
hasCallback_[taskId] = true;
|
||||
callbacks_[taskId] = [this, meta, bufferToTensor, bufferOffset, future] {
|
||||
for (int i = 0; i < meta->size; ++i) {
|
||||
meta->activeRanksTensor[i] = meta->activeRanks[i] ? 1 : 0;
|
||||
}
|
||||
bufferToTensor((void*)meta->segmentDescs[meta->rank]
|
||||
->buffers[bufferOffset + 2]
|
||||
.addr);
|
||||
future->markCompleted(c10::IValue());
|
||||
};
|
||||
|
||||
tasks_[taskId].active = true;
|
||||
++cpuTaskCount;
|
||||
++meta->taskCount;
|
||||
return c10::make_intrusive<MooncakeWorkCpu>(opType, future);
|
||||
}
|
||||
|
||||
c10::intrusive_ptr<c10d::Work> MooncakeWorker::putTaskCuda(
|
||||
c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
|
||||
TransferGroupMeta* meta, const at::cuda::CUDAStream& stream,
|
||||
const std::function<void(void* dst)>& tensorToBuffer,
|
||||
const std::function<void(void* src)>& bufferToTensor) {
|
||||
TORCH_CHECK(tensorSize * meta->size < kBufferSize, "Too large!");
|
||||
// Alternately use even-odd items to maintain tasks
|
||||
int taskId = cudaTaskCount % 2 + 2;
|
||||
int bufferOffset = meta->bufferBaseIndex + meta->taskCount % 2;
|
||||
tensorToBuffer(
|
||||
(void*)meta->segmentDescs[meta->rank]->buffers[bufferOffset].addr);
|
||||
|
||||
hasCallback_[taskId] = false;
|
||||
enqueueTaskKernel<<<1, 1, 0, stream>>>(
|
||||
opType, tensorSize, broadcastRoot, bufferOffset, meta, tasks_device_,
|
||||
meta->size, meta->activeRanksDevice,
|
||||
meta->activeRanksTensor.data_ptr<int>(), taskId);
|
||||
bufferToTensor(
|
||||
(void*)meta->segmentDescs[meta->rank]->buffers[bufferOffset + 2].addr);
|
||||
++cudaTaskCount;
|
||||
++meta->taskCount;
|
||||
auto event = std::make_shared<torch::Event>(torch::kCUDA);
|
||||
event->record(stream);
|
||||
return c10::make_intrusive<MooncakeWorkCuda>(opType, event);
|
||||
}
|
||||
|
||||
} // namespace mooncake
|
||||
|
|
@ -0,0 +1,192 @@
|
|||
#include <thread>
|
||||
#include <mooncake_worker.cuh>
|
||||
#include <transfer_engine.h>
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
enum WorkerTaskStatus {
|
||||
IDLE = 0,
|
||||
TRANSFERRED_1 = 1,
|
||||
SIGNALED_1 = 2,
|
||||
DONE = 3,
|
||||
};
|
||||
|
||||
void MooncakeWorker::startWorker() {
|
||||
running_ = true;
|
||||
std::thread([this] {
|
||||
std::atomic<WorkerTaskStatus> task_status[kNumTasks_];
|
||||
using clock = std::chrono::high_resolution_clock;
|
||||
clock::time_point activeTime[kNumTasks_];
|
||||
while (running_) {
|
||||
_mm_pause();
|
||||
for (size_t i = 0; i < kNumTasks_; ++i) {
|
||||
auto &task = tasks_[i];
|
||||
if (!task.active) {
|
||||
task_status[i].store(IDLE, std::memory_order_release);
|
||||
continue;
|
||||
}
|
||||
|
||||
auto group = (TransferGroupMeta *)task.transferGroupMeta;
|
||||
bool skipTransfer = (task.opType == c10d::OpType::BROADCAST &&
|
||||
group->rank != task.broadcastRoot) ||
|
||||
task.opType == c10d::OpType::BARRIER;
|
||||
if (task_status[i].load(std::memory_order_acquire) == IDLE) {
|
||||
if (skipTransfer) {
|
||||
task_status[i].store(TRANSFERRED_1,
|
||||
std::memory_order_release);
|
||||
continue;
|
||||
}
|
||||
std::vector<TransferRequest> entries;
|
||||
for (int j = 0; j < group->size; ++j) {
|
||||
uint64_t source = group->segmentDescs[group->rank]
|
||||
->buffers[task.bufferOffset]
|
||||
.addr;
|
||||
switch (task.opType) {
|
||||
case c10d::OpType::BROADCAST:
|
||||
case c10d::OpType::ALLREDUCE:
|
||||
case c10d::OpType::ALLGATHER:
|
||||
case c10d::OpType::_ALLGATHER_BASE:
|
||||
break;
|
||||
case c10d::OpType::REDUCE_SCATTER:
|
||||
case c10d::OpType::ALLTOALL_BASE:
|
||||
case c10d::OpType::ALLTOALL:
|
||||
source += j * task.tensorSize;
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
uint64_t target_offset =
|
||||
group->segmentDescs[j]
|
||||
->buffers[task.bufferOffset + 2]
|
||||
.addr;
|
||||
switch (task.opType) {
|
||||
case c10d::OpType::BROADCAST:
|
||||
break;
|
||||
case c10d::OpType::ALLREDUCE:
|
||||
case c10d::OpType::ALLGATHER:
|
||||
case c10d::OpType::_ALLGATHER_BASE:
|
||||
case c10d::OpType::REDUCE_SCATTER:
|
||||
case c10d::OpType::ALLTOALL_BASE:
|
||||
case c10d::OpType::ALLTOALL:
|
||||
target_offset += group->rank * task.tensorSize;
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
entries.push_back(TransferRequest{
|
||||
.opcode = TransferRequest::WRITE,
|
||||
.source = (void *)source,
|
||||
.target_id = group->segmentIDs[j],
|
||||
.target_offset = target_offset,
|
||||
.length = task.tensorSize,
|
||||
});
|
||||
}
|
||||
task.batchID =
|
||||
group->engine->allocateBatchID(entries.size());
|
||||
group->engine->submitTransfer(task.batchID, entries);
|
||||
task_status[i].store(TRANSFERRED_1,
|
||||
std::memory_order_release);
|
||||
} else if (task_status[i].load(std::memory_order_acquire) ==
|
||||
TRANSFERRED_1) {
|
||||
bool batch_done = true;
|
||||
TransferStatus status;
|
||||
|
||||
if (!skipTransfer) {
|
||||
for (int j = 0; j < group->size; ++j) {
|
||||
group->engine->getTransferStatus(task.batchID, j,
|
||||
status);
|
||||
if (group->activeRanks[j] &&
|
||||
status.s != TransferStatusEnum::COMPLETED) {
|
||||
if (status.s == TransferStatusEnum::FAILED) {
|
||||
LOG(ERROR)
|
||||
<< "Rank " << group->rank
|
||||
<< " marking peer " << j
|
||||
<< " as broken during transferring op "
|
||||
<< (int)task.opType;
|
||||
group->activeRanks[j] = false;
|
||||
} else {
|
||||
batch_done = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!batch_done) {
|
||||
continue;
|
||||
}
|
||||
auto source_ptr =
|
||||
(int32_t *)group->segmentDescs[group->rank]
|
||||
->buffers[task.bufferOffset + 4]
|
||||
.addr;
|
||||
std::vector<TransferRequest> entries;
|
||||
for (int j = 0; j < group->size; ++j) {
|
||||
if (!group->activeRanks[j]) {
|
||||
continue;
|
||||
}
|
||||
*source_ptr = 1;
|
||||
entries.push_back(TransferRequest{
|
||||
.opcode = TransferRequest::WRITE,
|
||||
.source = (void *)source_ptr,
|
||||
.target_id = group->segmentIDs[j],
|
||||
.target_offset =
|
||||
group->segmentDescs[j]
|
||||
->buffers[task.bufferOffset + 6]
|
||||
.addr +
|
||||
group->rank * sizeof(int32_t),
|
||||
.length = sizeof(int32_t),
|
||||
});
|
||||
}
|
||||
task.batchID =
|
||||
group->engine->allocateBatchID(entries.size());
|
||||
group->engine->submitTransfer(task.batchID, entries);
|
||||
activeTime[i] = clock::now();
|
||||
task_status[i].store(SIGNALED_1, std::memory_order_release);
|
||||
} else if (task_status[i].load(std::memory_order_acquire) ==
|
||||
SIGNALED_1) {
|
||||
bool all_received = true;
|
||||
auto signal_ptr =
|
||||
(int32_t *)group->segmentDescs[group->rank]
|
||||
->buffers[task.bufferOffset + 6]
|
||||
.addr;
|
||||
auto now = clock::now();
|
||||
auto diff =
|
||||
std::chrono::duration_cast<std::chrono::seconds>(
|
||||
now - activeTime[i]);
|
||||
for (int j = 0; j < group->size; ++j) {
|
||||
if (group->activeRanks[j] && signal_ptr[j] != 1) {
|
||||
TransferMetadata::NotifyDesc msg{"ping", "ping"};
|
||||
if (diff.count() > 1 &&
|
||||
group->engine->sendNotifyByName(
|
||||
group->segmentDescs[j]->name, msg)) {
|
||||
LOG(ERROR) << "Rank " << group->rank
|
||||
<< " marking peer " << j
|
||||
<< " as broken during syncing op "
|
||||
<< (int)task.opType;
|
||||
group->activeRanks[j] = false;
|
||||
} else {
|
||||
all_received = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (diff.count() > 1) {
|
||||
// reset timer
|
||||
activeTime[i] = clock::now();
|
||||
}
|
||||
if (all_received) {
|
||||
for (int j = 0; j < group->size; ++j) {
|
||||
signal_ptr[j] = 0;
|
||||
}
|
||||
task_status[i].store(DONE, std::memory_order_release);
|
||||
task.active = false;
|
||||
if (hasCallback_[i]) {
|
||||
callbacks_[i]();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}).detach();
|
||||
}
|
||||
|
||||
} // namespace mooncake
|
||||
|
|
@ -0,0 +1,95 @@
|
|||
|
||||
#include <mooncake_backend.h>
|
||||
#include <mooncake_ep_buffer.h>
|
||||
#include <pybind11/gil.h> // For GIL management
|
||||
#include <pybind11/stl.h>
|
||||
#include <pybind11/chrono.h>
|
||||
#include <pybind11/functional.h>
|
||||
#include <torch/csrc/utils/pybind.h>
|
||||
#include <torch/python.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
namespace py = pybind11;
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
c10::intrusive_ptr<c10d::Backend> createMooncakeBackend(
|
||||
c10d::DistributedBackendOptions distBackendOpts,
|
||||
c10::intrusive_ptr<MooncakeBackend::MooncakeBackendOptions>
|
||||
backendOptions) {
|
||||
return c10::make_intrusive<MooncakeBackend>(
|
||||
distBackendOpts.store, distBackendOpts.group_rank,
|
||||
distBackendOpts.group_size, backendOptions);
|
||||
}
|
||||
|
||||
c10::intrusive_ptr<c10d::Backend> createMooncakeCpuBackend(
|
||||
c10d::DistributedBackendOptions distBackendOpts,
|
||||
c10::intrusive_ptr<MooncakeBackend::MooncakeBackendOptions>
|
||||
backendOptions) {
|
||||
return c10::make_intrusive<MooncakeBackend>(
|
||||
distBackendOpts.store, distBackendOpts.group_rank,
|
||||
distBackendOpts.group_size, backendOptions, true);
|
||||
}
|
||||
|
||||
__attribute__((constructor)) static void MooncakeBackendConstructor() {
|
||||
auto version = py::module::import("torch")
|
||||
.attr("__version__")
|
||||
.attr("split")("+")
|
||||
.cast<std::vector<std::string>>()[0];
|
||||
TORCH_CHECK(version == "2.8.0", "Mooncake Backend requires torch==2.8.0");
|
||||
py::object module = py::module::import("torch.distributed");
|
||||
py::object register_backend =
|
||||
module.attr("Backend").attr("register_backend");
|
||||
py::dict kwargsCpu;
|
||||
kwargsCpu["devices"] = py::make_tuple("cpu");
|
||||
register_backend("mooncake-cpu", py::cpp_function(createMooncakeCpuBackend),
|
||||
/* extended_api */ true, **kwargsCpu);
|
||||
py::dict kwargsCuda;
|
||||
kwargsCuda["devices"] = py::make_tuple("cuda");
|
||||
register_backend("mooncake", py::cpp_function(createMooncakeBackend),
|
||||
/* extended_api */ true, **kwargsCuda);
|
||||
}
|
||||
|
||||
std::string getPreferredHca(c10::intrusive_ptr<c10d::Backend> backend,
|
||||
std::string location) {
|
||||
auto mooncakeBackend =
|
||||
c10::static_intrusive_pointer_cast<MooncakeBackend>(backend);
|
||||
return mooncakeBackend->getPreferredHca(location);
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(ep, m) {
|
||||
m.def("createMooncakeBackend", &createMooncakeBackend);
|
||||
m.def("createMooncakeCpuBackend", &createMooncakeCpuBackend);
|
||||
m.def("set_host_ip", &MooncakeBackend::setHostIp);
|
||||
m.def("set_device_filter", &MooncakeBackend::setDeviceFilter);
|
||||
m.def("get_preferred_hca", &getPreferredHca);
|
||||
|
||||
py::class_<MooncakeBackend::MooncakeBackendOptions,
|
||||
c10::intrusive_ptr<MooncakeBackend::MooncakeBackendOptions>>(
|
||||
m, "MooncakeBackendOptions")
|
||||
.def(py::init<at::Tensor>(), py::arg("active_ranks"));
|
||||
|
||||
m.def("get_ep_buffer_size_hint", &get_ep_buffer_size_hint);
|
||||
|
||||
py::class_<EventHandle>(m, "EventHandle")
|
||||
.def(py::init<>())
|
||||
.def("current_stream_wait", &EventHandle::current_stream_wait);
|
||||
|
||||
m.attr("MAX_QP_COUNT") = pybind11::int_(MAX_QP_COUNT);
|
||||
|
||||
py::class_<MooncakeEpBuffer>(m, "Buffer")
|
||||
.def(py::init<int, int, int64_t, int>())
|
||||
.def("is_roce", &MooncakeEpBuffer::is_roce)
|
||||
.def("sync_ib", &MooncakeEpBuffer::sync_ib)
|
||||
.def("sync_roce", &MooncakeEpBuffer::sync_roce)
|
||||
.def("get_mr_info", &MooncakeEpBuffer::get_mr_info)
|
||||
.def("get_gid", &MooncakeEpBuffer::get_gid)
|
||||
.def("get_local_qpns", &MooncakeEpBuffer::get_local_qpns)
|
||||
.def("get_local_lids", &MooncakeEpBuffer::get_local_lids)
|
||||
.def("dispatch", &MooncakeEpBuffer::dispatch)
|
||||
.def("combine", &MooncakeEpBuffer::combine)
|
||||
.def("get_next_combine_buffer",
|
||||
&MooncakeEpBuffer::get_next_combine_buffer);
|
||||
}
|
||||
|
||||
} // namespace mooncake
|
||||
|
|
@ -0,0 +1,158 @@
|
|||
import torch
|
||||
import torch.distributed as dist
|
||||
from typing import Any, Callable, List, Tuple, Optional, Union
|
||||
|
||||
# noinspection PyUnresolvedReferences
|
||||
from mooncake import ep
|
||||
|
||||
|
||||
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[ep.EventHandle] = None,
|
||||
extra_tensors: Optional[Tuple[torch.Tensor, ...]] = None) -> None:
|
||||
"""
|
||||
Initialize the class.
|
||||
|
||||
Arguments:
|
||||
event: the CUDA event captured.
|
||||
extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph.
|
||||
"""
|
||||
self.event = event
|
||||
|
||||
# NOTES: we use extra tensors to achieve stream recording, otherwise,
|
||||
# stream recording will be incompatible with CUDA graph.
|
||||
self.extra_tensors = extra_tensors
|
||||
|
||||
def current_stream_wait(self) -> None:
|
||||
"""
|
||||
The current stream `torch.cuda.current_stream()` waits for the event to be finished.
|
||||
"""
|
||||
assert self.event is not None
|
||||
self.event.current_stream_wait()
|
||||
|
||||
def __enter__(self) -> Any:
|
||||
"""
|
||||
Utility for overlapping and Python `with` syntax.
|
||||
|
||||
You can overlap the kernels on the current stream with the following example:
|
||||
```python
|
||||
event_overlap = event_after_all_to_all_kernels()
|
||||
with event_overlap():
|
||||
do_something_on_current_stream()
|
||||
# After exiting the `with` scope, the current stream with wait the event to be finished.
|
||||
```
|
||||
"""
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""
|
||||
Utility for overlapping and Python `with` syntax.
|
||||
|
||||
Please follow the example in the `__enter__` function.
|
||||
"""
|
||||
if self.event is not None:
|
||||
self.event.current_stream_wait()
|
||||
|
||||
|
||||
class Buffer:
|
||||
def __init__(self, group: dist.ProcessGroup, num_ep_buffer_bytes: int = 0):
|
||||
# Initialize the CPP runtime
|
||||
self.rank = group.rank()
|
||||
self.group_size = group.size()
|
||||
self.group = group
|
||||
self.num_ep_buffer_bytes = num_ep_buffer_bytes
|
||||
|
||||
# Get the index of the closest NIC
|
||||
backend = self.group._get_backend(torch.device('cuda'))
|
||||
preferred_hca = ep.get_preferred_hca(backend, f'cuda:{torch.cuda.current_device()}')
|
||||
nic_id = int(preferred_hca.split('_')[1])
|
||||
self.runtime = ep.Buffer(self.rank, self.group_size, num_ep_buffer_bytes, nic_id)
|
||||
|
||||
(raddr, rkey) = self.runtime.get_mr_info()
|
||||
|
||||
raddr = torch.tensor([raddr], dtype=torch.int64, device='cuda')
|
||||
raddrs = [torch.empty(1, dtype=torch.int64, device='cuda') for _ in range(self.group_size)]
|
||||
dist.all_gather(raddrs, raddr, group)
|
||||
raddrs = torch.cat(raddrs).tolist()
|
||||
|
||||
rkey = torch.tensor([rkey], dtype=torch.int32, device='cuda')
|
||||
rkeys = [torch.empty(1, dtype=torch.int32, device='cuda') for _ in range(self.group_size)]
|
||||
dist.all_gather(rkeys, rkey, group)
|
||||
rkeys = torch.cat(rkeys).tolist()
|
||||
|
||||
all_to_all_size = ep.MAX_QP_COUNT // self.group_size
|
||||
|
||||
local_qpns = self.runtime.get_local_qpns()
|
||||
local_qpns = list(torch.unbind(torch.tensor(local_qpns, dtype=torch.int32, device='cuda').view(-1, all_to_all_size)))
|
||||
remote_qpns = [torch.empty(all_to_all_size, dtype=torch.int32, device='cuda') for _ in range(self.group_size)]
|
||||
dist.all_to_all(remote_qpns, local_qpns, group)
|
||||
remote_qpns = torch.cat(remote_qpns).tolist()
|
||||
|
||||
if self.runtime.is_roce():
|
||||
(subnet_prefix, interface_id) = self.runtime.get_gid()
|
||||
|
||||
subnet_prefix = torch.tensor([subnet_prefix], dtype=torch.int64, device='cuda')
|
||||
subnet_prefixes = [torch.empty(1, dtype=torch.int64, device='cuda') for _ in range(self.group_size)]
|
||||
dist.all_gather(subnet_prefixes, subnet_prefix, group)
|
||||
subnet_prefixes = torch.cat(subnet_prefixes).tolist()
|
||||
|
||||
interface_id = torch.tensor([interface_id], dtype=torch.int64, device='cuda')
|
||||
interface_ids = [torch.empty(1, dtype=torch.int64, device='cuda') for _ in range(self.group_size)]
|
||||
dist.all_gather(interface_ids, interface_id, group)
|
||||
interface_ids = torch.cat(interface_ids).tolist()
|
||||
|
||||
self.runtime.sync_roce(raddrs, rkeys, remote_qpns, subnet_prefixes, interface_ids)
|
||||
else:
|
||||
|
||||
local_lids = self.runtime.get_local_lids()
|
||||
local_lids = list(torch.unbind(torch.tensor(local_lids, dtype=torch.int32, device='cuda').view(-1, all_to_all_size)))
|
||||
remote_lids = [torch.empty(all_to_all_size, dtype=torch.int32, device='cuda') for _ in range(self.group_size)]
|
||||
dist.all_to_all(remote_lids, local_lids, group)
|
||||
remote_lids = torch.cat(remote_lids).tolist()
|
||||
|
||||
self.runtime.sync_ib(raddrs, rkeys, remote_qpns, remote_lids)
|
||||
|
||||
@staticmethod
|
||||
def get_ep_buffer_size_hint(num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int) -> int:
|
||||
return ep.get_ep_buffer_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts)
|
||||
|
||||
# noinspection PyTypeChecker
|
||||
def dispatch(self, x: torch.Tensor, topk_idx: torch.Tensor, active_ranks: torch.Tensor,
|
||||
num_max_dispatch_tokens_per_rank: int, num_experts: int, timeout_us: int,
|
||||
use_fp8: bool = True, async_finish: bool = False, return_recv_hook: bool = False) -> \
|
||||
Tuple[Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor], torch.Tensor, Tuple, EventOverlap, Callable]:
|
||||
packed_recv_x, packed_recv_x_scales, packed_recv_count, packed_recv_src_info, packed_recv_layout_range, event, hook = \
|
||||
self.runtime.dispatch(x, topk_idx, active_ranks,
|
||||
num_max_dispatch_tokens_per_rank, num_experts, timeout_us,
|
||||
use_fp8, 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)
|
||||
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 combine(self, x: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor,
|
||||
active_ranks: torch.Tensor, timeout_us: int,
|
||||
handle: tuple, zero_copy: bool = False, async_finish: bool = False,
|
||||
return_recv_hook: bool = False, out: Optional[torch.Tensor] = None) -> \
|
||||
Tuple[torch.Tensor, EventOverlap, Callable]:
|
||||
src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
|
||||
combined_x, event, hook = self.runtime.combine(x, topk_idx, topk_weights, src_info, layout_range,
|
||||
active_ranks,
|
||||
num_max_dispatch_tokens_per_rank, num_experts, timeout_us,
|
||||
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 get_next_combine_buffer(self, handle: object):
|
||||
src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
|
||||
return self.runtime.get_next_combine_buffer(num_max_dispatch_tokens_per_rank, hidden, num_experts)
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
import os
|
||||
import sys
|
||||
import platform
|
||||
from setuptools import setup, Distribution
|
||||
|
|
@ -108,7 +109,40 @@ class CustomBdistWheel(bdist_wheel):
|
|||
# ---------------------------------------------------------------------------
|
||||
# setup()
|
||||
# ---------------------------------------------------------------------------
|
||||
setup(
|
||||
distclass=BinaryDistribution,
|
||||
cmdclass={"bdist_wheel": CustomBdistWheel},
|
||||
)
|
||||
if int(os.getenv("BUILD_WITH_EP", "0")):
|
||||
import torch
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
abi_flag = int(torch._C._GLIBCXX_USE_CXX11_ABI)
|
||||
current_dir = os.path.abspath(os.path.dirname(__file__))
|
||||
ext_modules = [
|
||||
CUDAExtension(
|
||||
name="mooncake.ep",
|
||||
include_dirs=[
|
||||
os.path.join(current_dir, "../mooncake-ep/include"),
|
||||
os.path.join(current_dir, "../mooncake-transfer-engine/include"),
|
||||
],
|
||||
sources=["../mooncake-integration/ep/ep_py.cpp"],
|
||||
extra_compile_args={
|
||||
"cxx": [f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}", "-std=c++20"],
|
||||
"nvcc": [f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}", "-std=c++20"],
|
||||
},
|
||||
libraries=["ibverbs", "mlx5"],
|
||||
extra_objects=[
|
||||
os.path.join(current_dir, "../build/mooncake-ep/src/libmooncake_ep.a"),
|
||||
os.path.join(current_dir, "mooncake/engine.so"),
|
||||
],
|
||||
)
|
||||
]
|
||||
setup(
|
||||
distclass=BinaryDistribution,
|
||||
cmdclass={
|
||||
"bdist_wheel": CustomBdistWheel,
|
||||
"build_ext": BuildExtension,
|
||||
},
|
||||
ext_modules=ext_modules,
|
||||
)
|
||||
else:
|
||||
setup(
|
||||
distclass=BinaryDistribution,
|
||||
cmdclass={"bdist_wheel": CustomBdistWheel},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -14,17 +14,17 @@ def init_dist(local_rank: int, num_local_ranks: int):
|
|||
node_rank = int(os.getenv('RANK', 0))
|
||||
assert (num_local_ranks < 8 and num_nodes == 1) or num_local_ranks == 8
|
||||
|
||||
torch.cuda.set_device(local_rank)
|
||||
dist.init_process_group(
|
||||
backend='nccl',
|
||||
backend='mooncake',
|
||||
init_method=f'tcp://{ip}:{port}',
|
||||
world_size=num_nodes * num_local_ranks,
|
||||
rank=node_rank * num_local_ranks + local_rank
|
||||
)
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
torch.set_default_device('cuda')
|
||||
torch.cuda.set_device(local_rank)
|
||||
|
||||
return dist.get_rank(), dist.get_world_size(), dist.new_group(list(range(num_local_ranks * num_nodes)))
|
||||
return dist.get_rank(), dist.get_world_size(), dist.new_group(list(range(num_local_ranks * num_nodes))), dist.new_group(list(range(num_local_ranks * num_nodes)), backend="mooncake-cpu")
|
||||
|
||||
|
||||
def calc_diff(x: torch.Tensor, y: torch.Tensor):
|
||||
|
|
@ -0,0 +1,72 @@
|
|||
import os
|
||||
import time
|
||||
import unittest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
from mooncake import ep
|
||||
|
||||
|
||||
def worker(rank, world_size, results, collective):
|
||||
torch.cuda.set_device(rank)
|
||||
dist.init_process_group(
|
||||
backend="mooncake",
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
pg_options=ep.MooncakeBackendOptions(torch.zeros((world_size,), dtype=torch.int32, device="cuda")),
|
||||
)
|
||||
|
||||
if collective == "all_reduce":
|
||||
tensor = torch.tensor([rank + 1], dtype=torch.int32, device="cuda")
|
||||
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
|
||||
results[rank] = tensor.item()
|
||||
|
||||
elif collective == "all_gather":
|
||||
tensor = torch.tensor([rank], device="cuda")
|
||||
gathered = [torch.zeros_like(tensor) for _ in range(world_size)]
|
||||
dist.all_gather(gathered, tensor)
|
||||
results[rank] = [t.item() for t in gathered]
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unsupported collective: {collective}")
|
||||
|
||||
while len(results) < world_size:
|
||||
time.sleep(1)
|
||||
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
class TestMooncakeBackend(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.world_size = torch.cuda.device_count()
|
||||
os.environ["MASTER_ADDR"] = "127.0.0.1"
|
||||
os.environ["MASTER_PORT"] = "29500"
|
||||
|
||||
def tearDown(self):
|
||||
pass
|
||||
|
||||
def _spawn_and_check(self, collective, expected_fn):
|
||||
mp_manager = mp.Manager()
|
||||
results = mp_manager.dict()
|
||||
mp.spawn(
|
||||
worker,
|
||||
args=(self.world_size, results, collective),
|
||||
nprocs=self.world_size,
|
||||
join=True,
|
||||
)
|
||||
|
||||
expected = expected_fn(self.world_size)
|
||||
for r in range(self.world_size):
|
||||
self.assertEqual(results[r], expected)
|
||||
|
||||
def test_allreduce(self):
|
||||
# Expected sum = 1 + 2 + 3 + 4 = 10
|
||||
self._spawn_and_check("all_reduce", lambda size: sum(range(1, size + 1)))
|
||||
|
||||
def test_allgather(self):
|
||||
# Expected gather = [0, 1, 2, 3]
|
||||
self._spawn_and_check("all_gather", lambda size: list(range(size)))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
@ -0,0 +1,107 @@
|
|||
import os
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import mooncake.ep
|
||||
import time
|
||||
import unittest
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
|
||||
def run_latency_test(rank, world_size, backend, device, collective, data_size, results, num_iterations=100):
|
||||
# Initialize the process group in each spawned process
|
||||
torch.cuda.set_device(rank)
|
||||
dist.init_process_group(backend=backend, rank=rank, world_size=world_size)
|
||||
|
||||
# Create a tensor for the collective operation
|
||||
tensor = torch.rand(data_size, device=device)
|
||||
|
||||
gathered = [torch.zeros_like(tensor) for _ in range(world_size)]
|
||||
|
||||
# Warm up
|
||||
for _ in range(num_iterations):
|
||||
if collective == 'broadcast':
|
||||
dist.broadcast(tensor, src=0)
|
||||
elif collective == 'allreduce':
|
||||
dist.all_reduce(tensor)
|
||||
elif collective == 'allgather':
|
||||
dist.all_gather(gathered, tensor)
|
||||
|
||||
# Synchronize before starting the test
|
||||
torch.cuda.synchronize()
|
||||
|
||||
start = time.perf_counter()
|
||||
for _ in range(num_iterations):
|
||||
if collective == 'broadcast':
|
||||
dist.broadcast(tensor, src=0)
|
||||
elif collective == 'allreduce':
|
||||
dist.all_reduce(tensor)
|
||||
elif collective == 'allgather':
|
||||
dist.all_gather(gathered, tensor)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
end = time.perf_counter()
|
||||
|
||||
# Calculate average time
|
||||
avg_time = (end - start) / num_iterations
|
||||
|
||||
# Store the result
|
||||
results[rank] = avg_time
|
||||
|
||||
while len(results) < world_size:
|
||||
time.sleep(1)
|
||||
|
||||
dist.destroy_process_group() # Destroy the process group after testing
|
||||
|
||||
class TestMooncakeBackendPerf(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.world_size = torch.cuda.device_count()
|
||||
os.environ["MASTER_ADDR"] = "127.0.0.1"
|
||||
os.environ["MASTER_PORT"] = "29500"
|
||||
|
||||
def tearDown(self):
|
||||
pass
|
||||
|
||||
def do_test(self, device, collective, data_size):
|
||||
# Use mp.spawn to call the latency test
|
||||
mp_manager = mp.Manager()
|
||||
# Test mooncake
|
||||
mooncake_results = mp_manager.dict()
|
||||
mp.spawn(run_latency_test, args=(self.world_size, 'mooncake-cpu' if device == 'cpu' else 'mooncake', device, collective, data_size, mooncake_results), nprocs=self.world_size, join=True)
|
||||
# Test baseline
|
||||
baseline_results = mp_manager.dict()
|
||||
mp.spawn(run_latency_test, args=(self.world_size, 'gloo' if device == 'cpu' else 'nccl', device, collective, data_size, baseline_results), nprocs=self.world_size, join=True)
|
||||
|
||||
# After all processes have completed, check the results
|
||||
mooncake_latency = max(mooncake_results[r] for r in mooncake_results)
|
||||
baseline_latency = max(baseline_results[r] for r in baseline_results)
|
||||
self.assertLessEqual(mooncake_latency, 10 * baseline_latency,
|
||||
f"Latency of mooncake({device}) for {collective} with size {data_size} exceeded 10 times the baseline.")
|
||||
|
||||
# cpu + allgather
|
||||
def test_cpu_allgather_1024(self):
|
||||
self.do_test("cpu", "allgather", 1024)
|
||||
|
||||
# cpu + allreduce
|
||||
def test_cpu_allreduce_1024(self):
|
||||
self.do_test("cpu", "allreduce", 1024)
|
||||
|
||||
# cpu + broadcast
|
||||
def test_cpu_broadcast_1024(self):
|
||||
self.do_test("cpu", "broadcast", 1024)
|
||||
|
||||
# cuda + allgather
|
||||
def test_cuda_allgather_1024(self):
|
||||
self.do_test("cuda", "allgather", 1024)
|
||||
|
||||
# cuda + allreduce
|
||||
def test_cuda_allreduce_1024(self):
|
||||
self.do_test("cuda", "allreduce", 1024)
|
||||
|
||||
# cuda + broadcast
|
||||
def test_cuda_broadcast_1024(self):
|
||||
self.do_test("cuda", "broadcast", 1024)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
@ -3,12 +3,12 @@ import torch
|
|||
import torch.distributed as dist
|
||||
from functools import partial
|
||||
|
||||
import mxa_ep
|
||||
from utils import init_dist, bench, bench_kineto, calc_diff, hash_tensor, per_token_cast_back
|
||||
from mooncake.mooncake_ep_buffer import Buffer
|
||||
from ep_test_utils import init_dist, bench, bench_kineto, calc_diff, hash_tensor, per_token_cast_back
|
||||
|
||||
|
||||
def test_main(num_tokens: int, hidden: int, num_experts: int, num_topk: int,
|
||||
rank: int, num_ranks: int, group: dist.ProcessGroup, buffer: mxa_ep.Buffer, seed: int = 0):
|
||||
rank: int, num_ranks: int, group: dist.ProcessGroup, cpu_group: dist.ProcessGroup, buffer: Buffer, seed: int = 0):
|
||||
torch.manual_seed(seed + rank)
|
||||
random.seed(seed + rank)
|
||||
|
||||
|
|
@ -32,13 +32,13 @@ def test_main(num_tokens: int, hidden: int, num_experts: int, num_topk: int,
|
|||
# Check dispatch correctness
|
||||
do_check = True
|
||||
hash_value, num_times = 0, 0
|
||||
broken_nodes = torch.zeros((num_tokens, ), dtype=torch.int32, device='cuda')
|
||||
active_ranks = torch.ones((num_tokens, ), dtype=torch.int32, device='cuda')
|
||||
for return_recv_hook in (False, True):
|
||||
for dispatch_use_fp8 in (False, True):
|
||||
num_times += 1
|
||||
for i in range((num_times % 2) + 1):
|
||||
packed_recv_x, packed_recv_count, handle, event, hook = \
|
||||
buffer.dispatch(x, topk_idx, broken_nodes, num_tokens, num_experts, 1000000, use_fp8=dispatch_use_fp8,
|
||||
buffer.dispatch(x, topk_idx, active_ranks, num_tokens, num_experts, -1, use_fp8=dispatch_use_fp8,
|
||||
async_finish=not return_recv_hook, return_recv_hook=return_recv_hook)
|
||||
hook() if return_recv_hook else event.current_stream_wait()
|
||||
packed_recv_x = (packed_recv_x[0], packed_recv_x[1].contiguous()) if dispatch_use_fp8 else packed_recv_x
|
||||
|
|
@ -78,8 +78,7 @@ def test_main(num_tokens: int, hidden: int, num_experts: int, num_topk: int,
|
|||
if zero_copy:
|
||||
buffer.get_next_combine_buffer(handle)[:, :, :] = simulated_gemm_x
|
||||
out = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
|
||||
gathered_experts = torch.zeros((num_experts, ), dtype=torch.int32, device='cuda')
|
||||
combined_x, event, hook = buffer.combine(simulated_gemm_x, topk_idx, topk_weights, gathered_experts, 1000000, handle,
|
||||
combined_x, event, hook = buffer.combine(simulated_gemm_x, topk_idx, topk_weights, active_ranks, -1, handle,
|
||||
async_finish=not return_recv_hook, zero_copy=zero_copy,
|
||||
return_recv_hook=return_recv_hook, out=out)
|
||||
hook() if return_recv_hook else event.current_stream_wait()
|
||||
|
|
@ -109,13 +108,12 @@ def test_main(num_tokens: int, hidden: int, num_experts: int, num_topk: int,
|
|||
# noinspection PyShadowingNames
|
||||
def test_func(zero_copy: bool, return_recv_hook: bool):
|
||||
recv_x, recv_count, handle, event, hook = \
|
||||
buffer.dispatch(x, topk_idx, broken_nodes, num_tokens, num_experts, 1000000,
|
||||
buffer.dispatch(x, topk_idx, active_ranks, num_tokens, num_experts, -1,
|
||||
async_finish=False, return_recv_hook=return_recv_hook)
|
||||
large_gemm_with_hook(hook) if return_recv_hook else None
|
||||
if zero_copy:
|
||||
buffer.get_next_combine_buffer(handle)[:, :, :] = simulated_gemm_x
|
||||
gathered_experts = torch.zeros((num_experts, ), dtype=torch.int32, device='cuda')
|
||||
combined_x, event, hook = buffer.combine(simulated_gemm_x, topk_idx, topk_weights, gathered_experts, 1000000, handle,
|
||||
combined_x, event, hook = buffer.combine(simulated_gemm_x, topk_idx, topk_weights, active_ranks, -1, handle,
|
||||
zero_copy=zero_copy, return_recv_hook=return_recv_hook)
|
||||
large_gemm_with_hook(hook) if return_recv_hook else None
|
||||
|
||||
|
|
@ -134,7 +132,7 @@ def test_main(num_tokens: int, hidden: int, num_experts: int, num_topk: int,
|
|||
|
||||
# Separate profiling
|
||||
for return_recv_hook in (False, True):
|
||||
group.barrier()
|
||||
cpu_group.barrier()
|
||||
dispatch_t, combine_t = bench_kineto(partial(test_func, zero_copy=True, return_recv_hook=return_recv_hook),
|
||||
kernel_names=('dispatch', 'combine'), barrier_comm_profiling=True,
|
||||
suppress_kineto_output=True)
|
||||
|
|
@ -150,22 +148,15 @@ def test_main(num_tokens: int, hidden: int, num_experts: int, num_topk: int,
|
|||
|
||||
# noinspection PyUnboundLocalVariable
|
||||
def test_loop(local_rank: int, num_local_ranks: int):
|
||||
rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
|
||||
rank, num_ranks, group, cpu_group = init_dist(local_rank, num_local_ranks)
|
||||
num_tokens, hidden, num_topk, num_experts = 128, 7168, 8, 288
|
||||
|
||||
bytes_reserved = 2000 * 3000 * num_ranks * 4
|
||||
num_mxa_bytes = mxa_ep.Buffer.get_mxa_size_hint(num_tokens, hidden, num_ranks, num_experts, bytes_reserved)
|
||||
num_ep_buffer_bytes = Buffer.get_ep_buffer_size_hint(num_tokens, hidden, num_ranks, num_experts)
|
||||
if local_rank == 0:
|
||||
print(f'Allocating buffer size: {num_mxa_bytes / 1e6} MB ...', flush=True)
|
||||
buffer = mxa_ep.Buffer(group, num_mxa_bytes=num_mxa_bytes, bytes_reserved=bytes_reserved)
|
||||
print(f'Allocating buffer size: {num_ep_buffer_bytes / 1e6} MB ...', flush=True)
|
||||
buffer = Buffer(group, num_ep_buffer_bytes=num_ep_buffer_bytes)
|
||||
|
||||
x = torch.full((2000, 3000), rank, dtype=torch.int32, device='cuda')
|
||||
broken_nodes = torch.zeros(num_ranks, dtype=torch.int32)
|
||||
buffer.all_reduce_without(broken_nodes, x)
|
||||
expected = torch.full((2000, 3000), sum(range(num_ranks)), dtype=torch.int32, device='cuda')
|
||||
assert torch.equal(x, expected), "All-reduce result is incorrect!"
|
||||
|
||||
test_main(num_tokens, hidden, num_experts, num_topk, rank, num_ranks, group, buffer, seed=1)
|
||||
test_main(num_tokens, hidden, num_experts, num_topk, rank, num_ranks, group, cpu_group, buffer, seed=1)
|
||||
|
||||
do_pressure_test = False
|
||||
for seed in range(int(1e9) if do_pressure_test else 0):
|
||||
|
|
@ -95,6 +95,7 @@ if [ "$PYTHON_VERSION" = "3.8" ]; then
|
|||
PATTERNS=(
|
||||
"libcurl.so*"
|
||||
"libibverbs.so*"
|
||||
"libmlx5.so*"
|
||||
"libnuma.so*"
|
||||
"libstdc++.so*"
|
||||
"libgcc_s.so*"
|
||||
|
|
@ -133,6 +134,12 @@ if [ "$PYTHON_VERSION" = "3.8" ]; then
|
|||
"libffi.so*"
|
||||
"libcuda.so*"
|
||||
"libcudart.so*"
|
||||
"libc10.so*"
|
||||
"libc10_cuda.so*"
|
||||
"libtorch.so*"
|
||||
"libtorch_cpu.so*"
|
||||
"libtorch_cuda.so*"
|
||||
"libtorch_python.so*"
|
||||
"libascendcl.so*"
|
||||
"libhccl.so*"
|
||||
"libmsprofiler.so*"
|
||||
|
|
@ -186,10 +193,15 @@ if [ "$PYTHON_VERSION" = "3.8" ]; then
|
|||
auditwheel repair ${OUTPUT_DIR}/*.whl $EXCLUDE_OPTS -w ${REPAIRED_DIR}/ --plat ${PLATFORM_TAG}
|
||||
else
|
||||
echo "Repairing wheel with auditwheel for platform: $PLATFORM_TAG"
|
||||
python -m build --wheel --outdir ${OUTPUT_DIR}
|
||||
if [ "$BUILD_WITH_EP" = "1" ]; then
|
||||
python -m build --wheel --outdir ${OUTPUT_DIR} --no-isolation
|
||||
else
|
||||
python -m build --wheel --outdir ${OUTPUT_DIR}
|
||||
fi
|
||||
auditwheel repair ${OUTPUT_DIR}/*.whl \
|
||||
--exclude libcurl.so* \
|
||||
--exclude libibverbs.so* \
|
||||
--exclude libmlx5.so* \
|
||||
--exclude libnuma.so* \
|
||||
--exclude libstdc++.so* \
|
||||
--exclude libgcc_s.so* \
|
||||
|
|
@ -228,6 +240,12 @@ else
|
|||
--exclude libffi.so* \
|
||||
--exclude libcuda.so* \
|
||||
--exclude libcudart.so* \
|
||||
--exclude libc10.so* \
|
||||
--exclude libc10_cuda.so* \
|
||||
--exclude libtorch.so* \
|
||||
--exclude libtorch_cpu.so* \
|
||||
--exclude libtorch_cuda.so* \
|
||||
--exclude libtorch_python.so* \
|
||||
--exclude libascendcl.so* \
|
||||
--exclude libhccl.so* \
|
||||
--exclude libmsprofiler.so* \
|
||||
|
|
|
|||
Loading…
Reference in New Issue