From 6a810383176a4e2a5b88f4a28ccf42cce2c1c538 Mon Sep 17 00:00:00 2001
From: "jacky.cheng" <yichiche@amd.com>
Date: Tue, 25 Aug 2026 17:49:57 +0800
Subject: [PATCH] [AMD] Drop redundant FP8 bpreshuffle scale transpose via
 fused AR kernel (#33021)

---
 .../srt/distributed/communication_op.py       | 13 ++++++++++-
 .../sglang/srt/distributed/parallel_state.py  |  6 +++++
 python/sglang/srt/layers/layernorm.py         | 22 ++++++++++++-------
 3 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/python/sglang/srt/distributed/communication_op.py b/python/sglang/srt/distributed/communication_op.py
index e8f91bd53..311a17dac 100644
--- a/python/sglang/srt/distributed/communication_op.py
+++ b/python/sglang/srt/distributed/communication_op.py
@@ -47,6 +47,7 @@ def tensor_model_parallel_fused_allreduce_rmsnorm_quant_per_group(
     eps: float,
     group_size: int = 128,
     emit_bf16: bool = False,
+    transpose_scale: bool = False,
 ) -> Optional[Tuple[torch.Tensor, ...]]:
     """Fused TP all-reduce + RMSNorm + per-group FP8 quant (ROCm/aiter).
 
@@ -57,9 +58,19 @@ def tensor_model_parallel_fused_allreduce_rmsnorm_quant_per_group(
     service the request (non-AMD, custom AR disabled, shape unsupported).
     Callers MUST handle ``None`` by falling back to the separate
     fused-AR-RMSNorm + per-group-quant path.
+
+    ``transpose_scale=True`` asks the kernel to emit the per-group scale in the
+    column-major layout the gfx95 bpreshuffle GEMM consumes, so callers can skip
+    the post-kernel ``materialize_bpreshuffle_fp8_scale`` transpose.
     """
     return get_tp_group().fused_allreduce_rmsnorm_quant_per_group(
-        input_, residual_inp_, weight_, eps, group_size, emit_bf16=emit_bf16
+        input_,
+        residual_inp_,
+        weight_,
+        eps,
+        group_size,
+        emit_bf16=emit_bf16,
+        transpose_scale=transpose_scale,
     )
 
 
diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py
index 6a0406e27..5aeb7a573 100644
--- a/python/sglang/srt/distributed/parallel_state.py
+++ b/python/sglang/srt/distributed/parallel_state.py
@@ -840,6 +840,7 @@ class GroupCoordinator:
         eps: float,
         group_size: int = 128,
         emit_bf16: bool = False,
+        transpose_scale: bool = False,
     ) -> Optional[Tuple[torch.Tensor, ...]]:
         """Attempt fused all-reduce + RMSNorm + per-group FP8 quant.
 
@@ -853,6 +854,10 @@ class GroupCoordinator:
         ``(fp8, residual_out, scale, bf16)`` — used by GDN-style layers that
         need both an FP8 projection and a bf16 gating projection without
         launching a separate per-group quant kernel.
+
+        When ``transpose_scale=True`` the kernel writes the per-group scale in
+        the column-major layout the gfx95 bpreshuffle GEMM consumes, so the
+        caller can skip the post-kernel scale transpose.
         """
         if not (is_hip() and is_gfx95_supported()):
             return None
@@ -899,6 +904,7 @@ class GroupCoordinator:
                 group_size,
                 use_1stage_ar,
                 emit_bf16=emit_bf16,
+                transpose_scale=transpose_scale,
             )
         except Exception:
             return None
diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py
index 6c4f92516..ef3828dfc 100644
--- a/python/sglang/srt/layers/layernorm.py
+++ b/python/sglang/srt/layers/layernorm.py
@@ -303,17 +303,24 @@ def _forward_with_allreduce_fusion_quant_per_group(
     if world_size <= 1:
         return None
 
-    # TODO: When ROCm/aiter#3652 is available in our bundled aiter, plumb
-    # transpose_scale=use_bpreshuffle into the fused AR+RMSNorm+quant kernel
-    # and drop this explicit post-kernel scale materialization.
+    # ``transpose_scale=use_bpreshuffle`` asks the fused kernel to emit the
+    # per-group scale directly in the column-major layout the gfx95 bpreshuffle
+    # GEMM consumes (identical to ``materialize_bpreshuffle_fp8_scale``), so the
+    # fused-success paths below need no post-kernel transpose. The separate
+    # per-group-quant fallbacks still materialize: ``per_1x128_quant``'s own
+    # ``transpose_scale`` byte-shuffles the scale into a *different* arrangement,
+    # not the column-major layout this GEMM expects.
     if not keep_bf16:
         result = tensor_model_parallel_fused_allreduce_rmsnorm_quant_per_group(
-            x, residual, weight, norm_module.variance_epsilon, group_size
+            x,
+            residual,
+            weight,
+            norm_module.variance_epsilon,
+            group_size,
+            transpose_scale=use_bpreshuffle,
         )
         if result is not None:
             fp8_out, residual_out, scale_out = result
-            if use_bpreshuffle:
-                scale_out = materialize_bpreshuffle_fp8_scale(scale_out)
             return (fp8_out, scale_out), residual_out
 
         # Fallback: fused AR+RMSNorm then separate per-group quant.
@@ -345,11 +352,10 @@ def _forward_with_allreduce_fusion_quant_per_group(
         norm_module.variance_epsilon,
         group_size,
         emit_bf16=True,
+        transpose_scale=use_bpreshuffle,
     )
     if result is not None and len(result) == 4:
         fp8_out, residual_out, scale_out, bf16_out = result
-        if use_bpreshuffle:
-            scale_out = materialize_bpreshuffle_fp8_scale(scale_out)
         return (bf16_out, fp8_out, scale_out), residual_out
 
     fused_result = tensor_model_parallel_fused_allreduce_rmsnorm(