forked from mooncake-track/Mooncake
[EP] Fallback impl of Mooncake EP when IBGDA is unavailable (#1002)
* Implement a fallback impl of EP's dispatch & combine when IBGDA is not available * Fix format * Fix * Fix * Add fault-tolerance support for the fallback path * Fix * Fix * Exit from `__init__` early, if using fallback impl of EP * Fix * Fix
This commit is contained in:
Xun Sun
GitHub
parent
866b5944d4
commit
77c16fea73
|
|
@ -69,6 +69,8 @@ class MooncakeBackend final : public ::c10d::Backend {
|
|||
return matrix[location].preferred_hca[0];
|
||||
}
|
||||
|
||||
at::Tensor getActiveRanksTensor() { return meta_.activeRanksTensor; }
|
||||
|
||||
private:
|
||||
static TransferEngine engine_;
|
||||
static Transport* transport_;
|
||||
|
|
|
|||
|
|
@ -81,6 +81,7 @@ struct MooncakeEpBuffer {
|
|||
void* qp_devctxs = nullptr;
|
||||
std::string device_name;
|
||||
bool is_roce_ = false;
|
||||
bool ibgda_disabled_ = false;
|
||||
|
||||
// Stream for communication
|
||||
at::cuda::CUDAStream comm_stream;
|
||||
|
|
@ -114,7 +115,9 @@ struct MooncakeEpBuffer {
|
|||
torch::Tensor get_next_combine_buffer(int num_max_dispatch_tokens_per_rank,
|
||||
int hidden, int num_experts);
|
||||
|
||||
void init_ibgda();
|
||||
int init_ibgda();
|
||||
|
||||
bool ibgda_disabled() { return ibgda_disabled_; }
|
||||
|
||||
bool is_roce() { return is_roce_; }
|
||||
|
||||
|
|
|
|||
|
|
@ -19,7 +19,13 @@ MooncakeEpBuffer::MooncakeEpBuffer(int rank, int num_ranks,
|
|||
CUDA_CHECK(cudaMalloc(&rkeys, num_ranks * sizeof(uint32_t)));
|
||||
CUDA_CHECK(
|
||||
cudaMalloc(&qp_devctxs, MAX_QP_COUNT * sizeof(mlx5gda_qp_devctx)));
|
||||
init_ibgda();
|
||||
int ret = init_ibgda();
|
||||
if (ret != 0) {
|
||||
LOG(WARNING) << "Failed to initialize IBGDA. "
|
||||
<< "Using fallback implementation. "
|
||||
<< "Performance will be degraded.";
|
||||
ibgda_disabled_ = true;
|
||||
}
|
||||
|
||||
// Create 32 MiB workspace
|
||||
CUDA_CHECK(cudaMalloc(&workspace, NUM_WORKSPACE_BYTES));
|
||||
|
|
@ -279,7 +285,7 @@ torch::Tensor MooncakeEpBuffer::get_next_combine_buffer(
|
|||
torch::TensorOptions().dtype(dtype).device(torch::kCUDA));
|
||||
}
|
||||
|
||||
void MooncakeEpBuffer::init_ibgda() {
|
||||
int MooncakeEpBuffer::init_ibgda() {
|
||||
int num_devices;
|
||||
ibv_device** dev_list = ibv_get_device_list(&num_devices);
|
||||
int nic_id = -1;
|
||||
|
|
@ -299,7 +305,7 @@ void MooncakeEpBuffer::init_ibgda() {
|
|||
ibv_context* ctx = ibv_open_device(dev_list[nic_id]);
|
||||
if (!ctx) {
|
||||
perror("Failed to open device");
|
||||
exit(1);
|
||||
return -1;
|
||||
}
|
||||
if (ibv_query_gid(ctx, 1, 3, &gid)) {
|
||||
perror("Failed to query gid");
|
||||
|
|
@ -309,7 +315,7 @@ void MooncakeEpBuffer::init_ibgda() {
|
|||
ibv_pd* pd = ibv_alloc_pd(ctx);
|
||||
if (!pd) {
|
||||
perror("Failed to allocate protection domain");
|
||||
exit(1);
|
||||
return -1;
|
||||
}
|
||||
mlx5dv_pd mpd;
|
||||
mlx5dv_obj dv_obj = {};
|
||||
|
|
@ -334,24 +340,24 @@ void MooncakeEpBuffer::init_ibgda() {
|
|||
fprintf(stderr,
|
||||
"If the error is `Bad address`, probably because your GPU "
|
||||
"does not support GPUDirect RDMA.\n");
|
||||
exit(1);
|
||||
return -1;
|
||||
}
|
||||
memheap* ctrl_buf_heap = memheap_create(CTRL_BUF_SIZE);
|
||||
if (!ctrl_buf_heap) {
|
||||
perror("Failed to create memory heap");
|
||||
exit(1);
|
||||
return -1;
|
||||
}
|
||||
for (int i = 0; i < MAX_QP_COUNT; ++i) {
|
||||
mlx5gda_qp* qp = mlx5gda_create_rc_qp(mpd, ctrl_buf, ctrl_buf_umem,
|
||||
ctrl_buf_heap, pd, 16384, 1);
|
||||
if (!qp) {
|
||||
perror("Failed to create QP");
|
||||
exit(1);
|
||||
return -1;
|
||||
}
|
||||
is_roce_ = qp->port_attr.link_layer == IBV_LINK_LAYER_ETHERNET;
|
||||
if (mlx5gda_modify_rc_qp_rst2init(qp, 0)) {
|
||||
perror("Failed to mlx5gda_modify_rc_qp_rst2init");
|
||||
exit(1);
|
||||
return -1;
|
||||
}
|
||||
mlx5gda_qp_devctx qp_devctx = {
|
||||
.qpn = qp->qpn,
|
||||
|
|
@ -365,6 +371,7 @@ void MooncakeEpBuffer::init_ibgda() {
|
|||
sizeof(mlx5gda_qp_devctx), cudaMemcpyHostToDevice);
|
||||
qps.push_back(qp);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
void MooncakeEpBuffer::sync_ib(const std::vector<int64_t>& remote_addrs,
|
||||
|
|
|
|||
|
|
@ -57,12 +57,19 @@ std::string getPreferredHca(c10::intrusive_ptr<c10d::Backend> backend,
|
|||
return mooncakeBackend->getPreferredHca(location);
|
||||
}
|
||||
|
||||
at::Tensor getActiveRanks(c10::intrusive_ptr<c10d::Backend> backend) {
|
||||
auto mooncakeBackend =
|
||||
c10::static_intrusive_pointer_cast<MooncakeBackend>(backend);
|
||||
return mooncakeBackend->getActiveRanksTensor();
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(ep, m) {
|
||||
m.def("createMooncakeBackend", &createMooncakeBackend);
|
||||
m.def("createMooncakeCpuBackend", &createMooncakeCpuBackend);
|
||||
m.def("set_host_ip", &MooncakeBackend::setHostIp);
|
||||
m.def("set_device_filter", &MooncakeBackend::setDeviceFilter);
|
||||
m.def("get_preferred_hca", &getPreferredHca);
|
||||
m.def("get_active_ranks", &getActiveRanks);
|
||||
|
||||
py::class_<MooncakeBackend::MooncakeBackendOptions,
|
||||
c10::intrusive_ptr<MooncakeBackend::MooncakeBackendOptions>>(
|
||||
|
|
@ -79,6 +86,7 @@ PYBIND11_MODULE(ep, m) {
|
|||
|
||||
py::class_<MooncakeEpBuffer>(m, "Buffer")
|
||||
.def(py::init<int, int, int64_t, std::string>())
|
||||
.def("ibgda_disabled", &MooncakeEpBuffer::ibgda_disabled)
|
||||
.def("is_roce", &MooncakeEpBuffer::is_roce)
|
||||
.def("sync_ib", &MooncakeEpBuffer::sync_ib)
|
||||
.def("sync_roce", &MooncakeEpBuffer::sync_roce)
|
||||
|
|
|
|||
|
|
@ -68,9 +68,15 @@ class Buffer:
|
|||
self.num_ep_buffer_bytes = num_ep_buffer_bytes
|
||||
|
||||
# Get the index of the closest NIC
|
||||
backend = self.group._get_backend(torch.device('cuda'))
|
||||
preferred_hca = ep.get_preferred_hca(backend, f'cuda:{torch.cuda.current_device()}')
|
||||
self.backend = self.group._get_backend(torch.device('cuda'))
|
||||
preferred_hca = ep.get_preferred_hca(self.backend, f'cuda:{torch.cuda.current_device()}')
|
||||
self.runtime = ep.Buffer(self.rank, self.group_size, num_ep_buffer_bytes, preferred_hca)
|
||||
# Fallback flag and buffers
|
||||
self._use_fallback = bool(self.runtime.ibgda_disabled())
|
||||
self._fallback_next_combine_buffer: Optional[torch.Tensor] = None
|
||||
|
||||
if self._use_fallback:
|
||||
return
|
||||
|
||||
(raddr, rkey) = self.runtime.get_mr_info()
|
||||
|
||||
|
|
@ -118,18 +124,26 @@ class Buffer:
|
|||
|
||||
@staticmethod
|
||||
def get_ep_buffer_size_hint(num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int) -> int:
|
||||
from mooncake import ep
|
||||
return ep.get_ep_buffer_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts)
|
||||
from mooncake.ep import get_ep_buffer_size_hint
|
||||
return get_ep_buffer_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts)
|
||||
|
||||
# noinspection PyTypeChecker
|
||||
def dispatch(self, x: torch.Tensor, topk_idx: torch.Tensor, active_ranks: torch.Tensor,
|
||||
num_max_dispatch_tokens_per_rank: int, num_experts: int, timeout_us: int,
|
||||
use_fp8: bool = True, async_finish: bool = False, return_recv_hook: bool = False) -> \
|
||||
Tuple[Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor], torch.Tensor, Tuple, EventOverlap, Callable]:
|
||||
packed_recv_x, packed_recv_x_scales, packed_recv_count, packed_recv_src_info, packed_recv_layout_range, event, hook = \
|
||||
self.runtime.dispatch(x, topk_idx, active_ranks,
|
||||
num_max_dispatch_tokens_per_rank, num_experts, timeout_us,
|
||||
use_fp8, async_finish, return_recv_hook)
|
||||
if self._use_fallback:
|
||||
from mooncake.ep import get_active_ranks
|
||||
packed_recv_x, packed_recv_x_scales, packed_recv_count, packed_recv_src_info, packed_recv_layout_range, event, hook = \
|
||||
self._fallback_dispatch(x, topk_idx, num_max_dispatch_tokens_per_rank, num_experts, use_fp8, return_recv_hook)
|
||||
backend_active_ranks = get_active_ranks(self.backend).to(device=active_ranks.device, dtype=active_ranks.dtype)
|
||||
if active_ranks.numel() == backend_active_ranks.numel():
|
||||
active_ranks.copy_(backend_active_ranks)
|
||||
else:
|
||||
packed_recv_x, packed_recv_x_scales, packed_recv_count, packed_recv_src_info, packed_recv_layout_range, event, hook = \
|
||||
self.runtime.dispatch(x, topk_idx, active_ranks,
|
||||
num_max_dispatch_tokens_per_rank, num_experts, timeout_us,
|
||||
use_fp8, async_finish, return_recv_hook)
|
||||
handle = (packed_recv_src_info, packed_recv_layout_range, num_max_dispatch_tokens_per_rank, x.size(1), num_experts)
|
||||
tensors_to_record = (x, topk_idx,
|
||||
packed_recv_x, packed_recv_x_scales, packed_recv_count,
|
||||
|
|
@ -144,13 +158,259 @@ class Buffer:
|
|||
return_recv_hook: bool = False, out: Optional[torch.Tensor] = None) -> \
|
||||
Tuple[torch.Tensor, EventOverlap, Callable]:
|
||||
src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
|
||||
combined_x, event, hook = self.runtime.combine(x, topk_idx, topk_weights, src_info, layout_range,
|
||||
active_ranks,
|
||||
num_max_dispatch_tokens_per_rank, num_experts, timeout_us,
|
||||
zero_copy, async_finish, return_recv_hook, out)
|
||||
if self._use_fallback:
|
||||
from mooncake.ep import get_active_ranks
|
||||
combined_x, event, hook = self._fallback_combine(x, topk_idx, topk_weights, src_info, layout_range,
|
||||
num_max_dispatch_tokens_per_rank, num_experts,
|
||||
zero_copy, return_recv_hook, out)
|
||||
backend_active_ranks = get_active_ranks(self.backend).to(device=active_ranks.device, dtype=active_ranks.dtype)
|
||||
if active_ranks.numel() == backend_active_ranks.numel():
|
||||
active_ranks.copy_(backend_active_ranks)
|
||||
else:
|
||||
combined_x, event, hook = self.runtime.combine(x, topk_idx, topk_weights, src_info, layout_range,
|
||||
active_ranks,
|
||||
num_max_dispatch_tokens_per_rank, num_experts, timeout_us,
|
||||
zero_copy, async_finish, return_recv_hook, out)
|
||||
tensors_to_record = (x, topk_idx, topk_weights, src_info, layout_range, combined_x)
|
||||
return combined_x, EventOverlap(event, tensors_to_record if async_finish else None), hook
|
||||
|
||||
def get_next_combine_buffer(self, handle: object):
|
||||
src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
|
||||
if self._use_fallback:
|
||||
if self._fallback_next_combine_buffer is None or \
|
||||
self._fallback_next_combine_buffer.shape != (num_experts // self.group_size, num_max_dispatch_tokens_per_rank * self.group_size, hidden):
|
||||
self._fallback_next_combine_buffer = torch.empty(
|
||||
(num_experts // self.group_size, num_max_dispatch_tokens_per_rank * self.group_size, hidden),
|
||||
dtype=torch.bfloat16, device='cuda')
|
||||
return self._fallback_next_combine_buffer
|
||||
return self.runtime.get_next_combine_buffer(num_max_dispatch_tokens_per_rank, hidden, num_experts)
|
||||
|
||||
# -----------------
|
||||
# Fallback helpers
|
||||
# -----------------
|
||||
class _DummyEvent:
|
||||
def current_stream_wait(self):
|
||||
torch.cuda.synchronize()
|
||||
|
||||
@staticmethod
|
||||
def _fp8_cast(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
assert x.dim() == 2 and x.size(1) % 128 == 0
|
||||
m, n = x.shape
|
||||
x_view = x.view(m, -1, 128)
|
||||
x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
|
||||
x_fp8 = (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view(m, n)
|
||||
x_scales = (x_amax / 448.0).view(m, -1)
|
||||
return x_fp8, x_scales
|
||||
|
||||
def _fallback_dispatch(self, x: torch.Tensor, topk_idx: torch.Tensor,
|
||||
num_max_dispatch_tokens_per_rank: int, num_experts: int,
|
||||
use_fp8: bool, return_recv_hook: bool):
|
||||
from mooncake.ep import get_active_ranks
|
||||
with torch.profiler.record_function('dispatch'):
|
||||
num_tokens, hidden = x.shape
|
||||
k = topk_idx.size(1)
|
||||
num_ranks = self.group_size
|
||||
num_local_experts = num_experts // num_ranks
|
||||
|
||||
# Gather sizes first to handle variable num_tokens per rank
|
||||
num_tokens_tensor = torch.tensor([num_tokens], dtype=torch.int64, device=x.device)
|
||||
num_tokens_list = [torch.empty(1, dtype=torch.int64, device=x.device) for _ in range(num_ranks)]
|
||||
dist.all_gather(num_tokens_list, num_tokens_tensor, group=self.group)
|
||||
num_tokens_per_rank = [t.item() for t in num_tokens_list]
|
||||
backend_active_ranks = get_active_ranks(self.backend).tolist()
|
||||
for i in range(num_ranks):
|
||||
if backend_active_ranks[i] == 0:
|
||||
num_tokens_per_rank[i] = 0
|
||||
max_num_tokens = max(num_tokens_per_rank)
|
||||
|
||||
# Pad inputs to max_num_tokens for all_gather (all ranks must have same shape)
|
||||
if num_tokens < max_num_tokens:
|
||||
pad_size = max_num_tokens - num_tokens
|
||||
x_padded = torch.cat([x, torch.zeros((pad_size, hidden), dtype=x.dtype, device=x.device)], dim=0)
|
||||
topk_padded = torch.cat([topk_idx, torch.full((pad_size, k), -1, dtype=topk_idx.dtype, device=x.device)], dim=0)
|
||||
else:
|
||||
x_padded = x
|
||||
topk_padded = topk_idx
|
||||
|
||||
num_max_dispatch_tokens = num_ranks * num_max_dispatch_tokens_per_rank
|
||||
|
||||
# Gather inputs from all ranks (all have same shape after padding)
|
||||
all_x = torch.empty((num_ranks, max_num_tokens, hidden), dtype=x.dtype, device=x.device)
|
||||
dist.all_gather_into_tensor(all_x, x_padded, group=self.group)
|
||||
all_topk = torch.empty((num_ranks, max_num_tokens, k), dtype=topk_idx.dtype, device=x.device)
|
||||
dist.all_gather_into_tensor(all_topk, topk_padded, group=self.group)
|
||||
|
||||
# Prepare outputs per local expert
|
||||
recv_x_list: List[torch.Tensor] = []
|
||||
recv_x_scales_list: List[torch.Tensor] = []
|
||||
recv_count = torch.zeros((num_local_experts,), dtype=torch.int32, device=x.device)
|
||||
recv_src_info = torch.full((num_local_experts, num_max_dispatch_tokens), -1, dtype=torch.int32, device=x.device)
|
||||
layout_range = torch.zeros((num_local_experts, num_ranks), dtype=torch.int64, device=x.device)
|
||||
|
||||
for le in range(num_local_experts):
|
||||
expert_id = self.rank * num_local_experts + le
|
||||
# Collect tokens from all ranks that route to this expert
|
||||
tokens_per_rank_tensors: List[torch.Tensor] = []
|
||||
for src_rank in range(num_ranks):
|
||||
src_num_tokens = num_tokens_per_rank[src_rank]
|
||||
src_topk = all_topk[src_rank, :src_num_tokens] # Only consider valid tokens
|
||||
# Find tokens that route to this expert
|
||||
pos = (src_topk == expert_id).any(dim=1).nonzero(as_tuple=False).view(-1)
|
||||
tokens_per_rank_tensors.append(pos)
|
||||
|
||||
# Build ordered list grouped by src_rank (matching CUDA kernel behavior)
|
||||
begin = 0
|
||||
ordered_src_ranks_list: List[torch.Tensor] = []
|
||||
ordered_token_indices_list: List[torch.Tensor] = []
|
||||
for src_rank, tokens in enumerate(tokens_per_rank_tensors):
|
||||
count = tokens.numel()
|
||||
if count > 0:
|
||||
layout_range[le, src_rank] = (begin << 32) | count
|
||||
ordered_src_ranks_list.append(torch.full_like(tokens, src_rank))
|
||||
ordered_token_indices_list.append(tokens)
|
||||
begin += count
|
||||
else:
|
||||
layout_range[le, src_rank] = 0
|
||||
|
||||
if ordered_src_ranks_list:
|
||||
ordered_src_ranks = torch.cat(ordered_src_ranks_list)
|
||||
ordered_token_indices = torch.cat(ordered_token_indices_list)
|
||||
else:
|
||||
ordered_src_ranks = torch.empty(0, dtype=topk_idx.dtype, device=x.device)
|
||||
ordered_token_indices = torch.empty(0, dtype=topk_idx.dtype, device=x.device)
|
||||
|
||||
num_valid = min(ordered_src_ranks.numel(), num_max_dispatch_tokens)
|
||||
recv_count[le] = num_valid
|
||||
|
||||
# Materialize data
|
||||
if num_valid > 0:
|
||||
gathered = all_x[ordered_src_ranks[:num_valid], ordered_token_indices[:num_valid]]
|
||||
src_meta = ordered_token_indices[:num_valid].to(dtype=torch.int32)
|
||||
else:
|
||||
gathered = torch.empty((num_valid, hidden), dtype=torch.bfloat16, device=x.device)
|
||||
src_meta = torch.empty((num_valid,), dtype=torch.int32, device=x.device)
|
||||
|
||||
# Pad to full size
|
||||
if use_fp8:
|
||||
pad = num_max_dispatch_tokens - num_valid
|
||||
if pad > 0:
|
||||
pad_tensor = torch.zeros((pad, hidden), dtype=torch.bfloat16, device=x.device)
|
||||
gathered = torch.cat([gathered, pad_tensor], dim=0)
|
||||
fp8, scales = self._fp8_cast(gathered)
|
||||
recv_x_list.append(fp8)
|
||||
recv_x_scales_list.append(scales)
|
||||
else:
|
||||
pad = num_max_dispatch_tokens - num_valid
|
||||
if pad > 0:
|
||||
pad_tensor = torch.zeros((pad, hidden), dtype=torch.bfloat16, device=x.device)
|
||||
gathered = torch.cat([gathered, pad_tensor], dim=0)
|
||||
recv_x_list.append(gathered)
|
||||
|
||||
# src info padded
|
||||
if num_valid > 0:
|
||||
recv_src_info[le, :num_valid] = src_meta
|
||||
|
||||
if use_fp8:
|
||||
packed_recv_x = torch.stack(recv_x_list, dim=0) if len(recv_x_list) > 0 else torch.empty((0, num_max_dispatch_tokens, hidden), dtype=torch.float8_e4m3fn, device=x.device)
|
||||
# Calculate scales shape correctly
|
||||
num_scales_per_token = hidden // 128
|
||||
packed_recv_x_scales = torch.stack(recv_x_scales_list, dim=0) if len(recv_x_scales_list) > 0 else torch.empty((0, num_max_dispatch_tokens, num_scales_per_token), dtype=torch.float32, device=x.device)
|
||||
else:
|
||||
packed_recv_x = torch.stack(recv_x_list, dim=0) if len(recv_x_list) > 0 else torch.empty((0, num_max_dispatch_tokens, hidden), dtype=torch.bfloat16, device=x.device)
|
||||
packed_recv_x_scales = None
|
||||
|
||||
# Allocate zero-copy buffer for next combine
|
||||
self._fallback_next_combine_buffer = torch.empty((num_local_experts, num_max_dispatch_tokens, hidden), dtype=torch.bfloat16, device=x.device)
|
||||
|
||||
hook = (lambda: None) if return_recv_hook else (lambda: None)
|
||||
event = Buffer._DummyEvent()
|
||||
return packed_recv_x, packed_recv_x_scales, recv_count, recv_src_info, layout_range, event, hook
|
||||
|
||||
def _fallback_combine(self, x: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor,
|
||||
src_info: torch.Tensor, layout_range: torch.Tensor,
|
||||
num_max_dispatch_tokens_per_rank: int, num_experts: int,
|
||||
zero_copy: bool, return_recv_hook: bool, out: Optional[torch.Tensor]):
|
||||
from mooncake.ep import get_active_ranks
|
||||
with torch.profiler.record_function('combine'):
|
||||
num_tokens = topk_idx.size(0)
|
||||
hidden = (x if not zero_copy else self._fallback_next_combine_buffer).size(-1)
|
||||
num_ranks = self.group_size
|
||||
num_local_experts = num_experts // num_ranks
|
||||
|
||||
# Gather sizes first to handle variable num_tokens per rank
|
||||
num_tokens_tensor = torch.tensor([num_tokens], dtype=torch.int64, device=topk_idx.device)
|
||||
num_tokens_list = [torch.empty(1, dtype=torch.int64, device=topk_idx.device) for _ in range(num_ranks)]
|
||||
dist.all_gather(num_tokens_list, num_tokens_tensor, group=self.group)
|
||||
num_tokens_per_rank = [t.item() for t in num_tokens_list]
|
||||
backend_active_ranks = get_active_ranks(self.backend).tolist()
|
||||
for i in range(num_ranks):
|
||||
if backend_active_ranks[i] == 0:
|
||||
num_tokens_per_rank[i] = 0
|
||||
max_num_tokens = max(num_tokens_per_rank)
|
||||
|
||||
# Gather routing info across ranks to fetch per-token weights
|
||||
k = topk_idx.size(1)
|
||||
# Pad to max_num_tokens for all_gather
|
||||
if num_tokens < max_num_tokens:
|
||||
pad_size = max_num_tokens - num_tokens
|
||||
topk_padded = torch.cat([topk_idx, torch.full((pad_size, k), -1, dtype=topk_idx.dtype, device=topk_idx.device)], dim=0)
|
||||
topk_w_padded = torch.cat([topk_weights, torch.zeros((pad_size, k), dtype=topk_weights.dtype, device=topk_weights.device)], dim=0)
|
||||
else:
|
||||
topk_padded = topk_idx
|
||||
topk_w_padded = topk_weights
|
||||
|
||||
all_topk_idx = torch.empty((num_ranks, max_num_tokens, k), dtype=topk_idx.dtype, device=topk_idx.device)
|
||||
dist.all_gather_into_tensor(all_topk_idx, topk_padded, group=self.group)
|
||||
all_topk_w = torch.empty((num_ranks, max_num_tokens, k), dtype=topk_weights.dtype, device=topk_weights.device)
|
||||
dist.all_gather_into_tensor(all_topk_w, topk_w_padded, group=self.group)
|
||||
|
||||
expert_buffers = self._fallback_next_combine_buffer if zero_copy else x
|
||||
# Ensure bf16 input for accumulation
|
||||
if expert_buffers.dtype != torch.bfloat16:
|
||||
# FP8 path should already have been cast back by caller before combine in tests
|
||||
expert_buffers = expert_buffers.to(torch.bfloat16)
|
||||
|
||||
# Build send buffer [num_ranks, max_num_tokens, hidden]
|
||||
send_buf = torch.zeros((num_ranks, max_num_tokens, hidden), dtype=torch.bfloat16, device=expert_buffers.device)
|
||||
|
||||
for le in range(num_local_experts):
|
||||
expert_id = self.rank * num_local_experts + le
|
||||
# layout_range[le, j]: upper 32 begin, lower 32 count
|
||||
for src_rank in range(num_ranks):
|
||||
entry = layout_range[le, src_rank]
|
||||
begin = (entry >> 32).item() & 0xFFFFFFFF
|
||||
count = (entry & ((1 << 32) - 1)).item()
|
||||
if count == 0:
|
||||
continue
|
||||
tokens = src_info[le, begin:begin + count].to(torch.long)
|
||||
contrib = expert_buffers[le, begin:begin + count]
|
||||
|
||||
# Get source rank's actual token count and validate tokens
|
||||
src_num_tokens = num_tokens_per_rank[src_rank]
|
||||
valid_mask = tokens < src_num_tokens
|
||||
|
||||
if valid_mask.any():
|
||||
tokens_valid = tokens[valid_mask]
|
||||
contrib_valid = contrib[valid_mask]
|
||||
|
||||
# Find the per-token weight for this expert on src_rank
|
||||
idx_rows = all_topk_idx[src_rank, tokens_valid] # [count_valid, k]
|
||||
w_rows = all_topk_w[src_rank, tokens_valid] # [count_valid, k]
|
||||
mask = (idx_rows == expert_id)
|
||||
weights = (w_rows * mask).sum(dim=1).view(-1, 1)
|
||||
send_buf[src_rank, tokens_valid] += (contrib_valid * weights)
|
||||
|
||||
# All-reduce then take local slice (only valid tokens)
|
||||
dist.all_reduce(send_buf, group=self.group)
|
||||
combined_x = send_buf[self.rank, :num_tokens]
|
||||
|
||||
# Write to out if provided
|
||||
if out is not None:
|
||||
out.copy_(combined_x)
|
||||
combined = out
|
||||
else:
|
||||
combined = combined_x
|
||||
|
||||
hook = (lambda: None) if return_recv_hook else (lambda: None)
|
||||
event = Buffer._DummyEvent()
|
||||
return combined, event, hook
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import torch
|
|||
import torch.distributed as dist
|
||||
from functools import partial
|
||||
|
||||
import mooncake.ep
|
||||
from mooncake.mooncake_ep_buffer import Buffer
|
||||
from ep_test_utils import init_dist, bench, bench_kineto, calc_diff, hash_tensor, per_token_cast_back
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue