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

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

@@ -588,6 +588,16 @@ 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):
+            if len(x) >= 3:
+                x = x[2]
+            else:
+                raise ValueError(
+                    f"_get_q_k_bf16 received a {len(x)}-tuple; expected a plain tensor "
+                    "or a 3-tuple (fp8, scale, bf16) from fused_rms_fp8_group_quant."
+                )
         weights_raw = None
         if enable_dual_stream:
             current_stream = torch.cuda.current_stream()
@@ -675,6 +685,16 @@ def _get_k_bf16(
         *,
         apply_hadamard: bool = True,
     ):
+        # 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):
+            if len(x) >= 3:
+                x = x[2]
+            else:
+                raise ValueError(
+                    f"_get_k_bf16 received a {len(x)}-tuple; expected a plain tensor "
+                    "or a 3-tuple (fp8, scale, bf16) from fused_rms_fp8_group_quant."
+                )
         # Non-fusion path only; self.wk does not exist when fusion is on.
         key, _ = self.wk(x)
         key = self.k_norm(key)

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

@@ -1988,6 +1988,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.
+    while x_s.dim() < x_q_channel.dim():
+        x_s = x_s.unsqueeze(-1)
     x_dq_channel = x_q_channel.to(torch.float32) * x_s
     x_q_tensor, scale = (
         scaled_fp8_quant(x_dq_channel)
@@ -2353,9 +2356,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

@@ -174,7 +174,14 @@ def apply_weights(
         x: torch.Tensor,
         bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-
+        # Activations must be a plain bf16 tensor: per-channel FP8 (weight_scale
+        # [N, 1]) is gated off the aiter fused RMSNorm+quant kernel upstream by
+        # _is_block_scale_fp8, so the per-token quant happens in apply_fp8_linear.
+        assert not isinstance(x, tuple), (
+            "quark W8A8 FP8 linear received a pre-quantized tuple; a fused "
+            "RMSNorm+quant producer was not gated off by _is_block_scale_fp8 "
+            "(per-channel FP8 must fall through to the plain bf16 path)."
+        )
         return apply_fp8_linear(
             x,
             layer.weight,

test/registered/unit/layers/quantization/test_quark_w8a8_fp8_ptpc.py

@@ -0,0 +1,65 @@
+"""CPU-only regression tests for the Quark PTPC-FP8 (W8A8 FP8) MLA attention fix.
+
+These guard the hardware-independent pieces of the fix that enables loading and
+running MLA models (e.g. GlmMoeDsaForCausalLM) with Quark-quantized PTPC FP8
+attention on ROCm/gfx95.
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=6, suite="base-a-test-cpu")
+
+import unittest
+
+import torch
+
+from sglang.srt.layers.quantization.fp8_utils import channel_quant_to_tensor_quant
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestChannelQuantToTensorQuantBroadcast(CustomTestCase):
+    """Regression for the 1D per-channel scale failing to broadcast against a 2D
+    weight.
+    """
+
+    def test_2d_weight_1d_scale_broadcasts_per_output_channel(self):
+        n, k = 8, 4
+        weight = torch.ones(n, k, dtype=torch.float8_e4m3fn)
+        # Distinct per-channel scales so a wrong-axis broadcast would be visible.
+        scale = torch.arange(1, n + 1, dtype=torch.float32)
+
+        q_tensor, tensor_scale = channel_quant_to_tensor_quant(weight, scale)
+
+        # Round-trip: dequantized tensor must equal weight * per-row scale. The
+        # requant to a single tensor-wide FP8 scale is lossy, so tolerance is
+        # loose — a wrong-axis broadcast would be off by whole factors (>>10%).
+        dq = q_tensor.to(torch.float32) * tensor_scale
+        expected = torch.ones(n, k, dtype=torch.float32) * scale.unsqueeze(-1)
+        torch.testing.assert_close(dq, expected, rtol=0.1, atol=0.5)
+
+    def test_mismatched_n_k_would_fail_without_unsqueeze(self):
+        # N != K: the pre-fix ``weight * scale`` raises; the fix must not.
+        n, k = 6, 3
+        weight = torch.ones(n, k, dtype=torch.float8_e4m3fn)
+        scale = torch.arange(1, n + 1, dtype=torch.float32)
+
+        q_tensor, tensor_scale = channel_quant_to_tensor_quant(weight, scale)
+
+        self.assertEqual(q_tensor.shape, torch.Size([n, k]))
+
+    def test_scale_already_matching_rank_is_unchanged(self):
+        # A [N, 1] scale must not be unsqueezed further (guards the while-loop
+        # condition from over-reshaping).
+        n, k = 5, 4
+        weight = torch.ones(n, k, dtype=torch.float8_e4m3fn)
+        scale = torch.arange(1, n + 1, dtype=torch.float32).unsqueeze(-1)
+
+        q_tensor, tensor_scale = channel_quant_to_tensor_quant(weight, scale)
+
+        dq = q_tensor.to(torch.float32) * tensor_scale
+        expected = torch.ones(n, k, dtype=torch.float32) * scale
+        torch.testing.assert_close(dq, expected, rtol=0.1, atol=0.5)
+
+
+if __name__ == "__main__":
+    unittest.main()