Skip to content

comm: lift the GIN scale-out ceiling by making the rail barrier a counting barrier #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
KeitaW wants to merge 1 commit into
base: main
Choose a base branch
from
Open

comm: lift the GIN scale-out ceiling by making the rail barrier a counting barrier #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions csrc/kernels/backend/nccl.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,9 +137,14 @@ NCCLSymmetricMemoryContext::NCCLSymmetricMemoryContext(const int64_t& nccl_comm,
gin_config.gin_indexed_signals_cnt = 0;
Comment thread
KeitaW marked this conversation as resolved.
}

EP_HOST_ASSERT(gin_config.gin_indexed_signals_cnt >= (num_rdma_ranks - 1) and
"GIN indexed-signal budget cannot give each peer rail team a dedicated "
"signal; reduce num_allocated_qps to raise the per-context signal count");
// The rail barrier is a counting barrier (`gin_barrier_wo_local_sync`, comm.cuh):
Comment thread
KeitaW marked this conversation as resolved.
// it costs `kNumReservedBarrierSignals` slots per context regardless of the team
// size. Single-domain runs take the NVLink barrier and consume none.
const int barrier_signal_slots =
scaleout_active ? elastic::gin_alloc::kNumReservedBarrierSignals : 0;
EP_HOST_ASSERT(gin_config.gin_indexed_signals_cnt >= barrier_signal_slots and
"GIN indexed-signal budget cannot host the barrier's counting signal; "
"reduce num_allocated_qps to raise the per-context signal count");

this->num_allocated_qps = gin_config.gin_context_cnt;

Expand Down Expand Up @@ -204,6 +209,11 @@ NCCLSymmetricMemoryContext::NCCLSymmetricMemoryContext(const int64_t& nccl_comm,
}
is_scaleup_nvlink = num_scaleup_ranks == num_nvl_ranks;

EP_HOST_ASSERT((is_scaleup_nvlink or num_scaleup_ranks <= 1 or num_scaleout_ranks <= 1) and
"A GIN scale-up barrier and a GIN scale-out barrier would share the "
"reserved barrier signal id; allocate a second reserved id before "
"allowing this combination");

// Create symmetric memory
// num_bytes = GPU + CPU, derive GPU portion
this->symmetric_memory = symmetric::alloc(
Expand Down
128 changes: 94 additions & 34 deletions deep_ep/include/deep_ep/common/comm.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
#include <nccl.h>
#include <nccl_device.h>

#include <deep_ep/common/gin_resource_alloc.cuh>
#include <deep_ep/common/handle.cuh>
#include <deep_ep/common/ptx.cuh>
#include <deep_ep/common/layout.cuh>
Expand Down Expand Up @@ -117,8 +118,10 @@ __device__ __forceinline__ std::pair<int, ncclGinResourceSharingMode> get_qp_mod
template <int kNumSMs, int kNumQPs, int kNumChannelsPerSM, bool kWithNotifyWarps = false>
__device__ __forceinline__ int get_qp_signal_id(
const int& sm_idx, const int& channel_in_sm_idx) {
return channel_to_signal_id<kNumSMs, kNumQPs, kNumChannelsPerSM, kWithNotifyWarps>(
sm_idx, channel_in_sm_idx);
// Shift past the barrier's reserved ids; `channel_to_signal_id` itself stays 0-based.
return elastic::gin_alloc::data_signal_id(
channel_to_signal_id<kNumSMs, kNumQPs, kNumChannelsPerSM, kWithNotifyWarps>(
sm_idx, channel_in_sm_idx));
}

// Per-part indexed-signal id: kNumParts contiguous ids under the channel's base id, one
Expand All @@ -128,8 +131,11 @@ template <int kNumSMs, int kNumQPs, int kNumChannelsPerSM, int kNumParts,
bool kWithNotifyWarps = false>
__device__ __forceinline__ int get_per_part_signal_id(
const int& sm_idx, const int& channel_in_sm_idx, const int& part_idx) {
return get_qp_signal_id<kNumSMs, kNumQPs, kNumChannelsPerSM, kWithNotifyWarps>(
sm_idx, channel_in_sm_idx) * kNumParts + part_idx;
// The reservation is added after the per-part multiply: going through
// `get_qp_signal_id` would scale the offset by `kNumParts`.
return elastic::gin_alloc::data_signal_id(
channel_to_signal_id<kNumSMs, kNumQPs, kNumChannelsPerSM, kWithNotifyWarps>(
sm_idx, channel_in_sm_idx) * kNumParts + part_idx);
}

template <int kNumRanks, int kNumSMs, int kNumThreads, int64_t kNumTimeoutCycles, int kTag = kDeviceBarrierTag>
Expand Down Expand Up @@ -208,36 +214,82 @@ __forceinline__ __device__ void gin_barrier_wo_local_sync(
ncclTeamWorld(nccl_dev_comm) : ncclTeamRail(nccl_dev_comm);
const ncclGin gin(nccl_dev_comm, 0, NCCL_GIN_RESOURCE_SHARING_CTA);

// Compact signal indexing: (kNumRanks - 1) signal slots per rank. Sender rank_idx
// writes to every peer i at the slot that identifies *itself* in the peer's
// enumeration:
// sig = (rank_idx < i) ? rank_idx : (rank_idx - 1)
// So on receiver R, each of the (kNumRanks - 1) slots gets exactly +1 from a
// distinct sender, and the wait side just iterates all slots looking for one
// increment per slot.
for (int i = thread_idx; i < kNumRanks; i += kNumThreads) {
if (i == rank_idx) continue;
const auto sig = static_cast<ncclGinSignal_t>((rank_idx < i) ? rank_idx : (rank_idx - 1));
gin.signal(team, i, ncclGin_SignalInc{sig});
}

for (int i = thread_idx; i < kNumRanks - 1; i += kNumThreads) {
const auto signal_idx = static_cast<ncclGinSignal_t>(i);
const auto shadow_ptr = gin.getSignalShadowPtr(signal_idx);
const auto target = ++(*shadow_ptr);

// TODO(NCCL): Using the official NCCL wait signal API, after they added timeout check.
timeout_while<kNumTimeoutCycles>([=](const bool& is_last_check) {
const auto signal = gin.readSignal(signal_idx, 64, cuda::memory_order_acquire);
if (signal >= target)
return true;

if (is_last_check) {
printf("DeepEP Gin barrier timeout, tag: %d, scaleout: %d, scaleup: %d, thread: %d, "
"signal: %lu, target: %lu\n", kTag, scaleout_rank_idx, scaleup_rank_idx, thread_idx, signal, target);
}
return false;
});
// The two team instantiations run different protocols. World keeps the per-peer
// barrier: its call sites pass `kFlushStores = true` and rely on per-peer arrival
// for release semantics. Rail is synchronisation-only (`kFlushStores = false` at
// every call site) and uses a counting barrier on one reserved signal id, so its
// per-context signal cost is constant in the team size.
if constexpr (std::is_same_v<team_t, ncclTeamTagWorld>) {
// Compact signal indexing: (kNumRanks - 1) signal slots per rank. Sender rank_idx
// writes to every peer i at the slot that identifies *itself* in the peer's
// enumeration:
// sig = (rank_idx < i) ? rank_idx : (rank_idx - 1)
// So on receiver R, each of the (kNumRanks - 1) slots gets exactly +1 from a
// distinct sender, and the wait side just iterates all slots looking for one
// increment per slot.
for (int i = thread_idx; i < kNumRanks; i += kNumThreads) {
if (i == rank_idx) continue;
const auto sig = static_cast<ncclGinSignal_t>((rank_idx < i) ? rank_idx : (rank_idx - 1));
gin.signal(team, i, ncclGin_SignalInc{sig});
}

for (int i = thread_idx; i < kNumRanks - 1; i += kNumThreads) {
const auto signal_idx = static_cast<ncclGinSignal_t>(i);
const auto shadow_ptr = gin.getSignalShadowPtr(signal_idx);
const auto target = ++(*shadow_ptr);

// TODO(NCCL): Using the official NCCL wait signal API, after they added timeout check.
timeout_while<kNumTimeoutCycles>([=](const bool& is_last_check) {
const auto signal = gin.readSignal(signal_idx, 64, cuda::memory_order_acquire);
if (signal >= target)
return true;

if (is_last_check) {
printf("DeepEP Gin barrier timeout, tag: %d, scaleout: %d, scaleup: %d, thread: %d, "
"signal: %lu, target: %lu\n", kTag, scaleout_rank_idx, scaleup_rank_idx, thread_idx, signal, target);
}
return false;
});
}
} else {
// Counting barrier: every sender adds 1 to the same reserved signal id on
// every peer, and the single waiter advances its shadow by the expected
// (kNumRanks - 1) increments. `SignalAdd{.., 1}` matches the accumulation
// pattern of the unordered data path (`hybrid_combine_unordered.cuh`).
constexpr auto kBarrierSignal =
static_cast<ncclGinSignal_t>(elastic::gin_alloc::kBarrierSignalId);
for (int i = thread_idx; i < kNumRanks; i += kNumThreads) {
if (i == rank_idx) continue;
gin.signal(team, i, ncclGin_SignalAdd{kBarrierSignal, 1ull});
}

// First sync: every thread has issued its sends before the single waiter
// starts polling. Second sync: no thread leaves the barrier before the wait
// completes.
__syncthreads();

if (thread_idx == 0) {
const auto shadow_ptr = gin.getSignalShadowPtr(kBarrierSignal);
const auto target = (*shadow_ptr += static_cast<uint64_t>(kNumRanks - 1));

// TODO(NCCL): Using the official NCCL wait signal API, after they added timeout check.
timeout_while<kNumTimeoutCycles>([=](const bool& is_last_check) {
const auto signal = gin.readSignal(kBarrierSignal, 64, cuda::memory_order_acquire);
if (signal >= target)
return true;

if (is_last_check) {
// A counting slot cannot identify the stalled peer; report the shortfall.
printf("DeepEP Gin barrier timeout, tag: %d, scaleout: %d, scaleup: %d, "
"signal_id: %d, observed: %lu, target: %lu, missing: %lu of %d\n",
kTag, scaleout_rank_idx, scaleup_rank_idx,
static_cast<int>(kBarrierSignal), signal, target,
target - signal, kNumRanks - 1);
}
return false;
});
}
__syncthreads();
}
}
}
Expand Down Expand Up @@ -291,6 +343,14 @@ __forceinline__ __device__ void gpu_barrier(const handle::NCCLGin& gin,
EP_STATIC_ASSERT(not kFlushStores, "No data to be flushed");
}

// World's per-peer slots and Rail's counting slot overlap in the same
// (context, signal) space, so the two GIN barriers must never be live concurrently.
// `NCCLSymmetricMemoryContext` enforces the same condition at init.
EP_STATIC_ASSERT(kIsScaleupNVLink or kNumScaleupRanks <= 1 or kNumScaleoutRanks <= 1,
"A GIN scale-up barrier and a GIN scale-out barrier would share the "
"reserved barrier signal id; allocate a second reserved id before "
"allowing this combination");

do_scaleout &= kNumScaleoutRanks > 1;
do_scaleup &= kNumScaleupRanks > 1;
if (do_scaleup and do_scaleout) {
Expand Down
84 changes: 64 additions & 20 deletions deep_ep/include/deep_ep/common/gin_resource_alloc.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,20 +64,32 @@ __forceinline__ __device__ __host__ constexpr GinResourceConfig make_gin_resourc
return GinResourceConfig{gin_indexed_signals_for(gin_context_cnt), gin_context_cnt};
}

// Each ScaleOut warp (== channel) needs its own dedicated indexed signal id, so the total
// signal budget (ctx * signals/ctx) must cover the worst-case warp count for EVERY legal
// context count, not just the default. The tightest points are ctx = 13 and ctx = 17, both at
// 221 against the 220-warp ceiling -- one signal of slack. Do not raise `kMaxSM` /
// `kMaxWarpsPerSM`, widen the context range, or lower `kTotalQPBudget` without re-checking.
__forceinline__ __host__ constexpr bool all_gin_context_counts_cover_warps() {
for (int ctx = kMinGinContextCnt; ctx <= kMaxGinContextCnt; ++ ctx)
if (ctx * gin_indexed_signals_for(ctx) < kMaxScaleoutWarps)
return false;
return true;
// Indexed-signal ids reserved for the rail counting barrier (`gin_barrier_wo_local_sync`,
// comm.cuh), taken off the bottom of every context's id space so no data channel can
// produce the barrier's id. Plain `int`, not an NCCL signal type: this header stays
// NCCL-free and host-compilable so the invariants below can be `static_assert`s.
static constexpr int kNumReservedBarrierSignals = 1;
static constexpr int kBarrierSignalId = 0;
static_assert(kBarrierSignalId >= 0 and kBarrierSignalId < kNumReservedBarrierSignals,
"the barrier's signal id must lie inside the reserved range");

// The single place the reservation offset is applied; both id derivations in `comm.cuh`
// route through it.
__forceinline__ __device__ __host__ constexpr int data_signal_id(int raw_offset) {
return kNumReservedBarrierSignals + raw_offset;
}

// Signals the data path may use: the provisioned per-context count minus the barrier
// reservation.
__forceinline__ __device__ __host__ constexpr int usable_signal_budget(const GinResourceConfig& cfg) {
return cfg.gin_indexed_signals_cnt > kNumReservedBarrierSignals
? cfg.gin_indexed_signals_cnt - kNumReservedBarrierSignals : 0;
}
static_assert(all_gin_context_counts_cover_warps(),
"GIN layout cannot give each ScaleOut warp a dedicated signal id "
"for every legal context count");

// Every legal context count must remain serviceable at the worst-case launch (`kMaxSM`
// SMs x `kMaxWarpsPerSM` ScaleOut warps): `compute_part_allocation` must return at least
// one channel per SM and at least one part. Defined below, after the math it checks.
__forceinline__ __host__ constexpr bool all_gin_context_counts_are_serviceable();

// Preferred (and workspace-sizing) maximum for per-part signalling.
static constexpr int kMaxParts = 4;
Expand Down Expand Up @@ -120,11 +132,13 @@ __forceinline__ __device__ __host__ constexpr int channels_per_context(
}

// Per-part signal allocation: pick the largest num_parts (up to kMaxParts) that fits
// channels_per_context(...) * num_parts <= gin_indexed_signals_cnt
// channels_per_context(...) * num_parts <= usable_signal_budget(cfg)
// at the requested channels/SM, then reduce channels_per_sm until the budget holds.
__forceinline__ __device__ __host__ constexpr GinPartAllocation compute_part_allocation(
// `_raw` is the diagnostics-free math so `static_assert`s can evaluate it;
// `compute_part_allocation` below adds the host-side warning and check.
__forceinline__ __device__ __host__ constexpr GinPartAllocation compute_part_allocation_raw(
const GinResourceConfig& cfg, int num_sms, int num_available_qps, int num_channels_per_sm) {
const int gin_signals = cfg.gin_indexed_signals_cnt;
const int gin_signals = usable_signal_budget(cfg);
const int channels_per_ctx = channels_per_context(num_sms, num_available_qps, num_channels_per_sm);
const int budget_parts = gin_signals / (channels_per_ctx > 1 ? channels_per_ctx : 1);
GinPartAllocation alloc{};
Expand All @@ -135,20 +149,50 @@ __forceinline__ __device__ __host__ constexpr GinPartAllocation compute_part_all
static_cast<long long>(channels_per_context(num_sms, num_available_qps,
alloc.num_channels_per_sm)) * alloc.num_parts > gin_signals)
--alloc.num_channels_per_sm;
return alloc;
}

// True when the allocation fits the usable budget (the reduction loop converged).
__forceinline__ __device__ __host__ constexpr bool part_allocation_fits(
const GinResourceConfig& cfg, int num_sms, int num_available_qps, const GinPartAllocation& alloc) {
return static_cast<long long>(channels_per_context(num_sms, num_available_qps,
alloc.num_channels_per_sm)) * alloc.num_parts
<= usable_signal_budget(cfg);
}

__forceinline__ __device__ __host__ constexpr GinPartAllocation compute_part_allocation(
const GinResourceConfig& cfg, int num_sms, int num_available_qps, int num_channels_per_sm) {
const GinPartAllocation alloc =
compute_part_allocation_raw(cfg, num_sms, num_available_qps, num_channels_per_sm);
#ifndef __CUDA_ARCH__
if (alloc.num_channels_per_sm < num_channels_per_sm)
printf("[WARN] DeepEP GIN signal budget reduced the number of channels per SM "
"from %d to %d\n", num_channels_per_sm, alloc.num_channels_per_sm);
#endif
#ifndef __CUDA_ARCH__
EP_HOST_ASSERT(static_cast<long long>(channels_per_context(num_sms, num_available_qps,
alloc.num_channels_per_sm)) * alloc.num_parts <= gin_signals and
EP_HOST_ASSERT(part_allocation_fits(cfg, num_sms, num_available_qps, alloc) and
"GIN signal budget cannot host even 1 part-signal per channel "
"at 1 channel/SM. Reduce --num-sms or num_allocated_qps.");
#endif
return alloc;
}

// The invariant declared above, now that the math it checks is in scope.
__forceinline__ __host__ constexpr bool all_gin_context_counts_are_serviceable() {
for (int ctx = kMinGinContextCnt; ctx <= kMaxGinContextCnt; ++ ctx) {
// The notify warp owns QP 0, so only `ctx - 1` contexts carry data channels.
const int avail = ctx > 1 ? ctx - 1 : 1;
const auto cfg = make_gin_resources(ctx);
const auto alloc = compute_part_allocation_raw(cfg, kMaxSM, avail, kMaxWarpsPerSM);
if (alloc.num_channels_per_sm < 1 or alloc.num_parts < 1)
return false;
if (not part_allocation_fits(cfg, kMaxSM, avail, alloc))
return false;
}
return true;
}
static_assert(all_gin_context_counts_are_serviceable(),
"some legal GIN context count cannot service the worst-case launch "
"(kMaxSM x kMaxWarpsPerSM) once the barrier reservation is taken out");

// Kernel-side entry points: derive the per-channel part count (and verify the launched
// channel count) as compile-time constants from the provisioned indexed-signal budget.
__forceinline__ __device__ __host__ constexpr int constexpr_num_parts(
Expand Down