DeepEP v2:Add experimental direct recv_x dispatch path.#690
Conversation
| } else if (stored_dst_slot_idx >= 0 and direct_recv_x_idx < 0) { | ||
| // Fallback: timeout on prefix sum, write SF/metadata only via RDMA | ||
| gin.put<team_t>( | ||
| recv_buffer.get_token_buffer(stored_dst_slot_idx).get_sf_ptr(), | ||
| send_buffer_token.get_sf_ptr(), | ||
| kNumSidecarBytes, stored_dst_rank_idx, |
There was a problem hiding this comment.
🔴 critical: Prefix-timeout/slot-invalid fallback silently drops the token's hidden payload. When the receiver-published base is missing/late/on a wrong peer, the kernel prints a timeout and then writes the sidecar only (or nothing), so the token's hidden content is lost — there is no abort and no fallback to the old staging-copy path. The experimental flag's correctness rests on this never happening, which is fragile in multi-scaleup/scaleout/hybrid routing. Same pattern in hybrid_dispatch.cuh:638-660 and 343.
🤖 v4
| if (stored_dst_slot_idx >= 0 and dst_x_ptr == nullptr and direct_recv_x_idx >= 0) { | ||
| gin.put<team_t>( | ||
| math::advance_ptr(direct_recv_x, static_cast<int64_t>(direct_recv_x_idx) * kNumHiddenBytes), | ||
| send_buffer_token.get_hidden_ptr(), | ||
| kNumHiddenBytes, stored_dst_rank_idx, | ||
| ncclGinOptFlagsAggregateRequests); | ||
| gin.put<team_t>( | ||
| recv_buffer.get_token_buffer(stored_dst_slot_idx).get_sf_ptr(), |
There was a problem hiding this comment.
🟡 warning: Even in the normal RDMA scaleout path, the hidden payload is written via a separate RDMA put into recv_x while the sidecar goes into staging. This is fine only if the published base is always valid here; see the timeout/drop concern above.
🤖 v4
| // Cached expand mode: read pre-computed dst_tensor_idx from metadata | ||
| dst_tensor_idx = recv_src_metadata[i * kMetadataStride + 2 + lane_idx]; | ||
| } else if (kDoExpand and not kCachedMode and dst_expert_idx >= 0) { | ||
| } else if (kDoExpand and dst_expert_idx >= 0) { |
There was a problem hiding this comment.
🔴 critical: Regression in the shared (non-experimental) epilogue: the expanded atomic-add branch lost its not kCachedMode guard (else if (kDoExpand and dst_expert_idx >= 0)). This epilogue is launched for cached/expanded dispatch even when Direct Recv X is off, so it can corrupt cached-expand dispatch and/or its per-expert counters on the baseline path.
🤖 v4
| workspace_layout.get_scaleup_rank_count_ptr<true>() + scaleup_rank_idx, | ||
| math::encode_decode_positive(static_cast<int64_t>(prefix)), i); | ||
| prefix += rank_count[i]; | ||
| } | ||
| __syncwarp(); | ||
| } |
There was a problem hiding this comment.
🟡 warning: The receiver-side prefix-publish loop uses a per-scaleup slot that is later polled with a timeout and a separate channel-signaled-tail mechanism (cf. dispatch.cuh:272/397/562). The two variants of the same protocol differ in which slot/peer they address; the design needs a single, race-free publication/decoder contract before the experimental flag can be trusted across scaleout+scaleup topologies.
🤖 v4
| 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; " |
There was a problem hiding this comment.
🟡 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
| 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, |
There was a problem hiding this comment.
🟡 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
🤖 ds-review-bot Code Reviewv6
v4The Direct Recv X PR is already fully implemented in the current working tree (committed at HEAD 0226cd2 on top of the real DeepEP base dd758ca). The seven changed files wire up and test the entire feature: a Direct Recv X window reserved at the tail of the symmetric NCCL comm buffer; recv_x returned as a view into that window; the dispatch kernel writing hidden states directly into recv_x using a receiver-published per-source prefix (both single-node/scaleup and hybrid/scaleout variants); the copy epilogue skipping the hidden copy while still producing metadata/top-k/SF; the experimental env flag and its guardrails (non-cached, non-expanded, no previous_event_before_epilogue). Since the feature is already written, the value of a fresh review is to surface the parts that are still fragile/unfinished before the experimental flag is promoted/merged — and the one part promised in the PR summary that is missing from the diff. v3
Files reviewed: 7 |
Reserve a symmetric receive window and let dispatch write recv_x payloads directly so the epilogue can skip the hidden-state copy in supported compact dispatch cases.
0226cd2 to
6757592
Compare
kzlxd
changed the title
2 participants
Add Direct Recv X dispatch path
Summary
This PR adds an experimental Direct Recv X path for compact EP dispatch.
The optimization removes the extra local hidden-state copy in the dispatch copy epilogue. In the original path, hidden states are first written into the NCCL symmetric staging buffer, then copied again into the final
recv_xtensor by the epilogue. With Direct Recv X enabled, dispatch writes hidden states directly into the final compactrecv_xwindow, while metadata and top-k tensors continue to use the existing staging + epilogue path.Direct Recv X is enabled by:
This PR also includes a small SM90 / H800 JIT compilation fix for an SM100-only shared-memory bulk-store instruction.
Background
The original dispatch path has two buffer layers:
stagingis NCCL symmetric / window memory and can be accessed by the communication path directly.recv_xis a normal PyTorch tensor. It is materialized after dispatch bydispatch_copy_epilogue.This means hidden states always pay one extra local copy:
The copy cost grows linearly with token count and becomes more visible with wider hidden sizes.
Direct Recv X design
The dispatch kernel already knows:
The receiver notify warp also knows how many tokens each source rank contributes. From those counts, it can compute the exclusive prefix base for each source rank in the compact
recv_xlayout.Direct Recv X publishes this receiver-side prefix information back to the source rank:
The hidden-state data path becomes:
The metadata / top-k path stays unchanged:
So the epilogue no longer copies hidden states, but it still materializes the metadata and top-k outputs.
Implementation
Reserve a Direct Recv X output window at the tail of the DeepEP symmetric communication buffer when
EP_V2_EXPERIMENTAL_DIRECT_RECV_X=1.Return
recv_xas a view into this symmetric-window tail in Direct Recv X mode.Extend dispatch runtime arguments and kernels so dispatch can write hidden states directly into the receiver-side
recv_xwindow.Let
dispatch_copy_epilogueskip the hidden-state copy when Direct Recv X is enabled.recv_topk_idxrecv_topk_weightsrecv_src_metadatachannel_linked_listSupport both compact dispatch variants:
Keep the feature behind one experimental env flag:
EP_V2_EXPERIMENTAL_DIRECT_RECV_X=1Add guardrails for unsupported modes:
previous_event_before_epilogueis rejected becauserecv_xis produced during dispatch instead of during the epilogue.H100 benchmark results
Configuration:
EP_V2_EXPERIMENTAL_DIRECT_RECV_X=1Parameter：
EP8
EP16
Result summary
Direct Recv X substantially reduces the copy epilogue cost by removing the hidden-state copy from staging to
recv_x.End-to-end dispatch + copy improvement:
In some BF16 cases, the dispatch kernel itself becomes slightly slower because it now performs direct remote
recv_xwrites. The removed epilogue hidden copy still more than offsets that overhead, so total latency improves.