Skip to content

perf: raise the GPU-to-NIC submission rate with channels, QPs, forward pairs #9

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
Xuan-1998 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

perf: raise the GPU-to-NIC submission rate with channels, QPs, forward pairs #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cdec521
perf(combine): schedule scale-out puts remote-first to mitigate node …
Xuan-1998 Aug 28, 2026
3c737dc
perf(combine): raise the GPU-to-NIC submission rate with channels, QP…
Xuan-1998 Aug 28, 2026
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
6 changes: 3 additions & 3 deletions csrc/elastic/buffer.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int>(num_channels_per_sm, 4);

@vladimiraerov vladimiraerov Aug 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You basically deprecated this flag and changed the customer facing API. I am fine with it, but now the perf base line that we compare against is not the same, saying that we reach the same performance as official DeepEP on IB is not correct as they set this flag to false.

We need to make it clear that this flag is no longer behave the same when we publish those results both internally and externally.

// 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.
Expand All @@ -945,7 +943,8 @@ class ElasticBuffer {
static_cast<int64_t>(2 * num_channels_per_sm) *
combine_token_layout.get_num_bytes<true, int64_t>() +
elastic::ProxyRingLayout::get_num_bytes(num_channels_per_sm,
elastic::kProxyRingDepthDefault) >
elastic::kProxyRingDepthDefault) +
/* cooperative-forward pair counters */ 2 * num_channels_per_sm * static_cast<int64_t>(sizeof(int)) >
num_smem_bytes)
-- num_channels_per_sm;
EP_HOST_ASSERT(num_channels_per_sm >= 1 and
Expand Down Expand Up @@ -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
Expand Down
20 changes: 16 additions & 4 deletions csrc/kernels/elastic/combine.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ class CombineRuntime final : public jit::LaunchRuntime<CombineRuntime> {
int num_topk;
int num_qps;
int64_t num_timeout_cycles;
int num_fw_warps_per_channel;

// Parameters
nv_bfloat16* x;
Expand Down Expand Up @@ -77,7 +78,7 @@ class CombineRuntime final : public jit::LaunchRuntime<CombineRuntime> {
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,
Expand All @@ -88,7 +89,9 @@ class CombineRuntime final : public jit::LaunchRuntime<CombineRuntime> {
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 <deep_ep/impls/{}.cuh>
Expand Down Expand Up @@ -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);
Expand All @@ -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");
Expand All @@ -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");
Expand All @@ -196,9 +205,11 @@ static void* launch_combine(void* x,
const int64_t tma_smem_bytes = static_cast<int64_t>(num_data_warps) * token_layout.get_num_bytes<true>();
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<int64_t>(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;
}
}

Expand All @@ -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<nv_bfloat16*>(x),
.topk_weights = static_cast<float*>(topk_weights),
.src_metadata = src_metadata,
Expand Down
4 changes: 2 additions & 2 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 @@ -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.

@vladimiraerov vladimiraerov Aug 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the biggest change, it change the sharing mode from _GPU to _CTA on 12SM configuration.
I remember it had some impact on the dispatch kernel and that 11 gin context are the max that does not impact it. was it "solved" because you also forced the "prefer_overlap_with_compute" to be true?

this need to be better highlight in the commit message.

// 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) {
Expand Down
Loading