[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: