From bd6555dcfd1b1f3cc392285b03a2a8c7f730faf4 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 22 Sep 2026 15:08:27 +0000 Subject: [PATCH 01/20] [AMD][gfx950] DSV4: fuse wo_b's GEMM with its all-reduce via mori cco `wo_b` is a RowParallelLinear: a block-scale fp8 GEMM followed by an NCCL all-reduce of the [M, hidden] result. mori'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-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 always the fallback rather than an error. Measured on 8x MI355X (gfx950), DeepSeek-V4-Pro, TP8, in 4096 / out 512: bs base fused fp8 gather 8 1136.1ms -3.36% -6.04% 16 2232.0ms -3.03% -6.29% 64 8890.9ms -3.26% -6.44% (TTFT; bs=1 is at the fusing threshold and is noise either way.) The layer itself goes 1460.3us -> 1145.1 at M=16384 N=7168 K=2048, i.e. -21.6%; the end-to-end share is smaller because wo_b is about 12% of the forward. The fused and unfused paths agree to relL2 2.72e-05 -- both quantise identically, so what is left is bf16 accumulation order. The optional fp8 all-gather leg (`..._FP8_GATHER=1`, also off by default) halves that leg's bytes for a further ~3%, consistently across batch sizes because the leg grows linearly with M. It costs accuracy: relL2 on the gather goes 2.35e-3 -> 2.49e-2, which is e4m3's 3-bit mantissa rather than a granularity that can be tuned, and the window grows 700 -> 812 MiB. Two call sites of this layer 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, so this checks rather than relies on the platform. `BUILD_CCO_SDMA=ON` is set where the image and CI build MORI. Without it cco has no SDMA queues, every put is a silent no-op, and `GemmAllReduceOp.self_test` refuses to serve a collective that moves no bytes -- so the flag would be accepted and then disable itself. 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. Co-Authored-By: Claude Opus 5 (1M context) --- docker/rocm.Dockerfile | 2 +- python/sglang/srt/environ.py | 26 ++ python/sglang/srt/layers/mori_gemm_ar.py | 392 +++++++++++++++++++++++ python/sglang/srt/models/deepseek_v4.py | 19 +- 4 files changed, 435 insertions(+), 4 deletions(-) create mode 100644 python/sglang/srt/layers/mori_gemm_ar.py diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index a48669104012..21b63e532fc8 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -398,7 +398,7 @@ ARG ENABLE_MORI=0 ARG NIC_BACKEND=none ARG MORI_REPO="https://github.com/ROCm/mori.git" -ARG MORI_COMMIT="879983bdbd8c65c52e9f79ad836a61cbffef98b6" +ARG MORI_COMMIT="066fd10e549e6e3bbd0a92b98963a9b3ec245414" # NIXL (upstream ai-dynamo/nixl) — KV transfer backend for prefill/decode disaggregation. # Built from source for ROCm; needs UCX built --with-rocm (built here from openucx). diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 62ca05fb7d86..32ed92aa7c71 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1630,6 +1630,32 @@ class Envs: # tuned batched_gemm_bf16 (gfx95). Off by default; see deepseek_v4.py # _apply_wo_a_bf16_matmul. SGLANG_OPT_USE_AITER_BATCHED_GEMM = EnvBool(False) + # ROCm gfx950 TP-only: replace wo_b's `GEMM then all-reduce` pair with mori + # cco's fused GEMM+reduce-scatter, which overlaps the scatter with the GEMM. + # Measured 1297us against 1617 per layer at [16384, 7168] K=2048 on 8 ranks. + # Off by default: it needs a ~700 MiB symmetric window outside torch's + # allocator, and it only engages for large M (ragged M is zero-padded up to + # a multiple of tp_size*128). + SGLANG_OPT_FUSED_WO_B_AR = EnvBool(False) + # Send the fused wo_b's all-gather leg as fp8 e4m3 with a per-row scale, + # halving its bytes. That leg is ~40% of a fused layer and already runs at + # the xGMI ceiling, so it is worth about -10% on the layer -- but it costs + # accuracy: relL2 goes from 2.35e-3 to 2.49e-2 at the model's shape, which + # is e4m3's 3-bit mantissa and not a granularity that can be tuned away. It + # also needs a larger window (812 MiB against 700) for the staging region. + SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER = EnvBool(False) + # Which transport carries the fused wo_b's all-gather leg, "lsa" or "sdma". + # LSA pulls, so with an fp8 gather it widens on the way in rather than in a + # second kernel, and it wins in both places (951us against SDMA's 1011 on + # the layer). Only read when the fused path is on. + SGLANG_OPT_FUSED_WO_B_AR_GATHER_TRANSPORT = EnvStr("lsa") + # Debug: log the distribution of M the fused wo_b is offered, and how much + # of it clears the fusing thresholds. M is the token count of one forward, + # not of one request, so this is the only way to see what the scheduler + # actually batches. + SGLANG_OPT_FUSED_WO_B_AR_SHAPE_LOG = EnvBool(False) + # Debug: run both the fused and the unfused wo_b and log their relative L2. + SGLANG_DEBUG_FUSED_WO_B_AR = EnvBool(False) SGLANG_OPT_BF16_FP32_GEMM_ALGO = EnvStr("cublas") SGLANG_OPT_FUSE_WQA_WKV = EnvBool(True) SGLANG_OPT_USE_MULTI_STREAM_OVERLAP = EnvBool(True) diff --git a/python/sglang/srt/layers/mori_gemm_ar.py b/python/sglang/srt/layers/mori_gemm_ar.py new file mode 100644 index 000000000000..f2c2c7e3a755 --- /dev/null +++ b/python/sglang/srt/layers/mori_gemm_ar.py @@ -0,0 +1,392 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused GEMM + all-reduce for DSV4 ``wo_b`` on ROCm gfx950, via mori cco. + +``wo_b`` is a ``RowParallelLinear``, so today it runs a block-scale fp8 GEMM and +then an NCCL all-reduce of the ``[M, hidden]`` result. The two can overlap: +mori's fused kernel writes each destination's row band straight into a symmetric +window and, as soon as a band's tiles are done, hands it to the SDMA copy +engines; the reduce and all-gather then run as their own kernels. + +At ``[16384, 7168]`` K=2048 on 8x MI355X the pair costs 1419.5us and the fused +kernel 1146.2, i.e. **-19.4%**. Essentially all of it is the overlap: the same +pipeline unfused is 1465.9us, and the GEMM alone is 369.4. In-server, over one +20000-token prefill captured with the same warm-up and flush on both sides, GPU +busy goes **1096.0ms -> 1070.3, -25.7ms / -2.3%**, and the fp8 gather takes it +to **1041.2, -5.0%**. The request's own wall time agrees: 1.1969s -> 1.1728 -> +1.1385. The layer-level win is much larger than the end-to-end one because +`wo_b` is about 12% of the profile. + +**Check that mori's cco SDMA path was compiled in before believing any +measurement of this path.** With it off every put silently does nothing: the +all-reduce returns mostly the local slice, the model still answers, and the +fused path looks *faster* than it is because it is not moving any data. Measured +that way it reads -6.8% instead of -2.3%, and the fp8 pull -- the only leg that +does not go through SDMA -- looks like the slowest of the three instead of the +fastest. Perplexity is what catches it: 862511 against 3.26 on the same text. +It is ON by default since ROCm/mori#710, so this is now a question only for a +build that passed `BUILD_CCO_SDMA=OFF` or predates that commit. + +Prerequisites, all of which this module checks rather than assumes: + +* a mori whose cco SDMA path was compiled in -- the default since + ROCm/mori#710. A build that opted out with ``BUILD_CCO_SDMA=OFF`` has no SDMA + queues, every put silently does nothing, and the all-reduce quietly produces + zeros. +* ``MORI_ENABLE_SDMA=1`` at process start. + +The kernel itself is ``mori.ops.gemm_ar.GemmAllReduceOp``, which owns the +symmetric window, the per-M compile cache and the row padding. What stays here +is the part that is sglang's: the TP communicator, how the window is sized, and +whether fusing is worth it at this shape. + +The fused epilogue pushes a whole ``BLOCK_M`` row band to one destination, so a +destination's row slice has to be a whole number of bands: M must be a multiple +of ``tp_size * 128``. Ragged M is zero-padded up to that, which costs GEMM rows +and buys the overlap -- see ``_MIN_PAD_FILL``. Decode and small batches still +fall back to the ordinary path, and that is expected rather than a failure. +""" + +from __future__ import annotations + +import logging +from typing import Optional + +import torch + +from sglang.srt.environ import envs + +logger = logging.getLogger(__name__) + +#: Smallest padded M worth fusing. Measured at ``[*, 7168] K=2048`` on 8 ranks +#: against the aiter GEMM + NCCL pair the model runs today: 1024 -2.6%, 2048 +#: -7.8%, 4096 -12.3%, 8192 -15.7%, 16384 -19.4%. At 1024 a destination gets a +#: single row band, so there is nothing to overlap and the fused kernel is 1.2% +#: *slower* than the same pipeline unfused. +_MIN_FUSED_M = 4096 +#: The fp8 gather used to need a higher floor: with the SDMA transport its two +#: conversion kernels were a fixed cost against a transfer that shrinks with M, +#: so it was 2.3% *slower* at M=4096 and only paid from 8192 up. The LSA pull +#: removed that -- it widens on the way in, so there is no second kernel and no +#: re-read of the landed fp8 -- and fp8 now wins wherever fusing does at all. +#: Measured on the fused path at ``[*, 7168] K=2048``, fp8/lsa against bf16: +#: M=4096 -7.3% (335.8 -> 311.3us), M=8192 -13.8% (593.7 -> 512.0). +#: +#: Keeping 8192 is not merely conservative, it is a loss: a chunk between the +#: two floors takes neither wire and falls back to NCCL entirely. +_MIN_FUSED_M_FP8_GATHER = _MIN_FUSED_M +#: Padding buys the overlap and costs GEMM rows, so the fill ratio has to clear +#: the gain at the padded size. 0.88 is break-even at ``_MIN_FUSED_M`` (12% +#: more rows against a 12.3% gain) and increasingly safe above it. Measured: +#: M=3616 padded to 4096 is 342.1us against 371.4 for the unfused pair. +_MIN_PAD_FILL = 0.88 +_state: Optional[_FusedWoB] = None +_disabled = False +_warned_layout = False +_warned_reject = False + +#: Log what M this layer is actually handed, as a periodic histogram. +#: +#: M is the token count of one forward, not of one request, so whether the fused +#: path engages depends on how the scheduler batches -- which is not something to +#: infer from the client's concurrency. Set +#: SGLANG_OPT_FUSED_WO_B_AR_SHAPE_LOG=1 to find out. +_shape_hist: dict = {} +_shape_calls = 0 +#: 61 wo_b calls make one forward, so this logs roughly every ten of them. +_SHAPE_EVERY = 610 + + +def _record_shape(m, eligible, world_size): + global _shape_calls + from mori.ops.gemm_ar import padded_m + + key = (m, padded_m(m, world_size), bool(eligible)) + _shape_hist[key] = _shape_hist.get(key, 0) + 1 + _shape_calls += 1 + if _shape_calls >= _SHAPE_EVERY: + _shape_calls = 0 + items = sorted(_shape_hist.items(), key=lambda kv: -kv[1]) + logger.info( + "wo_b shapes: %s", + " ".join(f"M={m}/pad{p}{'+' if e else '-'}x{c}" for (m, p, e), c in items), + ) + _shape_hist.clear() + + +_logged_config = False + + +def _shuffled_once(): + global _logged_config + if _logged_config: + return True + _logged_config = True + return False + + +def _fused_weight(layer): + """The weight to hand the op, or None if this layer cannot be fused. + + The op requires B preshuffled. sglang does that at load time, but only when + the aiter block-fp8 linear is what consumes the weight and the tuned triton + GEMM does not cover the shape, recording it as ``layer.aiter_bpreshuffled``. + Handing the row-major one to the op does not fail -- it returns an + uncorrelated result -- so check rather than assume. + """ + global _warned_layout + if getattr(layer, "aiter_bpreshuffled", False): + return layer.weight + if not _warned_layout: + _warned_layout = True + logger.warning( + "mori fused wo_b: this layer's weight is not aiter-preshuffled " + "(shape %s), which the fused kernel requires; not fusing.", + tuple(layer.weight.shape), + ) + return None + + +class _FusedWoB: + """The TP communicator and mori's op, one per rank. + + The window, the per-M compile cache and the row padding all live in + ``GemmAllReduceOp``; this holds the cco ``Communicator`` built from sglang's + TP group, which is the part mori cannot construct for us. + """ + + def __init__(self, m_max: int, n: int, k: int): + import torch.distributed as dist + from mori.cco import Communicator, UniqueId + from mori.ops.gemm_ar import GemmAllReduceOp + + from sglang.srt.distributed import get_tp_group + + tp = get_tp_group() + self.rank = tp.rank_in_group + self.world_size = tp.world_size + self.n, self.k, self.m_max = n, k, m_max + + payload = [bytes(Communicator.get_unique_id()) if self.rank == 0 else None] + dist.broadcast_object_list(payload, src=tp.ranks[0], group=tp.cpu_group) + uid = UniqueId.from_bytes(payload[0]) + + # The window is VMM memory outside torch's allocator, so the run has to + # leave room for it (lower --mem-fraction-static). Sized from m_max, and + # it cannot grow afterwards. + gather_dtype = ( + "fp8" if envs.SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER.get() else "bf16" + ) + window_bytes = GemmAllReduceOp.window_bytes_for( + self.world_size, m_max=m_max, n=n, gather_dtype=gather_dtype + ) + self._comm_ctx = Communicator.init( + self.world_size, + self.rank, + uid, + per_rank_vmm=4 * window_bytes + (512 << 20), + ) + self.comm = self._comm_ctx.__enter__() + try: + # Which way the fp8 gather moves. The LSA pull widens on the way in + # instead of in a second kernel, and it wins in both places: 951us + # against SDMA's 1011 on the layer, 1041.2ms of GPU busy against + # 1052.1 in the server. + self.op = GemmAllReduceOp( + self.comm, + n=n, + k=k, + m_max=m_max, + gather_dtype=gather_dtype, + gather_transport=envs.SGLANG_OPT_FUSED_WO_B_AR_GATHER_TRANSPORT.get(), + ) + # Prove the collective moves bytes before serving a single token. + # A mori built with BUILD_CCO_SDMA=OFF compiles the puts out: every + # kernel still launches, nothing moves, the all-reduce returns + # mostly the local slice, and the model still answers fluently while + # being wrong. + # It also measures *faster* that way, which is how a whole + # end-to-end campaign came out at -6.8% instead of -2.3%. This costs + # one collective at m_max, on kernels the first real call compiles + # anyway. + self.op.self_test() + except Exception: + # The communicator holds a VMM reservation the whole process pays + # for. Half-constructing this object and leaving it open makes the + # *fallback* path fail too, which is how a simple attribute error + # here once took the server down. + self._comm_ctx.__exit__(None, None, None) + raise + logger.info( + "mori fused wo_b: window %.0f MiB, tp=%d, M<=%d, N=%d, K=%d, gather=%s", + self.op.window_bytes / 2**20, + self.world_size, + m_max, + n, + k, + gather_dtype, + ) + + def pad_rows(self, x: torch.Tensor, m_pad: int) -> torch.Tensor: + return self.op.pad_rows(x, m_pad) + + def run(self, q_input, x_scale_raw, weight, weight_scale) -> torch.Tensor: + return self.op(q_input, weight, x_scale_raw, weight_scale) + + +def _window_m_max(m_pad: int, world_size: int) -> int: + """Rows the symmetric window is sized for. + + 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 full chunk + later -- the window cannot grow once allocated. + """ + from mori.ops.gemm_ar import padded_m + + from sglang.srt.runtime_context import get_schedule + + # The schedule bag, not the ServerArgs record: the record answers with what + # the operator typed, and chunked_prefill_size is one of the fields + # resolution fills in when it was left unset. + limit = get_schedule().chunked_prefill_size + if limit is not None and limit > 0: + return max(m_pad, padded_m(limit, world_size)) + return m_pad + + +def _eligible(m: int, n: int, k: int, world_size: int) -> bool: + """Whether fusing is both expressible and worth it at this shape. + + mori's ``supports`` answers the first; the thresholds here answer the + second, and stay on this side because they are measured against what the + model runs today rather than being a property of the kernel. + """ + from mori.ops.gemm_ar import padded_m, supports + + if not supports(m, n, k, world_size): + return False + floor = ( + _MIN_FUSED_M_FP8_GATHER + if envs.SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER.get() + else _MIN_FUSED_M + ) + m_pad = padded_m(m, world_size) + return m_pad >= floor and m >= _MIN_PAD_FILL * m_pad + + +def fused_wo_b_available() -> bool: + """Static gate, cheap enough to call per layer.""" + return not _disabled and envs.SGLANG_OPT_FUSED_WO_B_AR.get() + + +def fused_wo_b(layer, x: torch.Tensor) -> Optional[torch.Tensor]: + """wo_b's `GEMM + all-reduce`, fused. ``None`` means "use the normal path". + + ``x`` is the bf16 activation this rank holds, ``[M, K]``. The result is the + all-reduced ``[M, N]``, as a **view of the symmetric window** -- the next + layer's wo_b overwrites it, which is safe because the caller consumes it in + the residual add of the same layer. ``SGLANG_DEBUG_FUSED_WO_B_AR`` forces a copy + and cross-checks against the unfused path. + """ + global _state, _disabled + + if not fused_wo_b_available(): + return None + + import aiter + from mori.ops.gemm_ar import padded_m + + from sglang.srt.distributed import get_tp_group + from sglang.srt.layers.quantization.fp8_utils import aiter_per1x128_quant + + m, k = x.shape + n = layer.weight.shape[0] + world_size = get_tp_group().world_size + eligible = _eligible(m, n, k, world_size) + if envs.SGLANG_OPT_FUSED_WO_B_AR_SHAPE_LOG.get(): + _record_shape(m, eligible, world_size) + if not eligible: + return None + + m_pad = padded_m(m, world_size) + try: + if _state is None: + _state = _FusedWoB( + m_max=_window_m_max(m_pad, world_size), + n=n, + k=k, + ) + if m_pad > _state.m_max or n != _state.n or k != _state.k: + return None + x_in = x if m_pad == m else _state.pad_rows(x, m_pad) + # Same quantisation the unfused path does. transpose_scale=True makes + # the quantiser write the group scale in physical [K/128, M] order, + # which is exactly what the kernel indexes. + q_input, x_scale = aiter_per1x128_quant( + x_in, quant_dtype=aiter.dtypes.fp8, transpose_scale=True + ) + weight = _fused_weight(layer) + if weight is None: + return None + # Flat, explicitly. transpose_scale=True writes the group scale in + # K/128-major order but keeps the (M, K/128) shape, so the tensor's shape + # does not describe its contents and the op refuses to guess -- see + # _flatten_a_scale. Passing the 2-D tensor transposes it and returns an + # answer that is wrong by relL2 0.36, which is perplexity 862511 against + # the unfused path's 3.26. + out = _state.run(q_input, x_scale.reshape(-1), weight, layer.weight_scale_inv) + out = out[:m] + except ValueError as err: + # A shape or contract rejection is about *this call*, not about the + # path. Disabling the process on one would be a standing hazard: the op + # raises ValueError for an M it cannot serve, and a server's M changes + # with every batch, so one unlucky shape used to switch the whole + # optimisation off for good. + global _warned_reject + if not _warned_reject: + _warned_reject = True + logger.warning( + "mori fused wo_b declined a call and fell back for it; further " + "declines are silent: %s", + err, + ) + return None + except Exception as err: # noqa: BLE001 - anything else is not per-call + _disabled = True + logger.warning( + "mori fused wo_b failed and is disabled for this process; " + "falling back to the split path: %s", + err, + ) + return None + + if envs.SGLANG_DEBUG_FUSED_WO_B_AR.get(): + if not _shuffled_once(): + o = _state.op + logger.info( + "mori fused wo_b config: rank=%d/%d m_max=%d n=%d k=%d queues=%d " + "gather=%s/%s m=%d m_pad=%d", + o.rank, + o.world_size, + o.m_max, + o.n, + o.k, + o.sdma_queues, + o.gather_dtype, + o.gather_transport, + m, + m_pad, + ) + out = out.clone() + # Same inputs, immediately again: if the two disagree the server context + # is racing the collective; if they agree the op is deterministic here + # and merely disagrees with the reference. + again = _state.run( + q_input, x_scale.reshape(-1), weight, layer.weight_scale_inv + )[:m].clone() + self_rel = ((out.float() - again.float()).norm() / out.float().norm()).item() + ref, _ = layer(x) + rel = ((out.float() - ref.float()).norm() / ref.float().norm()).item() + logger.info( + "mori fused wo_b check: M=%d relL2=%.3e self_rel=%.3e", m, rel, self_rel + ) + return out diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index b5e2ce5c92b0..03208d3fd1e7 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -108,6 +108,7 @@ should_skip_post_experts_all_reduce, uses_per_rank_fused_shared_slots, ) +from sglang.srt.layers.mori_gemm_ar import fused_wo_b from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod from sglang.srt.layers.quantization.fp8_utils import ( Mxfp8DenseGemmBackend, @@ -2628,10 +2629,22 @@ def _project_wo_b(self, o, defer_all_reduce: bool = False): and not self.wo_b.use_decode_attn_tp and not should_skip_mlp_all_reduce() ) - o, _ = self.wo_b( - o.flatten(1) if isinstance(o, torch.Tensor) else o, - skip_all_reduce=defer_all_reduce, + o_in = o.flatten(1) if isinstance(o, torch.Tensor) else o + # ROCm gfx950 TP-only: wo_b's GEMM and its all-reduce can overlap. The + # helper returns None whenever it does not apply -- decode, ragged + # chunks, no SDMA -- and disables itself for the process on any failure. + # A deferred all-reduce (mHC owns it) and a non-tensor operand (wo_a + # fused an MXFP8 quantization into its epilogue) each keep the ordinary + # path, which owns the collective itself. + fused_o = ( + None + if (defer_all_reduce or not isinstance(o, torch.Tensor)) + else fused_wo_b(self.wo_b, o_in) ) + if fused_o is not None: + o = fused_o + else: + o, _ = self.wo_b(o_in, skip_all_reduce=defer_all_reduce) if defer_all_reduce: return mhc.AttnOutput(o) if self.attn_tp_size > 1 and self.attn_tp_size < get_parallel().tp_size: From b25005996018d805897b40994b9048d58f068fd3 Mon Sep 17 00:00:00 2001 From: xiangch Date: Fri, 9 Oct 2026 06:49:26 +0000 Subject: [PATCH 02/20] [AMD][gfx950] DSV4.1: take the fused wo_b to V4.1-Flash's mxfp8 quantisation The fused wo_b above is written for DeepSeek-V4-Pro, whose wo_b is a 128x128 block-scale fp8 linear preshuffled by the aiter block-fp8 route. V4.1-Flash quantises at 32-wide ue8m0 and takes a gfx950 mxfp8 route instead, so `aiter_bpreshuffled` was never set and the path declined every layer -- after reserving its symmetric window. - Quantise the way V4.1's wo_b does, and take the fp8-grid operand the gfx950 wo_a fork emits rather than requiring a plain tensor. `Mxfp8Activation`, the other form that fork emits, is pre-quantised with no bf16 left to quantise, so it keeps the ordinary path; only decode produces it, which the M floor declines anyway. - Serve both gfx950 mxfp8 routes, not just the native one. They differ in the weight and nothing else -- native rebinds it to the shuffled [N/16, K/128, 2048], aiter leaves the checkpoint's [N, K] -- and `--fp8-gemm-backend aiter`, which the MI355X cookbook pins, takes the second. mori's own `preshuffle_b` is defined on exactly the layout aiter leaves, so the conversion branches and nothing else does. `mxfp8_aiter_ready` records which route ran, because the weight's shape cannot: a wrong guess reads a permuted weight as plain and returns an uncorrelated result. - Re-derive the thresholds against what V4.1 actually runs: the bf16 wire only pays from a padded M of 8192, the fp8 one from 2048. The hooks follow where upstream now puts these routes: the mHC branch moved out of _project_wo_b into its caller's defer_all_reduce, and the native mxfp8 linear is _apply_native in fp8_hip. Measured on 2x MI355X with the cookbook's flags, 16384-token prefill: 21 of 49 NCCL all-reduces move to the fused path, and end-to-end TTFT is -6.9% (medians of 6 fused and 7 base server instances; a single pair does not resolve it, the instance-to-instance spread of aiter's custom all-reduce is of the same order). Co-Authored-By: Claude Opus 5 (1M context) --- docker/rocm.Dockerfile | 2 +- python/sglang/srt/environ.py | 26 +- python/sglang/srt/layers/mori_gemm_ar.py | 464 ++++++++++++------ .../sglang/srt/layers/quantization/fp8_hip.py | 3 + python/sglang/srt/models/deepseek_v4.py | 16 +- 5 files changed, 331 insertions(+), 180 deletions(-) diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 21b63e532fc8..a48669104012 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -398,7 +398,7 @@ ARG ENABLE_MORI=0 ARG NIC_BACKEND=none ARG MORI_REPO="https://github.com/ROCm/mori.git" -ARG MORI_COMMIT="066fd10e549e6e3bbd0a92b98963a9b3ec245414" +ARG MORI_COMMIT="879983bdbd8c65c52e9f79ad836a61cbffef98b6" # NIXL (upstream ai-dynamo/nixl) — KV transfer backend for prefill/decode disaggregation. # Built from source for ROCm; needs UCX built --with-rocm (built here from openucx). diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 32ed92aa7c71..eee80f6840c9 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1630,29 +1630,21 @@ class Envs: # tuned batched_gemm_bf16 (gfx95). Off by default; see deepseek_v4.py # _apply_wo_a_bf16_matmul. SGLANG_OPT_USE_AITER_BATCHED_GEMM = EnvBool(False) - # ROCm gfx950 TP-only: replace wo_b's `GEMM then all-reduce` pair with mori - # cco's fused GEMM+reduce-scatter, which overlaps the scatter with the GEMM. - # Measured 1297us against 1617 per layer at [16384, 7168] K=2048 on 8 ranks. - # Off by default: it needs a ~700 MiB symmetric window outside torch's - # allocator, and it only engages for large M (ragged M is zero-padded up to - # a multiple of tp_size*128). + # ROCm gfx950 TP-only: overlap wo_b's all-reduce with its GEMM through mori + # cco. Off by default: it reserves a symmetric window outside torch's + # allocator, so a run that enables it wants a lower --mem-fraction-static. SGLANG_OPT_FUSED_WO_B_AR = EnvBool(False) # Send the fused wo_b's all-gather leg as fp8 e4m3 with a per-row scale, - # halving its bytes. That leg is ~40% of a fused layer and already runs at - # the xGMI ceiling, so it is worth about -10% on the layer -- but it costs - # accuracy: relL2 goes from 2.35e-3 to 2.49e-2 at the model's shape, which - # is e4m3's 3-bit mantissa and not a granularity that can be tuned away. It - # also needs a larger window (812 MiB against 700) for the staging region. + # halving its bytes. Off by default: it costs accuracy, and that is e4m3's + # mantissa width rather than a granularity that can be tuned away. SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER = EnvBool(False) # Which transport carries the fused wo_b's all-gather leg, "lsa" or "sdma". # LSA pulls, so with an fp8 gather it widens on the way in rather than in a - # second kernel, and it wins in both places (951us against SDMA's 1011 on - # the layer). Only read when the fused path is on. + # second kernel. Only read when the fused path is on. SGLANG_OPT_FUSED_WO_B_AR_GATHER_TRANSPORT = EnvStr("lsa") - # Debug: log the distribution of M the fused wo_b is offered, and how much - # of it clears the fusing thresholds. M is the token count of one forward, - # not of one request, so this is the only way to see what the scheduler - # actually batches. + # Debug: log the distribution of M the fused wo_b is offered. M is the token + # count of one forward, not of one request, so this is the only way to see + # what the scheduler actually batches. SGLANG_OPT_FUSED_WO_B_AR_SHAPE_LOG = EnvBool(False) # Debug: run both the fused and the unfused wo_b and log their relative L2. SGLANG_DEBUG_FUSED_WO_B_AR = EnvBool(False) diff --git a/python/sglang/srt/layers/mori_gemm_ar.py b/python/sglang/srt/layers/mori_gemm_ar.py index f2c2c7e3a755..918f3c434aa2 100644 --- a/python/sglang/srt/layers/mori_gemm_ar.py +++ b/python/sglang/srt/layers/mori_gemm_ar.py @@ -1,95 +1,221 @@ # SPDX-License-Identifier: Apache-2.0 -"""Fused GEMM + all-reduce for DSV4 ``wo_b`` on ROCm gfx950, via mori cco. - -``wo_b`` is a ``RowParallelLinear``, so today it runs a block-scale fp8 GEMM and -then an NCCL all-reduce of the ``[M, hidden]`` result. The two can overlap: -mori's fused kernel writes each destination's row band straight into a symmetric -window and, as soon as a band's tiles are done, hands it to the SDMA copy -engines; the reduce and all-gather then run as their own kernels. - -At ``[16384, 7168]`` K=2048 on 8x MI355X the pair costs 1419.5us and the fused -kernel 1146.2, i.e. **-19.4%**. Essentially all of it is the overlap: the same -pipeline unfused is 1465.9us, and the GEMM alone is 369.4. In-server, over one -20000-token prefill captured with the same warm-up and flush on both sides, GPU -busy goes **1096.0ms -> 1070.3, -25.7ms / -2.3%**, and the fp8 gather takes it -to **1041.2, -5.0%**. The request's own wall time agrees: 1.1969s -> 1.1728 -> -1.1385. The layer-level win is much larger than the end-to-end one because -`wo_b` is about 12% of the profile. - -**Check that mori's cco SDMA path was compiled in before believing any -measurement of this path.** With it off every put silently does nothing: the -all-reduce returns mostly the local slice, the model still answers, and the -fused path looks *faster* than it is because it is not moving any data. Measured -that way it reads -6.8% instead of -2.3%, and the fp8 pull -- the only leg that -does not go through SDMA -- looks like the slowest of the three instead of the -fastest. Perplexity is what catches it: 862511 against 3.26 on the same text. -It is ON by default since ROCm/mori#710, so this is now a question only for a -build that passed `BUILD_CCO_SDMA=OFF` or predates that commit. - -Prerequisites, all of which this module checks rather than assumes: - -* a mori whose cco SDMA path was compiled in -- the default since - ROCm/mori#710. A build that opted out with ``BUILD_CCO_SDMA=OFF`` has no SDMA - queues, every put silently does nothing, and the all-reduce quietly produces - zeros. -* ``MORI_ENABLE_SDMA=1`` at process start. - -The kernel itself is ``mori.ops.gemm_ar.GemmAllReduceOp``, which owns the -symmetric window, the per-M compile cache and the row padding. What stays here -is the part that is sglang's: the TP communicator, how the window is sized, and -whether fusing is worth it at this shape. - -The fused epilogue pushes a whole ``BLOCK_M`` row band to one destination, so a -destination's row slice has to be a whole number of bands: M must be a multiple -of ``tp_size * 128``. Ragged M is zero-padded up to that, which costs GEMM rows -and buys the overlap -- see ``_MIN_PAD_FILL``. Decode and small batches still -fall back to the ordinary path, and that is expected rather than a failure. +"""Fused GEMM + all-reduce for DeepSeek-V4.1-Flash ``wo_b`` on ROCm gfx950, via mori cco. + +``wo_b`` is a ``RowParallelLinear``: an mxfp8 GEMM and then an all-reduce of the +``[M, hidden]`` result. mori's fused kernel writes each destination's row band +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; the reduce and +all-gather still run as their own kernels. + +A mori built with ``BUILD_CCO_SDMA=OFF`` compiles the puts out, and then the +all-reduce returns mostly the local slice while the model still answers -- and +measures *faster*, having moved no bytes. ``GemmAllReduceOp.self_test()`` is +what catches that; the flag has been ON by default since ROCm/mori#710. + +Needs ``MORI_ENABLE_SDMA=1`` at process start and a layer one of the two gfx950 +mxfp8 routes prepared -- see ``mxfp8_route``. + +The fused epilogue pushes a whole ``BLOCK_M`` row band to one destination, so M +must be a multiple of ``tp_size * 256``. That granule is mxfp8's, not a choice +-- see ``_BLOCK_M``. Ragged M is zero-padded up to it, which costs GEMM rows and +buys the overlap -- see ``_MIN_PAD_FILL``. Decode and small batches fall back to +the ordinary path, and that is expected rather than a failure. + +``mori.ops.gemm_ar.GemmAllReduceOp`` owns the symmetric window, the per-M +compile cache and the row padding. What stays here is sglang's: the TP +communicator, how the window is sized, the operand conversions, and whether +fusing is worth it at this shape. """ from __future__ import annotations import logging -from typing import Optional +import os import torch - +import triton +import triton.language as tl from sglang.srt.environ import envs logger = logging.getLogger(__name__) -#: Smallest padded M worth fusing. Measured at ``[*, 7168] K=2048`` on 8 ranks -#: against the aiter GEMM + NCCL pair the model runs today: 1024 -2.6%, 2048 -#: -7.8%, 4096 -12.3%, 8192 -15.7%, 16384 -19.4%. At 1024 a destination gets a -#: single row band, so there is nothing to overlap and the fused kernel is 1.2% -#: *slower* than the same pipeline unfused. -_MIN_FUSED_M = 4096 -#: The fp8 gather used to need a higher floor: with the SDMA transport its two -#: conversion kernels were a fixed cost against a transfer that shrinks with M, -#: so it was 2.3% *slower* at M=4096 and only paid from 8192 up. The LSA pull -#: removed that -- it widens on the way in, so there is no second kernel and no -#: re-read of the landed fp8 -- and fp8 now wins wherever fusing does at all. -#: Measured on the fused path at ``[*, 7168] K=2048``, fp8/lsa against bf16: -#: M=4096 -7.3% (335.8 -> 311.3us), M=8192 -13.8% (593.7 -> 512.0). -#: -#: Keeping 8192 is not merely conservative, it is a loss: a chunk between the -#: two floors takes neither wire and falls back to NCCL entirely. -_MIN_FUSED_M_FP8_GATHER = _MIN_FUSED_M -#: Padding buys the overlap and costs GEMM rows, so the fill ratio has to clear -#: the gain at the padded size. 0.88 is break-even at ``_MIN_FUSED_M`` (12% -#: more rows against a 12.3% gain) and increasingly safe above it. Measured: -#: M=3616 padded to 4096 is 342.1us against 371.4 for the unfused pair. -_MIN_PAD_FILL = 0.88 -_state: Optional[_FusedWoB] = None +# ---- operand conversion ---------------------------------------------- + +#: Rows one packed A-scale group spans, and the rows one quantiser program owns. +#: They are the same 64 by construction, which is what lets the quantiser write +#: mori's layout without any extra traffic -- see `_mxfp8_quant_packed_kernel`. +QUANT_BLOCK_M = 64 + + +def mxfp8_shape(layer) -> tuple[int, int]: + """The layer's logical ``(N, K)``, read off the ue8m0 scale. + + **Not** ``layer.weight.shape``: the native route rebinds the weight to its + shuffled ``[N/16, K/128, 2048]``, and reading that first axis as N + disqualifies every layer without a word. + """ + sn, sk = layer.weight_scale_mx_e8m0.shape + return sn * 32, sk * 32 + + +def mxfp8_route(layer) -> str | None: + """Which gfx950 mxfp8 route prepared this layer, or None if neither did. + + Both leave the ``[N/32, K/32]`` exponent bytes on + ``layer.weight_scale_mx_e8m0`` and differ only in the weight: ``native`` + rebinds it to the shuffled ``[N/16, K/128, 2048]``, ``aiter`` leaves plain + ``[N, K]``. Reading one as the other returns an uncorrelated result rather + than failing, so `mori_weight` branches on this instead of guessing. + """ + if not hasattr(layer, "weight_scale_mx_e8m0"): + return None + if getattr(layer, "mxfp8_native_ready", False): + return "native" + if getattr(layer, "mxfp8_aiter_ready", False): + return "aiter" + return None + + +def mxfp8_ready(layer) -> bool: + """Whether either gfx950 mxfp8 route already prepared this layer.""" + return mxfp8_route(layer) is not None + + +def mori_weight(layer): + """The layer's B operand and B scale in mori's layouts, cached on the layer. + + sglang and mori differ only in how a lane's two 64-wide K sub-blocks sit: + sglang interleaves them inside its 32 bytes, mori keeps them as two 16-byte + blocks. Cached because the permutation copies the whole weight. + """ + cached = getattr(layer, "_mori_b", None) + if cached is not None: + return cached + n, k = mxfp8_shape(layer) + w = layer.weight.data.contiguous().view(torch.uint8) + if mxfp8_route(layer) == "aiter": + # That route leaves the weight in its checkpoint [N, K] order, which is + # what preshuffle_b is defined on, so mori emits its own layout from the + # source rather than permuting sglang's. + from mori.ops.gemm_ar import preshuffle_b + + b = preshuffle_b(w.reshape(n, k).view(layer.weight.dtype)) + else: + # [N/16, K/128, 64 lanes, 2 sub-blocks, 16B] -> the two sub-blocks split out + b = ( + w.reshape(n // 16, k // 128, 64, 2, 16) + .permute(0, 1, 3, 2, 4) + .contiguous() + .reshape(n, k) + .view(layer.weight.dtype) + ) + b_scale = ( + layer.weight_scale_mx_e8m0.data.contiguous() + .view(torch.uint8) + .t() + .contiguous() + .to(torch.int32) + .reshape(-1) + ) + layer._mori_b = (b, b_scale) + return layer._mori_b + + +@triton.jit +def _mxfp8_quant_packed_kernel( + x_ptr, + xq_ptr, + s_ptr, + M, + K, + sxm, + sxk, + sqm, + sqk, + BLOCK_M: tl.constexpr, +): + """sglang's ``_mxfp8_quant_kernel`` writing mori's scale layout directly. + + A variant rather than a stride argument on the original, because the layout + is not expressible as strides: mori wants element ``(m, kb)`` at + ``kb*M + (m//64)*64 + (m%16)*4 + (m%64)//16``, which permutes *within* each + 64-row group so a lane's four M tiles land in one dword. + + Free here: the destination stays inside the 64 bytes this program already + owns, so only the store's order changes. + """ + pid_m = tl.program_id(0) + pid_b = tl.program_id(1) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_k = pid_b * 32 + tl.arange(0, 32) + m_mask = offs_m < M + x = tl.load( + x_ptr + offs_m[:, None] * sxm + offs_k[None, :] * sxk, + mask=m_mask[:, None], + other=0.0, + ).to(tl.float32) + amax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-30) + sb = tl.ceil(tl.log2(amax / 448.0)) + 127.0 + sb = tl.minimum(tl.maximum(sb, 0.0), 254.0) + descale = tl.exp2(sb - 127.0) + xq = tl.clamp(x / descale[:, None], -448.0, 448.0).to(xq_ptr.dtype.element_ty) + tl.store( + xq_ptr + offs_m[:, None] * sqm + offs_k[None, :] * sqk, + xq, + mask=m_mask[:, None], + ) + dst = pid_b * M + (offs_m // 64) * 64 + (offs_m % 16) * 4 + (offs_m % 64) // 16 + tl.store(s_ptr + dst, sb.to(tl.uint8), mask=m_mask) + + +def quantize_packed(x: torch.Tensor): + """bf16 ``[M, K]`` -> (fp8 e4m3 values, mori's packed ue8m0 A scale). + + ``M`` must be a multiple of `QUANT_BLOCK_M`; pad the bf16 input first, so + the scale is built over the padded M and its packed layout needs no repair. + """ + from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import MXFP8_VALUE_DTYPE + + m, k = x.shape + xq = torch.empty((m, k), dtype=MXFP8_VALUE_DTYPE, device=x.device) + scale = torch.empty((k // 32) * m, dtype=torch.uint8, device=x.device) + _mxfp8_quant_packed_kernel[(triton.cdiv(m, QUANT_BLOCK_M), k // 32)]( + x, + xq, + scale, + m, + k, + x.stride(0), + x.stride(1), + xq.stride(0), + xq.stride(1), + BLOCK_M=QUANT_BLOCK_M, + ) + return xq, scale.view(torch.int32) + + +#: Smallest padded M worth fusing on a **bf16** wire, measured against what +#: V4.1-Flash runs today. 5120 is not one answer: the fused cost is fixed by the +#: padded size while the baseline follows the true M, so what separates the +#: winning shapes from the losing ones is a floor rather than a fill ratio. +_MIN_FUSED_M = 8192 +#: The same for the **fp8** wire, a lower bar because that wire runs 15-25% +#: under the bf16 one. Only m_pad 1024 loses: one row band per destination +#: leaves the scatter nothing to overlap with. +_MIN_FUSED_M_FP8_GATHER = 2048 +#: Padding buys the overlap and costs GEMM rows, so a ragged M has to be full +#: enough that the gain at the padded size still clears it. Binds on the fp8 +#: wire only: at m_pad >= 8192 the granule already puts the worst fill at 0.875. +_MIN_PAD_FILL = 0.80 +_state: _FusedWoB | None = None _disabled = False _warned_layout = False _warned_reject = False +_mori_ok: bool | None = None -#: Log what M this layer is actually handed, as a periodic histogram. -#: -#: M is the token count of one forward, not of one request, so whether the fused -#: path engages depends on how the scheduler batches -- which is not something to -#: infer from the client's concurrency. Set -#: SGLANG_OPT_FUSED_WO_B_AR_SHAPE_LOG=1 to find out. +#: M is the token count of one forward, not of one request, so whether the +#: fused path engages depends on how the scheduler batches it. +_SHAPE_LOG = os.environ.get("SGLANG_OPT_FUSED_WO_B_AR_SHAPE_LOG") == "1" _shape_hist: dict = {} _shape_calls = 0 #: 61 wo_b calls make one forward, so this logs roughly every ten of them. @@ -98,9 +224,7 @@ def _record_shape(m, eligible, world_size): global _shape_calls - from mori.ops.gemm_ar import padded_m - - key = (m, padded_m(m, world_size), bool(eligible)) + key = (m, _padded_m(m, world_size), bool(eligible)) _shape_hist[key] = _shape_hist.get(key, 0) + 1 _shape_calls += 1 if _shape_calls >= _SHAPE_EVERY: @@ -113,6 +237,19 @@ def _record_shape(m, eligible, world_size): _shape_hist.clear() +#: mori's mxfp8 kernel packs a lane's four M tiles into one scale dword and +#: picks the byte with the MFMA's opsel, which is four tiles only at BLOCK_M 256 +#: -- so a destination's row band is 256 rows here, not blockscale's 128, and M +#: pads to a multiple of ``tp_size * 256``. +_BLOCK_M = 256 + + +def _padded_m(m: int, world_size: int) -> int: + from mori.ops.gemm_ar import padded_m + + return padded_m(m, world_size, _BLOCK_M) + + _logged_config = False @@ -125,40 +262,32 @@ def _shuffled_once(): def _fused_weight(layer): - """The weight to hand the op, or None if this layer cannot be fused. - - The op requires B preshuffled. sglang does that at load time, but only when - the aiter block-fp8 linear is what consumes the weight and the tuned triton - GEMM does not cover the shape, recording it as ``layer.aiter_bpreshuffled``. - Handing the row-major one to the op does not fail -- it returns an - uncorrelated result -- so check rather than assume. - """ + """The op's B operand and B scale, or None if this layer cannot be fused.""" global _warned_layout - if getattr(layer, "aiter_bpreshuffled", False): - return layer.weight - if not _warned_layout: - _warned_layout = True - logger.warning( - "mori fused wo_b: this layer's weight is not aiter-preshuffled " - "(shape %s), which the fused kernel requires; not fusing.", - tuple(layer.weight.shape), - ) - return None + if not mxfp8_ready(layer): + if not _warned_layout: + _warned_layout = True + logger.warning( + "mori fused wo_b: neither gfx950 mxfp8 route prepared this " + "layer (weight shape %s), which the fused kernel requires; " + "not fusing.", + tuple(layer.weight.shape), + ) + return None + return mori_weight(layer) class _FusedWoB: """The TP communicator and mori's op, one per rank. - The window, the per-M compile cache and the row padding all live in - ``GemmAllReduceOp``; this holds the cco ``Communicator`` built from sglang's - TP group, which is the part mori cannot construct for us. + Holds the cco ``Communicator`` built from sglang's TP group, which is the + part mori cannot construct for us. """ def __init__(self, m_max: int, n: int, k: int): import torch.distributed as dist from mori.cco import Communicator, UniqueId from mori.ops.gemm_ar import GemmAllReduceOp - from sglang.srt.distributed import get_tp_group tp = get_tp_group() @@ -177,7 +306,11 @@ def __init__(self, m_max: int, n: int, k: int): "fp8" if envs.SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER.get() else "bf16" ) window_bytes = GemmAllReduceOp.window_bytes_for( - self.world_size, m_max=m_max, n=n, gather_dtype=gather_dtype + self.world_size, + m_max=m_max, + n=n, + block_m=_BLOCK_M, + gather_dtype=gather_dtype, ) self._comm_ctx = Communicator.init( self.world_size, @@ -196,18 +329,17 @@ def __init__(self, m_max: int, n: int, k: int): n=n, k=k, m_max=m_max, + quant="mxfp8", gather_dtype=gather_dtype, - gather_transport=envs.SGLANG_OPT_FUSED_WO_B_AR_GATHER_TRANSPORT.get(), + gather_transport=os.environ.get( + "SGLANG_OPT_FUSED_WO_B_AR_GATHER_TRANSPORT", "lsa" + ), ) - # Prove the collective moves bytes before serving a single token. - # A mori built with BUILD_CCO_SDMA=OFF compiles the puts out: every - # kernel still launches, nothing moves, the all-reduce returns - # mostly the local slice, and the model still answers fluently while - # being wrong. - # It also measures *faster* that way, which is how a whole - # end-to-end campaign came out at -6.8% instead of -2.3%. This costs - # one collective at m_max, on kernels the first real call compiles - # anyway. + # Prove the collective moves bytes before serving a token. A mori + # built with BUILD_CCO_SDMA=OFF compiles the puts out: the kernels + # still launch, the all-reduce returns mostly the local slice, and + # the model answers fluently while being wrong -- and measures + # *faster*, having moved nothing. self.op.self_test() except Exception: # The communicator holds a VMM reservation the whole process pays @@ -240,16 +372,13 @@ def _window_m_max(m_pad: int, world_size: int) -> int: a short prompt arriving first cannot fix a window too small for a full chunk later -- the window cannot grow once allocated. """ - from mori.ops.gemm_ar import padded_m - + # `get_global_server_args()` is retired; the scheduling namespace carries + # the value in effect, which is what the other layers read. from sglang.srt.runtime_context import get_schedule - # The schedule bag, not the ServerArgs record: the record answers with what - # the operator typed, and chunked_prefill_size is one of the fields - # resolution fills in when it was left unset. limit = get_schedule().chunked_prefill_size if limit is not None and limit > 0: - return max(m_pad, padded_m(limit, world_size)) + return max(m_pad, _padded_m(limit, world_size)) return m_pad @@ -257,28 +386,55 @@ def _eligible(m: int, n: int, k: int, world_size: int) -> bool: """Whether fusing is both expressible and worth it at this shape. mori's ``supports`` answers the first; the thresholds here answer the - second, and stay on this side because they are measured against what the - model runs today rather than being a property of the kernel. + second, which is not a property of the kernel. """ - from mori.ops.gemm_ar import padded_m, supports + from mori.ops.gemm_ar import supports - if not supports(m, n, k, world_size): + if not supports(m, n, k, world_size, quant="mxfp8"): return False floor = ( _MIN_FUSED_M_FP8_GATHER if envs.SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER.get() else _MIN_FUSED_M ) - m_pad = padded_m(m, world_size) + m_pad = _padded_m(m, world_size) return m_pad >= floor and m >= _MIN_PAD_FILL * m_pad +def _mori_has_gemm_ar() -> bool: + """Whether the pinned mori carries ``ops.gemm_ar``, memoised. + + Checked rather than assumed because the import sits on the per-call path + below the switch: a mori predating the op would raise from inside a forward + rather than fall back. + """ + global _mori_ok, _disabled + if _mori_ok is None: + try: + import mori.ops.gemm_ar # noqa: F401 + + _mori_ok = True + except ImportError as err: + _mori_ok = False + _disabled = True + logger.warning( + "mori fused wo_b: this mori has no ops.gemm_ar, so the fused " + "path is off for this process: %s", + err, + ) + return _mori_ok + + def fused_wo_b_available() -> bool: """Static gate, cheap enough to call per layer.""" - return not _disabled and envs.SGLANG_OPT_FUSED_WO_B_AR.get() + return ( + not _disabled + and envs.SGLANG_OPT_FUSED_WO_B_AR.get() + and _mori_has_gemm_ar() + ) -def fused_wo_b(layer, x: torch.Tensor) -> Optional[torch.Tensor]: +def fused_wo_b(layer, x: torch.Tensor) -> torch.Tensor | None: """wo_b's `GEMM + all-reduce`, fused. ``None`` means "use the normal path". ``x`` is the bf16 activation this rank holds, ``[M, K]``. The result is the @@ -292,22 +448,36 @@ def fused_wo_b(layer, x: torch.Tensor) -> Optional[torch.Tensor]: if not fused_wo_b_available(): return None - import aiter - from mori.ops.gemm_ar import padded_m - + from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import Fp8GridActivation from sglang.srt.distributed import get_tp_group - from sglang.srt.layers.quantization.fp8_utils import aiter_per1x128_quant + + # wo_a hands wo_b either a plain bf16 activation or an Fp8GridActivation -- + # a bf16 tensor already rounded onto wo_b's fp8 grid, so quantising it below + # is lossless rather than a second rounding. Either way what the GEMM needs + # is the bf16 tensor. + if isinstance(x, Fp8GridActivation): + x = x.x + # wo_a's MXFP8 epilogue (SGLANG_HIP_WO_A_MXFP8) instead hands over fp8 plus + # a ue8m0 scale, with no bf16 left to quantise, and the op does its own + # quantisation. Only decode reaches that form, which _eligible declines on + # M anyway, so declining here costs no fusing. + if not isinstance(x, torch.Tensor): + return None m, k = x.shape - n = layer.weight.shape[0] + if not hasattr(layer, "weight_scale_mx_e8m0"): + return None + n, w_k = mxfp8_shape(layer) + if w_k != k: + return None world_size = get_tp_group().world_size eligible = _eligible(m, n, k, world_size) - if envs.SGLANG_OPT_FUSED_WO_B_AR_SHAPE_LOG.get(): + if _SHAPE_LOG: _record_shape(m, eligible, world_size) if not eligible: return None - m_pad = padded_m(m, world_size) + m_pad = _padded_m(m, world_size) try: if _state is None: _state = _FusedWoB( @@ -317,23 +487,17 @@ def fused_wo_b(layer, x: torch.Tensor) -> Optional[torch.Tensor]: ) if m_pad > _state.m_max or n != _state.n or k != _state.k: return None - x_in = x if m_pad == m else _state.pad_rows(x, m_pad) - # Same quantisation the unfused path does. transpose_scale=True makes - # the quantiser write the group scale in physical [K/128, M] order, - # which is exactly what the kernel indexes. - q_input, x_scale = aiter_per1x128_quant( - x_in, quant_dtype=aiter.dtypes.fp8, transpose_scale=True - ) - weight = _fused_weight(layer) - if weight is None: + prepared = _fused_weight(layer) + if prepared is None: return None - # Flat, explicitly. transpose_scale=True writes the group scale in - # K/128-major order but keeps the (M, K/128) shape, so the tensor's shape - # does not describe its contents and the op refuses to guess -- see - # _flatten_a_scale. Passing the 2-D tensor transposes it and returns an - # answer that is wrong by relL2 0.36, which is perplexity 862511 against - # the unfused path's 3.26. - out = _state.run(q_input, x_scale.reshape(-1), weight, layer.weight_scale_inv) + weight, b_scale = prepared + x_in = x if m_pad == m else _state.pad_rows(x, m_pad) + # The quantiser writes mori's packed scale layout itself, so there is no + # conversion pass after it. Zero-padded rows quantise to zero values + # (their scale is tiny but finite), and rows are independent in a GEMM, + # so the padding contributes nothing to any real row. + q_input, x_scale = quantize_packed(x_in) + out = _state.run(q_input, x_scale, weight, b_scale) out = out[:m] except ValueError as err: # A shape or contract rejection is about *this call*, not about the @@ -380,9 +544,7 @@ def fused_wo_b(layer, x: torch.Tensor) -> Optional[torch.Tensor]: # Same inputs, immediately again: if the two disagree the server context # is racing the collective; if they agree the op is deterministic here # and merely disagrees with the reference. - again = _state.run( - q_input, x_scale.reshape(-1), weight, layer.weight_scale_inv - )[:m].clone() + again = _state.run(q_input, x_scale, weight, b_scale)[:m].clone() self_rel = ((out.float() - again.float()).norm() / out.float().norm()).item() ref, _ = layer(x) rel = ((out.float() - ref.float()).norm() / ref.float().norm()).item() diff --git a/python/sglang/srt/layers/quantization/fp8_hip.py b/python/sglang/srt/layers/quantization/fp8_hip.py index 5223e3c170c9..b1533dc64c5a 100644 --- a/python/sglang/srt/layers/quantization/fp8_hip.py +++ b/python/sglang/srt/layers/quantization/fp8_hip.py @@ -40,6 +40,9 @@ def process_dense_weights(method, layer: torch.nn.Module, scale_u8) -> None: "weight_scale_mx_e8m0", ue8m0_weight_scale(layer.weight_scale_inv.data), ) + # This route leaves [N, K] where the native one rebinds to + # [N/16, K/128, 2048], and the shape does not say which ran. + layer.mxfp8_aiter_ready = True return assert backend.is_gfx95_mxfp8_native() n, k = layer.weight.shape diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 03208d3fd1e7..d5b692b2361c 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -2629,18 +2629,12 @@ def _project_wo_b(self, o, defer_all_reduce: bool = False): and not self.wo_b.use_decode_attn_tp and not should_skip_mlp_all_reduce() ) + # the gfx950 wo_a fork already returned [T, G * R], so only the plain + # tensor needs flattening o_in = o.flatten(1) if isinstance(o, torch.Tensor) else o - # ROCm gfx950 TP-only: wo_b's GEMM and its all-reduce can overlap. The - # helper returns None whenever it does not apply -- decode, ragged - # chunks, no SDMA -- and disables itself for the process on any failure. - # A deferred all-reduce (mHC owns it) and a non-tensor operand (wo_a - # fused an MXFP8 quantization into its epilogue) each keep the ordinary - # path, which owns the collective itself. - fused_o = ( - None - if (defer_all_reduce or not isinstance(o, torch.Tensor)) - else fused_wo_b(self.wo_b, o_in) - ) + # Fusing *performs* the all-reduce, where a deferred one is skipped so + # the mHC post can fold it in -- a correctness constraint, not a gap. + fused_o = None if defer_all_reduce else fused_wo_b(self.wo_b, o_in) if fused_o is not None: o = fused_o else: From 3dd6ca3368a9fe26396ece19801267c6d195de79 Mon Sep 17 00:00:00 2001 From: xiangch Date: Sat, 10 Oct 2026 10:01:19 +0000 Subject: [PATCH 03/20] [AMD][gfx950] Address review on the fused wo_b **DP attention: the fused path reduced over the wrong group, and could hang.** `wo_b` is declared `parallel_group="attn_tp"`, while the op always reduces over the full TP group. The call site gated only `defer_all_reduce`, and everything that made *that* false reached the fused path -- including `attn_tp_size == 1`, where the split path reduces nothing and fusing summed rows from unrelated DP ranks, and `1 < attn_tp_size < tp_size`, where `attn_tp_all_reduce` then ran on top of a full-TP sum. It could also hang: under DP attention each rank has its own M, so eligibility differed across ranks and mori's device-side barriers have no timeout. `_reduces_over_tp_group` now asks the layer, using `RowParallelLinear.forward`'s own conditions, and `fused_wo_b` checks it before anything collective. The env var's comment said "TP-only"; nothing enforced it. **A rank-local disable deadlocked the TP group.** `_disabled` was set only on the rank that raised, which then fell back to NCCL while its peers kept calling the op and spun on its barriers. Construction is now collective: every rank attempts it, the outcome is agreed over the CPU group, and ranks that succeeded tear their communicator back down if a peer did not. After construction there is no fallback left -- an OOM in `pad_rows` or `quantize_packed` is rank-local and `run` is collective, so returning None from either is a silent hang where dying is visible. Every remaining reason to decline is an explicit check that reads the same on all ranks. **A construction `ValueError` was retried on every call.** One `try` covered both construction and the per-call work, with a handler written for the latter: `_state` stayed None and each of the 61 eligible calls per forward repeated the broadcast, the VMM reservation and the teardown. Construction is attempted once and a failure is permanent. **The second copy of each `wo_b` weight is built before the profiler runs.** `mori_weight` re-lays the weight out -- a full second copy, since the split path still needs the original -- and did it lazily on the first fused call, after the memory profiler had sized the KV cache around its absence (61 layers, ~1.3 GiB at TP2). `prepare_wo_b_weight` now runs in `post_load_weights`. **The window is sized from the deployment's prefill ceiling.** With chunked prefill off, `m_max` fell back to the first eligible M -- the first-request sizing the docstring says it avoids -- and every larger M then declined silently for the life of the process. It now uses `max_prefill_buffer_tokens`, as the other buffer-sizing callers do, and logs once when a call exceeds the window. **Smaller** - `SGLANG_OPT_FUSED_WO_B_AR_SHAPE_LOG` and `..._GATHER_TRANSPORT` go through the `EnvField` API. The former was read at import and compared against "1", so the `EnvBool`-accepted `=true` was ignored and `override()` did nothing. - One ue8m0 scale rule: the quantiser calls `fp8_grid_quant` rather than restating it as `ceil(log2(amax / 448))`. The two are not the same function -- `fp8_grid_quant` reads the exponent off the IEEE bits and is exact at powers of two, where an approximate `log2` can land above the integer and `ceil` picks one exponent too high. Their agreeing is what the claim that re-quantising an `Fp8GridActivation` is lossless rests on. - The static gate memoises the env read, and `m_pad` is computed once per call instead of three times. Co-Authored-By: Claude Opus 5 (1M context) --- python/sglang/srt/layers/mori_gemm_ar.py | 331 +++++++++++++++++------ python/sglang/srt/models/deepseek_v4.py | 10 +- 2 files changed, 261 insertions(+), 80 deletions(-) diff --git a/python/sglang/srt/layers/mori_gemm_ar.py b/python/sglang/srt/layers/mori_gemm_ar.py index 918f3c434aa2..06b98ac8fec9 100644 --- a/python/sglang/srt/layers/mori_gemm_ar.py +++ b/python/sglang/srt/layers/mori_gemm_ar.py @@ -30,11 +30,15 @@ from __future__ import annotations import logging -import os import torch import triton import triton.language as tl + +# Module scope, not inside the kernel: Triton resolves a called `@triton.jit` +# function when it compiles the caller's AST, and a local import does not +# reliably reach it. +from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import fp8_grid_quant from sglang.srt.environ import envs logger = logging.getLogger(__name__) @@ -134,7 +138,7 @@ def _mxfp8_quant_packed_kernel( sqk, BLOCK_M: tl.constexpr, ): - """sglang's ``_mxfp8_quant_kernel`` writing mori's scale layout directly. + """``fp8_grid_quant`` writing mori's scale layout directly. A variant rather than a stride argument on the original, because the layout is not expressible as strides: mori wants element ``(m, kb)`` at @@ -143,6 +147,15 @@ def _mxfp8_quant_packed_kernel( Free here: the destination stays inside the 64 bytes this program already owns, so only the store's order changes. + + The scale rule is imported rather than restated. This used to carry its own + copy of it -- ``ceil(log2(amax / 448))`` -- which is not the same function: + ``fp8_grid_quant`` reads the exponent off the IEEE bits and is therefore + exact when ``amax / 448`` is a power of two, where an approximate ``log2`` + can land just above the integer and ``ceil`` then picks one exponent too + high. Both rules agreeing is exactly what ``fused_wo_b``'s claim that + re-quantising an ``Fp8GridActivation`` is lossless rests on, so there must + be one of them. """ pid_m = tl.program_id(0) pid_b = tl.program_id(1) @@ -154,11 +167,7 @@ def _mxfp8_quant_packed_kernel( mask=m_mask[:, None], other=0.0, ).to(tl.float32) - amax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-30) - sb = tl.ceil(tl.log2(amax / 448.0)) + 127.0 - sb = tl.minimum(tl.maximum(sb, 0.0), 254.0) - descale = tl.exp2(sb - 127.0) - xq = tl.clamp(x / descale[:, None], -448.0, 448.0).to(xq_ptr.dtype.element_ty) + xq, sb = fp8_grid_quant(x) tl.store( xq_ptr + offs_m[:, None] * sqm + offs_k[None, :] * sqk, xq, @@ -210,21 +219,19 @@ def quantize_packed(x: torch.Tensor): _state: _FusedWoB | None = None _disabled = False _warned_layout = False -_warned_reject = False +_warned_window = False _mori_ok: bool | None = None +_switch_on: bool | None = None -#: M is the token count of one forward, not of one request, so whether the -#: fused path engages depends on how the scheduler batches it. -_SHAPE_LOG = os.environ.get("SGLANG_OPT_FUSED_WO_B_AR_SHAPE_LOG") == "1" _shape_hist: dict = {} _shape_calls = 0 #: 61 wo_b calls make one forward, so this logs roughly every ten of them. _SHAPE_EVERY = 610 -def _record_shape(m, eligible, world_size): +def _record_shape(m, m_pad, eligible): global _shape_calls - key = (m, _padded_m(m, world_size), bool(eligible)) + key = (m, m_pad, bool(eligible)) _shape_hist[key] = _shape_hist.get(key, 0) + 1 _shape_calls += 1 if _shape_calls >= _SHAPE_EVERY: @@ -261,6 +268,28 @@ def _shuffled_once(): return False +def prepare_wo_b_weight(layer) -> None: + """Build this layer's fused B operand now, at weight-processing time. + + ``mori_weight`` re-lays the weight out, which is a full second copy of it + -- the original stays, because decode and every declined call still read + it. Doing that lazily on the first fused call allocated it *while serving*, + after the memory profiler had already sized the KV cache around its + absence: 61 layers of it, about 1.3 GiB at TP2. Preparing here puts it in + front of the profiler instead, and turns a mid-serving OOM into one at + startup. + + Best effort on purpose. A failure here leaves ``_mori_b`` unset, the per + call path falls back to building it, and that path is identical on every + rank -- so this is not a place where ranks can disagree. + """ + if not fused_wo_b_available() or not _reduces_over_tp_group(layer): + return + if not mxfp8_ready(layer) or not hasattr(layer, "weight_scale_mx_e8m0"): + return + mori_weight(layer) + + def _fused_weight(layer): """The op's B operand and B scale, or None if this layer cannot be fused.""" global _warned_layout @@ -331,9 +360,7 @@ def __init__(self, m_max: int, n: int, k: int): m_max=m_max, quant="mxfp8", gather_dtype=gather_dtype, - gather_transport=os.environ.get( - "SGLANG_OPT_FUSED_WO_B_AR_GATHER_TRANSPORT", "lsa" - ), + gather_transport=envs.SGLANG_OPT_FUSED_WO_B_AR_GATHER_TRANSPORT.get(), ) # Prove the collective moves bytes before serving a token. A mori # built with BUILD_CCO_SDMA=OFF compiles the puts out: the kernels @@ -364,25 +391,118 @@ def pad_rows(self, x: torch.Tensor, m_pad: int) -> torch.Tensor: def run(self, q_input, x_scale_raw, weight, weight_scale) -> torch.Tensor: return self.op(q_input, weight, x_scale_raw, weight_scale) + def close(self) -> None: + """Release the communicator's VMM reservation. + + Used when a peer failed to construct: this rank's op is sound but + useless, and the reservation is a per-process cost for the lifetime of + the server whether or not anything calls it. + """ + self.op = None + self._comm_ctx.__exit__(None, None, None) + + +def _construct_collectively(m_max: int, n: int, k: int) -> bool: + """Build the op on every rank, or disable the path on every rank. + + The decision has to be collective. mori's phases end in device-side + cross-rank barriers, so a rank that builds and uses the op while a peer + fell back to the split path spins on the GPU forever: no timeout, no + error, a hung server. That is what a rank-local ``_disabled`` produced. + + Rank-local failures are realistic -- an OOM in ``Communicator.init``'s VMM + reservation, or in ``self_test``'s buffers -- so this agrees on the outcome + over the CPU group before any rank is allowed to call the op, and the ranks + that *succeeded* tear their communicator down again when a peer did not. + + Construction is attempted once per process. A failure disables the path for + good rather than being retried: every retry is another broadcast, another + multi-hundred-MiB VMM reservation and another teardown, 61 times a forward. + """ + global _state, _disabled + + import torch.distributed as dist + + from sglang.srt.distributed import get_tp_group + + tp = get_tp_group() + state = None + err = None + try: + state = _FusedWoB(m_max=m_max, n=n, k=k) + except Exception as exc: # noqa: BLE001 - reported below, then disabled + err = exc + + ok = torch.tensor([0 if err is not None else 1], dtype=torch.int32) + dist.all_reduce(ok, op=dist.ReduceOp.MIN, group=tp.cpu_group) + if int(ok.item()) == 1: + _state = state + return True + + _disabled = True + if state is not None: + state.close() + if err is not None: + logger.warning( + "mori fused wo_b could not be constructed on rank %d; the fused " + "path is off for this process on every rank: %s", + tp.rank_in_group, + err, + ) + else: + logger.warning( + "mori fused wo_b could not be constructed on a peer rank; the " + "fused path is off for this process on every rank" + ) + return False + def _window_m_max(m_pad: int, world_size: int) -> int: """Rows the symmetric window is sized for. - 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 full chunk - later -- the window cannot grow once allocated. + Taken from the deployment's prefill ceiling rather than the first request + seen, because the window cannot grow once allocated: a short prompt + arriving first would otherwise fix a window too small for everything after + it, and every later call would silently fall back for the life of the + process. + + ``max_prefill_buffer_tokens`` is what the other buffer-sizing callers here + use, and it answers in the cases a bare ``chunked_prefill_size`` does not + -- chunked prefill disabled, and PP dynamic chunking, which probes above + the chunk size. Its own zero case falls through to ``max_prefill_tokens``, + the same way ``disaggregation/common/conn.py`` does it. """ - # `get_global_server_args()` is retired; the scheduling namespace carries - # the value in effect, which is what the other layers read. - from sglang.srt.runtime_context import get_schedule + from sglang.srt.runtime_context import get_schedule, max_prefill_buffer_tokens - limit = get_schedule().chunked_prefill_size - if limit is not None and limit > 0: - return max(m_pad, _padded_m(limit, world_size)) + limit = max_prefill_buffer_tokens() or get_schedule().max_prefill_tokens + if limit: + return max(m_pad, _padded_m(int(limit), world_size)) + # Nothing declares a ceiling. Sizing from this call is then the only option + # left, and `_warn_once_over_window` reports it if a later M outgrows it. return m_pad -def _eligible(m: int, n: int, k: int, world_size: int) -> bool: +def _warn_once_over_window(m_pad: int, m_max: int) -> None: + """Report the first M the window cannot hold. + + Declining is correct -- the window cannot grow -- but it is also permanent + for every M at least this large, so it should not be silent. One line, + because it then repeats on most prefills. + """ + global _warned_window + if _warned_window: + return + _warned_window = True + logger.warning( + "mori fused wo_b: M=%d exceeds the window's %d rows, falling back for " + "this and any larger M; size it with --chunked-prefill-size or " + "--max-prefill-tokens. Further declines are silent.", + m_pad, + m_max, + ) + + +def _eligible(m: int, m_pad: int, n: int, k: int, world_size: int) -> bool: """Whether fusing is both expressible and worth it at this shape. mori's ``supports`` answers the first; the thresholds here answer the @@ -397,10 +517,57 @@ def _eligible(m: int, n: int, k: int, world_size: int) -> bool: if envs.SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER.get() else _MIN_FUSED_M ) - m_pad = _padded_m(m, world_size) return m_pad >= floor and m >= _MIN_PAD_FILL * m_pad +def _reduces_over_tp_group(layer) -> bool: + """Whether this layer's own all-reduce is the one the op would perform. + + Fusing *performs* the all-reduce, so it is a substitute only where the + split path would have done the same one. The op always reduces over the + full TP group, while ``wo_b`` is declared ``parallel_group="attn_tp"`` -- + the same ranks only when attention is not data-parallel. Under + ``--attn-dp-size 8`` the attn-TP group is one rank, the split path reduces + nothing, and fusing would sum rows belonging to eight unrelated DP ranks. + + The conditions are ``RowParallelLinear.forward``'s own, read off the layer + rather than restated at the call site. That is where the previous version + got it wrong: the caller gated ``defer_all_reduce`` and let everything + making *that* false through to the fused path, including the cases where + the layer reduces over another group or does not reduce at all. + + Static -- true or false for the life of the process -- so weight prep can + ask it too. ``_substitutes_layer_all_reduce`` adds the per-forward flag. + + Not checked: ``quantize_communications`` would make the split path's reduce + numerically different, but it is rejected at startup on anything but NPU. + ``skip_all_reduce`` is the caller's own argument, and it passes False + whenever it reaches here. + """ + from sglang.srt.distributed import get_tp_group + from sglang.srt.distributed.utils import get_group_rank_size + + _, tp_size = get_group_rank_size(layer.tp_group) + if not (layer.reduce_results and tp_size > 1): + return False + # Both of these reduce over the attn-TP group instead. + if layer.use_decode_attn_tp or layer.use_dp_attention_reduce: + return False + # Equal size would do while attn_tp is a subgroup of tp, but comparing the + # ranks says what is meant and does not rest on that staying true. + return tuple(layer.tp_group.ranks) == tuple(get_tp_group().ranks) + + +def _substitutes_layer_all_reduce(layer) -> bool: + """``_reduces_over_tp_group`` plus the flags that vary per forward.""" + from sglang.srt.layers.moe.utils import should_skip_mlp_all_reduce + + # The mHC post folds the reduce in; the split path skips it here. + if should_skip_mlp_all_reduce(): + return False + return _reduces_over_tp_group(layer) + + def _mori_has_gemm_ar() -> bool: """Whether the pinned mori carries ``ops.gemm_ar``, memoised. @@ -426,12 +593,20 @@ def _mori_has_gemm_ar() -> bool: def fused_wo_b_available() -> bool: - """Static gate, cheap enough to call per layer.""" - return ( - not _disabled - and envs.SGLANG_OPT_FUSED_WO_B_AR.get() - and _mori_has_gemm_ar() - ) + """Static gate, called for every layer of every forward. + + The switch is memoised because it cannot change after startup and this sits + on the per-layer path of decode and of non-ROCm runs alike, where it was an + env lookup per call to return False 61 times a forward. ``_disabled`` stays + dynamic -- it is what a collective construction failure sets -- and is + checked first so the off path is one bool. + """ + global _switch_on + if _disabled: + return False + if _switch_on is None: + _switch_on = envs.SGLANG_OPT_FUSED_WO_B_AR.get() + return _switch_on and _mori_has_gemm_ar() def fused_wo_b(layer, x: torch.Tensor) -> torch.Tensor | None: @@ -451,6 +626,9 @@ def fused_wo_b(layer, x: torch.Tensor) -> torch.Tensor | None: from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import Fp8GridActivation from sglang.srt.distributed import get_tp_group + if not _substitutes_layer_all_reduce(layer): + return None + # wo_a hands wo_b either a plain bf16 activation or an Fp8GridActivation -- # a bf16 tensor already rounded onto wo_b's fp8 grid, so quantising it below # is lossless rather than a second rounding. Either way what the GEMM needs @@ -471,57 +649,52 @@ def fused_wo_b(layer, x: torch.Tensor) -> torch.Tensor | None: if w_k != k: return None world_size = get_tp_group().world_size - eligible = _eligible(m, n, k, world_size) - if _SHAPE_LOG: - _record_shape(m, eligible, world_size) + # Computed once and threaded through: it was recomputed in `_eligible`, in + # `_record_shape` and again below, and it is a call into mori. + m_pad = _padded_m(m, world_size) + eligible = _eligible(m, m_pad, n, k, world_size) + # M is the token count of one forward, not of one request, so whether the + # fused path engages depends on how the scheduler batches it. Read per call + # rather than at import, so `envs....override()` in a test is honoured. + if envs.SGLANG_OPT_FUSED_WO_B_AR_SHAPE_LOG.get(): + _record_shape(m, m_pad, eligible) if not eligible: return None - m_pad = _padded_m(m, world_size) - try: - if _state is None: - _state = _FusedWoB( - m_max=_window_m_max(m_pad, world_size), - n=n, - k=k, - ) - if m_pad > _state.m_max or n != _state.n or k != _state.k: - return None - prepared = _fused_weight(layer) - if prepared is None: - return None - weight, b_scale = prepared - x_in = x if m_pad == m else _state.pad_rows(x, m_pad) - # The quantiser writes mori's packed scale layout itself, so there is no - # conversion pass after it. Zero-padded rows quantise to zero values - # (their scale is tiny but finite), and rows are independent in a GEMM, - # so the padding contributes nothing to any real row. - q_input, x_scale = quantize_packed(x_in) - out = _state.run(q_input, x_scale, weight, b_scale) - out = out[:m] - except ValueError as err: - # A shape or contract rejection is about *this call*, not about the - # path. Disabling the process on one would be a standing hazard: the op - # raises ValueError for an M it cannot serve, and a server's M changes - # with every batch, so one unlucky shape used to switch the whole - # optimisation off for good. - global _warned_reject - if not _warned_reject: - _warned_reject = True - logger.warning( - "mori fused wo_b declined a call and fell back for it; further " - "declines are silent: %s", - err, - ) + if _state is None and not _construct_collectively( + _window_m_max(m_pad, world_size), n, k + ): return None - except Exception as err: # noqa: BLE001 - anything else is not per-call - _disabled = True - logger.warning( - "mori fused wo_b failed and is disabled for this process; " - "falling back to the split path: %s", - err, - ) + + # Every reason to decline a call is a check here rather than an exception + # caught below, and every one of them reads the same on all ranks -- the + # shape is identical across a TP group, and `_substitutes_layer_all_reduce` + # has already excluded the configurations where it is not. That matters + # more than it looks: declining on one rank and fusing on another hangs the + # group on mori's device-side barriers. + if m_pad > _state.m_max: + _warn_once_over_window(m_pad, _state.m_max) + return None + if n != _state.n or k != _state.k: + return None + prepared = _fused_weight(layer) + if prepared is None: return None + weight, b_scale = prepared + + # Past this point there is no fallback, deliberately. An OOM in pad_rows or + # quantize_packed is rank-local, and `run` is collective: returning None + # from either would put this rank on the split path while its peers sit in + # mori's barriers, which is a silent hang rather than a visible failure. A + # process that dies is the better outcome and is what the caller can see. + x_in = x if m_pad == m else _state.pad_rows(x, m_pad) + # The quantiser writes mori's packed scale layout itself, so there is no + # conversion pass after it. Zero-padded rows quantise to zero values (their + # scale is tiny but finite), and rows are independent in a GEMM, so the + # padding contributes nothing to any real row. + q_input, x_scale = quantize_packed(x_in) + out = _state.run(q_input, x_scale, weight, b_scale) + out = out[:m] if envs.SGLANG_DEBUG_FUSED_WO_B_AR.get(): if not _shuffled_once(): diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index d5b692b2361c..bd42a9691057 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -108,7 +108,7 @@ should_skip_post_experts_all_reduce, uses_per_rank_fused_shared_slots, ) -from sglang.srt.layers.mori_gemm_ar import fused_wo_b +from sglang.srt.layers.mori_gemm_ar import fused_wo_b, prepare_wo_b_weight from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod from sglang.srt.layers.quantization.fp8_utils import ( Mxfp8DenseGemmBackend, @@ -2634,6 +2634,9 @@ def _project_wo_b(self, o, defer_all_reduce: bool = False): o_in = o.flatten(1) if isinstance(o, torch.Tensor) else o # Fusing *performs* the all-reduce, where a deferred one is skipped so # the mHC post can fold it in -- a correctness constraint, not a gap. + # `defer_all_reduce` being false is not on its own a licence to fuse: + # it is also false when the layer reduces over another group or not at + # all. `fused_wo_b` checks that for itself against the layer. fused_o = None if defer_all_reduce else fused_wo_b(self.wo_b, o_in) if fused_o is not None: o = fused_o @@ -4795,6 +4798,11 @@ def post_load_weights(self, is_nextn=False, weight_names=None): ): self_attn.indexer.compressor.apply_ape_hotfix() layer.refresh_mhc_norm_weight_cache() + # Before the memory profiler sizes the KV cache: the fused wo_b + # keeps a second, re-laid-out copy of this weight, and building it + # on the first fused call would allocate it behind the profiler's + # back. + prepare_wo_b_weight(self_attn.wo_b) layers = self.model.layers for i, layer in enumerate(layers): if isinstance(layer, DeepseekV4DecoderLayer): From 18dd7fd47060a1f939683c930993f82f4c58cc5f Mon Sep 17 00:00:00 2001 From: xiangch Date: Sat, 10 Oct 2026 11:44:18 +0000 Subject: [PATCH 04/20] [AMD][gfx950] Prepare the fused wo_b weight where the mxfp8 layout exists `post_load_weights` runs before `process_weights_after_loading`, so `mxfp8_ready` was false there and `prepare_wo_b_weight` returned without allocating -- leaving the second weight copy to be built lazily during serving, which is the thing it was meant to stop. Measured: KV cache and available GPU memory were byte-identical with the switch on and off. `precompile_kernels_after_loading` runs after the quant method has laid the weight out and before the memory pool is sized, which is the window this needs. Co-Authored-By: Claude Opus 5 (1M context) --- python/sglang/srt/models/deepseek_v4.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index bd42a9691057..f7127da910b0 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -4798,11 +4798,6 @@ def post_load_weights(self, is_nextn=False, weight_names=None): ): self_attn.indexer.compressor.apply_ape_hotfix() layer.refresh_mhc_norm_weight_cache() - # Before the memory profiler sizes the KV cache: the fused wo_b - # keeps a second, re-laid-out copy of this weight, and building it - # on the first fused call would allocate it behind the profiler's - # back. - prepare_wo_b_weight(self_attn.wo_b) layers = self.model.layers for i, layer in enumerate(layers): if isinstance(layer, DeepseekV4DecoderLayer): @@ -4831,6 +4826,20 @@ def precompile_kernels_after_loading(self) -> None: build_mega_moe_shared_weights(module.shared_experts) ) + # The fused wo_b keeps a second, re-laid-out copy of each wo_b weight, + # and building it on the first fused call would allocate it after the + # memory profiler had already sized the KV cache around its absence. + # Here rather than in `post_load_weights`, which runs *before* + # `process_weights_after_loading` and so before the mxfp8 layout this + # reads even exists. + # + # Walked as modules, not as `self.model.layers`: the DSPARK draft + # wrapper borrows this method and has no `.model`. + for module in self.modules(): + wo_b = getattr(module, "wo_b", None) + if wo_b is not None: + prepare_wo_b_weight(wo_b) + @staticmethod def remap_weight_name_to_dpsk_hf_format( name: str, From bde0b40a9c84033a12b71b06ed3cc5f0ec9d59a3 Mon Sep 17 00:00:00 2001 From: xiangch Date: Mon, 12 Oct 2026 04:03:46 +0000 Subject: [PATCH 05/20] [AMD][gfx950] Address the second review round on the fused wo_b **The fused path served a stale wo_b after a weight update.** `mori_weight` returned `_mori_b` whenever it existed, and nothing invalidated it: `end_weight_update` calls `post_load_weights` only when asked, and `precompile_kernels_after_loading` -- where the copy is built -- not at all. So after `update_weights_from_*`, fused prefills multiplied by the weight the server started with while decode and every declined call used the new one. No error, just wrong output, and certain in an RL flow that loads dummy weights and then syncs real ones. The cache is now keyed on the source weight's identity and version counters, so a replaced *or* an in-place-written weight invalidates it without depending on which hook re-ran. **A rank failing around `self_test()` still hung its peers.** The CPU agreement sat after `_FusedWoB` was fully built, and `self_test` runs the scatter/reduce/gather phases, which end in cross-rank barriers on the GPU. A rank that failed between `register_window` and launching its phases -- a lost SDMA queue in `create_dev_comm`, the FlyDSL compile raising -- left the peers that did launch spinning in `synchronize()`, never reaching the agreement. `self_test` moves out of `__init__`, and construction now agrees twice: once that every rank holds an op, then once on the test's result. Rank 0 failing to make a unique id also broadcasts a None payload rather than going on to the agreement while its peers sit in `broadcast_object_list` on the same group. **The window was still allocated during a forward.** It is the larger of the two allocations (560-700 MiB per rank against ~840 MiB of weight copies across 40 layers), and since it is sized from the prefill ceiling rather than the first request there is nothing left to wait for. `prepare_fused_wo_b` now builds both at load time, so an OOM lands at startup and the collective construction runs outside any forward. The per-call path keeps a fallback for a caller that never prepared, and says so once. **Preparation was over-broad and never released.** Walking every module with a `wo_b` also prepared DSPARK's, which calls its linear directly, and the draft wrapper borrows this method -- copies nothing reads, ahead of the profiler. It now takes the `MQALayer` layers explicitly. `_can_ever_fuse` additionally declines a shape `supports()` rejects and an M floor the deployment's prefill ceiling cannot reach, and a collective disable releases every copy. An instance that declares a large ceiling but never prefills still pays for the window; narrowing that needs a signal this layer does not have. **`prepare_wo_b_weight`'s docstring contradicted itself** -- "best effort" with no `try`, next to a claim that it turns a serving OOM into a startup one. It fails at startup, and now says so. Dropped the `hasattr` that `mxfp8_ready` already covers. Co-Authored-By: Claude Opus 5 (1M context) --- python/sglang/srt/layers/mori_gemm_ar.py | 272 +++++++++++++++++++---- python/sglang/srt/models/deepseek_v4.py | 27 +-- 2 files changed, 237 insertions(+), 62 deletions(-) diff --git a/python/sglang/srt/layers/mori_gemm_ar.py b/python/sglang/srt/layers/mori_gemm_ar.py index 06b98ac8fec9..d07554a1cc39 100644 --- a/python/sglang/srt/layers/mori_gemm_ar.py +++ b/python/sglang/srt/layers/mori_gemm_ar.py @@ -85,16 +85,41 @@ def mxfp8_ready(layer) -> bool: return mxfp8_route(layer) is not None +def _weight_identity(layer): + """The source this layer's cached copy was built from. + + ``id`` catches a replaced tensor and ``_version`` an in-place write, which + is what the update paths do. Keying on a hook having re-run would not work: + ``end_weight_update`` calls ``post_load_weights`` only when asked, and + ``precompile_kernels_after_loading`` not at all. + """ + w, s = layer.weight, layer.weight_scale_mx_e8m0 + return (id(w), w._version, id(s), s._version) + + +def release_mori_weight(layer) -> None: + """Drop this layer's re-laid-out copy, if it has one.""" + layer._mori_b = None + + def mori_weight(layer): """The layer's B operand and B scale in mori's layouts, cached on the layer. sglang and mori differ only in how a lane's two 64-wide K sub-blocks sit: sglang interleaves them inside its 32 bytes, mori keeps them as two 16-byte blocks. Cached because the permutation copies the whole weight. + + Keyed on the source weight rather than on the cache's own existence: + returning it unconditionally meant that after ``update_weights_from_*`` the + fused path kept multiplying by the weight the server started with, while + decode and every declined call used the new one -- no error, just wrong + output, and certain in an RL flow that loads dummy weights then syncs real + ones. """ + src = _weight_identity(layer) cached = getattr(layer, "_mori_b", None) - if cached is not None: - return cached + if cached is not None and cached[0] == src: + return cached[1] n, k = mxfp8_shape(layer) w = layer.weight.data.contiguous().view(torch.uint8) if mxfp8_route(layer) == "aiter": @@ -121,8 +146,8 @@ def mori_weight(layer): .to(torch.int32) .reshape(-1) ) - layer._mori_b = (b, b_scale) - return layer._mori_b + layer._mori_b = (src, (b, b_scale)) + return layer._mori_b[1] @triton.jit @@ -220,6 +245,7 @@ def quantize_packed(x: torch.Tensor): _disabled = False _warned_layout = False _warned_window = False +_warned_late = False _mori_ok: bool | None = None _switch_on: bool | None = None @@ -268,26 +294,69 @@ def _shuffled_once(): return False -def prepare_wo_b_weight(layer) -> None: - """Build this layer's fused B operand now, at weight-processing time. +def _can_ever_fuse(layer) -> bool: + """Whether any batch this deployment can schedule would take the fused path. - ``mori_weight`` re-lays the weight out, which is a full second copy of it - -- the original stays, because decode and every declined call still read - it. Doing that lazily on the first fused call allocated it *while serving*, - after the memory profiler had already sized the KV cache around its - absence: 61 layers of it, about 1.3 GiB at TP2. Preparing here puts it in - front of the profiler instead, and turns a mid-serving OOM into one at - startup. + Everything here is fixed for the life of the process, so a False means the + allocations below would never be read. The prefill ceiling is the strongest + M available at load time: `_eligible`'s floor is monotonic in M, so failing + it at the ceiling means failing it at every M. - Best effort on purpose. A failure here leaves ``_mori_b`` unset, the per - call path falls back to building it, and that path is identical on every - rank -- so this is not a place where ranks can disagree. + It does not catch every case -- a decode-only instance can still declare a + large ceiling -- so an instance that never prefills may still pay for the + window. Narrowing that needs a signal this layer does not have. """ if not fused_wo_b_available() or not _reduces_over_tp_group(layer): + return False + if not mxfp8_ready(layer): + return False + from sglang.srt.distributed import get_tp_group + + world_size = get_tp_group().world_size + n, k = mxfp8_shape(layer) + ceiling = _prefill_ceiling() + if not ceiling: + return True + return _eligible(ceiling, _padded_m(ceiling, world_size), n, k, world_size) + + +def prepare_fused_wo_b(layers) -> None: + """Build everything the fused path allocates, before the profiler runs. + + Two allocations used to appear while serving, after the memory profiler had + sized the KV cache around their absence: the per-layer re-laid-out weight + (``mori_weight`` copies the whole thing, and the original stays because + decode and every declined call still read it -- about 840 MiB at TP2 over + 40 layers) and the symmetric window (another 560-700 MiB per rank). Both + are now sized from configuration alone, so both belong here. + + ``layers`` is the ``wo_b`` of each attention module that actually calls + ``fused_wo_b``. Walking every module with a ``wo_b`` instead would also + prepare the DSPARK draft's, which calls its linear directly. + + Not best effort: a failure here is a startup failure. The alternative -- + swallowing it and letting the per-call path retry -- puts the allocation + back where it was, which is what this exists to stop. + """ + candidates = [layer for layer in layers if _can_ever_fuse(layer)] + if not candidates: return - if not mxfp8_ready(layer) or not hasattr(layer, "weight_scale_mx_e8m0"): - return - mori_weight(layer) + for layer in candidates: + mori_weight(layer) + + # Collective, and it must run on every rank or none. `_can_ever_fuse` reads + # only configuration and the weight layout, which are the same everywhere, + # so the ranks agree on whether to be here at all. + layer = candidates[0] + n, k = mxfp8_shape(layer) + from sglang.srt.distributed import get_tp_group + + world_size = get_tp_group().world_size + m_max = _window_m_max(_padded_m(_prefill_ceiling() or 1, world_size), world_size) + if not _construct_collectively(m_max, n, k): + # Disabled for the process; the copies above have no reader left. + for layer in candidates: + release_mori_weight(layer) def _fused_weight(layer): @@ -324,8 +393,19 @@ def __init__(self, m_max: int, n: int, k: int): self.world_size = tp.world_size self.n, self.k, self.m_max = n, k, m_max - payload = [bytes(Communicator.get_unique_id()) if self.rank == 0 else None] + # Rank 0's id is produced inside the try so that a failure there still + # broadcasts -- raising before the broadcast would leave the peers + # blocked in it while this rank went on to the agreement below, which + # is two different collectives on the same gloo group. + try: + payload = [bytes(Communicator.get_unique_id()) if self.rank == 0 else None] + except Exception: + payload = [None] + dist.broadcast_object_list(payload, src=tp.ranks[0], group=tp.cpu_group) + raise dist.broadcast_object_list(payload, src=tp.ranks[0], group=tp.cpu_group) + if payload[0] is None: + raise RuntimeError("mori fused wo_b: rank 0 could not make a unique id") uid = UniqueId.from_bytes(payload[0]) # The window is VMM memory outside torch's allocator, so the run has to @@ -362,12 +442,6 @@ def __init__(self, m_max: int, n: int, k: int): gather_dtype=gather_dtype, gather_transport=envs.SGLANG_OPT_FUSED_WO_B_AR_GATHER_TRANSPORT.get(), ) - # Prove the collective moves bytes before serving a token. A mori - # built with BUILD_CCO_SDMA=OFF compiles the puts out: the kernels - # still launch, the all-reduce returns mostly the local slice, and - # the model answers fluently while being wrong -- and measures - # *faster*, having moved nothing. - self.op.self_test() except Exception: # The communicator holds a VMM reservation the whole process pays # for. Half-constructing this object and leaving it open makes the @@ -385,6 +459,33 @@ def __init__(self, m_max: int, n: int, k: int): gather_dtype, ) + def warmup(self) -> None: + """Compile the kernels, without running the collective. + + ``self_test`` compiles and *then* launches the device phases. Splitting + the compile out gives the ranks somewhere to agree in between, so a + FlyDSL failure on one rank is caught while its peers are still on the + host rather than already inside a cross-rank barrier. + + ``_compiled`` is mori-private; the public surface has no compile-only + entry point, and the alternative to reaching for it is a hang. + """ + self.op._compiled(self.m_max) + + def self_test(self) -> None: + """Prove the collective moves bytes before serving a token. + + A mori built with ``BUILD_CCO_SDMA=OFF`` compiles the puts out: the + kernels still launch, the all-reduce returns mostly the local slice, + and the model answers fluently while being wrong -- and measures + *faster*, having moved nothing. + + Outside ``__init__`` because it is a *device-side* collective, and so + may only run once every rank is known to hold an op. See + ``_construct_collectively``. + """ + self.op.self_test() + def pad_rows(self, x: torch.Tensor, m_pad: int) -> torch.Tensor: return self.op.pad_rows(x, m_pad) @@ -411,9 +512,22 @@ def _construct_collectively(m_max: int, n: int, k: int) -> bool: error, a hung server. That is what a rank-local ``_disabled`` produced. Rank-local failures are realistic -- an OOM in ``Communicator.init``'s VMM - reservation, or in ``self_test``'s buffers -- so this agrees on the outcome - over the CPU group before any rank is allowed to call the op, and the ranks - that *succeeded* tear their communicator down again when a peer did not. + reservation, ``create_dev_comm`` losing an SDMA queue, or the FlyDSL + compile inside the op raising on one rank. + + **Agreed after each stage that can fail locally**, because the stage after + it is collective. Agreeing only once, after ``_FusedWoB`` was fully built, + left the hole this closes: ``self_test`` runs the scatter/reduce/gather + phases, which end in cross-rank barriers on the GPU, so a rank that failed + before launching its phases left the peers that did launch spinning in + ``synchronize()``, never reaching the CPU agreement at all. + + What this cannot cover is a failure *inside* the device phases, once the + peers are already in the barrier -- no host-side agreement around a + collective can. Both failures the review named happen before that point + (``create_dev_comm`` losing an SDMA queue, the FlyDSL compile raising), and + the stages are cut so that each lands in a window that agrees. Beyond them + a watchdog is the only recourse. Construction is attempted once per process. A failure disables the path for good rather than being retried: every retry is another broadcast, another @@ -421,40 +535,78 @@ def _construct_collectively(m_max: int, n: int, k: int) -> bool: """ global _state, _disabled - import torch.distributed as dist - from sglang.srt.distributed import get_tp_group tp = get_tp_group() + state = None err = None try: state = _FusedWoB(m_max=m_max, n=n, k=k) except Exception as exc: # noqa: BLE001 - reported below, then disabled err = exc + if not _agreed(tp, err is None): + _disable(tp, state, err, "constructed") + return False - ok = torch.tensor([0 if err is not None else 1], dtype=torch.int32) - dist.all_reduce(ok, op=dist.ReduceOp.MIN, group=tp.cpu_group) - if int(ok.item()) == 1: - _state = state - return True + # Compiling is local. Agreeing here means a FlyDSL failure on one rank is + # caught while its peers are still on the host. + try: + state.warmup() + except Exception as exc: # noqa: BLE001 - reported below, then disabled + err = exc + if not _agreed(tp, err is None): + _disable(tp, state, err, "compiled") + return False + + # Every rank holds a compiled op, so the device phases are safe to launch. + try: + state.self_test() + except Exception as exc: # noqa: BLE001 - reported below, then disabled + err = exc + if not _agreed(tp, err is None): + _disable(tp, state, err, "self-tested") + return False + + _state = state + return True + + +def _agreed(tp, ok: bool) -> bool: + """Whether *every* rank reports ok, over the CPU group.""" + import torch.distributed as dist + + flag = torch.tensor([1 if ok else 0], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=tp.cpu_group) + return bool(flag.item()) + + +def _disable(tp, state, err, stage: str) -> None: + """Turn the path off here because some rank could not get this far. + + Releases this rank's op even when it built one: holding the window's + reservation for a path nothing may call costs the process its memory for + nothing. + """ + global _disabled _disabled = True if state is not None: state.close() if err is not None: logger.warning( - "mori fused wo_b could not be constructed on rank %d; the fused " - "path is off for this process on every rank: %s", + "mori fused wo_b could not be %s on rank %d; the fused path is off " + "for this process on every rank: %s", + stage, tp.rank_in_group, err, ) else: logger.warning( - "mori fused wo_b could not be constructed on a peer rank; the " - "fused path is off for this process on every rank" + "mori fused wo_b could not be %s on a peer rank; the fused path is " + "off for this process on every rank", + stage, ) - return False def _window_m_max(m_pad: int, world_size: int) -> int: @@ -472,16 +624,34 @@ def _window_m_max(m_pad: int, world_size: int) -> int: the chunk size. Its own zero case falls through to ``max_prefill_tokens``, the same way ``disaggregation/common/conn.py`` does it. """ - from sglang.srt.runtime_context import get_schedule, max_prefill_buffer_tokens - - limit = max_prefill_buffer_tokens() or get_schedule().max_prefill_tokens + limit = _prefill_ceiling() if limit: - return max(m_pad, _padded_m(int(limit), world_size)) + return max(m_pad, _padded_m(limit, world_size)) # Nothing declares a ceiling. Sizing from this call is then the only option # left, and `_warn_once_over_window` reports it if a later M outgrows it. return m_pad +def _prefill_ceiling() -> int: + """The most tokens one forward can carry, or 0 if nothing declares it.""" + from sglang.srt.runtime_context import get_schedule, max_prefill_buffer_tokens + + return int(max_prefill_buffer_tokens() or get_schedule().max_prefill_tokens or 0) + + +def _warn_once_late_window() -> None: + """Report a window built during a forward rather than at load time.""" + global _warned_late + if _warned_late: + return + _warned_late = True + logger.warning( + "mori fused wo_b: building the window on the first fused call. It is " + "outside the memory profiler's budget there; prepare_fused_wo_b() was " + "not called for this model." + ) + + def _warn_once_over_window(m_pad: int, m_max: int) -> None: """Report the first M the window cannot hold. @@ -661,10 +831,14 @@ def fused_wo_b(layer, x: torch.Tensor) -> torch.Tensor | None: if not eligible: return None - if _state is None and not _construct_collectively( - _window_m_max(m_pad, world_size), n, k - ): - return None + # Normally `prepare_fused_wo_b` already did this at load time. The fallback + # is for a caller that never ran it; it allocates the window mid-forward, + # which is what moving the preparation to load time exists to avoid, so it + # says so rather than doing it silently. + if _state is None: + _warn_once_late_window() + if not _construct_collectively(_window_m_max(m_pad, world_size), n, k): + return None # Every reason to decline a call is a check here rather than an exception # caught below, and every one of them reads the same on all ranks -- the diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index f7127da910b0..a1a2d7e7322b 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -108,7 +108,7 @@ should_skip_post_experts_all_reduce, uses_per_rank_fused_shared_slots, ) -from sglang.srt.layers.mori_gemm_ar import fused_wo_b, prepare_wo_b_weight +from sglang.srt.layers.mori_gemm_ar import fused_wo_b, prepare_fused_wo_b from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod from sglang.srt.layers.quantization.fp8_utils import ( Mxfp8DenseGemmBackend, @@ -4826,19 +4826,20 @@ def precompile_kernels_after_loading(self) -> None: build_mega_moe_shared_weights(module.shared_experts) ) - # The fused wo_b keeps a second, re-laid-out copy of each wo_b weight, - # and building it on the first fused call would allocate it after the - # memory profiler had already sized the KV cache around its absence. - # Here rather than in `post_load_weights`, which runs *before* - # `process_weights_after_loading` and so before the mxfp8 layout this - # reads even exists. + # The fused wo_b's weight copies and symmetric window are sized from + # configuration alone, so they are built here rather than on the first + # fused call, which allocated them after the memory profiler had sized + # the KV cache around their absence. Here rather than in + # `post_load_weights`, which runs *before* `process_weights_after_loading` + # and so before the mxfp8 layout this reads even exists. # - # Walked as modules, not as `self.model.layers`: the DSPARK draft - # wrapper borrows this method and has no `.model`. - for module in self.modules(): - wo_b = getattr(module, "wo_b", None) - if wo_b is not None: - prepare_wo_b_weight(wo_b) + # Only MQALayer routes wo_b through `fused_wo_b`. DSPARK's attention + # also has a `wo_b` and calls it directly, and the draft wrapper borrows + # this method, so a walk over every module with the attribute would + # prepare copies nothing reads. + prepare_fused_wo_b( + [m.wo_b for m in self.modules() if isinstance(m, MQALayer)] + ) @staticmethod def remap_weight_name_to_dpsk_hf_format( From 8925b982b131d346b9388e57e4f9d4b2caafd08f Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 12 Oct 2026 10:49:20 +0000 Subject: [PATCH 06/20] [AMD][gfx950] Refuse weight updates the fused copy cannot survive **The version-counter key did not see the write the loaders make.** `param.data.copy_()` is what `RowParallelLinear.weight_loader` and `default_weight_loader` do, and it moves nothing observable on the parameter: not `_version` (`.data` carries its own counter), not `data_ptr`, not the storage. Measured on torch 2.9.1+rocm7.2. So keying the cache on `(id, _version)` left the stale copy being served on the aiter route, where the weight stays the same `Parameter` -- the bug the key was added to fix. The previous test passed because it wrote through a *view*, which does bump the counter, so it exercised a mutation shape the real path never takes. Nothing on the tensor can tell a fresh weight from a stale one, so the update has to be refused rather than silently served. sglang already collects that constraint: `_unsupported_derived_weight_cache_error` walks the modules for `_derived_weight_cache_error`, and all four update paths check it before writing. Building a fused copy now publishes it, and releasing clears it. It is startup-determined and rank-uniform, as that mechanism requires -- `_can_ever_fuse` reads only configuration and the weight layout. **The NextN draft built a second communicator and window.** `DeepseekV4ForCausalLMNextN` subclasses `DeepseekV4ForCausalLM` without overriding `precompile_kernels_after_loading`, and its decoder builds an `MQALayer`, so with MTP this ran again in the same process: a second `_FusedWoB`, and `_state` overwritten without closing the first, leaking 560-700 MiB per rank for the life of the server. It now reuses the existing op when the shape fits and declines when it does not, since `fused_wo_b` would reject those layers anyway. Prepared layers are tracked across models, and `_disable` releases all of them. Previously a second model's construction failure turned the path off for the first model's layers while leaving their copies allocated. Co-Authored-By: Claude Opus 5 (1M context) --- python/sglang/srt/layers/mori_gemm_ar.py | 114 ++++++++++++++++------- 1 file changed, 79 insertions(+), 35 deletions(-) diff --git a/python/sglang/srt/layers/mori_gemm_ar.py b/python/sglang/srt/layers/mori_gemm_ar.py index d07554a1cc39..8c1bd2d12b55 100644 --- a/python/sglang/srt/layers/mori_gemm_ar.py +++ b/python/sglang/srt/layers/mori_gemm_ar.py @@ -85,21 +85,41 @@ def mxfp8_ready(layer) -> bool: return mxfp8_route(layer) is not None -def _weight_identity(layer): - """The source this layer's cached copy was built from. - - ``id`` catches a replaced tensor and ``_version`` an in-place write, which - is what the update paths do. Keying on a hook having re-run would not work: - ``end_weight_update`` calls ``post_load_weights`` only when asked, and - ``precompile_kernels_after_loading`` not at all. - """ - w, s = layer.weight, layer.weight_scale_mx_e8m0 - return (id(w), w._version, id(s), s._version) +#: Why an online weight update cannot run while a layer holds a fused copy. +#: +#: The loaders write through ``param.data.copy_()``, and nothing observable on +#: the tensor moves when they do -- not ``_version`` (``.data`` carries its own +#: counter), not ``data_ptr``, not the storage. So a cache keyed on the source +#: cannot tell a fresh weight from a stale one, and the update has to be +#: refused rather than silently served from the old copy. This is the same +#: constraint ``_unsupported_derived_weight_cache_error`` already collects for +#: the compensated mHC splits, and it is startup-determined and rank-uniform +#: for the same reason they are: `_can_ever_fuse` reads only configuration and +#: the weight layout. +_UPDATE_UNSUPPORTED = ( + "Online weight updates are not supported while the fused wo_b all-reduce " + "is active: it keeps a re-laid-out copy of each wo_b weight that in-place " + "loader writes do not refresh. Restart with " + "SGLANG_OPT_FUSED_WO_B_AR=0 to use online weight updates." +) def release_mori_weight(layer) -> None: - """Drop this layer's re-laid-out copy, if it has one.""" + """Drop this layer's re-laid-out copy, and the constraint that came with it.""" layer._mori_b = None + layer._derived_weight_cache_error = None + + +def release_all_prepared() -> None: + """Drop every copy this process built, across models. + + A worker can prepare more than one model -- the target and, with MTP, the + NextN draft. Disabling the path has to release both, since a disable from + either model's construction leaves the other's copies with no reader. + """ + for layer in _prepared: + release_mori_weight(layer) + _prepared.clear() def mori_weight(layer): @@ -109,17 +129,15 @@ def mori_weight(layer): sglang interleaves them inside its 32 bytes, mori keeps them as two 16-byte blocks. Cached because the permutation copies the whole weight. - Keyed on the source weight rather than on the cache's own existence: - returning it unconditionally meant that after ``update_weights_from_*`` the - fused path kept multiplying by the weight the server started with, while - decode and every declined call used the new one -- no error, just wrong - output, and certain in an RL flow that loads dummy weights then syncs real - ones. + Building it publishes ``_derived_weight_cache_error`` on the layer, which + is how a model-owned cache tells the weight updater it cannot survive an + in-place reload. Without it the fused path kept multiplying by the weight + the server started with while decode and every declined call used the new + one -- no error, just wrong output. """ - src = _weight_identity(layer) cached = getattr(layer, "_mori_b", None) - if cached is not None and cached[0] == src: - return cached[1] + if cached is not None: + return cached n, k = mxfp8_shape(layer) w = layer.weight.data.contiguous().view(torch.uint8) if mxfp8_route(layer) == "aiter": @@ -146,8 +164,9 @@ def mori_weight(layer): .to(torch.int32) .reshape(-1) ) - layer._mori_b = (src, (b, b_scale)) - return layer._mori_b[1] + layer._mori_b = (b, b_scale) + layer._derived_weight_cache_error = _UPDATE_UNSUPPORTED + return layer._mori_b @triton.jit @@ -248,6 +267,9 @@ def quantize_packed(x: torch.Tensor): _warned_late = False _mori_ok: bool | None = None _switch_on: bool | None = None +#: Every layer this process built a copy for, across models. A worker loads +#: the target and, with MTP, the NextN draft. +_prepared: list = [] _shape_hist: dict = {} _shape_calls = 0 @@ -341,22 +363,43 @@ def prepare_fused_wo_b(layers) -> None: candidates = [layer for layer in layers if _can_ever_fuse(layer)] if not candidates: return - for layer in candidates: - mori_weight(layer) - # Collective, and it must run on every rank or none. `_can_ever_fuse` reads - # only configuration and the weight layout, which are the same everywhere, - # so the ranks agree on whether to be here at all. - layer = candidates[0] - n, k = mxfp8_shape(layer) from sglang.srt.distributed import get_tp_group world_size = get_tp_group().world_size + n, k = mxfp8_shape(candidates[0]) m_max = _window_m_max(_padded_m(_prefill_ceiling() or 1, world_size), world_size) - if not _construct_collectively(m_max, n, k): - # Disabled for the process; the copies above have no reader left. + + # Called once per model, and a worker can load more than one: with MTP the + # NextN draft is a DeepseekV4ForCausalLM subclass that does not override + # `precompile_kernels_after_loading`, and its decoder builds an MQALayer, + # so this runs again in the same process. Building a second op would + # overwrite `_state` and leak the first window -- 560-700 MiB per rank for + # the life of the server. + if _state is not None: + if (n, k) != (_state.n, _state.k) or m_max > _state.m_max: + # `fused_wo_b` declines any layer whose shape does not match + # `_state`, so preparing these would allocate copies nothing reads. + logger.info( + "mori fused wo_b: not preparing a second shape (N=%d K=%d " + "M<=%d) against the existing window (N=%d K=%d M<=%d)", + n, k, m_max, _state.n, _state.k, _state.m_max, + ) + return for layer in candidates: - release_mori_weight(layer) + mori_weight(layer) + _prepared.extend(candidates) + return + + for layer in candidates: + mori_weight(layer) + _prepared.extend(candidates) + + # Collective, and it must run on every rank or none. `_can_ever_fuse` reads + # only configuration and the weight layout, which are the same everywhere, + # so the ranks agree on whether to be here at all. + # `_disable` releases every prepared layer on the way out, across models. + _construct_collectively(m_max, n, k) def _fused_weight(layer): @@ -584,15 +627,16 @@ def _agreed(tp, ok: bool) -> bool: def _disable(tp, state, err, stage: str) -> None: """Turn the path off here because some rank could not get this far. - Releases this rank's op even when it built one: holding the window's - reservation for a path nothing may call costs the process its memory for - nothing. + Releases this rank's op even when it built one, and every weight copy any + model prepared: holding the window's reservation and ~840 MiB of copies for + a path nothing may call costs the process that memory for nothing. """ global _disabled _disabled = True if state is not None: state.close() + release_all_prepared() if err is not None: logger.warning( "mori fused wo_b could not be %s on rank %d; the fused path is off " From 745a8ba38415535dca7c0fb7b5f3ed70cf544c47 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 13 Oct 2026 07:04:29 +0000 Subject: [PATCH 07/20] Apply pre-commit formatting isort's import grouping and ruff-format's argument wrapping. No semantic change; reproduced with `pre-commit run --files ...` and re-run to confirm it converges. --- python/sglang/srt/layers/mori_gemm_ar.py | 8 +++++++- python/sglang/srt/models/deepseek_v4.py | 4 +--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/mori_gemm_ar.py b/python/sglang/srt/layers/mori_gemm_ar.py index 8c1bd2d12b55..0861d47a627d 100644 --- a/python/sglang/srt/layers/mori_gemm_ar.py +++ b/python/sglang/srt/layers/mori_gemm_ar.py @@ -383,7 +383,12 @@ def prepare_fused_wo_b(layers) -> None: logger.info( "mori fused wo_b: not preparing a second shape (N=%d K=%d " "M<=%d) against the existing window (N=%d K=%d M<=%d)", - n, k, m_max, _state.n, _state.k, _state.m_max, + n, + k, + m_max, + _state.n, + _state.k, + _state.m_max, ) return for layer in candidates: @@ -429,6 +434,7 @@ def __init__(self, m_max: int, n: int, k: int): import torch.distributed as dist from mori.cco import Communicator, UniqueId from mori.ops.gemm_ar import GemmAllReduceOp + from sglang.srt.distributed import get_tp_group tp = get_tp_group() diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index f42d92892ae2..1c7974fb1130 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -5022,9 +5022,7 @@ def precompile_kernels_after_loading(self) -> None: # also has a `wo_b` and calls it directly, and the draft wrapper borrows # this method, so a walk over every module with the attribute would # prepare copies nothing reads. - prepare_fused_wo_b( - [m.wo_b for m in self.modules() if isinstance(m, MQALayer)] - ) + prepare_fused_wo_b([m.wo_b for m in self.modules() if isinstance(m, MQALayer)]) @staticmethod def remap_weight_name_to_dpsk_hf_format( From abfb6525e20ccbb2634308ffb5b6d7c994f66354 Mon Sep 17 00:00:00 2001 From: xiangch Date: Wed, 14 Oct 2026 02:45:11 +0000 Subject: [PATCH 08/20] Add gemm_ar unit tests and register them on gfx950 CI Three files, split by what each needs to run: - test_mori_gemm_ar.py stubs mori and covers the host-side gating (_reduces_over_tp_group, _eligible) and the weight-update refusal. - test_mori_gemm_ar_contract.py binds every mori symbol the adapter calls with the arguments it passes, and checks the A-scale and B-weight layouts against preshuffle_a_scale / preshuffle_b. - test_mori_gemm_ar_collective.py spawns the TP group: a failure on one rank must disable the path on all of them, and the fused result must match the split path. A new mori-gemm-ar job runs all three where mori is guaranteed present. Unlike the other jobs there it is not label-only: a change to the adapter or to the pinned MORI_COMMIT triggers it, since a mori whose gemm_ar signatures moved is exactly what this is meant to catch. MORI_ENABLE_SDMA=1 is set for that job; without it the cco puts move no bytes and self_test reports it. Measured on MI355X: 4s for the two single-GPU files, 84s for the collective one at TP8. Each case was checked against a mutation of the behaviour it guards -- all ten go red. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/pr-test-amd-mori.yml | 66 +++++ test/registered/amd/test_mori_gemm_ar.py | 241 +++++++++++++++++ .../amd/test_mori_gemm_ar_collective.py | 248 ++++++++++++++++++ .../amd/test_mori_gemm_ar_contract.py | 204 ++++++++++++++ 4 files changed, 759 insertions(+) create mode 100644 test/registered/amd/test_mori_gemm_ar.py create mode 100644 test/registered/amd/test_mori_gemm_ar_collective.py create mode 100644 test/registered/amd/test_mori_gemm_ar_contract.py diff --git a/.github/workflows/pr-test-amd-mori.yml b/.github/workflows/pr-test-amd-mori.yml index bd2d84a701d0..5dc12f01c7a9 100644 --- a/.github/workflows/pr-test-amd-mori.yml +++ b/.github/workflows/pr-test-amd-mori.yml @@ -30,6 +30,72 @@ concurrency: cancel-in-progress: ${{ github.event_name != 'pull_request' || github.event.action != 'labeled' }} jobs: + # The other jobs here are label-only. gemm-ar also runs unasked on a change to + # the adapter or to the pinned MORI_COMMIT, which is what it guards against: + # a mori whose gemm_ar signatures moved breaks the fused path silently, and a + # PR bumping the commit would otherwise never run the test that catches it. + check-changes: + runs-on: ubuntu-latest + outputs: + gemm_ar: ${{ steps.filter.outputs.gemm_ar }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.sha }} + + - name: Detect file changes + id: filter + uses: dorny/paths-filter@v3 + if: github.event_name == 'pull_request' + with: + filters: | + gemm_ar: + - "python/sglang/srt/layers/mori_gemm_ar.py" + - "test/registered/amd/test_mori_gemm_ar*.py" + - "docker/rocm.Dockerfile" + - ".github/workflows/pr-test-amd-mori.yml" + + mori-gemm-ar: + needs: [check-changes] + if: | + always() && ( + github.event_name == 'workflow_dispatch' || + needs.check-changes.outputs.gemm_ar == 'true' || + (github.event.action == 'labeled' && github.event.label.name == 'run-mori-gemm-ar') || + (github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-mori-gemm-ar')) + ) + runs-on: linux-mi35x-gpu-8 + timeout-minutes: 60 + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.sha }} + + - name: Ensure VRAM is clear + run: bash scripts/ci/amd/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm10 + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps + + # Without MORI_ENABLE_SDMA the cco puts move no bytes, the all-reduce + # returns mostly the local slice, and self_test is what reports it. + - name: Run gemm_ar tests + run: | + bash scripts/ci/amd/amd_ci_exec.sh \ + -e MORI_ENABLE_SDMA=1 \ + -w /sglang-checkout/test \ + python3 -m pytest -q \ + registered/amd/test_mori_gemm_ar.py \ + registered/amd/test_mori_gemm_ar_contract.py \ + registered/amd/test_mori_gemm_ar_collective.py + mori-pd: if: | github.event_name == 'workflow_dispatch' || diff --git a/test/registered/amd/test_mori_gemm_ar.py b/test/registered/amd/test_mori_gemm_ar.py new file mode 100644 index 000000000000..4e846c8f1d99 --- /dev/null +++ b/test/registered/amd/test_mori_gemm_ar.py @@ -0,0 +1,241 @@ +"""Host-side gating and weight-update refusal for the fused wo_b all-reduce. + +mori is stubbed here: what is under test is sglang's half -- which layers the +fused path may substitute for, which shapes are worth fusing, and that a layer +holding a re-laid-out copy refuses an online weight update. +``test_mori_gemm_ar_contract`` covers the mori side, and +``test_mori_gemm_ar_collective`` the cross-rank part. +""" + +import sys +import types +import unittest +from contextlib import contextmanager +from unittest.mock import patch + +import torch + +from sglang.srt.environ import envs +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import CustomTestCase + +register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd-mi35x") + +import sglang.srt.layers.mori_gemm_ar as mg # noqa: E402 + +N, K = 5120, 4096 + + +@contextmanager +def mori_stub(supports=True): + gemm_ar = types.ModuleType("mori.ops.gemm_ar") + gemm_ar.supports = lambda *a, **k: supports + mods = { + "mori": types.ModuleType("mori"), + "mori.ops": types.ModuleType("mori.ops"), + "mori.ops.gemm_ar": gemm_ar, + } + with patch.dict(sys.modules, mods): + yield + + +class _Group: + def __init__(self, ranks): + self.ranks = list(ranks) + self.world_size = len(self.ranks) + self.rank_in_group = 0 + + +class _Layer: + """The attributes ``_reduces_over_tp_group`` reads off a RowParallelLinear.""" + + def __init__( + self, + ranks=range(8), + reduce_results=True, + use_decode_attn_tp=False, + use_dp_attention_reduce=False, + ): + self.tp_group = _Group(ranks) + self.reduce_results = reduce_results + self.use_decode_attn_tp = use_decode_attn_tp + self.use_dp_attention_reduce = use_dp_attention_reduce + + +@contextmanager +def tp_group(ranks=range(8)): + group = _Group(ranks) + with ( + patch("sglang.srt.distributed.get_tp_group", lambda: group), + patch( + "sglang.srt.distributed.utils.get_group_rank_size", + lambda g: (0, g.world_size), + ), + ): + yield group + + +class TestReducesOverTpGroup(CustomTestCase): + """Fusing performs the all-reduce, so it may only replace the same one. + + Each False below is a configuration where the split path reduces over some + other group, or over nothing; fusing one of them sums rows belonging to + unrelated ranks instead of raising. + """ + + def test_full_tp_group_reducing(self): + with tp_group(): + self.assertTrue(mg._reduces_over_tp_group(_Layer())) + + def test_not_reducing(self): + with tp_group(): + self.assertFalse(mg._reduces_over_tp_group(_Layer(reduce_results=False))) + + def test_tp_size_one(self): + with tp_group(ranks=[0]): + self.assertFalse(mg._reduces_over_tp_group(_Layer(ranks=[0]))) + + def test_decode_attn_tp_reduces_over_another_group(self): + with tp_group(): + self.assertFalse(mg._reduces_over_tp_group(_Layer(use_decode_attn_tp=True))) + + def test_dp_attention_reduce_reduces_over_another_group(self): + with tp_group(): + self.assertFalse( + mg._reduces_over_tp_group(_Layer(use_dp_attention_reduce=True)) + ) + + def test_attn_tp_subgroup_is_not_the_tp_group(self): + # --attn-dp-size 2 at TP8: wo_b reduces over 4 ranks, the op over 8. + with tp_group(ranks=range(8)): + self.assertFalse(mg._reduces_over_tp_group(_Layer(ranks=range(4)))) + + def test_same_width_different_ranks(self): + # Equal size would pass a size comparison; the ranks are what is compared. + with tp_group(ranks=range(4)): + self.assertFalse(mg._reduces_over_tp_group(_Layer(ranks=range(4, 8)))) + + +class TestEligible(CustomTestCase): + """Expressible is mori's answer; worth it is this module's. + + The thresholds are measured crossovers (see the module's constants), so a + rewrite that reads as equivalent -- a fill ratio in place of the bf16 floor, + one floor for both wires -- moves at least one of these. + """ + + def _eligible(self, m, m_pad, fp8_gather=False, supports=True): + env = envs.SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER + with mori_stub(supports=supports), env.override(fp8_gather): + return mg._eligible(m, m_pad, N, K, 8) + + def test_inexpressible_shape_is_declined(self): + self.assertFalse(self._eligible(16384, 16384, supports=False)) + + def test_bf16_wire_floor(self): + self.assertTrue(self._eligible(8192, 8192)) + self.assertFalse(self._eligible(4096, 4096)) + + def test_fp8_wire_has_a_lower_floor(self): + self.assertFalse(self._eligible(2048, 2048)) + self.assertTrue(self._eligible(2048, 2048, fp8_gather=True)) + # m_pad 1024 still loses: one row band per destination leaves the + # scatter nothing to overlap with. + self.assertFalse(self._eligible(1024, 1024, fp8_gather=True)) + + def test_ragged_m_must_fill_its_padding(self): + self.assertTrue(self._eligible(7000, 8192)) # 0.854 + self.assertFalse(self._eligible(6000, 8192)) # 0.732 + + +class _WoB(torch.nn.Module): + """A wo_b prepared by the gfx950 native mxfp8 route. + + Native rather than aiter so the layout conversion stays pure torch and this + file needs no mori; the aiter route is covered in the contract test. + """ + + def __init__(self, n=N, k=K): + super().__init__() + w = torch.randint(0, 255, (n * k,), dtype=torch.uint8) + self.weight = torch.nn.Parameter( + w.view(torch.float8_e4m3fn), requires_grad=False + ) + self.weight_scale_mx_e8m0 = torch.nn.Parameter( + torch.randint(120, 135, (n // 32, k // 32), dtype=torch.uint8), + requires_grad=False, + ) + self.mxfp8_native_ready = True + + +def _update_error(model): + from sglang.srt.model_executor.model_runner_components.weight_updater import ( + _unsupported_derived_weight_cache_error, + ) + + return _unsupported_derived_weight_cache_error(model) + + +class TestWeightUpdateRefusal(CustomTestCase): + """A fused copy cannot survive an in-place reload, so the layer refuses one. + + Before this, a weight update left the fused path multiplying by the weight + the server started with while decode and every declined call used the new + one -- no error, just wrong output. + """ + + def setUp(self): + mg._prepared.clear() + self.layer = _WoB() + self.model = torch.nn.Module() + self.model.wo_b = self.layer + + def tearDown(self): + mg._prepared.clear() + + def test_no_copy_no_constraint(self): + self.assertIsNone(_update_error(self.model)) + + def test_building_a_copy_publishes_the_constraint(self): + mg.mori_weight(self.layer) + self.assertEqual(_update_error(self.model), mg._UPDATE_UNSUPPORTED) + + def test_release_lifts_it(self): + mg.mori_weight(self.layer) + mg._prepared.append(self.layer) + mg.release_all_prepared() + self.assertIsNone(_update_error(self.model)) + self.assertEqual(mg._prepared, []) + + def test_in_place_load_moves_nothing_observable(self): + """Why refusal is the mechanism and a cache key on the weight is not. + + The loaders write through ``param.data.copy_()``. A test that writes to + the parameter directly bumps ``_version`` and a version-keyed cache then + passes, while the real loader leaves it alone -- so the write here is + the loader's. A torch that does start bumping something turns this red, + which is when the constraint could be lifted. + """ + from sglang.srt.model_loader.weight_utils import default_weight_loader + + w = self.layer.weight + before = ( + w._version, + w.data._version, + w.data_ptr(), + w.untyped_storage().data_ptr(), + ) + fresh = torch.randint(0, 255, tuple(w.shape), dtype=torch.uint8).view( + torch.float8_e4m3fn + ) + default_weight_loader(w, fresh) + + self.assertTrue(torch.equal(w.view(torch.uint8), fresh.view(torch.uint8))) + self.assertEqual( + before, + (w._version, w.data._version, w.data_ptr(), w.untyped_storage().data_ptr()), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/amd/test_mori_gemm_ar_collective.py b/test/registered/amd/test_mori_gemm_ar_collective.py new file mode 100644 index 000000000000..b8a041f1866b --- /dev/null +++ b/test/registered/amd/test_mori_gemm_ar_collective.py @@ -0,0 +1,248 @@ +"""Cross-rank behaviour of the fused wo_b all-reduce. + +Two things only a multi-rank run can show. First, that a construction failure on +one rank turns the path off on all of them: mori's phases end in device-side +barriers, so a rank that falls back while its peers fuse spins on the GPU with +no timeout and no error. Second, that the collective actually computes the +all-reduce it replaces. +""" + +import os +import time +import unittest + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import CustomTestCase + +register_amd_ci(est_time=180, suite="stage-c-test-large-8-gpu-amd-mi35x") + +N, K = 5120, 4096 +PORT = 29731 +#: M per rank for the end-to-end case. A multiple of world_size * _BLOCK_M at +#: every world size this runs at, so no padding is involved. +M = 8192 +#: Long enough for a cold FlyDSL compile (~10s measured) with room to spare, and +#: short enough that a rank stranded in a collective fails instead of hanging CI. +SPAWN_TIMEOUT = 300 + + +class _Group: + def __init__(self, rank, world, cpu_group): + self.ranks = list(range(world)) + self.world_size = world + self.rank_in_group = rank + self.cpu_group = cpu_group + + +def _setup(rank, world): + """A fresh rank with sglang's TP group stood in for.""" + # mori's puts are compiled out without this, and then the all-reduce returns + # mostly the local slice while every kernel still launches. + os.environ.setdefault("MORI_ENABLE_SDMA", "1") + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", rank=rank, world_size=world, init_method=f"tcp://127.0.0.1:{PORT}" + ) + cpu_group = dist.new_group(ranks=list(range(world)), backend="gloo") + + import sglang.srt.distributed as dist_mod + + group = _Group(rank, world, cpu_group) + dist_mod.get_tp_group = lambda: group + return group + + +class _StubWoB: + """``_FusedWoB``'s surface, raising at whichever stage the scenario names. + + A stub rather than the real op: the agreement has to be exercised at stages + the real op cannot be made to fail at on demand. ``closed`` records whether + a rank that did build released it. + """ + + fail_rank = -1 + fail_stage = "" + closed = [] + + def __init__(self, m_max, n, k): + self.m_max, self.n, self.k = m_max, n, k + self._maybe_fail("construct") + + def _maybe_fail(self, stage): + if dist.get_rank() == self.fail_rank and self.fail_stage == stage: + raise RuntimeError(f"stub failure during {stage}") + + def warmup(self): + self._maybe_fail("warmup") + + def self_test(self): + self._maybe_fail("self_test") + + def close(self): + type(self).closed.append(dist.get_rank()) + + +def _agreement_worker(rank, world, fail_rank, fail_stage): + from unittest.mock import patch + + _setup(rank, world) + import sglang.srt.layers.mori_gemm_ar as mg + + _StubWoB.fail_rank, _StubWoB.fail_stage, _StubWoB.closed = fail_rank, fail_stage, [] + with patch.object(mg, "_FusedWoB", _StubWoB): + built = mg._construct_collectively(16384, N, K) + + _assert_agreed(mg, rank, built, expect=fail_rank < 0) + if fail_rank >= 0: + # A rank that got far enough to hold an op must release it: the window's + # VMM reservation is a per-process cost for a path nothing will call. + held = not (rank == fail_rank and fail_stage == "construct") + assert _StubWoB.closed == ([rank] if held else []), ( + f"rank {rank} failing at {fail_stage}: closed={_StubWoB.closed}" + ) + dist.destroy_process_group() + + +def _unique_id_worker(rank, world): + """Rank 0 cannot produce a unique id, which its peers are waiting on. + + ``_FusedWoB`` broadcasts from inside the ``except`` for exactly this: raising + before the broadcast leaves every peer blocked in it while this rank goes on + to the host agreement, which is two different collectives on one gloo group. + """ + from unittest.mock import patch + + import mori.cco + + _setup(rank, world) + import sglang.srt.layers.mori_gemm_ar as mg + + def boom(): + raise RuntimeError("stub failure producing the unique id") + + ctx = patch.object(mori.cco.Communicator, "get_unique_id", boom) + with ctx if rank == 0 else _nullcontext(): + built = mg._construct_collectively(16384, N, K) + + _assert_agreed(mg, rank, built, expect=False) + dist.destroy_process_group() + + +def _numerics_worker(rank, world): + _setup(rank, world) + from mori.ops.gemm_ar import preshuffle_b + + import sglang.srt.layers.mori_gemm_ar as mg + from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import mxfp8_e4m3_quantize + + torch.manual_seed(rank) + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") / 8 + w = (torch.randn(N, K, device="cuda") / 8).to(torch.float8_e4m3fn) + w_exps = torch.randint( + 120, 135, (N // 32, K // 32), dtype=torch.uint8, device="cuda" + ) + b_scale = w_exps.t().contiguous().to(torch.int32).reshape(-1) + + state = mg._FusedWoB(m_max=M, n=N, k=K) + try: + state.warmup() + state.self_test() + q, x_scale = mg.quantize_packed(x) + out = state.run(q, x_scale, preshuffle_b(w), b_scale).clone() + finally: + state.close() + + # Same operands the op was given, dequantised exactly -- ue8m0 scales are + # powers of two -- so the only difference left is summation order and the + # bf16 output rounding. The A scales come from the canonical quantiser, + # which the contract test pins to the packed one value for value. + _, x_exps = mxfp8_e4m3_quantize(x) + ref = (q.float() * _descale(x_exps)) @ ( + w.float() * _descale(w_exps).repeat_interleave(32, 0) + ).T + ref = ref.to(torch.bfloat16).float() + dist.all_reduce(ref) + + rel = ((out.float() - ref).norm() / ref.norm()).item() + assert rel < 2e-2, f"rank {rank}: relative L2 {rel:.3e} against the split path" + dist.destroy_process_group() + + +def _assert_agreed(mg, rank, built, expect): + assert built is expect, f"rank {rank}: built={built}, expected {expect}" + assert mg._disabled is not expect, f"rank {rank}: _disabled={mg._disabled}" + assert (mg._state is not None) is expect, f"rank {rank}: _state={mg._state}" + + +def _descale(exps): + """ue8m0 ``[*, K // 32]`` -> the fp32 multiplier for each of its K columns.""" + return torch.exp2(exps.float() - 127).repeat_interleave(32, -1) + + +class _nullcontext: + def __enter__(self): + return None + + def __exit__(self, *exc): + return False + + +class TestFusedWoBCollective(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.world = min(8, torch.cuda.device_count()) + if cls.world < 2: + raise unittest.SkipTest("needs at least 2 GPUs") + try: + import mori.ops.gemm_ar # noqa: F401 + except ImportError: + raise unittest.SkipTest("this build has no mori.ops.gemm_ar") + + def _spawn(self, fn, *args): + """Run `fn` on every rank, failing rather than hanging if one strands.""" + ctx = mp.start_processes( + fn, + args=(self.world, *args), + nprocs=self.world, + join=False, + start_method="spawn", + ) + deadline = time.monotonic() + SPAWN_TIMEOUT + while not ctx.join(timeout=1): + if time.monotonic() > deadline: + for proc in ctx.processes: + proc.kill() + self.fail(f"{fn.__name__} did not finish in {SPAWN_TIMEOUT}s") + + def test_every_rank_builds_when_none_fails(self): + self._spawn(_agreement_worker, -1, "") + + def test_a_failing_rank_disables_every_rank(self): + """One rank short of an op means none of them may use one. + + The stages are separate scenarios because ``_disable`` has different + work to do at each: before the op exists there is nothing to release, + after it there is a window reservation that nothing will read. + """ + for stage in ("construct", "warmup", "self_test"): + with self.subTest(stage=stage): + self._spawn(_agreement_worker, 1, stage) + + def test_unique_id_failure_does_not_strand_the_peers(self): + self._spawn(_unique_id_worker) + + def test_fused_result_matches_the_split_path(self): + """What signature conformance cannot show: that the bytes move. + + A mori built with ``BUILD_CCO_SDMA=OFF`` runs every kernel, returns + mostly the local slice, and measures faster for it. + """ + self._spawn(_numerics_worker) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/amd/test_mori_gemm_ar_contract.py b/test/registered/amd/test_mori_gemm_ar_contract.py new file mode 100644 index 000000000000..0ba602104ce7 --- /dev/null +++ b/test/registered/amd/test_mori_gemm_ar_contract.py @@ -0,0 +1,204 @@ +"""The mori surface the fused wo_b all-reduce binds to, on one GPU. + +Every mori symbol ``mori_gemm_ar`` calls is bound here with the exact arguments +it passes, and the two operand conversions are checked against mori's own +helpers. A mori whose signatures or layouts moved turns this red at the pinned +commit instead of inside a forward on an eight-rank run. +""" + +import inspect +import unittest + +import torch + +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import CustomTestCase + +register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd-mi35x") + +import sglang.srt.layers.mori_gemm_ar as mg # noqa: E402 + +N, K, M = 5120, 4096, 512 + +try: + import mori.ops.gemm_ar # noqa: F401 + + HAS_MORI = True +except ImportError: + HAS_MORI = False + +requires_mori = unittest.skipUnless(HAS_MORI, "this build has no mori.ops.gemm_ar") + + +@requires_mori +class TestMoriSignatures(CustomTestCase): + """Bind each call site's arguments without running it. + + ``_compiled`` is mori-private and the rest are keyword-only, so a rename or + a dropped default is a TypeError raised from inside a forward -- after the + window is allocated and the peers are in a barrier. + """ + + def test_module_level_helpers(self): + from mori.ops.gemm_ar import ( + padded_m, + preshuffle_a_scale, + preshuffle_b, + supports, + ) + + inspect.signature(padded_m).bind(16384, 8, mg._BLOCK_M) + inspect.signature(supports).bind(16384, N, K, 8, quant="mxfp8") + inspect.signature(preshuffle_b).bind(torch.empty(0)) + inspect.signature(preshuffle_a_scale).bind(torch.empty(0)) + + def test_communicator(self): + from mori.cco import Communicator, UniqueId + + inspect.signature(Communicator.init).bind(8, 0, None, per_rank_vmm=1 << 30) + self.assertTrue(callable(Communicator.get_unique_id)) + self.assertTrue(callable(UniqueId.from_bytes)) + + def test_window_bytes_for(self): + from mori.ops.gemm_ar import GemmAllReduceOp + + inspect.signature(GemmAllReduceOp.window_bytes_for).bind( + 8, m_max=16384, n=N, block_m=mg._BLOCK_M, gather_dtype="bf16" + ) + + def test_op_construction(self): + from mori.ops.gemm_ar import GemmAllReduceOp + + inspect.signature(GemmAllReduceOp.__init__).bind( + None, + None, + n=N, + k=K, + m_max=16384, + quant="mxfp8", + gather_dtype="bf16", + gather_transport="lsa", + ) + + def test_op_methods(self): + from mori.ops.gemm_ar import GemmAllReduceOp + + inspect.signature(GemmAllReduceOp._compiled).bind(None, 16384) + inspect.signature(GemmAllReduceOp.self_test).bind(None) + inspect.signature(GemmAllReduceOp.pad_rows).bind(None, torch.empty(0), 16384) + inspect.signature(GemmAllReduceOp.__call__).bind( + None, torch.empty(0), torch.empty(0), torch.empty(0), torch.empty(0) + ) + + def test_quant_and_transport_values_are_accepted(self): + """The strings this module passes, checked against mori's own tables.""" + from mori.ops.gemm_ar import supports + + for gather_dtype in ("bf16", "fp8"): + for transport in ("lsa", "sdma"): + supports( + 16384, + N, + K, + 8, + quant="mxfp8", + gather_dtype=gather_dtype, + gather_transport=transport, + ) + + +@requires_mori +class TestOperandConversions(CustomTestCase): + """The two layouts this module writes itself, against mori's helpers. + + Both are silent when wrong: a misplaced scale moves whole rows by a power of + two and a mislaid weight returns an uncorrelated result, neither of which + raises anywhere. + """ + + def test_packed_a_scale_matches_preshuffle_a_scale(self): + """``quantize_packed`` writes mori's A-scale layout without a second pass. + + Also pins that its scale rule is ``fp8_grid_quant``'s, which is what + ``fused_wo_b``'s claim that re-quantising an ``Fp8GridActivation`` is + lossless rests on. + """ + from mori.ops.gemm_ar import preshuffle_a_scale + + from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import ( + mxfp8_e4m3_quantize, + ) + + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + q, packed = mg.quantize_packed(x) + ref_q, ref_exps = mxfp8_e4m3_quantize(x) + + self.assertTrue(torch.equal(q.view(torch.uint8), ref_q.view(torch.uint8))) + self.assertTrue(torch.equal(packed, preshuffle_a_scale(ref_exps))) + + def test_native_route_weight_matches_preshuffle_b(self): + """sglang's scaled-MFMA weight order and mori's differ by one permute. + + The two keep a lane's 64-wide K sub-blocks differently; ``mori_weight`` + undoes that rather than re-shuffling from the checkpoint. + """ + from mori.ops.gemm_ar import preshuffle_b + + from sglang.kernels.ops.quantization.mxfp8_native_amd_gfx95 import ( + prepare_mxfp8_native_weight, + ) + + plain = torch.randint(0, 255, (N, K), dtype=torch.uint8, device="cuda").view( + torch.float8_e4m3fn + ) + scale = torch.full((N // 32, K // 32), 1.0, device="cuda") + shuffled, ue8m0 = prepare_mxfp8_native_weight(plain, scale, (32, 32)) + + layer = _layer(shuffled.view(torch.float8_e4m3fn), ue8m0, "native") + b, _ = mg.mori_weight(layer) + self.assertTrue( + torch.equal(b.view(torch.uint8), preshuffle_b(plain).view(torch.uint8)) + ) + + def test_aiter_route_weight_is_preshuffle_b(self): + from mori.ops.gemm_ar import preshuffle_b + + plain = torch.randint(0, 255, (N, K), dtype=torch.uint8, device="cuda").view( + torch.float8_e4m3fn + ) + ue8m0 = torch.randint( + 120, 135, (N // 32, K // 32), dtype=torch.uint8, device="cuda" + ) + layer = _layer(plain, ue8m0, "aiter") + b, _ = mg.mori_weight(layer) + self.assertTrue( + torch.equal(b.view(torch.uint8), preshuffle_b(plain).view(torch.uint8)) + ) + + def test_b_scale_is_k_major(self): + """The op reads the B scale transposed and flat; the shapes agree.""" + plain = torch.randint(0, 255, (N, K), dtype=torch.uint8, device="cuda").view( + torch.float8_e4m3fn + ) + ue8m0 = torch.randint( + 120, 135, (N // 32, K // 32), dtype=torch.uint8, device="cuda" + ) + layer = _layer(plain, ue8m0, "aiter") + _, b_scale = mg.mori_weight(layer) + self.assertEqual(b_scale.dtype, torch.int32) + self.assertTrue( + torch.equal(b_scale, ue8m0.t().contiguous().to(torch.int32).reshape(-1)) + ) + + +def _layer(weight, ue8m0, route): + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(weight, requires_grad=False) + layer.weight_scale_mx_e8m0 = torch.nn.Parameter(ue8m0, requires_grad=False) + layer.mxfp8_native_ready = route == "native" + layer.mxfp8_aiter_ready = route == "aiter" + return layer + + +if __name__ == "__main__": + unittest.main() From 3d213588307b6f10386b34f910f3d16bed6ba076 Mon Sep 17 00:00:00 2001 From: xiangch Date: Wed, 14 Oct 2026 11:36:14 +0000 Subject: [PATCH 09/20] Read the TP group through the runtime context main's accessor ratchet (test_runtime_context.py) forbids SRT code outside srt/distributed from calling get_tp_group(); this file was the only offender in the repo, at six call sites. Width-only reads become get_parallel().tp_size. The three that need the group handle itself -- for .ranks and .cpu_group -- become get_parallel().tp_group, which is what parallel_state.get_tp_group() returns anyway. The tests stubbed the old accessor by patching the module attribute; they now use get_parallel().override() instead, which is the supported seam and does not depend on where the call site imports from. Co-Authored-By: Claude Opus 5 (1M context) --- python/sglang/srt/layers/mori_gemm_ar.py | 23 ++++++------------- test/registered/amd/test_mori_gemm_ar.py | 3 ++- .../amd/test_mori_gemm_ar_collective.py | 6 +++-- 3 files changed, 13 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/layers/mori_gemm_ar.py b/python/sglang/srt/layers/mori_gemm_ar.py index 0861d47a627d..c0d8dc2836f5 100644 --- a/python/sglang/srt/layers/mori_gemm_ar.py +++ b/python/sglang/srt/layers/mori_gemm_ar.py @@ -40,6 +40,7 @@ # reliably reach it. from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import fp8_grid_quant from sglang.srt.environ import envs +from sglang.srt.runtime_context import get_parallel logger = logging.getLogger(__name__) @@ -332,9 +333,7 @@ def _can_ever_fuse(layer) -> bool: return False if not mxfp8_ready(layer): return False - from sglang.srt.distributed import get_tp_group - - world_size = get_tp_group().world_size + world_size = get_parallel().tp_size n, k = mxfp8_shape(layer) ceiling = _prefill_ceiling() if not ceiling: @@ -364,9 +363,7 @@ def prepare_fused_wo_b(layers) -> None: if not candidates: return - from sglang.srt.distributed import get_tp_group - - world_size = get_tp_group().world_size + world_size = get_parallel().tp_size n, k = mxfp8_shape(candidates[0]) m_max = _window_m_max(_padded_m(_prefill_ceiling() or 1, world_size), world_size) @@ -435,9 +432,7 @@ def __init__(self, m_max: int, n: int, k: int): from mori.cco import Communicator, UniqueId from mori.ops.gemm_ar import GemmAllReduceOp - from sglang.srt.distributed import get_tp_group - - tp = get_tp_group() + tp = get_parallel().tp_group self.rank = tp.rank_in_group self.world_size = tp.world_size self.n, self.k, self.m_max = n, k, m_max @@ -584,9 +579,7 @@ def _construct_collectively(m_max: int, n: int, k: int) -> bool: """ global _state, _disabled - from sglang.srt.distributed import get_tp_group - - tp = get_tp_group() + tp = get_parallel().tp_group state = None err = None @@ -764,7 +757,6 @@ def _reduces_over_tp_group(layer) -> bool: ``skip_all_reduce`` is the caller's own argument, and it passes False whenever it reaches here. """ - from sglang.srt.distributed import get_tp_group from sglang.srt.distributed.utils import get_group_rank_size _, tp_size = get_group_rank_size(layer.tp_group) @@ -775,7 +767,7 @@ def _reduces_over_tp_group(layer) -> bool: return False # Equal size would do while attn_tp is a subgroup of tp, but comparing the # ranks says what is meant and does not rest on that staying true. - return tuple(layer.tp_group.ranks) == tuple(get_tp_group().ranks) + return tuple(layer.tp_group.ranks) == tuple(get_parallel().tp_group.ranks) def _substitutes_layer_all_reduce(layer) -> bool: @@ -844,7 +836,6 @@ def fused_wo_b(layer, x: torch.Tensor) -> torch.Tensor | None: return None from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import Fp8GridActivation - from sglang.srt.distributed import get_tp_group if not _substitutes_layer_all_reduce(layer): return None @@ -868,7 +859,7 @@ def fused_wo_b(layer, x: torch.Tensor) -> torch.Tensor | None: n, w_k = mxfp8_shape(layer) if w_k != k: return None - world_size = get_tp_group().world_size + world_size = get_parallel().tp_size # Computed once and threaded through: it was recomputed in `_eligible`, in # `_record_shape` and again below, and it is a call into mori. m_pad = _padded_m(m, world_size) diff --git a/test/registered/amd/test_mori_gemm_ar.py b/test/registered/amd/test_mori_gemm_ar.py index 4e846c8f1d99..fc82927029bf 100644 --- a/test/registered/amd/test_mori_gemm_ar.py +++ b/test/registered/amd/test_mori_gemm_ar.py @@ -16,6 +16,7 @@ import torch from sglang.srt.environ import envs +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_amd_ci from sglang.test.test_utils import CustomTestCase @@ -66,7 +67,7 @@ def __init__( def tp_group(ranks=range(8)): group = _Group(ranks) with ( - patch("sglang.srt.distributed.get_tp_group", lambda: group), + get_parallel().override(tp_group=group, tp_size=group.world_size), patch( "sglang.srt.distributed.utils.get_group_rank_size", lambda g: (0, g.world_size), diff --git a/test/registered/amd/test_mori_gemm_ar_collective.py b/test/registered/amd/test_mori_gemm_ar_collective.py index b8a041f1866b..9e4d0d14f59d 100644 --- a/test/registered/amd/test_mori_gemm_ar_collective.py +++ b/test/registered/amd/test_mori_gemm_ar_collective.py @@ -49,10 +49,12 @@ def _setup(rank, world): ) cpu_group = dist.new_group(ranks=list(range(world)), backend="gloo") - import sglang.srt.distributed as dist_mod + from sglang.srt.runtime_context import get_parallel group = _Group(rank, world, cpu_group) - dist_mod.get_tp_group = lambda: group + # Permanent rather than scoped: the workers read it from several call + # depths and there is nothing to restore in a process that exits after. + get_parallel().override_permanently(tp_group=group, tp_size=world) return group From c7b0d8899081407de2e86d62f778095b4a37c878 Mon Sep 17 00:00:00 2001 From: xiangch Date: Fri, 16 Oct 2026 03:26:15 +0000 Subject: [PATCH 10/20] Bound the fused wo_b self-test so a stalled SDMA queue reports mori's self-test ends in synchronize() on phases that end in cross-rank barriers. A stack whose SDMA queues accept puts without ever completing them spins there with no timeout and no error: the server logs nothing, never opens its port, and the fused path never falls back. Run it on a side thread and give up at a deadline, so the rank still reaches the CPU-group agreement and every rank turns the path off together. Carry the device and stream across to the thread -- both are thread-local in torch, and a fresh thread gets device 0, which ran every rank's phases against the wrong GPU. A timeout does not make the process healthy: the stream stays wedged. What it buys is a named failure instead of a silent hang, so _disable keeps the window in that case rather than freeing memory a live kernel addresses. Co-Authored-By: Claude Opus 5 (1M context) --- python/sglang/srt/environ.py | 6 ++ python/sglang/srt/layers/mori_gemm_ar.py | 81 +++++++++++++++++++++--- 2 files changed, 77 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 0eee421bd151..f1ce451f6eb9 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1675,6 +1675,12 @@ class Envs: # LSA pulls, so with an fp8 gather it widens on the way in rather than in a # second kernel. Only read when the fused path is on. SGLANG_OPT_FUSED_WO_B_AR_GATHER_TRANSPORT = EnvStr("lsa") + # Seconds the fused wo_b's self-test may spend in its device phases. They end + # in cross-rank barriers with no timeout of their own, so a stack whose SDMA + # queues exist but never complete wedges every rank there; past this the path + # is turned off instead. A healthy self-test is well under a second -- the + # compile already happened in warmup -- so this is a loose upper bound. + SGLANG_OPT_FUSED_WO_B_AR_SELF_TEST_TIMEOUT_S = EnvFloat(60.0) # Debug: log the distribution of M the fused wo_b is offered. M is the token # count of one forward, not of one request, so this is the only way to see # what the scheduler actually batches. diff --git a/python/sglang/srt/layers/mori_gemm_ar.py b/python/sglang/srt/layers/mori_gemm_ar.py index c0d8dc2836f5..57529d8b7496 100644 --- a/python/sglang/srt/layers/mori_gemm_ar.py +++ b/python/sglang/srt/layers/mori_gemm_ar.py @@ -10,7 +10,10 @@ A mori built with ``BUILD_CCO_SDMA=OFF`` compiles the puts out, and then the all-reduce returns mostly the local slice while the model still answers -- and measures *faster*, having moved no bytes. ``GemmAllReduceOp.self_test()`` is -what catches that; the flag has been ON by default since ROCm/mori#710. +what catches that; the flag has been ON by default since ROCm/mori#710. That +check is run under a deadline here, because a stack whose SDMA queues accept +puts without ever completing them wedges it in a cross-rank barrier instead -- +see ``_FusedWoB.self_test``. Needs ``MORI_ENABLE_SDMA=1`` at process start and a layer one of the two gfx950 mxfp8 routes prepared -- see ``mxfp8_route``. @@ -527,8 +530,51 @@ def self_test(self) -> None: Outside ``__init__`` because it is a *device-side* collective, and so may only run once every rank is known to hold an op. See ``_construct_collectively``. + + Bounded, because mori's own check ends in ``synchronize()`` on phases + that end in cross-rank barriers: a stack whose SDMA queues exist but + whose copies never land spins there with no timeout and no error, and + the server never opens its port. Running it on a side thread leaves + this one free to give up -- ``synchronize()`` releases the GIL. + + **A timeout does not make the process healthy.** The device stream + stays wedged and the thread stays stuck; what this buys is a named + failure and a fused path that every rank agrees to turn off, instead of + a silent hang. The thread is left to leak on purpose: it holds the + stream this would otherwise free. """ - self.op.self_test() + import threading + + timeout = envs.SGLANG_OPT_FUSED_WO_B_AR_SELF_TEST_TIMEOUT_S.get() + failure: list[BaseException] = [] + # Both are thread-local in torch, and a fresh thread gets device 0 and + # its default stream. Carry this rank's over, or every rank but 0 runs + # its phases against the wrong GPU. + device = torch.cuda.current_device() + stream = torch.cuda.current_stream() + + def run(): + try: + torch.cuda.set_device(device) + with torch.cuda.stream(stream): + self.op.self_test() + except BaseException as exc: # noqa: BLE001 - re-raised below + failure.append(exc) + + thread = threading.Thread(target=run, daemon=True, name="mori-wo-b-selftest") + thread.start() + thread.join(timeout) + if thread.is_alive(): + raise TimeoutError( + f"mori fused wo_b self-test did not finish in {timeout:g}s. Its " + "device phases end in cross-rank barriers, so this is what an " + "SDMA queue that accepts puts but never completes them looks " + "like. Check MORI_ENABLE_SDMA=1, that mori was built with " + "BUILD_CCO_SDMA=ON, and the driver's SDMA support. This GPU " + "stream is wedged whatever happens next." + ) + if failure: + raise failure[0] def pad_rows(self, x: torch.Tensor, m_pad: int) -> torch.Tensor: return self.op.pad_rows(x, m_pad) @@ -566,12 +612,14 @@ def _construct_collectively(m_max: int, n: int, k: int) -> bool: before launching its phases left the peers that did launch spinning in ``synchronize()``, never reaching the CPU agreement at all. - What this cannot cover is a failure *inside* the device phases, once the - peers are already in the barrier -- no host-side agreement around a + What this cannot *recover* from is a failure inside the device phases, once + the peers are already in the barrier -- no host-side agreement around a collective can. Both failures the review named happen before that point (``create_dev_comm`` losing an SDMA queue, the FlyDSL compile raising), and - the stages are cut so that each lands in a window that agrees. Beyond them - a watchdog is the only recourse. + the stages are cut so that each lands in a window that agrees. The one that + does not, an SDMA queue that accepts puts and never completes them, is at + least reported rather than hung: ``self_test`` is bounded, and a rank that + gives up there still reaches this agreement over the CPU group. Construction is attempted once per process. A failure disables the path for good rather than being retried: every retry is another broadcast, another @@ -602,12 +650,17 @@ def _construct_collectively(m_max: int, n: int, k: int) -> bool: return False # Every rank holds a compiled op, so the device phases are safe to launch. + # A timeout here means they are still running, so the window stays. + wedged = False try: state.self_test() + except TimeoutError as exc: + err = exc + wedged = True except Exception as exc: # noqa: BLE001 - reported below, then disabled err = exc if not _agreed(tp, err is None): - _disable(tp, state, err, "self-tested") + _disable(tp, state, err, "self-tested", release=not wedged) return False _state = state @@ -623,21 +676,29 @@ def _agreed(tp, ok: bool) -> bool: return bool(flag.item()) -def _disable(tp, state, err, stage: str) -> None: +def _disable(tp, state, err, stage: str, release: bool = True) -> None: """Turn the path off here because some rank could not get this far. Releases this rank's op even when it built one, and every weight copy any model prepared: holding the window's reservation and ~840 MiB of copies for a path nothing may call costs the process that memory for nothing. + + ``release=False`` leaks the window on purpose, for the one caller whose + phases may still be running: freeing memory a live kernel addresses faults + the process. The prepared copies go either way -- the stranded phases read + the window and the dev-comm, never a weight. """ global _disabled _disabled = True - if state is not None: + if state is not None and release: state.close() release_all_prepared() if err is not None: - logger.warning( + # Not holding the window means its phases never ended, which outranks a + # fallback: nothing else on this stream will run either. + log = logger.warning if release else logger.error + log( "mori fused wo_b could not be %s on rank %d; the fused path is off " "for this process on every rank: %s", stage, From 26209c1ffaeaf555d02d3bc9ef3e1c3b4009afa8 Mon Sep 17 00:00:00 2001 From: xiangch Date: Fri, 16 Oct 2026 03:26:27 +0000 Subject: [PATCH 11/20] Gate mori-gemm-ar CI on mori's own SDMA put test The gemm_ar tests reach the collective only after mori's puts are known to move bytes. When they do not, the collective test is what times out, which reports the stack's problem as this PR's. mori's put test is bounded, needs no collective, and names the cause, so run it first. The job stays on rocm10: it is green there today, and the one report of ROCm 10 SDMA failing did not reproduce on another MI355X host with the same image and mori commit. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/pr-test-amd-mori.yml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/.github/workflows/pr-test-amd-mori.yml b/.github/workflows/pr-test-amd-mori.yml index 5dc12f01c7a9..7d586f168697 100644 --- a/.github/workflows/pr-test-amd-mori.yml +++ b/.github/workflows/pr-test-amd-mori.yml @@ -84,6 +84,17 @@ jobs: - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps + # mori's own put test, bounded and with no collective behind it. A stack + # whose SDMA queues exist but whose copies never land fails here in + # seconds, rather than as a timeout in the collective test below. + - name: Gate on mori's SDMA put + timeout-minutes: 5 + run: | + bash scripts/ci/amd/amd_ci_exec.sh \ + -e MORI_ENABLE_SDMA=1 \ + -w /sgl-workspace/mori \ + python3 -m pytest -q tests/python/cco/test_sdma_api.py -k alltoall + # Without MORI_ENABLE_SDMA the cco puts move no bytes, the all-reduce # returns mostly the local slice, and self_test is what reports it. - name: Run gemm_ar tests From aa08fc088939f46d1a47dc99e2db55c91bf183de Mon Sep 17 00:00:00 2001 From: xiangch Date: Fri, 16 Oct 2026 03:26:27 +0000 Subject: [PATCH 12/20] Cover the fp8 all-gather wire in the collective test The numerics case only ever ran the bf16 wire -- FP8_GATHER defaults off -- so its 2e-2 bound was guarding a leg that measures 2.4e-3, and the fp8 leg had no bound at all. Reading that 2e-2 as the fp8 wire's budget is what a reviewer did, and nothing in the file said otherwise. Run both wires with a bound each: 2e-2 for bf16, 3e-2 for fp8, which carries e4m3's mantissa width (2.3e-2 measured on an MI355X TP4 microbench, 2.4e-2 end to end). Co-Authored-By: Claude Opus 5 (1M context) --- .../amd/test_mori_gemm_ar_collective.py | 41 +++++++++++++------ 1 file changed, 29 insertions(+), 12 deletions(-) diff --git a/test/registered/amd/test_mori_gemm_ar_collective.py b/test/registered/amd/test_mori_gemm_ar_collective.py index 9e4d0d14f59d..54aa4ab6a3e9 100644 --- a/test/registered/amd/test_mori_gemm_ar_collective.py +++ b/test/registered/amd/test_mori_gemm_ar_collective.py @@ -18,7 +18,7 @@ from sglang.test.ci.ci_register import register_amd_ci from sglang.test.test_utils import CustomTestCase -register_amd_ci(est_time=180, suite="stage-c-test-large-8-gpu-amd-mi35x") +register_amd_ci(est_time=240, suite="stage-c-test-large-8-gpu-amd-mi35x") N, K = 5120, 4096 PORT = 29731 @@ -28,6 +28,11 @@ #: Long enough for a cold FlyDSL compile (~10s measured) with room to spare, and #: short enough that a rank stranded in a collective fails instead of hanging CI. SPAWN_TIMEOUT = 300 +#: Relative L2 each all-gather wire may reach, keyed by whether it is fp8. +#: The bf16 leg is exact up to summation order (2.4e-3 measured); the fp8 leg +#: carries e4m3's mantissa width, which is a fixed cost rather than a margin -- +#: 2.3e-2 on an MI355X TP4 microbench, 2.4e-2 end to end in review. +WIRE_REL_L2_BOUND = {False: 2e-2, True: 3e-2} class _Group: @@ -134,12 +139,13 @@ def boom(): dist.destroy_process_group() -def _numerics_worker(rank, world): +def _numerics_worker(rank, world, fp8_gather): _setup(rank, world) from mori.ops.gemm_ar import preshuffle_b import sglang.srt.layers.mori_gemm_ar as mg from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import mxfp8_e4m3_quantize + from sglang.srt.environ import envs torch.manual_seed(rank) x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") / 8 @@ -149,14 +155,16 @@ def _numerics_worker(rank, world): ) b_scale = w_exps.t().contiguous().to(torch.int32).reshape(-1) - state = mg._FusedWoB(m_max=M, n=N, k=K) - try: - state.warmup() - state.self_test() - q, x_scale = mg.quantize_packed(x) - out = state.run(q, x_scale, preshuffle_b(w), b_scale).clone() - finally: - state.close() + # Read once, in _FusedWoB.__init__, so the override has to wrap the build. + with envs.SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER.override(fp8_gather): + state = mg._FusedWoB(m_max=M, n=N, k=K) + try: + state.warmup() + state.self_test() + q, x_scale = mg.quantize_packed(x) + out = state.run(q, x_scale, preshuffle_b(w), b_scale).clone() + finally: + state.close() # Same operands the op was given, dequantised exactly -- ue8m0 scales are # powers of two -- so the only difference left is summation order and the @@ -170,7 +178,10 @@ def _numerics_worker(rank, world): dist.all_reduce(ref) rel = ((out.float() - ref).norm() / ref.norm()).item() - assert rel < 2e-2, f"rank {rank}: relative L2 {rel:.3e} against the split path" + bound = WIRE_REL_L2_BOUND[fp8_gather] + assert rel < bound, ( + f"rank {rank}: relative L2 {rel:.3e} against the split path, bound {bound:g}" + ) dist.destroy_process_group() @@ -242,8 +253,14 @@ def test_fused_result_matches_the_split_path(self): A mori built with ``BUILD_CCO_SDMA=OFF`` runs every kernel, returns mostly the local slice, and measures faster for it. + + Both wires, because they fail differently: the bf16 leg loses bytes + silently, while the fp8 leg is lossy by construction and so needs a + bound of its own before a regression can be told from e4m3. """ - self._spawn(_numerics_worker) + for fp8_gather in (False, True): + with self.subTest(wire="fp8" if fp8_gather else "bf16"): + self._spawn(_numerics_worker, fp8_gather) if __name__ == "__main__": From ea7d1c1b514538172c1460c454066b7f1686a3fe Mon Sep 17 00:00:00 2001 From: xiangch Date: Fri, 16 Oct 2026 03:55:48 +0000 Subject: [PATCH 13/20] Say what the fused wo_b does to numerics, per route Turning the path on substitutes the GEMM, so the output moves -- and the env var said nothing about that, while the fp8 gather's comment right below it called out its accuracy cost, which reads as though this one has none. It is exact only where wo_a hands wo_b an Fp8GridActivation, which needs the gfx950 native mxfp8 route. The MI355X recipes pin --fp8-gemm-backend aiter, where wo_b gets a plain bf16 tensor and the fused and unfused paths each quantise it: relative L2 7.5e-3 to 8.2e-3 at M=16384 on TP4, with --fp8-gemm-backend auto landing in the same range. The collective itself is deterministic there -- the op against itself is 0. Co-Authored-By: Claude Opus 5 (1M context) --- python/sglang/srt/environ.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index f1ce451f6eb9..39e5c8f8656f 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1666,6 +1666,11 @@ class Envs: # ROCm gfx950 TP-only: overlap wo_b's all-reduce with its GEMM through mori # cco. Off by default: it reserves a symmetric window outside torch's # allocator, so a run that enables it wants a lower --mem-fraction-static. + # It also substitutes the GEMM, so the output moves. Only where wo_a hands + # over an Fp8GridActivation -- the gfx950 native mxfp8 route -- do the two + # quantise identically and agree bit for bit. On --fp8-gemm-backend aiter, + # which the MI355X recipes pin, wo_b gets a plain bf16 tensor and each path + # rounds it itself: relative L2 7.5e-3 to 8.2e-3, measured at M=16384 TP4. SGLANG_OPT_FUSED_WO_B_AR = EnvBool(False) # Send the fused wo_b's all-gather leg as fp8 e4m3 with a per-row scale, # halving its bytes. Off by default: it costs accuracy, and that is e4m3's From bfc9acf238f9467c960623a277466b3dd9b4f814 Mon Sep 17 00:00:00 2001 From: xiangch Date: Mon, 26 Oct 2026 11:42:22 +0000 Subject: [PATCH 14/20] Guard fused wo_b against host-backed VMM windows Check kernel configuration before collective construction and validate settled per-GPU VRAM growth after every rank has mapped the window. Agree on preflight and placement failures before launching device collectives. Wait out deferred driver frees after weight preparation to avoid false declines. Cover configuration and counter handling, rank-local failures, real communicator cleanup, and the complete initialization path in numerical tests. --- python/sglang/srt/layers/mori_gemm_ar.py | 174 ++++++++++++++- test/registered/amd/test_mori_gemm_ar.py | 210 +++++++++++++++++- .../amd/test_mori_gemm_ar_collective.py | 143 +++++++++++- 3 files changed, 515 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/layers/mori_gemm_ar.py b/python/sglang/srt/layers/mori_gemm_ar.py index 57529d8b7496..6ccf12f663fe 100644 --- a/python/sglang/srt/layers/mori_gemm_ar.py +++ b/python/sglang/srt/layers/mori_gemm_ar.py @@ -18,6 +18,12 @@ Needs ``MORI_ENABLE_SDMA=1`` at process start and a layer one of the two gfx950 mxfp8 routes prepared -- see ``mxfp8_route``. +Peer-shared VMM memory must stay in VRAM. Some kernels lack +``CONFIG_DMABUF_MOVE_NOTIFY`` and amdgpu moves the window to host memory on +peer import: results stay correct but the fused path becomes much slower. +Check the kernel configuration before construction and the GPU's VRAM usage +after every rank has mapped the window, before compiling or running it. + The fused epilogue pushes a whole ``BLOCK_M`` row band to one destination, so M must be a multiple of ``tp_size * 256``. That granule is mxfp8's, not a choice -- see ``_BLOCK_M``. Ragged M is zero-padded up to it, which costs GEMM rows and @@ -423,6 +429,97 @@ def _fused_weight(layer): return mori_weight(layer) +_VMM_KERNEL_FIX = ( + "Use a kernel built with CONFIG_DMABUF_MOVE_NOTIFY=y, e.g. " + "linux-generic-hwe-22.04 (6.8) on Ubuntu 22.04; see Issue 1 in " + "https://github.com/ROCm/mori/blob/main/.claude/skills/known-issues/SKILL.md." +) + + +def _kernel_config_error() -> RuntimeError | None: + """Reject a running kernel whose config rules out a VRAM-backed window. + + Containers often expose neither source. None then means unknown, and the + placement check still runs. The build option, rather than the release + number, controls amdgpu's dma-buf pinning behavior. + """ + 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 config: + enabled = any( + line.strip() == "CONFIG_DMABUF_MOVE_NOTIFY=y" for line in config + ) + except (OSError, EOFError, UnicodeError): + continue + if enabled: + return None + return RuntimeError( + "the running kernel was built without CONFIG_DMABUF_MOVE_NOTIFY " + f"(read {path}), so amdgpu pins mori's peer-shared VMM window in " + "host memory. Results stay correct but every access crosses PCIe " + f"and the fused path can be several times slower. {_VMM_KERNEL_FIX}" + ) + return None + + +def _vram_used_bytes() -> int | None: + """Read this rank's physical GPU counter, including memory outside torch. + + Logical device IDs may be remapped by HIP_VISIBLE_DEVICES. Resolve the + PCI address through HIP's device properties instead of assuming cardN or + rank N names the same GPU. Missing properties/counters mean unknown. + """ + try: + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + bdf = ( + f"{props.pci_domain_id:04x}:" + f"{props.pci_bus_id:02x}:{props.pci_device_id:02x}.0" + ) + with open(f"/sys/bus/pci/devices/{bdf}/mem_info_vram_used") as counter: + used = int(counter.read()) + return used if used >= 0 else None + except (AttributeError, OSError, RuntimeError, TypeError, ValueError): + return None + + +def _settled_vram_used_bytes() -> int | None: + """Wait for deferred frees to stop changing the driver's VRAM counter. + + Synchronizing torch does not make amdgpu's accounting settle immediately. + During model loading, several GiB of old allocations can disappear while + the much smaller window is being built, hiding its growth. ROCm can retire + these allocations in batches hundreds of milliseconds apart. Require + 0.75 seconds of stable readings; a counter still changing after three + seconds is unknown, rather than evidence of a misplaced window. + """ + import time + + last = _vram_used_bytes() + if last is None: + return None + stable_since = time.monotonic() + deadline = stable_since + 3.0 + while time.monotonic() < deadline: + time.sleep(0.05) + current = _vram_used_bytes() + if current is None: + return None + now = time.monotonic() + if current != last: + last = current + stable_since = now + elif now - stable_since >= 0.75: + return current + return None + + class _FusedWoB: """The TP communicator and mori's op, one per rank. @@ -468,6 +565,7 @@ def __init__(self, m_max: int, n: int, k: int): block_m=_BLOCK_M, gather_dtype=gather_dtype, ) + self._vram_before = _settled_vram_used_bytes() self._comm_ctx = Communicator.init( self.world_size, self.rank, @@ -497,7 +595,8 @@ def __init__(self, m_max: int, n: int, k: int): self._comm_ctx.__exit__(None, None, None) raise logger.info( - "mori fused wo_b: window %.0f MiB, tp=%d, M<=%d, N=%d, K=%d, gather=%s", + "mori fused wo_b: allocated window %.0f MiB, tp=%d, " + "M<=%d, N=%d, K=%d, gather=%s", self.op.window_bytes / 2**20, self.world_size, m_max, @@ -506,6 +605,49 @@ def __init__(self, m_max: int, n: int, k: int): gather_dtype, ) + def check_placement(self) -> None: + """Check VRAM growth after *every* rank has mapped the window. + + The pinned mori backs and peer-maps the window in its constructor. + Its peers' imports can evict this rank's allocation, so the caller + must finish the construction agreement before taking this reading. + No device collective has launched yet; a failure can safely release + the window through the subsequent placement agreement. + + Half the window tolerates ordinary counter noise. This is a heuristic: + other processes allocating or releasing VRAM can affect the delta. + """ + torch.cuda.synchronize() + after = _settled_vram_used_bytes() + before = self._vram_before + if before is None or after is None: + logger.warning( + "mori fused wo_b: could not obtain stable VRAM counters; " + "window placement is unverified (before=%s, after=%s). " + "Continuing with the opt-in fused path.", + before, + after, + ) + return + window_bytes = self.op.window_bytes + grew = after - before + if grew < window_bytes // 2: + raise RuntimeError( + f"VRAM grew {grew / 2**20:.0f} MiB for a " + f"{window_bytes / 2**20:.0f} MiB window " + f"(before={before}, after={after} bytes); mori's peer-shared " + "VMM window may be in host memory, making the fused path " + "slower. Concurrent GPU allocations or frees can also affect " + f"this measurement. {_VMM_KERNEL_FIX}" + ) + logger.info( + "mori fused wo_b: VRAM placement check passed " + "(before=%d, after=%d, window=%d bytes)", + before, + after, + window_bytes, + ) + def warmup(self) -> None: """Compile the kernels, without running the collective. @@ -629,6 +771,24 @@ def _construct_collectively(m_max: int, n: int, k: int) -> bool: tp = get_parallel().tp_group + # This must precede the constructor's UID broadcast on every rank. A + # rank-local return from an availability gate would strand its peers. + try: + err = _kernel_config_error() + if err is None: + # Weight preparation leaves temporary allocations in torch's pool. + # Release them before measuring the window, and let the settled + # counter reader wait for the driver's deferred frees afterwards. + # These operations can fail locally, so they belong before this + # agreement, not inside the collective constructor. + torch.cuda.synchronize() + torch.cuda.empty_cache() + except Exception as exc: # noqa: BLE001 - agree before any rank constructs + err = exc + if not _agreed(tp, err is None): + _disable(tp, None, err, "preflight-checked") + return False + state = None err = None try: @@ -639,6 +799,17 @@ def _construct_collectively(m_max: int, n: int, k: int) -> bool: _disable(tp, state, err, "constructed") return False + # Every peer has now finished importing this rank's window. Checking in + # __init__ would rely on mori's internal synchronization to avoid an early + # reading, before a peer import can evict the allocation from VRAM. + try: + state.check_placement() + except Exception as exc: # noqa: BLE001 - safe to release before device phases + err = exc + if not _agreed(tp, err is None): + _disable(tp, state, err, "placement-checked") + return False + # Compiling is local. Agreeing here means a FlyDSL failure on one rank is # caught while its peers are still on the host. try: @@ -664,6 +835,7 @@ def _construct_collectively(m_max: int, n: int, k: int) -> bool: return False _state = state + logger.info("mori fused wo_b: enabled on all %d TP ranks", tp.world_size) return True diff --git a/test/registered/amd/test_mori_gemm_ar.py b/test/registered/amd/test_mori_gemm_ar.py index fc82927029bf..332725dbebe5 100644 --- a/test/registered/amd/test_mori_gemm_ar.py +++ b/test/registered/amd/test_mori_gemm_ar.py @@ -11,7 +11,7 @@ import types import unittest from contextlib import contextmanager -from unittest.mock import patch +from unittest.mock import mock_open, patch import torch @@ -149,6 +149,214 @@ def test_ragged_m_must_fill_its_padding(self): self.assertFalse(self._eligible(6000, 8192)) # 0.732 +class TestKernelConfig(CustomTestCase): + """The running kernel's options decide, including a rebuilt 5.15 kernel.""" + + def test_move_notify_enabled_on_an_old_kernel(self): + config = mock_open(read_data="CONFIG_DMABUF_MOVE_NOTIFY=y\n") + with ( + patch("platform.release", return_value="5.15.0-custom"), + patch("builtins.open", config), + patch("gzip.open") as compressed, + ): + self.assertIsNone(mg._kernel_config_error()) + self.assertEqual(config.call_args.args[0], "/boot/config-5.15.0-custom") + compressed.assert_not_called() + + def test_disabled_move_notify_reports_the_fix(self): + config = mock_open(read_data="# CONFIG_DMABUF_MOVE_NOTIFY is not set\n") + with patch("builtins.open", config): + error = mg._kernel_config_error() + self.assertIsInstance(error, RuntimeError) + self.assertIn("CONFIG_DMABUF_MOVE_NOTIFY", str(error)) + self.assertIn("linux-generic-hwe-22.04", str(error)) + + def test_proc_config_is_used_when_boot_is_unreadable(self): + compressed = mock_open(read_data="CONFIG_DMABUF_MOVE_NOTIFY=y\n") + with ( + patch("builtins.open", side_effect=FileNotFoundError), + patch("gzip.open", compressed), + ): + self.assertIsNone(mg._kernel_config_error()) + self.assertEqual(compressed.call_args.args[0], "/proc/config.gz") + + def test_disabled_proc_config_is_rejected(self): + with ( + patch("builtins.open", side_effect=PermissionError), + patch("gzip.open", mock_open(read_data="CONFIG_OTHER=y\n")), + ): + self.assertIsInstance(mg._kernel_config_error(), RuntimeError) + + def test_unreadable_config_leaves_the_placement_check_to_decide(self): + with ( + patch("builtins.open", side_effect=FileNotFoundError), + patch("gzip.open", side_effect=FileNotFoundError), + ): + self.assertIsNone(mg._kernel_config_error()) + + +class TestVramCounter(CustomTestCase): + def test_current_logical_device_is_resolved_to_its_pci_address(self): + # Device 2 need not be PCI device 2, especially with HIP_VISIBLE_DEVICES. + props = types.SimpleNamespace(pci_domain_id=1, pci_bus_id=0xAF, pci_device_id=3) + counter = mock_open(read_data="123456\n") + with ( + patch.object(torch.cuda, "current_device", return_value=2), + patch.object( + torch.cuda, "get_device_properties", return_value=props + ) as get, + patch("builtins.open", counter), + ): + self.assertEqual(mg._vram_used_bytes(), 123456) + get.assert_called_once_with(2) + self.assertEqual( + counter.call_args.args[0], + "/sys/bus/pci/devices/0001:af:03.0/mem_info_vram_used", + ) + + def test_missing_pci_properties_make_the_probe_unavailable(self): + with ( + patch.object(torch.cuda, "current_device", return_value=0), + patch.object( + torch.cuda, + "get_device_properties", + return_value=types.SimpleNamespace(), + ), + ): + self.assertIsNone(mg._vram_used_bytes()) + + def test_unreadable_or_invalid_counter_is_unknown(self): + props = types.SimpleNamespace(pci_domain_id=0, pci_bus_id=1, pci_device_id=0) + for contents in ("not a counter\n", "-1\n", None): + counter = mock_open(read_data=contents or "") + if contents is None: + counter.side_effect = FileNotFoundError + with ( + self.subTest(contents=contents), + patch.object(torch.cuda, "current_device", return_value=0), + patch.object(torch.cuda, "get_device_properties", return_value=props), + patch("builtins.open", counter), + ): + self.assertIsNone(mg._vram_used_bytes()) + + +class TestSettledVramCounter(CustomTestCase): + @contextmanager + def _clock(self): + """Advance probe time without delaying the test.""" + elapsed = [0.0] + + def sleep(seconds): + elapsed[0] = round(elapsed[0] + seconds, 8) + + with ( + patch("time.monotonic", side_effect=lambda: elapsed[0]), + patch("time.sleep", side_effect=sleep), + ): + yield elapsed + + def test_deferred_frees_finish_before_the_baseline_is_accepted(self): + # amdgpu can account for old allocations in batches, after torch's + # allocator already reports them released. Two equal reads, or even + # 250 ms of stability, would accept one of these stale baselines. + def vram_used(): + if elapsed[0] < 0.25: + return 20 << 30 + if elapsed[0] < 0.50: + return 12 << 30 + return 4 << 30 + + with ( + self._clock() as elapsed, + patch.object(mg, "_vram_used_bytes", side_effect=vram_used), + ): + self.assertEqual(mg._settled_vram_used_bytes(), 4 << 30) + self.assertGreaterEqual(elapsed[0], 1.25) + + def test_two_equal_readings_do_not_hide_a_later_large_free(self): + values = iter([20 << 30, 20 << 30, 2 << 30]) + with ( + self._clock() as elapsed, + patch.object( + mg, "_vram_used_bytes", side_effect=lambda: next(values, 2 << 30) + ), + ): + self.assertEqual(mg._settled_vram_used_bytes(), 2 << 30) + self.assertGreaterEqual(elapsed[0], 0.85) + + def test_a_continuously_changing_counter_times_out_as_unknown(self): + with ( + self._clock() as elapsed, + patch.object( + mg, "_vram_used_bytes", side_effect=lambda: int(elapsed[0] * 1_000_000) + ), + ): + self.assertIsNone(mg._settled_vram_used_bytes()) + self.assertGreaterEqual(elapsed[0], 3.0) + self.assertLessEqual(elapsed[0], 3.05) + + def test_unreadable_counter_is_unknown_even_after_a_valid_reading(self): + for values in ([None], [2 << 30, None]): + with ( + self.subTest(values=values), + self._clock() as elapsed, + patch.object(mg, "_vram_used_bytes", side_effect=values), + ): + self.assertIsNone(mg._settled_vram_used_bytes()) + self.assertLessEqual(elapsed[0], 0.05) + + +class TestWindowPlacement(CustomTestCase): + WINDOW = 64 << 20 + + def _state(self, before): + state = mg._FusedWoB.__new__(mg._FusedWoB) + state.rank = 0 + state._vram_before = before + state.op = types.SimpleNamespace(window_bytes=self.WINDOW) + return state + + def test_half_a_window_or_more_in_vram_is_accepted(self): + before = 4 << 30 + for growth in (self.WINDOW // 2, self.WINDOW): + with ( + self.subTest(growth=growth), + patch.object(torch.cuda, "synchronize"), + patch.object( + mg, "_settled_vram_used_bytes", return_value=before + growth + ), + ): + self._state(before).check_placement() + + def test_insufficient_vram_growth_reports_without_closing_the_op(self): + before = 4 << 30 + for growth in (self.WINDOW // 2 - 1, 0, -self.WINDOW): + state = self._state(before) + with ( + self.subTest(growth=growth), + patch.object(torch.cuda, "synchronize"), + patch.object( + mg, "_settled_vram_used_bytes", return_value=before + growth + ), + self.assertRaises(RuntimeError) as raised, + ): + state.check_placement() + self.assertIn("VRAM", str(raised.exception)) + self.assertIn("64", str(raised.exception)) + # Cleanup belongs to the collective caller, after every rank agrees. + self.assertIsNotNone(state.op) + + def test_unknown_counter_warns_and_keeps_the_opt_in_path(self): + for before, after in ((None, 100), (100, None), (None, None)): + with ( + self.subTest(before=before, after=after), + patch.object(torch.cuda, "synchronize"), + patch.object(mg, "_settled_vram_used_bytes", return_value=after), + self.assertLogs(mg.logger, level="WARNING"), + ): + self._state(before).check_placement() + + class _WoB(torch.nn.Module): """A wo_b prepared by the gfx950 native mxfp8 route. diff --git a/test/registered/amd/test_mori_gemm_ar_collective.py b/test/registered/amd/test_mori_gemm_ar_collective.py index 54aa4ab6a3e9..7eabad612ffb 100644 --- a/test/registered/amd/test_mori_gemm_ar_collective.py +++ b/test/registered/amd/test_mori_gemm_ar_collective.py @@ -74,15 +74,20 @@ class _StubWoB: fail_rank = -1 fail_stage = "" closed = [] + stages = [] def __init__(self, m_max, n, k): self.m_max, self.n, self.k = m_max, n, k self._maybe_fail("construct") def _maybe_fail(self, stage): + type(self).stages.append(stage) if dist.get_rank() == self.fail_rank and self.fail_stage == stage: raise RuntimeError(f"stub failure during {stage}") + def check_placement(self): + self._maybe_fail("placement") + def warmup(self): self._maybe_fail("warmup") @@ -99,8 +104,13 @@ def _agreement_worker(rank, world, fail_rank, fail_stage): _setup(rank, world) import sglang.srt.layers.mori_gemm_ar as mg - _StubWoB.fail_rank, _StubWoB.fail_stage, _StubWoB.closed = fail_rank, fail_stage, [] - with patch.object(mg, "_FusedWoB", _StubWoB): + _StubWoB.fail_rank, _StubWoB.fail_stage = fail_rank, fail_stage + _StubWoB.closed, _StubWoB.stages = [], [] + layer = _prepared_layer(mg) + with ( + patch.object(mg, "_kernel_config_error", return_value=None), + patch.object(mg, "_FusedWoB", _StubWoB), + ): built = mg._construct_collectively(16384, N, K) _assert_agreed(mg, rank, built, expect=fail_rank < 0) @@ -111,6 +121,106 @@ def _agreement_worker(rank, world, fail_rank, fail_stage): assert _StubWoB.closed == ([rank] if held else []), ( f"rank {rank} failing at {fail_stage}: closed={_StubWoB.closed}" ) + if fail_stage == "placement": + assert _StubWoB.stages == ["construct", "placement"] + _assert_released(mg, layer) + dist.destroy_process_group() + + +def _prepared_layer(mg): + """A copy that a collective disable must release, including its constraint.""" + layer = torch.nn.Module() + layer._mori_b = (torch.ones(16, device="cuda"),) + layer._derived_weight_cache_error = mg._UPDATE_UNSUPPORTED + mg._prepared.append(layer) + return layer + + +def _assert_released(mg, layer): + from sglang.srt.model_executor.model_runner_components.weight_updater import ( + _unsupported_derived_weight_cache_error, + ) + + assert layer._mori_b is None + assert mg._prepared == [] + assert _unsupported_derived_weight_cache_error(layer) is None + assert not mg.fused_wo_b_available() + + +def _host_check_worker(rank, world, fail_stage="kernel_config"): + """Preflight failure stops all ranks before the constructor's broadcast.""" + from unittest.mock import patch + + _setup(rank, world) + import sglang.srt.layers.mori_gemm_ar as mg + + layer = _prepared_layer(mg) + error = ( + RuntimeError("missing CONFIG_DMABUF_MOVE_NOTIFY") + if rank == 1 and fail_stage == "kernel_config" + else None + ) + original_empty_cache = torch.cuda.empty_cache + + def empty_cache(): + if rank == 1 and fail_stage == "empty_cache": + raise RuntimeError("injected cache flush failure") + original_empty_cache() + + with ( + patch.object(mg, "_kernel_config_error", return_value=error), + patch.object(torch.cuda, "empty_cache", side_effect=empty_cache), + patch.object(mg, "_FusedWoB") as construct, + ): + built = mg._construct_collectively(16384, N, K) + _assert_agreed(mg, rank, built, expect=False) + construct.assert_not_called() + _assert_released(mg, layer) + dist.destroy_process_group() + + +def _placement_worker(rank, world): + """Keep the real communicator; only one rank's VRAM readings are injected.""" + from unittest.mock import patch + + _setup(rank, world) + import sglang.srt.layers.mori_gemm_ar as mg + + layer = _prepared_layer(mg) + before = 4 << 30 + # The healthy rank reports ample VRAM growth; rank 1 reports none. A real + # constructor, placement check and close run on both ranks, before any + # fused kernel is allowed to launch. + after = before if rank == 1 else before + (4 << 30) + closed = [] + original_close = mg._FusedWoB.close + + def close(state): + original_close(state) + closed.append(state) + + with ( + patch.object(mg, "_kernel_config_error", return_value=None), + patch.object(mg, "_settled_vram_used_bytes", side_effect=[before, after]), + patch.object(mg._FusedWoB, "close", close), + patch.object(mg._FusedWoB, "warmup") as warmup, + patch.object(mg._FusedWoB, "self_test") as self_test, + ): + built = mg._construct_collectively(M, N, K) + + _assert_agreed(mg, rank, built, expect=False) + warmup.assert_not_called() + self_test.assert_not_called() + assert len(closed) == 1, f"rank {rank}: communicator closed {len(closed)} times" + assert closed[0].op is None + assert not closed[0].comm.is_valid + _assert_released(mg, layer) + + # A disabled flag alone does not establish that the fallback can run: its + # CUDA stream and the ordinary TP collective must still be usable. + value = torch.tensor([rank + 1], dtype=torch.float32, device="cuda") + dist.all_reduce(value) + assert value.item() == world * (world + 1) / 2 dist.destroy_process_group() @@ -132,7 +242,10 @@ def boom(): raise RuntimeError("stub failure producing the unique id") ctx = patch.object(mori.cco.Communicator, "get_unique_id", boom) - with ctx if rank == 0 else _nullcontext(): + with ( + patch.object(mg, "_kernel_config_error", return_value=None), + ctx if rank == 0 else _nullcontext(), + ): built = mg._construct_collectively(16384, N, K) _assert_agreed(mg, rank, built, expect=False) @@ -157,10 +270,10 @@ def _numerics_worker(rank, world, fp8_gather): # Read once, in _FusedWoB.__init__, so the override has to wrap the build. with envs.SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER.override(fp8_gather): - state = mg._FusedWoB(m_max=M, n=N, k=K) + built = mg._construct_collectively(M, N, K) + _assert_agreed(mg, rank, built, expect=True) + state = mg._state try: - state.warmup() - state.self_test() q, x_scale = mg.quantize_packed(x) out = state.run(q, x_scale, preshuffle_b(w), b_scale).clone() finally: @@ -215,12 +328,13 @@ def setUpClass(cls): except ImportError: raise unittest.SkipTest("this build has no mori.ops.gemm_ar") - def _spawn(self, fn, *args): + def _spawn(self, fn, *args, world=None): """Run `fn` on every rank, failing rather than hanging if one strands.""" + world = self.world if world is None else world ctx = mp.start_processes( fn, - args=(self.world, *args), - nprocs=self.world, + args=(world, *args), + nprocs=world, join=False, start_method="spawn", ) @@ -241,13 +355,22 @@ def test_a_failing_rank_disables_every_rank(self): work to do at each: before the op exists there is nothing to release, after it there is a window reservation that nothing will read. """ - for stage in ("construct", "warmup", "self_test"): + for stage in ("construct", "placement", "warmup", "self_test"): with self.subTest(stage=stage): self._spawn(_agreement_worker, 1, stage) def test_unique_id_failure_does_not_strand_the_peers(self): self._spawn(_unique_id_worker) + def test_host_check_failure_prevents_construction_on_every_rank(self): + self._spawn(_host_check_worker, world=2) + + def test_cache_flush_failure_prevents_construction_on_every_rank(self): + self._spawn(_host_check_worker, "empty_cache", world=2) + + def test_placement_failure_releases_real_windows_and_allows_fallback(self): + self._spawn(_placement_worker, world=2) + def test_fused_result_matches_the_split_path(self): """What signature conformance cannot show: that the bytes move. From 8c88d3b033b0abc7e634c8b1ec76e16e16611796 Mon Sep 17 00:00:00 2001 From: xiangch Date: Mon, 26 Oct 2026 13:11:47 +0000 Subject: [PATCH 15/20] Handle preshuffled AITER weights in fused wo_b Reuse AITER's (16, 16) MXFP8 weight layout when the loader marks it as preshuffled; its bytes already match MORI's B layout. Preserve conversion for plain and native weights, including the derived scale cache. Update the native helper import after the kernels directory refactor and cover both AITER layouts, including the production weight preparation path. --- python/sglang/srt/layers/mori_gemm_ar.py | 34 ++++++----- .../amd/test_mori_gemm_ar_contract.py | 61 +++++++++++++++++-- 2 files changed, 75 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/layers/mori_gemm_ar.py b/python/sglang/srt/layers/mori_gemm_ar.py index 6ccf12f663fe..9907feaf751d 100644 --- a/python/sglang/srt/layers/mori_gemm_ar.py +++ b/python/sglang/srt/layers/mori_gemm_ar.py @@ -76,10 +76,11 @@ def mxfp8_route(layer) -> str | None: """Which gfx950 mxfp8 route prepared this layer, or None if neither did. Both leave the ``[N/32, K/32]`` exponent bytes on - ``layer.weight_scale_mx_e8m0`` and differ only in the weight: ``native`` - rebinds it to the shuffled ``[N/16, K/128, 2048]``, ``aiter`` leaves plain - ``[N, K]``. Reading one as the other returns an uncorrelated result rather - than failing, so `mori_weight` branches on this instead of guessing. + ``layer.weight_scale_mx_e8m0``. ``native`` rebinds the weight to its shuffled + ``[N/16, K/128, 2048]``. ``aiter`` keeps the ``[N, K]`` shape, with either + checkpoint order or the (16, 16) shuffle recorded by + ``mxfp8_aiter_preshuffled``. Reading one layout as another returns an + uncorrelated result rather than failing, so `mori_weight` uses these markers. """ if not hasattr(layer, "weight_scale_mx_e8m0"): return None @@ -108,8 +109,8 @@ def mxfp8_ready(layer) -> bool: #: the weight layout. _UPDATE_UNSUPPORTED = ( "Online weight updates are not supported while the fused wo_b all-reduce " - "is active: it keeps a re-laid-out copy of each wo_b weight that in-place " - "loader writes do not refresh. Restart with " + "is active: it caches derived wo_b operands that in-place loader writes " + "do not fully refresh. Restart with " "SGLANG_OPT_FUSED_WO_B_AR=0 to use online weight updates." ) @@ -135,9 +136,10 @@ def release_all_prepared() -> None: def mori_weight(layer): """The layer's B operand and B scale in mori's layouts, cached on the layer. - sglang and mori differ only in how a lane's two 64-wide K sub-blocks sit: - sglang interleaves them inside its 32 bytes, mori keeps them as two 16-byte - blocks. Cached because the permutation copies the whole weight. + The native route interleaves a lane's two 64-wide K sub-blocks inside its + 32 bytes; mori keeps them as two 16-byte blocks. AITER's optional (16, 16) + shuffle already has mori's byte order and can be reused. Cached because + the other layouts require a weight copy and the scales are converted too. Building it publishes ``_derived_weight_cache_error`` on the layer, which is how a model-owned cache tells the weight updater it cannot survive an @@ -151,12 +153,14 @@ def mori_weight(layer): n, k = mxfp8_shape(layer) w = layer.weight.data.contiguous().view(torch.uint8) if mxfp8_route(layer) == "aiter": - # That route leaves the weight in its checkpoint [N, K] order, which is - # what preshuffle_b is defined on, so mori emits its own layout from the - # source rather than permuting sglang's. - from mori.ops.gemm_ar import preshuffle_b - - b = preshuffle_b(w.reshape(n, k).view(layer.weight.dtype)) + b = w.reshape(n, k).view(layer.weight.dtype) + # AITER's (16, 16) FP8 shuffle and mori's preshuffle_b both flatten to + # [N/16, K/16, 16, 16]. Re-shuffling an already prepared weight would + # silently change the matrix. Older/plain AITER layers have no marker. + if not getattr(layer, "mxfp8_aiter_preshuffled", False): + from mori.ops.gemm_ar import preshuffle_b + + b = preshuffle_b(b) else: # [N/16, K/128, 64 lanes, 2 sub-blocks, 16B] -> the two sub-blocks split out b = ( diff --git a/test/registered/amd/test_mori_gemm_ar_contract.py b/test/registered/amd/test_mori_gemm_ar_contract.py index 0ba602104ce7..f9a66462fbba 100644 --- a/test/registered/amd/test_mori_gemm_ar_contract.py +++ b/test/registered/amd/test_mori_gemm_ar_contract.py @@ -8,6 +8,7 @@ import inspect import unittest +from types import SimpleNamespace import torch @@ -144,7 +145,7 @@ def test_native_route_weight_matches_preshuffle_b(self): """ from mori.ops.gemm_ar import preshuffle_b - from sglang.kernels.ops.quantization.mxfp8_native_amd_gfx95 import ( + from sglang.kernels.ops.gemm.mxfp8_native_amd_gfx95 import ( prepare_mxfp8_native_weight, ) @@ -161,6 +162,7 @@ def test_native_route_weight_matches_preshuffle_b(self): ) def test_aiter_route_weight_is_preshuffle_b(self): + from aiter.ops.shuffle import shuffle_weight from mori.ops.gemm_ar import preshuffle_b plain = torch.randint(0, 255, (N, K), dtype=torch.uint8, device="cuda").view( @@ -169,11 +171,60 @@ def test_aiter_route_weight_is_preshuffle_b(self): ue8m0 = torch.randint( 120, 135, (N // 32, K // 32), dtype=torch.uint8, device="cuda" ) - layer = _layer(plain, ue8m0, "aiter") - b, _ = mg.mori_weight(layer) - self.assertTrue( - torch.equal(b.view(torch.uint8), preshuffle_b(plain).view(torch.uint8)) + expected_b = preshuffle_b(plain).view(torch.uint8) + expected_scale = ue8m0.t().contiguous().to(torch.int32).reshape(-1) + for preshuffled in (None, False, True): + with self.subTest(preshuffled=preshuffled): + weight = shuffle_weight(plain, (16, 16)) if preshuffled else plain + layer = _layer(weight, ue8m0, "aiter") + if preshuffled is not None: + layer.mxfp8_aiter_preshuffled = preshuffled + b, b_scale = mg.mori_weight(layer) + self.assertTrue(torch.equal(b.view(torch.uint8), expected_b)) + self.assertEqual(b_scale.dtype, torch.int32) + self.assertTrue(torch.equal(b_scale, expected_scale)) + + def test_processed_aiter_route_weight_matches_preshuffle_b(self): + """Follow the production loader so changes to its layout are covered.""" + from mori.ops.gemm_ar import preshuffle_b + + from sglang.srt.environ import envs + from sglang.srt.layers.quantization import fp8_hip + from sglang.srt.layers.quantization.fp8_utils import Mxfp8DenseGemmBackend + + plain = torch.randint(0, 255, (N, K), dtype=torch.uint8, device="cuda").view( + torch.float8_e4m3fn ) + ue8m0 = torch.randint( + 120, 135, (N // 32, K // 32), dtype=torch.uint8, device="cuda" + ) + scale = torch.exp2(ue8m0.to(torch.float32) - 127) + expected_b = preshuffle_b(plain).view(torch.uint8) + expected_scale = ue8m0.t().contiguous().to(torch.int32).reshape(-1) + method = SimpleNamespace( + mxfp8_dense_backend=Mxfp8DenseGemmBackend.GFX95_MXFP8_AITER, + weight_block_size=[32, 32], + ) + for preshuffled in (False, True): + with ( + self.subTest(preshuffled=preshuffled), + envs.SGLANG_ROCM_MXFP8_AITER_PRESHUFFLE.override(preshuffled), + envs.SGLANG_OPT_HIP_FLYDSL_MXFP8.override(False), + ): + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(plain.clone(), requires_grad=False) + layer.weight_scale_inv = torch.nn.Parameter( + scale.clone(), requires_grad=False + ) + fp8_hip.process_dense_weights(method, layer, None) + + self.assertTrue(layer.mxfp8_aiter_ready) + self.assertEqual(layer.mxfp8_aiter_preshuffled, preshuffled) + self.assertTrue(torch.equal(layer.weight_scale_mx_e8m0, ue8m0)) + b, b_scale = mg.mori_weight(layer) + self.assertTrue(torch.equal(b.view(torch.uint8), expected_b)) + self.assertEqual(b_scale.dtype, torch.int32) + self.assertTrue(torch.equal(b_scale, expected_scale)) def test_b_scale_is_k_major(self): """The op reads the B scale transposed and flat; the shapes agree.""" From ba5efe3f7a3755883997376595b5acdaf1f1a8c0 Mon Sep 17 00:00:00 2001 From: xiangch Date: Thu, 5 Nov 2026 02:55:11 +0000 Subject: [PATCH 16/20] Address review of the fused wo_b all-reduce - Decline under Megatron SP: RowParallelLinear reduce-scatters along the token dim there (row_parallel_gbar_matmul), so an all-reduced [M, N] is the wrong result and shape. sp_active is checked per forward in _substitutes_layer_all_reduce; it is rank-uniform. - mxfp8_shape reads N and K off the weight in its route's layout instead of the ue8m0 scale, whose ceil(N/32) rows overstated an N off the 32 grid and made the weight reshape fail at startup. - prepare_fused_wo_b copies only the layers the window can serve, the first candidate's (N, K); fused_wo_b declines every other shape. - Only a plain RowParallelLinear without a bias is fused: a LoRA wrapper moves the attributes under base_layer (startup raised AttributeError), and the fused GEMM has no bias term. - fused_wo_b_available() also requires is_gfx95_supported(), instead of relying on only the gfx950 routes setting the mxfp8 markers. - release_mori_weight restores the layer's previous _derived_weight_cache_error instead of clearing it. - fused_wo_b's docstring states how long the returned window view lives under DeepSeek-V4's op schedule, and what a prefill TBO yield would break. - environment_variables.mdx documents SGLANG_OPT_FUSED_WO_B_AR: the memory it reserves outside torch's allocator, that online weight updates are refused, and where it stays off. - The mori CI job also runs when the AMD install script changes. Tests: SP declines, a wrapped layer is declined rather than raising, N off the 32 grid, no copies for unserved shapes, the prior constraint is restored, and the gfx95 gate. The native-route test layer now has its real [N/16, K/128, 2048] shape. Co-Authored-By: Claude --- .github/workflows/pr-test-amd-mori.yml | 1 + .../docs/references/environment_variables.mdx | 5 ++ python/sglang/srt/layers/mori_gemm_ar.py | 69 ++++++++++++--- test/registered/amd/test_mori_gemm_ar.py | 88 ++++++++++++++++++- 4 files changed, 147 insertions(+), 16 deletions(-) diff --git a/.github/workflows/pr-test-amd-mori.yml b/.github/workflows/pr-test-amd-mori.yml index 7d586f168697..d735dd26a5ec 100644 --- a/.github/workflows/pr-test-amd-mori.yml +++ b/.github/workflows/pr-test-amd-mori.yml @@ -54,6 +54,7 @@ jobs: - "python/sglang/srt/layers/mori_gemm_ar.py" - "test/registered/amd/test_mori_gemm_ar*.py" - "docker/rocm.Dockerfile" + - "scripts/ci/amd/amd_ci_install_dependency.sh" - ".github/workflows/pr-test-amd-mori.yml" mori-gemm-ar: diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx index cac7cc287361..b418e697c726 100644 --- a/docs/docs/references/environment_variables.mdx +++ b/docs/docs/references/environment_variables.mdx @@ -480,6 +480,11 @@ SGLang supports various environment variables that can be used to configure its + + SGLANG_OPT_FUSED_WO_B_AR + ROCm gfx950, TP only: fuse DeepSeek-V4.1-Flash wo_b's GEMM with its all-reduce through mori cco (needs MORI_ENABLE_SDMA=1). It reserves memory outside torch's allocator, about 560-700 MiB of symmetric window per rank plus a re-laid-out copy of every wo_b weight (about 840 MiB at TP2), so lower --mem-fraction-static when enabling it. Online weight updates (update_weights_from_tensor / _from_distributed) are refused while it is on, because the copies would go stale. Off under Megatron SP, data-parallel attention and LoRA, where wo_b does not perform this all-reduce. + false + SGLANG_MORI_DISPATCH_DTYPE Override MoRI-EP dispatch quantization type. auto uses auto-detection from weight dtype; bf16/fp8/fp4 forces the specified type for all layers diff --git a/python/sglang/srt/layers/mori_gemm_ar.py b/python/sglang/srt/layers/mori_gemm_ar.py index 9907feaf751d..099b1a7fe508 100644 --- a/python/sglang/srt/layers/mori_gemm_ar.py +++ b/python/sglang/srt/layers/mori_gemm_ar.py @@ -62,14 +62,17 @@ def mxfp8_shape(layer) -> tuple[int, int]: - """The layer's logical ``(N, K)``, read off the ue8m0 scale. + """The layer's logical ``(N, K)``, read off the weight in its route's layout. - **Not** ``layer.weight.shape``: the native route rebinds the weight to its - shuffled ``[N/16, K/128, 2048]``, and reading that first axis as N - disqualifies every layer without a word. + The native route rebinds the weight to its shuffled ``[N/16, K/128, 2048]``, + so ``weight.shape[0]`` is N only on the aiter route. The ue8m0 scale is no + substitute: it has ``ceil(N/32)`` rows, which overstates any N that is not a + multiple of 32. """ - sn, sk = layer.weight_scale_mx_e8m0.shape - return sn * 32, sk * 32 + shape = layer.weight.shape + if mxfp8_route(layer) == "native": + return shape[0] * 16, shape[1] * 128 + return shape[0], shape[1] def mxfp8_route(layer) -> str | None: @@ -118,7 +121,9 @@ def mxfp8_ready(layer) -> bool: def release_mori_weight(layer) -> None: """Drop this layer's re-laid-out copy, and the constraint that came with it.""" layer._mori_b = None - layer._derived_weight_cache_error = None + # The attribute is an optional protocol (weight_updater reads it with getattr). + layer._derived_weight_cache_error = getattr(layer, "_mori_prior_cache_error", None) + layer._mori_prior_cache_error = None def release_all_prepared() -> None: @@ -179,6 +184,8 @@ def mori_weight(layer): .reshape(-1) ) layer._mori_b = (b, b_scale) + # Another owner of the layer may already have set it; release restores it. + layer._mori_prior_cache_error = getattr(layer, "_derived_weight_cache_error", None) layer._derived_weight_cache_error = _UPDATE_UNSUPPORTED return layer._mori_b @@ -342,7 +349,9 @@ def _can_ever_fuse(layer) -> bool: large ceiling -- so an instance that never prefills may still pay for the window. Narrowing that needs a signal this layer does not have. """ - if not fused_wo_b_available() or not _reduces_over_tp_group(layer): + if not fused_wo_b_available(): + return False + if not (_is_plain_row_parallel(layer) and _reduces_over_tp_group(layer)): return False if not mxfp8_ready(layer): return False @@ -379,6 +388,9 @@ def prepare_fused_wo_b(layers) -> None: world_size = get_parallel().tp_size n, k = mxfp8_shape(candidates[0]) m_max = _window_m_max(_padded_m(_prefill_ceiling() or 1, world_size), world_size) + # The window serves one (N, K) and `fused_wo_b` declines every other shape, + # so a copy of a mismatched layer would never be read. + candidates = [layer for layer in candidates if mxfp8_shape(layer) == (n, k)] # Called once per model, and a worker can load more than one: with MTP the # NextN draft is a DeepseekV4ForCausalLM subclass that does not override @@ -970,6 +982,18 @@ def _eligible(m: int, m_pad: int, n: int, k: int, world_size: int) -> bool: return m_pad >= floor and m >= _MIN_PAD_FILL * m_pad +def _is_plain_row_parallel(layer) -> bool: + """A ``RowParallelLinear`` itself, without a bias. + + A LoRA adapter wraps ``wo_b`` in ``RowParallelLinearWithLoRA``, which moves + the attributes read below under ``base_layer`` and adds its own delta; and + the fused GEMM has no bias term. ``wo_b`` is built with ``bias=False``. + """ + from sglang.srt.layers.linear import RowParallelLinear + + return isinstance(layer, RowParallelLinear) and layer.bias is None + + def _reduces_over_tp_group(layer) -> bool: """Whether this layer's own all-reduce is the one the op would perform. @@ -1010,11 +1034,17 @@ def _reduces_over_tp_group(layer) -> bool: def _substitutes_layer_all_reduce(layer) -> bool: """``_reduces_over_tp_group`` plus the flags that vary per forward.""" from sglang.srt.layers.moe.utils import should_skip_mlp_all_reduce + from sglang.srt.runtime_context import get_forward # The mHC post folds the reduce in; the split path skips it here. if should_skip_mlp_all_reduce(): return False - return _reduces_over_tp_group(layer) + # Under Megatron SP, RowParallelLinear reduce-scatters along the token dim + # instead (row_parallel_gbar_matmul); an all-reduced [M, N] is the wrong + # result and the wrong shape. sp_active is the same on every rank. + if get_forward().sp_active: + return False + return _is_plain_row_parallel(layer) and _reduces_over_tp_group(layer) def _mori_has_gemm_ar() -> bool: @@ -1054,7 +1084,11 @@ def fused_wo_b_available() -> bool: if _disabled: return False if _switch_on is None: - _switch_on = envs.SGLANG_OPT_FUSED_WO_B_AR.get() + from sglang.srt.utils.common import is_gfx95_supported + + # Stated rather than left to the mxfp8 markers, which only the gfx950 + # routes set today: mori's kernel and its tuning are gfx950's. + _switch_on = envs.SGLANG_OPT_FUSED_WO_B_AR.get() and is_gfx95_supported() return _switch_on and _mori_has_gemm_ar() @@ -1062,10 +1096,17 @@ def fused_wo_b(layer, x: torch.Tensor) -> torch.Tensor | None: """wo_b's `GEMM + all-reduce`, fused. ``None`` means "use the normal path". ``x`` is the bf16 activation this rank holds, ``[M, K]``. The result is the - all-reduced ``[M, N]``, as a **view of the symmetric window** -- the next - layer's wo_b overwrites it, which is safe because the caller consumes it in - the residual add of the same layer. ``SGLANG_DEBUG_FUSED_WO_B_AR`` forces a copy - and cross-checks against the unfused path. + all-reduced ``[M, N]``, as a **view of the symmetric window**, and the next + ``wo_b`` call overwrites it. The caller must consume it before then. + DeepSeek-V4 does, but through an op schedule this module does not own: + ``attn_boundary.finish()`` hands it to ``ResidualStream.record()``, and the + residual is applied at ``ffn_boundary.prepare()`` in + ``op_mhc_post_attn_pre_mlp``, still in the same layer. Under prefill TBO the + two microbatches share the window, so a ``YieldOperation`` between + ``op_attn`` and ``op_mhc_post_attn_pre_mlp`` would let one overwrite the + other's result without an error. Clone the result if that schedule changes. + ``SGLANG_DEBUG_FUSED_WO_B_AR`` forces a copy and cross-checks against the + unfused path. """ global _state, _disabled diff --git a/test/registered/amd/test_mori_gemm_ar.py b/test/registered/amd/test_mori_gemm_ar.py index 332725dbebe5..bed16edcb193 100644 --- a/test/registered/amd/test_mori_gemm_ar.py +++ b/test/registered/amd/test_mori_gemm_ar.py @@ -16,7 +16,7 @@ import torch from sglang.srt.environ import envs -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_forward, get_parallel from sglang.test.ci.ci_register import register_amd_ci from sglang.test.test_utils import CustomTestCase @@ -117,6 +117,29 @@ def test_same_width_different_ranks(self): self.assertFalse(mg._reduces_over_tp_group(_Layer(ranks=range(4, 8)))) +class TestPerForwardSubstitution(CustomTestCase): + def test_sequence_parallel_reduce_scatters_instead(self): + """Under Megatron SP the layer reduce-scatters along tokens, so an + all-reduced ``[M, N]`` would be the wrong result and the wrong shape.""" + with ( + tp_group(), + patch.object(mg, "_is_plain_row_parallel", return_value=True), + patch( + "sglang.srt.layers.moe.utils.should_skip_mlp_all_reduce", + return_value=False, + ), + ): + self.assertTrue(mg._substitutes_layer_all_reduce(_Layer())) + with get_forward().scoped(sp_active=True): + self.assertFalse(mg._substitutes_layer_all_reduce(_Layer())) + + def test_a_wrapped_layer_is_declined_rather_than_raising(self): + """A LoRA adapter wraps wo_b and moves its attributes under + ``base_layer``; startup must decline it, not raise AttributeError.""" + with tp_group(), patch.object(mg, "fused_wo_b_available", return_value=True): + self.assertFalse(mg._can_ever_fuse(torch.nn.Module())) + + class TestEligible(CustomTestCase): """Expressible is mori's answer; worth it is this module's. @@ -366,7 +389,8 @@ class _WoB(torch.nn.Module): def __init__(self, n=N, k=K): super().__init__() - w = torch.randint(0, 255, (n * k,), dtype=torch.uint8) + # The native route's shuffled layout, [N/16, K/128, 2048]. + w = torch.randint(0, 255, (n // 16, k // 128, 2048), dtype=torch.uint8) self.weight = torch.nn.Parameter( w.view(torch.float8_e4m3fn), requires_grad=False ) @@ -446,5 +470,65 @@ def test_in_place_load_moves_nothing_observable(self): ) +class TestShapeAndPreparation(CustomTestCase): + def tearDown(self): + mg._prepared.clear() + + def test_n_comes_from_the_weight_not_the_scale(self): + """The ue8m0 scale has ceil(N/32) rows; an N off the 32 grid was + overstated, and the weight reshape then failed at startup.""" + n = 5120 + 16 + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter( + torch.zeros(n, K, dtype=torch.uint8).view(torch.float8_e4m3fn), + requires_grad=False, + ) + layer.weight_scale_mx_e8m0 = torch.nn.Parameter( + torch.zeros(-(-n // 32), K // 32, dtype=torch.uint8), requires_grad=False + ) + layer.mxfp8_native_ready = False + layer.mxfp8_aiter_ready = True + self.assertEqual(mg.mxfp8_shape(layer), (n, K)) + self.assertEqual(mg.mxfp8_shape(_WoB()), (N, K)) + + def test_layers_the_window_cannot_serve_get_no_copy(self): + """The window serves the first candidate's (N, K); copying a layer of + another shape allocates a full weight nothing ever reads.""" + same, other = _WoB(), _WoB(n=N // 2) + with ( + tp_group(), + patch.object(mg, "_can_ever_fuse", return_value=True), + patch.object(mg, "_construct_collectively", return_value=True), + patch.object(mg, "_prefill_ceiling", return_value=0), + patch.object(mg, "_padded_m", lambda m, world_size: m), + patch.object(mg, "_state", None), + ): + mg.prepare_fused_wo_b([same, other]) + self.assertIsNotNone(same._mori_b) + self.assertIsNone(getattr(other, "_mori_b", None)) + self.assertEqual(mg._prepared, [same]) + + def test_release_restores_another_owners_constraint(self): + layer = _WoB() + layer._derived_weight_cache_error = "set by another cache" + mg.mori_weight(layer) + self.assertEqual(layer._derived_weight_cache_error, mg._UPDATE_UNSUPPORTED) + mg.release_mori_weight(layer) + self.assertEqual(layer._derived_weight_cache_error, "set by another cache") + + +class TestPlatformGate(CustomTestCase): + def test_off_on_anything_but_gfx95(self): + """The mxfp8 markers only gate by accident of which routes set them.""" + with ( + patch.object(mg, "_switch_on", None), + patch.object(mg, "_disabled", False), + patch.object(mg, "_mori_has_gemm_ar", return_value=True), + envs.SGLANG_OPT_FUSED_WO_B_AR.override(True), + patch("sglang.srt.utils.common.is_gfx95_supported", return_value=False), + ): + self.assertFalse(mg.fused_wo_b_available()) + + if __name__ == "__main__": unittest.main() From f02b7e66845706013da0826798043ace04b4389b Mon Sep 17 00:00:00 2001 From: xiangch Date: Thu, 5 Nov 2026 03:13:24 +0000 Subject: [PATCH 17/20] Import ForwardMode from the module that defines it test_kv_translate_ownership.py imported ForwardMode through forward_batch_info, which only re-imports it; check-indirect-imports rejects that, and the lint job runs it over every file. Co-Authored-By: Claude --- .../unit/layers/attention/test_kv_translate_ownership.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/registered/unit/layers/attention/test_kv_translate_ownership.py b/test/registered/unit/layers/attention/test_kv_translate_ownership.py index 4db3e3a9d2da..cc2678e0a3f5 100644 --- a/test/registered/unit/layers/attention/test_kv_translate_ownership.py +++ b/test/registered/unit/layers/attention/test_kv_translate_ownership.py @@ -37,7 +37,7 @@ MiniMaxHybridAttnBackend, ) from sglang.srt.layers.attention.tbo_backend import TboAttnBackend -from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.model_executor.forward_mode import ForwardMode from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase From 0aeb4903209c41e67146f58f50b5060945d9f86d Mon Sep 17 00:00:00 2001 From: xiangch Date: Mon, 9 Nov 2026 02:08:27 +0000 Subject: [PATCH 18/20] Import is_gfx95_supported from sglang.srt.utils.device main (#48781) split sglang.srt.utils.common into topic modules and moved is_gfx95_supported to sglang.srt.utils.device. The fused wo_b's platform gate imported it from common, and the platform-gate test patched it there, so both now point at the module that defines it. Co-Authored-By: Claude --- python/sglang/srt/layers/mori_gemm_ar.py | 2 +- test/registered/amd/test_mori_gemm_ar.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/mori_gemm_ar.py b/python/sglang/srt/layers/mori_gemm_ar.py index 099b1a7fe508..a4074b32b31f 100644 --- a/python/sglang/srt/layers/mori_gemm_ar.py +++ b/python/sglang/srt/layers/mori_gemm_ar.py @@ -1084,7 +1084,7 @@ def fused_wo_b_available() -> bool: if _disabled: return False if _switch_on is None: - from sglang.srt.utils.common import is_gfx95_supported + from sglang.srt.utils.device import is_gfx95_supported # Stated rather than left to the mxfp8 markers, which only the gfx950 # routes set today: mori's kernel and its tuning are gfx950's. diff --git a/test/registered/amd/test_mori_gemm_ar.py b/test/registered/amd/test_mori_gemm_ar.py index bed16edcb193..cd4c35bb6b03 100644 --- a/test/registered/amd/test_mori_gemm_ar.py +++ b/test/registered/amd/test_mori_gemm_ar.py @@ -525,7 +525,7 @@ def test_off_on_anything_but_gfx95(self): patch.object(mg, "_disabled", False), patch.object(mg, "_mori_has_gemm_ar", return_value=True), envs.SGLANG_OPT_FUSED_WO_B_AR.override(True), - patch("sglang.srt.utils.common.is_gfx95_supported", return_value=False), + patch("sglang.srt.utils.device.is_gfx95_supported", return_value=False), ): self.assertFalse(mg.fused_wo_b_available()) From 14b0d86762ec61e2ea60b94db5a6d916a168ff76 Mon Sep 17 00:00:00 2001 From: xiangch Date: Mon, 9 Nov 2026 02:14:58 +0000 Subject: [PATCH 19/20] Drop the removed SGLANG_OPT_HIP_FLYDSL_MXFP8 override from the contract test main (#48399) removed SGLANG_OPT_HIP_FLYDSL_MXFP8: the aiter route now prepares FlyDSL's weight scales whenever FlyDSL is installed. That adds weight_scale_mx_flydsl and leaves weight and weight_scale_mx_e8m0, which mori_weight reads, untouched, so the test follows the loader as it is. Co-Authored-By: Claude --- test/registered/amd/test_mori_gemm_ar_contract.py | 1 - 1 file changed, 1 deletion(-) diff --git a/test/registered/amd/test_mori_gemm_ar_contract.py b/test/registered/amd/test_mori_gemm_ar_contract.py index f9a66462fbba..a984d41c1de3 100644 --- a/test/registered/amd/test_mori_gemm_ar_contract.py +++ b/test/registered/amd/test_mori_gemm_ar_contract.py @@ -209,7 +209,6 @@ def test_processed_aiter_route_weight_matches_preshuffle_b(self): with ( self.subTest(preshuffled=preshuffled), envs.SGLANG_ROCM_MXFP8_AITER_PRESHUFFLE.override(preshuffled), - envs.SGLANG_OPT_HIP_FLYDSL_MXFP8.override(False), ): layer = torch.nn.Module() layer.weight = torch.nn.Parameter(plain.clone(), requires_grad=False) From 7f027240571b06783841b6cc1902510fca11d077 Mon Sep 17 00:00:00 2001 From: xiangch Date: Tue, 10 Nov 2026 02:38:19 +0000 Subject: [PATCH 20/20] Keep the fused wo_b behind _is_gfx95_supported in deepseek_v4 Move the fused wo_b's hooks into deepseek_v4_hip, the module deepseek_v4 binds as _hip only on ROCm, and call them from deepseek_v4 under _is_gfx95_supported. Other platforms no longer import mori_gemm_ar, and _project_wo_b's plain wo_b call is unchanged from main. The fused branch returns directly: mori_gemm_ar fuses only a reduce over the whole TP group, where the attention-TP all-reduce after wo_b does not apply. Co-Authored-By: Claude --- .../deepseek_common/amd/deepseek_v4_hip.py | 23 +++++++++++ python/sglang/srt/models/deepseek_v4.py | 39 +++++++------------ 2 files changed, 36 insertions(+), 26 deletions(-) diff --git a/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_hip.py b/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_hip.py index 0f739727c090..67e1e20ab62c 100644 --- a/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_hip.py +++ b/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_hip.py @@ -11,6 +11,7 @@ from sglang.kernels.ops.layernorm.mhc_post_split_h import mhc_post_split_h from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled from sglang.srt.environ import envs +from sglang.srt.layers import mori_gemm_ar from sglang.srt.layers.attention.hip_flash_mla import ( hip_attention_needs_head_pad, resolve_hip_flashmla_backend, @@ -166,6 +167,16 @@ def attention_inv_rope( return torch.view_as_real(attn.freqs_cis).flatten(-2), positions +def fused_wo_b(attn, o) -> Optional[torch.Tensor]: + """wo_b's GEMM fused with its all-reduce through mori (SGLANG_OPT_FUSED_WO_B_AR), or + None to run the plain wo_b. Only for an all-reduce that is not deferred, since fusing + performs it; mori_gemm_ar fuses only a reduce over the whole TP group, so the + attention-TP all-reduce that follows wo_b never applies to its result.""" + # the gfx950 wo_a fork already returned [T, G * R], so only the plain tensor needs flattening + o_in = o.flatten(1) if isinstance(o, torch.Tensor) else o + return mori_gemm_ar.fused_wo_b(attn.wo_b, o_in) + + # ---- DeepseekV4DecoderLayer ---- @@ -246,3 +257,15 @@ def input_norm( if coefficients is not None: coefficients.materialize() return layer.input_layernorm(hidden_states), None + + +# ---- DeepseekV4ForCausalLM ---- + + +def prepare_fused_wo_b(wo_b_layers) -> None: + """Build the fused wo_b's weight copies and symmetric window up front: they are sized + from configuration alone, so the memory profiler then sizes the KV cache around them. + Call after process_weights_after_loading, which lays out the mxfp8 weights they copy + (post_load_weights runs before it), with the wo_b of every MQALayer only: DSPARK's + attention also has a wo_b but calls it directly.""" + mori_gemm_ar.prepare_fused_wo_b(wo_b_layers) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 69f0d25c3af5..ac5a9c8c729e 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -112,7 +112,6 @@ should_skip_post_experts_all_reduce, uses_per_rank_fused_shared_slots, ) -from sglang.srt.layers.mori_gemm_ar import fused_wo_b, prepare_fused_wo_b from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod from sglang.srt.layers.quantization.fp8_utils import ( Mxfp8DenseGemmBackend, @@ -3173,19 +3172,14 @@ def _project_wo_b(self, o, defer_all_reduce: bool = False): ) and not should_skip_mlp_all_reduce() ) - # the gfx950 wo_a fork already returned [T, G * R], so only the plain - # tensor needs flattening - o_in = o.flatten(1) if isinstance(o, torch.Tensor) else o - # Fusing *performs* the all-reduce, where a deferred one is skipped so - # the mHC post can fold it in -- a correctness constraint, not a gap. - # `defer_all_reduce` being false is not on its own a licence to fuse: - # it is also false when the layer reduces over another group or not at - # all. `fused_wo_b` checks that for itself against the layer. - fused_o = None if defer_all_reduce else fused_wo_b(self.wo_b, o_in) - if fused_o is not None: - o = fused_o - else: - o, _ = self.wo_b(o_in, skip_all_reduce=defer_all_reduce) + if _is_gfx95_supported and not defer_all_reduce: + fused_o = _hip.fused_wo_b(self, o) + if fused_o is not None: + return fused_o + o, _ = self.wo_b( + o.flatten(1) if isinstance(o, torch.Tensor) else o, + skip_all_reduce=defer_all_reduce, + ) if defer_all_reduce: return mhc.AttnOutput(o) if self.attn_tp_size > 1 and self.attn_tp_size < get_parallel().tp_size: @@ -5265,18 +5259,11 @@ def precompile_kernels_after_loading(self) -> None: build_mega_moe_shared_weights(module.shared_experts) ) - # The fused wo_b's weight copies and symmetric window are sized from - # configuration alone, so they are built here rather than on the first - # fused call, which allocated them after the memory profiler had sized - # the KV cache around their absence. Here rather than in - # `post_load_weights`, which runs *before* `process_weights_after_loading` - # and so before the mxfp8 layout this reads even exists. - # - # Only MQALayer routes wo_b through `fused_wo_b`. DSPARK's attention - # also has a `wo_b` and calls it directly, and the draft wrapper borrows - # this method, so a walk over every module with the attribute would - # prepare copies nothing reads. - prepare_fused_wo_b([m.wo_b for m in self.modules() if isinstance(m, MQALayer)]) + if _is_gfx95_supported: + # only MQALayer routes wo_b through _hip.fused_wo_b + _hip.prepare_fused_wo_b( + [m.wo_b for m in self.modules() if isinstance(m, MQALayer)] + ) @staticmethod def remap_weight_name_to_dpsk_hf_format(