[yangyuhuiling]

What

wo_b in DeepSeek-V4.1-Flash is a RowParallelLinear: an mxfp8 GEMM followed
by an 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. The reduce and all-gather still run as their own kernels.

Opt-in and ROCm gfx950 only. fused_wo_b() returns None for everything it
does not serve — decode, small or ragged M, a layer neither mxfp8 route
prepared, no SDMA queues — and disables itself for the process on any failure,
so the ordinary path is the fallback rather than an error.

Supersedes #40757, which could not be reopened after its branch was
force-pushed. Same branch, rebased onto current main; the first commit is that
PR's content and the second takes it to V4.1-Flash.

Two switches, and they are not interchangeable

	TTFT, TP2	TTFT, TP4	numerics
SGLANG_OPT_FUSED_WO_B_AR=1	−2.4%	−2.0%	relL2 7.5e-3 – 8.2e-3
+ ..._FP8_GATHER=1	−5.8%	−7.0%	relL2 1.87e-2

The fp8 gather sends the all-gather leg as e4m3 with a per-row scale. It is
worth more than the fusing itself — and it is the half that costs accuracy, so
it stays off by default.

2x and 4x MI355X, DeepSeek-V4.1-Flash, 16384 in / 64 out, concurrency 1, the
MI355X cookbook's flags unchanged (including --fp8-gemm-backend aiter).
Medians over six fresh server instances per arm, arms interleaved. 48 instances,
24 pairs, every pair agreeing in sign.

Numerics, measured where the path engages

SGLANG_DEBUG_FUSED_WO_B_AR=1 reports two things per fused call: the op against
the unfused path, and the op against itself on the same inputs immediately
again. 42 samples per arm at M=16384:

Corrected after review (thanks @hubertlu-tw). This section claimed the bf16
wire was bit-identical on the cookbook flags. It is not, and the claim was the
wrong half of a real distinction: exactness needs wo_a to hand wo_b an
Fp8GridActivation, which wo_b_takes_fp8_grid() only returns for the gfx950
native mxfp8 route. The MI355X recipes pin --fp8-gemm-backend aiter, where
wo_b receives a plain bf16 tensor and the fused and unfused paths each
quantise it themselves.

Re-measured at M=16384, 4x MI355X, TP4, cookbook flags otherwise unchanged:

	vs unfused	vs itself
bf16 wire, --fp8-gemm-backend aiter	7.48e-3 .. 8.20e-3	0
bf16 wire, --fp8-gemm-backend auto	7.48e-3 .. 7.56e-3	0
fp8 wire	1.844e-2 .. 1.899e-2	0

The earlier 0 could not be reproduced on either backend value, and the
native route it would require is not one these flags select, so treat it as
withdrawn rather than conditional. The second column staying zero is the part
that holds: the collective is deterministic, without which "lossy but
deterministic" and "non-deterministic" would look the same.

The fp8 wire's spread is 3% of its own value, which is what a fixed mantissa
width looks like rather than a numerical hazard. What it does to downstream
quality is not measured here: GSM8K cannot answer it — its 5-shot prompts give
a padded M of ~512 against a floor of 8192, so every arm runs the same code, and
two passes of that same code scored 0.906 and 0.898, which puts its resolution
above 1% anyway.

Worth a reviewer's attention

• Both gfx950 mxfp8 routes are served, and the weight's shape cannot say which
  one ran. The native route rebinds the weight to a shuffled
  [N/16, K/128, 2048]; the aiter route — which --fp8-gemm-backend aiter
  takes, and which the MI355X cookbook pins — leaves the checkpoint's [N, K].
  mxfp8_aiter_ready records it, because reading one layout as the other
  returns an uncorrelated result rather than failing. mori's own preshuffle_b
  is defined on exactly what the aiter route leaves, so only the B conversion
  branches.

• The operand the gfx950 wo_a fork emits decides whether this can serve the
  call at all. An Fp8GridActivation unwraps to the bf16 underneath, but an
  Mxfp8Activation is already quantised with no bf16 left and the op does its
  own quantisation, so that form keeps the ordinary path. Only decode produces
  it, which the M floor declines anyway.

• The window is VMM memory outside torch's allocator, sized from
  get_schedule().chunked_prefill_size and unable to grow afterwards (560-640
  MiB at these shapes), so a run that enables this wants 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 leave a
  window too small for a later full chunk.

• The thresholds are on the SGLang side on purpose. _MIN_FUSED_M (8192 on
  the bf16 wire, 2048 on the fp8 one) answers "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.

• The MORI pin moves to 066fd10e so that export BUILD_CCO_SDMA=ON can
  come out of the Dockerfile and the CI install script: build(cco): compile the SDMA path by default ROCm/mori#710 made the
  cco SDMA path the default and normalised 1/TRUE/yes, which previously
  built a host library with SDMA and JITted a device wrapper without it. The
  guard in _FusedWoB stays — a build can still opt out, and the failure mode
  it catches is a collective that moves no bytes and therefore measures
  faster.

How these were measured, because a single pair does not resolve them

Instance-to-instance spread on these boxes is the same order as the effect,
while repeats inside one server instance land within ~1%. An earlier three-pair
run of the same comparison came out at −0.1% with two of its three pairs landing
on the baseline; six pairs of fresh instances put it at −5.8% with all six
agreeing. So: interleave the arms, take the median across instances, keep both
arms in one lease on one node, and record a signal that each arm is what it
claims to be (window present, declines zero) rather than trusting the switch.

Per-layer timings are deliberately not quoted. That harness synchronises
every call, so the transfer sits on the critical path with no neighbouring work
to hide behind — it cannot show an overlap benefit by construction, and its
baseline is the native route rather than the aiter one the server takes. On a
contended node it also inflates the NCCL baseline ~48% while leaving the SDMA
path untouched, which flips the sign.

Not in this PR

• The standalone mxfp8 GEMM. mori's GEMM without a collective attached is
  ~25-30% faster than aiter's at these shapes and serves layers this has nothing
  to fuse with, but it belongs in aiter rather than behind an SGLang switch, and
  is going there instead.
• Tests. The perf harnesses this was derived with are not included.
• V4-Pro. [AMD][gfx950] DSV4: fuse wo_b's GEMM with its all-reduce via mori cco #40757 measured V4-Pro, whose wo_b is a 128x128 block-scale
  linear preshuffled by the aiter block-fp8 route. V4.1-Flash quantises at
  32-wide ue8m0 and takes a gfx950 mxfp8 route, so the second commit replaces
  that path rather than adding to it — those numbers describe the first commit,
  not this PR's head.

---

CI States

Latest PR Test (Base): ❌ Run #41960774917
Latest PR Test (Extra): ❌ Run #41960774850
Latest PR Test (AMD ROCm 10): ❌ Run #41960774926

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[Duyi-Wang]

The main concerns are when the fused path may replace wo_b's own reduction, and how failures are handled across ranks:

Should be fixed before merge

1. DP attention: wrong reduction group, and a possible cross-rank hang (deepseek_v4.py)
2. A rank-local disable deadlocks the TP group on mori's device barriers (mori_gemm_ar.py)

Should fix

3. A construction-time ValueError is retried on every call
4. A lazily allocated second copy of each wo_b weight, outside the memory profiler
5. Window sizing when chunked prefill is off

Smaller

6. Raw os.environ access for registered env vars
7. Duplicated ue8m0 scale rule
8. Hot-path overhead of the gate

Details inline.

[yangyuhuiling]

Thanks — all eight were real, and two of them were wrong in ways I had documented in a comment and then not enforced. Fixed in 3dd6ca3 and 18dd7fd, verified on an 8×MI355X node.

1. DP attention: wrong reduction group, possible hang

You're right on both counts, and the second one follows from the first. The gate was on defer_all_reduce, which is a different question — everything that made that false reached the fused path, including attn_tp_size == 1 and 1 < attn_tp_size < tp_size. The env var's comment said "TP-only" and nothing enforced it.

Rather than gate at the call site I put it in fused_wo_b, as _reduces_over_tp_group(layer), using RowParallelLinear.forward's own conditions read off the layer — reduce_results and tp_size > 1, not use_decode_attn_tp, not use_dp_attention_reduce, and tuple(layer.tp_group.ranks) == tuple(get_tp_group().ranks). The call site is where this went wrong once; restating the rule there would leave two places to keep in agreement. should_skip_mlp_all_reduce() is in the per-forward half of it, which you'll also recognise as a case the old code let through: it made defer_all_reduce false, so the fused path performed a reduce the split path would have skipped.

Not gated: quantize_communications, which would make the split path's reduce numerically different — validation_hook.py:263 rejects it on anything but NPU.

Verified end to end: TP2 with --attn-dp-size 2 (so attn_tp_size == 1), switch on, an 8401-token prompt that is eligible on the TP-only run. Zero mori fused wo_b lines — it declines before constructing anything — and the model answers correctly. Plus a unit check of all six group configurations.

That also closes the hang, since eligibility can only diverge across ranks under DP attention, which no longer reaches _eligible.

2. A rank-local disable deadlocks the TP group

Agreed, and the architecture was wrong rather than the handler. Construction is now collective: _construct_collectively attempts it on every rank, agrees the outcome over tp.cpu_group, and the ranks that succeeded tear their communicator down when a peer did not — otherwise they hold the window's reservation for a path nothing can use.

After construction there is no fallback left. An OOM in pad_rows/quantize_packed is rank-local and run is collective, so returning None from either is exactly the silent hang you describe; dying is visible. Every remaining reason to decline is an explicit check above that point, and each reads the same on all ranks.

Verified with fault injection under torchrun, both --fail-rank 0 and --fail-rank 1:

rank 0: constructed=False  disabled=True  state_kept=False  torn_down=True
rank 1: constructed=False  disabled=True  state_kept=False  torn_down=False
all ranks agree: True

3. Construction ValueError retried every call

Same root cause — one try over both construction and the per-call work, with a handler written for the latter. Construction is attempted once and a failure is permanent (and collective).

4. Lazy second copy of each wo_b weight

prepare_wo_b_weight now runs at load time. Worth noting that my first attempt at this didn't work and the measurement caught it: I hooked post_load_weights, which loader.py:856 runs before process_weights_after_loading, so mxfp8_ready was false and it returned without allocating — KV cache and available memory came out byte-identical to the switch-off run. It is now in precompile_kernels_after_loading, which runs after the layout exists and before the pool is sized.

TP2, identical args, max_total_num_tokens:

	tokens
switch off	34,249,728
switch on	33,768,960

−480,768 tokens. With four MLA KV source layers that is roughly 1.1 GB, against 40 × 21 MiB ≈ 840 MiB of extra weight — the right order, and the point is that the profiler now sees it at all.

5. Window sizing without chunked prefill

Switched to max_prefill_buffer_tokens() with max_prefill_tokens as the fallback, which is what the other buffer-sizing callers here use and also covers PP dynamic chunking's 1.25× probe. Added the one-line warning when a call exceeds the window.

Verified with --chunked-prefill-size -1 --max-prefill-tokens 20480, first eligible request M=8401:

mori fused wo_b: window 700 MiB, tp=2, M<=20480, N=5120, K=4096, gather=bf16

Sized from the ceiling, not from that first request — which under the old code would have pinned it at ~8448 and declined everything larger for the life of the process.

6. Raw os.environ

Both go through the EnvField API now, and os is no longer imported. Your sub-point was the sharper half: SHAPE_LOG=true was silently ignored and override() did nothing, since it was compared against "1" at import time. Checked both in the unit test.

7. Duplicated ue8m0 scale rule

The quantiser calls fp8_grid_quant and keeps only the packed store. You're right that they are not the same function — fp8_grid_quant takes the exponent off the IEEE bits and is exact at powers of two, where ceil(tl.log2(.)) can land above the integer and pick one exponent higher.

Verified rather than reasoned: operands with whole groups placed on exact powers of two of 448, comparing against fp8_grid_quantize. Exponents and fp8 codes are now byte-identical; 0 differing groups.

8. Hot-path overhead

The switch is memoised (_disabled stays dynamic and is checked first, so the off path is one bool), and m_pad is computed once per call instead of three times. I left the call-site caching alone — it's a model-file change for what is now a memoised read.

---

Also re-ran the TP-only path end to end after all of this: fusion still engages at TP2 (window 560 MiB, M<=16384), both ranks agree, generation is correct.

[Duyi-Wang]

Thanks for the thorough follow-up. I re-reviewed 18dd7fd4 against mori's cco and ops/gemm_ar sources. Items 1, 3, 5, 6 and 7 from the earlier round look resolved, and the DP-attention gate is cleaner where it is now than at the call site.

A few things remain, details inline:

1. Collective construction: host-side failures look covered, but a rank failing between register_window and the end of self_test() still leaves its peers stuck on device barriers, before the CPU agreement. Also a question about where the fault was injected in the --fail-rank test.
2. Weight updates: the load-time _mori_b snapshot is never invalidated, so after update_weights_* the fused prefill serves the old wo_b.
3. Window allocation: still lazy, after the KV cache is sized. The reasoning that moved the weight copy to load time applies here too.
4. prepare_wo_b_weight scope: it also prepares DSPARK draft layers that never fuse, and copies that are never freed after a collective disable.
5. prepare_wo_b_weight docstring: it says "best effort", but there is no try.

[yangyuhuiling]

Thanks for the reproduction — that's a lot of careful work, and two of your points land. I've pushed four commits (c7b0d88..ea7d1c1) and corrected the PR description.

Point 4 first: you're right, and the claim is withdrawn

I pulled the checkpoint and reran this on 4x MI355X with the cookbook flags. wo_b receives a plain Tensor, not an Fp8GridActivation, with both --fp8-gemm-backend aiter and auto — exactly your reading. wo_b_takes_fp8_grid() only returns true for the gfx950 native mxfp8 route, and the recipes pin aiter, so the two paths each quantise the bf16 tensor themselves.

at M=16384, TP4	vs unfused	vs itself
bf16 wire, --fp8-gemm-backend aiter	7.48e-3 .. 8.20e-3	0
bf16 wire, --fp8-gemm-backend auto	7.48e-3 .. 7.56e-3	0

That matches your 7.7e-3 to 9.6e-3. The earlier 0 does not reproduce on either backend value, so I've withdrawn it in the description rather than narrowing it. What does hold is the second column: self_rel is 0 throughout, so the collective is deterministic — it was the "vs unfused" column that was wrong, not "vs itself".

ea7d1c1 writes this onto the env var, which previously said nothing about numerics while the fp8 gather's comment right below it called out its accuracy cost — which reads as though this switch has none.

Points 1 and 3: neither failure reproduces on another MI355X host

Your two images unmodified on an 8x MI355X box (gfx950, all XGMI), same mori f39fa765, FlyDSL 0.3.4.1, torch 2.11.0+rocm7.2 / HIP 7.2.26015 and 2.11.0+rocm10.0.0 / HIP 7.15.26333 — your table's stacks row for row. bench_gemm_ar.py, TP4, M=16384 N=5120 K=2048, mxfp8:

Mode	mori README	this host, rocm724	this host, rocm10	yours, rocm724
gemm-only	180.2 µs	179.9	182.0	3338.4
split-sdma bf16	1593.5	1588.3	1590.8	18364.0
fused-sdma bf16	1500.7	1500.9	1492.3	18409.2
fused-sdma fp8/LSA	1195.1	1193.3	1205.5	16606.2

All validated=True, relL2 identical across both stacks (bf16 2.37e-03, fp8 2.31e-02). On ROCm 10, mori's tests/python/cco/test_sdma_api.py -k alltoall gives 9 passed — the puts land.

On the CI prediction specifically: the mori-gemm-ar job has already run green on rocm10. Run 38566377447, job 115860094556 — container started with --rocm-version rocm10 (image v0.5.21-rocm10-mi35x-20261013), MORI_ENABLE_SDMA=1, and it ran test_mori_gemm_ar_collective.py, the file you expect to time out, to 29 passed, 3 subtests passed in 97.43s, job total 7m13s, zero skips. That file has no skip conditions: it spawns min(8, device_count()) ranks and asserts the fused result matches the split path, so inert puts would fail the assertion rather than pass it.

So I took your second option rather than the first. 26209c1 keeps the job on rocm10 and puts a bounded gate in front of it — mori's own put test, no collective behind it, ~50s. If a runner ever behaves like your host, that step goes red naming SDMA instead of the collective test timing out and pointing here.

The variable now looks like the host:

	yours	this host
kernel	Linux 5.15	6.8.0-65
amdgpu	6.16.13	6.14.14
runtime	docker	podman (rootless)

Could you post uname -r, cat /sys/module/amdgpu/version, and the full docker run flags? And one thing worth checking on the slow GEMM: your rocm724 gemm-only is ~3.3 ms at both K=2048 and K=4096, i.e. flat in the work. That is not the shape of a slow kernel, it is a fixed cost per call — FlyDSL is JIT-compiled, so watching host CPU during the run and confirming the JIT cache directory is writable would separate "compiling every call" from "codegen is bad on this stack". Our FlyDSL is the same 0.3.4.1.

Point 2: done

c7b0d88. self_test now runs under a deadline (SGLANG_OPT_FUSED_WO_B_AR_SELF_TEST_TIMEOUT_S, default 60s) on a side thread, so the rank still reaches the CPU-group agreement and every rank turns the path off together with a named error.

Being straight about the limit: a timeout does not make the process healthy. The stream stays wedged and the thread stays stuck. What it buys is a diagnosable failure instead of a server that logs nothing and never opens its port, which is the state you hit. _disable deliberately keeps the window in that case rather than freeing memory a live kernel still addresses.

Writing it turned up a bug worth naming: torch's current device and stream are thread-local, so the side thread defaulted to device 0 and ran every rank's phases against the wrong GPU. Caught by diffing against a baseline run rather than by the test passing. Both are carried across now.

Point 5: one correction, and you're right about the gap

The 2e-2 in test_mori_gemm_ar_collective.py:173 is asserted on the bf16 wire — SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER defaults off in that worker, and that leg measures 2.4e-3. So your 2.42e-2 was never compared against it.

But the gap underneath is real and worse than a wrong threshold: the fp8 wire had no coverage at all. aa08fc0 runs both wires with a bound each — 2e-2 for bf16, 3e-2 for fp8. I measure 2.31e-2 for fp8 on the TP4 microbench, next to your 2.42e-2; that is e4m3's mantissa width, a fixed cost rather than a margin, and the constant now says so.

Verification

8x MI355X, both of your images:

• rocm10, the exact command CI runs: 29 passed, 5 subtests, 103.69s (3 subtests before; the 2 new ones are the wire arms)
• rocm724, the collective file: 4 passed, 5 subtests, 128.19s
• the timeout path, deadline forced to 1 ms on 4 ranks: every rank gives up, agrees and disables, no hang

Your two side observations are both correct — the fused output is a view into the symmetric window and is only safe within the layer, and the window does take from the KV pool. Neither has a caller that breaks today, so I've left them; happy to assert the first if you'd rather not rely on that staying true.

[bmullick-amd]

Thanks @yangyuhuiling for reviewing Hubert's comment. I Additionally tested the latest PR on a mi355x machine. Similar to Hubert's observation, i too observed significant TTFT regression.

On TP2/EP1, with a 16,384-token uncached prefill, 64 output tokens, concurrency 1, and the bf16 gather wire, the fused path is about three times slower than the baseline, and the later PR update does not change that.

Initial PR (a8b0165): the baseline TTFT was 301.70 ms for one request and 299.90 ms across three. With SGLANG_OPT_FUSED_WO_B_AR=1, it was 944.19 ms and 945.34 ms, about +213% and +215%.

Updated PR (19a53ef): the same bf16-wire run measured 944.70 ms for one request and 944.38 ms across three, still about +213% and +215% versus that baseline.

I noted the difference in kernel and amdgpu version on my machine (kernel: Linux 5.15.0-70-generic , amdgpu: 6.16.13) v/s the one you noted. However, a 3x regression from kernel difference is unlikely.

[yangyuhuiling]

@hubertlu-tw @bmullick-amd Thanks for the detailed measurements. I think this is a known environment issue rather than something in the kernels or the integration. Both of your hosts run Linux 5.15, and the stock Ubuntu 22.04 5.15 kernel is built without CONFIG_DMABUF_MOVE_NOTIFY.

What happens. mori cco allocates its symmetric window through the HIP VMM API. hipMemSetAccess imports that memory into each peer as a dma-buf. On a kernel without move-notify support, amdgpu won't pin the buffer in VRAM: amdgpu_dma_buf_pin() drops the VRAM domain and pins it in host system memory instead. After that, every write into the window goes to host DRAM over the writing GPU's PCIe link, at roughly one link's bandwidth. Nothing errors, and results stay correct. It's documented in mori's known issues, Issue 1:
https://github.com/ROCm/mori/blob/main/.claude/skills/known-issues/SKILL.md (from ROCm/mori#595, where EPv2 dropped ~9x for the same reason).

Why it matches your numbers

• gemm-only writes C (16384×5120 bf16 = 168 MB) straight into the window. 168 MB / 3.3 ms ≈ 51 GB/s, which is about one link. That's also why the time stays flat from K=2048 to K=4096: the cost is the store, not the math.
• The reduce, gather and LSA phases all read and write the window too, so split-sdma and fused-sdma degrade far more than the GEMM alone.
• validated=True and normal relL2 are expected: the data is right, it just lives in the wrong place.
• torch GEMM, RCCL and peer copy don't use VMM-shared memory, so they're unaffected. Our hosts run Ubuntu's 6.8 HWE kernel, which has both options enabled, so the same images run normally there.

This also rules out the JIT-recompilation idea I raised earlier: graph replay doesn't recompile, and 3.3 ms is GPU time.

How to check (no GPU needed):

grep -E 'CONFIG_(PCI_P2PDMA|DMABUF_MOVE_NOTIFY)=' /boot/config-$(uname -r)

Both must print =y. On one of our boxes, the leftover 5.15 config shows the problem:

/boot/config-5.15.0-151-generic:  # CONFIG_DMABUF_MOVE_NOTIFY is not set
/boot/config-6.8.0-65-generic:    CONFIG_DMABUF_MOVE_NOTIFY=y

For a direct runtime check, vmm_peer_probe.cpp in the same skill directory reports VERDICT: AFFECTED if granting peer access moved the buffer out of VRAM.

Fix. Install the HWE kernel (sudo apt install linux-generic-hwe-22.04, which is 6.8), or rebuild 5.15 with DMABUF_MOVE_NOTIFY and PCI_P2PDMA enabled. The doc also has a one-line amdgpu DKMS patch. There's no application-level workaround: cco needs a flat symmetric VA across processes, and that only comes from VMM.

Could you run the check above, and if possible re-measure on a host with the options enabled? I'd also expect the ROCm 10 symptom (SDMA puts not landing) to be connected to the same window allocation path, but I haven't confirmed that, so it would be good to see whether it persists on a fixed kernel.

[hubertlu-tw]

@yangyuhuiling Thanks, that explains it. I checked the host I measured on:

It runs stock Ubuntu 5.15.0-70-generic, and its config has # CONFIG_PCI_P2PDMA is not set and # CONFIG_DMABUF_MOVE_NOTIFY is not set.
The installed amdgpu DKMS (6.16.13) has the if (!IS_ENABLED(CONFIG_DMABUF_MOVE_NOTIFY)) domains &= ~AMDGPU_GEM_DOMAIN_VRAM; branch in amdgpu_dma_buf_pin().
The linux-generic-hwe-22.04 package available to it (6.8.0-136) has both options set to =y.

We will re-measure after moving the host to the HWE kernel, and check whether the ROCm 10 SDMA symptom goes away too.

Because this fails silently (correct results, about 3x slower TTFT at 16K), I think the PR should detect it rather than leave it to a doc. Checking the kernel version isn't enough: the cause is a build option, and a rebuilt 5.15 is fine. Reading the config file also isn't enough on its own. Inside the rocm/sgl-dev containers, /boot is empty, and stock Ubuntu doesn't provide /proc/config.gz (CONFIG_IKCONFIG is not set). So for most users that check would return "unknown". The per-GPU sysfs counters are visible in the container, though, so I'd suggest two checks:

Host check before construction, put through _agreed() like the other stages. If the config is readable and the option is missing, every rank turns the fused path off with a message that names the fix.
Placement check in _FusedWoB.init. Compare this GPU's mem_info_vram_used before Communicator.init and after the op is built. On an affected kernel, granting peer access moves the window out of VRAM, so VRAM usage barely grows. This is the same signal mori's vmm_peer_probe.cpp uses, applied to this rank's own GPU.

Something like

_NO_MOVE_NOTIFY = (
    "the running kernel was built without CONFIG_DMABUF_MOVE_NOTIFY, so amdgpu "
    "pins mori's VMM window in host memory and every write to it crosses PCIe. "
    "Results stay correct but the fused path is several times slower than the "
    "split one. Use a kernel built with it, e.g. linux-generic-hwe-22.04 (6.8) "
    "on Ubuntu 22.04; see Issue 1 in ROCm/mori .claude/skills/known-issues."
)


def _kernel_config_error() -> RuntimeError | None:
    """The running kernel's config, if readable, rules out a VRAM window."""
    import gzip
    import platform

    sources = ((f"/boot/config-{platform.release()}", open), ("/proc/config.gz", gzip.open))
    for path, opener in sources:
        try:
            with opener(path, "rt") as f:
                if any(line.strip() == "CONFIG_DMABUF_MOVE_NOTIFY=y" for line in f):
                    return None
                return RuntimeError(_NO_MOVE_NOTIFY)
        except OSError:
            continue
    # Unreadable, as inside most containers; the placement check decides.
    return None


def _vram_used_bytes() -> int | None:
    """This rank's GPU VRAM in use, from amdgpu's sysfs counter."""
    props = torch.cuda.get_device_properties(torch.cuda.current_device())
    bdf = f"{props.pci_domain_id:04x}:{props.pci_bus_id:02x}:{props.pci_device_id:02x}.0"
    try:
        with open(f"/sys/bus/pci/devices/{bdf}/mem_info_vram_used") as f:
            return int(f.read())
    except (OSError, ValueError):
        return None

In _FusedWoB.init, read the counter before Communicator.init and check it once the op exists. A failure here is caught by the existing "constructed" agreement, so every rank still falls back together:

vram_before = _vram_used_bytes()
self._comm_ctx = Communicator.init(...)
# ... existing GemmAllReduceOp construction ...

torch.cuda.synchronize()
vram_after = _vram_used_bytes()
if (
    vram_before is not None
    and vram_after is not None
    and vram_after - vram_before < self.op.window_bytes // 2
):
    grew = (vram_after - vram_before) / 2**20
    self.close()
    raise RuntimeError(
        f"{_NO_MOVE_NOTIFY} (VRAM grew {grew:.0f} MiB for a "
        f"{self.op.window_bytes / 2**20:.0f} MiB window)"
    )

• When the window's memory is mapped. I'm assuming it's backed and peer-mapped by the time GemmAllReduceOp returns. If mori maps it lazily, the second read should move after warmup().
• Noise from other processes on the same GPU. The half-window threshold tolerates a little, but another process allocating on that GPU during init could in principle hide a real eviction.

This affects every cco user, EPv2 included, so the same check would also make sense in mori's Communicator.init.

cc: @bmullick-amd @HaiShaw

[yangyuhuiling]

@hubertlu-tw Added the two checks in bfc9acf238f, following your suggestion. They run during initialization, with _agreed() coordinating fallback across TP ranks before the fused device phases.

The kernel check reads /boot/config-<running release> and then /proc/config.gz. A readable config without CONFIG_DMABUF_MOVE_NOTIFY=y disables fusion through _agreed() before the constructor's UID broadcast. An unreadable config leaves the runtime placement check to decide. This checks the build option, so a rebuilt 5.15 kernel is not rejected by its version number.

For placement, the baseline is taken before Communicator.init(), using the current logical device's PCI BDF to read its own mem_info_vram_used. The final read is after the existing construction agreement, followed by a separate placement agreement, before warmup() and self_test(). The pinned MORI already allocates and peer-maps the window in its constructor; waiting for every rank to finish avoids depending on an implicit synchronization inside MORI. The threshold uses the actual op.window_bytes, not the VA reservation. Cleanup stays in _disable(), so it closes each successful state once and releases the prepared weight copies and their update restriction.

Real-model startup exposed another source of counter noise: deferred driver accounting after empty_cache(). The first implementation incorrectly disabled fusion on a healthy host: a 520 MiB window was built while the counter dropped by about 9.6–17.9 GiB. A short TP4 reproduction allocated and released two 8 GiB torch buffers per GPU. Torch's reserved memory became zero immediately, but sysfs retained those 16 GiB and dropped them in two batches, around 250 ms and 500 ms later.

The preflight now synchronizes and empties the torch cache after weight preparation, with failures covered by the same CPU agreement. Both VRAM readings require 0.75 s of stability, with a 3 s limit. A counter that is unreadable or keeps changing produces an explicit placement is unverified warning and preserves the opt-in behavior. A stable but insufficient increase disables fusion collectively. The check remains a device-wide heuristic: unrelated allocations or frees can still affect it. Successful checks log the before/after/window byte counts; enabled on all ... TP ranks is logged only after all initialization stages pass.

Validation on MI355X:

• 47 tests, 19 subtests passed across the gating, contract and collective files. Coverage includes config/gzip fallback, logical-device-to-BDF mapping, delayed releases, unknown counters, one-rank config/cache-flush failures, and a real TP2 communicator with one rank's placement readings injected. That last case verifies resource release and a subsequent ordinary RCCL all-reduce. Both wire numerics now go through _construct_collectively() and the real checks.
• DeepSeek-V4.1-Flash TP4/EP1, 16384 input / 64 output, concurrency 1 completed in all three arms. Both fused arms logged all four ranks passing placement and initialization. All nine measured requests had the exact lengths and zero cached tokens.

Arm	Mean TTFT	Median TTFT	Mean change vs off
Off	230.09 ms	228.94 ms	—
Fused, bf16 gather	222.10 ms	222.97 ms	-3.47%
Fused, fp8 gather	218.37 ms	217.83 ms	-5.10%

This is one fresh server per arm, three measured requests each, a quick regression check rather than the original six-pair measurement. Each instance had one 16K/32 warmup followed by an explicit cache flush. Debug and shape logging were off during timing. These numbers do not establish a repeatable percent-level gain or downstream quality. The host has a working kernel configuration; bad-kernel and bad-placement paths were fault-injected, not reproduced on an actual affected 5.15 host.

Stack: rootless Podman, rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014, torch 2.11.0+rocm10.0.0 / HIP 7.15.26333, FlyDSL 0.3.4.1, MORI f39fa765, Linux 6.8.0-65-generic, amdgpu 6.14.14.

TP4 launch and benchmark parameters

Run from the PR checkout in the image above, with the real checkpoint mounted at the model path. All three arms use the same settings below. The two fusion switches are (0, 0) for off, (1, 0) for bf16 gather and (1, 1) for fp8 gather.

One loading-only adjustment is explicit: all arms set SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1. This existing option copies CPU weight views into independent contiguous storage before H2D. It was absent from the retained PR launch script; 21 CPU-only cases, including real checkpoint TP4 weight/scale slices, verified that it preserves every byte. The shared memory fraction is 0.78.

export PYTHONPATH="$PWD/python"
export MODEL=/models/DeepSeek-V4.1-Flash
export HIP_VISIBLE_DEVICES=0,1,2,3
export MORI_ENABLE_SDMA=1 MORI_SOCKET_IFNAME=lo
export SGLANG_USE_AITER=1 SGLANG_USE_ROCM700A=0
export AITER_BF16_FP8_MOE_BOUND=0
export SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1
export SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE=1
export SGLANG_DSV41_ENGRAM_HOST_TABLE_LAYOUT=per_rank
export TORCH_BLAS_PREFER_HIPBLASLT=1 TRITON_HIP_USE_ASYNC_COPY=0
export ROCM_QUICK_REDUCE_QUANTIZATION=NONE
export HSA_NO_SCRATCH_RECLAIM=0 GPU_MAX_HW_QUEUES=2
export SGLANG_DEBUG_FUSED_WO_B_AR=0 SGLANG_OPT_FUSED_WO_B_AR_SHAPE_LOG=0
export SGLANG_OPT_FUSED_WO_B_AR=1 SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER=0

python3 -m sglang.launch_server \
  --model-path "$MODEL" --trust-remote-code \
  --tp 4 --ep-size 1 --mem-fraction-static 0.78 \
  --speculative-algorithm DSPARK --speculative-dspark-block-size 5 \
  --enforce-shared-experts-fusion --fp8-gemm-backend aiter \
  --cuda-graph-max-bs-decode 64 --enable-decoder-swa-bounded-replay \
  --cuda-graph-backend-prefill disabled --chunked-prefill-size 16384 \
  --reasoning-parser auto --tool-call-parser auto \
  --host 127.0.0.1 --port 31300

Client command, after readiness, from another shell with the same checkout and model path:

PYTHONPATH="$PWD/python" python3 -m sglang.benchmark.serving \
  --backend sglang --base-url http://127.0.0.1:31300 \
  --model /models/DeepSeek-V4.1-Flash \
  --dataset-name random-ids --tokenize-prompt \
  --random-input-len 16384 --random-output-len 64 --random-range-ratio 1 \
  --num-prompts 3 --max-concurrency 1 --request-rate inf \
  --warmup-requests 1 --flush-cache --cache-report --output-details \
  --seed 42 --disable-tqdm --output-file result.jsonl

The implementation is in the SGLang adapter for this PR; a common MORI check can follow separately.

[bmullick-amd]

@yangyuhuiling Thank you for the detailed explanation above. Following your comment I upgraded my machine from 5.15.0-70-generic to 6.8.0-138-generic.

Every run used PR head 8c88d3b033, DeepSeek-V4.1-Flash, 16,384 uncached input tokens, 64 output tokens, and concurrency 1.

Case ROCm 7.2.4, TP2 EP1:
The 3× slowdown from kernel 5.15 is gone, and a smaller regression remains. The BF16 gather rerun was +6.4% for one request and +5.1% for three (289.99 ms and 288.14 ms versus 272.59 ms and 274.13 ms).

The FP8 gather run, against a unfused baseline, was +1.7% and +2.0% (275.97 ms and 277.45 ms versus 271.43 ms and 272.09 ms). You claim speedups here: −2.4% for BF16 and −5.8% for FP8 at TP2.

Case ROCm 10, TP4 EP1:
Both wires are faster. BF16 gather was −2.6% (214.46 ms versus 220.19 ms). FP8 gather was −5.3% (208.47 ms versus 220.08 ms). Your TP4 numbers on the same image was −3.47% for BF16 (222.10 ms versus 230.09 ms) and −5.10% for FP8 (218.37 ms versus 230.09 ms).

So I was able to get performance improvement on rocm10 image but still saw regression on rocm 7.2.4

[yangyuhuiling]

Thanks for re-running after the kernel upgrade — and for posting both stacks side by side. The 3× slowdown disappearing on 6.8 confirms the 5.15 DMABUF_MOVE_NOTIFY diagnosis, and your ROCm 10 TP4 numbers (−2.6% bf16, −5.3% fp8) line up with ours (−3.47%, −5.10%).

I tried to reproduce the ROCm 7.2.4 TP2 regression and could not. PR head 8c88d3b033, rocm/sgl-dev:v0.5.21-rocm724-mi35x-20261014 (torch 2.11.0+rocm7.2, HIP 7.2.26015, the image's own MORI f39fa765, FlyDSL 0.3.4.1), Linux 6.8.12, amdgpu 6.16.13, DeepSeek-V4.1-Flash TP2/EP1, 16384 in / 64 out, concurrency 1. Six rounds of three arms, every arm a fresh server, order alternated per round; each instance one 16K warmup, a cache flush, then three measured requests:

Arm	Mean TTFT per instance (ms)	Median	vs off
Off	351.7, 350.3, 350.4, 350.7, 349.1, 350.0	350.4	—
Fused, bf16 gather	339.1, 338.6, 338.0, 340.0, 338.3, 341.0	338.9	−3.3%
Fused, fp8 gather	327.0, 325.4, 328.1, 324.4, 324.8, 324.9	325.2	−7.2%

Every round agrees in sign (bf16 −2.6% to −3.6%, fp8 −6.4% to −7.5%), and the instance-to-instance spread on this box was about 1%, so your +5–6% on bf16 is well outside noise and I'm treating it as real. Both fused arms logged enabled on all 2 TP ranks and the window allocation on every instance; the off arm logged neither.

What stands out is the TP2 baseline. On ROCm 10 TP4 our unfused TTFTs are within 4% of each other (your 220.1 ms, my 230.1 ms), but at TP2 yours is 29% faster than mine (272.6 vs 350.4 ms). Going from TP4 to TP2 doubles the per-GPU prefill work: my TTFT rises 52%, yours only 24%. That suggests your TP2 run differs from your TP4 run in configuration, not just in parallel size, and the regression may come from whatever that difference is. My TP2 launch is the TP4 command from my previous comment with --tp 2 and HIP_VISIBLE_DEVICES=0,1 (DSPARK speculative, --fp8-gemm-backend aiter, --chunked-prefill-size 16384, --mem-fraction-static 0.78, the env block as posted). Could you share:

1. your full TP2 launch command and environment variables, and whether they differ from your TP4 run;
2. the exact image tag, and how deepseek-recipe was installed in it (the 7.2.4 image doesn't ship it; I copied 0.2.0 from the 0.5.22 rocm10 image);
3. from a fused TP2 run's startup log, the mori fused wo_b: allocated window … and enabled on all … TP ranks lines.

With your configuration I can rerun the same six-round comparison and narrow down which setting turns the gain into a loss.

Issue has been resolved

[hubertlu-tw]

LGTM

[hubertlu-tw]

@HaiShaw could you please help review the PR?

[kkHuang-amd]

Reviewed the whole diff, with a focus on correctness of the gating and on the question of whether any non-gfx950 platform can reach this path.

Overall this is careful work — the collective-agreement staging in _construct_collectively, the bounded self_test, and the "disable on every rank or none" discipline all address failure modes that are easy to get wrong and hard to debug. The three-layer test split (host gating / mori contract / cross-rank) is the right shape. Below are the things I think are worth acting on.

Correctness

1. sp_active is not excluded, and under SP wo_b does not all-reduce

RowParallelLinear.forward has a branch before the all-reduce:

if (
    get_forward().sp_active
    and tp_size > 1
    and not skip_all_reduce
    and output_tensor is None
):
    output = layernorm_sp.row_parallel_gbar_matmul(self, input_parallel, bias_)
    output_bias = self.bias if self.skip_bias_add else None
    return output, output_bias

With Megatron SP active, this layer performs a reduce-scatter and returns an output that is sharded along the token dim for the next SP LayerNorm region. The fused path performs a full all-reduce and returns an unsharded [M, N], so both the semantics and the shape are wrong.

_reduces_over_tp_group() reads only static layer attributes, and _substitutes_layer_all_reduce() adds only should_skip_mlp_all_reduce(); neither can see this per-forward flag. DSV4 does reach SP — the sibling wo_b optimisation in the same method already guards on it:

fuse_mxfp8_quant = self.use_flashinfer_mxfp8_wo_b and not get_forward().sp_active

Suggest adding if get_forward().sp_active: return False to _substitutes_layer_all_reduce(), which is already the home for per-forward flags. sp_active is rank-uniform, so this introduces no rank divergence.

Everything else in _reduces_over_tp_group checks out against RowParallelLinear.forward: once use_dp_attention_reduce and use_decode_attn_tp are excluded and the rank tuples match, the layer reduces over the global TP group whether it takes the parallel_group == "tp" branch or tensor_model_parallel_all_reduce. The attn_tp subgroup case is declined conservatively, which is correct.

2. mxfp8_shape() overstates N when N is not a multiple of 32

mxfp8_shape() derives (sn * 32, sk * 32) from the scale, but weight_scale_mx_e8m0 has ceil(N / 32) rows. That non-divisible case is expected elsewhere in this file — _prepare_flydsl_mxfp8 does layer.weight_scale_mx_e8m0.repeat_interleave(32, dim=0)[:n] precisely to handle it.

The consequence is that mori_weight()'s w.reshape(n, k) raises, and since prepare_fused_wo_b is deliberately not best-effort, that is a startup failure rather than a fallback. K is protected by the w_k != k comparison against the activation; N has no equivalent.

DSV4.1-Flash's wo_b emits hidden_size, which divides, so this is not reachable today. But the aiter route leaves the weight as [N, K], so layer.weight.shape[0] is available directly; alternatively mxfp8_route() could require n % 32 == 0 and k % 32 == 0 and return None otherwise. Given how defensive the rest of the module is, this gap stands out.

3. prepare_fused_wo_b allocates copies for layers it will then decline

The window takes (n, k) from candidates[0], but mori_weight(layer) is then called for every candidate. fused_wo_b later declines any layer where n != _state.n or k != _state.k, so a mismatched layer's full weight copy is never read. Filtering candidates by (n, k) before the mori_weight loop would close it.

Worth discussing

The returned view lives longer than the docstring suggests. The comment says the caller consumes it in the residual add of the same layer. In practice attn_boundary.finish() passes it to ResidualStream.record(), which explicitly records the contribution without applying the residual update; it is applied at ffn_boundary.prepare() inside op_mhc_post_attn_pre_mlp. That is still within the same layer and before the next wo_b, so it is safe today — and TBO is safe too, because neither DSV4 prefill op list has a YieldOperation between op_attn and op_mhc_post_attn_pre_mlp. But the safety rests on an op schedule this module does not own, and prefill TBO is exactly the large-M regime where fusing engages: a yield inserted between those two ops would make two microbatches overwrite each other's window silently. Worth stating as an invariant in the module docstring, or checking under the debug flag.

Enabling this flag disables online weight updates. _UPDATE_UNSUPPORTED lands on every wo_b. That is a real trade-off for RL / online-update users and probably deserves a mention beyond the env comment. Minor related nit: release_mori_weight sets _derived_weight_cache_error = None rather than restoring a prior value. The only other writer is the Hopper-only path in fp8.py, so there is no collision on ROCm today.

The --mem-fraction-static requirement is only in the env comment and the PR body. A 560–700 MiB/rank window plus ~840 MiB of weight copies is enough that the first attempt to enable this will OOM for most users; it would help to have it in the docs.

Bias is silently dropped. The fused path ignores layer.bias. wo_b is constructed with bias=False, so this is not live, but an explicit layer.bias is None check would match the defensive style of the rest of the module.

LoRA would fail at startup. With a LoRA adapter, wo_b becomes RowParallelLinearWithLoRA and the attributes move under base_layer. _reduces_over_tp_group reads layer.tp_group / layer.reduce_results directly and _can_ever_fuse does not wrap it, so this raises AttributeError. Only affects runs that opt in, and one hasattr guard covers it.

Platform reachability

I checked whether anything other than ROCm gfx950 can reach the fused path. The answer is no, but the guarantee is implicit rather than stated in code.

fused_wo_b_available() checks only _disabled, the env switch, and that mori.ops.gemm_ar imports — there is no is_hip() or arch check anywhere in the module. What actually excludes other platforms is the mxfp8_route(layer) / weight_scale_mx_e8m0 test, via this chain: those markers are set only in fp8_hip.process_dense_weights, which is called only from the backend.is_gfx95() branch of _process_mxfp8_linear_weight_scale, and the GFX95_* backends are selected only under _is_hip and _is_gfx95_supported in resolve_block_fp8_mxfp8_backend().

Per platform:

• NVIDIA / CPU / XPU: mori is a ROCm-only package, so the import fails, _mori_has_gemm_ar() sets _disabled and warns once. Even with mori somehow present, no layer carries weight_scale_mx_e8m0 and every call stops at the hasattr.
• NPU: the markers are never set, and NPU has its own use_npu_arch35_mxfp8_wo_a route.
• gfx942 / MI300X: the closest case, because the ROCm image does ship mori, so fused_wo_b_available() returns True. _is_gfx95_supported is False, so no mxfp8 marker exists, _can_ever_fuse returns False and no window is allocated. Behaviour is correct; the cost is one attribute lookup per call.

So nothing is broken, but I would still add is_gfx95_supported() (already in sglang/srt/utils/common.py) to fused_wo_b_available(). The concern is not today's behaviour: the gate is currently "does this attribute exist" rather than "what hardware is this", so if any future backend reuses the weight_scale_mx_e8m0 / mxfp8_aiter_ready names — a later gfx revision, or a CUDA mxfp8 route — this path becomes eligible silently and fails inside mori's device-side barriers.

Imports are fine: deepseek_v4.py imports mori_gemm_ar unconditionally at module scope, which pulls in triton and mxfp8_amd_gfx95, but that module depends only on torch/triton/msgspec and deepseek_v4.py already imports several triton kernel modules at top level.

CI

The paths-filter covers docker/rocm.Dockerfile, and amd_ci_install_dependency.sh extracts the pin from that file's ARG MORI_COMMIT, so the claim that a commit bump triggers this job holds. The install script itself is not in the filter, though — a change to the install logic would not run the job. Worth adding.

The if: github.event_name == 'pull_request' on the filter step combined with the outer always() && (workflow_dispatch || ...) is right: on dispatch the output is empty and the first condition catches it.

[mintlify]

Preview approval required

This pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit f02b7e6.

Approve preview

[yangyuhuiling]

@kkHuang-amd thanks for the careful pass. It is addressed in ba5efe3, as one commit on top.

Correctness

1. SP. _substitutes_layer_all_reduce() now returns False when get_forward().sp_active is set, next to should_skip_mlp_all_reduce().
2. N off the 32 grid. mxfp8_shape() now reads N and K off the weight in its route's layout: [N/16, K/128, 2048] on the native route, [N, K] on aiter. It no longer derives them from the ue8m0 scale, so an N that is not a multiple of 32 is no longer overstated. For any N that is not a multiple of 256, mori's supports() already returns False, so the layer is declined rather than failing at startup.
3. Copies for declined layers. prepare_fused_wo_b filters candidates to the first candidate's (N, K) before mori_weight, so the window and the copies cover the same layers.

The other points

• Lifetime of the returned view. fused_wo_b's docstring now states the actual path: attn_boundary.finish() -> ResidualStream.record() -> applied at ffn_boundary.prepare() in op_mhc_post_attn_pre_mlp. It also states the invariant the two prefill TBO microbatches rely on: no YieldOperation between op_attn and that op.
• Weight updates and --mem-fraction-static. Both are now in environment_variables.mdx under MORI Configuration, along with where the flag stays off.
• Restoring a prior value. release_mori_weight now restores whatever _derived_weight_cache_error the layer held before, instead of setting None.
• Bias and LoRA. Only a plain RowParallelLinear with bias is None is fused. A LoRA wrapper is not a RowParallelLinear, so it is declined at startup instead of raising AttributeError. This is an isinstance check rather than hasattr, per the repo's rule on defensive attribute access.
• Platform. fused_wo_b_available() now requires is_gfx95_supported(), memoised with the switch.
• CI. scripts/ci/amd/amd_ci_install_dependency.sh is in the paths filter.

Tests. Each of SP, a wrapped layer, N off the 32 grid, no copies for unserved shapes, restoring the prior constraint, and the gfx95 gate has a case in test_mori_gemm_ar.py. I checked that all six fail against the previous commit and pass on this one. The native-route test layer now has its real [N/16, K/128, 2048] shape.

All three files (test_mori_gemm_ar.py, _contract.py, _collective.py) pass on 8x MI355X with the ROCm 10 image and the pinned mori f39fa765: 54 passed, plus 24 subtests. The SDMA gate also passes.