perf(hybrid): don't split a channel into more parts than it has tokens for - #2
Open
whn09 wants to merge 6 commits into
Open
perf(hybrid): don't split a channel into more parts than it has tokens for#2whn09 wants to merge 6 commits into
whn09 wants to merge 6 commits into
Conversation
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. The unordered pair (default, 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. The ordered pair is the upstream implementation, kept untouched under its original names (hybrid_dispatch.cuh / hybrid_combine.cuh), and runs under the upstream communication configuration (exclusive per-channel contexts, depth-1024 rings, barrier-only signal budget, VA/strong signals); it publishes a tail via a trailing signal and requires the backend to support strong signals and VA signals. Selection happens at JIT-generation time via EP_HYBRID_KERNEL: 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>
`hybrid_dispatch_unordered.cuh` gates the sub-part geometry behind `#ifndef` (`EP_NUM_SUB_PARTS` 2, `EP_MIN_SUB_TOKENS` 1, `EP_SM100_MIN_SUB_TOKENS` 15), but nothing in the tree sets those macros, so the only way to try a different split is to edit the header and reinstall. Forward the three names as JIT `-D` flags, following the `EP_NUM_TOPK_IDX_BITS` block immediately above (and its `EP_JIT_EXTRA_FLAGS` TODO). All three are device-only -- no host translation unit reads them -- so a JIT-only define cannot desync host and device sizing. `flags` is part of `kernel_signature`, so changing the env re-JITs instead of serving a cached cubin. Unset => no behaviour change.
…s for `kNumParts` -- how many `flush_part` puts a channel's tokens leave in -- is chosen today by `compute_part_allocation()` alone, which only ever caps the count from ABOVE when the GIN indexed-signal budget is tight. It is never lowered because the geometry asks for it: `kNumParts` is never compared against `kNumMaxTokensPerChannel`, and there is no minimum-tokens-per-part threshold. The budget is loosest exactly when a channel holds the fewest tokens (low `--num-sms`, small batch), so decode shapes land on `kMaxParts` -- the worst end of the axis -- with no way to opt out. At 128 tokens / 12 SMs a channel holds 3 tokens and is described as 4 parts x 1 token: the last part is always empty, and 3 tokens leave as three separate single-token puts instead of one 3-token put. Give parts the guard sub-parts already have. Sub-parts have both a clamp of `kNumSubParts` to `kBatchSize` and `EP_SM100_MIN_SUB_TOKENS` refusing to sub-split a part too small to be worth it; parts have neither. `kMinTokensPerPart` defaults to 15 (copied from the sub-token precedent in the same file) and is overridable by `EP_MIN_TOKENS_PER_PART`. `EP_MIN_TOKENS_PER_PART=1` short-circuits to the old value, so it is an exact in-image control rather than an approximation. `kNumMaxTokensPerChannel` moves above the part count in the template list; it depends only on already-declared parameters.
whn09
changed the title
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.
Re-post of Xuan-1998/DeepEP#43 against this repo, retargeted and re-measured on
main(ec623f3). Two things changed versus the old PR:ec623f3refactor split the hybrid kernel in two. The old patch touchedhybrid_dispatch.cuh, which is now the ordered (upstream) kernel and has no sub-parts at all. This version targetshybrid_dispatch_unordered.cuh, the kernelEP_HYBRID_KERNELselects by default and the one EFA actually runs.The problem
kNumParts— how manyflush_partputs a channel's tokens leave in — is chosen today bycompute_part_allocation()alone (common/gin_resource_alloc.cuh). That function only ever caps the count from above when the GIN indexed-signal budget is tight. It is never lowered because the geometry asks for it:kNumPartsis never compared againstkNumMaxTokensPerChannel, and there is no minimum-tokens-per-part threshold.The budget is loosest exactly when a channel holds the fewest tokens (low
--num-sms, small batch), so decode shapes settle onkMaxParts— the worst end of the axis — with no way to opt out. At 128 tokens / 12 SMs a channel holds 3 tokens and is described as 4 parts × 1 token: the last part is always empty, and 3 tokens leave as three separate single-token puts instead of one 3-token put.Sub-parts already have exactly this guard — a clamp of
kNumSubPartstokBatchSize, plusEP_SM100_MIN_SUB_TOKENSrefusing to sub-split a part too small to be worth it. Parts have neither.The change
kMinTokensPerPart, default 15 (copied from the sub-token precedent in the same file), overridable byEP_MIN_TOKENS_PER_PART.kNumPartsbecomesmin(budget_parts, tokens_per_channel / kMinTokensPerPart).EP_MIN_TOKENS_PER_PART=1short-circuits to the old value rather than dividing by one — so it is an exact in-image control, not an approximation. (kNumMaxTokensPerChannel / 1would still clamp whenever a channel holds fewer tokens than the budget allows parts, which is a different geometry from the old code and would not be a valid control.)kNumMaxTokensPerChannelmoves above the part count in the template parameter list; it depends only on already-declared parameters. Note the parentheses around the comparison in the newkNumGeomPartsexpression are load-bearing — an unparenthesized>inside a template parameter list closes the list instead of comparing, the same reason(kNumNotifyWarps > 0)above it is wrapped.Measurement
tests/elastic/test_ep.py, 2 ×p5en.48xlarge(8×H200 + 16 EFA each), EP8×2 = 16 ranks,--hidden=7168 --num-topk=8 --num-experts=256 --num-sms=12 --allow-hybrid-mode=1 --prefer-overlap-with-compute=0 --test-first-only. decode =--num-tokens=128, prefill =--num-tokens=8192.Method: one image, four env-selected variants, 3 reps, variants interleaved within each rep (never all-A-then-all-B), each variant on its own
EP_JIT_CACHE_DIRso a variant can never serve another's cubin, GPU memory asserted back to idle between rounds. Mean over all 16 ranks then over reps; ± is stdev across reps, not across ranks. All 48 rounds exited 0.stock=EP_MIN_TOKENS_PER_PART=1, i.e. the exact pre-patch geometry in the same binary.decode, 128 tokens — latency
stockEP_NUM_SUB_PARTS=1Latency is the right metric at this size: 5.9 MB per rank, ~5 GB/s scale-out. This shape is message-rate bound, which is exactly what merging three single-token puts into one addresses. Combine is untouched, as expected — the change is on the dispatch scale-out path only.
prefill, 8192 tokens — unchanged
stockEP_NUM_SUB_PARTS=1Everything is within ±2%, and the one cell that looks like a regression (dispatch +1.5%) carries a ±50.8 µs across-rep stdev — a single noisy rep, not a trend. At 8192 tokens a channel holds far more than 15 tokens, so the clamp is inactive and this is the expected no-op.
Per-rank bandwidth at prefill is likewise flat: dispatch 72–75 GB/s scale-out / 233–246 GB/s scale-up in both arms (399.8 MB per rank). Note this bench's scale-out figure includes intra-node traffic unless
--ignore-local-trafficis passed, so it is not a wire-rate number; it is quoted only to show the two arms match.Caveats
sm_90over EFA only. The default of 15 is a judgement call inherited fromEP_SM100_MIN_SUB_TOKENSin the same file, not something tuned per architecture — happy to make it arch-conditional, or to default it to 1 (opt-in) if you would rather not change behaviour for shapes nobody has measured.compute_part_allocation()'s cap tightens with rank count.