[AMD] DeepSeek-V4 MI355X: eliminate bpreshuffle fp8-scale copies at producer sites (MoE down, MLA o_proj bmm) (#33166)

Co-authored-by: kk <43161300+kkHuang-amd@users.noreply.github.com>
Co-authored-by: Thomas Wang <thomawan@amd.com>

python/sglang/srt/layers/quantization/fp8_utils.py

@@ -115,10 +115,21 @@ def materialize_bpreshuffle_fp8_scale(scale: torch.Tensor) -> torch.Tensor:
 
 
 def view_aiter_fused_rms_transposed_fp8_scale(scale: torch.Tensor) -> torch.Tensor:
-    """Expose AITER fused-RMS ``transpose_scale=True`` storage logically.
+    """Zero-copy view of a ``transpose_scale=True`` fp8 group scale.
 
-    The fused-RMS op returns transposed physical bytes through a row-major-looking
-    view. Restore logical ``[M, G]`` indexing without copying those bytes.
+    Producer-neutral counterpart of ``materialize_bpreshuffle_fp8_scale``. When an
+    AITER quant/fused-RMS kernel is asked for ``transpose_scale=True`` it writes the
+    per-token group scale directly in physical ``[num_groups, tokens]`` byte order
+    behind a row-major-looking ``[tokens, num_groups]`` tensor. Swapping the strides
+    restores logical ``[M, G]`` indexing over those same bytes -- i.e. the
+    column-major layout the gfx95 bpreshuffle GEMM consumes -- with no copy. Callers
+    that instead take the row-major (``transpose_scale=False``) path relayout via
+    ``materialize_bpreshuffle_fp8_scale``; this is the bit-identical no-copy path.
+
+    Only valid for M(tokens) >= 2. At M == 1 the ``[1, G]`` and ``[G, 1]`` byte
+    orders coincide, so producers keep ``transpose_scale=False`` and materialize;
+    the stride swap here would be a no-op on shape but is never taken at M == 1.
+    Non-2-D scales (e.g. per-tensor) pass through unchanged.
     """
     if scale.dim() != 2:
         return scale
@@ -136,6 +147,28 @@ def materialize_bpreshuffle_fp8_scale_tuple(
     )
 
 
+def view_aiter_fused_rms_transposed_fp8_scale_tuple(
+    value: Tuple[torch.Tensor, ...],
+) -> Tuple[torch.Tensor, ...]:
+    """Zero-copy scale reinterpret for FP8 ``(q_input, x_scale, ...)`` tuples."""
+    return (value[0], view_aiter_fused_rms_transposed_fp8_scale(value[1]), *value[2:])
+
+
+def emit_transposed_bpreshuffle_scale(m: int, *, on_bpreshuffle_gfx95: bool) -> bool:
+    """Whether a producer should emit its fp8 scale already transposed.
+
+    Producer sites choose between two equivalent gfx95 bpreshuffle scale layouts:
+    ``transpose_scale=True`` + zero-copy ``view_aiter_fused_rms_transposed_fp8_scale`` (this
+    predicate True), or row-major ``transpose_scale=False`` +
+    ``materialize_bpreshuffle_fp8_scale`` (this predicate False). The transposed
+    zero-copy path is only taken on gfx95 bpreshuffle and only for M(tokens) >= 2:
+    at M == 1 the ``[1, G]`` and ``[G, 1]`` byte orders coincide, so the transposed
+    emit buys nothing and the materialize path is used. Centralizes the gate shared
+    by the MoE-down and MLA o_proj producer sites.
+    """
+    return on_bpreshuffle_gfx95 and m >= 2
+
+
 def use_aiter_triton_gemm_w8a8_tuned_gfx950(n: int, k: int) -> bool:
     if _FORCE_CK_W8A8:
         return False

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py

@@ -31,7 +31,9 @@ from sglang.srt.layers.dcp import (
 )
 from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run
 from sglang.srt.layers.quantization.fp8_utils import (
+    emit_transposed_bpreshuffle_scale,
     materialize_bpreshuffle_fp8_scale_tuple,
+    view_aiter_fused_rms_transposed_fp8_scale_tuple,
 )
 from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
 from sglang.srt.lora.deepseek_mla_correction import (
@@ -233,13 +235,24 @@ def rocm_absorb_v_bmm(
         if attn.o_proj.weight.dtype == torch.uint8:
             attn_bmm_output = fused_flatten_mxfp4_quant(_bmm_buf)
         elif _is_block_scale_fp8(attn.o_proj):
+            # No-copy fp8 scale: emit the bpreshuffle scale already transposed and
+            # reinterpret it with a stride swap, instead of relaying out a copy.
+            # Falls back to the materialize (copy) path at M == 1 / non-gfx95.
+            _emit_bpre = emit_transposed_bpreshuffle_scale(
+                _bmm_buf.shape[0],
+                on_bpreshuffle_gfx95=_use_aiter_bpreshuffle_gfx95,
+            )
             attn_bmm_output = fused_flatten_fp8_group_quant(
                 _bmm_buf,
                 group_size=128,
                 dtype_quant=torch.float8_e4m3fn,
-                transpose_scale=False,
+                transpose_scale=_emit_bpre,
             )
-            if _use_aiter_bpreshuffle_gfx95:
+            if _emit_bpre:
+                attn_bmm_output = view_aiter_fused_rms_transposed_fp8_scale_tuple(
+                    attn_bmm_output
+                )
+            elif _use_aiter_bpreshuffle_gfx95:
                 attn_bmm_output = materialize_bpreshuffle_fp8_scale_tuple(
                     attn_bmm_output
                 )
@@ -250,13 +263,24 @@ def rocm_absorb_v_bmm(
         attn_bmm_output = fused_flatten_mxfp4_quant(attn_bmm_output)
     elif _is_block_scale_fp8(attn.o_proj):
         attn_bmm_output = attn_bmm_output.transpose(0, 1)
+        # No-copy fp8 scale: emit the bpreshuffle scale already transposed and
+        # reinterpret it with a stride swap, instead of relaying out a copy.
+        # Falls back to the materialize (copy) path at M == 1 / non-gfx95.
+        _emit_bpre = emit_transposed_bpreshuffle_scale(
+            attn_bmm_output.shape[0],
+            on_bpreshuffle_gfx95=_use_aiter_bpreshuffle_gfx95,
+        )
         attn_bmm_output = fused_flatten_fp8_group_quant(
             attn_bmm_output,
             group_size=128,
             dtype_quant=torch.float8_e4m3fn,
-            transpose_scale=False,
+            transpose_scale=_emit_bpre,
         )
-        if _use_aiter_bpreshuffle_gfx95:
+        if _emit_bpre:
+            attn_bmm_output = view_aiter_fused_rms_transposed_fp8_scale_tuple(
+                attn_bmm_output
+            )
+        elif _use_aiter_bpreshuffle_gfx95:
             attn_bmm_output = materialize_bpreshuffle_fp8_scale_tuple(attn_bmm_output)
     else:
         attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)

python/sglang/srt/models/deepseek_v2.py

@@ -123,7 +123,9 @@ from sglang.srt.layers.moe.utils import (
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.quantization.fp8 import Fp8Config
 from sglang.srt.layers.quantization.fp8_utils import (
+    emit_transposed_bpreshuffle_scale,
     materialize_bpreshuffle_fp8_scale,
+    view_aiter_fused_rms_transposed_fp8_scale,
 )
 from sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe import (
     maybe_fuse_routed_scale_and_shared_add,
@@ -419,14 +421,19 @@ class DeepseekV2MLP(nn.Module):
             if self._fused_clamp_use_fp8:
                 from aiter import dtypes
 
+                _emit_bpre = emit_transposed_bpreshuffle_scale(
+                    gate_up.shape[0], on_bpreshuffle_gfx95=_use_aiter_bpreshuffle_gfx95
+                )
                 x_fp8, x_scale = fused_clamp_act_mul(
                     gate_up,
                     swiglu_limit=self.swiglu_limit,
                     activation="silu",
                     dtype_quant=dtypes.fp8,
-                    transpose_scale=False,
+                    transpose_scale=_emit_bpre,
                 )
-                if _use_aiter_bpreshuffle_gfx95:
+                if _emit_bpre:
+                    x_scale = view_aiter_fused_rms_transposed_fp8_scale(x_scale)
+                elif _use_aiter_bpreshuffle_gfx95:
                     x_scale = materialize_bpreshuffle_fp8_scale(x_scale)
                 x = (x_fp8, x_scale)
             else:

test/registered/unit/layers/test_fp8_bpreshuffle_producer_mi35x.py

@@ -0,0 +1,156 @@
+"""Producer-level validation of the gfx95 bpreshuffle fp8-scale no-copy path.
+
+The CPU tests in ``test_fp8_bpreshuffle_scale.py`` only prove that the stride
+reinterpret recovers a *fabricated* ``transpose_scale=True`` layout. They cannot
+catch the real AITER producers ignoring or mis-implementing ``transpose_scale``.
+These tests invoke the two producers the optimization actually routes through --
+``fused_clamp_act_mul`` (MoE down) and ``fused_flatten_fp8_group_quant`` (MLA
+o_proj) -- and prove, on real gfx95 kernels, that the two producer paths are
+equivalent:
+
+    transpose_scale=True  + view_aiter_fused_rms_transposed_fp8_scale   (the optimized path)
+    transpose_scale=False + materialize_bpreshuffle_fp8_scale   (the row-major path)
+
+For M(tokens) >= 2 they must agree bit-for-bit on the quantized output and on the
+scale *values* after relayout, with the no-copy path landing on the bpreshuffle
+``(1, M)`` column-major stride and sharing the producer's storage. M == 1 is the
+materialize-only fallback (``emit_transposed_bpreshuffle_scale`` gates the
+transposed emit on M >= 2); we pin that its materialized scale is well-formed.
+
+Requires a gfx95 (MI35X) GPU with aiter; skips otherwise.
+"""
+
+import unittest
+
+import torch
+
+from sglang.srt.layers.quantization.fp8_utils import (
+    emit_transposed_bpreshuffle_scale,
+    materialize_bpreshuffle_fp8_scale,
+    view_aiter_fused_rms_transposed_fp8_scale,
+)
+from sglang.srt.utils.common import is_gfx95_supported, is_hip
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+_GROUP_SIZE = 128
+
+
+@unittest.skipUnless(
+    is_hip() and is_gfx95_supported(),
+    "bpreshuffle fp8-scale no-copy is a gfx95 (MI35X) + aiter optimization",
+)
+class TestBpreshuffleProducerScaleNoCopy(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        try:
+            from aiter import dtypes  # noqa: F401
+            from aiter.ops.triton.fused_fp8_quant import (  # noqa: F401
+                fused_flatten_fp8_group_quant,
+            )
+            from aiter.ops.triton.fusions.fused_clamp_act_mul import (  # noqa: F401
+                fused_clamp_act_mul,
+            )
+        except Exception as err:  # pragma: no cover - env-dependent
+            raise unittest.SkipTest(f"aiter producers unavailable: {err}")
+        cls.device = "cuda"  # torch maps "cuda" onto the ROCm HIP device
+
+    def setUp(self):
+        torch.manual_seed(0)
+
+    # --- producer adapters: run the producer and normalize to (q, scale) ------
+
+    def _run_fused_clamp_act_mul(self, m, transpose_scale):
+        from aiter import dtypes
+        from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul
+
+        inter = 4 * _GROUP_SIZE  # G = 4 groups
+        gate_up = torch.randn(m, 2 * inter, device=self.device, dtype=torch.bfloat16)
+        q, scale = fused_clamp_act_mul(
+            gate_up,
+            swiglu_limit=7.0,
+            activation="silu",
+            dtype_quant=dtypes.fp8,
+            transpose_scale=transpose_scale,
+        )
+        return q, scale, gate_up
+
+    def _run_fused_flatten_fp8_group_quant(self, m, transpose_scale):
+        from aiter.ops.triton.fused_fp8_quant import fused_flatten_fp8_group_quant
+
+        heads, dim = 8, _GROUP_SIZE  # heads*dim = 1024 -> G = 8 groups
+        buf = torch.randn(m, heads, dim, device=self.device, dtype=torch.bfloat16)
+        out = fused_flatten_fp8_group_quant(
+            buf,
+            group_size=_GROUP_SIZE,
+            dtype_quant=torch.float8_e4m3fn,
+            transpose_scale=transpose_scale,
+        )
+        return out[0], out[1], buf
+
+    def _assert_producer_paths_equivalent(self, run_producer, name):
+        for m in (1, 2, 8, 16):
+            with self.subTest(producer=name, m=m):
+                # Row-major path (transpose_scale=False) + materialize. This is
+                # the path M == 1 takes in production, so it must always be valid.
+                torch.manual_seed(m)
+                q_f, s_f, _ = run_producer(m, transpose_scale=False)
+                mat = materialize_bpreshuffle_fp8_scale(s_f)
+                self.assertEqual(mat.dim(), 2)
+                self.assertEqual(mat.shape[0], m)
+
+                if not emit_transposed_bpreshuffle_scale(m, on_bpreshuffle_gfx95=True):
+                    # M == 1: the transposed emit is skipped by design. At M == 1 the
+                    # [1, G] row-major and [G, 1] column-major byte orders coincide,
+                    # so materialize is a no-op that keeps the natural (G, 1) stride
+                    # (NOT the (1, M) column-major stride taken for M >= 2) while
+                    # sharing storage; the scale values must survive intact.
+                    self.assertEqual(m, 1)
+                    self.assertEqual(mat.stride(), (s_f.shape[1], 1))  # (G, 1)
+                    self.assertTrue(torch.equal(mat, s_f))
+                    continue
+
+                # M >= 2: materialize lands on the bpreshuffle (1, M) column-major
+                # stride.
+                self.assertEqual(mat.stride(), (1, m))
+
+                # Optimized path: same input, transpose_scale=True + no-copy view.
+                torch.manual_seed(m)
+                q_t, s_t, _ = run_producer(m, transpose_scale=True)
+                nocopy = view_aiter_fused_rms_transposed_fp8_scale(s_t)
+
+                # Quantized output is layout-independent: transpose_scale only
+                # changes the *scale* storage, never the quantized bytes.
+                self.assertEqual(
+                    q_t.dtype, q_f.dtype, "quant dtype differs between paths"
+                )
+                self.assertTrue(
+                    torch.equal(q_t.view(torch.uint8), q_f.view(torch.uint8)),
+                    "quantized output differs between transpose_scale paths",
+                )
+
+                # Scale values match the row-major + materialize path exactly...
+                self.assertEqual(nocopy.shape, mat.shape)
+                self.assertTrue(
+                    torch.equal(nocopy, mat),
+                    "no-copy scale values differ from materialized",
+                )
+                # ...on the bpreshuffle (1, M) stride, with no allocation.
+                self.assertEqual(nocopy.stride(), (1, m))
+                self.assertEqual(nocopy.data_ptr(), s_t.data_ptr())
+
+    def test_fused_clamp_act_mul_producer_paths_equivalent(self):
+        self._assert_producer_paths_equivalent(
+            self._run_fused_clamp_act_mul, "fused_clamp_act_mul"
+        )
+
+    def test_fused_flatten_fp8_group_quant_producer_paths_equivalent(self):
+        self._assert_producer_paths_equivalent(
+            self._run_fused_flatten_fp8_group_quant, "fused_flatten_fp8_group_quant"
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/test_fp8_bpreshuffle_scale.py

@@ -4,9 +4,11 @@ from unittest.mock import patch
 import torch
 
 from sglang.srt.layers.quantization.fp8_utils import (
+    emit_transposed_bpreshuffle_scale,
     materialize_bpreshuffle_fp8_scale,
     materialize_bpreshuffle_fp8_scale_tuple,
     view_aiter_fused_rms_transposed_fp8_scale,
+    view_aiter_fused_rms_transposed_fp8_scale_tuple,
 )
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
@@ -175,5 +177,65 @@ class TestBpreshuffleScaleFreshQuantNoCopy(CustomTestCase):
         self.assertTrue(materialized.t().is_contiguous())
 
 
+class TestBpreshuffleScaleProducerNoCopy(CustomTestCase):
+    """Producer-site (MoE down, MLA o_proj bmm) coverage for the shared no-copy
+    reinterpret that isn't exercised by the dense fresh-quant class above: the
+    guard that leaves non-2D scales untouched, and the tuple wrapper the producers
+    emit through (``view_aiter_fused_rms_transposed_fp8_scale_tuple``), which must
+    reinterpret only the scale slot and pass the rest through by identity."""
+
+    def test_nocopy_passthrough_for_non_2d_scale(self):
+        for scale in (
+            torch.arange(5, dtype=torch.float32),  # 1-D (per-tensor scale)
+            torch.arange(24, dtype=torch.float32).reshape(2, 3, 4),  # 3-D
+        ):
+            with self.subTest(dim=scale.dim()):
+                self.assertIs(view_aiter_fused_rms_transposed_fp8_scale(scale), scale)
+
+    def test_tuple_helper_reinterprets_only_the_scale_slot(self):
+        q_input = torch.ones((3, 8), dtype=torch.float8_e4m3fn)
+        values = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+        emitted = _simulate_transpose_scale_emit(values)
+        bf16_side = torch.ones((3, 8), dtype=torch.bfloat16)
+
+        q_out, scale_out, bf16_out = view_aiter_fused_rms_transposed_fp8_scale_tuple(
+            (q_input, emitted, bf16_side)
+        )
+
+        self.assertIs(q_out, q_input)
+        self.assertIs(bf16_out, bf16_side)
+        self.assertTrue(
+            torch.equal(scale_out, materialize_bpreshuffle_fp8_scale(values))
+        )
+        self.assertEqual(scale_out.stride(), (1, values.shape[0]))
+        self.assertEqual(scale_out.data_ptr(), emitted.data_ptr())
+
+
+class TestEmitTransposedBpreshuffleScaleGate(CustomTestCase):
+    """Pins the producer emit-gate shared by the MoE-down and MLA o_proj sites:
+    the transposed zero-copy path is taken only on gfx95 bpreshuffle and only for
+    M(tokens) >= 2; M == 1 must fall back to the materialize path. Guards the
+    ``>= 2`` boundary against being widened to `M >= 1` (which would send a
+    degenerate single-token scale down the stride-swap path)."""
+
+    def test_gate_false_off_gfx95_regardless_of_m(self):
+        for m in (1, 2, 8):
+            with self.subTest(m=m):
+                self.assertFalse(
+                    emit_transposed_bpreshuffle_scale(m, on_bpreshuffle_gfx95=False)
+                )
+
+    def test_gate_requires_m_ge_2_on_gfx95(self):
+        # M == 1 -> materialize fallback; M >= 2 -> transposed zero-copy path.
+        self.assertFalse(
+            emit_transposed_bpreshuffle_scale(1, on_bpreshuffle_gfx95=True)
+        )
+        for m in (2, 3, 16):
+            with self.subTest(m=m):
+                self.assertTrue(
+                    emit_transposed_bpreshuffle_scale(m, on_bpreshuffle_gfx95=True)
+                )
+
+
 if __name__ == "__main__":
     unittest.main()