[Example] EP v2 intranode dispatch/combine - #68
Draft
Rachmanino wants to merge 29 commits into
Draft
Conversation
A TileScale-native MoE all-to-all following DeepEP EPv2, scoped to intranode NVLink on SM100. Dispatch payload is bf16 or fp8 (per-token, per-128-element scales); combine is bf16. Dispatch runs three phases in one kernel -- count, count exchange, scatter -- and writes rows straight to their final compact index, so there is no copy epilogue. Combine stores back into a slot unique per (contributing rank, source token) and reduces locally, one block per source token. Measured on 8x B200, 8192 tokens/rank, hidden 7168, top-8, 256 experts, 64 SMs, against DeepEP tests/elastic/test_ep.py on the same machine. Per collective, including the epilogues DeepEP reports separately: dispatch (fp8) DeepEP 442+129us this 510-518us combine (bf16) DeepEP 839+156us this 965-974us total DeepEP ~1566us this 1475-1492us Also fixes address_of handling in LegalizeSafeMemoryAccess: bounds checking rewrote its argument into an if_then_else, which every consumer rejects, so callers had to clamp indices by hand and disable LoopUnswitching. Taking an address is not a dereference; the bound belongs to whoever dereferences it, exactly as for tl.access_ptr.
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
👋 Hi! Thank you for contributing to the TileScale project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
`topk_idx == -1` is DeepEPs no selection marker and reaches dispatchs dedup and slot claiming as a destination of -1. The guard was there but nothing exercised it: every test generated indices with `randint(0, num_experts)`. Adds a masked 8-GPU case and a `--masked-ratio` knob to both example scripts, and makes the reference exclude unselected entries. A token with no selections at all now combines back to zero, which the relative-error check has to special -case since the expected tensor is then zero. Also lists the DeepEP intranode features this port does not implement, so the boundary is explicit rather than inferred.
ruff-check: drop an unused idx local in dispatch.py's scatter loop (the per-destination idx_k inside the topk loop is what's actually used). ruff-format: apply formatter line-wrapping to buffer.py, dispatch.py, example_dispatch_combine_benchmark.py, example_dispatch_combine_correctness.py, reference.py and test_example_deepep_v2.py (the latter four needed it after rebasing onto the masked-topk commit, which landed unformatted). codespell: whitelist "agrs" (DeepEP's own AllGather-Reduce-Scatter term, confirmed against the DeepEP submodule) instead of letting it rewrite to "args".
`cp_warp_impl` reinterprets both pointers as `int4` and moves `N * sizeof(dtype_t) / sizeof(int4)` of them, so any element past the last whole 16-byte unit was never copied -- and when `N` is small enough that `N_int4 == 0`, neither the unrolled loop nor the drain loop ran at all and the transfer silently moved nothing. The shape that hit it: deepep_v2's FP8 dispatch moved its per-token scale with `put_warp(size=1)` on a `float32`, i.e. 4 bytes. `recv_x_scales` came back all zero on every rank, with no error anywhere -- the payload copy beside it was large and worked fine, so the kernel looked correct and merely inaccurate. Adds the remainder loop through the original, unreinterpreted pointers (`N` and `dtype_t` are compile-time, so the tail bounds are too), and a regression test covering `n=1,2,3` (under one `int4`) and `n=5` (one full `int4` plus a dropped element) for both `put_warp` and `get_warp`.
FP8 dispatch used to move each token's scale in a second `put_warp` beside the payload. A remote store's fixed cost (peer-address translation, warp setup, the NVLink round trip) is noise against 7168 payload bytes and dominates a 224-byte scale: 566.9us with the second store, 493.1us without, while splitting a *large* transfer into two calls costs nothing (901.8 vs 894.8us). So `per_token_cast_to_fp8` now returns one packed uint8 row -- payload, then the per-group fp32 scales -- and the scatter moves it in a single store. Quantisation had to write its output somewhere regardless. The packed row is padded to 16 bytes always (`put_warp`'s bulk path reinterprets it as `int4`; `hidden=128` packs to 132 raw bytes and faults without this) and to 512 -- 32 lanes x one `int4`, which keeps `cp_warp_impl`'s drain loop warp-uniform -- only when the padding costs under 5%. That threshold is measured, not guessed: 512-alignment is 8.9% slower at hidden=2048 and 6.8% slower at 4096, where it costs +21% and +9% padding, and 2.4% faster at 7168, where it costs +3.9%. Worth noting the trade is invisible in GB/s, which credits the padding bytes as useful work -- an isolated `put_warp` roofline reads 594 GB/s at 7392 bytes against 618 at 7680 while the wall time is identical. `kernels/dispatch.py` takes the resulting width as `row_bytes` instead of recomputing it, so the rule has one owner. Also: FP8 correctness tests at both the smoke and the V3 shape (these are what the `cp_warp` tail bug was found through), a `--fp8` flag on the benchmark, and a fix to combine's reported bandwidth, which divided by dispatch's byte count -- identical for bf16, wrong for FP8. All performance figures in the README re-measured, four samples each, every one taken in a verified-idle window: bf16 dispatch 686-688 GB/s, combine 623-625, FP8 dispatch 589-591 (~632 on the wire, counting the scales and padding the payload-only figure excludes).
Two DeepEP EP features that are additive to this port's layout. `dispatch(cumulative_local_expert_recv_stats=t)` accumulates this rank's received token count per local expert into a caller-owned uint32 tensor -- DeepEP's load-balance counter, accumulate-not-overwrite so the caller picks the window by choosing when to zero it. DeepEP derives it inside notify, from a per-expert count vector it already exchanges because its expanded layout needs per-expert offsets. This port has no such exchange, and widening the count matrix from `num_ranks` to `num_experts` entries per rank to avoid a local scan is a bad trade, so the count is taken on the receiver: `recv_topk_idx` already holds the local expert for every received (row, top-k slot), so it is a scan of something this rank was handed anyway, behind the exit barrier where every peer's stores have landed. That costs ~25us of dispatch's ~896 (2.8%) at the V3 shape -- enough that it compiles a separate kernel variant instead of being always-on, so the default path is byte-identical to before. It becomes free if the expanded layout ever lands, which would bring the per-expert exchange with it. `combine(bias=b)` or `combine(bias=(b0, b1))` adds DeepEP's `bias_0`/`bias_1` to the output. They seed the reduce accumulator rather than being added after it: the accumulator was going to be written once regardless, so seeding replaces the clear and costs nothing. It also means a token whose every selection was masked off comes back as its bias instead of zero, which is the behaviour a bias is for. Three tests: per-expert counts against an all-gathered torch reference over three accumulating dispatches with a quarter of the selections masked, and bias at one and two tensors, including the fully-masked-token case. Not implemented, and now documented as deliberate rather than missing: combine's `topk_weights` side output. DeepEP carries the weights back because its expanded layout reorders them; in this port's rank layout the source rank still holds them in the original order, so it would be a collective that returns its own input.
DeepEP's `do_expand`: one received row per (token, expert) instead of one per (token, rank), rows grouped by local expert so each expert's are the contiguous block a grouped GEMM consumes. `expert_alignment` rounds each segment up and `zero_padding` clears the gap, as DeepEP's `kDoZeroPadding` does. DeepEP expands in a receiver-side copy epilogue, which can afford an atomic bump per expert because it is already reading every staged row. This port has no epilogue -- rows land at their final index straight from the sender -- so the sender has to know that index. The only thing in the way was the granularity of the count exchange: phase 1 already tallies per destination, phase 2 turns tallies into a base, phase 3 adds a locally-claimed slot, and expanding just redefines destination from rank to expert. So `n_dst` is `num_experts` rather than `num_ranks` and the same three phases produce the expert-major layout. The one new term is the segment base: an exclusive prefix sum over the destination's local experts of their aligned counts, which every rank derives from the count matrix it already holds. Capacity is the one place the two layouts genuinely differ. Deduplicated, a peer sends at most `num_max_tokens_per_rank` rows, so `num_ranks` times that cannot be exceeded. Expanded, the bound is `min(topk, experts_per_rank)` times higher -- every one of a token's experts on one rank -- which at the V3 shape is 7 GiB of receive buffer against 0.88. Balanced routing needs only the 0.88, so `expand_factor` sizes it and phase 2 checks: every rank derives the same total from the same count matrix, so an overflow is known before a byte of payload moves and the destination is skipped rather than corrupted. `handle.expand_overflow` reports how many rows were needed. The rank layout's generated code is unchanged. Everything above is trace-time, including the two places where the generalisation would otherwise have put a warp-collective instruction in a per-token loop for an answer it already had: the rank mask's dedup (`dst` *is* `dst_rank` when not expanding) and the peer shuffle in the scatter. `combine` asserts rather than accepting an expanded dispatch. Its store-back slot is `comm_x[rank][src_token]`, unique only because dispatch deduplicated; expanded, a token with two experts on one rank has two rows that collide. The fix is DeepEP's `kDoExpandedSend`, summing a token's local-expert rows before sending, which needs a (src_rank, src_token) -> rows inversion dispatch does not record yet. That is the next piece, and until it lands the expanded layout is dispatch-only. Tested against `reference.expanded_layout` at 8 ranks: per-expert counts, the aligned segment offsets, the exact set of (src_rank, src_token) in every segment, the payload of every row (each token carries its own identity, so a misrouted row cannot look plausible), and that padding is zeroed and marked unoccupied. Also splits `test_combine_bias` in two. It called `_run` twice in one process, which tears the process group down and rebuilds it on the same port; that is intermittently refused, and it failed once here for that reason.
…ditions The async support was half a feature: `dispatch` had `async_finish`, `combine` had nothing and always joined the streams, the event came back as a bare `torch.cuda.Event` on the handle, and nothing in the tree exercised any of it. Since combine is the longer of the two collectives (986us against dispatch's 896), it was the half with more to overlap. Now both take `previous_event` / `async_finish` / `allocate_on_comm_stream` and both return an `EventOverlap`, DeepEP's wrapper, with `current_stream_wait` and the `with event:` context manager. It is returned either way -- synchronously it wraps `None` -- so callers do not branch on the mode they asked for. `allocate_on_comm_stream` is the one worth explaining. The old path called `Tensor.record_stream` on this call's temporaries to hand them to the compute stream, which is exactly what DeepEP avoids: `record_stream` is incompatible with CUDA graph capture, which is why its event carries `extra_tensors` instead. Setting the flag keeps the temporaries owned by the communication stream and alive through the returned event, so the call is capture-safe. `previous_event` requires it, as in DeepEP. `pipeline_depth=0` now disables the CPU run-ahead throttle. It blocks the *host*, so a bounded run-ahead and an asynchronous call are not the same thing, and a caller driving its own overlap should be able to opt out. The default stays 2 -- it exists because a rank queued several calls behind stalls every other rank inside the cross-rank barrier. Not ported: `combine`'s `combined_topk_weights`, for the reason already documented -- in this port's rank layout it would return the caller its own input. Two tests, since none of this was covered before: a round trip with both collectives asynchronous, real work enqueued inside the `with` block and the combine chained off the dispatch's event, which fails if the events do not actually order the streams; and the contract that a synchronous call still returns an `EventOverlap` wrapping `None`. Also three simplifications to what landed earlier today, from AGENTS.md's rules on dead code and special cases: - `dst = T.alloc_var(T.int32)` in both of dispatch's per-token loops allocated a variable the next line's rebinding threw away. - Phase 3's slot claim was split into expanded and deduplicated arms to keep one `send_base` load off the rank layout's path, on an assumption never measured. One condition covers both: deduplicated, `send_base` is never negative. - The handle carried `psum_recv_count` from the *previous* call in the expanded layout, which never computes it, behind a docstring warning. It is `None` there now. All 12 tests pass.
An AGENTS.md pass over what accumulated today. Three findings, all in `buffer.py`. `dispatch(num_sms=N)` silently did nothing. It resolved the value, stored it on the handle, and `combine` read it back at the top and never used it again -- both kernels were built from `self.num_sms`. A parameter that travels through two APIs and a handle field and changes no behaviour is worse than an absent one, and DeepEP's `num_sms` is a working knob. The kernel caches are keyed on it now and it is validated against the device's SM count, which both kernels need anyway since they rendezvous grid-wide. Worth having rather than deleting: the SM sweep measured 64 -> 128 blocks at about 4%. The three expanded-layout outputs were cloned twice -- once inside the comm-stream block and again into the handle -- and deduplicated, where nothing produces them, the first three were allocated, carried in `temporaries`, and thrown away on every call. Cloned once now, and only when there is something to snapshot. And `min(x + slack, hard + slack)` where the line above had already clamped `x` to `hard`. Left alone deliberately, since they are DeepEP's handle contract rather than accidents: `EPHandle.num_experts`, `num_max_tokens_per_rank`, `topk_idx` and `psum_recv_count` are unread today, but `topk_idx` is exactly what a cached `handle=` dispatch has to reuse. 12 tests pass. Separately verified that a per-call `num_sms` now compiles its own kernel and stays correct at two different values.
Phases 1-2 compute a layout and phase 3 moves the payload; they are two kernels now. The launch boundary hands phase 3 everything phase 2 published, which is exactly what the in-kernel spin on `exchange_done` was buying, so the split costs nothing: 881.7us against 883.7 fused at the V3 shape, the extra launch absorbed by the spinner it replaces. What it buys is a constraint removed. The fused kernel ends in `T.sync_grid()`, so every block has to be resident and the grid is pinned to `num_sms` -- and that pinned the scatter too, though the scatter needs no grid-wide rendezvous of its own until the reset. ncu had already found the cost: 64 blocks on a 148-SM device, 0.22 waves/SM, 25% occupancy, issuing on 1.51% of cycles. With the scatter free to be wider, `scatter_sms=128` gives 849.1us bf16 and 487.0 fp8 against 883.7 and 507.4 fused, about 4% each. Not DeepEP's split, and worth being clear about why. DeepEP cuts *after* the movement into a copy epilogue that compacts rows staged per sender; this port's scatter already writes rows at their final index, so there is nothing there to extract and adding one would cost the ~120us that is currently the margin over DeepEP. This cut is before the movement, and the piece it isolates -- 34.7us of layout work, constant across dtypes -- is what a cached `handle=` dispatch would skip outright. Fused, that was welded into the same kernel. `scatter_sms` defaults to `num_sms`, so the default path is unchanged: a caller who capped `num_sms` to leave the device for expert compute did not thereby ask for a scatter that takes it back. 13 tests, including a new one that runs the scatter on a wider grid than the notify -- `scatter_sms` has to stride the token loop, the alignment-padding loop and the stats scan, and any of the three left on the old stride drops or double-writes rows.
…worthy Every figure re-measured on the current tree. The default path is unchanged, as it should be -- the kernel split is free and `scatter_sms` defaults to `num_sms` -- so bf16 dispatch is still 686-687 GB/s and combine 623-624. New rows for the wider scatter grid: 703-716 GB/s bf16 and 613-617 fp8, about 3.4% and 4.0%. The gating note is rewritten because the old rule was wrong twice in one day. It waited for zero foreign GPU *processes*, which is neither necessary nor sufficient: an 8-way inference server holding 167 GB/GPU at 0% SM blocks it forever while disturbing nothing, and a job under this same account is invisible to a by-other-user check while pinning a GPU at 100%. Both happened, and both produced numbers 30-60% off that looked perfectly stable. The test is now whether any process that is not ours *used the SMs* during the run, sampled throughout with `nvidia-smi pmon`. The "DeepEP's figures are GPU time, this port's are whole Python calls" caveat is replaced by a like-for-like table. Measured the same way on the same idle machine, fp8 dispatch is 439us of movement plus 113-130us of compaction for DeepEP against 35us of layout plus ~470us of movement and no compaction here: DeepEP is about 7% faster at moving the bytes, this port about 9% faster overall for not having a second pass. Host overhead is 8-10us. Also records that `async_finish` changes neither side's latency -- 901.2us synchronous against 901.4 asynchronous with the event waited on -- and that timing an asynchronous call without waiting reports 3.0us, which is the launch.
Phases 1-2 depend only on the routing, so a dispatch whose `topk_idx` has not changed recomputes a layout it already has. `dispatch(x, handle=h)` skips it: `cached` traces the notify kernel away entirely -- verified in the generated source, one `__global__` instead of two -- leaving the scatter to read `send_base`, `send_rank_mask` and `num_recv` exactly as the previous call left them. None is touched by the end-of-call reset, which clears only what phases 1-2 consume. Splitting the kernel first is what makes this worth anything. Fused, notify was welded to the scatter and skipping one meant skipping the payload too. DeepEP has the same problem from the other side: its cached dispatch measures 441-444us against 437-439 uncached, saving host work rather than GPU time, because its notify sits in a kernel that still has to run. Here a whole 35us kernel goes. The entry barrier moves with it. It is not about the layout -- it stops a peer's next round landing on data this rank is still reading -- so it opens whichever kernel runs first. A stale handle is guarded because the failure is silent. `send_base` and `send_rank_mask` are updated in place, so a handle from before some other dispatch does not fail, it routes to the wrong slots and returns plausible numbers. `Buffer._layout_generation` increments whenever a layout is computed and a cached dispatch asserts the handle still matches. API follows DeepEP: `topk_idx` and `topk_weights` must be `None` when a handle is passed, and `handle.topk_idx` is what gets replayed -- the field flagged as unread three commits ago, now doing the job it was kept for. 14 tests. The new one checks fresh, cached and cached-twice agree, and that a handle whose layout has been overwritten is rejected.
The docstring claimed the whole 34us notify kernel disappears. It does not, and the estimate was about twice the truth. Measured, three clean samples per row: fp8 dispatch 514 -> 497us whole call (3.3%), bf16 888 -> 871 (2.0%). The gap is visible in the kernel times: the cached scatter runs ~481us against the uncached one's ~470. The entry barrier moves into it rather than disappearing -- it is about peers not overwriting data this rank is still reading, which holds however the layout was obtained -- so roughly 10us of notify's 36 is work that has to happen either way. Also fixes ep2_scratch/run8.sh, which is what wedged the host: it put a timeout on the wrapper only, so a rank blocked in a rendezvous that never completed outlived it, and eight stranded ranks per attempt eventually left sshd unable to fork. Ranks now carry their own timeout, run under setsid so the trap takes their children, and the launcher picks a port it has verified is free rather than a random one that collided with TIME_WAIT.
The expanded layout was dispatch-only: combine refused it, because its store-back slot `comm_x[rank][src_token]` is unique only when dispatch has deduplicated. Expanded, a token with two experts on one rank has two rows and they collide there. So `do_expand` could not be used end to end at all. DeepEP's answer is `kDoExpandedSend` -- sum a token's local-expert rows before sending -- and that is what this does, so one row per (rank, token) still crosses NVLink and `comm_x` and the reduce are untouched. It needs the inverse of what dispatch records: dispatch gives row -> (src_rank, src_token), and this needs (src_rank, src_token) -> rows. A third kernel builds it by bucketing each received row under its source, metadata only, never the payload. The store-back then runs one warp per group, elected by the group's first row. The common case is unchanged: with routing spread over ranks most groups hold a single row, sent straight from `x` with no summing and no staging. Only groups of two or more accumulate into `reduce_scratch` first. Dispatch had to change too, and this was the subtle part. Deduplicated, a row carries every local expert of its token and the epilogue weights it by their sum. Expanded, the token has one row *per* expert, so leaving that alone would weight every row of the token by the same sum and combine would then add the duplicates -- correct-looking rows, wrong total. An expanded row now marks only the expert it belongs to. Two tracing constraints worth recording, both already documented elsewhere in this example for other reasons: - `alloc_var` names bind for the whole traced function, not per `T.Kernel`, so the bucketing kernel's locals had to be named distinctly from the store-back's or they read as one immutable variable escaping its region. - the scratch row is indexed by the unrotated block/warp pair. `warp` folds in `my_rank`, whose range the compiler cannot prove, and the resulting bounds check wraps the index in an `if_then_else` that `address_of` rejects. 16 tests. Two new ones cover the expanded round trip, at top-8 over 8 ranks so collisions are common rather than incidental, and with half the selections masked so the grouping pass has to skip the alignment padding dispatch writes.
It hides about 30us, and that is host and launch overhead rather than any part of the collective. Beside a 673us GEMM the dispatch still costs 1538us against a 1569us serial, where real overlap would be ~892. Three explanations ruled out by measurement rather than argument: it is not the `wait_stream` dependency, because `previous_event` gives 1538.7 against 1538.4; it is not SM starvation, because 8, 16 and 32 SMs all hide the same ~35us; and it is not that the saving is proportional and small, because the same ~30us is hidden whether the GEMM is 99us or 673us. What is left is that both dispatch kernels end in `T.sync_grid()` and are therefore cooperative launches, which do not co-schedule with other work. That would make it structural rather than a bug in the API, and suggests a fix -- the only thing needing a grid-wide rendezvous is the end-of-call reset, so moving it to the front of the next call would let the scatter be an ordinary launch. Recorded as untested. The API itself is correct and DeepEP-shaped; this is about what it is worth today, which the tests could not have told us.
The README speculated that `async_finish` hides so little because both dispatch kernels end in `T.sync_grid()` and a cooperative grid does not co-schedule with other work, and offered a structural fix. Tested, and it is not the cause. The end-of-call reset really is the only thing needing a grid-wide rendezvous, and moving it to the front of the next call does make the scatter an ordinary launch -- correct on all 8 ranks, since peers only publish into `count_matrix` after the entry barrier and the reset now runs before it. Overlap went from 1538.4us to 1552.8: nothing. Discarded rather than landed, per the rule that a change which does not clear the threshold gets reverted. Dropping the rendezvous also permits a scatter grid larger than the device's SM count, which turns out to be worse as well: 450us at 128 blocks against 621 at 256 and 541 at 512. That makes four explanations measured and rejected -- the `wait_stream` dependency, SM starvation, proportionality, and now this. The remaining candidate, untested, is that a bandwidth-bound collective gated by its slowest rank simply has little to give when a large GEMM is slowing every rank's participation.
Run through the identical harness on the same idle machine, DeepEP hides 81% of a 673us GEMM behind its dispatch -- 1177.8us overlapped against a 1724.7us serial -- where this port hides 4%. So the previous conclusion, that a bandwidth-bound collective gated by its slowest rank has little to give, was wrong. The reversal is worth stating plainly: this port is faster serially (1566 against 1725) and 23% slower overlapped. The cause is `comm_stream.wait_stream(compute_stream)` at the top of `dispatch`, taken whenever no `previous_event` is given. It makes the collective wait for everything already queued on the caller's stream, including work it has no dependency on. Recording an event once before any compute and passing it in, so the comm stream never waits on the compute stream, drops the overlapped time to 675us -- the collective disappears behind the GEMM. That configuration is not usable as it stands, since it also lets successive dispatches overlap each other on shared buffers, but it isolates the mechanism. Closing the gap means `async_finish` no longer implying that wait, which is a semantic change rather than a bug fix: DeepEP puts the ordering on the caller via `previous_event` and `allocate_on_comm_stream`, and this port currently takes it on itself. Recorded, not yet done. Also corrects the earlier note: the cooperative-launch theory was tested and refuted, along with SM starvation, proportionality, and `previous_event` recorded inside the loop.
…lap gap The previous commit named `comm_stream.wait_stream(compute_stream)` as the cause. It is not. DeepEP's `stream_control_prologue` takes exactly the same dependency when no `previous_event` is given -- which is how it was measured -- and overlaps anyway. The reasoning was that removing the wait fixes it for us, which is true but does not make it the cause, and the same experiment had already been called unusable for letting successive dispatches overlap on shared buffers. Two further things measured since, both negative and both narrowing it: A notify-only dispatch, 63us against the full 892us, hides 28.2us where the full one hides 29.7. So the amount hidden is constant across a 14x range of dispatch length as well as a 7x range of compute. That is the signature of overlapping the enqueue gap and nothing else, and it rules out both the scatter specifically and barrier spinning -- a dispatch that is almost entirely barriers hides the same absolute amount as one moving 850us of payload. Six explanations now measured and rejected: SM starvation, proportionality, the cooperative launch, grids beyond the SM count, `previous_event` recorded inside the loop, and the stream dependency itself. The section says so rather than offering another guess.
…ther Best-motivated theory yet, and wrong. `sync_grid` is `cudaLaunchCooperativeKernel`, which requires every block co-resident and so cannot be co-scheduled with other work, and EPv2 has no `this_grid().sync()` in its intranode path at all -- only in the legacy `internode_ll`. That is a real structural difference between a dispatch that overlaps and one that does not. It had been tested twice, badly, and both times only half of it: the original arrangement leaves the scatter cooperative and the notify ordinary, and moving the reset to the front of the next call swaps which one is which. Every variant measured still contained exactly one cooperative launch, which is why "no change" meant nothing. Removed properly by moving the padding, stats and reset into a third kernel -- a kernel boundary supplies the same ordering the grid sync did, since `barrier_blocks` rendezvouses ranks rather than a rank's own blocks. Dispatch then has no cooperative launch anywhere, is correct on all 8 ranks, and hides 27.3us out of 1534.8. Indistinguishable from every other variant. Seven explanations now measured and rejected. The constant is striking: ~30us hidden regardless of dispatch length, compute length, SM count, launch mode or stream dependency, while removing the compute-stream dependency entirely gives full overlap. Further A/B is the wrong instrument; this needs a timeline.
Seven A/B experiments could only ever say "not this one". A timeline says what actually happens, and it is a strict alternation: notify, GEMM, scatter, GEMM, notify... with the GEMM starting 9-10us after the scatter ends, every iteration, and never running during it. The only concurrency is the next notify overlapping the GEMM's tail by 40-48us -- which is the constant ~30us that every measurement all day has been reporting. The dependency cycle accounts for it exactly. Iteration N+1's dispatch calls `wait_stream(comm, compute)` after iteration N's GEMM is already queued on the compute stream, so the communication waits for the previous GEMM; `with event:` puts the event wait on the compute stream, so the next GEMM waits for the previous dispatch. Each stream ends up waiting on the other's last item. DeepEP traced identically does not alternate: its GEMM starts 112us and 0.6us into two successive dispatch kernels and runs entirely inside them. Same stream plumbing on both sides -- verified in their source -- so the difference is in how the kernels are launched or admitted, not in the API. Records the numbers rather than another theory. This is where to pick it up.
The async path was already sound; it hid 4% of a concurrent GEMM because a private stream buys eligibility, not admission. The trace shows the collective becoming eligible the instant the previous scatter ends and then waiting 624us anyway, for an SM. A GEMM worth hiding behind is also large enough to hold the whole device -- 2048 blocks, each taking a full SM of registers and 213KB of shared memory, so 148 are resident and fourteen waves pending -- and block admission is greedy, so every SM that frees goes to the next GEMM block. The collective gets in only as the last wave drains, a fixed ~48us window, which is why the amount hidden was a constant ~30us regardless of dispatch length, compute length or SM count. The experiment that settled it: `torch.cuda._sleep` calibrated to the same 674us as the GEMM occupies the compute stream just as long but uses one block, and is hidden 100%. Same streams, same events, same dependencies, 147 SMs free instead of none. Occupancy, not ordering. Stream priority biases exactly that admission decision. 29.7us hidden -> 286.3 (4% -> 43%) at the V3 shape. Any raised priority does it -- -1, -2 and -3 all measure ~285us -- so the buffer takes whatever the device offers and a caller scheduling several streams can override it. `num_sms` matters now for the first time, being how many blocks are queued for admission: 158us at 16, 223.6 at 32, 285.3 at 64. Standalone performance is unchanged, which is the point: 898-900us bf16 dispatch, 984-985 combine, 522 fp8, all inside the spread of the existing numbers. 16 tests pass. Supersedes the six rejected explanations recorded in the previous commits; the README now carries all of them with their numbers.
The number was already there; this is the mechanism. Before the priority
change the GEMM started 9-10us after the scatter ended, every iteration, and
the two were never resident together. After, they are:
scatter [2916.8 -> 3770.6] GEMM [2605.8 -> 3505.6] 588.8us concurrent
scatter [4057.4 -> 4892.2] GEMM [3781.8 -> 4704.4] 647.0us concurrent
Also records that 64 SMs is already the peak -- 158us hidden at 16, 223.6 at
32, 285.3 at 64, 267.7 at 96, 96.9 at 128, where enough blocks are queued for
admission that the collective starves itself -- so the existing default needs
no change.
DeepEP still hides 81% against this port's 43%, and that difference is stated
as unexplained rather than guessed at. Its dispatch is one kernel where this
one is two, which gives the GEMM a seam to fall into.
Trace archived alongside the before/after pair in ep2_agent_docs/overlap-traces.
…tream Worth 144.6us of hidden time, 43% to 64%, for moving two lines above a wait. The admission stall that the stream-priority fix addressed is not paid uniformly. There is exactly one moment per iteration when it is free -- the ~2.5us window after the previous scatter releases its SMs and before the GEMM refills them -- and whichever operation is queued first on the communication stream gets it. Everything behind it pays 90-140us. That slot was going to `topk_idx.to(int32)`, a 3us elementwise kernel issued inside the communication-stream block, which won it 0.4us ahead of the GEMM and left the 850us collective behind it to pay 129.9us. Converting on the caller's stream puts the collective first in the queue; `wait_stream` is issued after the conversions, so the dependency still covers them. Isolated from the caller side first, by pre-converting the tensors so the buffer's cast becomes a no-op: 285.0us hidden with it, 425.7 without. Hoisting inside the buffer then measures 430.9 of 675.8. 16/16 tests pass and standalone dispatch is 890.9us, unchanged. Also corrects a claim in e4e519b, which is pushed and cannot be amended: EPv2's intranode path does use `this_grid().sync()`, at common/comm.cuh:233,251,269 from impls/dispatch.cuh:74,398. Cooperative launch was never the structural difference that commit said it was -- the test it motivated stands, since it measured our own kernels and found nothing.
Worth 126.6us of hidden time, 64% to 83%, which puts this port past DeepEP's 81%
on the same benchmark. Standalone is unchanged: 899.0us bf16 dispatch, 982.8
combine, 522.8 fp8.
With the stream priority and the hoisted cast in, one steady-state iteration:
cast 1170.8 -> 1173.4 2.6us
GEMM 1173.1 -> 2070.4 897.3us g=2048
notify 1299.3 -> 1325.9 26.6us 126.2us after the GEMM
scatter 1444.6 -> 2279.3 834.6us 118.7us after the notify
The scatter overruns the GEMM by 208.9us and that is all of what stays unhidden.
Not because it runs slowly while sharing the device -- 834.6us here against
844.1 solo -- but because it starts 271.5us late. Of that, 126.2us is the first
admission, which DeepEP pays too, and 118.7us is the gap between our own two
kernels. Solo that gap is 1.9us. Sixty-two times wider under contention, because
the notify releases all 64 SMs on the way out and the GEMM has twelve waves
queued to take them back.
Programmatic dependent launch is how DeepEP avoids a second admission -- its
epilogue seam measures -30.9us -- and TileLang exposes it, so it was tried
first: T.pdl_sync() on the dependent kernels, verified in the generated CUDA and
in the launch path, worth 8.2us. Pre-admission places the dependent grid while
the primary is resident, but a GEMM CTA holds 213KB of the SM's 228KB of shared
memory, so there is nothing for a comm block to be placed beside and the
placement still waits for an SM to drain.
Not giving the SMs up at all is the fix, and that means one kernel. The kernel
was one before 17123c6 split it for scatter_sms, and that commit's own numbers
put the fusion at 883.7us against 881.7 split -- free on an idle device, which
is where it was measured. The split survives for the case that motivated it, a
scatter grid wider than the notify's, which one sync_grid-bearing kernel cannot
express; scatter_sms=128 still gives 868.4us there.
Mechanically the phases become T.macros so both packagings share them. They sit
in dispatch_kernel rather than inside main: the eager builder rewrites a
prim_func body into a generated function and hoists nested defs out of its
scope, so a helper defined inside main is not in scope at its own call site.
16/16 tests pass, including the one that runs the scatter on a wider grid than
the notify, which is now the only caller of the split path.
kernels/dispatch.py 729 -> 583 lines, and the whole example 2767 -> 2606.
scatter_sms let the scatter run on a wider grid than the notify, worth about 4%
standalone. Fusing the two phases into one launch -- worth nineteen points of
overlap -- cannot express two grids, so the fusion kept both packagings: two
T.macros carrying thirty-two buffer parameters between them, a `fused` flag and
a duplicate launch site.
Measured, none of that earns its place. A fused kernel simply given the wider
grid beats the arrangement it was preserving:
num_sms=64 898.2 us
num_sms=96 873.6
num_sms=128 861.0 <- one grid, fused
num_sms=64, scatter_sms=128 867.1 <- two grids, split
So scatter_sms=128 is num_sms=128 spelled less clearly and 6.5us slower, and
what it bought the notify -- 64 blocks instead of 128 for 26us of work -- is
worth nothing. The knob goes, the split goes with it, and the kernel is one
grid with one stride throughout. The test that covered it now covers num_sms=128
on the fused grid, which exercises the same three strided loops.
Also, in passing and each measured only as unchanged:
- alloc_var declarations followed by a first assignment fold into init=, which
is what the rest of the file does
- combine's group reduce loops to the runtime count instead of unrolling to the
compile-time maximum and predicating back down; T.serial already takes a
runtime bound, as the token loop above it shows
- combine seeds the accumulator from bias_0 with T.copy rather than a Parallel
loop, since T.copy carries the cast
- dispatch and combine had twelve duplicated lines of stream handover; they are
_begin_comm and _end_comm now, which is also where the rule about what may be
queued first on the communication stream belongs, since that is easy to break
by adding a line in the wrong place. DeepEP factors the same pair.
16/16 tests pass. Standalone 898.1us bf16 dispatch, 984.5 combine, 521.3 fp8,
all within noise of before; overlap 82.8-83.0%, likewise.
The DeepEP column was bench_kineto GPU time, split into main kernel and
epilogue, sitting next to whole-Python-call figures on this side. That is not a
comparison: it charges this port for host overhead and launches that DeepEP's
column does not carry, while splitting DeepEP's dispatch into two rows neither
of which is what a caller pays.
Both columns now come from one harness -- same do_bench, same warmup and rep
counts, same shape, same 5s clock warm-up, same eight ranks, run back to back
on one idle machine. DeepEP runs with do_cpu_sync=False and do_handle_copy=False,
work this port has no equivalent of. Two rounds each, both reported.
DeepEP this port
dispatch bf16 1047.7 / 1049.2us 890.2 / 892.1 -15%
dispatch fp8 566.7 / 568.1 512.0 / 512.6 -10%
dispatch cached 1045.4 / 1049.2 874.6 / 875.8 -16%
combine bf16 988.9 / 992.1 954.4 / 954.9 -3.7%
overlapped (8192 GEMM) 1176.3 / 1176.2 1002.7 / 1006.0
hidden 546.6 / 548.2 82% 558.4 / 559.4 83%
A whole layer's collectives come to 2039us against 1846 in bf16 and 1558
against 1467 with an fp8 dispatch.
handle= is the row where the two differ in kind rather than degree: DeepEP
saves nothing measurable, because its layout work lives inside the main kernel,
where here it is a phase that can be skipped outright.
Also records what the numbers do not show and matters more than any of them --
DeepEP is a production multi-node implementation and this has no RDMA path --
along with where DeepEP is ahead (745 against 723 GB/s on the movement itself,
and a gentler slope as SMs are taken away), and the size of the two intranode
paths: ~2110 lines of CUDA against 828 of TileLang, 1107 lines of host Python
against 790. The two epilogues and the separate layout module are most of that
difference.
8e2ab99 was meant to replace the "Against DeepEP" section and instead duplicated 230 lines of the document, leaving the stale comparison it was replacing still in place below the new one. 712 lines where there should be 482, with "### Overlapping", "## Performance" and "### Against DeepEP" each appearing twice. The patch bounded the region to replace with s.index("## Overlap"), which prefix-matches the "### Overlapping" heading 165 lines *above* the section being replaced. End before start, so s[:start] + new + s[end:] re-appended everything in between. Nothing checked the result, which is the actual mistake -- the section boundaries had been printed a few commands earlier and showed exactly this. Same content as intended, bounded correctly this time.
…emeasure
Buffer defaults to 1024 threads for both collectives; the benchmark script still
asked for 512 and 256, so every figure in the README's Performance section
described a configuration the library does not use. A/B at 64 SMs, two rounds,
interleaved:
threads dispatch combine
512 / 256 897.4 / 897.9 985.4 / 984.3
1024 / 1024 892.3 / 891.9 955.1 / 955.4
1024 / 512 892.0 / 892.3 962.0 / 962.1
1024/1024 wins both, by 0.6% on dispatch and 3.0% on combine, so the script asks
for what the library already does and the tables are remeasured against it:
bf16 690.8 GB/s 892.0us combine 644.5-645.4 954.7-956.1
fp8 595.4-596.8 516.2-517.5 644.8-645.0 955.4-955.7
bf16 num_sms=128 718.2-718.3 857.9-858.0 659.0-659.6 934.1-935.1
bf16 num_sms=24 599.5-599.6 1027.7-1027.8 567.4-568.8 1083.3-1085.9
So the slope from 64 to 24 SMs is -13% / -12%, not -13% / -13%.
The per-phase breakdown is remeasured from a trace too. Dispatch is 865.9us of
kernel inside an 892.0us call. It is one kernel now, so the phases inside it no
longer have a launch boundary to be separated at, and the old decomposition into
scatter, barriers and count exchange cannot be reproduced -- what is left is the
total, the count exchange and the barriers. Combine is 866.6us of store-back and
114.6us of local reduce, which sum to more than its 955us call because the trace
isolates each call behind a full synchronise and so pays the cross-rank entry
barrier in full every iteration.
…ded dispatch The comparison section was all timings, while what is and is not supported sat split between the API section and Not implemented at opposite ends of the document -- so the one place a reader goes for "how does this compare" answered only half the question. Eleven rows now, pointing at both for the reasoning. Writing it down is what caught the stale part. The README still claimed combine cannot consume an expanded dispatch and asserts rather than returning a wrong answer. That stopped being true at ca1430d: there are two passing round-trip tests, a bucketing kernel in kernels/combine.py that builds the (src_rank, src_token) -> rows inversion the claim says is missing, and no assert anywhere that would block it. The Not implemented entry for kAllowMultipleReduction repeated the same error and is corrected with it -- the capability is covered, by summing on the sender rather than the receiver, and it is the flag rather than the feature that is absent.
Rachmanino
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
1 participant
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.
A TileScale-native MoE all-to-all following DeepEP EPv2's
impls/dispatch.cuhand
impls/combine.cuh, scoped to intranode NVLink on SM100. Dispatch payloadis bf16 or fp8 (per-token, per-128-element scales, DeepEP's
per_token_cast_to_fp8layout); combine is bf16.Design
dispatch runs three phases in one kernel with no host round trip: count
(dedup per (token, destination) with
T.match_any_sync, tally into sharedmemory with no atomics), exchange (publish this rank's count vector to every
peer so all ranks can derive
send_base[d]), scatter (one warp per token).Rows go straight to their final compact index, so there is no copy epilogue.
DeepEP needs one because its expand layout hides the final position at send
time; here the full count matrix makes
send_base[d] + slotfinal, which savesa whole local read+write of the payload and costs a ~19µs wait for the exchange.
combine stores back into a slot unique per (contributing rank, source
token) — no atomics — then reduces locally, one block per source token.
Both kernels bracket themselves with
tl::barrier_blockson private slots andrun on a private CUDA stream, so neither needs a
dist.barrier.Performance
8x B200, 8192 tokens/rank, hidden 7168, top-8, 256 experts, 64 SMs, bottleneck
rank. Compared against DeepEP's
tests/elastic/test_ep.pyon the same machine.DeepEP's headline covers
dispatch_impl/combine_implonly, so its epiloguesare added back here — a plain dispatch → expert → combine loop has nothing to
hide them behind:
DeepEP is ahead on the cross-rank movement itself (745 vs 723 GB/s on combine's
store-back, which is this port's
put_warproofline) and gives it back to theepilogues.
At 24 SMs: 589–599 GB/s dispatch, 557–562 combine — −13%/−12% against DeepEP's
−11%/−9% over the same range.
DeepEP's figures are GPU time only; this port's are whole Python calls.
Compiler change
LegalizeSafeMemoryAccessrewrote the argument ofaddress_ofintoif_then_else(cond, load, safe), which every consumer rejects with"address_of argument must be a BufferLoad". Callers had to work around it by
clamping every such index into range by hand and disabling
LoopUnswitching.Taking an address is not a dereference: the bound belongs to whoever
dereferences the pointer.
address_ofnow visits its argument's indices butkeeps it a
BufferLoad, exactly astl.access_ptralready did. No working IRcould have depended on the old behaviour, since it only ever produced IR that
failed downstream.
Testing
pytest examples/distributed/deepep_v2/test_example_deepep_v2.py(2 GPU smoke(rel_l2 0.003129 bf16 / 0.025973 fp8) and the compact list carries no duplicate
(src_rank, src_token) pairs — both checked again after thousands of
back-to-back calls, since every race this port has shipped was invisible on a
first call.