From cdec5214c33b81afce490d2af54562a52e10ca03 Mon Sep 17 00:00:00 2001 From: Xuan-1998 Date: Fri, 28 Aug 2026 03:00:28 +0000 Subject: [PATCH 1/2] perf(combine): schedule scale-out puts remote-first to mitigate node stagger The combine forward warps replayed tokens in the dispatch arrival order frozen into the handle, which interleaves local-bypass tokens (no RDMA) with remote tokens in a network-timing-dependent way. A node whose replay front-loads local tokens back-loads all of its scale-out puts so with the NIC near saturation the shift never drains and surfaces as the peer node's exit-wait tail, so one node runs ~62 GB/s SO while the other runs ~76. Here we make the schedule deterministic: the scale-up warps sweep the linked list twice (remote-destined tokens first, local-bypass second) and the forward warps replay in the same two-pass order, so count-based tail gating is unchanged and every rank issues its puts on the same schedule. Pass 0 flushes its TMA record and remainder batches before pass 1 so the wire keeps draining while locals are reduced. --- .../impls/hybrid_combine_unordered.cuh | 36 +++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/deep_ep/include/deep_ep/impls/hybrid_combine_unordered.cuh b/deep_ep/include/deep_ep/impls/hybrid_combine_unordered.cuh index e67b36a99..05755d449 100644 --- a/deep_ep/include/deep_ep/impls/hybrid_combine_unordered.cuh +++ b/deep_ep/include/deep_ep/impls/hybrid_combine_unordered.cuh @@ -188,6 +188,11 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, #pragma unroll for (int i = 0; i < kNumScaleupRanksPerLane; ++ i) stored_token_idx[i] = -1; + #pragma unroll 1 + for (int sweep = 0; sweep < 2; ++ sweep) { + #pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++ i) + stored_ll_idx[i] = 0; while (true) { // Load token indices in the list #pragma unroll @@ -211,10 +216,18 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, // Process tokens for all ranks together using bitmask to skip inactive ranks EP_STATIC_ASSERT(kNumScaleupRanks <= 64, "Too many scale-up ranks for 64-bit mask"); using mask_t = std::conditional_t<(kNumScaleupRanks <= 32), uint32_t, uint64_t>; + constexpr int kSweepMetadataStride = 2 + kNumTopk; mask_t wip_mask = 0; #pragma unroll - for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) - wip_mask |= static_cast(ptx::gather(stored_token_idx[j] >= 0)) << (j * 32); + for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) { + bool in_sweep = stored_token_idx[j] >= 0; + if (in_sweep) { + const auto src_global = __ldg(src_metadata + stored_token_idx[j] * kSweepMetadataStride); + const bool is_local = (src_global / (kNumMaxTokensPerRank * kNumScaleupRanks)) == scaleout_rank_idx; + in_sweep = (sweep == 0) != is_local; + } + wip_mask |= static_cast(ptx::gather(in_sweep)) << (j * 32); + } while (wip_mask) { // Find next active rank after `dst_scaleup_rank_idx` (round-robin) const auto start = (dst_scaleup_rank_idx + 1) % kNumScaleupRanks; @@ -382,6 +395,7 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, // Update for the unissued ones update_tails(true); + } } else if (warp_idx < kNumDataWarps) { const auto forward_warp_idx = warp_idx - kNumScaleupWarps; const auto channel_idx = sm_idx * kNumChannelsPerSM + forward_warp_idx; @@ -468,6 +482,8 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, // Replay the dispatch int stored_num_tokens_recv[kNumScaleupRanksPerLane] = {}, stored_cached_scaleup_tail[kNumScaleupRanksPerLane] = {}; + #pragma unroll 1 + for (int replay_pass = 0; replay_pass < 2; ++ replay_pass) { for (int i = 0; ; ++ i) { const auto src_token_global_idx = __ldg(token_metadata_at_forward + i * kNumForwardMetadataDims); const auto src_rank_idx = src_token_global_idx / kNumMaxTokensPerRank; @@ -480,6 +496,10 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, if (src_token_global_idx < 0) break; + // Two-pass schedule: deferred tokens are revisited by the other pass + if ((replay_pass == 0) == (src_scaleout_rank_idx == scaleout_rank_idx)) + continue; + // Scaleup rank mask EP_STATIC_ASSERT(kNumScaleupRanks <= 64, "Too many scale-up peers"); using mask_t = std::conditional_t; @@ -650,6 +670,18 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, } } + if (replay_pass == 0) { + flush_last_tma_and_record_batch(); + last_src_scaleout_rank_idx = -1; + if (ptx::elect_one_sync()) { + #pragma unroll + for (int dst = 0; dst < kNumScaleoutRanks; ++ dst) + issue_batched_rdma(dst); + } + __syncwarp(); + } + } + // Issue the last TMA and record operation for last RDMA if constexpr (kAllowMultipleReduction) flush_last_tma_and_record_batch(); From 3c737dcf0da5889ba7efd26e05b4808307cc38af Mon Sep 17 00:00:00 2001 From: Xuan Jiang Date: Fri, 28 Aug 2026 22:26:10 +0000 Subject: [PATCH 2/2] perf(combine): raise the GPU-to-NIC submission rate with channels, QPs, forward pairs Raise the unordered combine's submission rate three ways: drop the 4-channel/SM cap and let shared memory and the GIN signal budget decide (8/SM at 12 SM) for more parallel submitters, raise the default GIN contexts 11 -> 13 so every SM of the recommended 12-SM launch owns its QP instead of sharing WQE rings, and add cooperative forward warp pairs that split each token's hidden dim across two warps so tokens are reduced faster and therefore submitted faster. --- csrc/elastic/buffer.hpp | 6 +- csrc/kernels/elastic/combine.hpp | 20 +++- .../deep_ep/common/gin_resource_alloc.cuh | 4 +- .../impls/hybrid_combine_unordered.cuh | 106 ++++++++++++++---- 4 files changed, 106 insertions(+), 30 deletions(-) diff --git a/csrc/elastic/buffer.hpp b/csrc/elastic/buffer.hpp index a24907f2c..94d094f2d 100644 --- a/csrc/elastic/buffer.hpp +++ b/csrc/elastic/buffer.hpp @@ -928,8 +928,6 @@ class ElasticBuffer { get_num_notify_smem_bytes(nccl_context->num_ranks, num_experts) <= num_smem_bytes and "dispatch TMA pool exceeds the shared-memory budget"); - if (not prefer_overlap_with_compute) - num_channels_per_sm = std::min(num_channels_per_sm, 4); // Reduce the channel count to fit this launch's GIN indexed-signal budget. // `with_notify` is pinned (not `not cached_mode`) so a cached dispatch derives the // same count the handle was shaped with. @@ -945,7 +943,8 @@ class ElasticBuffer { static_cast(2 * num_channels_per_sm) * combine_token_layout.get_num_bytes() + elastic::ProxyRingLayout::get_num_bytes(num_channels_per_sm, - elastic::kProxyRingDepthDefault) > + elastic::kProxyRingDepthDefault) + + /* cooperative-forward pair counters */ 2 * num_channels_per_sm * static_cast(sizeof(int)) > num_smem_bytes) -- num_channels_per_sm; EP_HOST_ASSERT(num_channels_per_sm >= 1 and @@ -1429,6 +1428,7 @@ class ElasticBuffer { num_sms, jit::device_runtime->get_num_smem_bytes(), num_channels, use_expanded_layout, allow_multiple_reduction, + prefer_overlap_with_compute, comm_stream); // Allocate output tensors diff --git a/csrc/kernels/elastic/combine.hpp b/csrc/kernels/elastic/combine.hpp index 6a3f33c0f..84df7d458 100644 --- a/csrc/kernels/elastic/combine.hpp +++ b/csrc/kernels/elastic/combine.hpp @@ -40,6 +40,7 @@ class CombineRuntime final : public jit::LaunchRuntime { int num_topk; int num_qps; int64_t num_timeout_cycles; + int num_fw_warps_per_channel; // Parameters nv_bfloat16* x; @@ -77,7 +78,7 @@ class CombineRuntime final : public jit::LaunchRuntime { args.num_qps, args.num_timeout_cycles); } else { header_name = args.use_ordered_kernel ? "hybrid_combine" : "hybrid_combine_unordered"; - func_name = fmt::format("{}<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>", + func_name = fmt::format("{}<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}{}>", args.use_ordered_kernel ? "hybrid_combine_impl" : "hybrid_unordered_combine_impl", args.use_expanded_layout, args.allow_multiple_reduction, args.launch_args.grid_dim.first, @@ -88,7 +89,9 @@ class CombineRuntime final : public jit::LaunchRuntime { args.num_experts, args.num_topk, args.num_qps, - args.num_timeout_cycles); + args.num_timeout_cycles, + args.use_ordered_kernel ? std::string() : + fmt::format(", {}", args.num_fw_warps_per_channel)); } return fmt::format(R"( #include @@ -163,6 +166,7 @@ static void* launch_combine(void* x, const int& num_sms, const int& num_smem_bytes, const int& num_channels, const bool& use_expanded_layout, const bool& allow_multiple_reduction, + const bool& prefer_overlap_with_compute, const at::cuda::CUDAStream& stream) { // Maximize shared memory utilization const auto token_layout = get_combine_token_layout(hidden, sizeof(nv_bfloat16), num_topk); @@ -171,6 +175,7 @@ static void* launch_combine(void* x, // Decide warps const bool use_ordered_kernel = use_ordered_hybrid_kernel(); int num_scaleup_warps = 0, num_forward_warps = 0; + int args_num_fw_warps_per_channel = 1; if (num_scaleout_ranks > 1) { EP_HOST_ASSERT(num_channels % num_sms == 0 and "Invalid number of channels or SMs, you may use a different SM count than dispatch"); @@ -185,7 +190,11 @@ static void* launch_combine(void* x, "Invalid combine SM count, please try to match your dispatch config"); } else { const auto num_data_warps = num_scaleup_warps + num_forward_warps; - num_warps = num_data_warps + 1; + const int num_fw_warps_per_channel = + (allow_multiple_reduction and not use_expanded_layout and + not prefer_overlap_with_compute and + (num_scaleup_warps + 2 * num_forward_warps + 1) * 32 <= 1024) ? 2 : 1; + num_warps = num_scaleup_warps + num_forward_warps * num_fw_warps_per_channel + 1; EP_HOST_ASSERT(num_warps * 32 <= 1024 and "combine warp count (scale-up + forward + proxy) exceeds the " "1024-thread block limit; use at least num_channels / 15 SMs"); @@ -196,9 +205,11 @@ static void* launch_combine(void* x, const int64_t tma_smem_bytes = static_cast(num_data_warps) * token_layout.get_num_bytes(); const int64_t proxy_ring_bytes = deep_ep::elastic::ProxyRingLayout::get_num_bytes( num_forward_warps, deep_ep::elastic::kProxyRingDepthDefault); + const int64_t pair_sync_bytes = static_cast(2 * num_forward_warps) * sizeof(int); // The channel auto-tuner should prevent this assert from firing; leaving it as a sanity check. - EP_HOST_ASSERT(tma_smem_bytes + proxy_ring_bytes <= num_smem_bytes and + EP_HOST_ASSERT(tma_smem_bytes + proxy_ring_bytes + pair_sync_bytes <= num_smem_bytes and "Combine TMA buffers + proxy rings exceed per-block shared memory"); + args_num_fw_warps_per_channel = num_fw_warps_per_channel; } } @@ -216,6 +227,7 @@ static void* launch_combine(void* x, .num_experts = num_experts, .num_topk = num_topk, .num_qps = num_qps, .num_timeout_cycles = num_timeout_cycles, + .num_fw_warps_per_channel = args_num_fw_warps_per_channel, .x = static_cast(x), .topk_weights = static_cast(topk_weights), .src_metadata = src_metadata, diff --git a/deep_ep/include/deep_ep/common/gin_resource_alloc.cuh b/deep_ep/include/deep_ep/common/gin_resource_alloc.cuh index 7b2cf7ebf..98276a67f 100644 --- a/deep_ep/include/deep_ep/common/gin_resource_alloc.cuh +++ b/deep_ep/include/deep_ep/common/gin_resource_alloc.cuh @@ -47,13 +47,13 @@ struct GinResourceConfig { static constexpr int kMinGinContextCnt = 2; static constexpr int kMaxGinContextCnt = kMaxGinContextBudget; -// Default context count (== default QP count). 11 contexts -> 21 signals/context. +// Default context count (== default QP count). 13 contexts -> 17 signals/context. // Contexts and signals-per-context are inversely coupled through // `gin_indexed_signals_for`, so more QPs means a smaller per-context signal budget. The // equivalent alternatives are {5, 6, 7, 8, 9, 14}; everything else loses a part somewhere. // Notably 12, 15, 16 and 17 all drop to 3 parts at 12 SMs -- 17 (the provider maximum) leaves // only 13 signals/context, and its per-SM QP split puts 4 channels on the busiest QP. -static constexpr int kDefaultGinContextCnt = 11; +static constexpr int kDefaultGinContextCnt = 13; // Per-context indexed-signal budget, workaround for current limitations in provider. __forceinline__ __device__ __host__ constexpr int gin_indexed_signals_for(int gin_context_cnt) { diff --git a/deep_ep/include/deep_ep/impls/hybrid_combine_unordered.cuh b/deep_ep/include/deep_ep/impls/hybrid_combine_unordered.cuh index 05755d449..039a18846 100644 --- a/deep_ep/include/deep_ep/impls/hybrid_combine_unordered.cuh +++ b/deep_ep/include/deep_ep/impls/hybrid_combine_unordered.cuh @@ -28,6 +28,7 @@ template (token_layout, kNumDataWarps, 1, smem); - // The proxy warp never touches `tma_buffer`, clamp its index to 0. + const auto tma_buffer_layout = layout::BufferLayout(token_layout, kNumTmaWarps, 1, smem); + // Map each warp to its buffer; the proxy warp never touches `tma_buffer`, clamp to 0. + const auto tma_buffer_idx = warp_idx < kNumScaleupWarps ? + warp_idx : + (warp_idx < kNumDataWarps ? + kNumScaleupWarps + (warp_idx - kNumScaleupWarps) / kNumFwWarpsPerChannel : 0); const auto tma_buffer = tma_buffer_layout - .get_rank_buffer(warp_idx < kNumDataWarps ? warp_idx : 0).get_token_buffer(0); + .get_rank_buffer(tma_buffer_idx).get_token_buffer(0); // Proxy warp hand-off rings — placed in `smem[]` right after the TMA buffers. const auto proxy_ring_layout = ProxyRingLayout( kNumForwardWarps, kProxyRingDepth, tma_buffer_layout.get_buffer_end_ptr()); + auto* const pair_free_seq = reinterpret_cast(proxy_ring_layout.get_end_ptr()); + auto* const pair_half_done_seq = pair_free_seq + kNumForwardWarps; for (int f = thread_idx; f < kNumForwardWarps; f += kNumThreads) { *proxy_ring_layout.get_head(f) = 0; *proxy_ring_layout.get_tail(f) = 0; *proxy_ring_layout.get_done(f) = 0; + pair_free_seq[f] = 0; + pair_half_done_seq[f] = 0; } __syncthreads(); @@ -111,15 +123,20 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, // Init TMA for scale-up and forward warps ptx::arrival_phase phase = 0; const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); - if (ptx::elect_one_sync()) + const bool owns_tma_buffer = warp_idx < kNumScaleupWarps or + (warp_idx < kNumDataWarps and (warp_idx - kNumScaleupWarps) % kNumFwWarpsPerChannel == 0); + if (owns_tma_buffer and ptx::elect_one_sync()) ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); __syncwarp(); // NCCL Gin handle - // For data warp, each warp is a channel, for proxy warp each lane is a channel. - const auto gin_channel_in_sm = warp_idx < kNumDataWarps ? + // For data warp, each warp is a channel (a cooperative forward pair maps to the + // same channel), for proxy warp each lane is a channel. + const auto gin_channel_in_sm = warp_idx < kNumScaleupWarps ? (warp_idx % kNumChannelsPerSM) : - (lane_idx < kNumForwardWarps ? lane_idx : 0); + (warp_idx < kNumDataWarps ? + ((warp_idx - kNumScaleupWarps) / kNumFwWarpsPerChannel) % kNumChannelsPerSM : + (lane_idx < kNumForwardWarps ? lane_idx : 0)); const auto [qp_idx, sharing_mode] = comm::get_qp_mode(sm_idx, gin_channel_in_sm); const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, sharing_mode); @@ -132,7 +149,8 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, // Adjust register count at certain cases // TODO: support more cases, or try to make channel count more aligned - const bool kAdjustRegisters = (kNumChannelsPerSM == 4 or kNumChannelsPerSM == 8) and not kUseExpandedLayout; + const bool kAdjustRegisters = (kNumChannelsPerSM == 4 or kNumChannelsPerSM == 8) and + not kUseExpandedLayout and kNumFwWarpsPerChannel == 1; constexpr int kNumRegistersForScaleupWarps = 40; constexpr int kNumRegistersForForwardWarps = 256 - kNumRegistersForScaleupWarps; @@ -397,7 +415,13 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, update_tails(true); } } else if (warp_idx < kNumDataWarps) { - const auto forward_warp_idx = warp_idx - kNumScaleupWarps; + EP_STATIC_ASSERT(kNumFwWarpsPerChannel == 1 or + (kAllowMultipleReduction and not kUseExpandedLayout), + "Cooperative forward pairs are only supported on the " + "multiple-reduction, non-expanded path"); + const auto forward_warp_idx = (warp_idx - kNumScaleupWarps) / kNumFwWarpsPerChannel; + const auto fw_role = (warp_idx - kNumScaleupWarps) % kNumFwWarpsPerChannel; + const bool is_fw_leader = fw_role == 0; const auto channel_idx = sm_idx * kNumChannelsPerSM + forward_warp_idx; // Indexed-signal id owned by this channel within its GIN context. Uses the @@ -469,6 +493,7 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, int last_src_scaleout_rank_idx = -1; int last_slot_written = -1; + int token_seq = 0, tokens_written = 0; const auto flush_last_tma_and_record_batch = [&]() { if (last_src_scaleout_rank_idx >= 0 and ptx::elect_one_sync()) { ptx::tma_store_wait(); @@ -477,6 +502,10 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, if (last_src_scaleout_rank_idx != scaleout_rank_idx) record_slot_and_maybe_flush(last_src_scaleout_rank_idx, last_slot_written); } + if constexpr (kNumFwWarpsPerChannel > 1) { + if (ptx::elect_one_sync()) + ptx::st_release_cta(pair_free_seq + forward_warp_idx, tokens_written); + } __syncwarp(); }; @@ -548,6 +577,7 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, #pragma unroll for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) stored_num_tokens_recv[j] += static_cast(stored_is_scaleup_rank_needed[j]); + token_seq += 1; if constexpr (not kAllowMultipleReduction) { // Cases where multiple reduction is disabled. We need to forward all data from scaleup peers to scaleout peers @@ -619,23 +649,38 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, } ); - // Do reduce - constexpr int kUnrollFactor = get_max_unroll_factor(); - combine_reduce( - lane_idx, topk_slot_idx, static_cast(tma_buffer.get_base_ptr()), + // Do reduce. Cooperative pairs split the hidden dim: role r reduces + // vecs [r * kVecPerRole, (r + 1) * kVecPerRole) of every token into + // its half of the shared TMA buffer. The buffer-release wait is + // role-dependent: the leader waits its own outstanding TMA store (and + // publishes `free_seq`) and role 1 spins on `free_seq` until the previous + // token has left the buffer. + EP_STATIC_ASSERT(kHiddenVec % kNumFwWarpsPerChannel == 0, "Invalid hidden split"); + constexpr int kVecPerRole = kHiddenVec / kNumFwWarpsPerChannel; + constexpr int kUnrollFactor = get_max_unroll_factor(); + const int vec_off = fw_role * kVecPerRole; + combine_reduce( + lane_idx, topk_slot_idx, static_cast(tma_buffer.get_base_ptr()) + vec_off, /* Get source base */ [=](const int& slot_idx) { - return static_cast(scaleup_buffer.get_token_buffer(slot_idx, true).get_base_ptr()); + return static_cast(scaleup_buffer.get_token_buffer(slot_idx, true).get_base_ptr()) + vec_off; }, - /* Wait buffer release */ [=]() { - flush_last_tma_and_record_batch(); + /* Wait buffer release */ [&]() { + if (kNumFwWarpsPerChannel == 1 or is_fw_leader) { + flush_last_tma_and_record_batch(); + } else { + if (ptx::elect_one_sync()) + while (ptx::ld_acquire_cta(pair_free_seq + forward_warp_idx) < token_seq - 1) {} + __syncwarp(); + } } ); - // Merge topk weights + // Merge topk weights (leader only: the weights region is small and + // belongs to the leader's control state) // NOTES: the slot indices must follow the master lane stored_src_buffer_idx = ptx::exchange( stored_src_buffer_idx, ptx::get_master_lane_idx(ptx::match(stored_src_scaleup_rank_idx))); - if (stored_src_scaleup_rank_idx >= 0) { + if (is_fw_leader and stored_src_scaleup_rank_idx >= 0) { tma_buffer.get_topk_weights_ptr()[lane_idx] = scaleup_buffer.get_token_buffer(stored_src_buffer_idx, true) .get_topk_weights_ptr()[lane_idx]; @@ -643,6 +688,21 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, ptx::tma_store_fence(); __syncwarp(); // Necessary to let the leader lane see the writes + if constexpr (kNumFwWarpsPerChannel > 1) { + if (not is_fw_leader) { + // Upper half written and fenced: publish and move on. + if (ptx::elect_one_sync()) + ptx::st_release_cta(pair_half_done_seq + forward_warp_idx, token_seq); + __syncwarp(); + continue; + } + // Leader: wait for role 1's half before storing the whole token. + if (ptx::elect_one_sync()) + while (ptx::ld_acquire_cta(pair_half_done_seq + forward_warp_idx) < token_seq) {} + __syncwarp(); + ptx::tma_store_fence(); + } + // Assign send and receive buffers // NOTES: as we only have 1 destination, we will use "send" as "recv" for local transfer const int recv_slot = slot_of_per_channel[src_scaleout_rank_idx]; @@ -663,6 +723,7 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, ptx::tma_store_commit(); } __syncwarp(); + tokens_written += 1; // Record RDMA info to issue later last_src_scaleout_rank_idx = src_scaleout_rank_idx; @@ -670,7 +731,7 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, } } - if (replay_pass == 0) { + if (replay_pass == 0 and (kNumFwWarpsPerChannel == 1 or is_fw_leader)) { flush_last_tma_and_record_batch(); last_src_scaleout_rank_idx = -1; if (ptx::elect_one_sync()) { @@ -682,6 +743,9 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, } } + if (kNumFwWarpsPerChannel > 1 and not is_fw_leader) + return; + // Issue the last TMA and record operation for last RDMA if constexpr (kAllowMultipleReduction) flush_last_tma_and_record_batch();