Draft: cco + flydsl dispatch/combine#448
Draft
jhchouuu wants to merge 18 commits into
Draft
Conversation
Single-node EP8 MoE dispatch + combine built on mori-cco LSA (intra-node P2P over the flat symmetric VA) and FlyDSL device kernels. mori-parity algorithm with the shmem provider swapped for cco-LSA: peer addresses computed in-kernel via cco.Window(arena).lsa_ptr(pe, off), no host P2P tables. bf16 basic path. - SymmArena: one cco symmetric window carved into named sub-regions. - dispatch: ballot dedup + remote atomic slot alloc + 16B dual-issue scatter, grid barrier + per-peer count signal. - combine: P2P-read (gather) each token's k expert outputs back and reduce in f32; register-light i32 reads at high warp occupancy to hide xGMI read latency. - self-contained volatile/atomic spin-waits; counters self-reset (CUDAGraph-safe). - torchrun harness; correctness tests (vs CPU model) + eager/graph perf bench. EP8 hidden=7168 top-k=8: dispatch ~318 GB/s, combine ~318 GB/s per rank @8192 tok (CUDA-graph), near MI300X xGMI saturation. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Combine now reduces the K forwarded per-token weight vectors into out_weights[tok][e] = sum over valid experts of the forwarded weight, matching mori intranode's UseWeights path. Folded into the Stage-2 gather loop (once per token at part 0) so it overlaps hidden gather with no bandwidth loss (143.9 GB/s @128 tok, unchanged). Bench verifies out_weights == U[t]*wts[t] and adds a run helper. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Port mori/FlyDSL StdMoE to the cco-LSA v2 path as two local convert
kernels (the mori op-layer also exposes them standalone):
* convert_dispatch_output: per-(recv_tok,k) repack of received tokens
into per-local-expert buckets packed_x[localExpert,slot]; records
packed_cnt/packed_src and the i64 disp_tok_map slot map (ult local
test, -1 sentinel for non-local experts), matching the reference.
* convert_combine_input: weighted reduce of each recv token's local
expert outputs (sum_k wts[k]*packed_x[slot_k]) back into out_tok,
which the existing P2P-read combine then sums across dest PEs.
End to end with identity experts this telescopes to
comb_out[s] == (sum_k wts[s][k]) * input[s]; bench verifies EP8 at
128/512/2048 tok/rank. Fixes a latent bench bug: flyc.compile launches
the kernel once to specialize, so total_recv is zeroed before compiling
the convert kernels (else the compile-launch scribbles into out_tok).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Combine accumulators are now selected per transport dtype via _accum_funcs (mirrors the FlyDSL reference _to_accum/_from_accum/ _zero_accum branches): bf16 (i32=2bf16, v2f32), f32 (i32=1f32, scalar), fp8 (i32=4fp8, v4f32 via cvt_pk_f32_fp8). The dispatch token copy is already byte-generic, so f32 works by passing hidden_elem_size=4. Bench gains a DTYPE env (bf16|f32); EP8 f32 PASS @128/512 (254/202 GB/s). fp8 helpers (incl. fp8_direct_cast bf16-output stride) are in place on the combine side but untested end-to-end: the gather path has no Stage-1, so wire-fp8 needs dispatch to write fp8 into out_tok (TODO). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
mori-parity host op (cf. EpDispatchCombineOp / FlyDSLDispatchCombineIntraNodeOp) for the cco-LSA path: owns one SymmArena window + per-rank metadata, compiles dispatch/combine, and surfaces arena regions as torch tensors via from_gpu_ptr. out, total_recv, routing = op.dispatch(input, weights, indices) combined, combined_wts = op.combine(out, routing) # gather/UseP2PRead EpDispatchRoutingHandle caches (token,k)->dest map + recv count so combine runs without re-dispatching. EpDispatchCombineConfig carries the mori-parity geometry. test_op.py: EP8 identity round-trip PASS for bf16 and f32 (comb==U*input, out_weights==U*wts) at 128/512 tok/rank. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…p8DirectCast) Add make_combine_scatter: the mori useExternalInpBuffer/_nop2p combine — Stage 1 P2P-writes each post-expert token back to the origin rank's comb_inp[computing_rank*M + origin_lid] (origin decoded from tis), Stage 2 cross-device barrier, Stage 3 local-reads comb_inp[dest_pe*M + tok] per the token's k expert PEs and reduces. Two passes vs gather's one, but it provides the Stage-1 writer that fp8_direct_cast needs: under fp8_direct_cast the bf16 token is cast to fp8 on the wire (cvt_pk_fp8_f32) and dequantized to bf16 in Stage 3, halving transport bytes. Bench gains COMBINE=gather|scatter and QUANT=none|fp8_direct_cast. EP8 PASS: scatter bf16 (125-179 GB/s, 2-pass) and scatter fp8 (173-311 GB/s bf16-equiv, half-byte wire) at 128/512 tok/rank. Default gather path unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ccoMemAlloc's VMM windows now allocate as hipMemAllocationTypeUncached by default (was hipMemAllocationTypePinned / coarse-grained), matching mori-shmem's hipDeviceMallocUncached symmetric heap: P2P remote reads/writes stay coherent over the fabric without coarse-grained L2 coherence handling. CcoWindowAllocType() is shared by the granularity query and the hipMemCreate call so they stay in sync. Opt out with CCO_UNCACHED_WINDOW=0 (reverts to coarse-grained pinned). Validated 8-rank: cco test_host 8/8, LSA barrier/allreduce/memcheck 8/8, GDA put/get/barrier pass; ops v2 dispatch/combine (gather bf16/f32, scatter bf16/fp8) all correct. Bandwidth unchanged (gather 315/320 GB/s). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
EpDispatchCombineOp gains enable_std_moe (convert_dispatch_output / convert_combine_input methods reusing the standalone StdMoE kernels: dispatch -> repack per-expert -> [GEMM] -> weighted inverse -> combine) and max_total_recv_tokens (effective_max_recv caps recv slots + buffer sizing; over-cap tokens take the dup-sentinel drop path, no OOB). test_op.py: EP8 OP-STDMOE PASS @128/512 (sum-weighted identity), OP-CAP PASS (recv capped 2735->2048, no OOB). Standard op path unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
make_combine_scatter now supports enable_weights via a dedicated off_comb_wts region: Stage 1 forwards each recv slot's weights (out_wts[rt]) to the origin's comb_wts[computing_rank*M+lid]; Stage 3b reduces comb_wts[dest_pe*M+tok] over the token's k expert PEs into out_weights. Reusing out_wts would collide with dispatch's per-recv-slot layout. EP8 scatter bf16 + fp8 now PASS wts=ok. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add tuning_configs.lookup(world_size, hidden_dim, topk) -> measured-good
{dispatch_block_num, combine_block_num, warp_num_per_block} (mori
tuning_configs analogue; falls back to defaults). EpDispatchCombineConfig.tuned()
pulls geometry from the table unless explicitly overridden.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add the fp4 e2m1 accum path to _accum_funcs (i32 = 8 packed fp4 -> v8f32
via cvt_scalef32_pk_f32_fp4 / cvt_scalef32_pk_fp4_f32) and a fp4 flag on
make_combine (nbytes=hidden/2, n_i32=hidden/8). Faithful port of the FlyDSL
reference fp4 branch.
These cvt_scalef32_pk_*_fp4 intrinsics are gfx950 (MI350) only — codegen
fails "instruction not supported on this GPU" on this gfx942 (MI300X) node,
and torch cannot create float4_e2m1fn_x2 tensors here ("copy_kernel not
implemented"), so the path is opt-in and not e2e-testable on this hardware.
Default (non-fp4) path unchanged and still PASS.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
make_dispatch gains scale_dim/scale_type_size/off_out_scales: each published token's scale block (scale_n_i32 dwords) is forwarded verbatim to the dest peer's out_scales[dest_tok_id], lane-strided to cover scale_dim > one wavefront (mori intranode scales copy; combine consumes none). New addr_inp_scales kernel arg (callers pass 0 when disabled). Bench SCALE_DIM env + verify: each recv slot's scale block == its origin token's pattern (src_pe*100003+src_lid, origin from tis). EP8 SCALES PASS (scale_dim=32, recv=2735). Default path + op-layer unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…tuning Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Port mori's Fp8BlockwiseQuant to the scatter combine: Stage 1 quantizes the post-expert bf16 token to blockwise fp8 (block_elems = hidden/scale_dim; default scale_dim=hidden/128) with one block per lane (scale_dim<=64 -> serial amax, no warp reduce): scale = (amax>MAX)? amax/MAX : 1, quant = clamp(v*MAX/amax), per-block f32 scale scattered to the origin's comb_scales; token-level sign sentinel negates block-0 scale if any block scaled. Stage 3 dequantizes (fp8 * per-block scale, sb = elem//block_elems, block-0 abs) and accumulates to bf16. KEY: gfx942 native fp8 is e4m3fnuz (max finite 240, NOT OCP 448) -> clamp at 240 or cvt_pk_fp8_f32 yields NaN for values in (240,448]. EP8 PASS @128/512/2048 (median rel err ~2%). Quantize is lane-serial (correctness-first; slower than mori's subwarp-vec path). Bench QUANT=fp8_blockwise + BW_INSCALE. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace the lane-serial per-block quantize (lane l serially scans block l's 128 elems with stride-64 uncoalesced loads) with a fused coalesced path (block_elems==128): per block, all 64 lanes load the block coalesced (1 i32/lane = elems 2l,2l+1), a ds_bpermute butterfly gives every lane the block amax (readlane needs a uniform lane, ds_bpermute allows the per-lane butterfly index), then lane-pairs combine their 2 fp8 each into one i32 and even lanes store coalesced. _warp_amax helper added. EP8 blockwise combine: 2048 tok 28 -> 288 GB/s (eager), 8192 -> 296 GB/s; correctness unchanged (median rel ~2%). fp8_direct_cast + bf16 paths unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…e, routing handle, LEC, reset, replay Bring the intranode EpDispatchCombineOp up to mori parity and fix a scatter race: - dispatch(input, weights, scales, indices, *, routing, return_routing): forward per-token quant scales; return (out, out_weights, out_scales, out_indices, total_recv[, routing]); recv_weights/recv_indices/recv_scales. - EpDispatchRoutingHandle: full 5-tensor mori shape (+tensors()/from_tensors); inter-node maps are empty placeholders (intranode-only). - combine(input, weights, indices, *, routing): signature parity. - scatter combine + quant_type (fp8_direct_cast / fp8_blockwise) wired into the op (comb_inp/comb_wts/comb_scales regions, combine_mode gather|scatter). - local_expert_count() + kernel; reset() (mori LaunchReset); replay dispatch (make_dispatch replay= reuses cached tok_map, identical layout). - symm_arena: alloc_window -> alloc_mem + register_window (cco API drift). - fix: scatter combine missing fence_system_release after Stage-1 P2P writes caused a cross-device race (nondeterministic dropped contributions). EP8 verified (gfx942): gather/scatter bf16, fp8_blockwise, scales, StdMoE, local_expert_count, reset, replay all pass. Co-authored-by: Cursor <cursoragent@cursor.com>
jhchouuu
force-pushed
the
jiahzhou/cco-dispatch-combine
branch
from
f0a7790 to
5105c09
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist