diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 2d418f0142..d41968132b 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -1,3 +1,4 @@ +agrs cancelled dout HDA diff --git a/examples/distributed/deepep_v2/README.md b/examples/distributed/deepep_v2/README.md new file mode 100644 index 0000000000..9751279ad3 --- /dev/null +++ b/examples/distributed/deepep_v2/README.md @@ -0,0 +1,504 @@ +# DeepEP EPv2 port — intranode dispatch / combine + +A TileScale-native MoE all-to-all, following DeepEP EPv2's `impls/dispatch.cuh` +and `impls/combine.cuh`. Scoped to **intranode NVLink** on Blackwell (SM100): +no RDMA/scaleout, no low-latency decode path, no expert-alignment or expand +layout. + +Two collectives, both written in TileLang with no inline PTX: + +- **`dispatch`** — send each token to the ranks owning its top-k experts, and + land it directly at its final index in the receiver's compact buffer. +- **`combine`** — send each expert's contribution back and reduce, per source + token, into the original row order. + +## When to use this + +Suited to intranode expert parallelism where a whole MoE layer fits inside one +NVLink domain (EP ≤ 8 on a B200 node). Dispatch payload may be bf16 or fp8; +combine is always bf16, since it carries expert output. + +Not a drop-in for DeepEP: there is no RDMA path, so multi-node EP is out of +scope, and the low-latency decode kernels are not ported. + +## Design + +### dispatch — three phases in one kernel + +1. **Count.** Every warp scans this rank's `topk_idx` and deduplicates per + (token, destination rank) with `T.match_any_sync` — a token whose top-k + picks two experts on one rank is one row, not two. Tallies land in a + per-warp slice of shared memory; dedup makes the recording lanes hold + pairwise-distinct destinations, so the tally needs no atomics. +2. **Exchange.** One warp publishes this rank's count vector to every peer, so + all ranks hold the full `count_matrix[sender][destination]` and can derive + `send_base[d]` — where their rows begin in destination `d`'s output. +3. **Scatter.** One warp per token. Destinations claim slots with one round of + `atom_add` on a local counter, then the warp pushes the row, its scales (fp8 + only) and its metadata to `send_base[d] + slot` on peer `d`. + +**No copy epilogue.** DeepEP stages rows per sender and compacts them in a +second kernel, because its expand layout hides the final position at send time. +Here the full count matrix makes `send_base[d] + slot` the final compact index, +so rows land in place — saving a whole local read+write of the payload, at the +cost of the scatter having to wait for the exchange (~19µs). + +### combine — store-back, then local reduce + +One warp per compact row stores into `comm_x[my_rank][src_token]` on the source +rank; that slot is unique per (contributing rank, source token), so nothing +needs to be atomic. A second kernel then sums, per source token, the slots named +by the destination mask dispatch recorded. + +### Synchronisation + +Both kernels open and close with `tl::barrier_blocks` on private slots, so +neither needs a `dist.barrier` around it. Note that `barrier_blocks` +rendezvouses *ranks*, not a rank's own blocks — any single-block epilogue behind +it also needs `T.sync_grid()`. + +Communication runs on a private CUDA stream. `Buffer.pipeline_depth` bounds how +far the CPU may run ahead, which matters because a rank queued several calls +behind stalls every other rank inside the cross-rank barrier. + +## API + +```python +from buffer import Buffer + +buf = Buffer( + group=group, local_rank=local_rank, num_local_ranks=8, + num_max_tokens_per_rank=8192, hidden=7168, num_topk=8, num_experts=256, + dtype=torch.bfloat16, # or torch.float8_e4m3fn for fp8 dispatch + num_sms=64, +) + +# bf16: x is [num_tokens, hidden] +# fp8: x is (values, scales) from reference.per_token_cast_to_fp8 +recv_x, recv_topk_idx, recv_topk_weights, handle, event = buf.dispatch(x, topk_idx, topk_weights) + +n = handle.num_recv_tokens # device-to-host read; call outside timed regions +expert_out = my_expert_compute(recv_x[:n], recv_topk_idx[:n], recv_topk_weights[:n]) + +combined, event = buf.combine(expert_out, handle) # [num_tokens, hidden] bf16 +``` + +`dispatch` returns views over the **full receive capacity**; slice with +`handle.num_recv_tokens` when you need the compact rows. On the fp8 path the +first return value is `(values, scales)` and the caller casts back before the +expert computation — `reference.per_token_cast_back` does this. + +### Overlapping + +Both collectives return an `EventOverlap` as their last value, DeepEP's wrapper +around the communication-stream event. It comes back either way -- synchronously +it wraps `None` -- so a caller can write `with event:` without knowing which +mode it asked for. + +```python +recv_x, recv_topk_idx, recv_topk_weights, handle, event = buf.dispatch( + x, topk_idx, topk_weights, async_finish=True, allocate_on_comm_stream=True +) +with event: # runs on the compute stream, overlapping the dispatch + something_else() +# leaving the block, the current stream waits: recv_x is readable + +combined, event = buf.combine( + expert_out, handle, previous_event=event, async_finish=True, allocate_on_comm_stream=True +) +event.current_stream_wait() +``` + +`async_finish` leaves the caller's stream unjoined from the communication +stream; nothing returned may be read until the event is waited on. EPv2 spells +this `async_with_compute_stream` on dispatch; the name here follows its +`combine` and DeepEP's legacy buffer. + +`previous_event` starts the communication after one specific event instead of +after everything queued on the caller's stream. `allocate_on_comm_stream` keeps +this call's temporaries owned by the communication stream and alive through the +returned event, rather than `Tensor.record_stream`, which CUDA graph capture +does not permit -- the reason DeepEP carries `extra_tensors` on its event. As in +DeepEP, `previous_event` requires `allocate_on_comm_stream`. + +One asymmetry with DeepEP worth knowing: `Buffer.pipeline_depth` (default 2) +bounds how far the CPU may run ahead by blocking the *host* on an event a few +calls back. It exists because a rank queued several calls behind stalls every +other rank inside the kernel's cross-rank barrier, and it is orthogonal to +`async_finish` -- so an asynchronous call can still block the host. Pass +`pipeline_depth=0` to turn it off when driving the overlap yourself. + +DeepEP's `combine` returns a third value, `combined_topk_weights`, which has no +counterpart here for the reason given under *Not implemented*. + +`dispatch(x, handle=h)` reuses the layout `h` was built with and skips the +notify kernel outright -- DeepEP's cached dispatch. Phases 1-2 depend only on +the routing, so a call whose `topk_idx` has not changed is recomputing +something it already has. As in DeepEP, `topk_idx` and `topk_weights` must be +`None`; the handle replays its own copies. + +Measured, three clean samples each: fp8 dispatch 514 -> 497 µs (3.3%), bf16 888 +-> 871 µs (2.0%). Less than the 36 µs the notify kernel costs, because the +cached scatter runs about 10 µs slower than the uncached one: the entry barrier +moves into it rather than disappearing, being about peers not overwriting data +this rank is still reading, which holds however the layout was obtained. + +A handle is only good until the next layout-computing dispatch. `send_base` and +`send_rank_mask` are updated in place, so a stale handle would not fail, it +would route to the wrong slots and return plausible numbers -- dispatch +therefore tracks a layout generation and rejects a handle that no longer +matches. + +`dispatch(..., cumulative_local_expert_recv_stats=t)` adds this rank's received +token count per local expert into `t`, a `[num_experts // num_ranks]` uint32 +tensor -- DeepEP's load-balance counter. It accumulates rather than overwrites, +so the caller decides the window by choosing when to zero it. Costs ~25 µs of +dispatch's ~896 (2.8%) and compiles its own kernel variant, so the default path +does not pay for it. DeepEP gets the same number for free because its expanded +layout already exchanges per-expert counts; this port has no such exchange and +counts locally instead, which is why it is not free here. + +`dispatch` produces DeepEP's **expanded layout** when the buffer is built with +`do_expand=True`: one received row per (token, expert) instead of one per +(token, rank), and rows grouped by local expert, so each expert's rows are the +contiguous block a grouped GEMM wants. `handle.expert_offset` gives the segment +bounds and `handle.expert_count` how many rows in each are real; +`expert_alignment=n` rounds each segment up to a multiple of `n` and the gap is +zeroed unless `zero_padding=False`. + +DeepEP expands in a receiver-side copy epilogue. This port has none -- rows land +at their final index straight from the sender -- so instead the count exchange +runs at expert granularity and the sender derives the index itself. That also +makes the capacity check free: every rank computes the same layout from the same +count matrix, so an overflow is known before any payload moves. Deduplicated, +capacity cannot be exceeded; expanded it can, so `expand_factor` sizes the +receive buffer (default 1.0, right for balanced routing) and `handle.expand_overflow` +reports how many rows a call needed if it did not fit -- dispatch skips the rank +rather than writing past it. + +`combine` consumes an expanded dispatch. Its store-back slot is +`comm_x[rank][src_token]`, unique only because a deduplicated dispatch gives a +token one row per destination rank; expanded, a token with two experts here has +two rows that would collide. The answer is DeepEP's `kDoExpandedSend` -- sum a +token's local-expert rows before sending, so one row per (rank, token) still +crosses NVLink and the store-back and reduce are unchanged. That needs the +inverse of what dispatch records, so a third kernel buckets each received row +under its source; it touches metadata only, never the payload. With routing +spread over many ranks most groups hold a single row and are sent straight from +`x` with no summing. + +`combine(..., bias=b)` adds one tensor, or `bias=(b0, b1)` two, to the output -- +DeepEP's `bias_0`/`bias_1`, each `[num_tokens, hidden]`. They seed the reduce +accumulator instead of being added after it, so they cost nothing measurable, +and a token whose every selection was masked off still comes back as its bias. + +Knobs worth tuning: `num_sms`, `dispatch_threads`, `combine_threads`, and +`reduce_threads` (separate because the reduce wants `hidden / reduce_threads` to +be a whole number of 128-bit loads). The thread defaults are wide (1024) because +that measured at least as fast at every SM count tried: against 512/256 it is +worth 0.6% on dispatch and 3.0% on combine. + +## Performance + +8× B200, full NVLink mesh, 8192 tokens/rank, hidden 7168, top-8, 256 experts, +64 SMs. Bandwidth is the bottleneck rank's, over payload bytes that cross +NVLink. Three to four samples per row, each one gated on whether any process +that is not ours *used the SMs* at any point during the run, sampled throughout +with `nvidia-smi pmon`. Presence is not the test: an 8-way inference server +holding 167 GB/GPU at 0% utilisation blocks a presence-based gate 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; both produced +numbers 30–60% off. + +| | dispatch | combine | +|---|---|---| +| bf16 | **690.8 GB/s** (892.0 µs) | **644.5–645.4 GB/s** (954.7–956.1 µs) | +| fp8 | **595.4–596.8 GB/s** (516.2–517.5 µs) | 644.8–645.0 GB/s (955.4–955.7 µs) | +| bf16, `num_sms=128` | **718.2–718.3 GB/s** (857.9–858.0 µs) | **659.0–659.6 GB/s** (934.1–935.1 µs) | + +`num_sms` is the whole knob. Dispatch is one launch whose grid-wide rendezvous +needs every block resident, so it cannot exceed the device's 148 SMs, but +anything up to that is fair game: 898.1 µs at 64, 873.6 at 96, 861.0 at 128. +The default leaves the rest of the device for expert compute; a caller who +wants dispatch to have it says so. + +Combine is bf16 whichever dtype dispatch used, and measures the same either +way, as it should. + +FP8's rate is below bf16's partly as an accounting artifact: it counts the +7168 payload bytes a row carries, but the row that crosses NVLink is 7680 -- +the per-128 fp32 scales packed in alongside, plus alignment padding (see +`reference.packed_row_bytes`). On the wire that is ~632 GB/s. The rest of the +gap is the kernel's fixed phases -- notify, dedup, count exchange, metadata +stores -- costing the same in absolute terms against a payload half the size. + +From a trace, a bf16 dispatch is 865.9 µs of kernel inside an 892.0 µs call, so +about 26 µs is host and launch. It is one kernel now, so the phases inside it no +longer have a launch boundary to be separated at; the scatter is nearly all of +it and runs at the `put_warp` roofline, with the count exchange (~19 µs) and the +entry and exit barriers (~44 µs) the only other measurable pieces. + +Combine's two kernels are 866.6 µs of remote store-back and 114.6 µs of local +reduce (733 MB at ~6.4 TB/s), an 88/12 split. Those sum to more than the 955 µs +whole call because the trace isolates each call behind a full synchronise, so +the store-back pays its cross-rank entry barrier in full every iteration where +a pipelined benchmark does not. + +Fewer SMs, same shape: + +| #SMs | dispatch | combine | +|---|---|---| +| 64 | 690.8 GB/s | 644.5–645.4 GB/s | +| 24 | 599.5–599.6 GB/s | 567.4–568.8 GB/s | + +That is −13% / −12% against DeepEP's −11% / −9% over the same range. At 24 SMs +each block already carries 16 warps, past the point where `put_warp` saturates, +so widening blocks does not help (1024 threads against 512 is ~1% at either +end): what runs out is SM count itself. + +### Against DeepEP + +Both columns from one harness: same `do_bench`, same warmup and rep counts, +same shape, same 5 s clock warm-up, same eight ranks, run back to back on one +idle machine. Every figure is a whole Python call, so it carries the host side, +the launches, and every kernel in the call -- which is the only accounting under +which the two implementations are comparable, since DeepEP's dispatch is a main +kernel plus a copy epilogue and this port's is one kernel with no epilogue at +all. DeepEP runs with `do_cpu_sync=False` and `do_handle_copy=False`, work this +port has no equivalent of. Two rounds each; both are shown. + +| | DeepEP | this port | | +|---|---|---|---| +| dispatch bf16 | 1047.7 / 1049.2 µs | **890.2 / 892.1** | −15% | +| dispatch fp8 | 566.7 / 568.1 | **512.0 / 512.6** | −10% | +| dispatch cached (`handle=`) | 1045.4 / 1049.2 | **874.6 / 875.8** | −16% | +| combine bf16 | 988.9 / 992.1 | **954.4 / 954.9** | −3.7% | + +A whole layer's collectives -- dispatch → expert → combine, where neither side +has anything to hide an epilogue behind -- come to 2039 µs against 1846 (−9.5%) +in bf16, and 1558 against 1467 (−5.8%) with an fp8 dispatch. + +`handle=` is the one row where the two differ in kind rather than degree. +DeepEP saves nothing measurable from it, because its layout work lives inside +the main kernel; here it is a phase that can be skipped outright, worth 16 µs. + +**Overlap**, same harness, against an 8192-square bf16 GEMM: + +| | DeepEP | this port | +|---|---|---| +| compute | 670.3 / 670.5 µs | 674.4 / 672.5 | +| dispatch alone | 1049.6 / 1050.6 | 890.0 / 889.5 | +| serial | 1723.0 / 1724.4 | 1564.4 / 1562.1 | +| overlapped | 1176.3 / 1176.2 | **1002.7 / 1006.0** | +| hidden | 546.6 / 548.2 (82%) | **558.4 / 559.4 (83%)** | + +**Where DeepEP is ahead.** The cross-rank movement itself: 745 GB/s against 723 +on combine's store-back, about 7%, which is this port's roofline for +`put_warp`. It gives that back to the epilogues -- writing rows straight into +their final compact index costs a wait for the count exchange (~19 µs) and +saves a whole extra pass over the payload. It also degrades more gracefully as +SMs are taken away: from 64 to 24, this port loses 13% and 12% against +DeepEP's 11% and 9%. + +Note too that DeepEP's dispatch epilogue *produces* `recv_x`, so it cannot +overlap the expert computation that consumes it, and that its own headline +numbers report the main kernel and the epilogue separately. + +**Coverage.** Within the intranode path: + +| | DeepEP | this port | +|---|---|---| +| `async_finish` / `previous_event` / `allocate_on_comm_stream` | yes | yes | +| `handle=` cached dispatch | yes | yes | +| per-call `num_sms` | yes | yes | +| fp8 payload with per-128 scales | yes | yes, packed into the row | +| expanded layout, `expert_alignment`, zero padding | yes | yes | +| combine consuming an expanded dispatch | `kAllowMultipleReduction` | yes, summed on the sender | +| `cumulative_local_expert_recv_stats` | yes | yes | +| combine bias | yes | yes | +| `deterministic` mode | yes | **no** | +| `use_tma_aligned_col_major_sf` | yes | **no** | +| combine's `topk_weights` output | yes | deliberately absent | + +The API section above and *Not implemented* below give the reasoning for each. + +**Scope.** This is the comparison that matters most and no measurement shows +it: DeepEP is a production multi-node implementation. There is no RDMA path +here, so multi-node EP is out of reach, along with the low-latency decode path +and Engram/PP/CP. Everything above is one node. + +**Size**, for the intranode dispatch/combine path only: + +| | DeepEP | this port | +|---|---|---| +| kernels | ~2110 lines CUDA | **828 lines TileLang** | +| host | 1107 lines Python | **790** | + +DeepEP's kernel figure is `impls/dispatch.cuh` (411) plus its copy epilogue +(325), `impls/combine.cuh` (245) plus its reduce epilogue (145) and utils +(172), and the `common/` comm, layout and handle headers (272 + 313 + 230). The +two epilogues and the separate layout module are most of the difference, and +this port has neither. + +**What `async_finish` buys.** Measured against an 8192-square bf16 GEMM sized +to match the collective, on an idle machine: + +| | compute | dispatch | serial | overlapped | hidden | +|---|---|---|---|---|---| +| before | 672.4 us | 891.8 | 1568.9 | 1538.4 | 29.7 us (4%) | +| + stream priority | 671.0 us | 891.5 | 1566.2 | 1279.9 | 286.3 us (43%) | +| + cast hoisted | 675.8 us | 890.9 | 1566.7 | 1135.8 | 430.9 us (64%) | +| + fused launch | 673.5 us | 889.8 | 1563.3 | **1004.4** | **558.9 us (83%)** | + +Three fixes, DeepEP measures 81% on the same benchmark, and the reason each is +worth what it is worth is worth writing down -- six plausible theories were +measured and rejected before the first one. + +A private stream buys *eligibility*, not *admission*. The trace shows the +collective becoming eligible the instant the previous scatter ends, and then +waiting 624 us anyway. What it is waiting for is an SM. A GEMM large enough to +be worth hiding behind is also large enough to hold the whole device: 2048 +blocks, each needing a full SM's registers and 213 KB of shared memory, so 148 +are resident and fourteen waves are pending. Block admission is greedy, so +every SM that frees goes to the next GEMM block, and a collective that arrives +even microseconds later gets in only as the last wave drains -- a fixed ~48 us +window, which is why the amount hidden was a constant ~30 us regardless of +dispatch length, compute length or SM count. + +The experiment that settled it: `torch.cuda._sleep` calibrated to the same +674 us as the GEMM occupies the compute stream just as long but uses **one +block**, and it is hidden **100%**. Same streams, same events, same +dependencies, 147 SMs free instead of none. Not ordering -- occupancy. + +Stream priority biases precisely that admission decision. Any raised priority +works (-1, -2 and -3 all measure ~285 us), so the buffer takes whatever the +device offers. `num_sms` matters now for the first time, since it is how many +blocks the collective is trying to get admitted: 158 us hidden at 16 SMs, +223.6 at 32, 285.3 at 64. + +The timeline confirms the mechanism rather than just the number. Before, the +GEMM started 9-10 us *after* the scatter ended, every iteration, and the two +were never resident together. After: + +``` +scatter [2916.8 -> 3770.6] GEMM [2605.8 -> 3505.6] 588.8 us concurrent +scatter [4057.4 -> 4892.2] GEMM [3781.8 -> 4704.4] 647.0 us concurrent +``` + +`num_sms` is now the knob that matters, and 64 is already the peak: 158 us +hidden at 16, 223.6 at 32, 285.3 at 64, 267.7 at 96, and 96.9 at 128, where so +many blocks are queued for admission that the collective starves itself. + +**The second fix: do not spend the free admission slot on a dtype cast.** The +admission stall above is not paid uniformly. There is exactly one moment per +iteration when it is free -- the ~2.5 us window after the previous scatter +releases its SMs and before the GEMM's next wave refills them -- and whichever +operation is queued first on the communication stream gets it. Everything +behind it pays ~90-140 us. + +That slot was going to `topk_idx.to(int32)`, a 3 us elementwise kernel that the +buffer issued inside the communication-stream block, 0.4 us ahead of the GEMM. +The 850 us collective behind it then paid the full 129.9 us. Converting on the +caller's stream instead -- before `wait_stream`, so the dependency still covers +it -- puts the collective's own kernel first in the queue. Worth 144.6 us, or +43% to 64% hidden, for moving two lines above a `wait`. + +Isolated first from the caller side, by pre-converting the tensors so the +buffer's cast becomes a no-op: 285.0 us hidden with the cast, 425.7 without. + +**The third fix: do not hand the SMs back between our own two kernels.** With +the first two in, the timeline reads: + +``` +cast 1170.8 -> 1173.4 2.6 us +GEMM 1173.1 -> 2070.4 897.3 us g=2048 +notify 1299.3 -> 1325.9 26.6 us 126.2 us after the GEMM +scatter 1444.6 -> 2279.3 834.6 us 118.7 us after the notify +``` + +The scatter overruns the GEMM by 208.9 us and that overrun is all of what is +still unhidden. It is not that the scatter runs slowly while sharing the device +-- 834.6 us here against 844.1 solo, so it does not -- it is that it starts +271.5 us late. Of that, 126.2 us is the first admission, which DeepEP pays too, +and **118.7 us is the gap between our own two kernels**. Solo the same gap is +1.9 us: sixty-two times wider under contention, because the notify releases all +64 SMs when it exits and the GEMM has twelve waves queued to take them back. + +DeepEP does not pay that seam. Its epilogue is pre-admitted through programmatic +dependent launch, so the boundary between its two kernels measures *negative*, +-30.9 us. Tried that first, since TileLang exposes it: `T.pdl_sync()` on the +dependent kernels, verified in the generated CUDA and in the launch path, worth +**8.2 us**. Pre-admission lets the dependent grid be *placed* while the primary +is resident, but a GEMM CTA holds 213 KB of the SM's 228 KB of shared memory, so +there is nothing for a comm block to be placed beside, and the placement still +waits for a full SM to drain. + +Never giving the SMs up works instead, and that just means one kernel. The +kernel was one before it was split so the scatter could use a wider grid than +the notify; the split's own numbers put the fusion at 883.7 us against 881.7, +and the measurement above says the boundary costs 118.7 us of overlap. + +The wider scatter grid did not survive the fusion, and did not need to. Giving +the *whole* fused kernel the wider grid beats the split arrangement outright -- +861.0 us at `num_sms=128` against 867.1 for a 64-wide notify and a 128-wide +scatter -- so the knob that expressed it is gone and `num_sms` is the only one +left. Standalone is unchanged at 898.1 us bf16, 984.5 combine, 521.3 fp8. + +Standalone performance is unchanged -- 898-900 us bf16 dispatch, 984-985 +combine, 522 fp8, all within the spread of the numbers above -- so this costs +nothing when there is nothing to overlap with. + +Rejected along the way, each with numbers: SM starvation as a *count* problem +(8, 16 and 32 SMs all hid the same ~35 us, and at 8 SMs there are 140 free), +proportionality (99 us and 673 us of compute hid the same absolute amount), the +scatter and barrier spinning (a notify-only dispatch, 63 us and almost entirely +barriers, hid 28.2 us where the full 892 us one hid 29.7), the cooperative +launch (moving padding, stats and reset into a third kernel removes every +`sync_grid`, correct on 8 ranks, and changed nothing -- and note that the +motivation given in the commit for that test was itself wrong: EPv2's intranode +path *does* use `this_grid().sync()`, at `common/comm.cuh:233,251,269` from +`impls/dispatch.cuh:74,398`, so a cooperative launch was never the structural +difference it was claimed to be), grids beyond the SM +count (450 us at 128 blocks against 621 at 256), host run-ahead (both sides +queue thousands of microseconds ahead), and the stream dependency itself +(DeepEP's `stream_control_prologue` takes the same `wait_stream(comm, compute)` +when no `previous_event` is given, and overlaps anyway). + +Note that `async_finish` changes neither column. It moves who waits, not how +long the work takes: measured, the same dispatch is 901.2 µs synchronous and +901.4 µs asynchronous with the event waited on. Timing an asynchronous call +*without* waiting reports 3.0 µs, which is the launch and nothing else. + +## Not implemented + +DeepEP intranode features this port does not cover. None of them is blocked by +a design decision here -- each would be an additive parameter or output: + +| | | +|---|---| +| `kAllowMultipleReduction` | combine-side local sum across several experts on one rank. This port sums on the *sender* instead, DeepEP's `kDoExpandedSend`, which is what lets combine take an expanded dispatch -- so the capability is covered and the flag is not | +| `deterministic` mode | DeepEP has a separate prologue for it | +| `use_tma_aligned_col_major_sf` | column-major scale-factor layout for a downstream GEMM | + +Out of scope by design: RDMA/scaleout (`num_qps`), the low-latency decode path, +and Engram/PP/CP. + +Also deliberately absent: combine's `topk_weights` side output. DeepEP carries +the gate weights back through combine because its *expanded* layout reorders +them, so the source rank cannot reconstruct which weight went with which slot. +In this port's rank layout the source rank is the one that produced +`topk_weights` and still holds it in its original order, so the output would be +a collective that returns its own input. + +## Running + +```bash +python example_dispatch_combine_correctness.py --num-sms 64 +python example_dispatch_combine_benchmark.py --num-sms 64 +pytest test_example_deepep_v2.py +``` + +Both examples take `--tokens/--hidden/--topk/--experts` and the thread knobs. +Measurements need a quiet machine and warm clocks: an idle B200 sits at +1.3 GHz against a 1.965 GHz boost, and `do_bench` idles the GPU ~10ms per +iteration, so warm before every timed section rather than once per run. diff --git a/examples/distributed/deepep_v2/buffer.py b/examples/distributed/deepep_v2/buffer.py new file mode 100644 index 0000000000..60bfeb7d33 --- /dev/null +++ b/examples/distributed/deepep_v2/buffer.py @@ -0,0 +1,790 @@ +"""Buffer/EPHandle: DeepEP-EPv2-aligned API over the TileScale intranode port. + +Names and call shape mirror DeepEP EPv2's real ``ElasticBuffer``/``EPHandle`` +(see ``deep_ep/buffers/elastic.py`` in the DeepEP submodule) as closely as this +port's scope allows. Deliberately dropped: RDMA/hybrid mode, low-latency +decode path, Engram/PP/AGRS, expert-alignment and expand-layout (no fused +grouped GEMM in scope, so there's nothing to align or expand for), and +handle-cached "skip renotify" reuse (DeepEP's decode-replay optimization -- +a real simplification, not attempted here yet). + +What *is* aligned with DeepEP's real design (see ``kernels/dispatch.py`` and +``kernels/combine.py`` for the detailed mapping): per-(token, destination-rank) +dedup, GPU-side notify with a real cross-rank count exchange, local (not +remote) atomic slot claiming into a fixed per-sender receive slice sized +exactly `num_max_tokens_per_rank * num_ranks` (a hard capacity bound once +deduped, not a statistical headroom guess), and a combine that stores back +into unique per-(rank, token) slots and reduces locally instead of pushing +remote atomics. `dispatch_threads`/`combine_threads` control warps per block +independently; several warps per SM, like DeepEP itself, is what actually +uses NVLink's per-SM bandwidth. + +Dispatch writes straight into the compact output, so it needs no staging buffer +at all (see ``kernels/dispatch.py``); ``comm_x`` is combine's, holding one slot +per (contributing rank, source token) on the way back. +""" + +from collections import deque + +import torch +import torch.distributed as dist + +import tilelang +import tilelang.language as T +from tilelang.distributed.allocator import get_allocator + +from kernels.dispatch import dispatch_kernel +from kernels.combine import combine_kernel +from reference import packed_row_bytes + +_TL_DTYPES = { + torch.bfloat16: T.bfloat16, + torch.float16: T.float16, + torch.float32: T.float32, + torch.float8_e4m3fn: T.float8_e4m3fn, +} + +# FP8 is quantised per token over groups of this many elements, and carries one +# fp32 scale per group -- DeepEP's `per_token_cast_to_fp8` layout, which the +# caller is expected to produce. +FP8_GROUP = 128 + + +class EventOverlap: + """DeepEP's `EventOverlap`: a comm-stream event, plus what it must outlive. + + `dispatch` and `combine` return one whether or not they were asked to run + asynchronously; synchronously it wraps `None`, so `with event:` is a no-op + and callers do not branch. + + `extra_tensors` is DeepEP's mechanism and its reason is worth repeating: + the obvious way to keep a comm-stream tensor alive until the compute stream + is done with it is `Tensor.record_stream`, but that is incompatible with + CUDA graph capture. Holding a reference here instead ties the tensors' + lifetime to the event object, which the caller drops after waiting. + """ + + def __init__(self, event: torch.cuda.Event | None = None, extra_tensors: tuple = ()): + self.event = event + self.extra_tensors = extra_tensors + self._release_handle_by_call = False + + def current_stream_wait(self, release_handle: bool = False) -> None: + """Make the current stream wait for the comm kernels, without blocking the host.""" + assert self.event is not None, "no event: this call was not made with async_finish=True" + torch.cuda.current_stream().wait_event(self.event) + if release_handle: + self.event = None + self.extra_tensors = () + + def __call__(self, release_handle: bool = False) -> "EventOverlap": + self._release_handle_by_call = release_handle + return self + + def __enter__(self) -> "EventOverlap": + """Overlap whatever the block enqueues with the communication. + + ```python + recv_x, _, _, handle, event = buf.dispatch(x, topk_idx, topk_weights, async_finish=True) + with event: + unrelated_work_on_the_current_stream() + # leaving the block, the current stream waits: `recv_x` is now readable + ``` + """ + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + if self.event is not None: + self.current_stream_wait(release_handle=self._release_handle_by_call) + self._release_handle_by_call = False + + +class EPHandle: + """Communication handle returned by ``Buffer.dispatch``, consumed by ``Buffer.combine``. + + Attributes (named to match DeepEP's real ``EPHandle`` where a concept applies): + num_experts, num_max_tokens_per_rank, num_sms: as passed to dispatch. + topk_idx: this rank's own top-k expert indices, `[num_tokens, num_topk]`. + num_recv: total compacted tokens received by this rank, as a + one-element *device* tensor. `num_recv_tokens` reads it back to the + host, which synchronises -- call it outside a timed region. + finish_event: for `async_finish` dispatches, the event the caller must + wait on before reading anything the call returned; `None` otherwise. + psum_recv_count: inclusive prefix sum of deduplicated received counts per + sender rank, `[num_ranks]` -- DeepEP's `psum_num_recv_tokens_per_scaleup_rank`. + `None` in the expanded layout, which groups by expert rather than by + sender and so never computes it: use `expert_count`/`expert_offset`. + recv_src_rank, recv_src_token: per-compact-row source (rank, token) -- + a simplified, unpacked form of DeepEP's single encoded `recv_src_metadata`. + num_tokens: how many tokens this rank dispatched, i.e. how many rows + `combine` reduces back into. + """ + + def __init__( + self, + num_experts, + num_max_tokens_per_rank, + num_sms, + topk_idx, + num_recv, + psum_recv_count, + recv_src_rank, + recv_src_token, + num_tokens, + topk_weights=None, + layout_generation=-1, + finish_event=None, + expert_count=None, + expert_offset=None, + expand_overflow=None, + ): + self.num_experts = num_experts + self.num_max_tokens_per_rank = num_max_tokens_per_rank + self.num_sms = num_sms + self.topk_idx = topk_idx + self.num_recv = num_recv + self.psum_recv_count = psum_recv_count + self.recv_src_rank = recv_src_rank + self.recv_src_token = recv_src_token + self.num_tokens = num_tokens + # What a cached dispatch replays. `topk_idx`/`topk_weights` are this + # call's converted copies, so reusing them reproduces the previous + # dispatch's metadata exactly rather than approximately. + self.topk_weights = topk_weights + # Which layout the buffer held when this handle was made. A dispatch + # for different routing overwrites `send_base`/`send_rank_mask` in + # place, so an older handle is not merely stale, it is wrong -- and + # silently so. See `Buffer.dispatch`. + self.layout_generation = layout_generation + # Set only for `async_finish` dispatches: nothing the call returned may + # be read until this is waited on. + self.finish_event = finish_event + # Expanded layout only, `None` otherwise: local expert `e` owns rows + # `[expert_offset[e], expert_offset[e + 1])` of the received tensor, of + # which the first `expert_count[e]` are real and the rest is alignment + # padding, zeroed unless `zero_padding=False`. See kernels/dispatch.py. + self.expert_count = expert_count + self.expert_offset = expert_offset + self._expand_overflow = expand_overflow + self._num_recv_tokens = None + + @property + def expand_overflow(self) -> int: + """Rows the expanded layout needed beyond capacity, 0 if it fit. + + Non-zero means dispatch *skipped* this rank rather than writing past + its buffer, so everything the call returned is meaningless: re-create + the `Buffer` with `expand_factor` at least this over `recv_capacity`. + Reads back to the host, like `num_recv_tokens`. + """ + if self._expand_overflow is None: + return 0 + if self.finish_event is not None: + self.finish_event.synchronize() + return int(self._expand_overflow[0].item()) + + @property + def num_recv_tokens(self) -> int: + """The received row count, on the host. Synchronises on first read. + + `dispatch` deliberately does not do this itself: a device-to-host read + of `num_recv` cost ~33us of its ~970us, and nothing in the pipeline + needs the value -- `combine` takes the device tensor straight through to + its kernel. Callers that genuinely want the number (a benchmark's byte + count, an assertion, slicing for a reference implementation) pay for it + here, once, and outside whatever they are timing. + """ + if self._num_recv_tokens is None: + if self.finish_event is not None: + self.finish_event.synchronize() + self._num_recv_tokens = int(self.num_recv[0].item()) + return self._num_recv_tokens + + +class Buffer: + """Intranode (NVLink-only) MoE dispatch/combine buffer, DeepEP-EPv2-style.""" + + def __init__( + self, + group: dist.ProcessGroup, + local_rank: int, + num_local_ranks: int, + num_max_tokens_per_rank: int, + hidden: int, + num_topk: int, + num_experts: int, + dtype: torch.dtype = torch.bfloat16, + num_sms: int = 0, + dispatch_threads: int = 1024, + combine_threads: int = 1024, + reduce_threads: int = 256, + pipeline_depth: int = 2, + comm_stream_priority: int = None, + do_expand: bool = False, + expert_alignment: int = 1, + zero_padding: bool = True, + expand_factor: float = 1.0, + ): + self.group = group + self.rank_idx = local_rank + self.num_ranks = num_local_ranks + self.num_max_tokens_per_rank = num_max_tokens_per_rank + self.hidden = hidden + self.num_topk = num_topk + self.num_experts = num_experts + assert num_experts % num_local_ranks == 0, f"num_experts={num_experts} is not a multiple of {num_local_ranks} ranks" + self.experts_per_rank = num_experts // num_local_ranks + # `dtype` is the *dispatch payload* type. Combine always moves bf16: + # what it carries is the expert output, which is not quantised. + self.dtype = dtype + self.tl_dtype = _TL_DTYPES[dtype] + self.is_fp8 = dtype == torch.float8_e4m3fn + if self.is_fp8: + assert hidden % FP8_GROUP == 0, f"hidden={hidden} is not a multiple of {FP8_GROUP}" + self.scale_dim = hidden // FP8_GROUP if self.is_fp8 else 0 + # The row dispatch actually moves: for FP8, payload bytes followed by + # the per-group fp32 scale packed right after -- so the scatter needs + # one remote store per token-destination pair instead of two -- padded + # to `put_warp`'s preferred boundary. `reference.packed_row_bytes` owns + # the formula and the reasoning; kernels/dispatch.py mirrors it. + self.row_bytes = packed_row_bytes(hidden, FP8_GROUP) if self.is_fp8 else None + self.combine_dtype = torch.bfloat16 if self.is_fp8 else dtype + self.tl_combine_dtype = _TL_DTYPES[self.combine_dtype] + # Wide blocks at both ends of the SM range: at 64 SMs 1024 threads + # measured 688.0 GB/s against 680.8 for 512 (four readings each, back + # to back), and at 24 SMs the two are within noise. Fewer, wider blocks + # never lost here, so there is no SM-dependent rule to write. + self.dispatch_threads = dispatch_threads + self.combine_threads = combine_threads + # Separate from `combine_threads` because the reduce wants + # `hidden / reduce_threads` to be a whole number of 128-bit loads: at + # hidden=7168 that is 28 per thread at 256 against an awkward 14 at 512, + # worth 633.7 GB/s against 625.9. + self.reduce_threads = reduce_threads + + from tilelang.carver.arch import driver + + device_sms = driver.get_num_sms() + self.device_sms = device_sms + if num_sms == 0: + num_sms = device_sms + # Dispatch is a single persistent grid whose phases rendezvous through + # global counters, so every block has to be resident simultaneously. + # The kernel is compiled with `__launch_bounds__(threads, 1)`, which + # guarantees one block per SM fits, but nothing guarantees a grid larger + # than the device: that would spin forever instead of failing. + if num_sms > device_sms: + raise ValueError( + f"num_sms={num_sms} exceeds the device's {device_sms} SMs; dispatch's grid-wide " + "rendezvous requires every block to be co-resident and would deadlock" + ) + self.num_sms = num_sms + + self.cap = num_max_tokens_per_rank + self.total_capacity = self.cap * self.num_ranks + + # DeepEP's `do_expand`: one received row per (token, expert), grouped + # by local expert, which is the layout a grouped GEMM wants. See + # kernels/dispatch.py for how the sender computes the index and + # `reference.expanded_layout` for what the result should look like. + self.do_expand = do_expand + self.expert_alignment = expert_alignment if do_expand else 1 + self.zero_padding = zero_padding + # Deduplicated, `total_capacity` cannot be exceeded. Expanded it can: + # the hard bound is `min(num_topk, experts_per_rank)` times higher, for + # routing that puts 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 `expand_factor=1`, so that is the default and dispatch + # raises `expand_overflow` rather than corrupting memory if a call + # exceeds it. Raise the factor (up to the hard bound) for skewed + # routing. + self.expand_factor = expand_factor + if do_expand: + hard_bound = self.total_capacity * min(num_topk, self.experts_per_rank) + # Plus the alignment padding, which is not routing-dependent and so + # sits outside the bound rather than inside it. + aligned_slack = self.experts_per_rank * (self.expert_alignment - 1) + self.recv_capacity = min(int(self.total_capacity * expand_factor), hard_bound) + aligned_slack + else: + self.recv_capacity = self.total_capacity + + itemsize = torch.empty((), dtype=dtype).element_size() + comm_bytes = self.num_ranks * self.cap * hidden * 2 # combine is always bf16 + row_bytes = self.row_bytes if self.is_fp8 else hidden * itemsize + compact_bytes = self.recv_capacity * (row_bytes + 4 + 4 + num_topk * 4 + num_topk * 4) + combined_bytes = num_max_tokens_per_rank * (hidden * 2 + 4) + self.num_ranks * self.cap * 4 + if do_expand: + # Combine's grouping buffers and per-warp scratch; see below. + max_rows = min(num_topk, self.experts_per_rank) + combine_warps = (self.num_sms or 1) * 32 + combined_bytes += self.num_ranks * self.cap * 4 * (1 + max_rows) + combine_warps * hidden * 2 + total = comm_bytes + compact_bytes + combined_bytes + self.allocator = get_allocator( + # 5% of slack: every tensor is padded for alignment and the sum + # above does not model that, so a fixed margin does not scale. + size=max(int(total * 1.05) + (1 << 20), 1 << 24), + device=f"cuda:{local_rank}", + is_distributed=True, + local_rank=local_rank, + num_local_ranks=num_local_ranks, + group=group, + ) + + # One `num_ranks`-wide slot per barrier site: dispatch entry/exit, + # combine entry/exit. They must not share -- `tl::barrier_blocks` + # settles by having each peer's +TAG cancelled by its -TAG, so a second + # site on the same slot pushes a still-waiting block back above zero. + # That converges rather than deadlocks, which is why sharing cost 20x + # instead of hanging. DeepEP separates its tags for the same reason. + self.barrier = tilelang.tensor((4 * self.num_ranks,), torch.int32, allocator=self.allocator) + # uint32: atom_add's PTX intrinsic requires an unsigned target. + n_dst = num_experts if do_expand else self.num_ranks + self.send_count = tilelang.tensor((n_dst,), torch.uint32, allocator=self.allocator) + # Grid-wide rendezvous counters for dispatch's three phases. + self.notify_done = tilelang.tensor((1,), torch.uint32, allocator=self.allocator) + self.exchange_done = tilelang.tensor((1,), torch.uint32, allocator=self.allocator) + self.slot_counter = tilelang.tensor((n_dst,), torch.uint32, allocator=self.allocator) + # Stand-in for the degenerate `recv_expert_stats` argument when the + # caller did not ask for per-expert stats; see `dispatch`. + self._no_expert_stats = tilelang.tensor((1,), torch.uint32, allocator=self.allocator) + # Expanded-layout outputs. One element each unless expanding, matching + # the kernel's degenerate argument shapes. + self.expert_count = tilelang.tensor((self.experts_per_rank if self.do_expand else 1,), torch.int32, allocator=self.allocator) + self.expert_offset = tilelang.tensor((self.experts_per_rank + 1 if self.do_expand else 1,), torch.int32, allocator=self.allocator) + self.expand_overflow = tilelang.tensor((1,), torch.int32, allocator=self.allocator) + # Combine's (src_rank, src_token) -> rows inversion, and one scratch row + # per warp for tokens whose rows have to be summed before sending. Only + # the expanded layout can produce more than one row per pair, so these + # degenerate otherwise. uint32 for `atom_add`, as with `send_count`. + self.max_rows_per_token = min(num_topk, self.experts_per_rank) if do_expand else 1 + n_pairs = self.num_ranks * self.cap if do_expand else 1 + combine_warps = self.num_sms * (self.combine_threads // 32) + self.group_count = tilelang.tensor((n_pairs,), torch.uint32, allocator=self.allocator) + self.group_rows = tilelang.tensor((n_pairs * self.max_rows_per_token,), torch.int32, allocator=self.allocator) + self.reduce_scratch = tilelang.tensor((combine_warps * hidden if do_expand else 1,), self.combine_dtype, allocator=self.allocator) + # Likewise for combine's unused bias arguments; see `combine`. + self._no_bias = tilelang.tensor((1, hidden), self.combine_dtype, allocator=self.allocator) + # int32 (signed): -1 is the "not yet published" sentinel every rank + # spins on while the count matrix fills in. + self.count_matrix = tilelang.tensor((self.num_ranks * n_dst,), torch.int32, allocator=self.allocator) + self.send_base = tilelang.tensor((n_dst,), torch.int32, allocator=self.allocator) + self.psum_recv_count = tilelang.tensor((self.num_ranks,), torch.int32, allocator=self.allocator) + self.num_recv = tilelang.tensor((1,), torch.int32, allocator=self.allocator) + self.send_rank_mask = tilelang.tensor((num_max_tokens_per_rank,), torch.int32, allocator=self.allocator) + + # FP8: raw `row_bytes` per slot (payload followed by scale, packed -- + # see `row_bytes` above), opaque to the caller until unpacked with + # `reference.per_token_cast_back`. BF16: `hidden` elements of `dtype`, + # unchanged. + self.recv_x = tilelang.tensor( + (self.recv_capacity, self.row_bytes if self.is_fp8 else hidden), + torch.uint8 if self.is_fp8 else dtype, + allocator=self.allocator, + ) + self.recv_x_flat = self.recv_x.view(-1) + self.recv_src_rank = tilelang.tensor((self.recv_capacity,), torch.int32, allocator=self.allocator) + self.recv_src_token = tilelang.tensor((self.recv_capacity,), torch.int32, allocator=self.allocator) + self.recv_topk_idx = tilelang.tensor((self.recv_capacity, num_topk), torch.int32, allocator=self.allocator) + self.recv_topk_weights = tilelang.tensor((self.recv_capacity, num_topk), torch.float32, allocator=self.allocator) + + # Combine's staging buffer, one slot per (contributing rank, source + # token) -- the equivalent of DeepEP's `recv_buffer` on the way back. + self.comm_x = tilelang.tensor((self.num_ranks * self.cap * hidden,), self.combine_dtype, allocator=self.allocator) + + self.combined = tilelang.tensor((num_max_tokens_per_rank, hidden), self.combine_dtype, allocator=self.allocator) + + # `combine` only ever reads slots dispatch actually wrote, but zeroing + # once keeps a first-use read of never-written memory from producing + # NaNs in the (unused) tail of the compact output. + self.recv_x.zero_() + self.send_rank_mask.zero_() + # The kernel resets these at the end of every call; this is only about + # the first call finding them defined. `barrier` is the exception -- + # it returns to zero by itself, and re-zeroing it from the host is + # exactly what breaks a peer still inside it. + self.barrier.zero_() + self.send_count.zero_() + self.notify_done.zero_() + self.exchange_done.zero_() + self.slot_counter.zero_() + self.count_matrix.fill_(-1) + + # Communication runs on its own stream, DeepEP's `comm_stream`. On the + # caller's stream, everything this class does -- the topk conversions, + # the launch, the handle bookkeeping -- sits between GPU kernels with + # the GPU idle through it. On a private stream it overlaps with whatever + # the caller already has queued instead. + # + # The priority is what makes `async_finish` worth anything, and it is + # not a tuning knob but the difference between overlapping and not. A + # private stream buys eligibility, not admission: a GEMM large enough to + # be worth hiding behind is also large enough to hold every SM for many + # waves, and each SM that frees goes to the next of its blocks, so a + # collective that becomes eligible mid-GEMM waits for the whole thing. + # Measured at the V3 shape against an 8192-square bf16 GEMM: 29.4us of + # 674 hidden at the default priority, 285.3 at any raised one. The + # actual value does not matter -- -1, -2 and -3 measure the same -- so + # take whatever the device offers and let a caller who is scheduling + # several streams override it. + if comm_stream_priority is None: + comm_stream_priority = torch.cuda.Stream.priority_range()[1] + self.comm_stream_priority = comm_stream_priority + self.comm_stream = torch.cuda.Stream(priority=comm_stream_priority) + + # How far the CPU may run ahead, in calls -- not about ordering + # (`wait_stream` handles that) but about skew between ranks: the kernel + # ends in a cross-rank barrier, so a rank queued N calls behind stalls + # everyone inside it. 1, 2 and 4 all measure ~680-685 GB/s; 2 had no + # low outlier. + # 0 disables the throttle, which is what an `async_finish` caller + # driving its own overlap usually wants: the `synchronize` it does + # blocks the *host*, so a bounded run-ahead and a fully asynchronous + # call are not the same thing. + self.pipeline_depth = pipeline_depth + self._in_flight: deque = deque() + + # Bumped by every dispatch that recomputes a layout; a handle stays + # reusable only while it still matches. + self._layout_generation = 0 + self._dispatch_kernels = {} + self._combine_kernels = {} + + def _resolve_num_sms(self, num_sms: int) -> int: + """Per-call grid size, validated. Both kernels rendezvous grid-wide, so + every block has to be resident and the count cannot exceed the device.""" + num_sms = self.num_sms if num_sms == 0 else num_sms + assert 0 < num_sms <= self.device_sms, f"num_sms={num_sms} exceeds the device's {self.device_sms} SMs" + return num_sms + + def _get_dispatch_kernel(self, num_tokens: int, collect_expert_stats: bool = False, num_sms: int = 0, cached: bool = False): + num_sms = self._resolve_num_sms(num_sms) + key = (num_tokens, collect_expert_stats, num_sms, cached) + if key not in self._dispatch_kernels: + kernel = dispatch_kernel( + num_tokens, + self.num_ranks, + self.num_experts, + self.num_topk, + self.hidden, + self.num_max_tokens_per_rank, + num_sms, + self.dispatch_threads, + self.tl_dtype, + self.scale_dim, + self.row_bytes or 0, + collect_expert_stats, + self.do_expand, + self.expert_alignment, + self.zero_padding, + self.recv_capacity, + cached, + ) + kernel.compile_group = self.group + kernel.initialize(allocator=self.allocator) + self._dispatch_kernels[key] = kernel + return self._dispatch_kernels[key] + + def _get_combine_kernel(self, num_tokens: int, num_bias: int = 0, num_sms: int = 0): + num_sms = self._resolve_num_sms(num_sms) + key = (num_tokens, num_bias, num_sms, self.do_expand) + if key not in self._combine_kernels: + kernel = combine_kernel( + num_tokens, + self.num_ranks, + self.hidden, + self.num_max_tokens_per_rank, + self.total_capacity, + num_sms, + self.combine_threads, + self.reduce_threads, + self.tl_combine_dtype, + num_bias, + self.do_expand, + self.max_rows_per_token, + self.recv_capacity, + ) + kernel.compile_group = self.group + kernel.initialize(allocator=self.allocator) + self._combine_kernels[key] = kernel + return self._combine_kernels[key] + + def _begin_comm(self, previous_event, allocate_on_comm_stream): + """Join the communication stream to what this call must follow. + + DeepEP's `stream_control_prologue`. Anything the caller has to convert + or allocate belongs *before* this rather than inside the stream block it + opens: the first operation queued on the communication stream takes the + one block admission per iteration that a resident GEMM is not already + holding, and it should be the collective rather than a 3 us cast. See + kernels/dispatch.py on overlap. + """ + compute_stream = torch.cuda.current_stream() + if previous_event is not None: + assert allocate_on_comm_stream, "previous_event requires allocate_on_comm_stream" + self.comm_stream.wait_event(previous_event.event) + else: + self.comm_stream.wait_stream(compute_stream) + return compute_stream + + def _end_comm(self, compute_stream, temporaries, async_finish, allocate_on_comm_stream): + """Record the finish event, and keep the temporaries alive across the + handover: they were allocated on the communication stream and the caller + reads them on its own. + + `allocate_on_comm_stream` carries them on the event instead of calling + `record_stream`, which CUDA graph capture does not permit. + """ + finish_event = torch.cuda.Event() + finish_event.record(self.comm_stream) + if not async_finish: + compute_stream.wait_stream(self.comm_stream) + elif not allocate_on_comm_stream: + for t in temporaries: + t.record_stream(compute_stream) + event = EventOverlap( + finish_event if async_finish else None, + temporaries if async_finish and allocate_on_comm_stream else (), + ) + return finish_event, event + + def dispatch( + self, + x, + topk_idx: torch.Tensor = None, + topk_weights: torch.Tensor = None, + handle: "EPHandle" = None, + num_sms: int = 0, + previous_event: EventOverlap = None, + async_finish: bool = False, + allocate_on_comm_stream: bool = False, + cumulative_local_expert_recv_stats: torch.Tensor = None, + ): + """Scatter `x` to the ranks owning each token's top-k experts. + + `x` is the packed `(values, scale)` buffer `reference.per_token_cast_to_fp8` + produces when the buffer's dtype is FP8 (payload bytes followed by the + per-group fp32 scale, so the scatter moves both in one remote store per + token-destination pair -- see kernels/dispatch.py), and a plain tensor + otherwise. The first return value mirrors that. + + Returns `(recv_x, recv_topk_idx, recv_topk_weights, handle, event)`. + The `event` is an `EventOverlap`, as in DeepEP, and is returned either + way -- synchronously it wraps `None`, so callers need not branch. + + With `async_finish` the caller's stream is *not* joined to the + communication stream and nothing the call returns may be read until the + event is waited on (`with event:` or `event.current_stream_wait()`). + EPv2 spells this argument `async_with_compute_stream`; the name here + matches its `combine` and DeepEP's own legacy buffer. + + `previous_event` starts the communication after one specific event + rather than after everything queued on the caller's stream, and + `allocate_on_comm_stream` leaves this call's temporaries owned by the + communication stream -- keeping them alive through the returned event + instead of `record_stream`, which CUDA graph capture does not allow. As + in DeepEP the first requires the second. + + `cumulative_local_expert_recv_stats` is DeepEP's load-balance counter: + a `[num_experts // num_ranks]` uint32 tensor this rank's received token + count per local expert is *added into*. The caller owns it and decides + when to zero it, so it can accumulate over a step, a batch or a whole + run. Passing it compiles a separate kernel variant -- the tally is + absent from the default one, not merely skipped. + """ + if handle is not None: + # DeepEP's contract: passing a handle *is* the assertion that the + # routing is unchanged, so passing routing too would be ambiguous. + assert topk_idx is None and topk_weights is None, "pass either a handle or topk_idx/topk_weights, not both" + assert handle.layout_generation == self._layout_generation, ( + f"handle holds layout {handle.layout_generation}, buffer is on {self._layout_generation}: " + "another dispatch has overwritten it, so the cached layout no longer describes this buffer" + ) + topk_idx, topk_weights = handle.topk_idx, handle.topk_weights + else: + assert topk_idx is not None and topk_weights is not None, "dispatch needs topk_idx/topk_weights or a handle" + cached = handle is not None + collect_expert_stats = cumulative_local_expert_recv_stats is not None + if collect_expert_stats: + stats = cumulative_local_expert_recv_stats + assert stats.dtype == torch.uint32 and stats.shape == (self.experts_per_rank,), ( + f"expected a ({self.experts_per_rank},) uint32 tensor, got {tuple(stats.shape)} {stats.dtype}" + ) + else: + stats = self._no_expert_stats + if self.is_fp8: + assert x.dtype == torch.uint8 and x.shape[1] == self.row_bytes, ( + f"expected a packed (*, {self.row_bytes}) uint8 buffer from reference.per_token_cast_to_fp8, got {tuple(x.shape)} {x.dtype}" + ) + num_tokens = x.shape[0] + num_sms = self._resolve_num_sms(num_sms) + # Converted on the caller's stream, before the communication stream is + # joined to it: a 3 us cast queued first on the communication stream + # takes the admission the 850 us collective behind it then has to wait + # for, which measured 141 us of hidden time. See `_begin_comm`. + topk_idx_i32 = topk_idx if cached else topk_idx.to(torch.int32).contiguous() + topk_weights_f32 = topk_weights if cached else topk_weights.to(torch.float32).contiguous() + compute_stream = self._begin_comm(previous_event, allocate_on_comm_stream) + + with torch.cuda.stream(self.comm_stream): + # Nothing is reset here: the kernel does it at the end of every + # call, which measured 633 GB/s against 626 for six host-side + # `zero_()` launches -- consistently ahead across four interleaved + # rounds. It needs a `T.sync_grid()` to be correct; see + # kernels/dispatch.py. + + kernel = self._get_dispatch_kernel(num_tokens, collect_expert_stats, num_sms, cached) + # No `dist.barrier` on either side. The reset above is peer-visible + # state, so it does need ordering against peers -- but the kernel's + # own entry `barrier_blocks` provides it, which is why no collective + # is needed here. See kernels/dispatch.py. + kernel( + x, + topk_idx_i32, + topk_weights_f32, + self.notify_done, + self.exchange_done, + self.send_count, + self.count_matrix, + self.send_base, + self.psum_recv_count, + self.num_recv, + self.slot_counter, + self.send_rank_mask[:num_tokens], + self.barrier, + self.recv_x_flat, + self.recv_src_rank, + self.recv_src_token, + self.recv_topk_idx.view(-1), + self.recv_topk_weights.view(-1), + stats, + self.expert_count, + self.expert_offset, + self.expand_overflow, + ) + + # No device-to-host read of `num_recv`: it cost ~33us and nothing + # here needs the value. Overflow cannot happen by construction -- + # every peer sends at most `num_max_tokens_per_rank` rows and + # capacity is `num_ranks` times that. + num_recv = self.num_recv.clone() + psum_recv_count = self.psum_recv_count.clone() + # Only the expanded layout produces these, so deduplicated there is + # nothing to snapshot. + if self.do_expand: + expert_count = self.expert_count.clone() + expert_offset = self.expert_offset.clone() + expand_overflow = self.expand_overflow.clone() + else: + expert_count = expert_offset = expand_overflow = None + + temporaries = tuple( + t + for t in (topk_idx_i32, topk_weights_f32, num_recv, psum_recv_count, expert_count, expert_offset, expand_overflow) + if t is not None + ) + finish_event, event = self._end_comm(compute_stream, temporaries, async_finish, allocate_on_comm_stream) + + # Keep the CPU from running arbitrarily far ahead. See `pipeline_depth`. + if self.pipeline_depth: + self._in_flight.append(finish_event) + while len(self._in_flight) > self.pipeline_depth: + self._in_flight.popleft().synchronize() + + if not cached: + self._layout_generation += 1 + handle = EPHandle( + self.num_experts, + self.num_max_tokens_per_rank, + num_sms, + topk_idx_i32, + num_recv, + None if self.do_expand else psum_recv_count, + self.recv_src_rank, + self.recv_src_token, + num_tokens, + topk_weights_f32, + self._layout_generation, + finish_event if async_finish else None, + expert_count, + expert_offset, + expand_overflow, + ) + # FP8: the packed uint8 buffer, unpacked with `reference.per_token_cast_back`. + return (self.recv_x, self.recv_topk_idx, self.recv_topk_weights, handle, event) + + def combine( + self, + x: torch.Tensor, + handle: EPHandle, + num_sms: int = 0, + bias=None, + previous_event: EventOverlap = None, + async_finish: bool = False, + allocate_on_comm_stream: bool = False, + ): + """Reduce every rank's contribution back into this rank's token order. + + Returns `(combined, event)`. DeepEP's third element, + `combined_topk_weights`, has no counterpart here -- see the README on + why carrying the gate weights back would return the caller its own + input. `previous_event`, `async_finish` and `allocate_on_comm_stream` + mean exactly what they do on `dispatch`. + + `bias` is DeepEP's `bias_0`/`bias_1`: `None`, one `[num_tokens, hidden]` + tensor, or a pair of them, added to the output. As in DeepEP they are + added once per output token rather than once per contribution, and a + token with no contributions still receives them. Each distinct count + compiles its own kernel variant. + """ + if bias is None: + biases = () + elif torch.is_tensor(bias): + biases = (bias,) + else: + biases = tuple(bias) + assert len(biases) <= 2, f"DeepEP takes at most two bias tensors, got {len(biases)}" + for b in biases: + assert b.shape == (handle.num_tokens, self.hidden) and b.dtype == self.combine_dtype, ( + f"expected a ({handle.num_tokens}, {self.hidden}) {self.combine_dtype} bias, got {tuple(b.shape)} {b.dtype}" + ) + # The kernel always takes both arguments; unused ones are one-row + # stand-ins nothing in the generated code reads. See kernels/combine.py. + bias_args = tuple(biases) + tuple(self._no_bias for _ in range(2 - len(biases))) + + # Read the caller's contribution (see kernels/combine.py) in place -- + # `combine_kernel` takes its length as a symbolic extent, so no copy + # into a fixed-shape buffer is needed. + x_contig = x if x.is_contiguous() else x.contiguous() + x_flat = x_contig.reshape(-1) + num_sms = self._resolve_num_sms(num_sms or handle.num_sms) + num_tokens = handle.num_tokens + compute_stream = self._begin_comm(previous_event, allocate_on_comm_stream) + with torch.cuda.stream(self.comm_stream): + kernel = self._get_combine_kernel(num_tokens, len(biases), num_sms) + kernel( + x_flat, + self.recv_src_rank, + self.recv_src_token, + handle.num_recv, + self.send_rank_mask[:num_tokens], + self.barrier, + self.comm_x, + self.group_count, + self.group_rows, + self.reduce_scratch, + *bias_args, + self.combined[:num_tokens], + ) + # Only a `contiguous()` copy belongs to this call; a reshape view of + # the caller's own tensor does not, and the output is buffer-owned and + # outlives any single call. + temporaries = () if x_contig is x else (x_contig,) + # No pipeline bound here, unlike `dispatch`, and the asymmetry is not + # understood: bounding dispatch is worth 634 -> 530 -> 460 GB/s over + # successive runs, bounding this one costs 626 -> 610. Measured, not + # reasoned; re-measure over several runs before changing either. + _, event = self._end_comm(compute_stream, temporaries, async_finish, allocate_on_comm_stream) + return self.combined[:num_tokens], event + + def close(self): + self.allocator.close() diff --git a/examples/distributed/deepep_v2/example_dispatch_combine_benchmark.py b/examples/distributed/deepep_v2/example_dispatch_combine_benchmark.py new file mode 100644 index 0000000000..0936dd420b --- /dev/null +++ b/examples/distributed/deepep_v2/example_dispatch_combine_benchmark.py @@ -0,0 +1,191 @@ +"""Benchmark: dispatch/combine GB/s for the DeepEP-EPv2-aligned port. + + python example_dispatch_combine_benchmark.py --num-sms 24 + python example_dispatch_combine_benchmark.py --num-sms 64 + +**Warm the clocks or the numbers are noise.** These GPUs idle at 120 MHz against +a 1965 MHz boost ceiling, and `do_bench`'s `warmup` is a count of *iterations*, +not milliseconds -- the default handful of ~1ms iterations is nowhere near +enough to get off the idle clock, and each rep's `torch.cuda._sleep` spin draws +little power, so the clock can sag again mid-measurement. Runs taken without +`--clock-warmup-sec` varied by up to 1.5x on identical binaries and configs, +with dispatch and combine moving together (the signature of a clock effect, not +a kernel one). `_warm_clocks` runs real bf16 GEMMs on every rank first; locking +the clock outright (`nvidia-smi -lgc`) would be better but needs root. +""" + +import argparse +import time + +import torch +import torch.distributed as dist +import torch.multiprocessing + +from tilelang.distributed.host import init_dist +from tilelang.distributed.bench import do_bench + +from buffer import Buffer +import reference + + +def _warm_clocks(seconds: float, device: str) -> None: + """Drive the SMs hard enough, for long enough, to reach boost clocks.""" + if seconds <= 0: + return + a = torch.randn(8192, 8192, dtype=torch.bfloat16, device=device) + b = torch.randn(8192, 8192, dtype=torch.bfloat16, device=device) + c = torch.empty(8192, 8192, dtype=torch.bfloat16, device=device) + deadline = time.time() + seconds + while time.time() < deadline: + for _ in range(20): + torch.matmul(a, b, out=c) + torch.cuda.synchronize() + + +class _ClockProbe: + """Sample this rank's SM clock in a thread, so a measurement can report the + clock it actually ran at. + + Without this there is no way to tell a kernel regression from the GPU + having been at 120 MHz for half the run. + """ + + def __init__(self, device_index: int, period: float = 0.02): + self.period, self.samples, self._stop = period, [], None + try: + import pynvml + + pynvml.nvmlInit() + self._nvml, self._handle = pynvml, pynvml.nvmlDeviceGetHandleByIndex(device_index) + except Exception: + self._nvml = None + + def __enter__(self): + if self._nvml is None: + return self + import threading + + self._stop = threading.Event() + + def poll(): + while not self._stop.wait(self.period): + self.samples.append(self._nvml.nvmlDeviceGetClockInfo(self._handle, self._nvml.NVML_CLOCK_SM)) + + self._thread = threading.Thread(target=poll, daemon=True) + self._thread.start() + return self + + def __exit__(self, *exc): + if self._stop is not None: + self._stop.set() + self._thread.join(timeout=1.0) + return False + + def summary(self) -> str: + if not self.samples: + return "clock n/a" + s = sorted(self.samples) + return f"SM clock min/median/max {s[0]}/{s[len(s) // 2]}/{s[-1]} MHz over {len(s)} samples" + + +def main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank, num_ranks, group = init_dist(local_rank, num_local_ranks) + + torch.manual_seed(1234 + rank) + device = f"cuda:{local_rank}" + x = torch.randn(args.tokens, args.hidden, dtype=torch.bfloat16, device=device) + topk_idx, topk_weights = reference.make_topk(args.tokens, args.topk, args.experts, device, args.masked_ratio) + + dtype = torch.float8_e4m3fn if args.fp8 else torch.bfloat16 + # Quantising is the caller's job (see buffer.py's `dispatch` docstring); + # only the dispatch call itself sees fp8, everything downstream of the + # cast-back (expert compute, combine) stays bf16. + dispatch_x = reference.per_token_cast_to_fp8(x) if args.fp8 else x + + buf = Buffer( + group=group, + local_rank=local_rank, + num_local_ranks=num_ranks, + num_max_tokens_per_rank=args.tokens, + hidden=args.hidden, + num_topk=args.topk, + num_experts=args.experts, + dtype=dtype, + num_sms=args.num_sms, + dispatch_threads=args.dispatch_threads, + combine_threads=args.combine_threads, + ) + + itemsize = 1 if args.fp8 else 2 # fp8 payload byte, not counting the small per-128 scale + recv, recv_topk_idx, recv_topk_weights, handle, _ = buf.dispatch(dispatch_x, topk_idx, topk_weights) + # Outside every timed region: this is the one host read of the count. + num_recv_tokens = handle.num_recv_tokens + recv_topk_idx = recv_topk_idx[:num_recv_tokens] + recv_topk_weights = recv_topk_weights[:num_recv_tokens] + recv_x = reference.per_token_cast_back(recv[:num_recv_tokens], args.hidden) if args.fp8 else recv[:num_recv_tokens] + dispatch_bytes = num_recv_tokens * args.hidden * itemsize + # Combine always moves bf16 (see buffer.py) regardless of dispatch's dtype. + combine_bytes = num_recv_tokens * args.hidden * 2 + + expert_stats = torch.zeros(args.experts // num_local_ranks, dtype=torch.uint32, device=device) if args.expert_stats else None + + def run_dispatch(): + buf.dispatch(dispatch_x, topk_idx, topk_weights, cumulative_local_expert_recv_stats=expert_stats) + + _warm_clocks(args.clock_warmup_sec, device) + dist.barrier(group) + with _ClockProbe(local_rank) as probe: + dispatch_ms = do_bench(run_dispatch, warmup=args.warmup, rep=args.rep, group=group) + if rank == 0: + print( + f"dispatch: {dispatch_ms * 1000:.1f} us, {dispatch_bytes / (dispatch_ms * 1e-3) / 1e9:.1f} GB/s (recv-side, this rank) [{probe.summary()}]" + ) + + expert_out = reference.simulate_expert_compute(recv_x, recv_topk_idx, recv_topk_weights) + + def run_combine(): + buf.combine(expert_out, handle) + + _warm_clocks(args.clock_warmup_sec, device) + dist.barrier(group) + with _ClockProbe(local_rank) as probe: + combine_ms = do_bench(run_combine, warmup=args.warmup, rep=args.rep, group=group) + if rank == 0: + print( + f"combine: {combine_ms * 1000:.1f} us, {combine_bytes / (combine_ms * 1e-3) / 1e9:.1f} GB/s (send-side, this rank) [{probe.summary()}]" + ) + + buf.close() + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-processes", type=int, default=8) + # Fraction of top-k selections marked unselected (-1), DeepEP's marker. + parser.add_argument("--masked-ratio", type=float, default=0.0) + # Dispatch payload dtype; combine is always bf16 (see buffer.py). + parser.add_argument("--fp8", action="store_true") + # Accumulate DeepEP's per-local-expert receive counts during dispatch. + parser.add_argument("--expert-stats", action="store_true") + parser.add_argument("--tokens", type=int, default=8192) + parser.add_argument("--hidden", type=int, default=7168) + parser.add_argument("--topk", type=int, default=8) + parser.add_argument("--experts", type=int, default=256) + parser.add_argument("--num-sms", type=int, default=64) + # Neither collective stages rows through shared memory, so warps per block + # is a pure occupancy knob rather than something bounded by a shared-memory + # budget, and wide wins: against 512/256 these are worth 0.6% on dispatch + # and 3.0% on combine. Same as `Buffer`'s own defaults, deliberately -- when + # they drifted apart every number here described a configuration the library + # does not use. + parser.add_argument("--dispatch-threads", type=int, default=1024) + parser.add_argument("--combine-threads", type=int, default=1024) + # Iteration counts, not milliseconds. + parser.add_argument("--warmup", type=int, default=50) + parser.add_argument("--rep", type=int, default=50) + # Seconds of real GEMM load before each timed section -- see the module + # docstring. Set to 0 only if the clock is externally locked. + parser.add_argument("--clock-warmup-sec", type=float, default=5.0) + args = parser.parse_args() + torch.multiprocessing.spawn(main, args=(args.num_processes, args), nprocs=args.num_processes, join=True) diff --git a/examples/distributed/deepep_v2/example_dispatch_combine_correctness.py b/examples/distributed/deepep_v2/example_dispatch_combine_correctness.py new file mode 100644 index 0000000000..fd1909ae48 --- /dev/null +++ b/examples/distributed/deepep_v2/example_dispatch_combine_correctness.py @@ -0,0 +1,82 @@ +"""Standalone runnable correctness check for the DeepEP-EPv2-aligned dispatch/combine port. + + CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 MASTER_PORT=30071 \\ + python examples/distributed/deepep_v2/example_dispatch_combine_correctness.py + +For DeepEP's own headline shape: + + python examples/distributed/deepep_v2/example_dispatch_combine_correctness.py \\ + --tokens 8192 --hidden 7168 --topk 8 --experts 256 --num-processes 8 --num-sms 64 +""" + +import argparse + +import torch +import torch.distributed as dist +import torch.multiprocessing + +from tilelang.distributed.host import init_dist + +from buffer import Buffer +import reference + + +def main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank, num_ranks, group = init_dist(local_rank, num_local_ranks) + + torch.manual_seed(1234 + rank) + device = f"cuda:{local_rank}" + x = torch.randn(args.tokens, args.hidden, dtype=torch.bfloat16, device=device) + topk_idx, topk_weights = reference.make_topk(args.tokens, args.topk, args.experts, device, args.masked_ratio) + + buf = Buffer( + group=group, + local_rank=local_rank, + num_local_ranks=num_ranks, + num_max_tokens_per_rank=args.tokens, + hidden=args.hidden, + num_topk=args.topk, + num_experts=args.experts, + dtype=torch.bfloat16, + num_sms=args.num_sms, + dispatch_threads=args.dispatch_threads, + combine_threads=args.combine_threads, + ) + + recv_x, recv_topk_idx, recv_topk_weights, handle, _ = buf.dispatch(x, topk_idx, topk_weights) + # `dispatch` returns the full receive capacity; the reference only wants + # the rows that were actually written. Reading the count synchronises. + n = handle.num_recv_tokens + recv_x, recv_topk_idx, recv_topk_weights = recv_x[:n], recv_topk_idx[:n], recv_topk_weights[:n] + if rank == 0: + print(f"num_recv_tokens={handle.num_recv_tokens} total_capacity={buf.total_capacity}") + + expert_out = reference.simulate_expert_compute(recv_x, recv_topk_idx, recv_topk_weights) + combined, _ = buf.combine(expert_out, handle) + expected = reference.reference_combined(x, topk_weights, topk_idx) + err = (combined.float() - expected.float()).norm().item() + denom = expected.float().norm().item() + # `denom` is zero only when every selection was masked off. + rel_l2 = err / denom if denom > 0 else err + passed = rel_l2 < 0.05 + print(f"rank {rank}: rel_l2_error={rel_l2:.6f} passed={passed}") + assert passed, f"rank {rank}: mismatch, rel_l2_error={rel_l2}" + + buf.close() + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-processes", type=int, default=8) + # Fraction of top-k selections marked unselected (-1), DeepEP's marker. + parser.add_argument("--masked-ratio", type=float, default=0.0) + parser.add_argument("--tokens", type=int, default=8192) + parser.add_argument("--hidden", type=int, default=7168) + parser.add_argument("--topk", type=int, default=8) + parser.add_argument("--experts", type=int, default=256) + parser.add_argument("--num-sms", type=int, default=64) + parser.add_argument("--dispatch-threads", type=int, default=512) + parser.add_argument("--combine-threads", type=int, default=256) + args = parser.parse_args() + torch.multiprocessing.spawn(main, args=(args.num_processes, args), nprocs=args.num_processes, join=True) diff --git a/examples/distributed/deepep_v2/kernels/combine.py b/examples/distributed/deepep_v2/kernels/combine.py new file mode 100644 index 0000000000..9ee916cbc8 --- /dev/null +++ b/examples/distributed/deepep_v2/kernels/combine.py @@ -0,0 +1,243 @@ +"""Combine: remote store-back, then a local reduce. + +Mirrors DeepEP's ``combine_impl`` + ``combine_reduce_epilogue_impl`` as two +kernel launches: + +- **Store-back** -- one warp per compact row, ``put_warp`` into + ``comm_x[my_rank][src_token]`` on the source rank. That slot is unique per + (contributing rank, source token), so nothing needs to be atomic. +- **Reduce** -- one block per source token, summing the slots named by + ``send_rank_mask[token]``, the deduplicated destination set ``dispatch`` + recorded. + +Gate weights are *not* applied here. Like DeepEP, weighting and the local sum +across several experts on one rank belong to the expert epilogue; ``x`` is +already the per-compact-row contribution. + +**The expanded layout** (``do_expand``). The store-back slot +``comm_x[rank][src_token]`` is unique only because a deduplicated dispatch +gives a token one row per destination rank. Expanded, a token with two experts +here has two rows and they would collide. + +DeepEP's answer is ``kDoExpandedSend``: sum a token's local-expert rows before +sending, so one row per (rank, token) still crosses NVLink and ``comm_x`` and +the reduce stay exactly as they are. That needs the inverse of what dispatch +recorded -- dispatch gives row -> (src_rank, src_token), and this needs +(src_rank, src_token) -> rows -- so an extra kernel builds it, bucketing each +received row under its source. It touches metadata only, never the payload. + +The store-back then runs one warp per *group* rather than per row, elected by +the group's first row, and the common case is unchanged: with routing spread +over many ranks most groups hold a single row, which is sent straight from +``x`` with no summing and no staging. Only groups of two or more accumulate +into ``reduce_scratch`` first, one row of it per warp. + +Up to two bias tensors may be added to the output, DeepEP's ``bias_0``/ +``bias_1``. They seed the reduce accumulator rather than being added to it +afterwards, which costs nothing: the accumulator had to be written once +either way, and seeding replaces the clear. A token with no contributions at +all still gets its bias, which is what makes bias-only tokens behave. +""" + +import tilelang +import tilelang.language as T + + +@tilelang.jit(compile_once=True) +def combine_kernel( + num_tokens: int, + num_ranks: int, + hidden: int, + num_max_tokens_per_rank: int, + total_capacity: int, + num_sms: int, + threads: int = 256, + reduce_threads: int = 0, + dtype=T.bfloat16, + num_bias: int = 0, + do_expand: bool = False, + max_rows_per_token: int = 1, + recv_capacity: int = 0, +): + assert threads % 32 == 0 + assert 0 <= num_bias <= 2, f"DeepEP takes at most two bias tensors, got {num_bias}" + reduce_threads = reduce_threads or threads + assert reduce_threads % 32 == 0 + warps_per_cta = threads // 32 + total_warps = num_sms * warps_per_cta + cap = num_max_tokens_per_rank + + # Symbolic so the caller's contribution tensor can be read in place at + # whatever row count the last dispatch produced. + num_elems = T.symbolic("num_elems") + + # Trace-time, like dispatch's `scale_dim`: an unused bias degenerates to a + # one-row stand-in and every reference to it is absent from the generated + # code, not predicated in it. + n_bias_0 = num_tokens if num_bias >= 1 else 1 + n_bias_1 = num_tokens if num_bias >= 2 else 1 + # Deduplicated there is one row per (rank, token) and nothing to group, so + # the grouping buffers degenerate. + # Expanded, the receive buffers are wider than `total_capacity`; see + # buffer.py's `recv_capacity`. + n_recv_slots = recv_capacity or total_capacity + max_k = max_rows_per_token if do_expand else 1 + n_pairs = num_ranks * cap if do_expand else 1 + n_scratch = total_warps * hidden if do_expand else 1 + + @T.prim_func + def main( + x: T.Tensor((num_elems,), dtype), + recv_src_rank: T.Tensor((n_recv_slots,), T.int32), + recv_src_token: T.Tensor((n_recv_slots,), T.int32), + num_recv: T.Tensor((1,), T.int32), + send_rank_mask: T.Tensor((num_tokens,), T.int32), + barrier: T.Tensor((4 * num_ranks,), T.int32), + comm_x: T.Tensor((num_ranks * cap * hidden,), dtype), + # Expanded layout only: the (src_rank, src_token) -> rows inversion, + # and one scratch row per warp for groups that need summing. + group_count: T.Tensor((n_pairs,), T.uint32), + group_rows: T.Tensor((n_pairs * max_k,), T.int32), + reduce_scratch: T.Tensor((n_scratch,), dtype), + bias_0: T.Tensor((n_bias_0, hidden), dtype), + bias_1: T.Tensor((n_bias_1, hidden), dtype), + combined: T.Tensor((num_tokens, hidden), dtype), + ): + # ---------------- Bucket rows by source (expanded only) ---------------- + if do_expand: + with T.Kernel(num_sms, threads=threads) as bx: + tid = T.get_thread_binding() + n_recv = T.alloc_var(T.int32, init=num_recv[0]) + for i in T.serial(bx * threads + tid, n_pairs, num_sms * threads): + group_count[i] = 0 + # The fill below reads counters the zeroing above writes, and + # any block may touch any counter. + T.sync_grid() + # Names distinct from the store-back kernel's below: the + # tracer binds an `alloc_var` name for the whole traced + # function, so reusing one across the two kernels reads as the + # same immutable variable escaping its region. + for i in T.serial(bx * threads + tid, n_recv, num_sms * threads): + bucket_src = T.alloc_var(T.int32, init=recv_src_rank[i]) + # Alignment padding is marked -1 by dispatch and belongs to + # no token. + if bucket_src >= 0: + bucket_pair = T.alloc_var(T.int32, init=bucket_src * cap + recv_src_token[i]) + bucket_slot = T.alloc_var(T.int32, init=T.atom_add(group_count[bucket_pair], 1, scope="gpu")) + if bucket_slot < max_k: + group_rows[bucket_pair * max_k + bucket_slot] = i + + with T.Kernel(num_sms, threads=threads) as bx: + tid = T.get_thread_binding() + # Through a variable, not inline: see kernels/dispatch.py. + my_rank = T.alloc_var(T.int32, init=T.get_rank()) + + # Entry barrier. The exit barrier only says every store-back landed, + # not that every rank's *reduce* has read it, so without this a fast + # rank overwrites slots a slow one is still reducing. Buying that + # here is what lets `combine()` run with no collective around it. + T.barrier_blocks(barrier[2 * num_ranks]) + + # A contiguous chunk per warp, indexed warp-major -- DeepEP's + # `global_warp_idx`. Compact rows are grouped by source rank, so the + # obvious `bx * warps_per_cta + warp` gives one block a few hundred + # consecutive rows, all bound for one peer, and funnels the block + # down a single NVLink. Warp-major spreads a block's warps across + # the whole compact range instead; rotating the start by `my_rank` + # keeps the ranks from walking the peers in lockstep. Together worth + # 512 -> 609 GB/s, of which the rotation is 598 -> 609. + warp = ((tid // 32 + my_rank) % warps_per_cta) * num_sms + bx + # From device memory: as a scalar argument it would force `dispatch` + # to read the count back to the host, worth ~33us there. + n_recv = T.alloc_var(T.int32, init=num_recv[0]) + per_warp = T.alloc_var(T.int32, init=T.ceildiv(n_recv, total_warps)) + lane = tid % 32 + # A scratch row per warp, indexed by the *unrotated* block/warp + # pair. `warp` below folds in `my_rank`, whose range the compiler + # cannot prove, so using it here would bounds-wrap the index in an + # `if_then_else` that `address_of` rejects -- the same thing + # kernels/dispatch.py hit with `T.get_rank()` inline. + scratch_warp = bx * warps_per_cta + tid // 32 + for i in T.serial(warp * per_warp, T.min((warp + 1) * per_warp, n_recv)): + src_rank = T.alloc_var(T.int32, init=recv_src_rank[i]) + # Padding rows belong to nobody; only the expanded layout has any. + if src_rank >= 0: + slot = my_rank * cap + recv_src_token[i] + if do_expand: + pair = T.alloc_var(T.int32, init=src_rank * cap + recv_src_token[i]) + cnt = T.alloc_var(T.int32, init=group_count[pair]) + # One warp per group, not per row: the group's first row + # elects itself, the rest of the group does nothing. + if group_rows[pair * max_k] == i: + if cnt == 1: + # The common case with routing spread across + # ranks -- nothing to sum, send it where it lies. + T.put_warp( + src=T.address_of(x[i * hidden]), + dst=T.address_of(comm_x[slot * hidden]), + size=hidden, + dst_pe=src_rank, + ) + else: + # Lane-strided rather than `T.Parallel`: this is + # warp-scoped code, and a tile-level op here + # would have the compiler insert a block-wide + # barrier into it. + for e in T.serial(lane, hidden, 32): + acc = T.alloc_var(T.float32, init=0.0) + for j in T.serial(cnt): + acc = acc + T.Cast(T.float32, x[group_rows[pair * max_k + j] * hidden + e]) + reduce_scratch[scratch_warp * hidden + e] = T.Cast(dtype, acc) + T.sync_warp() + T.put_warp( + src=T.address_of(reduce_scratch[scratch_warp * hidden]), + dst=T.address_of(comm_x[slot * hidden]), + size=hidden, + dst_pe=src_rank, + ) + else: + T.put_warp( + src=T.address_of(x[i * hidden]), + dst=T.address_of(comm_x[slot * hidden]), + size=hidden, + dst_pe=src_rank, + ) + + T.barrier_blocks(barrier[3 * num_ranks]) + + # No PDL between the two kernels. PDL relaxes exactly the visibility + # an ordinary launch boundary provides, and what the reduce reads was + # written by *peer* GPUs, which PDL says nothing about; once the + # reduce got 3x faster, one rank in eight read stale slots. It saved + # no measurable time either. + + # One block per source token, not a persistent `num_sms` grid: the + # reduce has no cross-block state, and inheriting `num_sms=64` left most + # of the GPU's 148 SMs idle (551us for 733MB, 1.33 TB/s). + with T.Kernel(num_tokens, threads=reduce_threads) as token: + # Straight into a fragment, contribution by contribution. Two + # alternatives measured slower: register-staging every contribution + # first (1477us against 507us), and replacing the exit barrier with + # per-token arrival flags to overlap the reduce with the store-back + # (346-350 GB/s against 366-372) -- the chunks all finish together, + # so there is no slack to overlap into. + acc = T.alloc_fragment((hidden,), T.float32) + mask = send_rank_mask[token] + # Seed with the biases instead of clearing -- same number of + # writes to `acc`, and it keeps a token whose every selection was + # masked off from losing its bias. + if num_bias >= 1: + T.copy(bias_0[token, :], acc) + else: + T.clear(acc) + if num_bias >= 2: + for i in T.Parallel(hidden): + acc[i] += bias_1[token, i] + for r in range(num_ranks): + if ((mask >> r) & 1) == 1: + base = (r * cap + token) * hidden + for i in T.Parallel(hidden): + acc[i] += comm_x[base + i] + T.copy(acc, combined[token, :]) + + return main diff --git a/examples/distributed/deepep_v2/kernels/dispatch.py b/examples/distributed/deepep_v2/kernels/dispatch.py new file mode 100644 index 0000000000..d437b9c081 --- /dev/null +++ b/examples/distributed/deepep_v2/kernels/dispatch.py @@ -0,0 +1,583 @@ +"""Dispatch: whole-grid notify, then a deduplicated scatter into the compact output. + +Follows DeepEP EPv2's ``impls/dispatch.cuh``, intranode/NVLink only: no RDMA, +no expert alignment, no expand layout. Payload may be bf16 or fp8 (per-token, +per-128-element scales, DeepEP's ``per_token_cast_to_fp8`` layout, packed +right after the payload -- see below). + +Three phases in one kernel, no host round trip between them: + +1. **Count** -- every warp scans this rank's ``topk_idx``, deduplicates per + (token, destination rank) with ``T.match_any_sync``, and tallies into a + per-warp slice of shared memory. Dedup makes the recording lanes hold + pairwise-distinct destinations, so the tally needs no atomics; one global + ``atom_add`` per (block, destination) folds the slices into ``send_count``. +2. **Exchange** -- one warp publishes this rank's count vector to *every* peer, + so all ranks hold the full ``count_matrix[sender][destination]`` and can each + derive ``send_base[d]``, where their rows begin in destination ``d``'s output. +3. **Scatter** -- one warp per token. A token's deduplicated destinations claim + slots in one round of ``atom_add`` on a *local* counter, and the warp pushes + the row plus metadata to ``send_base[d] + slot`` on peer ``d``. + +**No copy epilogue.** DeepEP stages rows per sender and compacts them in a +second kernel, because expand layout and expert alignment hide the final +position at send time. Here the full count matrix makes ``send_base[d] + slot`` +the final compact index, so rows land in place -- saving a whole local +read+write of the payload, at the cost of the scatter having to wait for the +exchange. + +**Requires** the `address_of` fix in `src/transform/legalize_safe_memory_access.cc`: +without it, bounds checking rewrites `slot_counter[dst_rank]` into an +`if_then_else` and the `address_of` that `atom_add`/`st` build around it is +rejected at codegen. The workaround was an inline clamp on every such index +plus disabling `LoopUnswitching`; both are gone. + +**Why ``put_warp`` and not TMA.** The TMA path was built and never won (572.6 +against 563.5 GB/s at 512 threads, 518.9 against 520.4 at 256 -- inside the +noise). ``cp.async.bulk`` has no global-to-global form, so a TMA store must +stage through shared memory, trading that round trip for the issue slots +``put_warp`` spends. + +**Why FP8's scale is packed into the same row as the payload, not a second +buffer.** It used to be a second ``put_warp`` per token-destination pair. +Measured: 566.9us with it, 493.1us without -- the payload copy alone already +scales cleanly with bytes (493.1us is almost exactly half of bf16's ~895us, +matching fp8's half-size payload), and a *large* transfer split into two +``put_warp`` calls costs nothing extra (901.8us against a 894.8us single-call +baseline, same total bytes). So the ~74us is not "a second call" in general -- +it is a second call for something this small. Whatever fixed cost a remote +store pays regardless of size (peer-address translation, warp-level setup, an +NVLink round trip) is noise against a 7168-byte payload and dominates a +224-byte scale. Packing scale bytes right after the payload -- what +``reference.per_token_cast_to_fp8`` produces -- turns two stores into one and +costs nothing extra upstream: quantisation already has to write its output +somewhere. + +Fused, the real number is ~522us against the two-store 547-580us -- a win, +but not the full ~74us the isolated single-call-vs-two-call comparison above +suggested, because the fused call also moves the 224 scale bytes the no-scale +test did not (7392 against 7168) and then the row padding on top of that. +How wide to pad is its own measured tradeoff, and `reference.packed_row_bytes` +owns it; `row_bytes` arrives here already decided. + +**The expanded layout** (``do_expand``, DeepEP's ``kDoExpand``). Off, a token +that picks two experts on one rank is deduplicated into a single received row, +and rows are grouped by *sender*. On, it becomes one row per (token, expert), +and rows are grouped by *local expert* -- what a grouped GEMM needs, since each +expert's rows are then contiguous. + +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 at all -- rows land at their final index straight from the sender +-- so the sender has to know that index, and the only thing standing in the way +is the *granularity of the count exchange*. Everything else is unchanged: phase +1 tallies per destination, phase 2 turns the tallies into a base, phase 3 adds +a locally-claimed slot. Expanding only redefines "destination" from rank to +expert, so ``n_dst`` is ``num_experts`` instead of ``num_ranks`` and the same +three phases produce the expert-major layout. + +The one genuinely new term is the segment base: the destination groups its +output by local expert, so this rank's rows for expert ``e`` start at +``segment_base[e] + sender_base[e]``, where ``segment_base`` is the exclusive +prefix sum over the destination's local experts of their *aligned* received +counts, and ``sender_base`` is what phase 2 already computed. Every rank holds +the whole count matrix, so every rank computes the same segment bases without +talking to anyone. + +``expert_alignment`` rounds each segment up. The gap between an expert's real +count and its aligned one is never written by the scatter, so it holds +whatever the previous call left; ``zero_padding`` clears it. Turn it off only +if the consumer honours the unaligned counts. + +**Capacity.** Deduplicated, a peer can send at most ``num_max_tokens_per_rank`` +rows, so ``num_ranks`` times that is a bound no routing can exceed. Expanded, +the bound is ``min(topk, experts_per_rank)`` times higher -- every token +picking every one of its experts on the same rank -- which at the V3 shape is +7 GiB against 0.88. Balanced routing needs only the 0.88, so capacity is the +caller's ``expand_factor`` and phase 2 checks it: every rank derives the same +total from the same count matrix, so an overflow is known *before* any payload +moves, and the destination is skipped rather than corrupted. See +``buffer.py``'s ``expand_overflow``. + +**Reusing a layout** (``cached``, DeepEP's ``handle=``). Phases 1-2 depend only +on the routing, so a second dispatch with the same ``topk_idx`` recomputes a +layout it already has. ``cached`` traces the notify kernel away entirely and +leaves only the scatter, which reads ``send_base``/``send_rank_mask``/ +``num_recv`` exactly as the previous call left them -- none of them is touched +by the end-of-call reset, which clears only what phase 1-2 consume. + +Splitting the kernel is what makes this worth anything. Fused, notify was +welded to the scatter and skipping it meant skipping the payload too; DeepEP +has the same problem from the other direction and its cached dispatch measures +441-444us against 437-439 uncached -- i.e. it saves host work, not GPU time. + +Measured here, three clean samples each: fp8 dispatch 514 -> 497us whole call +(3.3%), bf16 888 -> 871 (2.0%). Not the whole 36us notify kernel, because the +cached scatter is about 10us slower than the uncached one -- the entry barrier +moves into it rather than disappearing, being about peers not overwriting data +this rank is still reading, which holds however the layout was obtained. + +The entry barrier moves with it. It stops a peer's next round from overwriting +data this rank has not finished reading, which is needed whether or not the +layout was recomputed, so when ``cached`` it opens the scatter instead. + +**One launch.** Phases 1-2 compute a layout and phase 3 moves the payload, with +``T.sync_grid()`` between them carrying what phase 2 wrote. They were two +kernels for a while, because a launch boundary carries it just as well and +costs nothing on an idle device -- 881.7us split against 883.7 fused. + +Behind a large GEMM it costs 118.7us, because the notify releases all of its +SMs on the way out and the GEMM's pending waves take them back before phase 3 +can ask for them. The same gap solo is 1.9us. That is nineteen points of +overlap for two microseconds of standalone throughput; see the section on +overlap. + +The split also let phase 3 run on a wider grid than the notify, which one +``sync_grid``-bearing kernel cannot express. That turned out not to be worth +keeping either: widening the whole grid beats it outright, 860.6us at +``num_sms=128`` against 867.1 for a 64-wide notify and a 128-wide scatter, and +the notify does not care whether it has 64 blocks or 128 for its 26us of work. + +It is not DeepEP's split. 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. The cut +here is before the movement instead, and the piece it isolates -- 34.7us of +layout computation, constant across dtypes -- is the piece a cached ``handle=`` +dispatch would skip outright. + +""" + +import tilelang +import tilelang.language as T + + +def _dedup_leader(value, lane): + """1 on the lowest-indexed lane holding `value` -- DeepEP's `ptx::deduplicate`. + + `T.match_any_sync` is warp-collective, so callers must evaluate this + unconditionally, never behind a lane-divergent guard or a short-circuiting + `and`. The int32 spelling of "bits below my lane" is deliberate: the uint32 + `(1u << lane) - 1` trips TVM's bounds analysis. + """ + match_mask = T.Cast("int32", T.match_any_sync(value)) + return T.if_then_else((match_mask & (~(-1 << lane))) == 0, 1, 0) + + +@tilelang.jit(compile_once=True) +def dispatch_kernel( + num_tokens: int, + num_ranks: int, + num_experts: int, + topk: int, + hidden: int, + num_max_tokens_per_rank: int, + num_sms: int, + threads: int = 256, + dtype=T.bfloat16, + scale_dim: int = 0, + row_bytes: int = 0, + collect_expert_stats: bool = False, + do_expand: bool = False, + expert_alignment: int = 1, + zero_padding: bool = True, + expand_capacity: int = 0, + cached: bool = False, +): + assert threads % 32 == 0 + assert num_experts % num_ranks == 0 + assert topk <= 32, "one lane per top-k entry" + assert expert_alignment >= 1 + assert do_expand or expert_alignment == 1, "expert_alignment only means anything in the expanded layout" + experts_per_rank = num_experts // num_ranks + # Rank or expert: the only thing `do_expand` really changes. See the + # module docstring. + n_dst = num_experts if do_expand else num_ranks + warps_per_cta = threads // 32 + total_warps = num_sms * warps_per_cta + cap = num_max_tokens_per_rank + total_capacity = cap * num_ranks + # Deduplicated, `total_capacity` is a bound no routing can exceed. Expanded + # it is not, so the caller sizes it and phase 2 checks. See the module + # docstring. + recv_capacity = expand_capacity if do_expand else total_capacity + assert not do_expand or recv_capacity > 0, "the expanded layout needs an explicit expand_capacity" + + # `scale_dim` is a Python int, so this branch is resolved while tracing. + # FP8's `x`/`recv_x` is the *packed* row `reference.per_token_cast_to_fp8` + # produces: `hidden` payload bytes followed by the per-group fp32 scale, + # then padding, moved as one opaque `uint8` region so the scatter needs + # one `put_warp` per token-destination pair instead of two -- see the + # module docstring. `row_bytes` is `reference.packed_row_bytes`'s result, + # passed in by `buffer.py` rather than recomputed here: which padding is + # fastest is a measured, size-dependent call and wants one owner (see + # that function's docstring for the numbers). + # bf16 (`scale_dim == 0`) is untouched: `row_width`/`row_dtype` degenerate + # to exactly what they were before this existed. + if scale_dim: + assert row_bytes >= hidden * (dtype.bits // 8) + scale_dim * 4, ( + f"row_bytes={row_bytes} cannot hold {hidden} payload + {scale_dim} scales" + ) + row_width = row_bytes + row_dtype = T.uint8 + else: + row_width = hidden + row_dtype = dtype + + # Another trace-time branch: off, the tally below is absent from the + # generated code and the argument degenerates to a one-element stand-in. + stats_dim = experts_per_rank if collect_expert_stats else 1 + # Likewise for the expanded layout's own outputs. Gated on + # `do_expand`, not on `collect_expert_stats` -- different features that + # happen to be per-local-expert. + count_dim = experts_per_rank if do_expand else 1 + seg_dim = experts_per_rank + 1 if do_expand else 1 + + @T.prim_func + def main( + x: T.Tensor((num_tokens, row_width), row_dtype), + topk_idx: T.Tensor((num_tokens, topk), T.int32), + topk_weights: T.Tensor((num_tokens, topk), T.float32), + # `notify_done` counts arriving blocks in phase 1; `exchange_done` + # releases every block once phase 2 has published the offsets. + notify_done: T.Tensor((1,), T.uint32), + exchange_done: T.Tensor((1,), T.uint32), + send_count: T.Tensor((n_dst,), T.uint32), + # `count_matrix[sender * n_dst + destination]`, -1 until published. + # Symmetric: every rank writes its own row into every peer's copy. + # "Destination" is a rank, or an expert when expanding. + count_matrix: T.Tensor((num_ranks * n_dst,), T.int32), + send_base: T.Tensor((n_dst,), T.int32), + psum_recv_count: T.Tensor((num_ranks,), T.int32), + num_recv: T.Tensor((1,), T.int32), + slot_counter: T.Tensor((n_dst,), T.uint32), + # Bit r set iff this token went to rank r. Free to produce here, and it + # is what lets `kernels/combine.py`'s reduce pass know which slots are + # live without redoing the dedup (DeepEP recomputes it there instead). + send_rank_mask: T.Tensor((num_tokens,), T.int32), + barrier: T.Tensor((4 * num_ranks,), T.int32), + # Flat 1D: `put_warp` takes a raw address, and `T.address_of(buf[row, 0])` + # on a buffer with a large row stride trips "Can't fetch the lanes of a + # scalable vector" in `StorageRewrite`; `st` only accepts single-index + # buffer loads at all. + recv_x: T.Tensor((recv_capacity * row_width,), row_dtype), + recv_src_rank: T.Tensor((recv_capacity,), T.int32), + recv_src_token: T.Tensor((recv_capacity,), T.int32), + recv_topk_idx: T.Tensor((recv_capacity * topk,), T.int32), + recv_topk_weights: T.Tensor((recv_capacity * topk,), T.float32), + # DeepEP's `cumulative_local_expert_recv_stats`: tokens received per + # *local* expert, accumulated across calls. uint32 for `atom_add`, as + # with `send_count`. Never reset here -- the caller owns the window + # it is accumulating over. + recv_expert_stats: T.Tensor((stats_dim,), T.uint32), + # Expanded layout only. `expert_offset` is the exclusive prefix sum of + # this rank's *aligned* per-expert counts, so local expert `e` owns + # `[expert_offset[e], expert_offset[e + 1])` and `expert_count[e]` of + # those rows are real. `expand_overflow` is set if the aligned total + # does not fit the capacity the caller sized for. + expert_count: T.Tensor((count_dim,), T.int32), + expert_offset: T.Tensor((seg_dim,), T.int32), + expand_overflow: T.Tensor((1,), T.int32), + ): + # A single persistent grid: the phases rendezvous through global + # counters, so every block has to be resident at once. `num_sms` + # defaults to the device's SM count (see buffer.py) and one block per SM + # always fits. + with T.Kernel(num_sms, threads=threads) as bx: + # Through a variable, not `T.get_rank()` inline: inline leaves the + # index range unknown, so bounds checking wraps every + # `buf[my_rank ...]` in an `if_then_else` that `address_of` rejects. + tid = T.get_thread_binding() + lane = tid % 32 + local_warp = tid // 32 + warp = bx * warps_per_cta + local_warp + my_rank = T.alloc_var(T.int32, init=T.get_rank()) + + if not cached: + # Entry barrier, DeepEP's `kDispatchTag0`: stops a peer's *next* + # round from writing this rank's `count_matrix` before this round's + # reset has landed. The exit barrier cannot -- it only says a peer + # reached it, after which the peer may finish and relaunch. Its own + # slot, disjoint from the exit barrier's; see buffer.py. + T.barrier_blocks(barrier[0]) + + # ---------------- Phase 1: count ---------------- + blk_count = T.alloc_shared((warps_per_cta * n_dst,), "int32") + # Flat, one row per warp, and deliberately untouched by any tile-level + # op: keeping `T.copy`/`T.atomic_add` away from it is what keeps the + # compiler from inserting block-wide barriers into warp-scoped code. + for i in T.serial(tid, warps_per_cta * n_dst, threads): + blk_count[i] = 0 + T.sync_threads() + + for token in T.serial(warp, num_tokens, total_warps): + expert = T.alloc_var(T.int32, init=-1) + dst_rank = T.alloc_var(T.int32, init=-1) + if lane < topk: + expert = topk_idx[token, lane] + if expert >= 0: + dst_rank = expert // experts_per_rank + # Expanding, the destination *is* the expert, and the dedup + # below is a no-op: DeepEP asserts a token's top-k entries are + # distinct experts, so the lanes are already pairwise-distinct. + # The counter stays atomic-free for the same reason either way. + dst = expert if do_expand else dst_rank + leader = T.alloc_var(T.int32, init=_dedup_leader(dst, lane)) + if leader == 1 and dst >= 0: + # No atomic: after dedup the lanes reaching here hold + # pairwise-distinct destinations, and each warp owns its own + # slice, so no two threads ever touch the same counter. + at = local_warp * n_dst + dst + blk_count[at] = blk_count[at] + 1 + # Record the destination set here rather than in the scatter: + # combine needs it, and gathering it costs a round of shuffles + # that has no business being on the data path. + # Rank granularity even when expanding: combine reduces over + # ranks either way. Expanding, `leader` is per-expert and this + # needs a dedup of its own; not expanding, `dst` *is* `dst_rank`, + # so reuse it rather than issue a second warp-collective + # `T.match_any_sync` per token for the same answer. + if do_expand: + rank_leader = T.alloc_var(T.int32, init=_dedup_leader(dst_rank, lane)) + else: + rank_leader = leader + rank_mask = T.alloc_var(T.int32, init=0) + for k in range(topk): + dst_k = T.alloc_var(T.int32, init=T.shfl_sync(dst_rank, k)) + lead_k = T.alloc_var(T.int32, init=T.shfl_sync(rank_leader, k)) + if lead_k == 1 and dst_k >= 0: + rank_mask = rank_mask + (1 << dst_k) + if lane == 0: + send_rank_mask[token] = rank_mask + T.sync_threads() + + for d in T.serial(tid, n_dst, threads): + folded = T.alloc_var(T.int32, init=0) + for w in range(warps_per_cta): + folded = folded + blk_count[w * n_dst + d] + if folded > 0: + T.atom_add(send_count[d], folded, scope="gpu") + T.sync_threads() + if tid == 0: + T.atom_add(notify_done[0], 1, scope="gpu", sem="release") + + # ---------------- Phase 2: exchange ---------------- + if bx == 0 and local_warp == 0: + if lane == 0: + T.wait_ge(notify_done[0], num_sms, scope=T.WaitScope.GPU, semantics=T.WaitSemantics.ACQUIRE) + T.sync_warp() + + # Publish this rank's count vector to every peer, so all ranks + # hold the same `count_matrix`. + # + # One fence, then relaxed stores. Per-store `sem="release"` cost + # 4.4%: each carries a `fence.release.sys` that waits for prior + # writes to cross NVLink, serialising eight round trips to + # publish 64 integers. Peers read `count_matrix` directly and + # infer nothing else from it, so the stores need no ordering + # against each other. + T.fence_sys() + # `n_dst` is `num_ranks` unless expanding, where it is + # `num_experts` and the row no longer fits one lane each. + for p in range(num_ranks): + for c in T.serial(lane, n_dst, 32): + T.st(count_matrix[my_rank * n_dst + c], send_count[c], scope="sys", sem="relaxed", dst_pe=p) + T.sync_warp() + for s in range(num_ranks): + for c in T.serial(lane, n_dst, 32): + T.wait_ge(count_matrix[s * n_dst + c], 0, scope=T.WaitScope.SYS, semantics=T.WaitSemantics.ACQUIRE) + T.sync_warp() + + # Lane d: where my rows start inside destination d's output. + # Expanding, `d` is an expert and this is only the offset + # *within* that expert's segment; the segment base is added + # below, once every destination's aligned layout is known. + for d in T.serial(lane, n_dst, 32): + base = T.alloc_var(T.int32, init=0) + for s in T.serial(my_rank): + base = base + count_matrix[s * n_dst + d] + send_base[d] = base + T.sync_warp() + + if do_expand: + # Every destination lays its output out the same way, from + # the same count matrix, so each rank can reconstruct all of + # them and nobody has to be told. Lane p handles peer p. + if lane < num_ranks: + seg = T.alloc_var(T.int32, init=0) + for e in range(experts_per_rank): + total = T.alloc_var(T.int32, init=0) + for s in range(num_ranks): + total = total + count_matrix[s * n_dst + lane * experts_per_rank + e] + # My rows for this expert sit at the segment base + # plus the offset among earlier senders. + send_base[lane * experts_per_rank + e] = seg + send_base[lane * experts_per_rank + e] + # Aligned, so the next segment starts on a boundary + # a grouped GEMM can consume. + seg = seg + T.ceildiv(total, expert_alignment) * expert_alignment + if lane == my_rank: + expert_count[e] = total + expert_offset[e + 1] = seg + # `seg` is now this destination's aligned total. Known + # to every rank before a single payload byte has moved, + # which is what makes the capacity check free. + if lane == my_rank: + expert_offset[0] = 0 + num_recv[0] = seg + if seg > recv_capacity: + expand_overflow[0] = seg + if seg > recv_capacity: + # Skip this destination entirely rather than write + # past its buffer. It raises its own flag. + for e in range(experts_per_rank): + send_base[lane * experts_per_rank + e] = -1 + else: + # Lane 0: my own receive prefix sum (DeepEP's + # `psum_num_recv_tokens_per_scaleup_rank`) and total. + if lane == 0: + psum = T.alloc_var(T.int32, init=0) + for s in range(num_ranks): + psum = psum + count_matrix[s * num_ranks + my_rank] + psum_recv_count[s] = psum + num_recv[0] = psum + T.sync_warp() + if lane == 0: + T.atom_add(exchange_done[0], 1, scope="gpu", sem="release") + + # Everything phase 2 published has to be visible to phase 3, and + # `barrier_blocks` rendezvouses ranks rather than this rank's own + # blocks, so the grid sync is what carries it now that the launch + # boundary is gone. Same guarantee the spin on `exchange_done` and + # then the launch boundary each gave in turn. + T.sync_grid() + + # With no notify kernel ahead of it, the scatter owns the entry + # barrier. It is not about the layout -- it stops a peer's next + # round from landing on data this rank is still reading -- so it is + # needed either way, just in whichever kernel goes first. + if cached: + T.barrier_blocks(barrier[0]) + + # ---------------- Phase 3: scatter ---------------- + for token in T.serial(warp, num_tokens, total_warps): + expert = T.alloc_var(T.int32, init=-1) + dst_rank = T.alloc_var(T.int32, init=-1) + slot = T.alloc_var(T.int32, init=-1) + if lane < topk: + expert = topk_idx[token, lane] + if expert >= 0: + dst_rank = expert // experts_per_rank + dst = expert if do_expand else dst_rank + leader = T.alloc_var(T.int32, init=_dedup_leader(dst, lane)) + # Every destination of this token claims its slot at once, one + # atomic per owning lane. `send_base[dst] < 0` marks a + # destination phase 2 found would overflow, which only the + # expanded layout can hit -- deduplicated, capacity is a hard + # bound -- but the condition is the same either way and + # `send_base` is eight hot integers. + if leader == 1 and dst >= 0 and send_base[dst] >= 0: + slot = T.atom_add(slot_counter[dst], 1, scope="gpu") + + for k in range(topk): + dst_k = T.alloc_var(T.int32, init=T.shfl_sync(dst, k)) + slot_k = T.alloc_var(T.int32, init=T.shfl_sync(slot, k)) + # Expanding, the destination is an expert and the peer that + # owns it has to be shuffled separately; otherwise they are + # the same value and the second shuffle is pure cost. + if do_expand: + pe_k = T.shfl_sync(dst_rank, k) + else: + pe_k = dst_k + if slot_k >= 0: + # The final index on the destination -- known here only + # because phase 2 handed every rank the full count + # matrix. Expanding, it is already inside the right + # expert's segment. + idx_k = T.alloc_var(T.int32, init=send_base[dst_k] + slot_k) + T.put_warp( + src=T.address_of(x[token, 0]), + dst=T.address_of(recv_x[idx_k * row_width]), + size=row_width, + dst_pe=pe_k, + ) + if lane == 0: + T.st(recv_src_rank[idx_k], my_rank, scope="sys", sem="relaxed", dst_pe=pe_k) + T.st(recv_src_token[idx_k], token, scope="sys", sem="relaxed", dst_pe=pe_k) + if lane < topk: + local_expert = T.alloc_var(T.int32, init=-1) + if do_expand: + # This row *is* one (token, expert) pair, and + # `dst_k` is that expert. Marking the token's + # other local experts here too would have the + # epilogue weight every one of the token's rows + # by the same sum, and combine would add them up. + if expert == dst_k: + local_expert = expert - pe_k * experts_per_rank + elif pe_k * experts_per_rank <= expert and expert < (pe_k + 1) * experts_per_rank: + local_expert = expert - pe_k * experts_per_rank + T.st(recv_topk_idx[idx_k * topk + lane], local_expert, scope="sys", sem="relaxed", dst_pe=pe_k) + T.st(recv_topk_weights[idx_k * topk + lane], topk_weights[token, lane], scope="sys", sem="relaxed", dst_pe=pe_k) + + T.barrier_blocks(barrier[num_ranks]) + + # Reset for the next call. The `T.sync_grid()` is load-bearing: + # `tl::barrier_blocks` rendezvouses *ranks*, not this rank's own + # blocks, so behind it alone block 0 zeroes `slot_counter` while + # other blocks are still claiming slots -- which silently produces + # duplicate (src_rank, src_token) rows. DeepEP's `gpu_barrier` uses + # `this_grid().sync()` for exactly this. + T.sync_grid() + + # ---------------- Alignment padding ---------------- + # The scatter writes only real rows, so the gap between an expert's + # count and its aligned segment end still holds whatever the last + # call put there. A grouped GEMM reading the aligned segment would + # consume it, so clear it -- DeepEP's `kDoZeroPadding`, and for the + # same reason. Rows, not bytes: one warp per padding row, strided + # over the whole grid. + if do_expand and zero_padding: + for e in range(experts_per_rank): + pad_begin = T.alloc_var(T.int32, init=expert_offset[e] + expert_count[e]) + pad_end = T.alloc_var(T.int32, init=expert_offset[e + 1]) + for row in T.serial(pad_begin + warp, pad_end, total_warps): + for i in T.serial(lane, row_width, 32): + recv_x[row * row_width + i] = T.Cast(row_dtype, 0) + if lane == 0: + # A padding row belongs to no sender and no token; + # combine skips anything marked this way. + recv_src_rank[row] = -1 + recv_src_token[row] = -1 + if lane < topk: + recv_topk_idx[row * topk + lane] = -1 + recv_topk_weights[row * topk + lane] = T.Cast(T.float32, 0) + + # ---------------- Per-expert receive stats (optional) ---------------- + # Counted here, on the receiver, rather than exchanged like the rank + # counts: `recv_topk_idx` already holds the local expert for every + # received (row, top-k slot) and -1 elsewhere, so the answer is a + # local scan of something this rank was going to be handed anyway. + # DeepEP derives it in notify instead, from a per-expert count + # vector it already exchanges -- it has one because the expanded + # layout needs per-expert offsets. This port's layout does not, and + # widening the exchange from `num_ranks` to `num_experts` entries + # per rank to avoid a scan that costs microseconds is a bad trade. + # + # Behind the exit barrier and `sync_grid`, so every peer's stores + # have landed and this rank's own blocks are past the scatter. + if collect_expert_stats: + n_recv = T.alloc_var(T.int32, init=num_recv[0]) + for i in T.serial(bx * threads + tid, n_recv * topk, num_sms * threads): + local_expert = T.alloc_var(T.int32, init=recv_topk_idx[i]) + if local_expert >= 0: + T.atom_add(recv_expert_stats[local_expert], 1, scope="gpu") + + if bx == 0: + for d in T.serial(tid, n_dst, threads): + send_count[d] = 0 + slot_counter[d] = 0 + for c in T.serial(tid, num_ranks * n_dst, threads): + count_matrix[c] = -1 + if tid == 0: + notify_done[0] = 0 + exchange_done[0] = 0 + + return main diff --git a/examples/distributed/deepep_v2/reference.py b/examples/distributed/deepep_v2/reference.py new file mode 100644 index 0000000000..057da41a72 --- /dev/null +++ b/examples/distributed/deepep_v2/reference.py @@ -0,0 +1,225 @@ +"""Correctness reference for the dispatch/combine round trip. + +There is no fused GEMM in this port's scope. DeepEP's own non-expand +``combine`` does not re-apply gate weights or sum multiple local-expert +contributions itself -- that happens in the grouped-GEMM epilogue *before* +`combine` is called (see ``kernels/combine.py`` for the detailed reasoning). +``simulate_expert_compute`` stands in for that epilogue: identity "expert +compute" (the input row, unchanged) scaled by the sum of this rank's valid +top-k gate weights for that compact row (0 if this rank doesn't actually own +any of the token's experts, >1 term summed if it owns more than one). This +still exercises every routing/addressing/accumulation path in dispatch and +combine -- a misrouted or dropped/duplicated contribution changes the sum a +token's original weights should reconstruct to. +""" + +from __future__ import annotations + +import torch + + +def simulate_expert_compute(recv_x: torch.Tensor, recv_topk_idx: torch.Tensor, recv_topk_weights: torch.Tensor) -> torch.Tensor: + """Identity-compute stand-in for the missing grouped-GEMM epilogue.""" + valid = recv_topk_idx >= 0 + weight_sum = torch.where(valid, recv_topk_weights, torch.zeros_like(recv_topk_weights)).sum(dim=-1, keepdim=True) + return recv_x * weight_sum.to(recv_x.dtype) + + +def reference_combined(x: torch.Tensor, topk_weights: torch.Tensor, topk_idx: torch.Tensor | None = None) -> torch.Tensor: + """Expected combine output: every token's top-k weight sum applied once. + + Entries with `topk_idx < 0` are unselected and contribute nothing, so a + token with no selections at all comes back as zero. + """ + w = topk_weights + if topk_idx is not None: + w = torch.where(topk_idx >= 0, w, torch.zeros_like(w)) + return x * w.sum(dim=-1, keepdim=True).to(x.dtype) + + +def make_topk(num_tokens: int, topk: int, num_experts: int, device, masked_ratio: float = 0.0): + """Routing inputs, with `masked_ratio` of the selections marked unselected. + + `-1` is DeepEP's "no selection" marker and reaches dispatch's dedup and + slot-claiming as a destination of -1, so it needs exercising rather than + assuming. + """ + idx = torch.randint(0, num_experts, (num_tokens, topk), device=device) + weights = torch.rand(num_tokens, topk, device=device) + if masked_ratio > 0: + idx = idx.masked_fill(torch.rand_like(idx, dtype=torch.float) < masked_ratio, -1) + return idx, weights + + +# --------------------------------------------------------------------------- +# FP8 dispatch +# +# The layout is DeepEP's `per_token_cast_to_fp8`: quantise per token over +# groups of `group` elements, keeping one fp32 scale per group. Values and +# scales are packed into one uint8 row -- payload bytes, then the per-group +# fp32 scales immediately after -- rather than returned as two separate +# tensors. Dispatch's scatter does one remote store per token-destination +# pair either way; keeping values and scales apart would mean a *second* +# store for the scale alone, and a store's fixed per-call cost (peer-address +# translation, warp sync) turns out to swamp a transfer that small -- +# measured at ~74us out of ~567us end to end, for 224 bytes of scale against +# 7168 bytes of payload. Packing them contiguously here costs nothing extra: +# quantisation already has to write its output somewhere, and writing the +# scale right after the payload it belongs to is no more expensive than +# writing it elsewhere. +# +# The row is padded, at minimum to 16 bytes: `put_warp`'s bulk path +# reinterprets the row as `int4`, so a row whose *stride* is not 16-byte +# aligned is a `"misaligned address"` CUDA error outright (`hidden=128` packs +# to 132 raw bytes, not even a multiple of 16). +# +# 512 = 32 lanes * 16 bytes is a second, coarser boundary worth reaching when +# it is cheap. `cp_warp_impl`'s *drain* loop -- the tail past whole +# `UNROLL_FACTOR`-sized chunks -- is `for i = drain_start + lane; i < N_int4; +# i += 32`, so unless the leftover `int4` count is a multiple of 32 some lanes +# run one more iteration than the rest, which idle through it. At hidden=7168 +# the fused row is 7392 bytes: a multiple of 16 but not of 512 (462 int4, +# remainder 78 -- 14 lanes take a 3rd drain iteration, 18 take 2). +# +# But reaching 512 costs the padding bytes, and those cross NVLink too. Which +# way that trades depends entirely on how much padding it takes, so +# `packed_row_bytes` only pays for 512 when it is nearly free. Measured on the +# real dispatch kernel (8xB200, 8192 tokens/rank, top-8, 256 experts, 64 SMs, +# min of 3-4 samples each, all taken in verified-idle windows): +# +# hidden 16-aligned 512-aligned padding result +# 2048 2112 -> 211us 2560 -> 230us +21.2% 512 is 8.9% slower +# 4096 4224 -> 325us 4608 -> 348us +9.1% 512 is 6.8% slower +# 7168 7392 -> 535us 7680 -> 522us +3.9% 512 is 2.4% faster +# +# So a uniform drain is worth roughly 6% gross, and `_ROW_ALIGN_MAX_PAD` +# takes it only when the padding costs less than that. Note this is invisible +# in GB/s -- an isolated `put_warp` roofline reads 594 GB/s at 7392 against +# 618 at 7680 and 512 looks like a clear win, but the wall time is identical +# (203.8 against 203.7us) because the extra bytes eat exactly what the +# uniformity buys. Compare times, not rates, whenever the byte count moves. +# +# `buffer.py`'s `row_bytes` calls this; `kernels/dispatch.py` is handed the +# result rather than recomputing it. +# --------------------------------------------------------------------------- + +_FP8_MAX = 448.0 + +# `int4`, the unit `put_warp`'s bulk path moves -- a hard requirement. +_ROW_ALIGN = 16 +# 32 lanes * one `int4` each: one full drain iteration, warp-uniform. +_ROW_ALIGN_WIDE = 512 +# How much padding reaching `_ROW_ALIGN_WIDE` may cost before it stops paying +# for itself. See the measurements in the module docstring above. +_ROW_ALIGN_MAX_PAD = 0.05 + + +def _cdiv(a: int, b: int) -> int: + return (a + b - 1) // b + + +def align_up(x: int, align: int) -> int: + return _cdiv(x, align) * align + + +def packed_row_bytes(hidden: int, group: int = 128) -> int: + """FP8's packed row width in bytes: `hidden` payload bytes plus one fp32 + scale per `group`, padded to whichever of `_ROW_ALIGN`/`_ROW_ALIGN_WIDE` + is faster at this size. See the module docstring above.""" + raw = hidden + (hidden // group) * 4 + narrow = align_up(raw, _ROW_ALIGN) + wide = align_up(raw, _ROW_ALIGN_WIDE) + return wide if wide <= narrow * (1 + _ROW_ALIGN_MAX_PAD) else narrow + + +def per_token_cast_to_fp8(x: torch.Tensor, group: int = 128) -> torch.Tensor: + """`[m, n]` bf16 -> `[m, packed_row_bytes(n, group)]` uint8 (fp8 payload, then fp32 scales, then padding).""" + assert x.dim() == 2 and x.shape[1] % group == 0, f"hidden {x.shape[1]} is not a multiple of {group}" + m, n = x.shape + grouped = x.view(m, -1, group) + amax = grouped.abs().float().amax(dim=2).clamp(1e-4) + q = (grouped.float() * (_FP8_MAX / amax.unsqueeze(2))).to(torch.float8_e4m3fn).view(m, n) + scales = (amax / _FP8_MAX).contiguous() + scale_bytes = scales.shape[1] * 4 + + packed = torch.empty((m, packed_row_bytes(n, group)), dtype=torch.uint8, device=x.device) + packed[:, :n] = q.view(torch.uint8) + packed[:, n : n + scale_bytes].view(torch.float32).copy_(scales) + return packed + + +def per_token_cast_back(packed: torch.Tensor, hidden: int, group: int = 128) -> torch.Tensor: + """Inverse of `per_token_cast_to_fp8`, back to bf16.""" + m, _ = packed.shape + scale_dim = hidden // group + values = packed[:, :hidden].view(torch.float8_e4m3fn) + scales = packed[:, hidden : hidden + scale_dim * 4].view(torch.float32) + return (values.view(m, -1, group).float() * scales.view(m, -1, 1)).view(m, hidden).to(torch.bfloat16) + + +# --------------------------------------------------------------------------- +# Expanded layout +# +# DeepEP's `do_expand`: one received row per (token, expert) rather than one +# per (token, destination rank), and rows grouped by *local expert* so a +# grouped GEMM can consume each expert's segment as a contiguous block. +# +# DeepEP expands in its receiver-side copy epilogue, which is free to bump an +# atomic 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 the index, which means the count exchange has to +# carry per-expert counts instead of per-rank ones. The whole layout then falls +# out of the same three numbers dispatch already computes, one granularity +# down: +# +# segment_base[e] exclusive prefix sum over local experts of the +# *aligned* received counts -- where expert e's rows +# start in the destination's output +# sender_base[s][e] sum over senders before s of count[s][e] -- where +# sender s's rows start inside that segment +# slot a local bump counter per (destination expert) +# +# and the destination index is the sum of the three. Each sender owns a +# disjoint sub-range of each expert's segment, so nothing needs a remote +# atomic, exactly as in the rank layout. +# +# `expert_alignment` rounds each expert's segment up, so the gap between an +# expert's real count and its aligned one holds whatever the last call left +# there. `zero_padding` clears it; without it a grouped GEMM would consume +# stale rows. +# --------------------------------------------------------------------------- + + +def expanded_layout(topk_idx_per_rank, num_experts: int, num_ranks: int, expert_alignment: int = 1): + """The layout every rank's dispatch should produce, from every rank's routing. + + `topk_idx_per_rank[s]` is sender `s`'s `[num_tokens, topk]` selection. Returns + `(rows, counts, offsets)`: + + - `rows[r][e]` -- the `(src_rank, src_token)` pairs destination `r`'s local + expert `e` should receive, in the order the layout puts them: grouped by + sender, senders in rank order. Within one sender the order is whatever + its slot counter hands out, so compare these as sets per sender block. + - `counts[r][e]` -- unaligned received count. + - `offsets[r]` -- `[experts_per_rank + 1]` exclusive prefix sum of the + *aligned* counts: expert `e`'s segment is `[offsets[e], offsets[e + 1])`, + of which the first `counts[e]` rows are real and the rest is padding. + """ + experts_per_rank = num_experts // num_ranks + rows = [[[] for _ in range(experts_per_rank)] for _ in range(num_ranks)] + for src_rank, topk_idx in enumerate(topk_idx_per_rank): + for token in range(topk_idx.shape[0]): + # A token reaches an expert at most once: DeepEP asserts the top-k + # entries of a token are distinct experts, and so does this port. + for expert in sorted({int(e) for e in topk_idx[token].tolist() if e >= 0}): + rows[expert // experts_per_rank][expert % experts_per_rank].append((src_rank, token)) + + counts = [[len(rows[r][e]) for e in range(experts_per_rank)] for r in range(num_ranks)] + offsets = [] + for r in range(num_ranks): + off, acc = [0], 0 + for e in range(experts_per_rank): + acc += align_up(counts[r][e], expert_alignment) + off.append(acc) + offsets.append(off) + return rows, counts, offsets diff --git a/examples/distributed/deepep_v2/test_example_deepep_v2.py b/examples/distributed/deepep_v2/test_example_deepep_v2.py new file mode 100644 index 0000000000..bf0da0681e --- /dev/null +++ b/examples/distributed/deepep_v2/test_example_deepep_v2.py @@ -0,0 +1,495 @@ +"""Correctness tests for the DeepEP-EPv2-aligned intranode dispatch/combine port. + +Small (nprocs=2) smoke plus a full 8-GPU run at DeepEP's own headline shape +(8K tokens, hidden=7168, top-8, 256 experts), for both dispatch payload dtypes +this port supports (bf16, fp8). Bf16 accumulates rounding error across topk +contributions, so correctness is judged by relative L2 error against the +identity-compute reference (see ``reference.py``), not a strict per-element +allclose -- a few outlier elements near a sign-cancelling zero can have a +large *relative* error while the overall reconstruction is fine. FP8 adds the +per-128-element quantisation step on top, hence the wider threshold. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.distributed as dist + +import tilelang.testing +from testing.python.distributed._utils import distributed_test + +from tilelang.distributed.host import init_dist + +from buffer import Buffer +import reference + +_BF16_REL_L2_THRESHOLD = 0.05 +# FP8 (e4m3, ~2 mantissa bits) quantises `x` before dispatch ever sees it, on +# top of the same bf16 accumulation error the plain threshold covers. +# Measured rel_l2 is a tight ~0.026 at every shape tried (smoke, masked, and +# the full V3 shape at both 2 and 8 GPUs, std well under 1e-3) -- consistent +# with a quantisation-dominated, not noise-dominated, error, so ~2x that +# measured value is real margin, not padding for run-to-run variance. +_FP8_REL_L2_THRESHOLD = 0.05 + + +def _run( + local_rank: int, + num_ranks: int, + num_tokens: int, + hidden: int, + topk: int, + num_experts: int, + num_sms: int, + masked_ratio: float = 0.0, + dtype: torch.dtype = torch.bfloat16, + num_bias: int = 0, + do_expand: bool = False, +): + rank, num_ranks, group = init_dist(local_rank, num_ranks) + + torch.manual_seed(1234 + rank) + device = f"cuda:{local_rank}" + x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) + topk_idx, topk_weights = reference.make_topk(num_tokens, topk, num_experts, device, masked_ratio) + + is_fp8 = dtype == torch.float8_e4m3fn + # Quantising is the caller's job (see buffer.py's `dispatch` docstring); + # `x` itself stays bf16 for the reference computation below. + dispatch_x = reference.per_token_cast_to_fp8(x) if is_fp8 else x + + buf = Buffer( + group=group, + local_rank=local_rank, + num_local_ranks=num_ranks, + num_max_tokens_per_rank=num_tokens, + hidden=hidden, + num_topk=topk, + num_experts=num_experts, + dtype=dtype, + num_sms=num_sms, + do_expand=do_expand, + expand_factor=float(min(topk, num_experts // num_ranks)) if do_expand else 1.0, + ) + try: + recv, recv_topk_idx, recv_topk_weights, handle, _ = buf.dispatch(dispatch_x, topk_idx, topk_weights) + n = handle.num_recv_tokens + recv_topk_idx, recv_topk_weights = recv_topk_idx[:n], recv_topk_weights[:n] + recv_x = reference.per_token_cast_back(recv[:n], hidden) if is_fp8 else recv[:n] + expert_out = reference.simulate_expert_compute(recv_x, recv_topk_idx, recv_topk_weights) + # Deliberately the same magnitude as the combined output, so a dropped + # or double-applied bias moves rel_l2 well past the threshold. + biases = [torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) for _ in range(num_bias)] + combined, _ = buf.combine(expert_out, handle, bias=biases or None) + + expected = reference.reference_combined(x, topk_weights, topk_idx) + for b in biases: + expected = expected + b + err = (combined.float() - expected.float()).norm().item() + denom = expected.float().norm().item() + # `denom` is zero only when every selection was masked off. + rel_l2 = err / denom if denom > 0 else err + threshold = _FP8_REL_L2_THRESHOLD if is_fp8 else _BF16_REL_L2_THRESHOLD + assert rel_l2 < threshold, f"rank {rank}: rel_l2_error={rel_l2} exceeds {threshold}" + finally: + buf.close() + dist.destroy_process_group() + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=2) +def test_dispatch_combine_smoke(local_rank: int, num_ranks: int): + _run(local_rank, num_ranks, num_tokens=64, hidden=128, topk=2, num_experts=8, num_sms=2) + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=2) +def test_dispatch_combine_smoke_fp8(local_rank: int, num_ranks: int): + _run(local_rank, num_ranks, num_tokens=64, hidden=128, topk=2, num_experts=8, num_sms=2, dtype=torch.float8_e4m3fn) + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=8) +def test_dispatch_combine_v3_shape(local_rank: int, num_ranks: int): + """DeepEP's own headline benchmark shape: 8K tokens, hidden=7168, top-8, 256 experts.""" + _run(local_rank, num_ranks, num_tokens=8192, hidden=7168, topk=8, num_experts=256, num_sms=64) + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=8) +def test_dispatch_combine_v3_shape_fp8(local_rank: int, num_ranks: int): + """DeepEP's own headline benchmark shape, FP8 dispatch payload.""" + _run(local_rank, num_ranks, num_tokens=8192, hidden=7168, topk=8, num_experts=256, num_sms=64, dtype=torch.float8_e4m3fn) + + +# One `_run` per test, like every other case here: `_run` owns the process +# group's whole lifetime, so looping over parameters inside one test tears it +# down and rebuilds it on the same port, which is intermittently refused. +@tilelang.testing.requires_cuda +@distributed_test(nprocs=2) +def test_combine_bias_single(local_rank: int, num_ranks: int): + """DeepEP's `bias_0`, on its own.""" + _run(local_rank, num_ranks, num_tokens=64, hidden=128, topk=2, num_experts=8, num_sms=2, num_bias=1) + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=2) +def test_combine_bias_pair(local_rank: int, num_ranks: int): + """DeepEP's `bias_0` and `bias_1` together.""" + _run(local_rank, num_ranks, num_tokens=64, hidden=128, topk=2, num_experts=8, num_sms=2, num_bias=2) + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=8) +def test_combine_bias_masked(local_rank: int, num_ranks: int): + """Bias with half the selections masked off: a token with no contributions + at all must still come back as its bias, not as zero.""" + _run(local_rank, num_ranks, num_tokens=1024, hidden=1024, topk=8, num_experts=256, num_sms=32, masked_ratio=0.5, num_bias=2) + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=8) +def test_wide_grid(local_rank: int, num_ranks: int): + """A grid wider than the shape needs, which is how a caller buys the last + few percent of dispatch throughput: 860.6 us at 128 SMs against 898.2 at 64. + + Every strided loop -- the token loop, the alignment-padding loop, the stats + scan -- has to agree on the grid it is striding over, and one left on a + stale stride drops or double-writes rows. + """ + _run(local_rank, num_ranks, num_tokens=1024, hidden=1024, topk=8, num_experts=256, num_sms=128) + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=8) +def test_cached_dispatch(local_rank: int, num_ranks: int): + """DeepEP's `handle=`: reuse a layout and skip the notify kernel. + + A cached dispatch must produce exactly what a fresh one does, twice over -- + the second reuse proves the first did not consume the layout. And a handle + whose layout has since been overwritten must be rejected: `send_base` and + `send_rank_mask` are updated in place, so a stale handle does not fail, it + quietly routes to the wrong slots. + """ + rank, num_ranks, group = init_dist(local_rank, num_ranks) + device = f"cuda:{local_rank}" + torch.manual_seed(1234 + rank) + + num_tokens, hidden, topk, num_experts = 1024, 1024, 8, 256 + x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) + topk_idx, topk_weights = reference.make_topk(num_tokens, topk, num_experts, device, 0.25) + expected = reference.reference_combined(x, topk_weights, topk_idx) + + buf = Buffer( + group=group, + local_rank=local_rank, + num_local_ranks=num_ranks, + num_max_tokens_per_rank=num_tokens, + hidden=hidden, + num_topk=topk, + num_experts=num_experts, + dtype=torch.bfloat16, + num_sms=16, + ) + + def roundtrip(reuse=None): + if reuse is None: + recv, recv_idx, recv_w, handle, _ = buf.dispatch(x, topk_idx, topk_weights) + else: + recv, recv_idx, recv_w, handle, _ = buf.dispatch(x, handle=reuse) + n = handle.num_recv_tokens + expert_out = reference.simulate_expert_compute(recv[:n], recv_idx[:n], recv_w[:n]) + combined, _ = buf.combine(expert_out, handle) + err = (combined.float() - expected.float()).norm().item() + return handle, err / expected.float().norm().item() + + try: + fresh, rel_fresh = roundtrip() + once, rel_once = roundtrip(fresh) + _, rel_twice = roundtrip(once) + for tag, rel in (("fresh", rel_fresh), ("cached", rel_once), ("cached twice", rel_twice)): + assert rel < _BF16_REL_L2_THRESHOLD, f"rank {rank}: {tag} rel_l2={rel} exceeds {_BF16_REL_L2_THRESHOLD}" + + roundtrip() # recomputes the layout, so `fresh` no longer describes the buffer + try: + buf.dispatch(x, handle=fresh) + except AssertionError: + pass + else: + raise AssertionError(f"rank {rank}: a handle whose layout was overwritten was accepted") + finally: + buf.close() + dist.destroy_process_group() + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=8) +def test_expanded_round_trip(local_rank: int, num_ranks: int): + """The expanded layout end to end, dispatch through combine. + + A token with two experts on one rank has two received rows, and combine has + to sum them before sending -- one row per (rank, token) is all + `comm_x` has slots for. topk=8 over 8 ranks makes collisions common, so + the summing path is exercised rather than incidentally avoided. + """ + _run(local_rank, num_ranks, num_tokens=1024, hidden=1024, topk=8, num_experts=256, num_sms=16, do_expand=True) + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=8) +def test_expanded_round_trip_masked(local_rank: int, num_ranks: int): + """Same, with half the selections unset and the segments aligned, so the + padding rows dispatch writes have to be skipped by the grouping pass.""" + _run( + local_rank, + num_ranks, + num_tokens=512, + hidden=512, + topk=8, + num_experts=256, + num_sms=16, + masked_ratio=0.5, + do_expand=True, + ) + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=8) +def test_dispatch_combine_masked(local_rank: int, num_ranks: int): + """Half the selections unset: dispatch must route -1 nowhere, and a token + with no selections at all must combine back to zero.""" + _run(local_rank, num_ranks, num_tokens=1024, hidden=1024, topk=8, num_experts=256, num_sms=32, masked_ratio=0.5) + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=8) +def test_cumulative_local_expert_recv_stats(local_rank: int, num_ranks: int): + """Per-expert receive counts, against an all-gathered torch reference. + + Three dispatches into the same counter, and a quarter of the top-k entries + masked off, so the test covers the accumulate-don't-overwrite contract and + the -1 entries the scan has to skip. + """ + rank, num_ranks, group = init_dist(local_rank, num_ranks) + device = torch.device(f"cuda:{local_rank}") + torch.manual_seed(1234 + rank) + + num_tokens, hidden, topk, num_experts, num_calls = 512, 512, 4, 32, 3 + experts_per_rank = num_experts // num_ranks + x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) + topk_idx, topk_weights = reference.make_topk(num_tokens, topk, num_experts, device, 0.25) + + buf = Buffer( + group=group, + local_rank=local_rank, + num_local_ranks=num_ranks, + num_max_tokens_per_rank=num_tokens, + hidden=hidden, + num_topk=topk, + num_experts=num_experts, + dtype=torch.bfloat16, + num_sms=8, + ) + try: + stats = torch.zeros(experts_per_rank, dtype=torch.uint32, device=device) + for _ in range(num_calls): + buf.dispatch(x, topk_idx, topk_weights, cumulative_local_expert_recv_stats=stats) + + # What every rank sent to each expert, summed -- my counters are the + # columns belonging to my own experts. + sent = torch.zeros(num_experts, dtype=torch.int64, device=device) + selected = topk_idx[topk_idx >= 0] + sent.scatter_add_(0, selected.long(), torch.ones_like(selected, dtype=torch.int64)) + per_rank = [torch.zeros_like(sent) for _ in range(num_ranks)] + dist.all_gather(per_rank, sent, group) + expected = torch.stack(per_rank).sum(0)[rank * experts_per_rank : (rank + 1) * experts_per_rank] * num_calls + + got = stats.to(torch.int64) + assert torch.equal(got, expected), f"rank {rank}: got {got.tolist()}, expected {expected.tolist()}" + finally: + buf.close() + dist.destroy_process_group() + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=8) +def test_dispatch_expanded_layout(local_rank: int, num_ranks: int): + """DeepEP's `do_expand`: one row per (token, expert), grouped by expert. + + Checked against `reference.expanded_layout` -- per-expert counts, the + aligned segment offsets, the exact set of (src_rank, src_token) in each + segment, the payload of every row, and that the alignment padding is + zeroed and marked unoccupied. + """ + rank, num_ranks, group = init_dist(local_rank, num_ranks) + device = torch.device(f"cuda:{local_rank}") + torch.manual_seed(1234 + rank) + + num_tokens, hidden, topk, num_experts, alignment = 128, 256, 4, 32, 8 + experts_per_rank = num_experts // num_ranks + # Distinct experts per token, which is what DeepEP asserts and real top-k + # routing produces; -1 marks an unselected slot. + idx = torch.stack([torch.randperm(num_experts, device=device)[:topk] for _ in range(num_tokens)]).int() + idx = idx.masked_fill(torch.rand_like(idx, dtype=torch.float) < 0.25, -1) + weights = torch.rand(num_tokens, topk, device=device) + # A token's identity, so a misrouted row is obvious rather than plausible. + x = torch.arange(num_tokens, device=device, dtype=torch.bfloat16).view(-1, 1).repeat(1, hidden) + x = x + rank * 1000 + + buf = Buffer( + group=group, + local_rank=local_rank, + num_local_ranks=num_ranks, + num_max_tokens_per_rank=num_tokens, + hidden=hidden, + num_topk=topk, + num_experts=num_experts, + dtype=torch.bfloat16, + num_sms=8, + do_expand=True, + expert_alignment=alignment, + expand_factor=float(min(topk, experts_per_rank)), + ) + try: + all_idx = [torch.zeros_like(idx) for _ in range(num_ranks)] + dist.all_gather(all_idx, idx, group) + exp_rows, exp_counts, exp_offsets = reference.expanded_layout([t.cpu() for t in all_idx], num_experts, num_ranks, alignment) + + recv_x, _, _, handle, _ = buf.dispatch(x, idx, weights) + assert handle.expand_overflow == 0, f"rank {rank}: overflowed by {handle.expand_overflow}" + + counts = handle.expert_count.cpu().tolist() + offsets = handle.expert_offset.cpu().tolist() + assert counts == exp_counts[rank], f"rank {rank}: counts {counts} != {exp_counts[rank]}" + assert offsets == exp_offsets[rank], f"rank {rank}: offsets {offsets} != {exp_offsets[rank]}" + + src_rank = buf.recv_src_rank.cpu() + src_token = buf.recv_src_token.cpu() + rows = recv_x.cpu() + for e in range(experts_per_rank): + begin, end = offsets[e], offsets[e] + counts[e] + got = sorted(zip(src_rank[begin:end].tolist(), src_token[begin:end].tolist())) + assert got == sorted(exp_rows[rank][e]), f"rank {rank} expert {e}: {got} != {sorted(exp_rows[rank][e])}" + # Every row carries its origin, so the payload pins down routing. + for i in range(begin, end): + want = float(src_token[i]) + float(src_rank[i]) * 1000 + assert torch.all(rows[i] == want), f"rank {rank} row {i}: payload {rows[i][0]} != {want}" + # Alignment padding: zeroed, and owned by nobody. + for i in range(end, offsets[e + 1]): + assert torch.all(rows[i] == 0), f"rank {rank} pad row {i} not zeroed" + assert src_rank[i] == -1 and src_token[i] == -1, f"rank {rank} pad row {i} marked occupied" + finally: + buf.close() + dist.destroy_process_group() + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=8) +def test_async_finish_round_trip(local_rank: int, num_ranks: int): + """The same round trip run asynchronously must give the same answer. + + Both collectives run with `async_finish=True` and real work is enqueued on + the compute stream inside the `with event:` block, so the test fails if the + event does not actually order the two streams -- reading `recv_x` before + the dispatch landed gives whatever the previous iteration left. + + `allocate_on_comm_stream=True` on the dispatch exercises the path that + keeps this call's temporaries alive through the event rather than + `record_stream`, and `previous_event` chains the combine behind that same + event instead of behind the whole compute stream. + """ + rank, num_ranks, group = init_dist(local_rank, num_ranks) + device = f"cuda:{local_rank}" + torch.manual_seed(1234 + rank) + + num_tokens, hidden, topk, num_experts = 512, 512, 4, 32 + x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) + topk_idx, topk_weights = reference.make_topk(num_tokens, topk, num_experts, device, 0.25) + + buf = Buffer( + group=group, + local_rank=local_rank, + num_local_ranks=num_ranks, + num_max_tokens_per_rank=num_tokens, + hidden=hidden, + num_topk=topk, + num_experts=num_experts, + dtype=torch.bfloat16, + num_sms=8, + # Otherwise every third dispatch blocks the host, which is not what a + # caller driving its own overlap asked for. See `Buffer.pipeline_depth`. + pipeline_depth=0, + ) + try: + recv, recv_topk_idx, recv_topk_weights, handle, dispatch_event = buf.dispatch( + x, topk_idx, topk_weights, async_finish=True, allocate_on_comm_stream=True + ) + # Something real on the compute stream, overlapping the dispatch. + filler = torch.randn(1024, 1024, device=device, dtype=torch.bfloat16) + with dispatch_event: + filler = filler @ filler + + n = handle.num_recv_tokens + expert_out = reference.simulate_expert_compute(recv[:n], recv_topk_idx[:n], recv_topk_weights[:n]) + combined, combine_event = buf.combine( + expert_out, handle, previous_event=dispatch_event, async_finish=True, allocate_on_comm_stream=True + ) + combine_event.current_stream_wait() + + expected = reference.reference_combined(x, topk_weights, topk_idx) + err = (combined.float() - expected.float()).norm().item() + denom = expected.float().norm().item() + rel_l2 = err / denom if denom > 0 else err + assert rel_l2 < _BF16_REL_L2_THRESHOLD, f"rank {rank}: rel_l2_error={rel_l2} exceeds {_BF16_REL_L2_THRESHOLD}" + assert filler.isfinite().all(), f"rank {rank}: overlapped work was corrupted" + finally: + buf.close() + dist.destroy_process_group() + + +@tilelang.testing.requires_cuda +@distributed_test(nprocs=2) +def test_event_overlap_is_returned_when_synchronous(local_rank: int, num_ranks: int): + """Synchronous calls still return an `EventOverlap`, wrapping `None`. + + That is what lets a caller write `with event:` without knowing which mode + it asked for, so it is part of the contract rather than an accident. + """ + rank, num_ranks, group = init_dist(local_rank, num_ranks) + device = f"cuda:{local_rank}" + torch.manual_seed(1234 + rank) + + num_tokens, hidden, topk, num_experts = 64, 128, 2, 8 + x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) + topk_idx, topk_weights = reference.make_topk(num_tokens, topk, num_experts, device) + + buf = Buffer( + group=group, + local_rank=local_rank, + num_local_ranks=num_ranks, + num_max_tokens_per_rank=num_tokens, + hidden=hidden, + num_topk=topk, + num_experts=num_experts, + dtype=torch.bfloat16, + num_sms=2, + ) + try: + recv, recv_topk_idx, recv_topk_weights, handle, event = buf.dispatch(x, topk_idx, topk_weights) + assert event.event is None + with event: # a no-op, but it must not raise + pass + n = handle.num_recv_tokens + expert_out = reference.simulate_expert_compute(recv[:n], recv_topk_idx[:n], recv_topk_weights[:n]) + _, combine_event = buf.combine(expert_out, handle) + assert combine_event.event is None + with pytest.raises(AssertionError): + combine_event.current_stream_wait() + finally: + buf.close() + dist.destroy_process_group() + + +if __name__ == "__main__": + tilelang.testing.main() diff --git a/src/tl_templates/cuda/distributed/copy.h b/src/tl_templates/cuda/distributed/copy.h index 15a8c2bd58..7dc830bed8 100644 --- a/src/tl_templates/cuda/distributed/copy.h +++ b/src/tl_templates/cuda/distributed/copy.h @@ -44,6 +44,18 @@ TL_DEVICE void cp_warp_impl(dtype_t const *const dst_addr, for (int __i = (N_int4 / kLoopStride) * kLoopStride + lane_id; __i < N_int4; __i += 32) __dst[__i] = __src[__i]; + // Tail: `N` elements of `dtype_t` are not in general a whole number of + // int4s -- e.g. a single fp32 scale (`N=1`, 4 bytes) has `N_int4 == 0`, so + // both loops above run zero iterations and silently copy nothing at all. + // `threadgroup_cp` (the block-level copy just below) already falls back + // through progressively narrower element types down to a byte; here `N` + // and `dtype_t` are compile-time, so the remainder is just the last + // `N - N_int4 * kElemsPerInt4` elements, copied directly through the + // original (unreinterpreted) pointers. + constexpr int kElemsPerInt4 = sizeof(int4) / sizeof(dtype_t); + constexpr int kTailStart = N_int4 * kElemsPerInt4; + for (int __i = kTailStart + lane_id; __i < N; __i += 32) + const_cast(dst_addr)[__i] = src_addr[__i]; } template op.same_as(tl::access_ptr())) { return VisitAccessPtrCall(op); } + if (op->op.same_as(builtin::address_of())) { + return VisitAddressOfCall(op); + } PrimExpr expr = IRMutatorWithAnalyzer::VisitExpr_(op); const auto *call_node = expr.as(); @@ -320,6 +323,21 @@ class SafeMemorysRewriter : public IRMutatorWithAnalyzer { return if_then_else(CombineConditions(conditions), call, safe_value); } + // `address_of` yields a pointer, not a value. Rewriting its argument into + // `if_then_else(cond, load, safe)` would ask for the address of a + // conditional, which every consumer rejects ("address_of argument must be a + // BufferLoad"). Visit the indices so nested rewrites still happen, but keep + // the argument a BufferLoad; whoever dereferences the pointer is responsible + // for the bound, exactly as for `tl.access_ptr` above. + PrimExpr VisitAddressOfCall(const CallNode *op) { + ICHECK_EQ(op->args.size(), 1U) << "address_of expects 1 arg"; + auto visit_expr = [this](const PrimExpr &expr) { + return this->VisitExpr(expr); + }; + Array args{detail::VisitAccessPtrBase(op->args[0], visit_expr)}; + return Call(op->dtype, op->op, args, op->annotations, op->span); + } + PrimExpr VisitAccessPtrCall(const CallNode *op) { ICHECK_EQ(op->args.size(), 3U) << "tl.access_ptr expects 3 args: (BufferLoad, extent, rw_mask)"; diff --git a/testing/python/distributed/primitives/comm/test_put_get.py b/testing/python/distributed/primitives/comm/test_put_get.py index 521d2fe03d..44f302f4c1 100644 --- a/testing/python/distributed/primitives/comm/test_put_get.py +++ b/testing/python/distributed/primitives/comm/test_put_get.py @@ -212,6 +212,78 @@ def test_put_accepts_constant_size(): assert "tl::cp_warp<128," in source +def _kernel_put_warp_small(n): + @T.prim_func + def main(dst: T.Tensor((n,), "float32"), src: T.Tensor((n,), "float32")): + with T.Kernel(1, threads=32): + rank = T.alloc_local((1,), "uint64") + rank[0] = T.get_rank() + T.put_warp(T.address_of(src[0]), T.address_of(dst[0]), n, dst_pe=rank[0] ^ 1) + + return main + + +def _kernel_get_warp_small(n): + @T.prim_func + def main(dst: T.Tensor((n,), "float32"), src: T.Tensor((n,), "float32")): + with T.Kernel(1, threads=32): + rank = T.alloc_local((1,), "uint64") + rank[0] = T.get_rank() + T.get_warp(T.address_of(src[0]), T.address_of(dst[0]), n, src_pe=rank[0] ^ 1) + + return main + + +@tilelang.testing.requires_cuda_compute_version_ge(9, 0) +@distributed_test(nprocs=2) +def test_put_get_warp_small_size(local_rank: int, num_ranks: int): + """`cp_warp`'s bulk loop only ever moves whole 16-byte (`int4`) units. + + `n=1..3` (`float32`, so 4-12 bytes total) used to hit `N_int4 == 0`: both + the unrolled and the drain loop run zero iterations, so the whole transfer + -- not just a misaligned tail -- silently copied nothing. `n=5` (20 bytes) + exercises the more general form of the same bug: one full `int4` copied, + the trailing element dropped. This is the exact shape deepep_v2's FP8 + per-token scale (`scale_dim=1`, a single `float32`) hit: `recv_x_scales` + came back all zero, silently, with no error anywhere. + """ + from tilelang.distributed.host import init_dist + + rank, num_ranks, group = init_dist(local_rank, num_ranks) + allocator = tilelang.get_allocator( + size=2**20, + device="cuda", + is_distributed=True, + local_rank=local_rank, + num_local_ranks=num_ranks, + group=group, + ) + + for n in (1, 2, 3, 5): + for name, kernel_fn in [("put_warp", _kernel_put_warp_small), ("get_warp", _kernel_get_warp_small)]: + kernel = tilelang.compile(kernel_fn(n), compile_once=True, compile_group=group) + kernel.initialize(allocator=allocator) + + src = tilelang.tensor((n,), torch.float32, allocator=allocator).uniform_(1.0, 2.0) + dst = tilelang.tensor((n,), torch.float32, allocator=allocator).zero_() + + torch.cuda.synchronize() + dist.barrier(group) + kernel(dst, src) + torch.cuda.synchronize() + dist.barrier(group) + + dst_refs = [torch.empty_like(src) for _ in range(num_ranks)] + dist.all_gather(dst_refs, src, group) + expected = dst_refs[local_rank ^ 1] + assert torch.allclose(expected, dst, atol=1e-6, rtol=1e-6), ( + f"rank {local_rank}: {name} n={n} mismatch, got {dst.tolist()}, expected {expected.tolist()}" + ) + + allocator.close() + dist.destroy_process_group() + + def _peer_none_kernel(): @T.prim_func def main(A: T.Tensor((16,), "float32"), B: T.Tensor((16,), "float32"), flag: T.Tensor((1,), "uint32")):