perf(hybrid): don't split a channel into more parts than it has tokens for - #2
perf(hybrid): don't split a channel into more parts than it has tokens for#2whn09 wants to merge 2 commits into
Conversation
whn09
changed the title
bde11bd to
b097b03
Compare
|
Rebased onto The old branch showed a 3,000-line diff over 20 files, which was an artefact, not this change. It was cut from Measurement caveat: the numbers in the description were taken on |
b097b03 to
5a594a5
Compare
| // NOTES: the parentheses around the comparison are load-bearing -- an unparenthesized | ||
| // `>` inside a template parameter list closes the list instead of comparing (same | ||
| // reason `(kNumNotifyWarps > 0)` above is wrapped) | ||
| int kNumGeomParts = kMinTokensPerPart <= 1 ? kNumBudgetParts |
There was a problem hiding this comment.
The geometry clamp is based on kNumMaxTokensPerChannel, i.e. buffer capacity, not the actual batch. If a deployment sizes one buffer for prefill and pushes decode batches through it, the clamp stays inactive and the many-tiny-parts problem remains. Sub-parts handle this at runtime via num_sub_parts_at(part_tokens). Is it possible for us to handle it at runtime too?
There was a problem hiding this comment.
Good catch — the clamp keys off kNumMaxTokensPerChannel (buffer capacity), not the per-call token count, so a caller that sizes one buffer for prefill and reuses that instantiation for decode does stay on the many-tiny-parts side. What this PR covers is the case where the callsite instantiates a kernel per shape (the bench, and callers that size one buffer per mode); for those it lands the correct geometry at JIT time.
Doing it at runtime is more invasive than num_sub_parts_at(), because kNumParts is a compile-time constant used for stack arrays (int64_t rx_part_base[kNumParts], ~L796), workspace layout (constexpr int kNumSlotsPerChannel = kNumParts * kBatchSize, L253), EP_STATIC_ASSERT(kNumParts <= 32, "One lane drains one part") at L1081, and several unrolled for (p = 0; p < kNumParts; ++p) loops. kNumSubParts can go runtime cheaply because it is a within-part ceiling and the workspace layout does not change; a runtime effective_parts would need kNumParts to stay as the allocation ceiling and every part loop / flush / signal path to key off the runtime value and skip the unused tail.
Happy to do that as a follow-up — kept out of this PR because the change surface (loops + sync + tail-skip) is much larger and warrants its own testing. Would you rather we merge this once the other comments are settled, or hold it and land runtime + compile-time together?
| // restore the previous behaviour exactly (a value of 1 must SHORT-CIRCUIT rather than divide by | ||
| // one: `kNumMaxTokensPerChannel / 1` still clamps whenever a channel holds fewer tokens than the | ||
| // budget allows parts, which is a different geometry from the old code and not a control). | ||
| #ifndef EP_MIN_TOKENS_PER_PART |
There was a problem hiding this comment.
This should go to readme's env var list and clarify its default
There was a problem hiding this comment.
Done in the amended head — added EP_MIN_TOKENS_PER_PART (default 15, 1 disables the clamp) to the General list.
| // device-only (no | ||
| // host caller reads them), so forwarding them as JIT defines cannot desync host and device, | ||
| // and `flags` is part of `kernel_signature` below, so a change re-JITs rather than | ||
| // reusing a stale cubin. Tuning them per network/arch currently requires editing the |
There was a problem hiding this comment.
after your change we don't need to edit the header and reinstall right?
There was a problem hiding this comment.
Right — with this PR any of the four (EP_NUM_SUB_PARTS, EP_MIN_SUB_TOKENS, EP_SM100_MIN_SUB_TOKENS, EP_MIN_TOKENS_PER_PART) is overridable via env var and the #ifndef still supplies the default when unset or 0. Adjusted the comment tense ("currently requires ..." → "previously required ...") to match.
`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. (cherry picked from commit 5118c2e)
…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. (cherry picked from commit b097b03)
5a594a5 to
bfbdd15
Compare
2 participants
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.Second architecture: B300 /
sm_103over EFAIndependently reproduced on Blackwell Ultra. The clamp holds, and on this
architecture it is the whole effect —
EP_NUM_SUB_PARTS=1adds nothing here(see #1).
Setup. 2 ×
p6-b300(8×B300 SXM6 + 16 EFA gen-3 each, EFA installer 1.50.0,efa.ko3.3.0g, driver 595.91.07,compute_cap=10.3), EP8×2 = 16 ranks, NCCL2.31.2 GIN with
NCCL_GIN_TYPE=5(EFA-GDA). Two images: base8e7b42evsthis branch
5a594a5—compare/8e7b42e...5a594a5isahead 2, behind 0touching only
csrc/jit/compiler.hppanddeep_ep/include/deep_ep/impls/hybrid_dispatch_unordered.cuh, so the delta isexactly #1 + #2 with no base skew. 3 reps per cell, dtype order rotated across
reps, GPU memory asserted idle before every run, all 57 rounds exited 0.
One JIT cache directory per image, not per env variant only. The cache key is
name$$compiler$$flags$$codeandcodeis the generated wrapper — the implheaders are only
#included, so their content is not in the key, and withEP_MIN_TOKENS_PER_PARTunsetflagsis byte-identical between the two trees. Ashared cache dir therefore serves the base cubin to the patched image and the
patch measures as a no-op. (Env-selected variants within one image are safe, since
those do land in
flags— that is the property #1 relies on.)decode, 128 tokens,
--num-sms=24—tests/elastic/test_ep.pyKernel+epilogue, median over the 16 ranks then over 3 reps, µs:
8e7b42eEP_NUM_SUB_PARTS=1Same attribution as on H200: dispatch −37%, combine flat within 0.7%. The
change is on the dispatch scale-out path and nothing else moves.
End-to-end, and an in-image control
The same cells through a second, independent harness — a wall-clock
prepare + dispatch + combineregion, slowest-rank CUDA elapsed, median of 100iterations after 20 warmups, median over 3 reps. This one also carries the
in-image control,
EP_MIN_TOKENS_PER_PART=1, which #1 makes possible:8e7b42eEP_NUM_SUB_PARTS=1EP_MIN_TOKENS_PER_PART=1(clamp off)The clamp-off arm lands on the base image to within 0.3% / 0.2%, in the same
binary, so the win is the clamp and not a build or environment difference.
Run-to-run CV is ≤2% on every arm above except two cells at 8–10%.
Kernel-only and wall-clock move together, which is the check that the effect is
really in the kernel: bf16 kernel 464.9 → 359.9 µs (−105.0) against wall
504.4 → 392.9 (−111.5), with the harness overhead flat at 39.4 → 33.0 µs.
The clamp also removes the SM dependence
--num-smsUnpatched, 12 SMs costs 11% over 24. Patched, the two are identical — 12 SMs is
enough and the rest can go to compute. (On H200 the optimum moved down to 12;
on B300 the curve just flattens.)
prefill, 4096 tokens,
--num-sms=24— unchanged8e7b42eEP_NUM_SUB_PARTS=1Run-to-run CV ≤0.9%, so these are genuinely flat rather than noise-covered: at
4096 tokens a channel holds far more than 15 tokens and the clamp is inactive.
Caveats on this architecture
H200/GIN type 5, where the decode-dispatch win drops to about −8% — the
indexed-signal budget tightens with rank count, exactly as the caveat above
predicted. Nothing here should be extrapolated past 2 nodes.
on H200. The clamp-off control ran at the end of each block and still matched
the base image, which bounds any drift over a block.
sm_103; it is simply inactive at prefillshapes and active at decode shapes, which is the intent.