diff --git a/mooncake-ep/include/mooncake_backend.h b/mooncake-ep/include/mooncake_backend.h index b012eb6f..951a2ac5 100644 --- a/mooncake-ep/include/mooncake_backend.h +++ b/mooncake-ep/include/mooncake_backend.h @@ -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_; diff --git a/mooncake-ep/include/mooncake_ep_buffer.h b/mooncake-ep/include/mooncake_ep_buffer.h index e3420e93..afddf149 100644 --- a/mooncake-ep/include/mooncake_ep_buffer.h +++ b/mooncake-ep/include/mooncake_ep_buffer.h @@ -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_; } diff --git a/mooncake-ep/src/mooncake_ep_buffer.cpp b/mooncake-ep/src/mooncake_ep_buffer.cpp index 71de0249..168736cf 100644 --- a/mooncake-ep/src/mooncake_ep_buffer.cpp +++ b/mooncake-ep/src/mooncake_ep_buffer.cpp @@ -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& remote_addrs, diff --git a/mooncake-integration/ep/ep_py.cpp b/mooncake-integration/ep/ep_py.cpp index 22fe6683..e9470e1c 100644 --- a/mooncake-integration/ep/ep_py.cpp +++ b/mooncake-integration/ep/ep_py.cpp @@ -57,12 +57,19 @@ std::string getPreferredHca(c10::intrusive_ptr backend, return mooncakeBackend->getPreferredHca(location); } +at::Tensor getActiveRanks(c10::intrusive_ptr backend) { + auto mooncakeBackend = + c10::static_intrusive_pointer_cast(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_>( @@ -79,6 +86,7 @@ PYBIND11_MODULE(ep, m) { py::class_(m, "Buffer") .def(py::init()) + .def("ibgda_disabled", &MooncakeEpBuffer::ibgda_disabled) .def("is_roce", &MooncakeEpBuffer::is_roce) .def("sync_ib", &MooncakeEpBuffer::sync_ib) .def("sync_roce", &MooncakeEpBuffer::sync_roce) diff --git a/mooncake-wheel/mooncake/mooncake_ep_buffer.py b/mooncake-wheel/mooncake/mooncake_ep_buffer.py index 5ba549de..22227d2d 100644 --- a/mooncake-wheel/mooncake/mooncake_ep_buffer.py +++ b/mooncake-wheel/mooncake/mooncake_ep_buffer.py @@ -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 diff --git a/mooncake-wheel/tests/test_mooncake_ep.py b/mooncake-wheel/tests/test_mooncake_ep.py index 2893928d..b73e9557 100644 --- a/mooncake-wheel/tests/test_mooncake_ep.py +++ b/mooncake-wheel/tests/test_mooncake_ep.py @@ -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