[yangyuhuiling]

What

wo_b in DeepSeek-V4 is a RowParallelLinear: a block-scale fp8 GEMM followed
by an NCCL all-reduce of the [M, hidden] result. The two can overlap.
mori.ops.gemm_ar's fused kernel writes each destination's row band straight
into a symmetric window and hands it to the SDMA copy engines as soon as that
band's tiles are done, so the transfer runs inside the GEMM instead of after
it.

Opt-in and ROCm gfx950 only. SGLANG_OPT_FUSED_WO_B_AR=1 turns it on.
fused_wo_b() returns None for everything it does not serve — decode, small
or ragged M, a weight that is not aiter-preshuffled, no SDMA queues — and
disables itself for the process on any failure, so the ordinary path is the
fallback rather than an error.

Numbers

8x MI355X (gfx950), DeepSeek-V4-Pro, TP8, input 4096 / output 512,
bench_one_batch_server, last_ttft:

bs	base	fused	fused + fp8 gather
8	1136.1 ms	−3.36%	−6.04%
16	2232.0 ms	−3.03%	−6.29%
64	8890.9 ms	−3.26%	−6.44%

bs=1 sits at the fusing threshold (one row band per destination, nothing to
overlap) and is noise in either direction.

The layer itself goes 1460.3us → 1145.1 (−21.6%) at M=16384 N=7168 K=2048
on 8 ranks. End-to-end is smaller because wo_b is about 12% of the forward —
TTFT is the metric that moves, since the fused path only engages in prefill;
decode runs at M=batch_size and never fuses.

Fused and unfused agree to relL2 2.72e-05. Both quantise identically, so
what is left is bf16 accumulation order.

The optional fp8 gather, and what it costs

SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER=1 (also off by default) sends the
all-gather leg as e4m3 with a per-row scale. That leg grows linearly with M,
which is why the extra gain is near-constant across batch sizes (−2.8 / −3.4 /
−3.3% on top of the bf16 wire).

It is not free: relL2 on that leg goes 2.35e-3 → 2.49e-2, which is e4m3's
3-bit mantissa rather than a granularity that can be tuned away, and the
symmetric window grows 700 → 812 MiB. Off by default for that reason.

Worth a reviewer's attention

• Two call sites own the all-reduce themselves and the fused path stands
  aside for both: mHC post-fusion passes skip_all_reduce=True, and a
  Mxfp8SwizzledInput arrives when wo_a fuses an MXFP8 quantization into its
  epilogue. Neither is reachable on ROCm today (mHC needs Blackwell and
  tp_size == attn_tp_size == 4), so this checks rather than relying on the
  platform.
• BUILD_CCO_SDMA=ON is set where the image and CI build MORI. Without it
  cco has no SDMA queues and every put is a silent no-op;
  GemmAllReduceOp.self_test() then refuses to serve a collective that moves
  no bytes, so the flag would be accepted and immediately disable itself. The
  Dockerfile and amd_ci_install_dependency.sh must change together — the CI
  step rebuilds MORI over the image's copy and would otherwise undo it.
• The window is VMM memory outside torch's allocator, sized from
  get_schedule().chunked_prefill_size and unable to grow afterwards, so a run
  that enables this needs a correspondingly lower --mem-fraction-static. It
  is taken from the chunked-prefill limit rather than the first request seen,
  so a short prompt arriving first cannot fix a window too small for a later
  full chunk.
• The thresholds are on the SGLang side on purpose. _MIN_FUSED_M and
  _MIN_PAD_FILL answer "is fusing worth it against what the model runs
  today", which is not a property of the kernel; mori.ops.gemm_ar.supports()
  answers the separate question of whether the shape is expressible.

Not in this PR

A three-arm nightly perf suite (base / fused / fp8 + GSM8K) exists and runs,
but it is a 6-hour 8-GPU job for a feature that is off by default, so it is
held back rather than added to nightly-test-amd.yml here. Happy to bring it
in if maintainers would rather have the regression coverage up front.

Dependency

Needs mori with mori.ops.gemm_ar (ROCm/mori#662) built with
BUILD_CCO_SDMA=ON, and MORI_ENABLE_SDMA=1 at process start. SGLang already
depends on mori for moriep and the disaggregation transport; this is the
first use of it inside the model forward.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #37813744419
Latest PR Test (Extra): ❌ Run #37813743891
Latest PR Test (AMD ROCm 10): ❌ Run #37813744414