[AMD] Fix Load and Inference of MLA models with Quark PTPC FP8 attention on ROCm

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -444,6 +444,10 @@ def _get_q_k_bf16(
         enable_dual_stream: bool,
         forward_batch: ForwardBatch,
     ):
+        # fused_rms_fp8_group_quant with output_unquantized_inp1=True produces a
+        # (fp8, scale, bf16) 3-tuple; extract the bf16 for the unquantized wk layer.
+        if isinstance(x, tuple) and len(x) >= 3:
+            x = x[2]
         if enable_dual_stream:
             current_stream = torch.cuda.current_stream()
             self.alt_stream.wait_stream(current_stream)
@@ -534,6 +538,10 @@ def _get_k_bf16(
         positions: torch.Tensor,
         enable_dual_stream: bool,
     ):
+        # fused_rms_fp8_group_quant with output_unquantized_inp1=True produces a
+        # (fp8, scale, bf16) 3-tuple; extract the bf16 for the unquantized wk layer.
+        if isinstance(x, tuple) and len(x) >= 3:
+            x = x[2]
         # Compute only key, skip query
         key, _ = self.wk(x)
         key = self.k_norm(key)

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

@@ -1428,6 +1428,9 @@ def channel_quant_to_tensor_quant(
     x_q_channel: torch.Tensor,
     x_s: torch.Tensor,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
+    # Reshape per-output-channel scale [N] -> [N, 1, ...] to broadcast over K dims.
+    if x_s.dim() < x_q_channel.dim():
+        x_s = x_s.view(x_s.shape + (1,) * (x_q_channel.dim() - x_s.dim()))
     x_dq_channel = x_q_channel.to(torch.float32) * x_s
     x_q_tensor, scale = (
         scaled_fp8_quant(x_dq_channel)
@@ -1693,9 +1696,10 @@ def apply_fp8_ptpc_linear(
     compressed_tensor_quant: bool = False,
 ) -> torch.Tensor:
     """FP8 per-token per-channel linear. Only used with the aiter (ROCm) backend."""
-    # Handle pre-quantized (fp8_tensor, scale) tuple from fused RMSNorm+Quant
+    # Handle pre-quantized (fp8_tensor, scale[, bf16]) tuple from fused RMSNorm+Quant.
+    # The optional 3rd element is the unquantized bf16 tensor kept for DSA; ignore it.
     if isinstance(input, tuple):
-        q_input, x_scale = input
+        q_input, x_scale = input[0], input[1]
         q_input = q_input.view(-1, q_input.shape[-1])
         output_shape = [*q_input.shape[:-1], weight.shape[0]]
         output = aiter.gemm_a8w8_bpreshuffle(

python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py

@@ -12,6 +12,7 @@
 )
 from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.layers.quantization.fp8_utils import (
+    _use_aiter_bpreshuffle_gfx95,
     apply_fp8_linear,
     cutlass_fp8_supported,
     normalize_e4m3fn_to_e4m3fnuz,
@@ -174,6 +175,27 @@ def apply_weights(
         x: torch.Tensor,
         bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
+        # fused_rms_fp8_group_quant (aiter gfx95) passes activations as a tuple.
+        # 3-tuple (fp8, scale, bf16): use the bf16 copy produced by output_unquantized_inp1=True.
+        # 2-tuple (fp8, group_scale): dequantize to bf16; scale layout mirrors
+        #   transpose_scale=_use_aiter_bpreshuffle_gfx95 used at quantization time.
+        if isinstance(x, tuple):
+            if len(x) >= 3:
+                x = x[2]
+            else:
+                fp8_q, gs = x[0], x[1]
+                m, k = fp8_q.shape
+                if _use_aiter_bpreshuffle_gfx95:
+                    gs = gs.t().contiguous()
+                k_groups = gs.shape[1]
+                group_size = k // k_groups
+                x = (
+                    fp8_q.to(torch.float32)
+                    .view(m, k_groups, group_size)
+                    .mul_(gs.to(torch.float32).unsqueeze(-1))
+                    .view(m, k)
+                    .to(torch.bfloat16)
+                )
 
         return apply_fp8_linear(
             x,