a2a_gemm examples with flydsl + mori

[zjing14]

Self-contained example demonstrating how to overlap an All-to-All collective with
the bf16 GEMM that consumes it, using device-initiated SDMA (mori putmem_nbi_signal)
so the transfer runs on the CU-free DMA copy engines while the GEMM runs on the CUs.

Models a sequence-parallel "A2A -> projection" step (P = ranks): each rank holds
[world*sw, Hp], an All-to-All redistributes it to chunk-major a_full[P, sw, Hp],
then a split-K projection GEMM computes out = sum_i a_full[i] @ W[i].T. Compares
unfused (nccl), seq (SDMA, no overlap), and fused (SDMA + per-source gate+GEMM
overlap), with a correctness check.

The FlyDSL SDMA producer + per-source gate kernels are bundled in the single file
(no external project deps beyond the flydsl/mori/aiter runtime). Generic dims only.

Ran 4-rank on a gfx950 node via the flydsl/mori/aiter runtime:

  HIP_VISIBLE_DEVICES=0,1,2,3 LD_LIBRARY_PATH=/opt/rocm/lib MORI_ENABLE_SDMA=1 \
    MORI_SHMEM_HEAP_SIZE=16G MORI_SOCKET_IFNAME=lo torchrun --nproc_per_node=4 \
    scripts/jizhan/a2a_gemm/a2a_gemm_example.py --hidden 8192 --seqs 8192 16384

Output (correct=True both sizes):

  8192  | unfused_nccl 0.658 | unfused_mori 0.786 | fused 0.700 | fused/mori 1.12x | fused/nccl 0.94x
  16384 | unfused_nccl 1.261 | unfused_mori 1.332 | fused 1.285 | fused/mori 1.04x | fused/nccl 0.98x

python3 -m py_compile passes; arc lint -a clean (BLACK).

unfused_nccl_gemm        1.267 ms
unfused_sdma_gemm        1.304 ms
overlapped_sdma_gemm        1.292 ms  (0.98× nccl)
unfused_sdma_single_gemm 1.172 ms  (1.08×)
unfused_nccl_single_gemm 1.122 ms  (1.13×)  ← fastest

[zjing14]

The SDMA+GEMM overlap rests on one hardware fact: on gfx950 the SDMA copy engines are separate hardware from the compute units (CUs). So a data transfer can run on the DMA engine at the same time as a GEMM runs on the CUs — if you can (a) launch the transfer without occupying CUs, and (b) let the GEMM start consuming data before the whole transfer finishes.

1. CU-free transfer (the enabler)

putmem_nbi_signal_block (mori) issues an SDMA copy: a device kernel enqueues a copy descriptor + a signal (flag increment) onto the SDMA queue, then returns immediately (nbi = non-blocking). The actual bytes move on the DMA engine, not the CUs. So after the producer kernel exits, the CUs are free — the transfer keeps going in the background. (nccl all_to_all, by contrast, runs a copy on the CUs, so it can't overlap a CU-bound GEMM.)

2. The flag handshake (correctness + pipelining)

The transfer and the GEMM are decoupled, so they synchronize through flags:

• The SDMA queue does copy-then-signal in order → a flag is set only after its data has landed in HBM.
• The GEMM waits on that flag (uint64_wait_until_greater_than, SYSTEM-scope read since a remote SDMA wrote it). Once it passes, the data is guaranteed present.
• Flags are per-chunk (per source, or per row-block), so the GEMM can start on chunk i the instant it lands — while chunks i+1… are still transferring.

That per-chunk granularity is what turns "transfer, then compute" into an overlapped pipeline:

DMA engine:  chunk0 → flag0 │ chunk1 → flag1 │ chunk2 → flag2 …   (background, CU-free)
CUs (GEMM):             wait0→compute0 │ wait1→compute1 │ …
                         └── compute(i) overlaps the in-flight transfer of chunk i+1.. ──┘

3. Two ways to express the wait

• Two-kernel gate (unfused_mori / the old fused): a tiny separate chunk_gate kernel spins on the flag on the stream, then the GEMM kernel runs. Works with a stock GEMM, but the gate is a separate launch.
• In-kernel gate (fused_ingrid): the wait is compiled inside the GEMM (gated_hgemm) — each M-tile spins on its own row-block flag before its A-load. One kernel, per-tile gating, no separate launches. This is the tight form.

Either way, the overlap is the DMA engine ∥ CUs, not two kernels racing.

4. When it actually pays (the limits we measured)

Overlap hides min(comm, compute) and only the middle of the pipeline (the first chunk's arrival and the last tile's compute are always exposed). So it wins when:

• comm is CU-free (SDMA, ✓) — otherwise comm and compute fight over CUs;
• compute ≥ comm so the GEMM can absorb the transfer (bf16 ✓; fp4 shrinks the GEMM below comm → little to hide);
• the transfer is big enough that SDMA beats nccl's low fixed overhead (crossover ~25 MB — A2A sits right at it, so overlap only ties nccl there; MoE dispatch is far bigger, so it pays).

In one sentence

Issue the transfer on the CU-free SDMA engine and return; the GEMM computes each tile as its per-chunk flag fires, so compute runs on the CUs while the remaining data streams in on the DMA engine — bounded by whether the comm is CU-free, whether the GEMM is big enough to hide it, and whether the transfer is large enough to beat nccl's overhead.

[jhchouuu]

Hi @zjing14 , just to comfirm, do you based on this PR #445 to enable SDMA use putmem with signal？

[jhchouuu]

Thanks for the great work! Using device-initiated SDMA (putmem_nbi_signal) to run the A2A on the CU-free copy engines while the GEMM keeps the CUs busy is a really nice idea, and the in-kernel per-tile gate is a clean way to express it.

On why the overlap gain looks marginal here — I think three things compound:

1. Small A2A volume: only 4 GPUs in a full-mesh (no switch) topology, so bandwidth is mesh-limited and the total transfer sits below the SDMA/nccl crossover — nccl already wins the comm.
2. Split-K dilutes the per-tile gate (P small GEMMs instead of one large gated GEMM).
3. Small per-copy size (~1MB, serialized on one queue) keeps SDMA below peak BW. Our SDMA engine has a relatively long issue latency.

A large-GEMM or ≥8-GPU data point would help show where it actually pays off.

One impl note: the SDMA transport is currently selected at runtime (transportTypes[pe] branch), not compile-time — for a fine-grained per-tile push loop that runtime dispatch may add hot-path overhead. On our side we plan to use CCO (shmem v2) for explicit compile-time calls to remove this kind of overhead.

I think this is a valuable reference for the fused-comm pattern.
cc @carlushuang