wehub-resource-sync
3618 lines
147 KiB
C++
3618 lines
147 KiB
C++
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
// The file has been adapted from DeepSeek DeepEP project
|
|
// Copyright (c) 2025 DeepSeek
|
|
// Licensed under the MIT License -
|
|
// https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
|
|
|
|
#include <cuda_runtime.h>
|
|
#include <atomic>
|
|
#include <chrono>
|
|
#include <memory>
|
|
|
|
#include "paddle/fluid/distributed/collective/deep_ep/deep_ep.hpp"
|
|
#include "paddle/fluid/distributed/collective/deep_ep/kernels/api.cuh"
|
|
#include "paddle/fluid/distributed/collective/deep_ep/kernels/configs.cuh"
|
|
|
|
#include "paddle/fluid/distributed/collective/deep_ep/include/CUDADataType.h"
|
|
#include "paddle/fluid/distributed/collective/deep_ep/include/ScalarType.h"
|
|
#include "paddle/fluid/distributed/collective/process_group_nccl.h"
|
|
#include "paddle/phi/api/include/api.h"
|
|
#include "paddle/phi/api/include/tensor_utils.h"
|
|
#include "paddle/phi/common/data_type.h"
|
|
#include "paddle/phi/common/place.h"
|
|
#include "paddle/phi/core/distributed/utils.h"
|
|
#include "paddle/phi/core/memory/allocation/allocator_facade.h"
|
|
|
|
COMMON_DECLARE_int64(deep_ep_comm_prealloc_in_mb);
|
|
|
|
namespace deep_ep {
|
|
std::once_flag pre_alloc_once_flag;
|
|
|
|
namespace detail {
|
|
void SetAllocatorStreamForGPUContext(cudaStream_t stream,
|
|
phi::GPUContext* ctx) {
|
|
ctx->SetAllocator(paddle::memory::allocation::AllocatorFacade::Instance()
|
|
.GetAllocator(ctx->GetPlace(), stream)
|
|
.get());
|
|
}
|
|
} // namespace detail
|
|
|
|
void PreAlloc(paddle::Tensor tensor, cudaStream_t stream) {
|
|
int64_t numel = tensor.numel();
|
|
auto alloc_size = FLAGS_deep_ep_comm_prealloc_in_mb * 1000000;
|
|
std::cout << "alloc once here, size: " << alloc_size << " numel: " << numel
|
|
<< std::endl;
|
|
std::cout << tensor.place() << "\t" << stream << std::endl;
|
|
paddle::memory::allocation::AllocatorFacade::Instance()
|
|
.GetAllocator(tensor.place(), stream)
|
|
->Allocate(alloc_size);
|
|
}
|
|
|
|
Buffer::Buffer(int rank,
|
|
int num_ranks,
|
|
int64_t num_nvl_bytes,
|
|
int64_t num_rdma_bytes,
|
|
bool low_latency_mode,
|
|
int context_ring_id)
|
|
: rank(rank),
|
|
num_ranks(num_ranks),
|
|
num_nvl_bytes(num_nvl_bytes),
|
|
num_rdma_bytes(num_rdma_bytes),
|
|
low_latency_mode(low_latency_mode) {
|
|
CUDA_CHECK(cudaGetDevice(&device_id));
|
|
auto map = paddle::distributed::ProcessGroupMapFromGid::getInstance();
|
|
paddle::distributed::ProcessGroup* pg = map->get(context_ring_id);
|
|
const auto& place = phi::GPUPlace(device_id);
|
|
comm_ctx =
|
|
reinterpret_cast<paddle::distributed::ProcessGroupNCCL*>(pg)
|
|
->GetOrCreateCommContext(place, phi::distributed::CommType::ALLTOALL);
|
|
comm_stream = comm_ctx->GetStream();
|
|
calc_ctx = reinterpret_cast<phi::GPUContext*>(
|
|
reinterpret_cast<paddle::distributed::ProcessGroupNCCL*>(pg)
|
|
->GetDeviceContext(place, true));
|
|
// Task fifo memory
|
|
int64_t fifo_bytes = sizeof(int) * NUM_MAX_FIFO_SLOTS;
|
|
int64_t buffer_ptr_bytes = sizeof(void*) * NUM_MAX_NVL_PEERS;
|
|
int64_t task_ptr_bytes = sizeof(int*) * NUM_MAX_NVL_PEERS;
|
|
|
|
// Common checks
|
|
EP_HOST_ASSERT(
|
|
num_nvl_bytes % NUM_BUFFER_ALIGNMENT_BYTES == 0 &&
|
|
((low_latency_mode || num_nvl_bytes <= std::numeric_limits<int>::max()) ||
|
|
num_rdma_bytes == 0));
|
|
EP_HOST_ASSERT(
|
|
num_rdma_bytes % NUM_BUFFER_ALIGNMENT_BYTES == 0 &&
|
|
(low_latency_mode || num_rdma_bytes <= std::numeric_limits<int>::max()));
|
|
EP_HOST_ASSERT(0 <= rank && rank < num_ranks &&
|
|
(num_ranks <= NUM_MAX_NVL_PEERS * NUM_MAX_RDMA_PEERS ||
|
|
low_latency_mode));
|
|
EP_HOST_ASSERT(num_ranks < NUM_MAX_NVL_PEERS ||
|
|
num_ranks % NUM_MAX_NVL_PEERS == 0);
|
|
if (num_rdma_bytes > 0)
|
|
EP_HOST_ASSERT(num_ranks > NUM_MAX_NVL_PEERS || low_latency_mode);
|
|
|
|
// Get ranks
|
|
// CUDA_CHECK(cudaGetDevice(&device_id));
|
|
rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS;
|
|
num_rdma_ranks = std::max(1, num_ranks / NUM_MAX_NVL_PEERS),
|
|
num_nvl_ranks = std::min(num_ranks, NUM_MAX_NVL_PEERS);
|
|
|
|
// Get device info
|
|
cudaDeviceProp device_prop = {};
|
|
CUDA_CHECK(cudaGetDeviceProperties(&device_prop, device_id));
|
|
|
|
if (num_nvl_bytes > 0) {
|
|
// Local IPC: alloc local memory and set local IPC handle
|
|
CUDA_CHECK(cudaMalloc(
|
|
&buffer_ptrs[nvl_rank],
|
|
num_nvl_bytes + fifo_bytes + buffer_ptr_bytes + task_ptr_bytes));
|
|
CUDA_CHECK(
|
|
cudaIpcGetMemHandle(&ipc_handles[nvl_rank], buffer_ptrs[nvl_rank]));
|
|
buffer_ptrs_gpu = reinterpret_cast<void**>(
|
|
reinterpret_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes +
|
|
fifo_bytes);
|
|
|
|
// Set task fifo
|
|
EP_HOST_ASSERT(NUM_MAX_FIFO_SLOTS % num_nvl_ranks == 0);
|
|
task_fifo_ptrs[nvl_rank] = reinterpret_cast<int*>(
|
|
reinterpret_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes);
|
|
task_fifo_ptrs_gpu = reinterpret_cast<int**>(
|
|
reinterpret_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes +
|
|
fifo_bytes + buffer_ptr_bytes);
|
|
|
|
// No need to synchronize, will do a full device sync during `sync`
|
|
CUDA_CHECK(cudaMemsetAsync(
|
|
buffer_ptrs[nvl_rank],
|
|
0,
|
|
num_nvl_bytes + fifo_bytes + buffer_ptr_bytes + task_ptr_bytes,
|
|
comm_stream));
|
|
}
|
|
|
|
// Create 32 MiB workspace
|
|
// Note(ZKK): here we allocate more(2 * M2N_NUM_WORKSPACE) to support M2N!
|
|
// Later we will optimize here!
|
|
CUDA_CHECK(
|
|
cudaMalloc(&workspace, 2 * M2N_NUM_WORKSPACE * NUM_WORKSPACE_BYTES));
|
|
CUDA_CHECK(cudaMemsetAsync(
|
|
workspace, 0, 2 * M2N_NUM_WORKSPACE * NUM_WORKSPACE_BYTES, comm_stream));
|
|
|
|
// MoE counter
|
|
CUDA_CHECK(
|
|
cudaMallocHost(&moe_recv_counter, sizeof(int64_t), cudaHostAllocMapped));
|
|
CUDA_CHECK(cudaHostGetDevicePointer(
|
|
&moe_recv_counter_mapped, const_cast<int*>(moe_recv_counter), 0));
|
|
*moe_recv_counter = -1;
|
|
|
|
// MoE expert-level counter
|
|
CUDA_CHECK(cudaMallocHost(&moe_recv_expert_counter,
|
|
sizeof(int) * NUM_MAX_LOCAL_EXPERTS,
|
|
cudaHostAllocMapped));
|
|
CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_expert_counter_mapped,
|
|
const_cast<int*>(moe_recv_expert_counter),
|
|
0));
|
|
for (int i = 0; i < NUM_MAX_LOCAL_EXPERTS; ++i)
|
|
moe_recv_expert_counter[i] = -1;
|
|
|
|
// MoE RDMA-level counter
|
|
if (num_rdma_ranks > 0) {
|
|
CUDA_CHECK(cudaMallocHost(
|
|
&moe_recv_rdma_counter, sizeof(int), cudaHostAllocMapped));
|
|
CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_rdma_counter_mapped,
|
|
const_cast<int*>(moe_recv_rdma_counter),
|
|
0));
|
|
*moe_recv_rdma_counter = -1;
|
|
}
|
|
}
|
|
|
|
Buffer::~Buffer() noexcept(false) {
|
|
// Synchronize
|
|
CUDA_CHECK(cudaDeviceSynchronize());
|
|
printf("Buffer::~Buffer begin!!!\n");
|
|
if (num_nvl_bytes > 0) {
|
|
// Barrier
|
|
intranode::barrier(
|
|
task_fifo_ptrs_gpu, head, nvl_rank, num_nvl_ranks, comm_stream);
|
|
move_fifo_slots();
|
|
CUDA_CHECK(cudaDeviceSynchronize());
|
|
|
|
// Close remote IPC
|
|
if (is_available()) {
|
|
for (int i = 0; i < num_nvl_ranks; ++i)
|
|
if (i != nvl_rank) CUDA_CHECK(cudaIpcCloseMemHandle(buffer_ptrs[i]));
|
|
}
|
|
|
|
// Free local buffer and error flag
|
|
CUDA_CHECK(cudaFree(buffer_ptrs[nvl_rank]));
|
|
}
|
|
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
// Free NVSHMEM
|
|
if (num_rdma_bytes > 0) {
|
|
CUDA_CHECK(cudaDeviceSynchronize());
|
|
internode::barrier();
|
|
internode::free(rdma_buffer_ptr);
|
|
internode::finalize();
|
|
}
|
|
#endif
|
|
|
|
// Free cuBLAS handle, workspace and MoE counter
|
|
CUDA_CHECK(cudaFree(workspace));
|
|
CUDA_CHECK(cudaFreeHost(const_cast<int*>(moe_recv_counter)));
|
|
|
|
// Free chunked mode staffs
|
|
CUDA_CHECK(cudaFreeHost(const_cast<int*>(moe_recv_expert_counter)));
|
|
}
|
|
|
|
void Buffer::move_fifo_slots(int num_slots) {
|
|
head = (head + num_ranks * num_slots) % NUM_MAX_FIFO_SLOTS;
|
|
}
|
|
|
|
bool Buffer::is_available() const { return available; }
|
|
|
|
bool Buffer::is_internode_available() const {
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
return is_available() && num_ranks > NUM_MAX_NVL_PEERS;
|
|
#else
|
|
return false;
|
|
#endif
|
|
}
|
|
|
|
int Buffer::get_num_rdma_ranks() const { return num_rdma_ranks; }
|
|
|
|
int Buffer::get_rdma_rank() const { return rdma_rank; }
|
|
|
|
int Buffer::get_root_rdma_rank(bool global) const {
|
|
return global ? nvl_rank : 0;
|
|
}
|
|
|
|
int Buffer::get_local_device_id() const { return device_id; }
|
|
|
|
cudaStream_t Buffer::get_comm_stream() const { return comm_stream; }
|
|
|
|
#ifndef PADDLE_NO_PYTHON
|
|
pybind11::bytearray Buffer::get_local_ipc_handle() const {
|
|
return {ipc_handles[nvl_rank].reserved, CUDA_IPC_HANDLE_SIZE};
|
|
}
|
|
|
|
pybind11::bytearray Buffer::get_local_nvshmem_unique_id() const {
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
EP_HOST_ASSERT(rdma_rank == 0 &&
|
|
"Only RDMA rank 0 can get NVSHMEM unique ID");
|
|
auto unique_id = internode::get_unique_id();
|
|
#else
|
|
LOG(ERROR) << "NVSHMEM is not enabled. You can enable it by setting cmake "
|
|
"option WITH_NVSHMEM=ON.";
|
|
std::vector<uint8_t> unique_id;
|
|
#endif
|
|
return {reinterpret_cast<const char*>(unique_id.data()), unique_id.size()};
|
|
}
|
|
|
|
void Buffer::sync(
|
|
const std::vector<int>& device_ids,
|
|
const std::vector<std::optional<pybind11::bytearray>>& all_gathered_handles,
|
|
const std::optional<pybind11::bytearray>& root_unique_id_opt) {
|
|
EP_HOST_ASSERT(!is_available());
|
|
|
|
// Sync IPC handles
|
|
if (num_nvl_bytes > 0) {
|
|
EP_HOST_ASSERT(num_ranks == static_cast<int64_t>(device_ids.size()));
|
|
EP_HOST_ASSERT(device_ids.size() == all_gathered_handles.size());
|
|
for (int i = 0, offset = rdma_rank * num_nvl_ranks; i < num_nvl_ranks;
|
|
++i) {
|
|
EP_HOST_ASSERT(all_gathered_handles[offset + i].has_value());
|
|
auto handle_str = std::string(all_gathered_handles[offset + i].value());
|
|
EP_HOST_ASSERT(handle_str.size() == CUDA_IPC_HANDLE_SIZE);
|
|
if (offset + i != rank) {
|
|
std::memcpy(
|
|
ipc_handles[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE);
|
|
CUDA_CHECK(cudaIpcOpenMemHandle(
|
|
&buffer_ptrs[i], ipc_handles[i], cudaIpcMemLazyEnablePeerAccess));
|
|
task_fifo_ptrs[i] = reinterpret_cast<int*>(
|
|
reinterpret_cast<uint8_t*>(buffer_ptrs[i]) + num_nvl_bytes);
|
|
} else {
|
|
EP_HOST_ASSERT(std::memcmp(ipc_handles[i].reserved,
|
|
handle_str.c_str(),
|
|
CUDA_IPC_HANDLE_SIZE) == 0);
|
|
}
|
|
}
|
|
|
|
// Copy all buffer and task pointers to GPU
|
|
CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu,
|
|
buffer_ptrs,
|
|
sizeof(void*) * NUM_MAX_NVL_PEERS,
|
|
cudaMemcpyHostToDevice));
|
|
CUDA_CHECK(cudaMemcpy(task_fifo_ptrs_gpu,
|
|
task_fifo_ptrs,
|
|
sizeof(int*) * NUM_MAX_NVL_PEERS,
|
|
cudaMemcpyHostToDevice));
|
|
CUDA_CHECK(cudaDeviceSynchronize());
|
|
}
|
|
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
// Sync NVSHMEM handles and allocate memory
|
|
if (num_rdma_bytes > 0) {
|
|
// Initialize NVSHMEM
|
|
EP_HOST_ASSERT(root_unique_id_opt.has_value());
|
|
std::vector<uint8_t> root_unique_id(root_unique_id_opt->size());
|
|
auto root_unique_id_str = root_unique_id_opt->cast<std::string>();
|
|
std::memcpy(root_unique_id.data(),
|
|
root_unique_id_str.c_str(),
|
|
root_unique_id_opt->size());
|
|
auto nvshmem_rank = low_latency_mode ? rank : rdma_rank;
|
|
auto num_nvshmem_ranks = low_latency_mode ? num_ranks : num_rdma_ranks;
|
|
EP_HOST_ASSERT(nvshmem_rank == internode::init(root_unique_id,
|
|
nvshmem_rank,
|
|
num_nvshmem_ranks,
|
|
low_latency_mode));
|
|
internode::barrier();
|
|
|
|
// Allocate
|
|
rdma_buffer_ptr =
|
|
internode::alloc(num_rdma_bytes, NUM_BUFFER_ALIGNMENT_BYTES);
|
|
|
|
// Clean buffer (mainly for low-latency mode)
|
|
CUDA_CHECK(cudaMemset(rdma_buffer_ptr, 0, num_rdma_bytes));
|
|
|
|
// Barrier
|
|
internode::barrier();
|
|
CUDA_CHECK(cudaDeviceSynchronize());
|
|
}
|
|
#endif
|
|
|
|
// Ready to use
|
|
available = true;
|
|
}
|
|
#endif
|
|
|
|
std::tuple<deep_ep::detail::Tensor,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
std::optional<EventHandle>>
|
|
Buffer::get_dispatch_layout(const deep_ep::detail::Tensor& topk_idx,
|
|
int num_experts,
|
|
std::optional<EventHandle>& previous_event,
|
|
bool async,
|
|
bool allocate_on_comm_stream) {
|
|
EP_HOST_ASSERT(topk_idx.dim() == 2);
|
|
EP_HOST_ASSERT(topk_idx.is_contiguous());
|
|
EP_HOST_ASSERT(num_experts > 0);
|
|
|
|
// Allocate all tensors on comm stream if set
|
|
// NOTES: do not allocate tensors upfront!
|
|
auto compute_stream = calc_ctx->stream();
|
|
if (allocate_on_comm_stream) {
|
|
EP_HOST_ASSERT(previous_event.has_value() && async);
|
|
deep_ep::detail::SetAllocatorStreamForGPUContext(comm_stream, calc_ctx);
|
|
}
|
|
|
|
// Wait previous tasks to be finished
|
|
if (previous_event.has_value()) {
|
|
stream_wait(comm_stream, previous_event.value());
|
|
} else {
|
|
stream_wait(comm_stream, compute_stream);
|
|
}
|
|
|
|
auto num_tokens = static_cast<int>(topk_idx.size(0)),
|
|
num_topk = static_cast<int>(topk_idx.size(1));
|
|
auto num_tokens_per_rank =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_ranks}, phi::DataType::INT32, phi::GPUPlace(device_id)));
|
|
auto num_tokens_per_rdma_rank = std::optional<deep_ep::detail::Tensor>();
|
|
auto num_tokens_per_expert =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_experts}, phi::DataType::INT32, phi::GPUPlace(device_id)));
|
|
auto is_token_in_rank = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_tokens, num_ranks},
|
|
phi::DataType::BOOL,
|
|
phi::GPUPlace(device_id)));
|
|
if (is_internode_available())
|
|
num_tokens_per_rdma_rank =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_rdma_ranks}, phi::DataType::INT32, phi::GPUPlace(device_id)));
|
|
|
|
// get_dispatch_layout is used for both intranode and internode.
|
|
internode::get_dispatch_layout(
|
|
topk_idx.data_ptr<int64_t>(),
|
|
num_tokens_per_rank.data_ptr<int>(),
|
|
num_tokens_per_rdma_rank.has_value()
|
|
? num_tokens_per_rdma_rank.value().data_ptr<int>()
|
|
: nullptr,
|
|
num_tokens_per_expert.data_ptr<int>(),
|
|
is_token_in_rank.data_ptr<bool>(),
|
|
num_tokens,
|
|
num_topk,
|
|
num_ranks,
|
|
num_experts,
|
|
comm_stream);
|
|
|
|
// Wait streams
|
|
std::optional<EventHandle> event;
|
|
if (async) {
|
|
event = EventHandle(comm_stream);
|
|
for (auto& t : {topk_idx,
|
|
num_tokens_per_rank,
|
|
num_tokens_per_expert,
|
|
is_token_in_rank}) {
|
|
t.record_stream(comm_stream);
|
|
if (allocate_on_comm_stream) t.record_stream(compute_stream);
|
|
}
|
|
for (auto& to : {num_tokens_per_rdma_rank}) {
|
|
to.has_value() ? to->record_stream(comm_stream) : void();
|
|
if (allocate_on_comm_stream)
|
|
to.has_value() ? to->record_stream(compute_stream) : void();
|
|
}
|
|
} else {
|
|
stream_wait(compute_stream, comm_stream);
|
|
}
|
|
|
|
// Switch back compute stream
|
|
if (allocate_on_comm_stream) {
|
|
deep_ep::detail::SetAllocatorStreamForGPUContext(compute_stream, calc_ctx);
|
|
}
|
|
|
|
return {num_tokens_per_rank,
|
|
num_tokens_per_rdma_rank,
|
|
num_tokens_per_expert,
|
|
is_token_in_rank,
|
|
event};
|
|
}
|
|
|
|
std::tuple<deep_ep::detail::Tensor,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
std::vector<int>,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
std::optional<EventHandle>>
|
|
Buffer::intranode_dispatch(
|
|
const deep_ep::detail::Tensor& x,
|
|
const std::optional<deep_ep::detail::Tensor>& x_scales,
|
|
const std::optional<deep_ep::detail::Tensor>& topk_idx,
|
|
const std::optional<deep_ep::detail::Tensor>& topk_weights,
|
|
const std::optional<deep_ep::detail::Tensor>& num_tokens_per_rank,
|
|
const deep_ep::detail::Tensor& is_token_in_rank,
|
|
const std::optional<deep_ep::detail::Tensor>& num_tokens_per_expert,
|
|
int cached_num_recv_tokens,
|
|
const std::optional<deep_ep::detail::Tensor>& cached_rank_prefix_matrix,
|
|
const std::optional<deep_ep::detail::Tensor>& cached_channel_prefix_matrix,
|
|
int expert_alignment,
|
|
const Config& config,
|
|
std::optional<EventHandle>& previous_event, // NOLINT
|
|
bool async,
|
|
bool allocate_on_comm_stream) {
|
|
bool cached_mode = cached_rank_prefix_matrix.has_value();
|
|
|
|
// One channel use two blocks, even-numbered blocks for sending, odd-numbered
|
|
// blocks for receiving.
|
|
EP_HOST_ASSERT(config.num_sms % 2 == 0);
|
|
int num_channels = config.num_sms / 2;
|
|
if (cached_mode) {
|
|
EP_HOST_ASSERT(cached_rank_prefix_matrix.has_value());
|
|
EP_HOST_ASSERT(cached_channel_prefix_matrix.has_value());
|
|
} else {
|
|
EP_HOST_ASSERT(num_tokens_per_rank.has_value());
|
|
EP_HOST_ASSERT(num_tokens_per_expert.has_value());
|
|
}
|
|
|
|
// Type checks
|
|
EP_HOST_ASSERT(is_token_in_rank.scalar_type() == deep_ep::detail::kBool);
|
|
if (cached_mode) {
|
|
EP_HOST_ASSERT(cached_rank_prefix_matrix->scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(cached_channel_prefix_matrix->scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
} else {
|
|
EP_HOST_ASSERT(num_tokens_per_expert->scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(num_tokens_per_rank->scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
}
|
|
|
|
// Shape and contiguous checks
|
|
EP_HOST_ASSERT(x.dim() == 2 && x.is_contiguous());
|
|
EP_HOST_ASSERT((x.size(1) * x.element_size()) % sizeof(int4) == 0);
|
|
EP_HOST_ASSERT(is_token_in_rank.dim() == 2 &&
|
|
is_token_in_rank.is_contiguous());
|
|
EP_HOST_ASSERT(is_token_in_rank.size(0) == x.size(0) &&
|
|
is_token_in_rank.size(1) == num_ranks);
|
|
if (cached_mode) {
|
|
EP_HOST_ASSERT(cached_rank_prefix_matrix->dim() == 2 &&
|
|
cached_rank_prefix_matrix->is_contiguous());
|
|
EP_HOST_ASSERT(cached_rank_prefix_matrix->size(0) == num_ranks &&
|
|
cached_rank_prefix_matrix->size(1) == num_ranks);
|
|
EP_HOST_ASSERT(cached_channel_prefix_matrix->dim() == 2 &&
|
|
cached_channel_prefix_matrix->is_contiguous());
|
|
EP_HOST_ASSERT(cached_channel_prefix_matrix->size(0) == num_ranks &&
|
|
cached_channel_prefix_matrix->size(1) == num_channels);
|
|
} else {
|
|
EP_HOST_ASSERT(num_tokens_per_expert->dim() == 1 &&
|
|
num_tokens_per_expert->is_contiguous());
|
|
EP_HOST_ASSERT(num_tokens_per_expert->size(0) % num_ranks == 0);
|
|
EP_HOST_ASSERT(num_tokens_per_expert->size(0) / num_ranks <=
|
|
NUM_MAX_LOCAL_EXPERTS);
|
|
EP_HOST_ASSERT(num_tokens_per_rank->dim() == 1 &&
|
|
num_tokens_per_rank->is_contiguous());
|
|
EP_HOST_ASSERT(num_tokens_per_rank->size(0) == num_ranks);
|
|
}
|
|
|
|
auto num_tokens = static_cast<int>(x.size(0)),
|
|
hidden = static_cast<int>(x.size(1));
|
|
auto num_experts =
|
|
cached_mode ? 0 : static_cast<int>(num_tokens_per_expert->size(0)),
|
|
num_local_experts = num_experts / num_ranks;
|
|
|
|
// Top-k checks
|
|
int num_topk = 0;
|
|
int64_t* topk_idx_ptr = nullptr;
|
|
float* topk_weights_ptr = nullptr;
|
|
EP_HOST_ASSERT(topk_idx.has_value() == topk_weights.has_value());
|
|
if (topk_idx.has_value()) {
|
|
num_topk = static_cast<int>(topk_idx->size(1));
|
|
EP_HOST_ASSERT(num_experts > 0);
|
|
EP_HOST_ASSERT(topk_idx->dim() == 2 && topk_idx->is_contiguous());
|
|
EP_HOST_ASSERT(topk_weights->dim() == 2 && topk_weights->is_contiguous());
|
|
EP_HOST_ASSERT(num_tokens == topk_idx->size(0) &&
|
|
num_tokens == topk_weights->size(0));
|
|
EP_HOST_ASSERT(num_topk == topk_weights->size(1));
|
|
EP_HOST_ASSERT(topk_weights->scalar_type() == deep_ep::detail::kFloat32);
|
|
topk_idx_ptr = topk_idx->data_ptr<int64_t>();
|
|
topk_weights_ptr = topk_weights->data_ptr<float>();
|
|
}
|
|
|
|
// FP8 scales checks
|
|
float* x_scales_ptr = nullptr;
|
|
int num_scales = 0;
|
|
if (x_scales.has_value()) {
|
|
EP_HOST_ASSERT(x.element_size() == 1);
|
|
EP_HOST_ASSERT(x_scales->scalar_type() == deep_ep::detail::kFloat32);
|
|
EP_HOST_ASSERT(x_scales->dim() > 0 && x_scales->dim() < 3 &&
|
|
x_scales->is_contiguous());
|
|
EP_HOST_ASSERT(x_scales->size(0) == num_tokens);
|
|
num_scales = x_scales->dim() == 1 ? 1 : static_cast<int>(x_scales->size(1));
|
|
x_scales_ptr = x_scales->data_ptr<float>();
|
|
}
|
|
|
|
// Allocate all tensors on comm stream if set
|
|
// NOTES: do not allocate tensors upfront!
|
|
auto compute_stream = calc_ctx->stream();
|
|
if (allocate_on_comm_stream) {
|
|
EP_HOST_ASSERT(previous_event.has_value() && async);
|
|
deep_ep::detail::SetAllocatorStreamForGPUContext(comm_stream, calc_ctx);
|
|
if (FLAGS_deep_ep_comm_prealloc_in_mb > 0)
|
|
std::call_once(
|
|
pre_alloc_once_flag, PreAlloc, x.raw_tensor(), comm_stream);
|
|
}
|
|
|
|
// Wait previous tasks to be finished
|
|
if (previous_event.has_value()) {
|
|
stream_wait(comm_stream, previous_event.value());
|
|
} else {
|
|
stream_wait(comm_stream, compute_stream);
|
|
}
|
|
|
|
// Create handles (only return for non-cached mode)
|
|
int num_recv_tokens = -1;
|
|
auto rank_prefix_matrix = deep_ep::detail::Tensor();
|
|
auto channel_prefix_matrix = deep_ep::detail::Tensor();
|
|
std::vector<int> num_recv_tokens_per_expert_list;
|
|
|
|
// Barrier or send sizes
|
|
// To clean: channel start/end offset, head and tail
|
|
int num_memset_int = num_channels * num_ranks * 4;
|
|
if (cached_mode) {
|
|
num_recv_tokens = cached_num_recv_tokens;
|
|
rank_prefix_matrix = cached_rank_prefix_matrix.value();
|
|
channel_prefix_matrix = cached_channel_prefix_matrix.value();
|
|
|
|
// Copy rank prefix matrix and clean flags
|
|
intranode::cached_notify_dispatch(rank_prefix_matrix.data_ptr<int>(),
|
|
num_memset_int,
|
|
buffer_ptrs_gpu,
|
|
task_fifo_ptrs_gpu,
|
|
head,
|
|
rank,
|
|
num_ranks,
|
|
comm_stream);
|
|
move_fifo_slots(2);
|
|
} else {
|
|
rank_prefix_matrix = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_ranks, num_ranks},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
channel_prefix_matrix = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_ranks, num_channels},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
|
|
// Send sizes
|
|
// Meta information:
|
|
// - Size prefix by ranks, shaped as `[num_ranks, num_ranks]`
|
|
// - Size prefix by experts (not used later), shaped as `[num_ranks,
|
|
// num_local_experts]`
|
|
// NOTES: no more token dropping in this version
|
|
*moe_recv_counter = -1;
|
|
for (int i = 0; i < num_local_experts; ++i) moe_recv_expert_counter[i] = -1;
|
|
EP_HOST_ASSERT(num_ranks * (num_ranks + num_local_experts) *
|
|
static_cast<int64_t>(sizeof(int)) <=
|
|
num_nvl_bytes);
|
|
intranode::notify_dispatch(num_tokens_per_rank->data_ptr<int>(),
|
|
moe_recv_counter_mapped,
|
|
num_ranks,
|
|
num_tokens_per_expert->data_ptr<int>(),
|
|
moe_recv_expert_counter_mapped,
|
|
num_experts,
|
|
num_tokens,
|
|
is_token_in_rank.data_ptr<bool>(),
|
|
channel_prefix_matrix.data_ptr<int>(),
|
|
rank_prefix_matrix.data_ptr<int>(),
|
|
num_memset_int,
|
|
expert_alignment,
|
|
buffer_ptrs_gpu,
|
|
task_fifo_ptrs_gpu,
|
|
head,
|
|
rank,
|
|
comm_stream,
|
|
num_channels);
|
|
move_fifo_slots(3);
|
|
|
|
// Synchronize total received tokens and tokens per expert
|
|
auto start_time = std::chrono::high_resolution_clock::now();
|
|
while (true) {
|
|
// Read total count
|
|
num_recv_tokens = static_cast<int>(*moe_recv_counter);
|
|
|
|
// Read per-expert count
|
|
bool ready = (num_recv_tokens >= 0);
|
|
for (int i = 0; i < num_local_experts && ready; ++i)
|
|
ready &= moe_recv_expert_counter[i] >= 0;
|
|
|
|
if (ready) break;
|
|
|
|
// Timeout check
|
|
if (std::chrono::duration_cast<std::chrono::seconds>(
|
|
std::chrono::high_resolution_clock::now() - start_time)
|
|
.count() > NUM_CPU_TIMEOUT_SECS)
|
|
throw std::runtime_error("DeepEP error: CPU recv timeout");
|
|
}
|
|
num_recv_tokens_per_expert_list = std::vector<int>(
|
|
moe_recv_expert_counter, moe_recv_expert_counter + num_local_experts);
|
|
}
|
|
|
|
// Allocate new tensors
|
|
auto recv_x = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_recv_tokens, hidden}, x.dtype(), x.place()));
|
|
auto recv_src_idx =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_recv_tokens}, phi::DataType::INT32, phi::GPUPlace(device_id)));
|
|
auto recv_topk_idx = std::optional<deep_ep::detail::Tensor>(),
|
|
recv_topk_weights = std::optional<deep_ep::detail::Tensor>(),
|
|
recv_x_scales = std::optional<deep_ep::detail::Tensor>();
|
|
auto recv_channel_prefix_matrix = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_ranks, num_channels},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
auto send_head = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_tokens, num_ranks},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
|
|
// Assign pointers
|
|
int64_t* recv_topk_idx_ptr = nullptr;
|
|
float* recv_topk_weights_ptr = nullptr;
|
|
float* recv_x_scales_ptr = nullptr;
|
|
if (topk_idx.has_value()) {
|
|
recv_topk_idx =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_recv_tokens, num_topk}, topk_idx->dtype(), topk_idx->place()));
|
|
recv_topk_weights = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_recv_tokens, num_topk},
|
|
topk_weights->dtype(),
|
|
topk_idx->place()));
|
|
recv_topk_idx_ptr = recv_topk_idx->data_ptr<int64_t>();
|
|
recv_topk_weights_ptr = recv_topk_weights->data_ptr<float>();
|
|
}
|
|
if (x_scales.has_value()) {
|
|
recv_x_scales =
|
|
x_scales->dim() == 1
|
|
? ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_recv_tokens}, x_scales->dtype(), x_scales->place()))
|
|
: ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_recv_tokens, num_scales},
|
|
x_scales->dtype(),
|
|
x_scales->place()));
|
|
|
|
recv_x_scales_ptr = recv_x_scales->data_ptr<float>();
|
|
}
|
|
|
|
// Dispatch
|
|
EP_HOST_ASSERT(
|
|
num_ranks * num_ranks *
|
|
static_cast<int64_t>(sizeof(int)) + // prefix matrix
|
|
num_channels * num_ranks *
|
|
static_cast<int64_t>(sizeof(int)) + // Channel start offset
|
|
num_channels * num_ranks *
|
|
static_cast<int64_t>(sizeof(int)) + // Channel end offset
|
|
num_channels * num_ranks * static_cast<int64_t>(sizeof(int)) *
|
|
2 + // Queue head and tail
|
|
num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens *
|
|
hidden * recv_x.element_size() + // Data buffer
|
|
num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens *
|
|
static_cast<int64_t>(sizeof(int)) + // Source index buffer
|
|
num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens *
|
|
num_topk *
|
|
static_cast<int64_t>(sizeof(int64_t)) + // Top-k index buffer
|
|
num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens *
|
|
num_topk *
|
|
static_cast<int64_t>(sizeof(float)) + // Top-k weight buffer
|
|
num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens *
|
|
static_cast<int64_t>(sizeof(float)) *
|
|
num_scales // FP8 scale buffer
|
|
<= num_nvl_bytes);
|
|
intranode::dispatch(
|
|
recv_x.data_ptr(),
|
|
recv_x_scales_ptr,
|
|
recv_src_idx.data_ptr<int>(),
|
|
recv_topk_idx_ptr,
|
|
recv_topk_weights_ptr,
|
|
recv_channel_prefix_matrix.data_ptr<int>(),
|
|
send_head.data_ptr<int>(),
|
|
x.data_ptr(),
|
|
x_scales_ptr,
|
|
topk_idx_ptr,
|
|
topk_weights_ptr,
|
|
is_token_in_rank.data_ptr<bool>(),
|
|
channel_prefix_matrix.data_ptr<int>(),
|
|
num_tokens,
|
|
static_cast<int>(hidden * recv_x.element_size() / sizeof(int4)),
|
|
num_topk,
|
|
num_experts,
|
|
num_scales,
|
|
buffer_ptrs_gpu,
|
|
rank,
|
|
num_ranks,
|
|
comm_stream,
|
|
config.num_sms,
|
|
config.num_max_nvl_chunked_send_tokens,
|
|
config.num_max_nvl_chunked_recv_tokens);
|
|
|
|
// Wait streams
|
|
std::optional<EventHandle> event;
|
|
if (async) {
|
|
event = EventHandle(comm_stream);
|
|
for (auto& t : {x,
|
|
is_token_in_rank,
|
|
rank_prefix_matrix,
|
|
channel_prefix_matrix,
|
|
recv_x,
|
|
recv_src_idx,
|
|
recv_channel_prefix_matrix,
|
|
send_head}) {
|
|
t.record_stream(comm_stream);
|
|
if (allocate_on_comm_stream) t.record_stream(compute_stream);
|
|
}
|
|
for (auto& to : {x_scales,
|
|
topk_idx,
|
|
topk_weights,
|
|
num_tokens_per_rank,
|
|
num_tokens_per_expert,
|
|
cached_channel_prefix_matrix,
|
|
cached_rank_prefix_matrix,
|
|
recv_topk_idx,
|
|
recv_topk_weights,
|
|
recv_x_scales}) {
|
|
to.has_value() ? to->record_stream(comm_stream) : void();
|
|
if (allocate_on_comm_stream)
|
|
to.has_value() ? to->record_stream(compute_stream) : void();
|
|
}
|
|
} else {
|
|
stream_wait(compute_stream, comm_stream);
|
|
}
|
|
|
|
// Switch back compute stream
|
|
if (allocate_on_comm_stream) {
|
|
deep_ep::detail::SetAllocatorStreamForGPUContext(compute_stream, calc_ctx);
|
|
}
|
|
|
|
// Return values
|
|
return {recv_x,
|
|
recv_x_scales,
|
|
recv_topk_idx,
|
|
recv_topk_weights,
|
|
num_recv_tokens_per_expert_list,
|
|
rank_prefix_matrix,
|
|
channel_prefix_matrix,
|
|
recv_channel_prefix_matrix,
|
|
recv_src_idx,
|
|
send_head,
|
|
event};
|
|
}
|
|
|
|
std::tuple<deep_ep::detail::Tensor,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
std::optional<EventHandle>>
|
|
Buffer::intranode_combine(
|
|
const deep_ep::detail::Tensor& x,
|
|
const std::optional<deep_ep::detail::Tensor>& topk_weights,
|
|
const deep_ep::detail::Tensor& src_idx,
|
|
const deep_ep::detail::Tensor& rank_prefix_matrix,
|
|
const deep_ep::detail::Tensor& channel_prefix_matrix,
|
|
const deep_ep::detail::Tensor& send_head,
|
|
const Config& config,
|
|
std::optional<EventHandle>& previous_event,
|
|
bool async,
|
|
bool allocate_on_comm_stream) {
|
|
EP_HOST_ASSERT(x.dim() == 2 && x.is_contiguous());
|
|
EP_HOST_ASSERT(src_idx.dim() == 1 && src_idx.is_contiguous() &&
|
|
src_idx.scalar_type() == deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(send_head.dim() == 2 && send_head.is_contiguous() &&
|
|
send_head.scalar_type() == deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(rank_prefix_matrix.dim() == 2 &&
|
|
rank_prefix_matrix.is_contiguous() &&
|
|
rank_prefix_matrix.scalar_type() == deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(channel_prefix_matrix.dim() == 2 &&
|
|
channel_prefix_matrix.is_contiguous() &&
|
|
channel_prefix_matrix.scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
|
|
// One channel use two blocks, even-numbered blocks for sending, odd-numbered
|
|
// blocks for receiving.
|
|
EP_HOST_ASSERT(config.num_sms % 2 == 0);
|
|
int num_channels = config.num_sms / 2;
|
|
|
|
auto num_tokens = static_cast<int>(x.size(0)),
|
|
hidden = static_cast<int>(x.size(1));
|
|
auto num_recv_tokens = static_cast<int>(send_head.size(0));
|
|
EP_HOST_ASSERT(src_idx.size(0) == num_tokens);
|
|
EP_HOST_ASSERT(send_head.size(1) == num_ranks);
|
|
EP_HOST_ASSERT(rank_prefix_matrix.size(0) == num_ranks &&
|
|
rank_prefix_matrix.size(1) == num_ranks);
|
|
EP_HOST_ASSERT(channel_prefix_matrix.size(0) == num_ranks &&
|
|
channel_prefix_matrix.size(1) == num_channels);
|
|
EP_HOST_ASSERT((hidden * x.element_size()) % sizeof(int4) == 0);
|
|
|
|
// Allocate all tensors on comm stream if set
|
|
// NOTES: do not allocate tensors upfront!
|
|
auto compute_stream = calc_ctx->stream();
|
|
if (allocate_on_comm_stream) {
|
|
EP_HOST_ASSERT(previous_event.has_value() && async);
|
|
deep_ep::detail::SetAllocatorStreamForGPUContext(comm_stream, calc_ctx);
|
|
}
|
|
|
|
// Wait previous tasks to be finished
|
|
if (previous_event.has_value()) {
|
|
stream_wait(comm_stream, previous_event.value());
|
|
} else {
|
|
stream_wait(comm_stream, compute_stream);
|
|
}
|
|
|
|
int num_topk = 0;
|
|
auto recv_topk_weights = std::optional<deep_ep::detail::Tensor>();
|
|
float* topk_weights_ptr = nullptr;
|
|
float* recv_topk_weights_ptr = nullptr;
|
|
if (topk_weights.has_value()) {
|
|
EP_HOST_ASSERT(topk_weights->dim() == 2 && topk_weights->is_contiguous());
|
|
EP_HOST_ASSERT(topk_weights->size(0) == num_tokens);
|
|
EP_HOST_ASSERT(topk_weights->scalar_type() == deep_ep::detail::kFloat32);
|
|
num_topk = static_cast<int>(topk_weights->size(1));
|
|
topk_weights_ptr = topk_weights->data_ptr<float>();
|
|
recv_topk_weights = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_recv_tokens, num_topk},
|
|
topk_weights->dtype(),
|
|
topk_weights->place()));
|
|
recv_topk_weights_ptr = recv_topk_weights->data_ptr<float>();
|
|
}
|
|
|
|
// Launch barrier and reset queue head and tail
|
|
EP_HOST_ASSERT(num_channels * num_ranks * static_cast<int64_t>(sizeof(int)) *
|
|
2 <=
|
|
num_nvl_bytes);
|
|
intranode::cached_notify_combine(buffer_ptrs_gpu,
|
|
send_head.data_ptr<int>(),
|
|
num_channels,
|
|
num_recv_tokens,
|
|
num_channels * num_ranks * 2,
|
|
task_fifo_ptrs_gpu,
|
|
head,
|
|
rank,
|
|
num_ranks,
|
|
comm_stream);
|
|
|
|
// NOTES: this function uses two FIFO slots (barrier before and after)
|
|
move_fifo_slots(2);
|
|
|
|
// Combine data
|
|
auto recv_x = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_recv_tokens, hidden}, x.dtype(), x.place()));
|
|
EP_HOST_ASSERT(
|
|
num_channels * num_ranks * static_cast<int64_t>(sizeof(int)) *
|
|
2 + // Queue head and tail
|
|
num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens *
|
|
hidden * x.element_size() + // Data buffer
|
|
num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens *
|
|
static_cast<int64_t>(sizeof(int)) + // Source index buffer
|
|
num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens *
|
|
num_topk *
|
|
static_cast<int64_t>(sizeof(float)) // Top-k weight buffer
|
|
<= num_nvl_bytes);
|
|
intranode::combine(deep_ep::detail::ScalarTypeToCudaDataType(x.scalar_type()),
|
|
recv_x.data_ptr(),
|
|
recv_topk_weights_ptr,
|
|
x.data_ptr(),
|
|
topk_weights_ptr,
|
|
src_idx.data_ptr<int>(),
|
|
rank_prefix_matrix.data_ptr<int>(),
|
|
channel_prefix_matrix.data_ptr<int>(),
|
|
send_head.data_ptr<int>(),
|
|
num_tokens,
|
|
num_recv_tokens,
|
|
hidden,
|
|
num_topk,
|
|
buffer_ptrs_gpu,
|
|
rank,
|
|
num_ranks,
|
|
comm_stream,
|
|
config.num_sms,
|
|
config.num_max_nvl_chunked_send_tokens,
|
|
config.num_max_nvl_chunked_recv_tokens);
|
|
|
|
// Wait streams
|
|
std::optional<EventHandle> event;
|
|
if (async) {
|
|
event = EventHandle(comm_stream);
|
|
for (auto& t : {x,
|
|
src_idx,
|
|
send_head,
|
|
rank_prefix_matrix,
|
|
channel_prefix_matrix,
|
|
recv_x}) {
|
|
t.record_stream(comm_stream);
|
|
if (allocate_on_comm_stream) t.record_stream(compute_stream);
|
|
}
|
|
for (auto& to : {topk_weights, recv_topk_weights}) {
|
|
to.has_value() ? to->record_stream(comm_stream) : void();
|
|
if (allocate_on_comm_stream)
|
|
to.has_value() ? to->record_stream(compute_stream) : void();
|
|
}
|
|
} else {
|
|
stream_wait(compute_stream, comm_stream);
|
|
}
|
|
|
|
// Switch back compute stream
|
|
if (allocate_on_comm_stream) {
|
|
deep_ep::detail::SetAllocatorStreamForGPUContext(compute_stream, calc_ctx);
|
|
}
|
|
|
|
return {recv_x, recv_topk_weights, event};
|
|
}
|
|
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
std::tuple<deep_ep::detail::Tensor,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
std::vector<int>,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
deep_ep::detail::Tensor,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
deep_ep::detail::Tensor,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
std::optional<EventHandle>>
|
|
Buffer::internode_dispatch(
|
|
const deep_ep::detail::Tensor& x,
|
|
const std::optional<deep_ep::detail::Tensor>& x_scales,
|
|
const std::optional<deep_ep::detail::Tensor>& topk_idx,
|
|
const std::optional<deep_ep::detail::Tensor>& topk_weights,
|
|
const std::optional<deep_ep::detail::Tensor>& num_tokens_per_rank,
|
|
const std::optional<deep_ep::detail::Tensor>& num_tokens_per_rdma_rank,
|
|
const deep_ep::detail::Tensor& is_token_in_rank,
|
|
const std::optional<deep_ep::detail::Tensor>& num_tokens_per_expert,
|
|
int cached_num_recv_tokens,
|
|
int cached_num_rdma_recv_tokens,
|
|
const std::optional<deep_ep::detail::Tensor>&
|
|
cached_rdma_channel_prefix_matrix,
|
|
const std::optional<deep_ep::detail::Tensor>&
|
|
cached_recv_rdma_rank_prefix_sum,
|
|
const std::optional<deep_ep::detail::Tensor>&
|
|
cached_gbl_channel_prefix_matrix,
|
|
const std::optional<deep_ep::detail::Tensor>&
|
|
cached_recv_gbl_rank_prefix_sum,
|
|
int expert_alignment,
|
|
const Config& config,
|
|
std::optional<EventHandle>& previous_event, // NOLINT
|
|
bool async,
|
|
bool allocate_on_comm_stream) {
|
|
// In dispatch, CPU will busy-wait until GPU receive tensor size metadata from
|
|
// other ranks, which can be quite long. If users of DeepEP need to execute
|
|
// other Python code on other threads, such as KV transfer, their code will
|
|
// get stuck due to GIL unless we release GIL here.
|
|
// pybind11::gil_scoped_release release;
|
|
|
|
const int num_channels = config.num_sms / 2;
|
|
EP_HOST_ASSERT(config.num_sms % 2 == 0);
|
|
EP_HOST_ASSERT(0 < get_num_rdma_ranks() &&
|
|
get_num_rdma_ranks() <= NUM_MAX_RDMA_PEERS);
|
|
|
|
bool cached_mode = cached_rdma_channel_prefix_matrix.has_value();
|
|
if (cached_mode) {
|
|
EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix.has_value());
|
|
EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum.has_value());
|
|
EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix.has_value());
|
|
EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum.has_value());
|
|
} else {
|
|
EP_HOST_ASSERT(num_tokens_per_rank.has_value());
|
|
EP_HOST_ASSERT(num_tokens_per_rdma_rank.has_value());
|
|
EP_HOST_ASSERT(num_tokens_per_expert.has_value());
|
|
}
|
|
|
|
// Type checks
|
|
if (cached_mode) {
|
|
EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix->scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum->scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix->scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum->scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
} else {
|
|
EP_HOST_ASSERT(num_tokens_per_rank->scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(num_tokens_per_rdma_rank->scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(num_tokens_per_expert->scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
}
|
|
|
|
// Shape and contiguous checks
|
|
EP_HOST_ASSERT(x.dim() == 2 && x.is_contiguous());
|
|
EP_HOST_ASSERT((x.size(1) * x.element_size()) % sizeof(int4) == 0);
|
|
if (cached_mode) {
|
|
EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix->dim() == 2 &&
|
|
cached_rdma_channel_prefix_matrix->is_contiguous());
|
|
EP_HOST_ASSERT(cached_rdma_channel_prefix_matrix->size(0) ==
|
|
num_rdma_ranks &&
|
|
cached_rdma_channel_prefix_matrix->size(1) == num_channels);
|
|
EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum->dim() == 1 &&
|
|
cached_recv_rdma_rank_prefix_sum->is_contiguous());
|
|
EP_HOST_ASSERT(cached_recv_rdma_rank_prefix_sum->size(0) == num_rdma_ranks);
|
|
EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix->dim() == 2 &&
|
|
cached_gbl_channel_prefix_matrix->is_contiguous());
|
|
EP_HOST_ASSERT(cached_gbl_channel_prefix_matrix->size(0) == num_ranks &&
|
|
cached_gbl_channel_prefix_matrix->size(1) == num_channels);
|
|
EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum->dim() == 1 &&
|
|
cached_recv_gbl_rank_prefix_sum->is_contiguous());
|
|
EP_HOST_ASSERT(cached_recv_gbl_rank_prefix_sum->size(0) == num_ranks);
|
|
} else {
|
|
EP_HOST_ASSERT(num_tokens_per_rank->dim() == 1 &&
|
|
num_tokens_per_rank->is_contiguous());
|
|
EP_HOST_ASSERT(num_tokens_per_rdma_rank->dim() == 1 &&
|
|
num_tokens_per_rdma_rank->is_contiguous());
|
|
EP_HOST_ASSERT(num_tokens_per_expert->dim() == 1 &&
|
|
num_tokens_per_expert->is_contiguous());
|
|
EP_HOST_ASSERT(num_tokens_per_rank->size(0) == num_ranks);
|
|
EP_HOST_ASSERT(num_tokens_per_rdma_rank->size(0) == num_rdma_ranks);
|
|
EP_HOST_ASSERT(num_tokens_per_expert->size(0) % num_ranks == 0);
|
|
EP_HOST_ASSERT(num_tokens_per_expert->size(0) / num_ranks <=
|
|
NUM_MAX_LOCAL_EXPERTS);
|
|
}
|
|
|
|
auto num_tokens = static_cast<int>(x.size(0)),
|
|
hidden = static_cast<int>(x.size(1)),
|
|
hidden_int4 =
|
|
static_cast<int>(x.size(1) * x.element_size() / sizeof(int4));
|
|
auto num_experts =
|
|
cached_mode ? 0 : static_cast<int>(num_tokens_per_expert->size(0)),
|
|
num_local_experts = num_experts / num_ranks;
|
|
|
|
// Top-k checks
|
|
int num_topk = 0;
|
|
int64_t* topk_idx_ptr = nullptr;
|
|
float* topk_weights_ptr = nullptr;
|
|
EP_HOST_ASSERT(topk_idx.has_value() == topk_weights.has_value());
|
|
if (topk_idx.has_value()) {
|
|
num_topk = static_cast<int>(topk_idx->size(1));
|
|
EP_HOST_ASSERT(num_experts > 0);
|
|
EP_HOST_ASSERT(topk_idx->dim() == 2 && topk_idx->is_contiguous());
|
|
EP_HOST_ASSERT(topk_weights->dim() == 2 && topk_weights->is_contiguous());
|
|
EP_HOST_ASSERT(num_tokens == topk_idx->size(0) &&
|
|
num_tokens == topk_weights->size(0));
|
|
EP_HOST_ASSERT(num_topk == topk_weights->size(1));
|
|
EP_HOST_ASSERT(topk_weights->scalar_type() == deep_ep::detail::kFloat32);
|
|
topk_idx_ptr = topk_idx->data_ptr<int64_t>();
|
|
topk_weights_ptr = topk_weights->data_ptr<float>();
|
|
}
|
|
|
|
// FP8 scales checks
|
|
float* x_scales_ptr = nullptr;
|
|
int num_scales = 0;
|
|
if (x_scales.has_value()) {
|
|
EP_HOST_ASSERT(x.element_size() == 1);
|
|
EP_HOST_ASSERT(x_scales->scalar_type() == deep_ep::detail::kFloat32);
|
|
EP_HOST_ASSERT(x_scales->dim() > 0 && x_scales->dim() < 3 &&
|
|
x_scales->is_contiguous());
|
|
EP_HOST_ASSERT(x_scales->size(0) == num_tokens);
|
|
num_scales = x_scales->dim() == 1 ? 1 : static_cast<int>(x_scales->size(1));
|
|
x_scales_ptr = x_scales->data_ptr<float>();
|
|
}
|
|
|
|
// Allocate all tensors on comm stream if set
|
|
// NOTES: do not allocate tensors upfront!
|
|
auto compute_stream = calc_ctx->stream();
|
|
if (allocate_on_comm_stream) {
|
|
EP_HOST_ASSERT(previous_event.has_value() && async);
|
|
deep_ep::detail::SetAllocatorStreamForGPUContext(comm_stream, calc_ctx);
|
|
if (FLAGS_deep_ep_comm_prealloc_in_mb > 0)
|
|
std::call_once(
|
|
pre_alloc_once_flag, PreAlloc, x.raw_tensor(), comm_stream);
|
|
}
|
|
|
|
// Wait previous tasks to be finished
|
|
if (previous_event.has_value()) {
|
|
stream_wait(comm_stream, previous_event.value());
|
|
} else {
|
|
stream_wait(comm_stream, compute_stream);
|
|
}
|
|
|
|
// Create handles (only return for non-cached mode)
|
|
int num_recv_tokens = -1, num_rdma_recv_tokens = -1;
|
|
auto rdma_channel_prefix_matrix = deep_ep::detail::Tensor();
|
|
auto recv_rdma_rank_prefix_sum = deep_ep::detail::Tensor();
|
|
auto gbl_channel_prefix_matrix = deep_ep::detail::Tensor();
|
|
auto recv_gbl_rank_prefix_sum = deep_ep::detail::Tensor();
|
|
std::vector<int> num_recv_tokens_per_expert_list;
|
|
|
|
// Barrier or send sizes
|
|
if (cached_mode) {
|
|
num_recv_tokens = cached_num_recv_tokens;
|
|
num_rdma_recv_tokens = cached_num_rdma_recv_tokens;
|
|
rdma_channel_prefix_matrix = cached_rdma_channel_prefix_matrix.value();
|
|
recv_rdma_rank_prefix_sum = cached_recv_rdma_rank_prefix_sum.value();
|
|
gbl_channel_prefix_matrix = cached_gbl_channel_prefix_matrix.value();
|
|
recv_gbl_rank_prefix_sum = cached_recv_gbl_rank_prefix_sum.value();
|
|
|
|
// Just a barrier and clean flags
|
|
internode::cached_notify(
|
|
hidden_int4,
|
|
num_scales,
|
|
num_topk,
|
|
num_topk,
|
|
num_ranks,
|
|
num_channels,
|
|
0,
|
|
nullptr,
|
|
nullptr,
|
|
nullptr,
|
|
nullptr,
|
|
rdma_buffer_ptr,
|
|
config.num_max_rdma_chunked_recv_tokens,
|
|
buffer_ptrs_gpu,
|
|
config.num_max_nvl_chunked_recv_tokens,
|
|
task_fifo_ptrs_gpu,
|
|
head,
|
|
rank,
|
|
comm_stream,
|
|
config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks),
|
|
num_nvl_bytes,
|
|
true,
|
|
low_latency_mode);
|
|
move_fifo_slots(2);
|
|
} else {
|
|
rdma_channel_prefix_matrix = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_rdma_ranks, num_channels},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
recv_rdma_rank_prefix_sum =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_rdma_ranks}, phi::DataType::INT32, phi::GPUPlace(device_id)));
|
|
gbl_channel_prefix_matrix = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_ranks, num_channels},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
recv_gbl_rank_prefix_sum =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_ranks}, phi::DataType::INT32, phi::GPUPlace(device_id)));
|
|
|
|
// Send sizes
|
|
*moe_recv_counter = -1, *moe_recv_rdma_counter = -1;
|
|
for (int i = 0; i < num_local_experts; ++i) moe_recv_expert_counter[i] = -1;
|
|
internode::notify_dispatch(
|
|
num_tokens_per_rank->data_ptr<int>(),
|
|
moe_recv_counter_mapped,
|
|
num_ranks,
|
|
num_tokens_per_rdma_rank->data_ptr<int>(),
|
|
moe_recv_rdma_counter_mapped,
|
|
num_tokens_per_expert->data_ptr<int>(),
|
|
moe_recv_expert_counter_mapped,
|
|
num_experts,
|
|
is_token_in_rank.data_ptr<bool>(),
|
|
num_tokens,
|
|
num_channels,
|
|
hidden_int4,
|
|
num_scales,
|
|
num_topk,
|
|
expert_alignment,
|
|
rdma_channel_prefix_matrix.data_ptr<int>(),
|
|
recv_rdma_rank_prefix_sum.data_ptr<int>(),
|
|
gbl_channel_prefix_matrix.data_ptr<int>(),
|
|
recv_gbl_rank_prefix_sum.data_ptr<int>(),
|
|
rdma_buffer_ptr,
|
|
config.num_max_rdma_chunked_recv_tokens,
|
|
buffer_ptrs_gpu,
|
|
config.num_max_nvl_chunked_recv_tokens,
|
|
task_fifo_ptrs_gpu,
|
|
head,
|
|
rank,
|
|
comm_stream,
|
|
config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks),
|
|
num_nvl_bytes,
|
|
low_latency_mode);
|
|
move_fifo_slots(3);
|
|
|
|
// Synchronize total received tokens and tokens per expert
|
|
auto start_time = std::chrono::high_resolution_clock::now();
|
|
while (true) {
|
|
// Read total count
|
|
num_recv_tokens = static_cast<int>(*moe_recv_counter);
|
|
num_rdma_recv_tokens = static_cast<int>(*moe_recv_rdma_counter);
|
|
|
|
// Read per-expert count
|
|
bool ready = (num_recv_tokens >= 0) && (num_rdma_recv_tokens >= 0);
|
|
for (int i = 0; i < num_local_experts && ready; ++i)
|
|
ready &= moe_recv_expert_counter[i] >= 0;
|
|
|
|
if (ready) break;
|
|
|
|
// Timeout check
|
|
if (std::chrono::duration_cast<std::chrono::seconds>(
|
|
std::chrono::high_resolution_clock::now() - start_time)
|
|
.count() > NUM_CPU_TIMEOUT_SECS) {
|
|
LOG(INFO) << "Global rank: " << rank
|
|
<< ", num_recv_tokens: " << num_recv_tokens
|
|
<< ", num_rdma_recv_tokens: " << num_rdma_recv_tokens;
|
|
for (int i = 0; i < num_local_experts; ++i)
|
|
LOG(INFO) << "moe_recv_expert_counter[" << i
|
|
<< "]: " << moe_recv_expert_counter[i];
|
|
throw std::runtime_error("DeepEP error: timeout (dispatch CPU)");
|
|
}
|
|
}
|
|
num_recv_tokens_per_expert_list = std::vector<int>(
|
|
moe_recv_expert_counter, moe_recv_expert_counter + num_local_experts);
|
|
}
|
|
|
|
// Allocate new tensors
|
|
auto recv_x = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_recv_tokens, hidden}, x.dtype(), x.place()));
|
|
auto recv_topk_idx = std::optional<deep_ep::detail::Tensor>(),
|
|
recv_topk_weights = std::optional<deep_ep::detail::Tensor>(),
|
|
recv_x_scales = std::optional<deep_ep::detail::Tensor>();
|
|
auto recv_src_meta = std::optional<deep_ep::detail::Tensor>();
|
|
auto recv_rdma_channel_prefix_matrix =
|
|
std::optional<deep_ep::detail::Tensor>();
|
|
auto recv_gbl_channel_prefix_matrix =
|
|
std::optional<deep_ep::detail::Tensor>();
|
|
auto send_rdma_head = std::optional<deep_ep::detail::Tensor>();
|
|
auto send_nvl_head = std::optional<deep_ep::detail::Tensor>();
|
|
if (!cached_mode) {
|
|
recv_src_meta =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_recv_tokens, internode::get_source_meta_bytes()},
|
|
phi::DataType::INT8,
|
|
phi::GPUPlace(device_id)));
|
|
recv_rdma_channel_prefix_matrix = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_rdma_ranks, num_channels},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
recv_gbl_channel_prefix_matrix = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_ranks, num_channels},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
send_rdma_head = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_tokens, num_rdma_ranks},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
send_nvl_head = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_rdma_recv_tokens, NUM_MAX_NVL_PEERS},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
}
|
|
|
|
// Assign pointers
|
|
int64_t* recv_topk_idx_ptr = nullptr;
|
|
float* recv_topk_weights_ptr = nullptr;
|
|
float* recv_x_scales_ptr = nullptr;
|
|
if (topk_idx.has_value()) {
|
|
recv_topk_idx =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_recv_tokens, num_topk}, topk_idx->dtype(), topk_idx->place()));
|
|
recv_topk_weights = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_recv_tokens, num_topk},
|
|
topk_weights->dtype(),
|
|
topk_weights->place()));
|
|
recv_topk_idx_ptr = recv_topk_idx->data_ptr<int64_t>();
|
|
recv_topk_weights_ptr = recv_topk_weights->data_ptr<float>();
|
|
}
|
|
if (x_scales.has_value()) {
|
|
recv_x_scales =
|
|
x_scales->dim() == 1
|
|
? ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_recv_tokens}, x_scales->dtype(), x_scales->place()))
|
|
: ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_recv_tokens, num_scales},
|
|
x_scales->dtype(),
|
|
x_scales->place()));
|
|
recv_x_scales_ptr = recv_x_scales->data_ptr<float>();
|
|
}
|
|
|
|
// Launch data dispatch
|
|
// NOTES: the buffer size checks are moved into the `.cu` file
|
|
internode::dispatch(
|
|
recv_x.data_ptr(),
|
|
recv_x_scales_ptr,
|
|
recv_topk_idx_ptr,
|
|
recv_topk_weights_ptr,
|
|
cached_mode ? nullptr : recv_src_meta->data_ptr(),
|
|
x.data_ptr(),
|
|
x_scales_ptr,
|
|
topk_idx_ptr,
|
|
topk_weights_ptr,
|
|
cached_mode ? nullptr : send_rdma_head->data_ptr<int>(),
|
|
cached_mode ? nullptr : send_nvl_head->data_ptr<int>(),
|
|
cached_mode ? nullptr : recv_rdma_channel_prefix_matrix->data_ptr<int>(),
|
|
cached_mode ? nullptr : recv_gbl_channel_prefix_matrix->data_ptr<int>(),
|
|
rdma_channel_prefix_matrix.data_ptr<int>(),
|
|
recv_rdma_rank_prefix_sum.data_ptr<int>(),
|
|
gbl_channel_prefix_matrix.data_ptr<int>(),
|
|
recv_gbl_rank_prefix_sum.data_ptr<int>(),
|
|
num_tokens,
|
|
hidden_int4,
|
|
num_scales,
|
|
num_topk,
|
|
num_experts,
|
|
is_token_in_rank.data_ptr<bool>(),
|
|
rdma_buffer_ptr,
|
|
config.num_max_rdma_chunked_send_tokens,
|
|
config.num_max_rdma_chunked_recv_tokens,
|
|
buffer_ptrs_gpu,
|
|
config.num_max_nvl_chunked_send_tokens,
|
|
config.num_max_nvl_chunked_recv_tokens,
|
|
rank,
|
|
num_ranks,
|
|
cached_mode,
|
|
comm_stream,
|
|
num_channels,
|
|
low_latency_mode);
|
|
|
|
// Wait streams
|
|
std::optional<EventHandle> event;
|
|
if (async) {
|
|
event = EventHandle(comm_stream);
|
|
for (auto& t : {x,
|
|
is_token_in_rank,
|
|
recv_x,
|
|
rdma_channel_prefix_matrix,
|
|
recv_rdma_rank_prefix_sum,
|
|
gbl_channel_prefix_matrix,
|
|
recv_gbl_rank_prefix_sum}) {
|
|
t.record_stream(comm_stream);
|
|
if (allocate_on_comm_stream) t.record_stream(compute_stream);
|
|
}
|
|
for (auto& to : {x_scales,
|
|
topk_idx,
|
|
topk_weights,
|
|
num_tokens_per_rank,
|
|
num_tokens_per_rdma_rank,
|
|
num_tokens_per_expert,
|
|
cached_rdma_channel_prefix_matrix,
|
|
cached_recv_rdma_rank_prefix_sum,
|
|
cached_gbl_channel_prefix_matrix,
|
|
cached_recv_gbl_rank_prefix_sum,
|
|
recv_topk_idx,
|
|
recv_topk_weights,
|
|
recv_x_scales,
|
|
recv_rdma_channel_prefix_matrix,
|
|
recv_gbl_channel_prefix_matrix,
|
|
send_rdma_head,
|
|
send_nvl_head,
|
|
recv_src_meta}) {
|
|
to.has_value() ? to->record_stream(comm_stream) : void();
|
|
if (allocate_on_comm_stream)
|
|
to.has_value() ? to->record_stream(compute_stream) : void();
|
|
}
|
|
} else {
|
|
stream_wait(compute_stream, comm_stream);
|
|
}
|
|
|
|
// Switch back compute stream
|
|
if (allocate_on_comm_stream) {
|
|
deep_ep::detail::SetAllocatorStreamForGPUContext(compute_stream, calc_ctx);
|
|
}
|
|
|
|
// Return values
|
|
return {recv_x,
|
|
recv_x_scales,
|
|
recv_topk_idx,
|
|
recv_topk_weights,
|
|
num_recv_tokens_per_expert_list,
|
|
rdma_channel_prefix_matrix,
|
|
gbl_channel_prefix_matrix,
|
|
recv_rdma_channel_prefix_matrix,
|
|
recv_rdma_rank_prefix_sum,
|
|
recv_gbl_channel_prefix_matrix,
|
|
recv_gbl_rank_prefix_sum,
|
|
recv_src_meta,
|
|
send_rdma_head,
|
|
send_nvl_head,
|
|
event};
|
|
}
|
|
|
|
std::tuple<deep_ep::detail::Tensor,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
std::optional<EventHandle>>
|
|
Buffer::internode_combine(
|
|
const deep_ep::detail::Tensor& x,
|
|
const std::optional<deep_ep::detail::Tensor>& topk_weights,
|
|
const deep_ep::detail::Tensor& src_meta,
|
|
const deep_ep::detail::Tensor& is_combined_token_in_rank,
|
|
const deep_ep::detail::Tensor& rdma_channel_prefix_matrix,
|
|
const deep_ep::detail::Tensor& rdma_rank_prefix_sum,
|
|
const deep_ep::detail::Tensor& gbl_channel_prefix_matrix,
|
|
const deep_ep::detail::Tensor& combined_rdma_head,
|
|
const deep_ep::detail::Tensor& combined_nvl_head,
|
|
const Config& config,
|
|
std::optional<EventHandle>& previous_event, // NOLINT
|
|
bool async,
|
|
bool allocate_on_comm_stream) {
|
|
const int num_channels = config.num_sms / 2;
|
|
EP_HOST_ASSERT(config.num_sms % 2 == 0);
|
|
|
|
// Shape and contiguous checks
|
|
EP_HOST_ASSERT(x.dim() == 2 && x.is_contiguous());
|
|
EP_HOST_ASSERT(src_meta.dim() == 2 && src_meta.is_contiguous() &&
|
|
src_meta.scalar_type() == deep_ep::detail::kByte);
|
|
EP_HOST_ASSERT(is_combined_token_in_rank.dim() == 2 &&
|
|
is_combined_token_in_rank.is_contiguous() &&
|
|
is_combined_token_in_rank.scalar_type() ==
|
|
deep_ep::detail::kBool);
|
|
EP_HOST_ASSERT(rdma_channel_prefix_matrix.dim() == 2 &&
|
|
rdma_channel_prefix_matrix.is_contiguous() &&
|
|
rdma_channel_prefix_matrix.scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(rdma_rank_prefix_sum.dim() == 1 &&
|
|
rdma_rank_prefix_sum.is_contiguous() &&
|
|
rdma_rank_prefix_sum.scalar_type() == deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(gbl_channel_prefix_matrix.dim() == 2 &&
|
|
gbl_channel_prefix_matrix.is_contiguous() &&
|
|
gbl_channel_prefix_matrix.scalar_type() ==
|
|
deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(combined_rdma_head.dim() == 2 &&
|
|
combined_rdma_head.is_contiguous() &&
|
|
combined_rdma_head.scalar_type() == deep_ep::detail::kInt32);
|
|
EP_HOST_ASSERT(combined_nvl_head.dim() == 2 &&
|
|
combined_nvl_head.is_contiguous() &&
|
|
combined_nvl_head.scalar_type() == deep_ep::detail::kInt32);
|
|
|
|
auto num_tokens = static_cast<int>(x.size(0)),
|
|
hidden = static_cast<int>(x.size(1)),
|
|
hidden_int4 =
|
|
static_cast<int>(x.size(1) * x.element_size() / sizeof(int4));
|
|
auto num_combined_tokens =
|
|
static_cast<int>(is_combined_token_in_rank.size(0));
|
|
EP_HOST_ASSERT((hidden * x.element_size()) % sizeof(int4) == 0);
|
|
EP_HOST_ASSERT(src_meta.size(1) == internode::get_source_meta_bytes());
|
|
EP_HOST_ASSERT(is_combined_token_in_rank.size(1) == num_ranks);
|
|
EP_HOST_ASSERT(rdma_channel_prefix_matrix.size(0) == num_rdma_ranks &&
|
|
rdma_channel_prefix_matrix.size(1) == num_channels);
|
|
EP_HOST_ASSERT(rdma_rank_prefix_sum.size(0) == num_rdma_ranks);
|
|
EP_HOST_ASSERT(gbl_channel_prefix_matrix.size(0) == num_ranks &&
|
|
gbl_channel_prefix_matrix.size(1) == num_channels);
|
|
EP_HOST_ASSERT(combined_rdma_head.dim() == 2 &&
|
|
combined_rdma_head.size(0) == num_combined_tokens &&
|
|
combined_rdma_head.size(1) == num_rdma_ranks);
|
|
EP_HOST_ASSERT(combined_nvl_head.dim() == 2 &&
|
|
combined_nvl_head.size(1) == NUM_MAX_NVL_PEERS);
|
|
|
|
// Allocate all tensors on comm stream if set
|
|
// NOTES: do not allocate tensors upfront!
|
|
auto compute_stream = calc_ctx->stream();
|
|
if (allocate_on_comm_stream) {
|
|
EP_HOST_ASSERT(previous_event.has_value() && async);
|
|
deep_ep::detail::SetAllocatorStreamForGPUContext(comm_stream, calc_ctx);
|
|
}
|
|
|
|
// Wait previous tasks to be finished
|
|
if (previous_event.has_value()) {
|
|
stream_wait(comm_stream, previous_event.value());
|
|
} else {
|
|
stream_wait(comm_stream, compute_stream);
|
|
}
|
|
|
|
// Top-k checks
|
|
int num_topk = 0;
|
|
auto combined_topk_weights = std::optional<deep_ep::detail::Tensor>();
|
|
float* topk_weights_ptr = nullptr;
|
|
float* combined_topk_weights_ptr = nullptr;
|
|
if (topk_weights.has_value()) {
|
|
EP_HOST_ASSERT(topk_weights->dim() == 2 && topk_weights->is_contiguous());
|
|
EP_HOST_ASSERT(topk_weights->size(0) == num_tokens);
|
|
EP_HOST_ASSERT(topk_weights->scalar_type() == deep_ep::detail::kFloat32);
|
|
num_topk = static_cast<int>(topk_weights->size(1));
|
|
topk_weights_ptr = topk_weights->data_ptr<float>();
|
|
combined_topk_weights = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_combined_tokens, num_topk},
|
|
topk_weights->dtype(),
|
|
topk_weights->place()));
|
|
combined_topk_weights_ptr = combined_topk_weights->data_ptr<float>();
|
|
}
|
|
|
|
// Extra check for avoid-dead-lock design
|
|
EP_HOST_ASSERT(config.num_max_nvl_chunked_recv_tokens % num_rdma_ranks == 0);
|
|
EP_HOST_ASSERT(config.num_max_nvl_chunked_send_tokens <=
|
|
config.num_max_nvl_chunked_recv_tokens / num_rdma_ranks);
|
|
|
|
// Launch barrier and reset queue head and tail
|
|
internode::cached_notify(
|
|
hidden_int4,
|
|
0,
|
|
0,
|
|
num_topk,
|
|
num_ranks,
|
|
num_channels,
|
|
num_combined_tokens,
|
|
combined_rdma_head.data_ptr<int>(),
|
|
rdma_channel_prefix_matrix.data_ptr<int>(),
|
|
rdma_rank_prefix_sum.data_ptr<int>(),
|
|
combined_nvl_head.data_ptr<int>(),
|
|
rdma_buffer_ptr,
|
|
config.num_max_rdma_chunked_recv_tokens,
|
|
buffer_ptrs_gpu,
|
|
config.num_max_nvl_chunked_recv_tokens,
|
|
task_fifo_ptrs_gpu,
|
|
head,
|
|
rank,
|
|
comm_stream,
|
|
config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks),
|
|
num_nvl_bytes,
|
|
false,
|
|
low_latency_mode);
|
|
move_fifo_slots(2);
|
|
|
|
// Launch data combine
|
|
auto combined_x =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_combined_tokens, hidden}, x.dtype(), x.place()));
|
|
internode::combine(deep_ep::detail::ScalarTypeToCudaDataType(x.scalar_type()),
|
|
combined_x.data_ptr(),
|
|
combined_topk_weights_ptr,
|
|
is_combined_token_in_rank.data_ptr<bool>(),
|
|
x.data_ptr(),
|
|
topk_weights_ptr,
|
|
combined_rdma_head.data_ptr<int>(),
|
|
combined_nvl_head.data_ptr<int>(),
|
|
src_meta.data_ptr(),
|
|
rdma_channel_prefix_matrix.data_ptr<int>(),
|
|
rdma_rank_prefix_sum.data_ptr<int>(),
|
|
gbl_channel_prefix_matrix.data_ptr<int>(),
|
|
num_tokens,
|
|
num_combined_tokens,
|
|
hidden,
|
|
num_topk,
|
|
rdma_buffer_ptr,
|
|
config.num_max_rdma_chunked_send_tokens,
|
|
config.num_max_rdma_chunked_recv_tokens,
|
|
buffer_ptrs_gpu,
|
|
config.num_max_nvl_chunked_send_tokens,
|
|
config.num_max_nvl_chunked_recv_tokens,
|
|
rank,
|
|
num_ranks,
|
|
comm_stream,
|
|
num_channels,
|
|
low_latency_mode);
|
|
|
|
// Wait streams
|
|
std::optional<EventHandle> event;
|
|
if (async) {
|
|
event = EventHandle(comm_stream);
|
|
for (auto& t : {x,
|
|
src_meta,
|
|
is_combined_token_in_rank,
|
|
rdma_channel_prefix_matrix,
|
|
rdma_rank_prefix_sum,
|
|
gbl_channel_prefix_matrix,
|
|
combined_x,
|
|
combined_rdma_head,
|
|
combined_nvl_head}) {
|
|
t.record_stream(comm_stream);
|
|
if (allocate_on_comm_stream) t.record_stream(compute_stream);
|
|
}
|
|
for (auto& to : {topk_weights, combined_topk_weights}) {
|
|
to.has_value() ? to->record_stream(comm_stream) : void();
|
|
if (allocate_on_comm_stream)
|
|
to.has_value() ? to->record_stream(compute_stream) : void();
|
|
}
|
|
} else {
|
|
stream_wait(compute_stream, comm_stream);
|
|
}
|
|
|
|
// Switch back compute stream
|
|
if (allocate_on_comm_stream) {
|
|
deep_ep::detail::SetAllocatorStreamForGPUContext(compute_stream, calc_ctx);
|
|
}
|
|
|
|
// Return values
|
|
return {combined_x, combined_topk_weights, event};
|
|
}
|
|
#endif // PADDLE_WITH_NVSHMEM
|
|
|
|
void Buffer::clean_low_latency_buffer(int num_max_dispatch_tokens_per_rank,
|
|
int hidden,
|
|
int num_experts) {
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
EP_HOST_ASSERT(low_latency_mode);
|
|
|
|
auto layout = LowLatencyLayout(rdma_buffer_ptr,
|
|
num_max_dispatch_tokens_per_rank,
|
|
hidden,
|
|
num_ranks,
|
|
num_experts);
|
|
auto clean_meta_0 = layout.buffers[0].clean_meta();
|
|
auto clean_meta_1 = layout.buffers[1].clean_meta();
|
|
|
|
auto check_boundary = [=](void* ptr, size_t num_bytes) {
|
|
auto offset = reinterpret_cast<int64_t>(ptr) -
|
|
reinterpret_cast<int64_t>(rdma_buffer_ptr);
|
|
EP_HOST_ASSERT(0 <= offset &&
|
|
offset + static_cast<int64_t>(num_bytes) <= num_rdma_bytes);
|
|
};
|
|
check_boundary(clean_meta_0.first, clean_meta_0.second * sizeof(int));
|
|
check_boundary(clean_meta_1.first, clean_meta_1.second * sizeof(int));
|
|
|
|
internode_ll::clean_low_latency_buffer(clean_meta_0.first,
|
|
clean_meta_0.second,
|
|
clean_meta_1.first,
|
|
clean_meta_1.second,
|
|
calc_ctx->stream());
|
|
#else
|
|
LOG(ERROR) << "NVSHMEM is not enabled. You can enable it by setting cmake "
|
|
"option WITH_NVSHMEM=ON.";
|
|
#endif
|
|
}
|
|
|
|
void Buffer::clean_low_latency_two_stage_buffer(
|
|
int num_max_dispatch_tokens_per_rank,
|
|
int hidden,
|
|
int num_experts,
|
|
int num_topk,
|
|
int num_ranks,
|
|
bool use_fp8) {
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
EP_HOST_ASSERT(low_latency_mode);
|
|
|
|
const int num_local_experts = num_experts / num_ranks;
|
|
const int num_rdma_experts = num_local_experts * NUM_MAX_NVL_PEERS;
|
|
const int num_scales = hidden / 128;
|
|
const int num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS;
|
|
const size_t dispatch_num_bytes_per_msg =
|
|
sizeof(int4) + (use_fp8 ? (hidden + num_scales * sizeof(float))
|
|
: (hidden * sizeof(nv_bfloat16)));
|
|
auto dispatch_nvl_num_bytes = num_local_experts * num_ranks *
|
|
num_max_dispatch_tokens_per_rank *
|
|
dispatch_num_bytes_per_msg;
|
|
const size_t combine_num_bytes_per_msg = hidden * sizeof(nv_bfloat16);
|
|
auto combine_nvl_num_bytes = num_rdma_experts * num_rdma_ranks *
|
|
num_max_dispatch_tokens_per_rank *
|
|
combine_num_bytes_per_msg;
|
|
const size_t signal_bytes = (num_local_experts * num_ranks * sizeof(int) +
|
|
NUM_BUFFER_ALIGNMENT_BYTES - 1) /
|
|
NUM_BUFFER_ALIGNMENT_BYTES *
|
|
NUM_BUFFER_ALIGNMENT_BYTES;
|
|
auto max_nvl_num_bytes =
|
|
(std::max(dispatch_nvl_num_bytes, combine_nvl_num_bytes) +
|
|
NUM_BUFFER_ALIGNMENT_BYTES - 1) /
|
|
NUM_BUFFER_ALIGNMENT_BYTES * NUM_BUFFER_ALIGNMENT_BYTES;
|
|
|
|
auto layout = LowLatencyTwoStageLayout(rdma_buffer_ptr,
|
|
num_max_dispatch_tokens_per_rank,
|
|
hidden,
|
|
num_ranks,
|
|
num_experts,
|
|
num_topk);
|
|
auto clean_meta_0 = layout.buffers[0].clean_meta();
|
|
auto clean_meta_1 = layout.buffers[1].clean_meta();
|
|
|
|
auto check_boundary = [=](void* ptr, size_t num_bytes) {
|
|
auto offset = reinterpret_cast<int64_t>(ptr) -
|
|
reinterpret_cast<int64_t>(rdma_buffer_ptr);
|
|
EP_HOST_ASSERT(0 <= offset &&
|
|
offset + static_cast<int64_t>(num_bytes) <= num_rdma_bytes);
|
|
};
|
|
check_boundary(clean_meta_0.first, clean_meta_0.second * sizeof(int));
|
|
check_boundary(clean_meta_1.first, clean_meta_1.second * sizeof(int));
|
|
|
|
internode_ll_two_stage::clean_low_latency_buffer_two_stage(
|
|
buffer_ptrs_gpu,
|
|
max_nvl_num_bytes,
|
|
signal_bytes,
|
|
nvl_rank,
|
|
num_experts,
|
|
clean_meta_0.first,
|
|
clean_meta_0.second,
|
|
clean_meta_1.first,
|
|
clean_meta_1.second,
|
|
calc_ctx->stream());
|
|
#else
|
|
LOG(ERROR) << "NVSHMEM is not enabled. You can enable it by setting cmake "
|
|
"option WITH_NVSHMEM=ON.";
|
|
#endif
|
|
}
|
|
|
|
void Buffer::barrier_all() {
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
internode_ll::barrier_all(calc_ctx->stream());
|
|
#else
|
|
LOG(ERROR) << "NVSHMEM is not enabled. You can enable it by setting cmake "
|
|
"option WITH_NVSHMEM=ON.";
|
|
#endif
|
|
}
|
|
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
std::tuple<deep_ep::detail::Tensor,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
std::optional<EventHandle>,
|
|
std::optional<std::function<void()>>>
|
|
Buffer::low_latency_dispatch(
|
|
const deep_ep::detail::Tensor& x,
|
|
const deep_ep::detail::Tensor& topk_idx,
|
|
const std::optional<deep_ep::detail::Tensor>& expertwise_scale,
|
|
int num_max_dispatch_tokens_per_rank,
|
|
int num_experts,
|
|
bool use_fp8,
|
|
bool async,
|
|
bool return_recv_hook,
|
|
int num_per_channel) {
|
|
EP_HOST_ASSERT(low_latency_mode);
|
|
|
|
// Tensor checks
|
|
// By default using `ptp128c` FP8 cast
|
|
EP_HOST_ASSERT(x.dim() == 2 && x.is_contiguous() &&
|
|
x.scalar_type() == deep_ep::detail::kBFloat16);
|
|
EP_HOST_ASSERT(x.size(1) % sizeof(int4) == 0 && x.size(1) % 128 == 0);
|
|
EP_HOST_ASSERT(topk_idx.dim() == 2 && topk_idx.is_contiguous());
|
|
EP_HOST_ASSERT(x.size(0) == topk_idx.size(0) &&
|
|
x.size(0) <= num_max_dispatch_tokens_per_rank);
|
|
EP_HOST_ASSERT(topk_idx.scalar_type() == deep_ep::detail::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 = num_per_channel == -1 ? 1 : hidden / 128,
|
|
num_topk = static_cast<int>(topk_idx.size(1));
|
|
int num_local_experts = num_experts / num_ranks;
|
|
|
|
// Buffer control
|
|
LowLatencyLayout layout(rdma_buffer_ptr,
|
|
num_max_dispatch_tokens_per_rank,
|
|
hidden,
|
|
num_ranks,
|
|
num_experts);
|
|
EP_HOST_ASSERT(static_cast<int64_t>(layout.total_bytes) <= num_rdma_bytes);
|
|
auto buffer = layout.buffers[low_latency_buffer_idx];
|
|
auto next_buffer = layout.buffers[low_latency_buffer_idx ^= 1];
|
|
|
|
// Wait previous tasks to be finished
|
|
// NOTES: the hook mode will always use the default stream
|
|
auto compute_stream = calc_ctx->stream();
|
|
auto launch_stream = return_recv_hook ? compute_stream : comm_stream;
|
|
EP_HOST_ASSERT(!(async && return_recv_hook));
|
|
if (!return_recv_hook) stream_wait(launch_stream, compute_stream);
|
|
|
|
auto return_x_dtype = phi::DataType::BFLOAT16;
|
|
if (use_fp8) {
|
|
if (expertwise_scale.has_value()) {
|
|
EP_HOST_ASSERT(expertwise_scale.value().size(0) == num_experts);
|
|
}
|
|
return_x_dtype = phi::DataType::FLOAT8_E4M3FN;
|
|
} else if (expertwise_scale.has_value()) {
|
|
EP_HOST_ASSERT(expertwise_scale.value().size(0) == num_experts);
|
|
return_x_dtype = phi::DataType::INT8;
|
|
}
|
|
|
|
// Allocate packed tensors
|
|
auto packed_recv_x = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_local_experts,
|
|
num_ranks * num_max_dispatch_tokens_per_rank,
|
|
hidden},
|
|
return_x_dtype,
|
|
x.place()));
|
|
auto packed_recv_src_info =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
auto packed_recv_layout_range = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_local_experts, num_ranks},
|
|
phi::DataType::INT64,
|
|
phi::GPUPlace(device_id)));
|
|
auto packed_recv_count =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_local_experts}, phi::DataType::INT32, phi::GPUPlace(device_id)));
|
|
|
|
// Allocate column-majored scales
|
|
auto packed_recv_x_scales = std::optional<deep_ep::detail::Tensor>();
|
|
|
|
float* packed_recv_x_scales_ptr = nullptr;
|
|
|
|
if (use_fp8 && !expertwise_scale.has_value()) {
|
|
EP_HOST_ASSERT((num_ranks * num_max_dispatch_tokens_per_rank) % 4 == 0 &&
|
|
"TMA requires the number of tokens to be multiple of 4");
|
|
packed_recv_x_scales =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_local_experts,
|
|
num_scales,
|
|
num_ranks * num_max_dispatch_tokens_per_rank},
|
|
phi::DataType::FLOAT32,
|
|
phi::GPUPlace(device_id)));
|
|
packed_recv_x_scales =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::transpose(
|
|
ConvertDetailTensorToPaddleTensor(packed_recv_x_scales.value()),
|
|
std::vector<int>{0, 2, 1}));
|
|
packed_recv_x_scales_ptr = packed_recv_x_scales.value().data_ptr<float>();
|
|
}
|
|
|
|
float* expertwise_scale_ptr = nullptr;
|
|
if (expertwise_scale.has_value()) {
|
|
expertwise_scale_ptr = expertwise_scale.value().data_ptr<float>();
|
|
}
|
|
|
|
// Kernel launch
|
|
auto next_clean_meta = next_buffer.clean_meta();
|
|
auto launcher = [=](int phases) {
|
|
internode_ll::dispatch(packed_recv_x.data_ptr(),
|
|
packed_recv_x_scales_ptr,
|
|
packed_recv_src_info.data_ptr<int>(),
|
|
packed_recv_layout_range.data_ptr<int64_t>(),
|
|
packed_recv_count.data_ptr<int>(),
|
|
buffer.dispatch_rdma_recv_data_buffer,
|
|
buffer.dispatch_rdma_recv_count_buffer,
|
|
buffer.dispatch_rdma_send_buffer,
|
|
x.data_ptr(),
|
|
topk_idx.data_ptr<int64_t>(),
|
|
expertwise_scale_ptr,
|
|
next_clean_meta.first,
|
|
next_clean_meta.second,
|
|
num_tokens,
|
|
hidden,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_topk,
|
|
num_experts,
|
|
rank,
|
|
num_ranks,
|
|
use_fp8,
|
|
workspace,
|
|
launch_stream,
|
|
phases,
|
|
num_per_channel);
|
|
};
|
|
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 (!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<deep_ep::detail::Tensor,
|
|
std::optional<EventHandle>,
|
|
std::optional<std::function<void()>>>
|
|
Buffer::low_latency_combine(const deep_ep::detail::Tensor& x,
|
|
const deep_ep::detail::Tensor& topk_idx,
|
|
const deep_ep::detail::Tensor& topk_weights,
|
|
const deep_ep::detail::Tensor& src_info,
|
|
const deep_ep::detail::Tensor& layout_range,
|
|
int num_max_dispatch_tokens_per_rank,
|
|
int num_experts,
|
|
bool zero_copy,
|
|
bool async,
|
|
bool return_recv_hook,
|
|
const std::optional<deep_ep::detail::Tensor>& out) {
|
|
EP_HOST_ASSERT(low_latency_mode);
|
|
|
|
// Tensor checks
|
|
EP_HOST_ASSERT(x.dim() == 3 && x.is_contiguous() &&
|
|
x.scalar_type() == deep_ep::detail::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 && x.size(2) % 128 == 0);
|
|
EP_HOST_ASSERT(topk_idx.dim() == 2 && topk_idx.is_contiguous());
|
|
EP_HOST_ASSERT(topk_idx.size(0) == topk_weights.size(0) &&
|
|
topk_idx.size(1) == topk_weights.size(1));
|
|
EP_HOST_ASSERT(topk_idx.scalar_type() == deep_ep::detail::kInt64);
|
|
EP_HOST_ASSERT(topk_weights.dim() == 2 && topk_weights.is_contiguous());
|
|
EP_HOST_ASSERT(topk_weights.size(0) <= num_max_dispatch_tokens_per_rank);
|
|
EP_HOST_ASSERT(topk_weights.scalar_type() == deep_ep::detail::kFloat32);
|
|
EP_HOST_ASSERT(src_info.dim() == 2 && src_info.is_contiguous());
|
|
EP_HOST_ASSERT(src_info.scalar_type() == deep_ep::detail::kInt32 &&
|
|
x.size(0) == src_info.size(0));
|
|
EP_HOST_ASSERT(layout_range.dim() == 2 && layout_range.is_contiguous());
|
|
EP_HOST_ASSERT(layout_range.scalar_type() == deep_ep::detail::kInt64);
|
|
EP_HOST_ASSERT(layout_range.size(0) == num_experts / num_ranks &&
|
|
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));
|
|
(void)num_local_experts;
|
|
auto num_combined_tokens = static_cast<int>(topk_weights.size(0));
|
|
|
|
// Buffer control
|
|
LowLatencyLayout layout(rdma_buffer_ptr,
|
|
num_max_dispatch_tokens_per_rank,
|
|
hidden,
|
|
num_ranks,
|
|
num_experts);
|
|
EP_HOST_ASSERT(static_cast<int64_t>(layout.total_bytes) <= num_rdma_bytes);
|
|
auto buffer = layout.buffers[low_latency_buffer_idx];
|
|
auto next_buffer = layout.buffers[low_latency_buffer_idx ^= 1];
|
|
|
|
// Wait previous tasks to be finished
|
|
// NOTES: the hook mode will always use the default stream
|
|
auto compute_stream = calc_ctx->stream();
|
|
auto launch_stream = return_recv_hook ? compute_stream : comm_stream;
|
|
EP_HOST_ASSERT(!(async && return_recv_hook));
|
|
if (!return_recv_hook) stream_wait(launch_stream, compute_stream);
|
|
|
|
// Allocate output tensor
|
|
deep_ep::detail::Tensor combined_x;
|
|
if (out.has_value()) {
|
|
EP_HOST_ASSERT(out->dim() == 2 && out->is_contiguous());
|
|
EP_HOST_ASSERT(out->size(0) == num_combined_tokens &&
|
|
out->size(1) == hidden);
|
|
EP_HOST_ASSERT(out->scalar_type() == x.scalar_type());
|
|
combined_x = out.value();
|
|
} else {
|
|
combined_x = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_combined_tokens, hidden}, x.dtype(), x.place()));
|
|
}
|
|
|
|
// Kernel launch
|
|
auto next_clean_meta = next_buffer.clean_meta();
|
|
auto launcher = [=](int phases) {
|
|
internode_ll::combine(combined_x.data_ptr(),
|
|
buffer.combine_rdma_recv_data_buffer,
|
|
buffer.combine_rdma_recv_flag_buffer,
|
|
buffer.combine_rdma_send_buffer,
|
|
x.data_ptr(),
|
|
topk_idx.data_ptr<int64_t>(),
|
|
topk_weights.data_ptr<float>(),
|
|
src_info.data_ptr<int>(),
|
|
layout_range.data_ptr<int64_t>(),
|
|
next_clean_meta.first,
|
|
next_clean_meta.second,
|
|
num_combined_tokens,
|
|
hidden,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_topk,
|
|
num_experts,
|
|
rank,
|
|
num_ranks,
|
|
workspace,
|
|
launch_stream,
|
|
phases,
|
|
zero_copy);
|
|
};
|
|
launcher(return_recv_hook
|
|
? LOW_LATENCY_SEND_PHASE
|
|
: (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE));
|
|
|
|
// Wait streams
|
|
std::optional<EventHandle> event;
|
|
if (async) {
|
|
// NOTES: we must ensure the all tensors will not be deallocated before the
|
|
// stream-wait happens, so in Python API, we must wrap all tensors into the
|
|
// event handle.
|
|
event = EventHandle(launch_stream);
|
|
} else if (!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 std::tuple<deep_ep::detail::Tensor,
|
|
std::optional<EventHandle>,
|
|
std::optional<std::function<void()>>>{
|
|
deep_ep::detail::Tensor{combined_x}, event, recv_hook};
|
|
}
|
|
|
|
std::tuple<deep_ep::detail::Tensor,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
std::optional<EventHandle>,
|
|
std::optional<std::function<void()>>>
|
|
Buffer::low_latency_dispatch_two_stage(
|
|
const deep_ep::detail::Tensor& x,
|
|
const deep_ep::detail::Tensor& topk_idx,
|
|
const deep_ep::detail::Tensor& topk_weights,
|
|
int num_max_dispatch_tokens_per_rank,
|
|
int num_experts,
|
|
bool use_fp8,
|
|
bool async,
|
|
bool return_recv_hook,
|
|
int num_per_channel) {
|
|
EP_HOST_ASSERT(low_latency_mode);
|
|
|
|
// Tensor checks
|
|
EP_HOST_ASSERT(x.dim() == 2 && x.is_contiguous() &&
|
|
x.scalar_type() == deep_ep::detail::kBFloat16);
|
|
EP_HOST_ASSERT(x.size(1) % sizeof(int4) == 0 && x.size(1) % 128 == 0);
|
|
EP_HOST_ASSERT(topk_idx.dim() == 2 && topk_idx.is_contiguous());
|
|
EP_HOST_ASSERT(x.size(0) == topk_idx.size(0) &&
|
|
x.size(0) <= num_max_dispatch_tokens_per_rank);
|
|
EP_HOST_ASSERT(topk_idx.scalar_type() == deep_ep::detail::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 = num_per_channel == -1 ? 1 : hidden / 128,
|
|
num_topk = static_cast<int>(topk_idx.size(1));
|
|
int num_local_experts = num_experts / num_ranks;
|
|
|
|
// Buffer control
|
|
LowLatencyTwoStageLayout layout(rdma_buffer_ptr,
|
|
num_max_dispatch_tokens_per_rank,
|
|
hidden,
|
|
num_ranks,
|
|
num_experts,
|
|
num_topk);
|
|
EP_HOST_ASSERT(layout.total_bytes <= num_rdma_bytes);
|
|
// fixed buffer, 0 for dispatch, 1 for combine
|
|
auto buffer = layout.buffers[low_latency_buffer_idx];
|
|
auto next_buffer = layout.buffers[low_latency_buffer_idx ^= 1];
|
|
|
|
// Wait previous tasks to be finished
|
|
auto compute_stream = calc_ctx->stream();
|
|
auto launch_stream = async ? comm_stream : compute_stream;
|
|
EP_HOST_ASSERT(!(async && return_recv_hook));
|
|
|
|
auto return_x_dtype = phi::DataType::BFLOAT16;
|
|
if (use_fp8) {
|
|
return_x_dtype = phi::DataType::FLOAT8_E4M3FN;
|
|
}
|
|
|
|
// Allocate packed tensors
|
|
auto packed_recv_x = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_local_experts,
|
|
num_ranks * num_max_dispatch_tokens_per_rank,
|
|
hidden},
|
|
return_x_dtype,
|
|
x.place()));
|
|
auto rdma_send_flags = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_tokens, num_ranks / NUM_MAX_NVL_PEERS},
|
|
phi::DataType::BOOL,
|
|
phi::GPUPlace(device_id)));
|
|
auto packed_recv_src_info =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
auto packed_recv_layout_range = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_local_experts, num_ranks},
|
|
phi::DataType::INT64,
|
|
phi::GPUPlace(device_id)));
|
|
auto packed_recv_count =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_local_experts}, phi::DataType::INT32, phi::GPUPlace(device_id)));
|
|
auto packed_rdma_recv_count = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_ranks / NUM_MAX_NVL_PEERS},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
const size_t num_bytes_per_msg =
|
|
sizeof(int4) +
|
|
(num_ranks / NUM_MAX_NVL_PEERS * (num_topk * 3 + 1) * sizeof(int) +
|
|
sizeof(int4) - 1) /
|
|
sizeof(int4) * sizeof(int4) +
|
|
(use_fp8 ? (hidden + (num_scales + 3) / 4 * 4 * sizeof(float))
|
|
: (hidden * sizeof(nv_bfloat16)));
|
|
auto packed_rdma_recv_x = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_ranks / NUM_MAX_NVL_PEERS,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_bytes_per_msg},
|
|
phi::DataType::UINT8,
|
|
phi::GPUPlace(device_id)));
|
|
|
|
// Allocate column-majored scales
|
|
auto packed_recv_x_scales = std::optional<deep_ep::detail::Tensor>();
|
|
float* packed_recv_x_scales_ptr = nullptr;
|
|
if (use_fp8) {
|
|
EP_HOST_ASSERT((num_ranks * num_max_dispatch_tokens_per_rank) % 4 == 0 &&
|
|
"TMA requires the number of tokens to be multiple of 4");
|
|
packed_recv_x_scales =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_local_experts,
|
|
num_scales,
|
|
num_ranks * num_max_dispatch_tokens_per_rank},
|
|
phi::DataType::FLOAT32,
|
|
phi::GPUPlace(device_id)));
|
|
packed_recv_x_scales =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::transpose(
|
|
ConvertDetailTensorToPaddleTensor(packed_recv_x_scales.value()),
|
|
std::vector<int>{0, 2, 1}));
|
|
packed_recv_x_scales_ptr = packed_recv_x_scales.value().data_ptr<float>();
|
|
}
|
|
|
|
// Kernel launch
|
|
auto next_clean_meta = next_buffer.clean_meta();
|
|
auto launcher = [=](int phases) {
|
|
internode_ll_two_stage::dispatch(
|
|
packed_recv_x.data_ptr(),
|
|
packed_recv_x_scales_ptr,
|
|
packed_rdma_recv_x.data_ptr(),
|
|
packed_recv_src_info.data_ptr<int>(),
|
|
packed_recv_layout_range.data_ptr<int64_t>(),
|
|
packed_recv_count.data_ptr<int>(),
|
|
packed_rdma_recv_count.data_ptr<int>(),
|
|
rdma_send_flags.data_ptr<bool>(),
|
|
buffer.dispatch_rdma_recv_data_buffer,
|
|
buffer.dispatch_rdma_recv_count_buffer,
|
|
buffer.dispatch_rdma_send_buffer,
|
|
buffer_ptrs_gpu,
|
|
x.data_ptr(),
|
|
topk_idx.data_ptr<int64_t>(),
|
|
topk_weights.data_ptr<float>(),
|
|
next_clean_meta.first,
|
|
next_clean_meta.second,
|
|
num_tokens,
|
|
hidden,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_topk,
|
|
num_experts,
|
|
rank,
|
|
num_ranks,
|
|
use_fp8,
|
|
workspace,
|
|
launch_stream,
|
|
phases,
|
|
low_latency_buffer_idx,
|
|
num_per_channel);
|
|
};
|
|
launcher(return_recv_hook
|
|
? LOW_LATENCY_SEND_PHASE
|
|
: (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE));
|
|
// Async event
|
|
std::optional<EventHandle> event;
|
|
if (async) {
|
|
event = EventHandle(launch_stream);
|
|
}
|
|
std::optional<std::function<void()>> recv_hook = std::nullopt;
|
|
if (return_recv_hook) recv_hook = [=]() { launcher(LOW_LATENCY_RECV_PHASE); };
|
|
return {packed_recv_x,
|
|
packed_recv_x_scales,
|
|
packed_rdma_recv_x,
|
|
packed_recv_count,
|
|
packed_rdma_recv_count,
|
|
packed_recv_src_info,
|
|
packed_recv_layout_range,
|
|
rdma_send_flags,
|
|
event,
|
|
recv_hook};
|
|
}
|
|
|
|
std::tuple<deep_ep::detail::Tensor,
|
|
std::optional<EventHandle>,
|
|
std::optional<std::function<void()>>>
|
|
Buffer::low_latency_combine_two_stage(
|
|
const deep_ep::detail::Tensor& x,
|
|
const deep_ep::detail::Tensor& rdma_recv_x,
|
|
const deep_ep::detail::Tensor& topk_idx,
|
|
const deep_ep::detail::Tensor& topk_weights,
|
|
const deep_ep::detail::Tensor& src_info,
|
|
const deep_ep::detail::Tensor& layout_range,
|
|
const deep_ep::detail::Tensor& rdma_send_flags,
|
|
const deep_ep::detail::Tensor& dispatch_rdma_recv_count,
|
|
int num_max_dispatch_tokens_per_rank,
|
|
int num_experts,
|
|
bool dispatch_use_fp8,
|
|
bool async,
|
|
bool return_recv_hook,
|
|
int num_per_channel,
|
|
const std::optional<deep_ep::detail::Tensor>& out) {
|
|
EP_HOST_ASSERT(low_latency_mode);
|
|
|
|
// Tensor checks
|
|
EP_HOST_ASSERT(x.dim() == 3 && x.is_contiguous() &&
|
|
x.scalar_type() == deep_ep::detail::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 && x.size(2) % 128 == 0);
|
|
EP_HOST_ASSERT(topk_idx.dim() == 2 && topk_idx.is_contiguous());
|
|
EP_HOST_ASSERT(topk_idx.size(0) == topk_weights.size(0) &&
|
|
topk_idx.size(1) == topk_weights.size(1));
|
|
EP_HOST_ASSERT(topk_idx.scalar_type() == deep_ep::detail::kInt64);
|
|
EP_HOST_ASSERT(topk_weights.dim() == 2 && topk_weights.is_contiguous());
|
|
EP_HOST_ASSERT(topk_weights.size(0) <= num_max_dispatch_tokens_per_rank);
|
|
EP_HOST_ASSERT(topk_weights.scalar_type() == deep_ep::detail::kFloat32);
|
|
EP_HOST_ASSERT(src_info.dim() == 2 && src_info.is_contiguous());
|
|
EP_HOST_ASSERT(src_info.scalar_type() == deep_ep::detail::kInt32 &&
|
|
x.size(0) == src_info.size(0));
|
|
EP_HOST_ASSERT(layout_range.dim() == 2 && layout_range.is_contiguous());
|
|
EP_HOST_ASSERT(layout_range.scalar_type() == deep_ep::detail::kInt64);
|
|
EP_HOST_ASSERT(layout_range.size(0) == num_experts / num_ranks &&
|
|
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
|
|
LowLatencyTwoStageLayout layout(rdma_buffer_ptr,
|
|
num_max_dispatch_tokens_per_rank,
|
|
hidden,
|
|
num_ranks,
|
|
num_experts,
|
|
num_topk);
|
|
EP_HOST_ASSERT(layout.total_bytes <= num_rdma_bytes);
|
|
|
|
auto buffer = layout.buffers[low_latency_buffer_idx];
|
|
auto next_buffer = layout.buffers[low_latency_buffer_idx ^= 1];
|
|
|
|
auto compute_stream = calc_ctx->stream();
|
|
auto launch_stream = async ? comm_stream : compute_stream;
|
|
EP_HOST_ASSERT(!(async && return_recv_hook));
|
|
|
|
// Allocate output tensor
|
|
deep_ep::detail::Tensor combined_x;
|
|
if (out.has_value()) {
|
|
EP_HOST_ASSERT(out->dim() == 2 && out->is_contiguous());
|
|
EP_HOST_ASSERT(out->size(0) == num_combined_tokens &&
|
|
out->size(1) == hidden);
|
|
EP_HOST_ASSERT(out->scalar_type() == x.scalar_type());
|
|
combined_x = out.value();
|
|
} else {
|
|
combined_x = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_combined_tokens, hidden}, x.dtype(), x.place()));
|
|
}
|
|
|
|
// Kernel launch
|
|
auto next_clean_meta = next_buffer.clean_meta();
|
|
auto launcher = [=](int phases) {
|
|
internode_ll_two_stage::combine(combined_x.data_ptr(),
|
|
buffer.combine_rdma_recv_data_buffer,
|
|
buffer.combine_rdma_recv_flag_buffer,
|
|
buffer.combine_rdma_send_buffer,
|
|
rdma_recv_x.data_ptr(),
|
|
dispatch_rdma_recv_count.data_ptr<int>(),
|
|
buffer_ptrs_gpu,
|
|
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>(),
|
|
rdma_send_flags.data_ptr<bool>(),
|
|
next_clean_meta.first,
|
|
next_clean_meta.second,
|
|
num_combined_tokens,
|
|
hidden,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_topk,
|
|
num_experts,
|
|
rank,
|
|
num_ranks,
|
|
workspace,
|
|
launch_stream,
|
|
phases,
|
|
dispatch_use_fp8,
|
|
low_latency_buffer_idx,
|
|
num_per_channel);
|
|
};
|
|
launcher(return_recv_hook
|
|
? LOW_LATENCY_SEND_PHASE
|
|
: (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE));
|
|
// Async event
|
|
std::optional<EventHandle> event;
|
|
if (async) {
|
|
event = EventHandle(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};
|
|
}
|
|
|
|
std::tuple<deep_ep::detail::Tensor,
|
|
std::optional<deep_ep::detail::Tensor>,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
deep_ep::detail::Tensor,
|
|
std::optional<EventHandle>,
|
|
std::optional<std::function<EventHandle()>>>
|
|
Buffer::m2n_low_latency_dispatch_two_stage(
|
|
const deep_ep::detail::Tensor& x,
|
|
const deep_ep::detail::Tensor& topk_idx,
|
|
const deep_ep::detail::Tensor& topk_weights,
|
|
int num_max_dispatch_tokens_per_rank,
|
|
int num_experts,
|
|
int a_start_rank,
|
|
int a_num_ranks,
|
|
int e_start_rank,
|
|
int e_num_ranks,
|
|
bool use_fp8,
|
|
bool async,
|
|
bool return_recv_hook) {
|
|
EP_HOST_ASSERT(low_latency_mode);
|
|
|
|
// Tensor checks
|
|
EP_HOST_ASSERT(x.dim() == 2 && x.is_contiguous() &&
|
|
x.scalar_type() == deep_ep::detail::kBFloat16);
|
|
EP_HOST_ASSERT(x.size(1) % sizeof(int4) == 0 && x.size(1) % 128 == 0);
|
|
EP_HOST_ASSERT(topk_idx.dim() == 2 && topk_idx.is_contiguous());
|
|
EP_HOST_ASSERT(x.size(0) == topk_idx.size(0) &&
|
|
x.size(0) <= num_max_dispatch_tokens_per_rank);
|
|
EP_HOST_ASSERT(topk_idx.scalar_type() == deep_ep::detail::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
|
|
LowLatencyTwoStageLayout layout(rdma_buffer_ptr,
|
|
num_max_dispatch_tokens_per_rank,
|
|
hidden,
|
|
num_ranks,
|
|
num_experts,
|
|
num_topk);
|
|
EP_HOST_ASSERT(layout.total_bytes <= num_rdma_bytes);
|
|
// fixed buffer, 0 for dispatch, 1 for combine
|
|
auto buffer = layout.buffers[0];
|
|
auto next_buffer = layout.buffers[1];
|
|
auto dispatch_workspace = reinterpret_cast<void*>(
|
|
reinterpret_cast<uint8_t*>(workspace) +
|
|
m2n_ll_dispatch_workspace_idx * NUM_WORKSPACE_BYTES);
|
|
m2n_ll_dispatch_workspace_idx =
|
|
(m2n_ll_dispatch_workspace_idx + 1) % M2N_NUM_WORKSPACE;
|
|
auto dispatch_rdma_recv_complete =
|
|
buffer.dispatch_rdma_recv_complete_buffer +
|
|
m2n_ll_dispatch_recv_complete_idx * num_ranks;
|
|
m2n_ll_dispatch_recv_complete_idx =
|
|
(m2n_ll_dispatch_recv_complete_idx + 1) % M2N_NUM_MAX_MICRO_BATCHES;
|
|
|
|
// Wait previous tasks to be finished
|
|
// NOTES: the hook mode will always use the default stream
|
|
// auto compute_stream = calc_ctx->stream();
|
|
// auto launch_stream = return_recv_hook ? compute_stream : comm_stream;
|
|
// EP_HOST_ASSERT(!(async && return_recv_hook));
|
|
// if (!return_recv_hook) stream_wait(launch_stream, compute_stream);
|
|
|
|
auto compute_stream = calc_ctx->stream();
|
|
auto launch_stream = comm_stream;
|
|
if (rank >= a_start_rank && rank < a_start_rank + a_num_ranks) {
|
|
stream_wait(launch_stream, compute_stream);
|
|
}
|
|
|
|
if (rank >= a_start_rank && rank < a_start_rank + a_num_ranks) {
|
|
stream_wait(compute_stream, launch_stream);
|
|
}
|
|
|
|
auto return_x_dtype = phi::DataType::BFLOAT16;
|
|
if (use_fp8) {
|
|
return_x_dtype = phi::DataType::FLOAT8_E4M3FN;
|
|
}
|
|
|
|
// Allocate packed tensors
|
|
auto packed_recv_x = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_local_experts,
|
|
num_ranks * num_max_dispatch_tokens_per_rank,
|
|
hidden},
|
|
return_x_dtype,
|
|
x.place()));
|
|
auto rdma_send_flags = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_tokens, num_ranks / NUM_MAX_NVL_PEERS},
|
|
phi::DataType::BOOL,
|
|
phi::GPUPlace(device_id)));
|
|
auto packed_recv_src_info =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
auto packed_recv_layout_range = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_local_experts, num_ranks},
|
|
phi::DataType::INT64,
|
|
phi::GPUPlace(device_id)));
|
|
auto packed_recv_count =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_local_experts}, phi::DataType::INT32, phi::GPUPlace(device_id)));
|
|
auto packed_rdma_recv_count = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_ranks / NUM_MAX_NVL_PEERS},
|
|
phi::DataType::INT32,
|
|
phi::GPUPlace(device_id)));
|
|
|
|
const size_t num_bytes_per_msg =
|
|
sizeof(int4) +
|
|
(num_ranks / NUM_MAX_NVL_PEERS * (num_topk * 3 + 1) * sizeof(int) +
|
|
sizeof(int4) - 1) /
|
|
sizeof(int4) * sizeof(int4) +
|
|
(use_fp8 ? (hidden + num_scales * sizeof(float))
|
|
: (hidden * sizeof(nv_bfloat16)));
|
|
auto packed_rdma_recv_x = ConvertPaddleTensorToDetailTensor(
|
|
paddle::experimental::empty({num_ranks / NUM_MAX_NVL_PEERS,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_bytes_per_msg},
|
|
phi::DataType::UINT8,
|
|
phi::GPUPlace(device_id)));
|
|
|
|
// Allocate column-majored scales
|
|
auto packed_recv_x_scales = std::optional<deep_ep::detail::Tensor>();
|
|
float* packed_recv_x_scales_ptr = nullptr;
|
|
if (use_fp8) {
|
|
EP_HOST_ASSERT((num_ranks * num_max_dispatch_tokens_per_rank) % 4 == 0 &&
|
|
"TMA requires the number of tokens to be multiple of 4");
|
|
packed_recv_x_scales =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_local_experts,
|
|
num_scales,
|
|
num_ranks * num_max_dispatch_tokens_per_rank},
|
|
phi::DataType::FLOAT32,
|
|
phi::GPUPlace(device_id)));
|
|
packed_recv_x_scales =
|
|
ConvertPaddleTensorToDetailTensor(paddle::experimental::transpose(
|
|
ConvertDetailTensorToPaddleTensor(packed_recv_x_scales.value()),
|
|
std::vector<int>{0, 2, 1}));
|
|
packed_recv_x_scales_ptr = packed_recv_x_scales.value().data_ptr<float>();
|
|
}
|
|
|
|
// Kernel launch
|
|
auto next_clean_meta = next_buffer.clean_meta();
|
|
auto launcher = [=](int phases) {
|
|
m2n_ll_two_stage::dispatch(packed_recv_x.data_ptr(),
|
|
packed_recv_x_scales_ptr,
|
|
packed_rdma_recv_x.data_ptr(),
|
|
packed_recv_src_info.data_ptr<int>(),
|
|
packed_recv_layout_range.data_ptr<int64_t>(),
|
|
packed_recv_count.data_ptr<int>(),
|
|
packed_rdma_recv_count.data_ptr<int>(),
|
|
rdma_send_flags.data_ptr<bool>(),
|
|
buffer.dispatch_rdma_recv_data_buffer,
|
|
buffer.dispatch_rdma_recv_count_buffer,
|
|
dispatch_rdma_recv_complete,
|
|
buffer.dispatch_rdma_send_buffer,
|
|
buffer_ptrs_gpu,
|
|
x.data_ptr(),
|
|
topk_idx.data_ptr<int64_t>(),
|
|
topk_weights.data_ptr<float>(),
|
|
next_clean_meta.first,
|
|
next_clean_meta.second,
|
|
num_tokens,
|
|
hidden,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_topk,
|
|
num_experts,
|
|
rank,
|
|
num_ranks,
|
|
a_start_rank,
|
|
a_num_ranks,
|
|
e_start_rank,
|
|
e_num_ranks,
|
|
use_fp8,
|
|
dispatch_workspace,
|
|
launch_stream,
|
|
phases);
|
|
};
|
|
|
|
// TODO(Zhenyu Li): supports async/return_recv_hook
|
|
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 (!return_recv_hook) {
|
|
// stream_wait(compute_stream, launch_stream);
|
|
// }
|
|
|
|
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);
|
|
}
|
|
// // stream_wait(launch_stream, compute_stream);
|
|
// if (rank >= a_start_rank && rank < a_start_rank + a_num_ranks) {
|
|
// stream_wait(compute_stream, launch_stream);
|
|
// }
|
|
|
|
// Receiver callback
|
|
std::optional<std::function<EventHandle()>> recv_hook = std::nullopt;
|
|
if (return_recv_hook)
|
|
recv_hook = [=]() {
|
|
// stream_wait(launch_stream, compute_stream);
|
|
launcher(LOW_LATENCY_RECV_PHASE);
|
|
// stream_wait(compute_stream, launch_stream);
|
|
|
|
// if (rank >= e_start_rank && rank < e_start_rank + e_num_ranks) {
|
|
// stream_wait(compute_stream, launch_stream);
|
|
// }
|
|
return EventHandle(launch_stream);
|
|
};
|
|
|
|
return {packed_recv_x,
|
|
packed_recv_x_scales,
|
|
packed_rdma_recv_x,
|
|
packed_recv_count,
|
|
packed_rdma_recv_count,
|
|
packed_recv_src_info,
|
|
packed_recv_layout_range,
|
|
rdma_send_flags,
|
|
event,
|
|
recv_hook};
|
|
}
|
|
|
|
std::tuple<deep_ep::detail::Tensor,
|
|
std::optional<EventHandle>,
|
|
std::optional<std::function<EventHandle()>>>
|
|
Buffer::m2n_low_latency_combine_two_stage(
|
|
const deep_ep::detail::Tensor& x,
|
|
const deep_ep::detail::Tensor& rdma_recv_x,
|
|
const deep_ep::detail::Tensor& topk_idx,
|
|
const deep_ep::detail::Tensor& topk_weights,
|
|
const deep_ep::detail::Tensor& src_info,
|
|
const deep_ep::detail::Tensor& layout_range,
|
|
const deep_ep::detail::Tensor& rdma_send_flags,
|
|
const deep_ep::detail::Tensor& dispatch_rdma_recv_count,
|
|
int num_max_dispatch_tokens_per_rank,
|
|
int num_experts,
|
|
int a_start_rank,
|
|
int a_num_ranks,
|
|
int e_start_rank,
|
|
int e_num_ranks,
|
|
bool dispatch_use_fp8,
|
|
bool async,
|
|
bool return_recv_hook,
|
|
const std::optional<deep_ep::detail::Tensor>& out) {
|
|
EP_HOST_ASSERT(low_latency_mode);
|
|
|
|
// Tensor checks
|
|
EP_HOST_ASSERT(x.dim() == 3 && x.is_contiguous() &&
|
|
x.scalar_type() == deep_ep::detail::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 && x.size(2) % 128 == 0);
|
|
EP_HOST_ASSERT(topk_idx.dim() == 2 && topk_idx.is_contiguous());
|
|
EP_HOST_ASSERT(topk_idx.size(0) == topk_weights.size(0) &&
|
|
topk_idx.size(1) == topk_weights.size(1));
|
|
EP_HOST_ASSERT(topk_idx.scalar_type() == deep_ep::detail::kInt64);
|
|
EP_HOST_ASSERT(topk_weights.dim() == 2 && topk_weights.is_contiguous());
|
|
EP_HOST_ASSERT(topk_weights.size(0) <= num_max_dispatch_tokens_per_rank);
|
|
EP_HOST_ASSERT(topk_weights.scalar_type() == deep_ep::detail::kFloat32);
|
|
EP_HOST_ASSERT(src_info.dim() == 2 && src_info.is_contiguous());
|
|
EP_HOST_ASSERT(src_info.scalar_type() == deep_ep::detail::kInt32 &&
|
|
x.size(0) == src_info.size(0));
|
|
EP_HOST_ASSERT(layout_range.dim() == 2 && layout_range.is_contiguous());
|
|
EP_HOST_ASSERT(layout_range.scalar_type() == deep_ep::detail::kInt64);
|
|
EP_HOST_ASSERT(layout_range.size(0) == num_experts / num_ranks &&
|
|
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
|
|
LowLatencyTwoStageLayout layout(rdma_buffer_ptr,
|
|
num_max_dispatch_tokens_per_rank,
|
|
hidden,
|
|
num_ranks,
|
|
num_experts,
|
|
num_topk);
|
|
EP_HOST_ASSERT(layout.total_bytes <= num_rdma_bytes);
|
|
// fixed buffer, 0 for dispatch, 1 for combine
|
|
auto dispatch_buffer = layout.buffers[0];
|
|
auto buffer = layout.buffers[1];
|
|
auto next_buffer = layout.buffers[0];
|
|
auto combine_workspace = reinterpret_cast<void*>(
|
|
reinterpret_cast<uint8_t*>(workspace) +
|
|
(M2N_NUM_WORKSPACE + m2n_ll_combine_workspace_idx) * NUM_WORKSPACE_BYTES);
|
|
m2n_ll_combine_workspace_idx =
|
|
(m2n_ll_combine_workspace_idx + 1) % M2N_NUM_WORKSPACE;
|
|
auto combine_rdma_recv_complete =
|
|
buffer.combine_rdma_recv_complete_buffer +
|
|
m2n_ll_combine_recv_complete_idx * num_ranks;
|
|
m2n_ll_combine_recv_complete_idx =
|
|
(m2n_ll_combine_recv_complete_idx + 1) % M2N_NUM_MAX_MICRO_BATCHES;
|
|
|
|
// Wait previous tasks to be finished
|
|
// NOTES: the hook mode will always use the default stream
|
|
// auto compute_stream = calc_ctx->stream();
|
|
// auto launch_stream = return_recv_hook ? compute_stream : comm_stream;
|
|
// EP_HOST_ASSERT(!(async && return_recv_hook));
|
|
// if (!return_recv_hook) stream_wait(launch_stream, compute_stream);
|
|
|
|
auto compute_stream = calc_ctx->stream();
|
|
auto launch_stream = comm_stream;
|
|
if (rank >= e_start_rank && rank < e_start_rank + e_num_ranks) {
|
|
stream_wait(launch_stream, compute_stream);
|
|
}
|
|
|
|
if (rank >= e_start_rank && rank < e_start_rank + e_num_ranks) {
|
|
stream_wait(compute_stream, launch_stream);
|
|
}
|
|
|
|
// Allocate output tensor
|
|
deep_ep::detail::Tensor combined_x;
|
|
if (out.has_value()) {
|
|
EP_HOST_ASSERT(out->dim() == 2 && out->is_contiguous());
|
|
EP_HOST_ASSERT(out->size(0) == num_combined_tokens &&
|
|
out->size(1) == hidden);
|
|
EP_HOST_ASSERT(out->scalar_type() == x.scalar_type());
|
|
combined_x = out.value();
|
|
} else {
|
|
combined_x = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
|
|
{num_combined_tokens, hidden}, x.dtype(), x.place()));
|
|
}
|
|
|
|
// Kernel launch
|
|
auto next_clean_meta = next_buffer.clean_meta();
|
|
auto launcher = [=](int phases) {
|
|
m2n_ll_two_stage::combine(combined_x.data_ptr(),
|
|
buffer.combine_rdma_recv_data_buffer,
|
|
buffer.combine_rdma_recv_flag_buffer,
|
|
buffer.combine_rdma_send_buffer,
|
|
combine_rdma_recv_complete,
|
|
rdma_recv_x.data_ptr(),
|
|
dispatch_rdma_recv_count.data_ptr<int>(),
|
|
buffer_ptrs_gpu,
|
|
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>(),
|
|
rdma_send_flags.data_ptr<bool>(),
|
|
next_clean_meta.first,
|
|
next_clean_meta.second,
|
|
num_combined_tokens,
|
|
hidden,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_topk,
|
|
num_experts,
|
|
rank,
|
|
num_ranks,
|
|
a_start_rank,
|
|
a_num_ranks,
|
|
e_start_rank,
|
|
e_num_ranks,
|
|
combine_workspace,
|
|
launch_stream,
|
|
phases,
|
|
dispatch_use_fp8);
|
|
};
|
|
// TODO(Zhenyu Li): supports async/return_recv_hook
|
|
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 (!return_recv_hook) {
|
|
// stream_wait(compute_stream, launch_stream);
|
|
// }
|
|
|
|
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);
|
|
}
|
|
// // stream_wait(launch_stream, compute_stream);
|
|
// if (rank >= e_start_rank && rank < e_start_rank + e_num_ranks) {
|
|
// stream_wait(compute_stream, launch_stream);
|
|
// }
|
|
// Receiver callback
|
|
std::optional<std::function<EventHandle()>> recv_hook = std::nullopt;
|
|
if (return_recv_hook)
|
|
recv_hook = [=]() {
|
|
// stream_wait(launch_stream, compute_stream);
|
|
launcher(LOW_LATENCY_RECV_PHASE);
|
|
// stream_wait(compute_stream, launch_stream);
|
|
// stream_wait(launch_stream, compute_stream);
|
|
// if (rank >= a_start_rank && rank < a_start_rank + a_num_ranks) {
|
|
// stream_wait(compute_stream, launch_stream);
|
|
// }
|
|
return EventHandle(launch_stream);
|
|
};
|
|
|
|
// Return values
|
|
return {combined_x, event, recv_hook};
|
|
}
|
|
|
|
#endif // PADDLE_WITH_NVSHMEM
|
|
|
|
std::tuple<paddle::Tensor,
|
|
std::optional<paddle::Tensor>,
|
|
std::optional<paddle::Tensor>,
|
|
std::optional<paddle::Tensor>,
|
|
std::vector<int>,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
std::optional<paddle::Tensor>,
|
|
paddle::Tensor,
|
|
std::optional<paddle::Tensor>,
|
|
paddle::Tensor,
|
|
std::optional<paddle::Tensor>,
|
|
std::optional<paddle::Tensor>,
|
|
std::optional<paddle::Tensor>,
|
|
std::optional<EventHandle>>
|
|
Buffer::internode_dispatch_api(
|
|
const paddle::Tensor& x,
|
|
const std::optional<paddle::Tensor>& x_scales,
|
|
const std::optional<paddle::Tensor>& topk_idx,
|
|
const std::optional<paddle::Tensor>& topk_weights,
|
|
const std::optional<paddle::Tensor>& num_tokens_per_rank,
|
|
const std::optional<paddle::Tensor>& num_tokens_per_rdma_rank,
|
|
const paddle::Tensor& is_token_in_rank,
|
|
const std::optional<paddle::Tensor>& num_tokens_per_expert,
|
|
int cached_num_recv_tokens,
|
|
int cached_num_rdma_recv_tokens,
|
|
const std::optional<paddle::Tensor>& cached_rdma_channel_prefix_matrix,
|
|
const std::optional<paddle::Tensor>& cached_recv_rdma_rank_prefix_sum,
|
|
const std::optional<paddle::Tensor>& cached_gbl_channel_prefix_matrix,
|
|
const std::optional<paddle::Tensor>& cached_recv_gbl_rank_prefix_sum,
|
|
int expert_alignment,
|
|
const Config& config,
|
|
std::optional<EventHandle>& previous_event, // NOLINT
|
|
bool async,
|
|
bool allocate_on_comm_stream) {
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
const auto& x_ = ConvertPaddleTensorToDetailTensor(x);
|
|
std::optional<deep_ep::detail::Tensor> x_scales_ =
|
|
ConvertOptionalPaddleTensorToDetailTensor(x_scales);
|
|
|
|
std::optional<deep_ep::detail::Tensor> topk_idx_ =
|
|
ConvertOptionalPaddleTensorToDetailTensor(topk_idx);
|
|
std::optional<deep_ep::detail::Tensor> topk_weights_ =
|
|
ConvertOptionalPaddleTensorToDetailTensor(topk_weights);
|
|
std::optional<deep_ep::detail::Tensor> num_tokens_per_rank_ =
|
|
ConvertOptionalPaddleTensorToDetailTensor(num_tokens_per_rank);
|
|
std::optional<deep_ep::detail::Tensor> num_tokens_per_rdma_rank_ =
|
|
ConvertOptionalPaddleTensorToDetailTensor(num_tokens_per_rdma_rank);
|
|
|
|
const auto& is_token_in_rank_ =
|
|
ConvertPaddleTensorToDetailTensor(is_token_in_rank);
|
|
std::optional<deep_ep::detail::Tensor> num_tokens_per_expert_ =
|
|
ConvertOptionalPaddleTensorToDetailTensor(num_tokens_per_expert);
|
|
|
|
std::optional<deep_ep::detail::Tensor> cached_rdma_channel_prefix_matrix_ =
|
|
ConvertOptionalPaddleTensorToDetailTensor(
|
|
cached_rdma_channel_prefix_matrix);
|
|
std::optional<deep_ep::detail::Tensor> cached_recv_rdma_rank_prefix_sum_ =
|
|
ConvertOptionalPaddleTensorToDetailTensor(
|
|
cached_recv_rdma_rank_prefix_sum);
|
|
std::optional<deep_ep::detail::Tensor> cached_gbl_channel_prefix_matrix_ =
|
|
ConvertOptionalPaddleTensorToDetailTensor(
|
|
cached_gbl_channel_prefix_matrix);
|
|
std::optional<deep_ep::detail::Tensor> cached_recv_gbl_rank_prefix_sum_ =
|
|
ConvertOptionalPaddleTensorToDetailTensor(
|
|
cached_recv_gbl_rank_prefix_sum);
|
|
|
|
auto res = internode_dispatch(x_,
|
|
x_scales_,
|
|
topk_idx_,
|
|
topk_weights_,
|
|
num_tokens_per_rank_,
|
|
num_tokens_per_rdma_rank_,
|
|
is_token_in_rank_,
|
|
num_tokens_per_expert_,
|
|
cached_num_recv_tokens,
|
|
cached_num_rdma_recv_tokens,
|
|
cached_rdma_channel_prefix_matrix_,
|
|
cached_recv_rdma_rank_prefix_sum_,
|
|
cached_gbl_channel_prefix_matrix_,
|
|
cached_recv_gbl_rank_prefix_sum_,
|
|
expert_alignment,
|
|
config,
|
|
previous_event,
|
|
async,
|
|
allocate_on_comm_stream);
|
|
|
|
auto recv_x_ = ConvertDetailTensorToPaddleTensor(std::get<0>(res));
|
|
std::optional<paddle::Tensor> recv_x_scales_ =
|
|
ConvertOptionalDetailTensorToPaddleTensor(std::get<1>(res));
|
|
|
|
std::optional<paddle::Tensor> recv_topk_idx_ =
|
|
ConvertOptionalDetailTensorToPaddleTensor(std::get<2>(res));
|
|
std::optional<paddle::Tensor> recv_topk_weights_ =
|
|
ConvertOptionalDetailTensorToPaddleTensor(std::get<3>(res));
|
|
|
|
const auto& num_recv_tokens_per_expert_list = std::get<4>(res);
|
|
|
|
auto rdma_channel_prefix_matrix_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<5>(res));
|
|
|
|
auto gbl_channel_prefix_matrix_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<6>(res));
|
|
|
|
std::optional<paddle::Tensor> recv_rdma_channel_prefix_matrix_ =
|
|
ConvertOptionalDetailTensorToPaddleTensor(std::get<7>(res));
|
|
auto recv_rdma_rank_prefix_sum_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<8>(res));
|
|
|
|
std::optional<paddle::Tensor> recv_gbl_channel_prefix_matrix_ =
|
|
ConvertOptionalDetailTensorToPaddleTensor(std::get<9>(res));
|
|
auto recv_gbl_rank_prefix_sum_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<10>(res));
|
|
|
|
std::optional<paddle::Tensor> recv_src_meta_ =
|
|
ConvertOptionalDetailTensorToPaddleTensor(std::get<11>(res));
|
|
|
|
std::optional<paddle::Tensor> send_rdma_head_ =
|
|
ConvertOptionalDetailTensorToPaddleTensor(std::get<12>(res));
|
|
std::optional<paddle::Tensor> send_nvl_head_ =
|
|
ConvertOptionalDetailTensorToPaddleTensor(std::get<13>(res));
|
|
|
|
const auto& event = std::get<14>(res);
|
|
|
|
return {recv_x_,
|
|
recv_x_scales_,
|
|
recv_topk_idx_,
|
|
recv_topk_weights_,
|
|
num_recv_tokens_per_expert_list,
|
|
rdma_channel_prefix_matrix_,
|
|
gbl_channel_prefix_matrix_,
|
|
recv_rdma_channel_prefix_matrix_,
|
|
recv_rdma_rank_prefix_sum_,
|
|
recv_gbl_channel_prefix_matrix_,
|
|
recv_gbl_rank_prefix_sum_,
|
|
recv_src_meta_,
|
|
send_rdma_head_,
|
|
send_nvl_head_,
|
|
event};
|
|
#else
|
|
LOG(ERROR) << "NVSHMEM is not enabled. You can enable it by setting cmake "
|
|
"option WITH_NVSHMEM=ON.";
|
|
return {};
|
|
#endif
|
|
}
|
|
|
|
std::tuple<paddle::Tensor,
|
|
std::optional<paddle::Tensor>,
|
|
std::optional<EventHandle>>
|
|
Buffer::internode_combine_api(
|
|
const paddle::Tensor& x,
|
|
const std::optional<paddle::Tensor>& topk_weights,
|
|
const paddle::Tensor& src_meta,
|
|
const paddle::Tensor& is_combined_token_in_rank,
|
|
const paddle::Tensor& rdma_channel_prefix_matrix,
|
|
const paddle::Tensor& rdma_rank_prefix_sum,
|
|
const paddle::Tensor& gbl_channel_prefix_matrix,
|
|
const paddle::Tensor& combined_rdma_head,
|
|
const paddle::Tensor& combined_nvl_head,
|
|
const Config& config,
|
|
std::optional<EventHandle>& previous_event, // NOLINT
|
|
bool async,
|
|
bool allocate_on_comm_stream) {
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
const auto& x_ = ConvertPaddleTensorToDetailTensor(x);
|
|
|
|
std::optional<deep_ep::detail::Tensor> topk_weights_ =
|
|
ConvertOptionalPaddleTensorToDetailTensor(topk_weights);
|
|
|
|
const auto& src_meta_ = ConvertPaddleTensorToDetailTensor(src_meta);
|
|
const auto& is_combined_token_in_rank_ =
|
|
ConvertPaddleTensorToDetailTensor(is_combined_token_in_rank);
|
|
|
|
const auto& rdma_channel_prefix_matrix_ =
|
|
ConvertPaddleTensorToDetailTensor(rdma_channel_prefix_matrix);
|
|
const auto& rdma_rank_prefix_sum_ =
|
|
ConvertPaddleTensorToDetailTensor(rdma_rank_prefix_sum);
|
|
const auto& gbl_channel_prefix_matrix_ =
|
|
ConvertPaddleTensorToDetailTensor(gbl_channel_prefix_matrix);
|
|
|
|
const auto& combined_rdma_head_ =
|
|
ConvertPaddleTensorToDetailTensor(combined_rdma_head);
|
|
const auto& combined_nvl_head_ =
|
|
ConvertPaddleTensorToDetailTensor(combined_nvl_head);
|
|
|
|
auto res = internode_combine(x_,
|
|
topk_weights_,
|
|
src_meta_,
|
|
is_combined_token_in_rank_,
|
|
rdma_channel_prefix_matrix_,
|
|
rdma_rank_prefix_sum_,
|
|
gbl_channel_prefix_matrix_,
|
|
combined_rdma_head_,
|
|
combined_nvl_head_,
|
|
config,
|
|
previous_event,
|
|
async,
|
|
allocate_on_comm_stream);
|
|
|
|
auto combined_x_ = ConvertDetailTensorToPaddleTensor(std::get<0>(res));
|
|
std::optional<paddle::Tensor> combined_topk_weights_ =
|
|
ConvertOptionalDetailTensorToPaddleTensor(std::get<1>(res));
|
|
|
|
const auto& event = std::get<2>(res);
|
|
|
|
return {combined_x_, combined_topk_weights_, event};
|
|
#else
|
|
LOG(ERROR) << "NVSHMEM is not enabled. You can enable it by setting cmake "
|
|
"option WITH_NVSHMEM=ON.";
|
|
return {};
|
|
#endif
|
|
}
|
|
|
|
std::tuple<paddle::Tensor,
|
|
std::optional<paddle::Tensor>,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
std::optional<EventHandle>,
|
|
std::optional<std::function<void()>>>
|
|
Buffer::low_latency_dispatch_api(
|
|
const paddle::Tensor& x,
|
|
const paddle::Tensor& topk_idx,
|
|
const std::optional<paddle::Tensor>& expertwise_scale,
|
|
int num_max_dispatch_tokens_per_rank,
|
|
int num_experts,
|
|
bool use_fp8,
|
|
bool async,
|
|
bool return_recv_hook,
|
|
int num_per_channel) {
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
const auto& x_ = ConvertPaddleTensorToDetailTensor(x);
|
|
const auto& topk_idx_ = ConvertPaddleTensorToDetailTensor(topk_idx);
|
|
|
|
std::optional<deep_ep::detail::Tensor> expertwise_scale_;
|
|
if (expertwise_scale.has_value()) {
|
|
expertwise_scale_ =
|
|
ConvertPaddleTensorToDetailTensor(expertwise_scale.value());
|
|
}
|
|
|
|
auto res = low_latency_dispatch(x_,
|
|
topk_idx_,
|
|
expertwise_scale_,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_experts,
|
|
use_fp8,
|
|
async,
|
|
return_recv_hook,
|
|
num_per_channel);
|
|
|
|
auto packed_recv_x_ = ConvertDetailTensorToPaddleTensor(std::get<0>(res));
|
|
|
|
std::optional<paddle::Tensor> packed_recv_x_scales_;
|
|
if (std::get<1>(res).has_value()) {
|
|
packed_recv_x_scales_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<1>(res).value());
|
|
}
|
|
|
|
auto packed_recv_count_ = ConvertDetailTensorToPaddleTensor(std::get<2>(res));
|
|
auto packed_recv_src_info_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<3>(res));
|
|
auto packed_recv_layout_range_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<4>(res));
|
|
|
|
const auto& event = std::get<5>(res);
|
|
auto recv_hook = std::get<6>(res);
|
|
|
|
return {packed_recv_x_,
|
|
packed_recv_x_scales_,
|
|
packed_recv_count_,
|
|
packed_recv_src_info_,
|
|
packed_recv_layout_range_,
|
|
event,
|
|
recv_hook};
|
|
#else
|
|
LOG(ERROR) << "NVSHMEM is not enabled. You can enable it by setting cmake "
|
|
"option WITH_NVSHMEM=ON.";
|
|
return {};
|
|
#endif
|
|
}
|
|
|
|
std::tuple<paddle::Tensor,
|
|
std::optional<EventHandle>,
|
|
std::optional<std::function<void()>>>
|
|
Buffer::low_latency_combine_api(const paddle::Tensor& x,
|
|
const paddle::Tensor& topk_idx,
|
|
const paddle::Tensor& topk_weights,
|
|
const paddle::Tensor& src_info,
|
|
const paddle::Tensor& layout_range,
|
|
int num_max_dispatch_tokens_per_rank,
|
|
int num_experts,
|
|
bool zero_copy,
|
|
bool async,
|
|
bool return_recv_hook,
|
|
const std::optional<paddle::Tensor>& out) {
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
const auto& x_ = ConvertPaddleTensorToDetailTensor(x);
|
|
const auto& topk_idx_ = ConvertPaddleTensorToDetailTensor(topk_idx);
|
|
const auto& topk_weights_ = ConvertPaddleTensorToDetailTensor(topk_weights);
|
|
const auto& src_info_ = ConvertPaddleTensorToDetailTensor(src_info);
|
|
const auto& layout_range_ = ConvertPaddleTensorToDetailTensor(layout_range);
|
|
std::optional<deep_ep::detail::Tensor> out_ = std::nullopt;
|
|
if (out.has_value()) {
|
|
out_ = ConvertOptionalPaddleTensorToDetailTensor(out.value());
|
|
}
|
|
|
|
auto res = low_latency_combine(x_,
|
|
topk_idx_,
|
|
topk_weights_,
|
|
src_info_,
|
|
layout_range_,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_experts,
|
|
zero_copy,
|
|
async,
|
|
return_recv_hook,
|
|
out_);
|
|
|
|
auto combined_x_ = ConvertDetailTensorToPaddleTensor(std::get<0>(res));
|
|
const auto& event = std::get<1>(res);
|
|
auto recv_hook = std::get<2>(res);
|
|
|
|
return {combined_x_, event, recv_hook};
|
|
#else
|
|
LOG(ERROR) << "NVSHMEM is not enabled. You can enable it by setting cmake "
|
|
"option WITH_NVSHMEM=ON.";
|
|
return {};
|
|
#endif
|
|
}
|
|
|
|
std::tuple<paddle::Tensor,
|
|
std::optional<paddle::Tensor>,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
std::optional<EventHandle>,
|
|
std::optional<std::function<void()>>>
|
|
Buffer::low_latency_dispatch_two_stage_api(const paddle::Tensor& x,
|
|
const paddle::Tensor& topk_idx,
|
|
const paddle::Tensor& topk_weights,
|
|
int num_max_dispatch_tokens_per_rank,
|
|
int num_experts,
|
|
bool use_fp8,
|
|
bool async,
|
|
bool return_recv_hook,
|
|
int num_per_channel) {
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
const auto& x_ = ConvertPaddleTensorToDetailTensor(x);
|
|
const auto& topk_idx_ = ConvertPaddleTensorToDetailTensor(topk_idx);
|
|
const auto& topk_weights_ = ConvertPaddleTensorToDetailTensor(topk_weights);
|
|
|
|
auto res = low_latency_dispatch_two_stage(x_,
|
|
topk_idx_,
|
|
topk_weights_,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_experts,
|
|
use_fp8,
|
|
async,
|
|
return_recv_hook,
|
|
num_per_channel);
|
|
|
|
auto packed_recv_x_ = ConvertDetailTensorToPaddleTensor(std::get<0>(res));
|
|
|
|
std::optional<paddle::Tensor> packed_recv_x_scales_;
|
|
if (std::get<1>(res).has_value()) {
|
|
packed_recv_x_scales_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<1>(res).value());
|
|
}
|
|
auto packed_recv_rdma_x_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<2>(res));
|
|
|
|
auto packed_recv_count_ = ConvertDetailTensorToPaddleTensor(std::get<3>(res));
|
|
auto packed_rdma_recv_count_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<4>(res));
|
|
auto packed_recv_src_info_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<5>(res));
|
|
auto packed_recv_layout_range_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<6>(res));
|
|
auto rdma_send_flags_ = ConvertDetailTensorToPaddleTensor(std::get<7>(res));
|
|
|
|
const auto& event = std::get<8>(res);
|
|
auto recv_hook = std::get<9>(res);
|
|
|
|
return {packed_recv_x_,
|
|
packed_recv_x_scales_,
|
|
packed_recv_rdma_x_,
|
|
packed_recv_count_,
|
|
packed_rdma_recv_count_,
|
|
packed_recv_src_info_,
|
|
packed_recv_layout_range_,
|
|
rdma_send_flags_,
|
|
event,
|
|
recv_hook};
|
|
#else
|
|
LOG(ERROR) << "NVSHMEM is not enabled. You can enable it by setting cmake "
|
|
"option WITH_NVSHMEM=ON.";
|
|
return {};
|
|
#endif
|
|
}
|
|
|
|
std::tuple<paddle::Tensor,
|
|
std::optional<EventHandle>,
|
|
std::optional<std::function<void()>>>
|
|
Buffer::low_latency_combine_two_stage_api(
|
|
const paddle::Tensor& x,
|
|
const paddle::Tensor& rdma_recv_x,
|
|
const paddle::Tensor& topk_idx,
|
|
const paddle::Tensor& topk_weights,
|
|
const paddle::Tensor& src_info,
|
|
const paddle::Tensor& layout_range,
|
|
const paddle::Tensor& rdma_send_flags,
|
|
const paddle::Tensor& dispatch_rdma_recv_count,
|
|
int num_max_dispatch_tokens_per_rank,
|
|
int num_experts,
|
|
bool dispatch_use_fp8,
|
|
bool async,
|
|
bool return_recv_hook,
|
|
int num_per_channel,
|
|
const std::optional<paddle::Tensor>& out) {
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
const auto& x_ = ConvertPaddleTensorToDetailTensor(x);
|
|
const auto& rdma_recv_x_ = ConvertPaddleTensorToDetailTensor(rdma_recv_x);
|
|
const auto& topk_idx_ = ConvertPaddleTensorToDetailTensor(topk_idx);
|
|
const auto& topk_weights_ = ConvertPaddleTensorToDetailTensor(topk_weights);
|
|
const auto& src_info_ = ConvertPaddleTensorToDetailTensor(src_info);
|
|
const auto& layout_range_ = ConvertPaddleTensorToDetailTensor(layout_range);
|
|
const auto& rdma_send_flags_ =
|
|
ConvertPaddleTensorToDetailTensor(rdma_send_flags);
|
|
const auto& dispatch_rdma_recv_count_ =
|
|
ConvertPaddleTensorToDetailTensor(dispatch_rdma_recv_count);
|
|
|
|
std::optional<deep_ep::detail::Tensor> out_ = std::nullopt;
|
|
if (out.has_value()) {
|
|
out_ = ConvertOptionalPaddleTensorToDetailTensor(out.value());
|
|
}
|
|
|
|
auto res = low_latency_combine_two_stage(x_,
|
|
rdma_recv_x_,
|
|
topk_idx_,
|
|
topk_weights_,
|
|
src_info_,
|
|
layout_range_,
|
|
rdma_send_flags_,
|
|
dispatch_rdma_recv_count_,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_experts,
|
|
dispatch_use_fp8,
|
|
async,
|
|
return_recv_hook,
|
|
num_per_channel,
|
|
out_);
|
|
|
|
auto combined_x_ = ConvertDetailTensorToPaddleTensor(std::get<0>(res));
|
|
const auto& event = std::get<1>(res);
|
|
auto recv_hook = std::get<2>(res);
|
|
|
|
return {combined_x_, event, recv_hook};
|
|
#else
|
|
LOG(ERROR) << "NVSHMEM is not enabled. You can enable it by setting cmake "
|
|
"option WITH_NVSHMEM=ON.";
|
|
return {};
|
|
#endif
|
|
}
|
|
|
|
std::tuple<paddle::Tensor,
|
|
std::optional<paddle::Tensor>,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
std::optional<EventHandle>,
|
|
std::optional<std::function<EventHandle()>>>
|
|
Buffer::m2n_low_latency_dispatch_two_stage_api(
|
|
const paddle::Tensor& x,
|
|
const paddle::Tensor& topk_idx,
|
|
const paddle::Tensor& topk_weights,
|
|
int num_max_dispatch_tokens_per_rank,
|
|
int num_experts,
|
|
int a_start_rank,
|
|
int a_num_ranks,
|
|
int e_start_rank,
|
|
int e_num_ranks,
|
|
bool use_fp8,
|
|
bool async,
|
|
bool return_recv_hook) {
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
const auto& x_ = ConvertPaddleTensorToDetailTensor(x);
|
|
const auto& topk_idx_ = ConvertPaddleTensorToDetailTensor(topk_idx);
|
|
const auto& topk_weights_ = ConvertPaddleTensorToDetailTensor(topk_weights);
|
|
|
|
auto res =
|
|
m2n_low_latency_dispatch_two_stage(x_,
|
|
topk_idx_,
|
|
topk_weights_,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_experts,
|
|
a_start_rank,
|
|
a_num_ranks,
|
|
e_start_rank,
|
|
e_num_ranks,
|
|
use_fp8,
|
|
async,
|
|
return_recv_hook);
|
|
|
|
auto packed_recv_x_ = ConvertDetailTensorToPaddleTensor(std::get<0>(res));
|
|
|
|
std::optional<paddle::Tensor> packed_recv_x_scales_;
|
|
if (std::get<1>(res).has_value()) {
|
|
packed_recv_x_scales_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<1>(res).value());
|
|
}
|
|
auto packed_recv_rdma_x_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<2>(res));
|
|
auto packed_recv_count_ = ConvertDetailTensorToPaddleTensor(std::get<3>(res));
|
|
auto packed_rdma_recv_count_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<4>(res));
|
|
auto packed_recv_src_info_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<5>(res));
|
|
auto packed_recv_layout_range_ =
|
|
ConvertDetailTensorToPaddleTensor(std::get<6>(res));
|
|
auto rdma_send_flags_ = ConvertDetailTensorToPaddleTensor(std::get<7>(res));
|
|
|
|
const auto& event = std::get<8>(res);
|
|
auto recv_hook = std::get<9>(res);
|
|
|
|
return {packed_recv_x_,
|
|
packed_recv_x_scales_,
|
|
packed_recv_rdma_x_,
|
|
packed_recv_count_,
|
|
packed_rdma_recv_count_,
|
|
packed_recv_src_info_,
|
|
packed_recv_layout_range_,
|
|
rdma_send_flags_,
|
|
event,
|
|
recv_hook};
|
|
#else
|
|
LOG(ERROR) << "NVSHMEM is not enabled. You can enable it by setting cmake "
|
|
"option WITH_NVSHMEM=ON.";
|
|
return {};
|
|
#endif
|
|
}
|
|
|
|
std::tuple<paddle::Tensor,
|
|
std::optional<EventHandle>,
|
|
std::optional<std::function<EventHandle()>>>
|
|
Buffer::m2n_low_latency_combine_two_stage_api(
|
|
const paddle::Tensor& x,
|
|
const paddle::Tensor& rdma_recv_x,
|
|
const paddle::Tensor& topk_idx,
|
|
const paddle::Tensor& topk_weights,
|
|
const paddle::Tensor& src_info,
|
|
const paddle::Tensor& layout_range,
|
|
const paddle::Tensor& rdma_send_flags,
|
|
const paddle::Tensor& dispatch_rdma_recv_count,
|
|
int num_max_dispatch_tokens_per_rank,
|
|
int num_experts,
|
|
int a_start_rank,
|
|
int a_num_ranks,
|
|
int e_start_rank,
|
|
int e_num_ranks,
|
|
bool dispatch_use_fp8,
|
|
bool async,
|
|
bool return_recv_hook,
|
|
const std::optional<paddle::Tensor>& out) {
|
|
#ifdef PADDLE_WITH_NVSHMEM
|
|
const auto& x_ = ConvertPaddleTensorToDetailTensor(x);
|
|
const auto& rdma_recv_x_ = ConvertPaddleTensorToDetailTensor(rdma_recv_x);
|
|
const auto& topk_idx_ = ConvertPaddleTensorToDetailTensor(topk_idx);
|
|
const auto& topk_weights_ = ConvertPaddleTensorToDetailTensor(topk_weights);
|
|
const auto& src_info_ = ConvertPaddleTensorToDetailTensor(src_info);
|
|
const auto& layout_range_ = ConvertPaddleTensorToDetailTensor(layout_range);
|
|
const auto& rdma_send_flags_ =
|
|
ConvertPaddleTensorToDetailTensor(rdma_send_flags);
|
|
const auto& dispatch_rdma_recv_count_ =
|
|
ConvertPaddleTensorToDetailTensor(dispatch_rdma_recv_count);
|
|
|
|
std::optional<deep_ep::detail::Tensor> out_ = std::nullopt;
|
|
if (out.has_value()) {
|
|
out_ = ConvertOptionalPaddleTensorToDetailTensor(out.value());
|
|
}
|
|
|
|
auto res = m2n_low_latency_combine_two_stage(x_,
|
|
rdma_recv_x_,
|
|
topk_idx_,
|
|
topk_weights_,
|
|
src_info_,
|
|
layout_range_,
|
|
rdma_send_flags_,
|
|
dispatch_rdma_recv_count_,
|
|
num_max_dispatch_tokens_per_rank,
|
|
num_experts,
|
|
a_start_rank,
|
|
a_num_ranks,
|
|
e_start_rank,
|
|
e_num_ranks,
|
|
dispatch_use_fp8,
|
|
async,
|
|
return_recv_hook,
|
|
out_);
|
|
|
|
auto combined_x_ = ConvertDetailTensorToPaddleTensor(std::get<0>(res));
|
|
const auto& event = std::get<1>(res);
|
|
auto recv_hook = std::get<2>(res);
|
|
|
|
return {combined_x_, event, recv_hook};
|
|
#else
|
|
LOG(ERROR) << "NVSHMEM is not enabled. You can enable it by setting cmake "
|
|
"option WITH_NVSHMEM=ON.";
|
|
return {};
|
|
#endif
|
|
}
|
|
|
|
std::tuple<paddle::Tensor,
|
|
std::optional<paddle::Tensor>,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
std::optional<EventHandle>>
|
|
Buffer::get_dispatch_layout_api(const paddle::Tensor& topk_idx,
|
|
int num_experts,
|
|
std::optional<EventHandle>& previous_event,
|
|
bool async,
|
|
bool allocate_on_comm_stream) {
|
|
const auto& topk_idx_ = ConvertPaddleTensorToDetailTensor(topk_idx);
|
|
auto res = get_dispatch_layout(
|
|
topk_idx_, num_experts, previous_event, async, allocate_on_comm_stream);
|
|
const auto& num_tokens_per_rank = std::get<0>(res);
|
|
const auto& num_tokens_per_rdma_rank = std::get<1>(res);
|
|
const auto& num_tokens_per_expert = std::get<2>(res);
|
|
const auto& is_token_in_rank = std::get<3>(res);
|
|
const auto& event = std::get<4>(res);
|
|
auto num_tokens_per_rank_ =
|
|
ConvertDetailTensorToPaddleTensor(num_tokens_per_rank);
|
|
std::optional<paddle::Tensor> num_tokens_per_rdma_rank_ = std::nullopt;
|
|
if (num_tokens_per_rdma_rank.has_value()) {
|
|
num_tokens_per_rdma_rank_ =
|
|
ConvertDetailTensorToPaddleTensor(num_tokens_per_rdma_rank.value());
|
|
}
|
|
auto num_tokens_per_expert_ =
|
|
ConvertDetailTensorToPaddleTensor(num_tokens_per_expert);
|
|
auto is_token_in_rank_ = ConvertDetailTensorToPaddleTensor(is_token_in_rank);
|
|
return {num_tokens_per_rank_,
|
|
num_tokens_per_rdma_rank_,
|
|
num_tokens_per_expert_,
|
|
is_token_in_rank_,
|
|
event};
|
|
}
|
|
|
|
std::tuple<paddle::Tensor,
|
|
std::optional<paddle::Tensor>,
|
|
std::optional<paddle::Tensor>,
|
|
std::optional<paddle::Tensor>,
|
|
std::vector<int>,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
paddle::Tensor,
|
|
std::optional<EventHandle>>
|
|
Buffer::intranode_dispatch_api(
|
|
const paddle::Tensor& x,
|
|
const std::optional<paddle::Tensor>& x_scales,
|
|
const std::optional<paddle::Tensor>& topk_idx,
|
|
const std::optional<paddle::Tensor>& topk_weights,
|
|
const std::optional<paddle::Tensor>& num_tokens_per_rank,
|
|
const paddle::Tensor& is_token_in_rank,
|
|
const std::optional<paddle::Tensor>& num_tokens_per_expert,
|
|
int cached_num_recv_tokens,
|
|
const std::optional<paddle::Tensor>& cached_rank_prefix_matrix,
|
|
const std::optional<paddle::Tensor>& cached_channel_prefix_matrix,
|
|
int expert_alignment,
|
|
const Config& config,
|
|
std::optional<EventHandle>& previous_event, // NOLINT
|
|
bool async,
|
|
bool allocate_on_comm_stream) {
|
|
const auto& x_ = ConvertPaddleTensorToDetailTensor(x);
|
|
std::optional<deep_ep::detail::Tensor> x_scales_;
|
|
if (x_scales.has_value()) {
|
|
x_scales_ = ConvertPaddleTensorToDetailTensor(x_scales.value());
|
|
}
|
|
std::optional<deep_ep::detail::Tensor> topk_idx_;
|
|
if (topk_idx.has_value()) {
|
|
topk_idx_ = ConvertPaddleTensorToDetailTensor(topk_idx.value());
|
|
}
|
|
std::optional<deep_ep::detail::Tensor> topk_weights_;
|
|
if (topk_weights.has_value()) {
|
|
topk_weights_ = ConvertPaddleTensorToDetailTensor(topk_weights.value());
|
|
}
|
|
std::optional<deep_ep::detail::Tensor> num_tokens_per_rank_;
|
|
if (num_tokens_per_rank.has_value()) {
|
|
num_tokens_per_rank_ =
|
|
ConvertPaddleTensorToDetailTensor(num_tokens_per_rank.value());
|
|
}
|
|
const auto& is_token_in_rank_ =
|
|
ConvertPaddleTensorToDetailTensor(is_token_in_rank);
|
|
std::optional<deep_ep::detail::Tensor> num_tokens_per_expert_;
|
|
if (num_tokens_per_expert.has_value()) {
|
|
num_tokens_per_expert_ =
|
|
ConvertPaddleTensorToDetailTensor(num_tokens_per_expert.value());
|
|
}
|
|
std::optional<deep_ep::detail::Tensor> cached_rank_prefix_matrix_;
|
|
if (cached_rank_prefix_matrix.has_value()) {
|
|
cached_rank_prefix_matrix_ =
|
|
ConvertPaddleTensorToDetailTensor(cached_rank_prefix_matrix.value());
|
|
}
|
|
std::optional<deep_ep::detail::Tensor> cached_channel_prefix_matrix_;
|
|
if (cached_channel_prefix_matrix.has_value()) {
|
|
cached_channel_prefix_matrix_ =
|
|
ConvertPaddleTensorToDetailTensor(cached_channel_prefix_matrix.value());
|
|
}
|
|
|
|
auto res = intranode_dispatch(x_,
|
|
x_scales_,
|
|
topk_idx_,
|
|
topk_weights_,
|
|
num_tokens_per_rank_,
|
|
is_token_in_rank_,
|
|
num_tokens_per_expert_,
|
|
cached_num_recv_tokens,
|
|
cached_rank_prefix_matrix_,
|
|
cached_channel_prefix_matrix_,
|
|
expert_alignment,
|
|
config,
|
|
previous_event,
|
|
async,
|
|
allocate_on_comm_stream);
|
|
|
|
const auto& recv_x = std::get<0>(res);
|
|
const auto& recv_x_scales = std::get<1>(res);
|
|
const auto& recv_topk_idx = std::get<2>(res);
|
|
const auto& recv_topk_weights = std::get<3>(res);
|
|
const auto& num_recv_tokens_per_expert_list = std::get<4>(res);
|
|
const auto& rank_prefix_matrix = std::get<5>(res);
|
|
const auto& channel_prefix_matrix = std::get<6>(res);
|
|
const auto& recv_channel_prefix_matrix = std::get<7>(res);
|
|
const auto& recv_src_idx = std::get<8>(res);
|
|
const auto& send_head = std::get<9>(res);
|
|
const auto& event = std::get<10>(res);
|
|
|
|
auto recv_x_ = ConvertDetailTensorToPaddleTensor(recv_x);
|
|
std::optional<paddle::Tensor> recv_x_scales_;
|
|
if (recv_x_scales.has_value()) {
|
|
recv_x_scales_ = ConvertDetailTensorToPaddleTensor(recv_x_scales.value());
|
|
}
|
|
std::optional<paddle::Tensor> recv_topk_idx_;
|
|
if (recv_topk_idx.has_value()) {
|
|
recv_topk_idx_ = ConvertDetailTensorToPaddleTensor(recv_topk_idx.value());
|
|
}
|
|
std::optional<paddle::Tensor> recv_topk_weights_;
|
|
if (recv_topk_weights.has_value()) {
|
|
recv_topk_weights_ =
|
|
ConvertDetailTensorToPaddleTensor(recv_topk_weights.value());
|
|
}
|
|
auto rank_prefix_matrix_ =
|
|
ConvertDetailTensorToPaddleTensor(rank_prefix_matrix);
|
|
auto channel_prefix_matrix_ =
|
|
ConvertDetailTensorToPaddleTensor(channel_prefix_matrix);
|
|
auto recv_channel_prefix_matrix_ =
|
|
ConvertDetailTensorToPaddleTensor(recv_channel_prefix_matrix);
|
|
auto recv_src_idx_ = ConvertDetailTensorToPaddleTensor(recv_src_idx);
|
|
auto send_head_ = ConvertDetailTensorToPaddleTensor(send_head);
|
|
return {recv_x_,
|
|
recv_x_scales_,
|
|
recv_topk_idx_,
|
|
recv_topk_weights_,
|
|
num_recv_tokens_per_expert_list,
|
|
rank_prefix_matrix_,
|
|
channel_prefix_matrix_,
|
|
recv_channel_prefix_matrix_,
|
|
recv_src_idx_,
|
|
send_head_,
|
|
event};
|
|
}
|
|
|
|
std::tuple<paddle::Tensor,
|
|
std::optional<paddle::Tensor>,
|
|
std::optional<EventHandle>>
|
|
Buffer::intranode_combine_api(const paddle::Tensor& x,
|
|
const std::optional<paddle::Tensor>& topk_weights,
|
|
const paddle::Tensor& src_idx,
|
|
const paddle::Tensor& rank_prefix_matrix,
|
|
const paddle::Tensor& channel_prefix_matrix,
|
|
const paddle::Tensor& send_head,
|
|
const Config& config,
|
|
std::optional<EventHandle>& previous_event,
|
|
bool async,
|
|
bool allocate_on_comm_stream) {
|
|
const auto& x_ = ConvertPaddleTensorToDetailTensor(x);
|
|
std::optional<deep_ep::detail::Tensor> topk_weights_;
|
|
if (topk_weights.has_value()) {
|
|
topk_weights_ = ConvertPaddleTensorToDetailTensor(topk_weights.value());
|
|
}
|
|
const auto& src_idx_ = ConvertPaddleTensorToDetailTensor(src_idx);
|
|
const auto& rank_prefix_matrix_ =
|
|
ConvertPaddleTensorToDetailTensor(rank_prefix_matrix);
|
|
const auto& channel_prefix_matrix_ =
|
|
ConvertPaddleTensorToDetailTensor(channel_prefix_matrix);
|
|
const auto& send_head_ = ConvertPaddleTensorToDetailTensor(send_head);
|
|
|
|
auto res = intranode_combine(x_,
|
|
topk_weights_,
|
|
src_idx_,
|
|
rank_prefix_matrix_,
|
|
channel_prefix_matrix_,
|
|
send_head_,
|
|
config,
|
|
previous_event,
|
|
async,
|
|
allocate_on_comm_stream);
|
|
|
|
const auto& recv_x = std::get<0>(res);
|
|
const auto& recv_topk_weights = std::get<1>(res);
|
|
const auto& event = std::get<2>(res);
|
|
|
|
auto recv_x_ = ConvertDetailTensorToPaddleTensor(recv_x);
|
|
std::optional<paddle::Tensor> recv_topk_weights_;
|
|
if (recv_topk_weights.has_value()) {
|
|
recv_topk_weights_ =
|
|
ConvertDetailTensorToPaddleTensor(recv_topk_weights.value());
|
|
}
|
|
auto event_ = event;
|
|
return {recv_x_, recv_topk_weights_, event_};
|
|
}
|
|
|
|
deep_ep::detail::Tensor ConvertPaddleTensorToDetailTensor(
|
|
const paddle::Tensor& tensor) {
|
|
deep_ep::detail::Tensor res(tensor);
|
|
return res;
|
|
}
|
|
|
|
paddle::Tensor ConvertDetailTensorToPaddleTensor(
|
|
const deep_ep::detail::Tensor& tensor) {
|
|
return tensor.raw_tensor();
|
|
}
|
|
|
|
std::optional<deep_ep::detail::Tensor>
|
|
ConvertOptionalPaddleTensorToDetailTensor(
|
|
const std::optional<paddle::Tensor>& tensor) {
|
|
std::optional<deep_ep::detail::Tensor> res;
|
|
if (tensor.has_value()) {
|
|
res = ConvertPaddleTensorToDetailTensor(tensor.value());
|
|
}
|
|
return res;
|
|
}
|
|
|
|
std::optional<paddle::Tensor> ConvertOptionalDetailTensorToPaddleTensor(
|
|
const std::optional<deep_ep::detail::Tensor>& tensor) {
|
|
std::optional<paddle::Tensor> res;
|
|
if (tensor.has_value()) {
|
|
res = ConvertDetailTensorToPaddleTensor(tensor.value());
|
|
}
|
|
return res;
|
|
}
|
|
|
|
} // namespace deep_ep
|