[Perf] [AMD] DSV4.1: Fuse wo_b's GEMM with its all-reduce via mori cco on gfx950

.github/workflows/pr-test-amd-mori.yml

@@ -30,6 +30,83 @@ concurrency:
   cancel-in-progress: ${{ github.event_name != 'pull_request' || github.event.action != 'labeled' }}
 
 jobs:
+  # The other jobs here are label-only. gemm-ar also runs unasked on a change to
+  # the adapter or to the pinned MORI_COMMIT, which is what it guards against:
+  # a mori whose gemm_ar signatures moved breaks the fused path silently, and a
+  # PR bumping the commit would otherwise never run the test that catches it.
+  check-changes:
+    runs-on: ubuntu-latest
+    outputs:
+      gemm_ar: ${{ steps.filter.outputs.gemm_ar }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.ref || github.sha }}
+
+      - name: Detect file changes
+        id: filter
+        uses: dorny/paths-filter@v3
+        if: github.event_name == 'pull_request'
+        with:
+          filters: |
+            gemm_ar:
+              - "python/sglang/srt/layers/mori_gemm_ar.py"
+              - "test/registered/amd/test_mori_gemm_ar*.py"
+              - "docker/rocm.Dockerfile"
+              - ".github/workflows/pr-test-amd-mori.yml"
+
+  mori-gemm-ar:
+    needs: [check-changes]
+    if: |
+      always() && (
+        github.event_name == 'workflow_dispatch' ||
+        needs.check-changes.outputs.gemm_ar == 'true' ||
+        (github.event.action == 'labeled' && github.event.label.name == 'run-mori-gemm-ar') ||
+        (github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-mori-gemm-ar'))
+      )
+    runs-on: linux-mi35x-gpu-8
+    timeout-minutes: 60
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.ref || github.sha }}
+
+      - name: Ensure VRAM is clear
+        run: bash scripts/ci/amd/ensure_vram_clear.sh rocm
+
+      - name: Start CI container
+        run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm10
+        env:
+          GITHUB_WORKSPACE: ${{ github.workspace }}
+
+      - name: Install dependencies
+        run: bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps
+
+      # 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' ||

python/sglang/srt/environ.py

@@ -1870,6 +1870,35 @@ class Envs:
     # tuned batched_gemm_bf16 (gfx95). Off by default; see deepseek_v4.py
     # _apply_wo_a_bf16_matmul.
     SGLANG_OPT_USE_AITER_BATCHED_GEMM = EnvBool(False)
+    # ROCm gfx950 TP-only: 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