perf(ep): InterNodeV1LL low-token latency on MI300X + CX7 - #614
Draft
jhchouuu wants to merge 38 commits into
Draft
perf(ep): InterNodeV1LL low-token latency on MI300X + CX7#614jhchouuu wants to merge 38 commits into
jhchouuu wants to merge 38 commits into
Conversation
The winner threshold was an absolute 1.0 GB/s. That constant cannot work across the operating range: at 4 tokens the whole sweep lives in 2.0-3.2 GB/s, so `bw > best + 1.0` never fires and the tuner returns candidate [1/N] verbatim for every shape; at large token counts the same 1.0 GB/s is under 2% and lets noise through. Dispatch happened to be fine because candidate 1 is genuinely near-optimal for it, but combine wants the opposite end of the space and was losing ~13% on EP16/MI300X. Four changes, together aimed at "the saved config should be the best one *and* should stop churning between re-tunes": - relative margin (_BW_REL_MARGIN = 2%) instead of absolute GB/s - score each config by the median over 9 rounds, not the mean over 5, so a single stalled round cannot decide a winner - break ties inside the margin deterministically on block_num rather than letting whichever sample landed higher win -- two configs that are statistically indistinguishable must not alternate between runs - require the same 2% margin before an existing rule is overwritten on save, so a re-tune that lands higher purely on noise leaves the checked-in JSON alone Rounds per config go 5 -> 9, so a full sweep takes ~1.8x longer. With this the tuner selects genuinely different geometries for dispatch and combine, which it structurally could not do before.
dispatch() and combine() rebuilt six torch views over persistent shmem buffers on every call. The addresses come from _dispatch_out_ptrs / _combine_out_ptrs, which __init__ already snapshots once from handle-owned symmetric memory, and the shapes are fixed by the config -- so every rebuild produced an identical view, yet each one still walked __cuda_array_interface__ and queried torch.cuda.current_device(). At small token counts that is not bookkeeping, it is latency. On EP16 (2x8 MI300X, hidden 6144, 4 tokens) host-side enqueue costs ~210us per dispatch+combine round against ~222us of wall time, i.e. the host is what paces the GPU, and cProfile attributes ~60us of that to this wrapping. The regime is overhead-bound rather than data-bound: latency is flat from hidden 1024 to 6144 and flat across 4/8/16 tokens. Same-session A/B, 5 trials each: tokens combine wall/round host enqueue/round 4 150 -> 88 us 222 -> 164 us 211 -> 146 us 8 146 -> 86 us 229 -> 162 us 16 140 -> 93 us 232 -> 184 us Cross-checked with the DeepEP-style bench (test_low_latency.py at 4 tokens): end-to-end 214 -> 157 us, while its kineto kernel-time metric is unchanged -- that metric deliberately excludes CPU launch overhead, which is why this went unnoticed. Safety notes: - The cache key includes the pointer, so a future reallocation misses the lookup and builds a correct view instead of returning a stale one. - Entries are bounded by the distinct (shape, dtype) an op sees -- one for a fixed model. They alias handle-owned memory and share the op's lifetime, so nothing is retained beyond it. - Callers get the same object rather than an equal one. Both alias the same buffer and the kernels overwrite it in place, so the only visible change is object identity. - Stable output objects are also what makes the op capturable into a HIP graph; rebuilding views per call puts host logic inside the capture. - MORI_EP_DISABLE_VIEW_CACHE=1 restores the old behaviour. It is read once in __init__ -- probing the environment per call would reintroduce the cost this removes. Verified with --cmd test at 4 and 16 tokens, 500 rounds, 0 errors on all 16 ranks, and HIP graph replay bit-exact against eager on all 16 ranks.
test_low_latency.py had num_tokens / hidden / num_topk / num_experts
hardcoded in test_loop(), so measuring any shape other than the one baked
in meant editing the file. Expose them as arguments, along with the rank
count and the pressure-test loop that was a hardcoded `False`.
Defaults reproduce exactly the shape the test has always run
(128 / 7168 / 8 / 288), so existing invocations are unaffected.
python examples/ops/dispatch_combine/test_low_latency.py \
--num-tokens 4 --hidden 6144
Generated with MORI_TUNING_SCOPE=full (75 configs) on a 2x8 MI300X pair
for the small-token regime (4/8/16 tokens, hidden 6144, fp8_e4m3_fnuz
dispatch / bf16 combine, topk 8, num_qp 1).
dispatch combine
4 tokens 32 / 4 / 21 128 / 4 / 64 (block/warp/rdma_block)
8 tokens 32 / 8 / 16 32 / 4 / 8
16 tokens 32 / 8 / 16 32 / 4 / 16
Produced after the tuning-margin fix, so the two phases carry genuinely
different geometries -- previously every shape collapsed onto the first
swept candidate.
Known limitation, worth reading before trusting the 4-token row. The
tuner measures in the same host-bound regime this series is about, and
there its metric cannot separate good geometries: benchmarking these
against 32/4/8 and against 32/8/16 + 64/4/32 at 4 tokens gives 158.9 /
158.9 / 155.6 us, i.e. all three are inside the run-to-run spread. Take
the host out of the critical path and they separate clearly, with 32/4/8
ahead of what the tuner picked here. So these are a reproducible starting
point rather than a proven optimum, and the 4-token entry in particular
should be revisited once tuning can measure without the host in the way.
They only take effect with MORI_EP_LAUNCH_CONFIG_MODE=AUTO, and a
non-editable install reads the copy under site-packages rather than this
one.
The previous commit replaced an unusable absolute 1.0 GB/s threshold with a 2% relative one. 2% is the right default for keeping a checked-in config stable, but it is the wrong default for someone tuning to find the fastest config: an improvement of 1% is discarded, and the margin also blocks re-tunes from persisting a better result. Default both margins to 0 -- take any improvement -- and expose them as MORI_EP_TUNING_MARGIN for the cases that want stability back: MORI_EP_TUNING_MARGIN=0.02 <tuning command> The two margins move together because they answer the same question at different moments: which config wins inside a sweep, and whether a re-tune may overwrite the saved rule. The tradeoff is real and worth stating. Run-to-run spread on this workload is 10-20%, so with no margin the winner is partly whichever config drew the luckiest sample, and every re-tune is free to rewrite the JSON. The median over 9 rounds introduced alongside the relative margin is the primary defence against that and is unchanged here; the margin was only ever a second, optional one. Exact ties still keep the earlier (smaller block_num) candidate, so the sweep remains deterministic when two configs measure identically.
…igurable
Customer data on a production MI308 rig showed the median-scored, no-margin
tuner (introduced two commits ago) regressing Worst and Average by 30-45%
at bs=16/32, while Best stayed essentially flat and the winning geometry
trended toward larger, more parallel configs (up to block=80 warp=16
rdma=53). That signature is not environment noise -- noise would not spare
Best while inflating Worst for one selection method only. It is what you
get when the scoring statistic cannot see a config's worst case: median
only asks "what does a typical round look like", so a geometry with a fine
median and a heavy tail (more blocks/warps contending for the same SMs and
RDMA queues) can win outright, whereas mean is pulled up hard by a single
bad round and would have penalized exactly that tail.
Revert the default to mean, matching what --cmd bench itself reports and
what main always used. Keep both knobs configurable without editing the
file:
MORI_EP_TUNING_ROUNDS=<n> # rounds per candidate, default 9
MORI_EP_TUNING_STAT=mean|median # scoring statistic, default mean
A full 75-config sweep at 4/8/16/32 tokens on an idle 2x8 MI300X pair (no
contention, different SKU from the customer's box) did not reproduce the
same magnitude -- mean's and median's picks came out within normal
run-to-run noise of each other there, sometimes in either direction. That
does not contradict the mechanism, it just means an idle, uncontended rig
does not stress the tail the way production load does; treat the fix as
warranted by the mechanism and the customer signature, not as bounded by
this sanity check.
Also regenerates the shipped MI300X EP16 configs with the new mean
default so the checked-in JSON matches what a fresh tune actually
produces now:
dispatch combine
4 tokens 32/4/21 128/4/32
8 tokens 64/4/32 128/6/64
16 tokens 64/4/32 256/4/64
Verified with --cmd test at 4 and 16 tokens, 0 errors on all 16 ranks.
… history Comments-only. The customer incident narrative belongs in the previous commit's message, not as permanent inline prose everyone re-reads. Verified the mean-branch computation is untouched and still matches main's formula exactly: kept.mean(dim=0) -> same column indexing -> same * ll_scale -> same .min().item(), for both dispatch and combine.
Selection compared a slowest-rank XGMI bandwidth figure (min-across-ranks of each rank's own mean/median bandwidth) while the "Average" column the final table -- and every re-run of --cmd bench -- actually shows is a different quantity: the grand mean of raw latency across all ranks and all rounds. A config could clear the selection bar on the first metric while looking worse than another candidate on the second, because the two were never mathematically required to agree. That gap is what let a customer ask "tuning says this config won, so why does the printed bench table show it's worse" -- a reasonable question, because the two numbers were never the same thing to begin with. Make the winner literally the config with the lowest disp_stats["lat"][2] / comb_stats["lat"][2] -- the exact scalar _build_phase_stats computes and _print_phase_table prints as "Average". Verified directly: across a full 75-candidate sweep, the winning config's printed Average (dispatch 49.19us, combine 73.54us) equals the minimum seen across every candidate in the sweep (49.2us / 73.5us at the sweep's own print precision) -- selection and the printed number are now the same value by construction, not just usually close. Consequences of tying selection to this specific quantity: - Retires MORI_EP_TUNING_STAT (mean vs median). _build_phase_stats always reports a grand mean, never a median, so a "median" selection mode could no longer mean anything coherent once selection has to match what gets printed -- keeping the knob around would have meant it silently did nothing for this decision. - _BW_REL_MARGIN is renamed _TUNING_MARGIN and now gates a lower-is-better latency comparison instead of a higher-is-better bandwidth one. MORI_EP_TUNING_MARGIN keeps its name and default (0) as a relative margin; only the quantity it's a margin *of* changed. - The saved JSON's top-level bandwidth_gbps no longer comes from the selection variable (which is now a latency, not a bandwidth); it's computed post-hoc from the winning config's own already-computed stats (_headline_bw), so the field keeps meaning what it always meant -- LL bandwidth for LL/AsyncLL kernels, RDMA bandwidth otherwise -- decoupled from whatever criterion picked the winner. Verified with --cmd test at 4 and 16 tokens, 0 errors on all 16 ranks.
…uner
The previous commit changes what the tuner selects (grand-mean latency,
matching the printed Average, instead of a slowest-rank bandwidth figure
that could diverge from it). The checked-in configs were tuned under the
old selection metric, so they no longer reflect what this branch's own
tuner would pick if run today -- regenerate them so the shipped JSON is
consistent with the code that ships alongside it.
dispatch combine
4 tokens 64/4/42 64/4/32 (block/warp/rdma_block)
8 tokens 64/4/42 32/4/21
16 tokens 64/4/42 32/4/21
Verified with --cmd test at 4 and 16 tokens through the real
MORI_EP_LAUNCH_CONFIG_MODE=AUTO load path, 0 errors on all 16 ranks.
9 was a leftover from an abandoned median-scoring experiment (which wanted extra samples to resist a single stalled round); that experiment was reverted back to mean scoring, but the rounds default was never reverted with it, silently paying a ~1.8x longer sweep for a defense mean scoring doesn't need. Main has always used repeat=5. Verified on an idle 2x8 MI300X pair: a full 75-candidate sweep at 4 tokens with the new default finished in 12s and picked the same dispatch geometry (48.2 -> 47.6us, within noise) and an equally good combine geometry (69.0 -> 69.4us, within noise) as the previous 9-round default.
… add 32-token
Full 75-candidate sweeps on an idle 2x8 MI300X pair (per-phase, per-token-count),
now running under the 5-round default instead of 9:
dispatch combine
4 tokens 32/4/21 -> 48.04us 64/4/32 -> 69.04us (unchanged)
8 tokens 64/4/42 -> 48.10us (unchanged) 64/4/32 -> 69.11us
16 tokens 64/4/42 -> 49.10us 32/4/21 -> 69.16us (unchanged)
32 tokens 128/4/64 -> 52.44us (new) 64/4/32 -> 74.80us (new)
Every changed number is within this workload's documented run-to-run noise of
its 9-round predecessor; several entries didn't change at all because the
sweep re-picked the exact same geometry. Verified with --cmd test at 4/16/32
tokens, 0 errors across all 16 ranks.
… rounds/margin guards)
Four issues from review, all confirmed live:
- _headline_bw saved the grand-mean bandwidth (stats["ll"/"rdma"][2], the
same number this file prints as "Average") as the JSON's bandwidth_gbps.
save_tuning_result gates overwriting an existing rule on new_bw > old_bw,
and every rule on disk -- including ones this PR never touches, e.g. other
GPU models/kernel types -- was written under main's original definition,
slowest-rank mean bandwidth. Grand mean is systematically >= slowest-rank
mean, so the first re-tune with this code would silently overwrite a
still-good rule with zero real improvement. Restored the slowest-rank
semantics via two new _build_phase_stats fields (rdma_worst_rank,
ll_worst_rank) so bandwidth_gbps means what it always meant; avg_latency_us
and the avg_*_bandwidth_gbps fields (already grand-mean, and what _beats
actually selects on) are untouched.
- _beats' "ties break on smaller block_num" branch was dead: the sweep visits
block_num in ascending order, so a later candidate's block_num is never
smaller than the incumbent's once one has won. It looked like a deliberate
tie-break but was actually just first-found-wins via iteration order, which
silently stops being deterministic if the sweep is ever reordered/
parallelized. Switched to comparing the full (block_num, warp_per_block,
rdma_block_num) tuple, which is a real tie-break independent of visitation
order.
- MORI_EP_TUNING_ROUNDS had no lower bound. tuning_dispatch_combine drops
round 0 as in-loop warmup (kept = all_data[1:]), so ROUNDS=1 leaves an
empty tensor and _compute_stats' .min()/.max()/.mean() crash with a shape
error that gives no hint the real cause is too few rounds. Now rejected
at import time with a clear message.
- MORI_EP_TUNING_MARGIN="" (set but empty, not unset) reached float("") and
crashed at import in both this file and tuning_config.py, since
os.environ.get(k, default) only substitutes default when the key is
absent. Both now fall back through `or "0.0"` so empty is treated the same
as unset.
Regenerated the shipped MI300X EP16 JSON (dispatch+combine, 4/8/16/32
tokens) since bandwidth_gbps values shift under the corrected semantics;
avg_latency_us per entry is within noise of the prior commit's numbers.
Verified with --cmd test at all four token counts, 0 errors across all 16
ranks.
…gbps semantics Companion to the previous commit's _headline_bw fix -- bandwidth_gbps now reports slowest-rank mean bandwidth instead of the grand mean, so every entry's number drops even though nothing about the winning geometry or its actual latency changed. Regenerated fresh (not via a re-tune diff, which would have been blocked from correcting the stale entries by the very bug being fixed) on an idle 2x8 MI300X pair; avg_latency_us matches the prior commit within noise. Verified with --cmd test at 4/8/16/32 tokens, 0 errors across all 16 ranks.
Debug aid, no effect on selection or saved results: when set, prints every candidate's full PrettyTable (the same _print_phase_table bench uses) instead of just the one-line "disp sel=... comb sel=..." summary. Lets a candidate's raw Best/Worst/Average be diffed directly against a --cmd bench run of that same block/warp/rdma, which is what's needed to pin down why tuning's self-reported latency for a config doesn't match bench's for the same config.
The per-candidate progress header printed "block_num=X, warp=Y, rdma_block_num=Z" while every table title (both the final Tuning Result and the new MORI_EP_TUNING_VERBOSE per-candidate tables) prints "block=X warp=Y rdma=Z" -- same fields, different names/punctuation, made grepping a specific config across a sweep's output inconsistent. No behavior change.
The sweep's block_num candidates started at 32 (32, 64, 128, ... up to
sm_count) with no lower-range coverage. Measured directly on MI300X EP16
v1_ll:
- --cmd test at block=2/4/6/8/16 (tok=4): all correct, 0 errors -- the
kernel has no hard minimum block count, this was purely an unexplored
region of the sweep.
- --cmd bench across tok=4/8/16/32: block=2/4/6/8 lose everywhere (e.g.
93.94us dispatch at block=2 vs ~48us for the eventual winner, at 4
tokens). block=16 ties the winner at 4 tokens but is clearly worse from
8 tokens up (56.29 vs 49.01us dispatch at 8 tokens).
Added only 8 and 16 (not 2/4/6), generically via sm_count rather than
hardcoded per architecture, so a re-tune on different hardware (e.g.
MI308's 80 CUs) empirically finds its own answer instead of assuming this
MI300X result transfers. A fresh full sweep with the wider candidate list
already found a new best for 4-token dispatch: block=16/warp=4/rdma=10 at
47.64us average, edging out the previous 64/4/42 at 48.46us. Verified: 0
crashes across a 105-candidate sweep (up from 75).
Re-tuned with the widened block_num candidate list (8/16 added). Adopted
only where an independent --cmd bench re-check confirmed the new pick
holds up:
- 4 tokens: block 64->16 (dispatch), 128->16 (combine). Latency is a wash
within noise (dispatch 47.99 vs 48.14us, combine 71.02 vs 70.93us on
direct re-check) but block=16 uses a quarter to an eighth of the CUs
for the same performance -- strictly better use of resources on a
workload this small, with headroom to spare if the GPU is shared with
other concurrent kernels.
- 32 tokens: rdma_block_num 85->64 (dispatch), same geometry with
refreshed numbers (combine). Not a real change, just a closer number
from a repeat sweep.
8 and 16 tokens are deliberately left untouched: the widened sweep's own
picks for both (dispatch 16->51.69us at 16 tokens vs the existing 49.27us;
combine 32->76.75us at 16 tokens vs the existing 70.05us) turned out to be
a single noisy sample winning under 0-margin selection, not a real
improvement -- confirmed by re-benching both the old and new picks
independently and finding the existing config still faster. Re-tuning the
same hardware isn't guaranteed to reproduce or improve on a previous
result when the search space grows; each candidate here was verified
against the currently shipped config before being adopted, not taken from
the sweep's own printout.
Verified with --cmd test at 4/8/16/32 tokens, 0 errors across all 16 ranks.
…sion Two bugs in bench's per-round debug print (--cmd bench, before the Best/Worst/Average table): - The print loop was phase-outer, round-inner: all dispatch rounds printed first (each correctly under its own "Round i"), then all combine rounds printed with no "Round i" header at all, since the header was gated on `cols is _labels[0][1]` (only true for dispatch). Combine's round 0 data ended up visually attached to dispatch's last round in the output, with no way to tell which combine line belonged to which round. Restructured to round-outer, phase-inner so one "Round i" header covers both phases' lines for that round, matching what the output already looked like it meant. - The raw per-rank list used `.int().tolist()`, truncating every value to an integer (bandwidth/duration in the 1-10 range loses almost all its precision this way) even though the "avg" on the same line already kept 2 decimals. Switched to round(v, 2) so the raw list and its average are consistent.
--cmd bench (repeat=10, hardcoded) and --cmd tuning (repeat=_TUNING_ROUNDS, MORI_EP_TUNING_ROUNDS, default 5) measured on different round counts, which was one of the two confirmed structural differences behind tuning's self-reported latency for a config not matching an independent bench run of that same config (the other being GPU DVFS/clock-ramp state, which is an environment property, not something this change addresses). Renamed to MORI_EP_ROUNDS, shared by both call sites, default 10 (bench's old value). Tuning's default goes from 5 back to 10, doubling sweep time again, but now the two are measured on equal footing by construction instead of by coincidence. Verified MORI_EP_ROUNDS=1 still raises the expected "must be >= 2" error (single-process check, no GPU needed -- the validation fires at import time before any distributed setup).
Mlx5QpContainer::CreateQueuePair allocates a UAR with
devx_alloc_uar(MLX5DV_UAR_ALLOC_TYPE_NC) and then hipHostRegister()s it for
every queue pair. rdma-core does not hand out a fresh UAR for that flag:
// providers/mlx5/verbs.c, _mlx5dv_devx_alloc_uar()
if (flags & MLX5_IB_UAPI_UAR_ALLOC_TYPE_NC)
return mlx5_get_singleton_nc_uar(context);
It returns the ibv_context's singleton, i.e. the identical reg_addr on every
call (rdma-core e50a2af03f, v33+, so every modern distro). The second QP's
registration therefore fails with hipErrorHostMemoryAlreadyRegistered and
HIP_RUNTIME_CHECK exits the process.
The teardown path already assumed the sharing -- "Multiple qp may share the
same uar address, only unregister once" -- and guarded itself by probing the
pointer. Only the register side was missing a guard, and bnxt has had one all
along (BnxtDeviceContext::TryRegisterUar). Port it: Mlx5DeviceContext owns the
registered-UAR set, CreateQueuePair registers only on first use, and teardown
unregisters through the same bookkeeping instead of hipPointerGetAttributes --
that probe returns an error once the first QP has unregistered, which would
have exited(-1) during destruction.
Two ordering details the bookkeeping forces, both of which bnxt also gets right:
the set and its mutex are declared BEFORE the pools, because members die in
reverse declaration order and ~Mlx5QpContainer calls back into
TryUnregisterUar; and ~Mlx5DeviceContext clears the pools explicitly rather
than leaving it to implicit member destruction. Getting either wrong makes the
unregister lock a destroyed mutex and mutate a destroyed std::set, which shows
up as glibc heap corruption at shmem_finalize -- once per rank, long after the
last measurement, with no other symptom.
Found bringing EP16 up on 2x8 MI300X with ConnectX-7: every rank died at QP 2
of 32 with
[.../providers/mlx5/mlx5.cpp:294] hip failed with part or all of the
requested memory range is already mapped
With this, EP16 dispatch/combine runs clean on that cluster -- no abort, no
malloc_consolidate diagnostics, at 4 and 64 tokens per rank.
The block ladder starts at 32 and doubles, which never asks whether a few-token shape wants fewer. At 4 tokens the LL send path has exactly one warp with work to do, so 32 blocks is a guess, not a floor. MORI_TUNING_BLOCKS / MORI_TUNING_WARPS make the sweep reachable from a script without editing it.
generate_profiler_bindings.py groups slots by source-file stem, so the three Dispatch* slots used by intranode_ll.hpp land in IntranodeLlSlot -- but the file asks for INTRANODE_PROFILER_INIT_CONTEXT, which binds IntranodeSlot. Any ENABLE_PROFILER=1 build therefore fails to compile: intranode_ll.hpp:267:30: error: no member named DispatchSendTokens in mori::profiler::dispatch_combine::IntranodeSlot intranode_1250x.hpp has the same mismatch. It happens to compile because its slot names coincide with intranode.hpp, but the ids it emits then come from the wrong enum and the trace decodes to the wrong names.
DispatchIntraNode partitioned work by block over tokens:
tokenPerBlock = ceil(curRankNumToken / xgmiBlockNum)
so whenever curRankNumToken < xgmiBlockNum most blocks got an empty range and
did nothing. At 4 tokens per rank with 16 XGMI blocks, 12 blocks idled and 16
warps carried all 32 token x expert copies. Measured with the kernel profiler
on 2x8 MI355X (EP16, hidden 7168 bf16, topk 8, 32/16/4): warps 64-79 spent
15-32 us in dispatch_intra while warps 80+ measured 0.6 us.
Stride the same work items flatly across every XGMI warp in the grid instead.
The set of (tokenId, expertOffset, inTokenExpertId) triples is unchanged --
the old block-local index i maps to startTokenIdx * K + i, which is exactly
the flat index -- so routing, dedup and replay are untouched. dispatch_intra
drops to 14 us.
Also unroll the intra-node hidden-payload peer write by 4, matching what the
inter-node recv path at :457/:567 already does.
This does NOT reduce end-to-end 4-token dispatch latency: at that size the
intra-node path finishes at ~19 us, which is when the cross-node data arrives,
so it is not on the critical path. Over 4 runs each at 32/16/4 the change is
inside the noise (orig 49.59 us sd 1.2, new 49.65 us sd 2.9). It is kept as a
starvation fix, not a latency win. No regression at 64/1024/4096 tokens.
DispatchInterNodeLLSend's routing scan read one topk index per loop iteration and fed it straight into two dependent integer divides and a conditional store. numExpertPerToken is a runtime value, so the compiler cannot unroll that loop, and this is the one warp on the critical path to the doorbell -- at 4 tokens per rank block 0 warp 1 is the only sender, and there is no other resident warp to hide the load latency behind. Issue all MAX_EXPERTS_PER_TOKEN loads into registers first, then consume them, and collapse the two divides into one against a loop-invariant numExpertPerRank * gpuPerNode. (x/a/b == x/(a*b) for non-negative x, and a negative sentinel expert truncates to 0 under both forms, so routing is unchanged.) Measured with the kernel profiler on 2x8 MI355X, EP16, 4 tokens/rank, hidden 7168 bf16, 32/16/4: the sending warp's scan span drops 3.10 us -> 2.28 us, moving the doorbell earlier by the same amount. This is a span-level result, not an end-to-end one. The dispatch ends in an 8-GPU rendezvous that waits for every local rank, and it absorbs roughly 85% of any purely local saving -- so ~0.8 us here is worth ~0.1 us of the ~48 us dispatch, which is far below the +-3 us run-to-run spread at this size. Taken because it is strictly less work on the critical path with no tradeoff, not because the end-to-end number moved. Also measured and rejected in the same pass: raising the receive-side hidden WarpCopy from Unroll=4 to 8. It is slower (LLRecvCopy 8.34 -> 10.4 us) because 14336 B splits as 1x8192 + 6x1024 = 7 steps at Unroll=8 against 3x4096 + 2x1024 = 5 steps at Unroll=4, on top of 32 more VGPRs held live.
CombineSync is launched with one block per CU (304 on MI300X) but partitions by token: tokenPerBlock = ceil(totalRecvTokenNum / blockNum). At 4 tokens/rank EP16 a rank receives ~28 tokens, so tokenPerBlock is 1, only warp 0 of the first 28 blocks enters the loop, and 28 warps out of 2432 carry every 14 KB copy. Partition by (token, hidden slice) across the grid instead, which is what DispatchIntraNode already had to be changed to do.
DispatchInterNodeLLSend distributes destination nodes as `for (i = warpId; i < nNodes; i += warpNum)`, so with warpNum >= nNodes only the first nNodes warps of a block ever enter it -- but every warp then had to arrive at interNodeBlocksBarrier[1] before the last one could send the token count to the proxy. At 2 nodes and 8 warps/block that is 6 useless arrivals in 8 on a contended atomic, and it is the most likely reason the geometry sweep keeps preferring warp_per_block=2 for dispatch at every token count. Count min(warpNum, nNodes) warps per block instead and let the rest go straight on to the resolve pass.
At 4 tokens/rank the relay has 32 (token, expert) items and the tuned geometry gives it 84 RDMA warps, so one warp per item leaves most of them idle while the busy ones each push a full hidden vector across the XGMI peer aperture -- p90 8.2us at hidden 6144, and it lands on the critical path of whichever rank arrives last. Give each block a contiguous run of items instead of grid-striding them, resolve one per warp into shared memory, __syncthreads(), then let every warp of the block take a slice of every item. The sync is a block barrier, not a grid barrier: an earlier version of this idea synchronised the whole grid and cost more than the copy it parallelised. Only taken when a block owns at most kMaxBlockRelayItems items, so large batches keep the grid-strided path that pipelines chunk arrivals.
…end barrier Block b of the LL send path owns tokens [blockChunkNum*b*warpSize, ...), so every block past ceil(curRankNumToken / (blockChunkNum*warpSize)) has an empty range and posts nothing. At one chunk per block -- every shape up to warpSize tokens per rank, which is the whole low-latency regime -- that is exactly one block, yet all rdmaBlockNum of them were counted. With the tuned geometry that is ~84 arrivals on one contended address sitting directly in front of the token-count signal the receiving node is waiting for.
warpsPerItem was ceil(globalWarpNum / numItems) with no floor on the slice, so at low token counts the split is decided by how wide the grid happens to be rather than by the work. EpCombineAll runs 608 warps over 4 tokens and handed each warp 40 elements of a 6144-wide vector; EpDispatchCopyToStaging handed out 24 bytes. The size to stop at was measured on the LL combine relay, which splits the same way: 8 warps per token beat 4 and beat 16, so 512 elements is a floor with margin on the side that was shown to help.
get_registered_combine_input_buffer() hands the caller combineInp so it can write expert outputs straight in, and IntraNode combine already honours useExternalInpBuffer that way. InterNodeV1/V1LL did not: CombineSync staged inpTokenBuf over the top unconditionally, which both wasted the copy and overwrote whatever the caller had put there -- so the accessor was unusable on the internode path. Gate the staging loop on the flag, as XG-zheng does on zxg/low_latency_ep, and give the bench MORI_ZERO_COPY_COMBINE=1 to exercise it.
The branch tunes fp8 dispatch and bf16 combine at hidden 6144, which is the serving shape. These fill in what is missing next to it: bf16 dispatch (6144 at 4 tokens, 7168 at 4-64) and bf16 combine at 7168. Where both have a rule -- bf16 combine at 6144 -- the branch keeps its own. Swept and then A/B-verified against the incumbent, which matters here: a straight sweep pick lost to the shipped default twice on this box, once by 2.3us at 64 tokens and once by 3.9us at 4.
The rule on this branch is 16/10/4. Measured against 64/42/2 on p19 (2x8 MI300X + CX7, EP16, interleaved A/B): combine 71.1 -> 62.2us, 5/5 pairs with no overlap, dispatch unchanged. Not a criticism of the sweep that produced 16/10/4 -- the optimum moved underneath it. Widening the relay split (4 -> 8 warps per token) and putting a floor under MultiWarpIter both change how much work a warp is worth, and both land in this same combine path.
… clamp down TuningConfigManager picks the tightest ceiling match and clamps to the largest recorded rule when num_tokens runs past it. Adding a lone bf16/6144 4-token dispatch rule therefore captured every larger count as well: measured at 64 tokens on p19, dispatch went 180 -> 202us against the branch with no bf16 rule at all, because 64 tokens was being run at the 4-token geometry. Fill 8/16/32/64 with the geometries already A/B-validated at hidden 7168, which is the nearest measured shape. A lone low-token rule is a trap in this table format -- it is not additive, it shadows everything above it.
fp8 dispatch + bf16 combine at hidden 6144 is the shape this branch targets, so this is the rule that matters most. Measured on p19 (2x8 MI300X + CX7, EP16, interleaved A/B): dispatch 44.7 -> 41.9us, 4/4 pairs with no overlap. Same geometry the bf16 rule at this shape uses. The kernel changes in this PR move the optimum for both dtypes the same way, since the relay split and the MultiWarpIter floor are dtype-agnostic.
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.
Draft, stacked on
perf/ep-small-token-host-overhead(#604) — please merge that first.Kernel-side latency work on
InterNodeV1LLat low token counts, measured on 2x8 MI300X + ConnectX-7 InfiniBand (p19), EP16, hidden 6144 bf16.Result
Against this branch's head (plus the mlx5 fix below, without which nothing runs on CX7), interleaved same-session A/B, medians:
perf/ep-small-token-host-overheadAt 64 tokens/rank: dispatch 181.7 -> 164.0 (-9.7%), combine 201.1 -> 196.7 (-2.2%).
Splitting kernel work from tuning work with a three-way round robin (same tables in the middle row, so the last step is the kernel changes alone):
The two columns are not strictly additive: the tables were swept on the changed kernels, so A1 runs stock kernels at a geometry that is not their own optimum. The defensible statement is that with the shipped geometry held fixed, the kernel work is worth -15.6 / -6.8 us.
What is in it
Two of the wins are the same bug in different clothes — a work split sized against something that is not the real token count:
CombineInterNodeLLTypedrelay split, 4 -> 8 warps per token (-5.5 us combine). The dynamic formula commented out above the constant divides bynodeCount, but the sender signalsnumChunks * warpSize + 1, sonodeCountis 64 whenever a rank sends one chunk — i.e. every shape in the low-latency regime, no matter that only ~8 tokens are real. The formula was measuring the padding, which is why it lost to a constant.MultiWarpIter's slice (-2.9 us at 4 tokens, -4.2 at 64).warpsPerItem = ceil(globalWarpNum / numItems)had no lower bound, so the split was decided by how wide the grid happens to be:EpCombineAllruns 608 warps over 4 tokens and handed each warp 40 elements of a 6144-wide vector;EpDispatchCopyToStaginghanded out 24 bytes. 512 elements is a measured optimum — 768 and 1024 are both worse.Dispatch-side critical path, from profiler sub-spans:
proxyPe = node * gpuPerNode + (rank % gpuPerNode)pairs each rank with exactly one remote rank, andflagSlotIdcomes from a per-node counter on the sender.AMO_ADDmakes the NIC issue a read-modify-write on the responder;AMO_SETis an 8-byte inline write ordered behind the payload for free.__syncthreads(), not a grid barrier), flat XGMI stride inDispatchIntraNode, batched topk loads, and narrowing the send barrier to the blocks that actually send.Also included:
fix(rdma/mlx5)— mori cannot run multi-QP on ConnectX-7 at all without this: rdma-core returns the ibv_context's singleton NC UAR from everydevx_alloc_uar(MLX5DV_UAR_ALLOC_TYPE_NC), so the second QP'shipHostRegisteraborts the process. Also pushed standalone asjiahzhou/mlx5-uar-single-registerif you would rather take it separately.fix(profiler)—ENABLE_PROFILER=ONdoes not compile on main:intranode_ll.hppbindsIntranodeSlotbut its slots are generated intoIntranodeLlSlot.perf(ep): honour the zero-copy combine input—get_registered_combine_input_buffer()is silently unusable on the internode path today;CombineSyncstaged over the caller's data. IntraNode already honoursuseExternalInpBuffer.Two things worth knowing about the tuning table
num_tokensruns past it, so a lone low-token rule shadows every larger count. Adding a single bf16/6144 4-token dispatch rule cost 22 us at 64 tokens before I noticed. A rule is not additive in this format.More generally, on this box a straight sweep pick lost to the incumbent three times (by 2.3 us at 64 tokens, 3.9 us at 4, and 8.9 us here). The tuner's spread is larger than the differences it is ranking at 4 tokens. Every rule here was A/B'd against the incumbent before being written down.
Validation
--cmd test, 1000 rounds x 16 ranks, 0 errors at both 4 and 64 tokens per rank, hidden 6144 and 7168.Not done
kMinDimPerWarpof 512) were tuned here.zero_copymake rule selection nondeterministic — worth a look independently of this PR.Mlx5PostWrite). Five attempts to shrink the first and three to shrink the second all measured zero or worse.