Add unordered hybrid kernels for transports without ordered delivery (opt-in via EP_HYBRID_KERNEL) - #732
Open
Xuan-1998 wants to merge 4 commits into
Open
Add unordered hybrid kernels for transports without ordered delivery (opt-in via EP_HYBRID_KERNEL) #732Xuan-1998 wants to merge 4 commits into
Xuan-1998 wants to merge 4 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
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
2 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.