diff --git a/csrc/standalone/custom_ar/moe_allgather.cu b/csrc/standalone/custom_ar/moe_allgather.cu index a2704047bb9..16bf945e9cb 100644 --- a/csrc/standalone/custom_ar/moe_allgather.cu +++ b/csrc/standalone/custom_ar/moe_allgather.cu @@ -1,25 +1,27 @@ -// Standalone fused MoE all-gather kernel for EP dispatch. +// Lamport-based MoE all-gather kernel for EP dispatch. // -// JIT-compilable via torch.utils.cpp_extension — no vLLM build required. +// Replaces the flag-barrier approach with a Lamport sentinel protocol +// (inspired by FlashInfer's trtllm_allreduce_fusion). // -// Gathers the MoE dispatch tensors from all EP ranks in a single kernel: -// - topk_ids [N, topk] int32 -// - topk_weights [N, topk] float32 / bfloat16 -// - hidden_states [N, D_h] uint8 (NVFP4) / bfloat16 -// - quant_scales [N, D_s] (optional) any dtype +// Key advantages over the flag-barrier approach: +// - No explicit barriers (sentinels provide per-element synchronization). +// - Push model: NVLink writes (fire-and-forget) instead of NVLink reads. +// - Triple buffering: no end barrier needed. // -// Each tensor is packed into a pre-registered IPC buffer at 16-byte-aligned -// offsets. The kernel: -// 1. Copies inputs into the IPC buffer (local SM write). -// 2. Barrier — all ranks' writes become visible via NVLink. -// 3. Gathers from all peers' buffers into separate output tensors. -// 4. Barrier — done. +// Gathers the MoE dispatch tensors from all EP ranks: +// - topk_ids [N, topk] int32 +// - topk_weights [N, topk] float32 / bfloat16 +// - hidden_states [N, D_h] uint8 (NVFP4) / bfloat16 +// - quant_scales [N, D_s] (optional) // -// Under CUDA graphs, the input tensor addresses are fixed and the buffer -// copies are captured. The total overhead is dominated by the two barrier -// round-trips (~5µs each on NVLink). +// Double-buffer layout in each rank's IPC buffer: +// [Segment 0][Segment 1] +// Each segment: [Rank 0 slot][Rank 1 slot]...[Rank N-1 slot] +// Each rank slot: packed tensors at 16-byte aligned offsets. // -// All data movement uses 128-bit (int4) loads/stores. +// Sentinel: 0x80000000 (negative-zero in float32). The writer replaces +// any data word matching the sentinel with 0 before pushing. The reader +// spin-loads (volatile) until no sentinel words remain in the vector. #include #include @@ -28,193 +30,166 @@ #include #include -// --------------------------------------------------------------------------- -// Flag-based barrier (from custom_all_reduce.cuh, standalone) -// --------------------------------------------------------------------------- - -constexpr int kMaxBlocks = 36; - -using FlagType = uint32_t; - -struct Signal { - alignas(128) FlagType start[kMaxBlocks][8]; - alignas(128) FlagType end[kMaxBlocks][8]; - alignas(128) FlagType _flag[kMaxBlocks]; -}; - -struct __align__(16) RankData { - const void* ptrs[8]; -}; - -struct __align__(16) RankSignals { - Signal* signals[8]; -}; - #define DINLINE __device__ __forceinline__ -static DINLINE void st_flag_volatile(FlagType* addr, FlagType val) { - asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(val), "l"(addr)); -} +constexpr uint32_t SENTINEL = 0x80000000u; +constexpr int kMaxBlocks = 36; -static DINLINE FlagType ld_flag_volatile(FlagType* addr) { - FlagType v; - asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(v) : "l"(addr)); +// --------------------------------------------------------------------------- +// Volatile 128-bit load/store and sentinel helpers +// --------------------------------------------------------------------------- + +static DINLINE int4 ld128v(const void* addr) { + int4 v; + asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(addr)); return v; } -template -DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, - int rank) { - FlagType flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - st_flag_volatile(&sg.signals[threadIdx.x]->start[blockIdx.x][rank], flag); - while (ld_flag_volatile(&self_sg->start[blockIdx.x][threadIdx.x]) != flag); - } - __syncthreads(); - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +static DINLINE bool has_sentinel(int4 v) { + return reinterpret_cast(v.x) == SENTINEL | + reinterpret_cast(v.y) == SENTINEL | + reinterpret_cast(v.z) == SENTINEL | + reinterpret_cast(v.w) == SENTINEL; } -template -DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { - __syncthreads(); - FlagType flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - st_flag_volatile(&sg.signals[threadIdx.x]->end[blockIdx.x][rank], flag); - while (ld_flag_volatile(&self_sg->end[blockIdx.x][threadIdx.x]) != flag); - } - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +static DINLINE int4 remove_sentinel(int4 v) { + if (reinterpret_cast(v.x) == SENTINEL) v.x = 0; + if (reinterpret_cast(v.y) == SENTINEL) v.y = 0; + if (reinterpret_cast(v.z) == SENTINEL) v.z = 0; + if (reinterpret_cast(v.w) == SENTINEL) v.w = 0; + return v; } // --------------------------------------------------------------------------- -// Fused MoE dispatch all-gather kernel +// Lamport all-gather kernel // --------------------------------------------------------------------------- - -// The kernel has 4 phases: -// 1. Pack local inputs into IPC buffer. -// 2. Barrier (all ranks' data visible). -// 3. Gather: one tight contiguous read from each peer → flat staging buffer. -// 4. Barrier (gather done, safe to overwrite IPC buffer next iteration). -// 5. Scatter: redistribute the flat staging into per-tensor outputs (local -// L2). // -// Phase 3 is NVLink-critical: one loop, no conditionals, all peers pipelined. -// Phase 5 is local-memory only (L2 speed), runs AFTER the end barrier. - -// has_scales: compile-time flag for the optional 4th tensor (quant_scales). -// -// Future optimization (TODO): register hidden_states via IPC (like custom -// allreduce does in graph mode) to skip its Phase 1 copy. This would save -// ~3µs by eliminating the hidden_states copy + reducing the scatter. -// Requires CUDA graph integration to register the hidden_states tensor -// address during capture. +// Phase 1 — PUSH: each rank writes its packed data to ALL peers' current +// segment via regular stores (NVLink push, fire-and-forget). +// Phase 2 — CLEAR: each rank writes sentinels to the OLDEST segment of +// its own buffer, preparing it for reuse. +// Phase 3 — POLL + SCATTER: each rank volatile-loads from its own current +// segment, spinning until sentinels disappear, then scatters +// directly to per-tensor output arrays. +// Phase 4 — ADVANCE: one thread advances the triple-buffer ring counter. template -__global__ void __launch_bounds__(512, 1) - moe_allgather_kernel(RankData* _dp, RankSignals sg, Signal* self_sg, - int rank, - // up to 4 inputs - const void* inp0, const void* inp1, const void* inp2, - const void* inp3, int off0, int sz0, int off1, int sz1, - int off2, int sz2, int off3, int sz3, void* out0, - void* out1, void* out2, void* out3, void* staging, - int total_sz) { +__global__ void __launch_bounds__(512, 1) moe_allgather_lamport_kernel( + int64_t* buf_ptrs, // [ngpus] IPC buffer base addresses (device) + int* counters, // [0] = unused, [1] = ring (0/1/2), [2] = prev total_sz + int rank, + int seg_capacity, // bytes per segment + int rank_stride, // bytes per rank-slot within a segment + int total_sz, // int4 units of actual packed data per rank + // inputs (up to 4) + const void* inp0, const void* inp1, const void* inp2, const void* inp3, + int off0, int sz0, int off1, int sz1, int off2, int sz2, int off3, int sz3, + // outputs (up to 4) + void* out0, void* out1, void* out2, void* out3) { using V = int4; + const int tid = blockIdx.x * blockDim.x + threadIdx.x; + const int stride = gridDim.x * blockDim.x; + + // Read segment index and previous clear size. + const int seg = counters[1]; // 0 or 1 + const int prev_total_sz = counters[2]; // set by previous invocation + const int cur_seg = seg; + const int old_seg = 1 - seg; + + char* bufs[ngpus]; +#pragma unroll + for (int r = 0; r < ngpus; r++) + bufs[r] = reinterpret_cast(buf_ptrs[r]) + cur_seg * seg_capacity; + + // Sentinel vector for clearing. + V sent; + sent.x = sent.y = sent.z = sent.w = static_cast(SENTINEL); + + // ---- Phase 1: PUSH local data to ALL peers ---- + // Write to peer_r's buffer at [rank * rank_stride + off_i]. + +#define PUSH(idx, inp_ptr, off_val, sz_val) \ + if constexpr (nbufs > (idx)) { \ + const V* src = reinterpret_cast(inp_ptr); \ + for (int i = tid; i < (sz_val); i += stride) { \ + V val = remove_sentinel(src[i]); \ + _Pragma("unroll") for (int r = 0; r < ngpus; r++) { \ + reinterpret_cast(bufs[r] + rank * rank_stride + (off_val))[i] = \ + val; \ + } \ + } \ + } + + PUSH(0, inp0, off0, sz0) + PUSH(1, inp1, off1, sz1) + PUSH(2, inp2, off2, sz2) + PUSH(3, inp3, off3, sz3) +#undef PUSH + + // ---- Phase 2: CLEAR only the previously-written data in oldest segment ---- + // Only clear what the previous invocation actually wrote (per rank-slot). + if (prev_total_sz > 0) { + char* clr_base = + reinterpret_cast(buf_ptrs[rank]) + old_seg * seg_capacity; +#pragma unroll + for (int r = 0; r < ngpus; r++) { + V* clr = reinterpret_cast(clr_base + r * rank_stride); + for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent; + } + } + + // ---- Phase 3: POLL + SCATTER ---- + // Volatile-load from own buffer; spin until sentinel gone; scatter to output. + char* my = bufs[rank]; + +#define POLL(idx, out_ptr, off_val, sz_val) \ + if constexpr (nbufs > (idx)) { \ + for (int i = tid; i < (sz_val); i += stride) { \ + _Pragma("unroll") for (int s = 0; s < ngpus; s++) { \ + V val; \ + do { \ + val = ld128v( \ + reinterpret_cast(my + s * rank_stride + (off_val)) + i); \ + } while (has_sentinel(val)); \ + reinterpret_cast(out_ptr)[s * (sz_val) + i] = val; \ + } \ + } \ + } + + POLL(0, out0, off0, sz0) + POLL(1, out1, off1, sz1) + POLL(2, out2, off2, sz2) + POLL(3, out3, off3, sz3) +#undef POLL + + // ---- Phase 4: ADVANCE ring counter + store clear size for next call ---- + // Stream serialization ensures the next kernel sees these updates. + if (blockIdx.x == 0 && threadIdx.x == 0) { + counters[1] = 1 - seg; + counters[2] = total_sz; + } +} + +// --------------------------------------------------------------------------- +// Sentinel initialization kernel +// --------------------------------------------------------------------------- + +__global__ void lamport_init_kernel(uint32_t* buf, int n) { int tid = blockIdx.x * blockDim.x + threadIdx.x; int stride = gridDim.x * blockDim.x; - auto dp = *_dp; - char* my_buf = (char*)dp.ptrs[rank]; - - // Phase 1: pack local inputs into IPC buffer. - if constexpr (nbufs > 0) { - const V* s = (const V*)inp0; - V* d = (V*)(my_buf + off0); - for (int i = tid; i < sz0; i += stride) d[i] = s[i]; - } - if constexpr (nbufs > 1) { - const V* s = (const V*)inp1; - V* d = (V*)(my_buf + off1); - for (int i = tid; i < sz1; i += stride) d[i] = s[i]; - } - if constexpr (nbufs > 2) { - const V* s = (const V*)inp2; - V* d = (V*)(my_buf + off2); - for (int i = tid; i < sz2; i += stride) d[i] = s[i]; - } - if constexpr (nbufs > 3) { - const V* s = (const V*)inp3; - V* d = (V*)(my_buf + off3); - for (int i = tid; i < sz3; i += stride) d[i] = s[i]; - } - - __threadfence_system(); - - // Phase 2: barrier. - barrier_at_start(sg, self_sg, rank); - - // Phase 3: single contiguous gather into staging buffer. - { - const V* peers[ngpus]; -#pragma unroll - for (int s = 0; s < ngpus; s++) peers[s] = (const V*)dp.ptrs[s]; - - for (int i = tid; i < total_sz; i += stride) { -#pragma unroll - for (int s = 0; s < ngpus; s++) - ((V*)staging)[s * total_sz + i] = peers[s][i]; - } - } - - // Phase 4: end barrier. - barrier_at_end(sg, self_sg, rank); - - // Phase 5: scatter from staging to per-tensor outputs (local L2). - { - const V* stg = (const V*)staging; - - if constexpr (nbufs > 0) { - const int b0 = off0 / (int)sizeof(V); - for (int i = tid; i < sz0; i += stride) { -#pragma unroll - for (int s = 0; s < ngpus; s++) - ((V*)out0)[s * sz0 + i] = stg[s * total_sz + b0 + i]; - } - } - if constexpr (nbufs > 1) { - const int b1 = off1 / (int)sizeof(V); - for (int i = tid; i < sz1; i += stride) { -#pragma unroll - for (int s = 0; s < ngpus; s++) - ((V*)out1)[s * sz1 + i] = stg[s * total_sz + b1 + i]; - } - } - if constexpr (nbufs > 2) { - const int b2 = off2 / (int)sizeof(V); - for (int i = tid; i < sz2; i += stride) { -#pragma unroll - for (int s = 0; s < ngpus; s++) - ((V*)out2)[s * sz2 + i] = stg[s * total_sz + b2 + i]; - } - } - if constexpr (nbufs > 3) { - const int b3 = off3 / (int)sizeof(V); - for (int i = tid; i < sz3; i += stride) { -#pragma unroll - for (int s = 0; s < ngpus; s++) - ((V*)out3)[s * sz3 + i] = stg[s * total_sz + b3 + i]; - } - } - } + for (int i = tid; i < n; i += stride) buf[i] = SENTINEL; } // --------------------------------------------------------------------------- // Host launcher // --------------------------------------------------------------------------- -// Compute 16-byte-aligned offset and int4-unit size for one tensor. struct TensorDesc { void* inp; - int off; // byte offset in IPC buffer - int sz; // size in int4 (16-byte) units + int off; + int sz; int64_t nbytes; }; @@ -233,9 +208,16 @@ static TensorDesc make_desc(torch::Tensor& inp, int64_t& cursor) { return d; } -void moe_all_gather(int64_t rank_data_ptr, int64_t signals_ptr, - int64_t self_signal_ptr, int64_t rank, int64_t world_size, - std::vector& inputs, +void lamport_init(int64_t buf_ptr, int64_t nbytes) { + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + int n = static_cast(nbytes / 4); + lamport_init_kernel<<<256, 256, 0, stream>>>( + reinterpret_cast(buf_ptr), n); +} + +void moe_all_gather(int64_t buf_ptrs_ptr, int64_t counters_ptr, int64_t rank, + int64_t world_size, int64_t seg_capacity, + int64_t rank_stride, std::vector& inputs, std::vector& outputs) { auto stream = c10::cuda::getCurrentCUDAStream().stream(); int n = static_cast(inputs.size()); @@ -247,17 +229,16 @@ void moe_all_gather(int64_t rank_data_ptr, int64_t signals_ptr, for (int i = 0; i < n; i++) descs[i] = make_desc(inputs[i], cursor); TORCH_CHECK(cursor % 16 == 0); int total_sz = static_cast(cursor / 16); + TORCH_CHECK(cursor <= rank_stride, "packed data (", cursor, + " bytes) exceeds rank_stride (", rank_stride, " bytes)"); + int ws = static_cast(world_size); for (int i = 0; i < n; i++) { TORCH_CHECK(outputs[i].is_contiguous()); - TORCH_CHECK(outputs[i].numel() == inputs[i].numel() * world_size); + TORCH_CHECK(outputs[i].numel() == inputs[i].numel() * ws); } - auto* ptrs = reinterpret_cast(rank_data_ptr); - RankSignals sg = *reinterpret_cast(signals_ptr); - auto* self_sg = reinterpret_cast(self_signal_ptr); int r = static_cast(rank); - int threads = 512; int blocks = std::max(1, std::min(kMaxBlocks, (total_sz + threads - 1) / threads)); @@ -271,32 +252,33 @@ void moe_all_gather(int64_t rank_data_ptr, int64_t signals_ptr, outs[i] = outputs[i].data_ptr(); } - auto staging = torch::empty( - {(int64_t)world_size * total_sz * (int64_t)sizeof(int4)}, - torch::TensorOptions().dtype(torch::kUInt8).device(inputs[0].device())); + auto* bp = reinterpret_cast(buf_ptrs_ptr); + auto* ct = reinterpret_cast(counters_ptr); + int sc = static_cast(seg_capacity); + int rs = static_cast(rank_stride); -#define KL(ngpus, nb) \ - moe_allgather_kernel<<>>( \ - ptrs, sg, self_sg, r, inps[0], inps[1], inps[2], inps[3], offs[0], \ - szs[0], offs[1], szs[1], offs[2], szs[2], offs[3], szs[3], outs[0], \ - outs[1], outs[2], outs[3], staging.data_ptr(), total_sz); +#define KL(ng, nb) \ + moe_allgather_lamport_kernel<<>>( \ + bp, ct, r, sc, rs, total_sz, inps[0], inps[1], inps[2], inps[3], \ + offs[0], szs[0], offs[1], szs[1], offs[2], szs[2], offs[3], szs[3], \ + outs[0], outs[1], outs[2], outs[3]); -#define GPU_CASE(ngpus) \ - case ngpus: \ - switch (n) { \ - case 2: \ - KL(ngpus, 2); \ - break; \ - case 3: \ - KL(ngpus, 3); \ - break; \ - case 4: \ - KL(ngpus, 4); \ - break; \ - } \ +#define GPU_CASE(ng) \ + case ng: \ + switch (n) { \ + case 2: \ + KL(ng, 2); \ + break; \ + case 3: \ + KL(ng, 3); \ + break; \ + case 4: \ + KL(ng, 4); \ + break; \ + } \ break; - switch (world_size) { + switch (ws) { GPU_CASE(2) GPU_CASE(4) GPU_CASE(6) @@ -313,6 +295,7 @@ void moe_all_gather(int64_t rank_data_ptr, int64_t signals_ptr, // --------------------------------------------------------------------------- PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("moe_all_gather", &moe_all_gather, - "Fused MoE dispatch all-gather with in-kernel scatter"); + m.def("moe_all_gather", &moe_all_gather, "Lamport MoE all-gather"); + m.def("lamport_init", &lamport_init, + "Initialize Lamport buffer with sentinels"); } diff --git a/csrc/standalone/custom_ar/moe_allgather.py b/csrc/standalone/custom_ar/moe_allgather.py index f0248aa5bde..0a97ec19eb9 100644 --- a/csrc/standalone/custom_ar/moe_allgather.py +++ b/csrc/standalone/custom_ar/moe_allgather.py @@ -1,12 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -Standalone fused MoE all-gather for EP dispatch. +Lamport-based fused MoE all-gather for EP dispatch. JIT-compiles the CUDA kernel on first use (cached afterwards). -The kernel gathers contiguously from all peers (NVLink-optimal), -then scatters to per-tensor outputs inside the kernel (local L2). -Zero Python-side post-processing. +Uses a Lamport sentinel protocol (push writes + per-element sync) +with triple-buffered IPC regions — no explicit barriers. Usage: ag = MoeAllGather(custom_allreduce) @@ -41,28 +40,43 @@ def _load_lib(): class MoeAllGather: - """Fused MoE dispatch all-gather backed by a one-shot flag-barrier kernel.""" + """Lamport-based MoE dispatch all-gather with triple buffering.""" def __init__(self, ca_comm): self.rank = ca_comm.rank self.world_size = ca_comm.world_size self.device = ca_comm.device - self.meta_ptrs = ca_comm.meta_ptrs self.buffer_ptrs = ca_comm.buffer_ptrs self.max_size = ca_comm.max_size - self._rank_signals = torch.zeros(8, dtype=torch.int64) - for i in range(self.world_size): - self._rank_signals[i] = self.meta_ptrs[i] - self._rank_signals_ptr = self._rank_signals.data_ptr() - self._self_signal_ptr = self.meta_ptrs[self.rank] + ws = self.world_size + # Double-buffer layout: 2 segments, each with ws rank-slots. + # Safe because kernels in the same stream are serialized, and the + # Lamport poll ensures all cross-GPU pushes complete before the + # kernel returns. + # seg_capacity and rank_stride are 16-byte aligned. + self.seg_capacity = (self.max_size // 2) & ~15 + self.rank_stride = (self.seg_capacity // ws) & ~15 + self.max_per_rank = self.rank_stride # max packed bytes per rank - self._rank_data = torch.zeros( + # Buffer pointer array on device. + self._buf_ptrs = torch.zeros( 8, dtype=torch.int64, device=f"cuda:{self.device.index}" ) - for i in range(self.world_size): - self._rank_data[i] = self.buffer_ptrs[i] - self._rank_data_ptr = self._rank_data.data_ptr() + for i in range(ws): + self._buf_ptrs[i] = self.buffer_ptrs[i] + self._buf_ptrs_ptr = self._buf_ptrs.data_ptr() + + # Counters on device: [0]=unused, [1]=ring index (0/1/2), [2]=prev_total_sz. + self._counters = torch.zeros( + 3, dtype=torch.int32, device=f"cuda:{self.device.index}" + ) + self._counters_ptr = self._counters.data_ptr() + + # Initialize ALL segments with sentinel values. + lib = _load_lib() + lib.lamport_init(self.buffer_ptrs[self.rank], self.max_size) + torch.accelerator.synchronize(self.device) def gather( self, @@ -84,11 +98,12 @@ class MoeAllGather: ] lib.moe_all_gather( - self._rank_data_ptr, - self._rank_signals_ptr, - self._self_signal_ptr, + self._buf_ptrs_ptr, + self._counters_ptr, self.rank, self.world_size, + self.seg_capacity, + self.rank_stride, inputs, outputs, ) diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 6130c3ea327..e0637e9c56e 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -103,11 +103,13 @@ def moe_dispatch( total_bytes = (total_bytes + 15) & ~15 total_bytes += nbytes - max_bytes = min(1 * 1024 * 1024, moe_ag.max_size) + max_bytes = moe_ag.max_per_rank if ok and total_bytes <= max_bytes: outputs = [ torch.empty( - (t.shape[0] * ws, *t.shape[1:]), dtype=t.dtype, device=t.device + (t.shape[0] * ws, *t.shape[1:]), + dtype=t.dtype, + device=t.device, ) for t in tensors ] @@ -118,11 +120,12 @@ def moe_dispatch( lib = _load_lib() lib.moe_all_gather( - moe_ag._rank_data_ptr, - moe_ag._rank_signals_ptr, - moe_ag._self_signal_ptr, + moe_ag._buf_ptrs_ptr, + moe_ag._counters_ptr, moe_ag.rank, ws, + moe_ag.seg_capacity, + moe_ag.rank_stride, tensors, outputs, ) @@ -385,10 +388,10 @@ class AgRsAll2AllManager(All2AllManagerBase): total_bytes = (total_bytes + 15) & ~15 total_bytes += nbytes - if total_bytes > moe_ag.max_size: + if total_bytes > moe_ag.max_per_rank: if _dbg < 10: _log.warning( - f"MoE AG: too large {total_bytes} > {moe_ag.max_size} " + f"MoE AG: too large {total_bytes} > {moe_ag.max_per_rank} " f"n={n} shapes={[t.shape for t in tensors]}" ) self._ag_dbg_count = _dbg + 1 @@ -406,11 +409,12 @@ class AgRsAll2AllManager(All2AllManagerBase): lib = _load_lib() lib.moe_all_gather( - moe_ag._rank_data_ptr, - moe_ag._rank_signals_ptr, - moe_ag._self_signal_ptr, + moe_ag._buf_ptrs_ptr, + moe_ag._counters_ptr, moe_ag.rank, ws, + moe_ag.seg_capacity, + moe_ag.rank_stride, tensors, outputs, )