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
`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.
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 |
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.