Skip to content

DeepEP v2:Add experimental direct recv_x dispatch path. #690

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
kzlxd wants to merge 1 commit into
base: main
Choose a base branch
from
Open

DeepEP v2:Add experimental direct recv_x dispatch path. #690

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
85 changes: 74 additions & 11 deletions csrc/elastic/buffer.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -649,6 +649,32 @@ class ElasticBuffer {
}
}

static int64_t get_comm_buffer_size(const int& num_max_tokens_per_rank,
const int& hidden, const int& num_sf_packs, const int& num_topk,
const int& elem_size,
const int& num_scaleout_ranks, const int& num_scaleup_ranks,
const bool& is_scaleup_nvlink,
const bool& allow_multiple_reduction) {
const auto num_dispatch_bytes = get_dispatch_buffer_size(
num_max_tokens_per_rank, hidden, num_sf_packs, num_topk, elem_size,
num_scaleout_ranks, num_scaleup_ranks,
is_scaleup_nvlink);
const auto num_combine_bytes = get_combine_buffer_size(
num_max_tokens_per_rank, hidden, num_topk,
num_scaleout_ranks, num_scaleup_ranks,
is_scaleup_nvlink, allow_multiple_reduction);
return math::align(std::max(num_dispatch_bytes, num_combine_bytes), symmetric::kNumAlignmentBytes);
}

static int64_t get_direct_recv_x_buffer_size(const int& num_max_tokens_per_rank,
const int& hidden, const int& elem_size,
const int& num_scaleout_ranks, const int& num_scaleup_ranks) {
const auto num_ranks = num_scaleup_ranks * num_scaleout_ranks;
return math::align<int64_t>(
static_cast<int64_t>(num_ranks) * num_max_tokens_per_rank * hidden * elem_size,
symmetric::kNumAlignmentBytes);
}

static int64_t calculate_buffer_size(const int64_t& nccl_comm,
const int& num_max_tokens_per_rank, const int& hidden,
int num_topk, const bool& use_fp8_dispatch,
Comment on lines +666 to 680

@ds-review-bot ds-review-bot Jul 14, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

🟡 warning: The reserved Direct Recv X window is appended to the max(dispatch, combine) size only when the env flag is on. But the recv_x view (buffer.hpp:1130) is taken at a fixed offset equal to the whole comm buffer size, while the epilogue/combine layout still reads the staging area within the comm buffer. If num_bytes is sized manually/by hint without the extra window, the host assert catches it — good — but the buffer-size hint path (get_buffer_size_hint) and the construction path must stay consistent. Worth a unit-level size-accounting check before merging.

🤖 v4

Expand All @@ -670,19 +696,21 @@ class ElasticBuffer {
// Dispatch size
const auto elem_size = use_fp8_dispatch ? sizeof(__nv_fp8_e4m3) : sizeof(nv_bfloat16);
const auto num_sf_packs = use_fp8_dispatch ? math::ceil_div(hidden, 32) : 0; // An approximation for number of SF packs
const auto num_dispatch_bytes = get_dispatch_buffer_size(
num_max_tokens_per_rank, hidden, num_sf_packs, num_topk, elem_size,
num_scaleout_ranks, num_scaleup_ranks,
is_scaleup_nvlink);

// Combine layout
const auto num_combine_bytes = get_combine_buffer_size(
num_max_tokens_per_rank, hidden, num_topk,
auto num_comm_bytes = get_comm_buffer_size(
num_max_tokens_per_rank, hidden, num_sf_packs, num_topk,
elem_size,
num_scaleout_ranks, num_scaleup_ranks,
is_scaleup_nvlink, allow_multiple_reduction);

// Return the maximum of those layouts, aligned to 2 MB
return math::align(std::max(num_dispatch_bytes, num_combine_bytes), symmetric::kNumAlignmentBytes);
if (get_env<int>("EP_V2_EXPERIMENTAL_DIRECT_RECV_X", 0) != 0) {
num_comm_bytes += get_direct_recv_x_buffer_size(
num_max_tokens_per_rank, hidden, elem_size,
num_scaleout_ranks, num_scaleup_ranks);
}

// Return the maximum of dispatch/combine layouts, aligned to 2 MB.
// Direct recv_x appends a window-backed output region.
return num_comm_bytes;
}

static symmetric::cpu_handle_t create_cpu_handle(const int64_t& num_cpu_bytes) {
Expand Down Expand Up @@ -964,6 +992,28 @@ class ElasticBuffer {
num_max_tokens_per_rank, hidden, num_sf_packs, num_topk, x.element_size(),
nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks,
nccl_context->is_scaleup_nvlink) <= num_buffer_bytes);
const bool use_direct_recv_x = get_env<int>("EP_V2_EXPERIMENTAL_DIRECT_RECV_X", 0) != 0;
if (use_direct_recv_x) {
EP_HOST_ASSERT(not cached_mode and
"EP_V2_EXPERIMENTAL_DIRECT_RECV_X currently supports non-cached dispatch only");
EP_HOST_ASSERT(not do_expand and
"EP_V2_EXPERIMENTAL_DIRECT_RECV_X currently supports non-expanded dispatch only");
EP_HOST_ASSERT(not previous_event_before_epilogue.has_value() and
"EP_V2_EXPERIMENTAL_DIRECT_RECV_X writes recv_x in dispatch and cannot honor previous_event_before_epilogue");
}
const auto direct_recv_x_buffer_offset = use_direct_recv_x ? get_comm_buffer_size(
num_max_tokens_per_rank, hidden, num_sf_packs, num_topk, x.element_size(),
nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks,
nccl_context->is_scaleup_nvlink, allow_multiple_reduction) : 0;
void* direct_recv_x_ptr = use_direct_recv_x ? math::advance_ptr(buffer, direct_recv_x_buffer_offset) : nullptr;
if (use_direct_recv_x) {
const auto direct_recv_x_bytes = get_direct_recv_x_buffer_size(
num_max_tokens_per_rank, hidden, x.element_size(),
nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks);
EP_HOST_ASSERT(direct_recv_x_buffer_offset + direct_recv_x_bytes <= num_gpu_buffer_bytes and
"EP_V2_EXPERIMENTAL_DIRECT_RECV_X requires extra GPU buffer bytes; "
Comment on lines +995 to +1014

@ds-review-bot ds-review-bot Jul 14, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

🟡 warning: Guardrails are asserted on the C++ side for cached/expanded/previous_event and the window-fit/size, but the deterministic-dispatch restriction exists only in the Python runtime/elastic.py:938 and is only skipped (not asserted) in the test. If the C++ kernel is templated with deterministic + Direct Recv X it has no guard. Recommend asserting determinism alongside the others in buffer.hpp too, or documenting why it is safe.

🤖 v4

"pass a larger num_bytes if the buffer was sized manually");
}

// Ready and clean host workspace for this round
const auto host_workspace_layout = layout::WorkspaceLayout(
Expand Down Expand Up @@ -1000,6 +1050,8 @@ class ElasticBuffer {
num_smem_bytes,
num_qps, num_gpu_timeout_cycles,
cached_mode, do_cpu_sync,
use_direct_recv_x,
direct_recv_x_ptr,
comm_stream);

// Received token counters
Expand Down Expand Up @@ -1073,7 +1125,17 @@ class ElasticBuffer {
// Allocate received tensors
// `recv_src_metadata` includes source token indices and buffer slot indices
const auto num_allocated_tokens = do_expand ? num_expanded_tokens : num_recv_tokens;
auto recv_x = torch::empty({num_allocated_tokens, hidden}, x.options());
torch::Tensor recv_x;
if (use_direct_recv_x) {
// Experimental: `recv_x` is carved from the registered NCCL symmetric window tail.
// Dispatch writes hidden states here directly and uses the epilogue only for metadata/top-k.
recv_x = torch::from_blob(
direct_recv_x_ptr,
{num_allocated_tokens, hidden},
x.options());
} else {
recv_x = torch::empty({num_allocated_tokens, hidden}, x.options());
}
auto recv_sf = std::optional<torch::Tensor>();
auto recv_topk_idx = std::optional<torch::Tensor>();
auto recv_topk_weights = std::optional<torch::Tensor>();
Expand Down Expand Up @@ -1141,6 +1203,7 @@ class ElasticBuffer {
jit::device_runtime->get_num_sms(),
jit::device_runtime->get_num_smem_bytes(),
num_channels,
use_direct_recv_x,
do_expand, cached_mode,
do_zero_padding,
comm_stream);
Expand Down
21 changes: 18 additions & 3 deletions csrc/kernels/elastic/dispatch.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ class DispatchRuntime final : public jit::LaunchRuntime<DispatchRuntime> {
bool is_scaleup_nvlink;
bool do_cpu_sync;
bool reuse_slot_indices;
bool direct_recv_x;
int num_notify_warps;
int num_dispatch_warps; // For hybrid dispatch
int num_scaleout_warps, num_forward_warps; // For direct dispatch
Expand All @@ -37,6 +38,7 @@ class DispatchRuntime final : public jit::LaunchRuntime<DispatchRuntime> {
int* num_unaligned_recv_tokens_per_expert;
int* dst_buffer_slot_idx;
int* token_metadata_at_forward;
void* direct_recv_x_ptr;
int num_tokens;
int sf_token_stride, sf_hidden_stride;
jit::NoRefPtr nccl_dev_comm;
Expand All @@ -52,10 +54,11 @@ class DispatchRuntime final : public jit::LaunchRuntime<DispatchRuntime> {
std::string header_name, func_name;
if (args.num_scaleout_ranks == 1) {
header_name = "dispatch";
func_name = fmt::format("dispatch_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>",
func_name = fmt::format("dispatch_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>",
args.is_scaleup_nvlink,
args.do_cpu_sync,
args.reuse_slot_indices,
args.direct_recv_x,
args.launch_args.grid_dim.first,
args.num_notify_warps, args.num_dispatch_warps,
args.num_scaleup_ranks,
Expand All @@ -65,9 +68,10 @@ class DispatchRuntime final : public jit::LaunchRuntime<DispatchRuntime> {
args.num_qps, args.num_timeout_cycles);
} else {
header_name = "hybrid_dispatch";
func_name = fmt::format("hybrid_dispatch_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>",
func_name = fmt::format("hybrid_dispatch_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>",
args.do_cpu_sync,
args.reuse_slot_indices,
args.direct_recv_x,
args.launch_args.grid_dim.first,
args.num_notify_warps, args.num_scaleout_warps, args.num_forward_warps,
args.num_scaleout_ranks, args.num_scaleup_ranks,
Expand Down Expand Up @@ -99,6 +103,7 @@ static void __instantiate_kernel() {{
args.psum_num_recv_tokens_per_expert,
args.num_unaligned_recv_tokens_per_expert,
args.dst_buffer_slot_idx,
args.direct_recv_x_ptr,
args.num_tokens,
args.sf_token_stride, args.sf_hidden_stride,
args.nccl_dev_comm, args.nccl_window,
Expand All @@ -116,6 +121,7 @@ static void __instantiate_kernel() {{
args.num_unaligned_recv_tokens_per_expert,
args.dst_buffer_slot_idx,
args.token_metadata_at_forward,
args.direct_recv_x_ptr,
args.num_tokens,
args.sf_token_stride, args.sf_hidden_stride,
args.nccl_dev_comm, args.nccl_window,
Expand Down Expand Up @@ -162,10 +168,13 @@ static void launch_dispatch(void* x, void* sf,
const int& num_qps, const int64_t& num_timeout_cycles,
const bool& cached_mode,
const bool& do_cpu_sync,
const bool& direct_recv_x,
void* direct_recv_x_ptr,
const at::cuda::CUDAStream& stream) {
// Cached mode does not support expert token counting
if (cached_mode)
EP_HOST_ASSERT(cumulative_local_expert_recv_stats == nullptr);
EP_HOST_ASSERT(not (cached_mode and direct_recv_x));

// Utils
const auto num_ranks = num_scaleout_ranks * num_scaleup_ranks;
Expand Down Expand Up @@ -200,6 +209,7 @@ static void launch_dispatch(void* x, void* sf,
.is_scaleup_nvlink = is_scaleup_nvlink,
.do_cpu_sync = do_cpu_sync,
.reuse_slot_indices = reuse_slot_indices,
.direct_recv_x = direct_recv_x,
.num_notify_warps = num_notify_warps,
.num_dispatch_warps = num_dispatch_warps,
.num_scaleout_warps = num_scaleout_warps, .num_forward_warps = num_forward_warps,
Expand All @@ -216,6 +226,7 @@ static void launch_dispatch(void* x, void* sf,
.num_unaligned_recv_tokens_per_expert = num_unaligned_recv_tokens_per_expert,
.dst_buffer_slot_idx = dst_buffer_slot_idx,
.token_metadata_at_forward = token_metadata_at_forward,
.direct_recv_x_ptr = direct_recv_x_ptr,
.num_tokens = num_tokens,
.sf_token_stride = sf_token_stride, .sf_hidden_stride = sf_hidden_stride,
.nccl_dev_comm = nccl_dev_comm, .nccl_window = nccl_window,
Expand All @@ -234,6 +245,7 @@ class DispatchCopyEpilogueRuntime final : public jit::LaunchRuntime<DispatchCopy
struct Args {
// Templated arguments
bool do_expand, cached_mode, do_zero_padding;
bool skip_recv_x_copy;
int num_channels;
int num_warps;
int num_scaleout_ranks, num_scaleup_ranks;
Expand Down Expand Up @@ -264,10 +276,11 @@ class DispatchCopyEpilogueRuntime final : public jit::LaunchRuntime<DispatchCopy
using namespace deep_ep::elastic;

static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&dispatch_copy_epilogue_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>);
auto ptr = reinterpret_cast<void*>(&dispatch_copy_epilogue_impl<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>);
}}
)",
args.do_expand, args.cached_mode, args.do_zero_padding,
args.skip_recv_x_copy,
args.launch_args.grid_dim.first, args.num_channels, args.num_warps,
args.num_scaleout_ranks, args.num_scaleup_ranks,
args.num_hidden_bytes, args.num_sf_packs,
Expand Down Expand Up @@ -306,6 +319,7 @@ static void launch_dispatch_copy_epilogue(void* buffer, void* workspace,
const int& num_scaleout_ranks, const int& num_scaleup_ranks,
const int& num_sms, const int& num_smem_bytes,
const int& num_channels,
const bool& skip_recv_x_copy,
const bool& do_expand, const bool& cached_mode,
const bool& do_zero_padding,
const at::cuda::CUDAStream& stream) {
Expand All @@ -317,6 +331,7 @@ static void launch_dispatch_copy_epilogue(void* buffer, void* workspace,
// Generate, build and launch
const DispatchCopyEpilogueRuntime::Args args = {
.do_expand = do_expand, .cached_mode = cached_mode, .do_zero_padding = do_zero_padding,
.skip_recv_x_copy = skip_recv_x_copy,
.num_channels = num_channels, .num_warps = num_warps,
.num_scaleout_ranks = num_scaleout_ranks, .num_scaleup_ranks = num_scaleup_ranks,
.num_hidden_bytes = num_hidden_bytes, .num_sf_packs = num_sf_packs,
Expand Down
27 changes: 26 additions & 1 deletion deep_ep/buffers/elastic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,6 +251,8 @@ def __init__(self,
group: the communication group.
num_bytes: the total buffer size in bytes (GPU + CPU, excludes workspace), if set, overrides MoE-based calculation.
Must be aligned to 2 MB (``get_elastic_buffer_alignment()``).
The V2 direct-recv-x experiment reserves an additional receive tensor window after
the communication buffer when ``EP_V2_EXPERIMENTAL_DIRECT_RECV_X=1`` is enabled before construction.
num_cpu_bytes: the number of CPU buffer bytes (e.g. for Engram storage). Must be aligned to 2 MB.
num_max_tokens_per_rank: the maximum number of tokens per rank, used for buffer size calculation.
hidden: the hidden dimension of each token.
Expand Down Expand Up @@ -920,8 +922,20 @@ def dispatch(self,
recv_topk_weights: received expert weights (`None` if `topk_weights` was not provided).
handle: the returned communication handle.
event: the event after executing the kernel (valid only if `async_with_compute_stream` is set).

Experimental:
If ``EP_V2_EXPERIMENTAL_DIRECT_RECV_X=1`` is set, non-expanded dispatch returns ``recv_x`` from
the tail of DeepEP's NCCL symmetric window and writes hidden states directly into that output window.
The epilogue only materializes metadata/top-k/SF tensors. The compact path supports BF16 and FP8 for
both single-node and hybrid dispatch in the experimental kernel path. Cached and deterministic dispatch
are not supported in direct recv_x mode.
"""
check_torch_deterministic()
use_direct_recv_x = os.environ.get('EP_V2_EXPERIMENTAL_DIRECT_RECV_X', '0') == '1'
if use_direct_recv_x and handle is not None:
raise RuntimeError('EP_V2_EXPERIMENTAL_DIRECT_RECV_X does not support cached dispatch')
if use_direct_recv_x and self.deterministic:
raise RuntimeError('EP_V2_EXPERIMENTAL_DIRECT_RECV_X does not support deterministic dispatch')

# Automatic decide SM and QP count
num_topk = (handle.topk_idx if topk_idx is None else topk_idx).shape[1]
Expand Down Expand Up @@ -1088,6 +1102,17 @@ def combine(self,
assert num_qps <= self.num_allocated_qps, f'Allocated QPs are not enough'

bias_0, bias_1 = ElasticBuffer._unpack_bias(bias)

extra_tensors = None
if x.shape[0] == 0:
x_backing = torch.empty((1, x.shape[1]), device=x.device, dtype=x.dtype)
x = x_backing[:0]
extra_tensors = (x_backing, x)
if topk_weights is not None and topk_weights.shape[0] == 0:
topk_weights_backing = torch.empty((1, topk_weights.shape[1]), device=topk_weights.device, dtype=topk_weights.dtype)
topk_weights = topk_weights_backing[:0]
extra_tensors = (x_backing, x, topk_weights_backing, topk_weights)

combined_x, combined_topk_weights, event = \
self.runtime.combine(x, topk_weights,
bias_0, bias_1,
Expand All @@ -1104,4 +1129,4 @@ def combine(self,
async_with_compute_stream,
allocate_on_comm_stream,
handle.do_expand)
return combined_x, combined_topk_weights, EventOverlap(event)
return combined_x, combined_topk_weights, EventOverlap(event, extra_tensors)
Loading