Add unordered hybrid kernels for transports without ordered delivery (opt-in via EP_HYBRID_KERNEL) - #732
Open
Xuan-1998 wants to merge 7 commits into
Open
Add unordered hybrid kernels for transports without ordered delivery (opt-in via EP_HYBRID_KERNEL) #732Xuan-1998 wants to merge 7 commits into
Xuan-1998 wants to merge 7 commits into
Conversation
Collaborator
🤖 ds-review-bot Code Reviewv6该补丁会破坏非 hybrid/direct 模式的默认 QP 与信号配置，并且部分合法 unordered 配置会因共享内存预算不足而无法启动 combine kernel。 v5v4p本 MR 为无有序投递/无 VA 与强信号支持的传输（如 AWS EFA SRD）新增了 EP_HYBRID_KERNEL=unordered 的混合 dispatch/combine 内核对，通过分批 put、带内 header 和计数信号替代尾信号同步；默认仍为 ordered，使用原有内核。整体实现完整度较高，但 GIN 资源配置的分支重构影响了默认 ordered + allow_hybrid_mode=false 的直接模式路径，存在会导致运行失败或信号越界的回归；此外 unordered combine 在特定 channel/SM 组合下会超过线程上限，需修正。 Files reviewed: 19 |
Infrastructure the unordered dispatch/combine kernels build on: a unified GIN context/signal layout with per-peer barrier signal indexing (gin_resource_alloc, qp_mapping), the NCCL device-comm setup reconciled with upstream's runtime-version probe, and the elastic buffer / Python plumbing that resolves the QP budget at construction. One GIN context supplies one QP; the context count (default 11, explicit via num_allocated_qps within [2, 17]) sets the per-context indexed-signal budget that bounds the per-channel part count and the ScaleOut rank count. Also the ptx/layout/comm helpers these paths need. Co-authored-by: Vladimir Aerov <vaerov@amazon.com> Signed-off-by: Xuan Jiang <xuanj@amazon.com>
Split each channel's token stream into parts and sub-parts, each sent as one batched put that carries an in-band header (iteration, token count, continuation flag) and completes a per-part counting signal. The receiver treats the signal purely as a completion count and validates each batch through its header, so correctness holds no matter what order the puts land in. Front-load a smaller first part to cut time-to-first-forward, share one signal across a part's sub-parts to stay inside the per-channel signal budget, and record a per-(token, k) recv-slot map that combine later returns partials through. Co-authored-by: Vladimir Aerov <vaerov@amazon.com> Signed-off-by: Xuan Jiang <xuanj@amazon.com>
Pack the scale-out return partials contiguously per channel and send them as batched puts, each fused with a signal add on a shared per-channel accumulator; the receiver gates on the accumulated count instead of on arrival order. A dedicated proxy warp takes put issuance off the data warps' critical path through a shared-memory hand-off ring. The reduce epilogue locates partials through the per-(token, k) recv map recorded at dispatch rather than assuming in-place token slots. Carries upstream's expanded-send weight handling (kDoExpandedSend) through the slot walk. Co-authored-by: Vladimir Aerov <vaerov@amazon.com> Signed-off-by: Xuan Jiang <xuanj@amazon.com>
The hybrid (scale-out) path now carries two kernel pairs, selected at JIT-generation time via EP_HYBRID_KERNEL. The ordered pair (default) is the existing implementation, kept untouched under its original names (hybrid_dispatch.cuh / hybrid_combine.cuh) and running under the existing communication configuration; it publishes a tail via a trailing signal and requires the backend to support strong signals and VA signals. The unordered pair (EP_HYBRID_KERNEL=unordered, hybrid_dispatch_unordered.cuh / hybrid_combine_unordered.cuh) synchronizes through in-band headers and counting signals, so strong signal is not required, and weak signal is enough; use it on transports without ordered delivery or VA/strong signal support (e.g. EFA SRD). The generated source names the variant's header and kernel, so the JIT cache keys the two apart automatically. The combine reduce epilogue is shared between the two paths via a kOrderedLayout template parameter. Signed-off-by: Xuan Jiang <xuanj@amazon.com>
Xuan-1998
force-pushed
the
unordered-kernels
branch
from
August 21, 2026 23:51
6408ad8 to
48ddfff
Compare
On a single node `scaleout_active` is false, and the else arm copied
`num_allocated_qps` verbatim instead of calling `resolve_gin_context_cnt()`.
When the caller leaves that value at its default of 0 and the unordered hybrid
kernels are in use, the Python auto-fill is deliberately skipped so that C++
resolves the count, so a single-node run kept 0 all the way into the kernels:
nccl.cu else arm -> gin_context_cnt = 0
elastic.py get_theoretical_num_qps -> min(num_sms * 16 + 1, 0) = 0
comm.cuh get_qp_mode -> kNumQPs == 1 fast path skipped,
kNumSMs <= kNumAvailableQPs false
qp_mapping.cuh -> balanced_partition(idx, 512, 0) -> n / 0
Nothing rejects q == 0. Compiled for the host that specialization exits on
SIGFPE; in a GPU run it produced corrupted output and no reported fault: the
combine reported 623 MB in 1.148 us, and its output differed from the reference
in 99.98% of elements, surfacing as 'AssertionError: Diff: nan' in
tests/elastic/test_ep.py.
Resolve the count on both branches and write it back unconditionally, since
Python reads it via get_num_allocated_qps() and caps the per-launch QP count
with it. Add a static_assert so an illegal specialization fails to compile
rather than miscomputing, and assert the delegation invariant in Python at the
point the contract is broken.
An explicit num_allocated_qps outside the feasible GIN context range used to fail the EP_HOST_ASSERT in resolve_gin_context_cnt. Clamp it instead and warn, following the channels-per-SM budget reduction warning in gin_resource_alloc.cuh: - Unordered: clamped to [kMinGinContextCnt, kMaxGinContextCnt] (currently [2, 17]) in both the Python and C++ layers, each printing a [WARN] line with the requested and effective value. The Python clamp runs first, so the C++ warning only fires for direct C++ users. - Ordered: passed through without restriction (matches upstream). - 0: auto default (C++ resolves 11 for unordered; Python resolves 65/129 for ordered hybrid, 17 for direct mode) which is unchanged. The constants are exposed to Python as _C.min_unordered_gin_qps, _C.max_unordered_gin_qps and _C.default_unordered_gin_qps.
Xuan-1998
force-pushed
the
unordered-kernels
branch
from
August 25, 2026 23:09
ed5c2ff to
00b0697
Compare
Add a section on the ordered and unordered hybrid kernel pair selected by EP_HYBRID_KERNEL: what each variant assumes about network delivery order, which GIN signal types each requests, and why unordered is the default and the only correct variant on backends that can deliver a signal before the data it trails, such as the EFA GPU-initiated one. Add a Running on AWS EFA section describing the two GIN backends (EFA-GDA and CPU proxy) and the requirements: EFA installer 1.50.0 or later, NCCL 2.31 or later for the GIN device API, GPUDirect RDMA (GDR) support on both backend paths, and the gdrcopy kernel module, with an explicit note that gdrdrv loading does not by itself confirm GDR support. Link the EFA public guide and the awsome-distributed-ai DeepEP V2 benchmark for the ready-to-run setup. Register EP_HYBRID_KERNEL in the environment variable list.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The hybrid (scale-out) dispatch/combine kernels synchronize through a trailing tail signal: the receiver assumes that when the tail arrives, all data written before it has landed. This holds on transports with ordered delivery and VA/strong signal support (e.g. IB RC), but not on unordered transports such as AWS EFA (SRD), it can be overcome with proxy thread but will hurt the performance. This PR adds an alternative kernel pair that only needs weak signals, so DeepEP's hybrid mode can run on those transports.
The default behavior is unchanged:
EP_HYBRID_KERNELdefaults toordered, which runs the existing kernels under the existing communication configuration, untouched and under their original file names. SettingEP_HYBRID_KERNEL=unorderedselects the new pair.