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 e67b36a99..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; @@ -188,6 +206,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 +234,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,8 +413,15 @@ 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; + 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 @@ -455,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(); @@ -463,11 +502,17 @@ 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(); }; // 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 +525,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; @@ -528,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 @@ -599,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]; @@ -623,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]; @@ -643,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; @@ -650,6 +731,21 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, } } + 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()) { + #pragma unroll + for (int dst = 0; dst < kNumScaleoutRanks; ++ dst) + issue_batched_rdma(dst); + } + __syncwarp(); + } + } + + 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();