diff --git a/.github/workflows/pr-test-amd-mori.yml b/.github/workflows/pr-test-amd-mori.yml index bd2d84a701d0..d735dd26a5ec 100644 --- a/.github/workflows/pr-test-amd-mori.yml +++ b/.github/workflows/pr-test-amd-mori.yml @@ -30,6 +30,84 @@ 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" + - "scripts/ci/amd/amd_ci_install_dependency.sh" + - ".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 + + # 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 + 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/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx index 6cc89f95ec32..2ada5d113875 100644 --- a/docs/docs/references/environment_variables.mdx +++ b/docs/docs/references/environment_variables.mdx @@ -485,6 +485,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/environ.py b/python/sglang/srt/environ.py index 2aa290c99839..5d4c5e1e3554 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -2051,6 +2051,35 @@ class Envs: # CUDA: fuse inverse RoPE into the UE8M0 FP8 wo_a activation quantization. # Off by default; requires the FP8 wo_a path. SGLANG_OPT_FP8_WO_A_FUSED_INVROPE = EnvBool(False) + # 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 + # 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. 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. + 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) # DSpark decode: build the verify window and the draft anchor column in one 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..a4074b32b31f --- /dev/null +++ b/python/sglang/srt/layers/mori_gemm_ar.py @@ -0,0 +1,1220 @@ +# SPDX-License-Identifier: Apache-2.0 +"""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. 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``. + +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 +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 + +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 +from sglang.srt.runtime_context import get_parallel + +logger = logging.getLogger(__name__) + +# ---- 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 weight in its route's layout. + + 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. + """ + 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: + """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``. ``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 + 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 + + +#: 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 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." +) + + +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 + # 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: + """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): + """The layer's B operand and B scale in mori's layouts, cached on the layer. + + 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 + 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. + """ + 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": + 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 = ( + 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) + # 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 + + +@triton.jit +def _mxfp8_quant_packed_kernel( + x_ptr, + xq_ptr, + s_ptr, + M, + K, + sxm, + sxk, + sqm, + sqk, + BLOCK_M: tl.constexpr, +): + """``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 + ``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. + + 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) + 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) + xq, sb = fp8_grid_quant(x) + 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_window = False +_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 +#: 61 wo_b calls make one forward, so this logs roughly every ten of them. +_SHAPE_EVERY = 610 + + +def _record_shape(m, m_pad, eligible): + global _shape_calls + key = (m, m_pad, 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() + + +#: 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 + + +def _shuffled_once(): + global _logged_config + if _logged_config: + return True + _logged_config = True + return False + + +def _can_ever_fuse(layer) -> bool: + """Whether any batch this deployment can schedule would take the fused path. + + 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. + + 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(): + return False + if not (_is_plain_row_parallel(layer) and _reduces_over_tp_group(layer)): + return False + if not mxfp8_ready(layer): + return False + world_size = get_parallel().tp_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 + + 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 + # `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: + 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): + """The op's B operand and B scale, or None if this layer cannot be fused.""" + global _warned_layout + 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) + + +_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. + + 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 + + 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 + + # 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 + # 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, + 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, + 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, + quant="mxfp8", + gather_dtype=gather_dtype, + gather_transport=envs.SGLANG_OPT_FUSED_WO_B_AR_GATHER_TRANSPORT.get(), + ) + 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: 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, + n, + k, + 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. + + ``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``. + + 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. + """ + 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) + + 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, ``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 *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. 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 + multi-hundred-MiB VMM reservation and another teardown, 61 times a forward. + """ + global _state, _disabled + + 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: + 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 + + # 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: + 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. + # 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", release=not wedged) + return False + + _state = state + logger.info("mori fused wo_b: enabled on all %d TP ranks", tp.world_size) + 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, 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 and release: + state.close() + release_all_prepared() + if err is not None: + # 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, + tp.rank_in_group, + err, + ) + else: + logger.warning( + "mori fused wo_b could not be %s on a peer rank; the fused path is " + "off for this process on every rank", + stage, + ) + + +def _window_m_max(m_pad: int, world_size: int) -> int: + """Rows the symmetric window is sized for. + + 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. + """ + limit = _prefill_ceiling() + if limit: + 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. + + 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 + second, which is not a property of the kernel. + """ + from mori.ops.gemm_ar import supports + + 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 + ) + 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. + + 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.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_parallel().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 + 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 + # 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: + """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, 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: + 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. + _switch_on = envs.SGLANG_OPT_FUSED_WO_B_AR.get() and is_gfx95_supported() + return _switch_on and _mori_has_gemm_ar() + + +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**, 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 + + if not fused_wo_b_available(): + return None + + from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import Fp8GridActivation + + 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 + # 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 + 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_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) + 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 + + # 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 + # 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(): + 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, 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() + 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/layers/quantization/fp8_hip.py b/python/sglang/srt/layers/quantization/fp8_hip.py index b5ea4ad32989..7f54f826e224 100644 --- a/python/sglang/srt/layers/quantization/fp8_hip.py +++ b/python/sglang/srt/layers/quantization/fp8_hip.py @@ -51,6 +51,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 _prepare_flydsl_mxfp8(layer) from aiter.ops.shuffle import shuffle_weight 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 01f8bf101f5f..ac5a9c8c729e 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -3172,6 +3172,10 @@ def _project_wo_b(self, o, defer_all_reduce: bool = False): ) and not should_skip_mlp_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, @@ -5255,6 +5259,12 @@ def precompile_kernels_after_loading(self) -> None: build_mega_moe_shared_weights(module.shared_experts) ) + 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( name: str, 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..cd4c35bb6b03 --- /dev/null +++ b/test/registered/amd/test_mori_gemm_ar.py @@ -0,0 +1,534 @@ +"""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 mock_open, patch + +import torch + +from sglang.srt.environ import envs +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 + +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 ( + 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), + ), + ): + 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 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. + + 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 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. + + 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__() + # 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 + ) + 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()), + ) + + +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.device.is_gfx95_supported", return_value=False), + ): + self.assertFalse(mg.fused_wo_b_available()) + + +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..7eabad612ffb --- /dev/null +++ b/test/registered/amd/test_mori_gemm_ar_collective.py @@ -0,0 +1,390 @@ +"""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=240, 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 +#: 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: + 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") + + from sglang.srt.runtime_context import get_parallel + + group = _Group(rank, world, cpu_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 + + +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 = [] + 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") + + 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 = 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) + 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}" + ) + 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() + + +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 ( + 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) + dist.destroy_process_group() + + +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 + 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) + + # 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): + built = mg._construct_collectively(M, N, K) + _assert_agreed(mg, rank, built, expect=True) + state = mg._state + try: + 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() + 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() + + +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, 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=(world, *args), + nprocs=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", "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. + + 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. + """ + 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__": + 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..a984d41c1de3 --- /dev/null +++ b/test/registered/amd/test_mori_gemm_ar_contract.py @@ -0,0 +1,254 @@ +"""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 +from types import SimpleNamespace + +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.gemm.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 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( + torch.float8_e4m3fn + ) + ue8m0 = torch.randint( + 120, 135, (N // 32, K // 32), dtype=torch.uint8, device="cuda" + ) + 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), + ): + 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.""" + 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()