[ROCM] Optimized deepseek-r1 model with rmsnorm + fp8 quant fusion (#12689)

should be clean after https://github.com/sgl-project/sglang/pull/13017 landed

python/sglang/srt/layers/communicator.py

@@ -66,6 +66,8 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip()
 _is_gfx95_supported = is_gfx95_supported()
 
 if _use_aiter and _is_gfx95_supported:
+    from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
+
     from sglang.srt.layers.quantization.rocm_mxfp4_utils import fused_rms_mxfp4_quant
 
 FUSE_ALLREDUCE_MAX_BATCH_SIZE = 2048
@@ -248,7 +250,7 @@ class LayerCommunicator:
         hidden_states: torch.Tensor,
         residual: torch.Tensor,
         forward_batch: ForwardBatch,
-        qaunt_format: str = "",
+        quant_format: str = "",
     ):
         if hidden_states.shape[0] == 0:
             residual = hidden_states
@@ -267,7 +269,7 @@ class LayerCommunicator:
                 if residual is None:
                     residual = hidden_states
 
-                    if _use_aiter and _is_gfx95_supported and ("mxfp4" in qaunt_format):
+                    if _use_aiter and _is_gfx95_supported and ("mxfp4" in quant_format):
                         hidden_states, *_, _ = fused_rms_mxfp4_quant(
                             hidden_states,
                             self.input_layernorm.weight,
@@ -277,10 +279,26 @@ class LayerCommunicator:
                             None,
                             None,
                         )
+                    elif _use_aiter and _is_gfx95_supported and ("fp8" in quant_format):
+
+                        hidden_states, _, _, _res = fused_rms_fp8_group_quant(
+                            hidden_states,
+                            self.input_layernorm.weight,
+                            self.input_layernorm.variance_epsilon,
+                            inp2=None,
+                            inp2_weight=None,
+                            inp2_epsilon=None,
+                            group_size=128,
+                            dtype_quant=torch.float8_e4m3fn,
+                            res1=None,
+                            output_unquantized_inp1=False,
+                        )
+
                     else:
                         hidden_states = self.input_layernorm(hidden_states)
                 else:
-                    if _use_aiter and _is_gfx95_supported and ("mxfp4" in qaunt_format):
+
+                    if _use_aiter and _is_gfx95_supported and ("mxfp4" in quant_format):
                         hidden_states, *_, residual = fused_rms_mxfp4_quant(
                             hidden_states,
                             self.input_layernorm.weight,
@@ -290,6 +308,23 @@ class LayerCommunicator:
                             None,
                             residual,
                         )
+                    elif _use_aiter and _is_gfx95_supported and ("fp8" in quant_format):
+                        # RMSNorm + FP8 per-group quant
+                        # return hidden_states：
+                        #   out_fp8  : FP8 activation →  a8w8 GEMM
+                        #   out_bs   : block-scale →  gemm_a8w8_blockscale.x_scale
+                        hidden_states, _, _, residual = fused_rms_fp8_group_quant(
+                            hidden_states,
+                            self.input_layernorm.weight,
+                            self.input_layernorm.variance_epsilon,
+                            inp2=None,
+                            inp2_weight=None,
+                            inp2_epsilon=None,
+                            group_size=128,
+                            dtype_quant=torch.float8_e4m3fn,
+                            res1=residual,
+                            output_unquantized_inp1=False,
+                        )
                     else:
                         hidden_states, residual = self.input_layernorm(
                             hidden_states, residual

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

@@ -488,6 +488,16 @@ class Fp8LinearMethod(LinearMethodBase):
                     True,  # is_vnni
                 )
 
+            if isinstance(x, tuple):
+                return self.w8a8_block_fp8_linear(
+                    input=x[0],
+                    weight=layer.weight,
+                    block_size=self.quant_config.weight_block_size,
+                    weight_scale=layer.weight_scale_inv,
+                    input_scale=x[1],
+                    bias=bias,
+                )
+
             return self.w8a8_block_fp8_linear(
                 input=x,
                 weight=layer.weight,

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

@@ -263,19 +263,32 @@ def aiter_w8a8_block_fp8_linear(
     input_scale: Optional[torch.Tensor] = None,
     bias: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
-    assert input_scale is None
+    # assert input_scale is None
     input_2d = input.view(-1, input.shape[-1])
     output_shape = [*input.shape[:-1], weight.shape[0]]
 
+    # if input_scale not None, input is quanted
+    if input_scale is not None:
+        q_input = input_2d
+        x_scale = input_scale
+
+    else:
         q_input, x_scale = aiter_per1x128_quant(input_2d, quant_dtype=aiter.dtypes.fp8)
+
     output = gemm_a8w8_blockscale(
-        q_input, weight, x_scale, weight_scale, dtype=input.dtype
+        q_input,
+        weight,
+        x_scale,
+        weight_scale,
+        dtype=torch.bfloat16 if input_scale is not None else input.dtype,
     )
 
     if bias is not None:
         output += bias
 
-    return output.to(dtype=input_2d.dtype).view(*output_shape)
+    return output.to(
+        dtype=torch.bfloat16 if input_scale is not None else input_2d.dtype
+    ).view(*output_shape)
 
 
 def triton_w8a8_block_fp8_linear(

python/sglang/srt/models/deepseek_v2.py

@@ -153,6 +153,8 @@ _is_gfx95_supported = is_gfx95_supported()
 _use_aiter_gfx95 = _use_aiter and _is_gfx95_supported
 
 if _use_aiter_gfx95:
+    from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
+
     from sglang.srt.layers.quantization.quark.utils import quark_post_load_weights
     from sglang.srt.layers.quantization.rocm_mxfp4_utils import (
         batched_gemm_afp4wfp4_pre_quant,
@@ -1507,13 +1509,14 @@ class DeepseekV2AttentionMLA(nn.Module):
             q, latent_cache = self.fused_qkv_a_proj_with_mqa(hidden_states)[0].split(
                 [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
             )
+
+            # NSA Indexer: cache quantized keys, auto-skip topk for sequences <= nsa_index_topk
+
+            if self.use_nsa:
                 q_lora = self.q_a_layernorm(q)
                 q = self.q_b_proj(q_lora)[0].view(
                     -1, self.num_local_heads, self.qk_head_dim
                 )
-
-            # NSA Indexer: cache quantized keys, auto-skip topk for sequences <= nsa_index_topk
-            if self.use_nsa:
                 _ = self.indexer(
                     x=hidden_states,
                     q_lora=q_lora,
@@ -1522,6 +1525,26 @@ class DeepseekV2AttentionMLA(nn.Module):
                     layer_id=self.layer_id,
                     return_indices=False,
                 )
+
+            elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
+
+                q, _, _, _ = fused_rms_fp8_group_quant(
+                    q,
+                    self.q_a_layernorm.weight,
+                    self.q_a_layernorm.variance_epsilon,
+                    None,
+                    None,
+                    None,
+                    group_size=128,
+                    dtype_quant=torch.float8_e4m3fn,
+                    res1=None,
+                    output_unquantized_inp1=False,
+                )
+                q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+            else:
+                q = self.q_a_layernorm(q)
+                q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+
         else:
             q = self.q_proj(hidden_states)[0].view(
                 -1, self.num_local_heads, self.qk_head_dim
@@ -1531,7 +1554,31 @@ class DeepseekV2AttentionMLA(nn.Module):
         _, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
         kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
         latent_cache = latent_cache.unsqueeze(1)
+
+        if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
+
+            kv_a_quanted, kv_a, _, _ = fused_rms_fp8_group_quant(
+                kv_a,
+                self.kv_a_layernorm.weight,
+                self.kv_a_layernorm.variance_epsilon,
+                None,
+                None,
+                None,
+                group_size=128,
+                dtype_quant=torch.float8_e4m3fn,
+                res1=None,
+                output_unquantized_inp1=True,  # return unqaunt kv_a
+            )
+            kv = self.kv_b_proj(
+                kv_a_quanted,
+            )[0]
+
+        else:
             kv_a = self.kv_a_layernorm(kv_a)
+            kv = self.kv_b_proj(kv_a)[0]
+
+        # kv_a = self.kv_a_layernorm(kv_a)
+
         k_pe = latent_cache[:, :, self.kv_lora_rank :]
         if self.rotary_emb is not None:
             q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
@@ -1616,6 +1663,25 @@ class DeepseekV2AttentionMLA(nn.Module):
                         self.kv_a_layernorm.weight,
                         self.kv_a_layernorm.variance_epsilon,
                     )
+                else:
+                    if (
+                        _use_aiter_gfx95
+                        and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
+                    ):
+
+                        q, _, k_nope, _ = fused_rms_fp8_group_quant(
+                            q,
+                            self.q_a_layernorm.weight,
+                            self.q_a_layernorm.variance_epsilon,
+                            k_nope,
+                            self.kv_a_layernorm.weight,
+                            self.kv_a_layernorm.variance_epsilon,
+                            group_size=128,
+                            dtype_quant=torch.float8_e4m3fn,
+                            res1=None,
+                            output_unquantized_inp1=False,
+                        )
+
                     else:
                         q = self.q_a_layernorm(q)
                         k_nope = self.kv_a_layernorm(k_nope)
@@ -1945,6 +2011,25 @@ class DeepseekV2AttentionMLA(nn.Module):
                         self.kv_a_layernorm.weight,
                         self.kv_a_layernorm.variance_epsilon,
                     )
+                else:
+                    if (
+                        _use_aiter_gfx95
+                        and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
+                    ):
+
+                        q, _, k_nope, _ = fused_rms_fp8_group_quant(
+                            q,
+                            self.q_a_layernorm.weight,
+                            self.q_a_layernorm.variance_epsilon,
+                            k_nope,
+                            self.kv_a_layernorm.weight,
+                            self.kv_a_layernorm.variance_epsilon,
+                            group_size=128,
+                            dtype_quant=torch.float8_e4m3fn,
+                            res1=None,
+                            output_unquantized_inp1=False,
+                        )
+
                     else:
                         q = self.q_a_layernorm(q)
                         k_nope = self.kv_a_layernorm(k_nope)
@@ -2673,13 +2758,30 @@ class DeepseekV2DecoderLayer(nn.Module):
     ) -> torch.Tensor:
         quant_format = (
             "mxfp4"
-            if _is_gfx95_supported
+            if (
-            and getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None) is not None
+                _is_gfx95_supported
+                and getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None)
+                is not None
                 and getattr(self.self_attn.fused_qkv_a_proj_with_mqa, "weight", None)
                 is not None
                 and self.self_attn.fused_qkv_a_proj_with_mqa.weight.dtype == torch.uint8
+            )
+            else (
+                "fp8"
+                if (
+                    _is_gfx95_supported
+                    and getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None)
+                    is not None
+                    and getattr(
+                        self.self_attn.fused_qkv_a_proj_with_mqa, "weight", None
+                    )
+                    is not None
+                    and self.self_attn.fused_qkv_a_proj_with_mqa.weight.dtype
+                    == getattr(torch, "float8_e4m3fn", None)
+                )
                 else ""
             )
+        )
 
         hidden_states, residual = self.layer_communicator.prepare_attn(
             hidden_states,

test/srt/quant/test_fused_rms_fp8_group_quant.py

@@ -0,0 +1,148 @@
+# test/srt/quant/test_fused_rms_fp8_group_quant.py
+import itertools
+import unittest
+
+import torch
+import torch.nn.functional as F
+
+from sglang.test.test_utils import CustomTestCase
+
+
+def _fp8_available() -> bool:
+    # requirement：1) GPU；2) ROCm；3) torch support float8_e4m3fn
+    if not torch.cuda.is_available():
+        return False
+    if getattr(torch.version, "hip", None) is None:
+        return False
+    return hasattr(torch, "float8_e4m3fn")
+
+
+def _rmsnorm(x, weight, eps=1e-6):
+    # row-wise RMSNorm
+    row_norm = (x * x).sum(dim=-1)
+    norm = torch.rsqrt(row_norm / x.shape[1] + eps)
+    return x * norm[:, None] * weight[None, :]
+
+
+def _per_token_fp8_group_quant(x, dtype_quant, group_size=128):
+    """per token、group-size quant, return (quantized, scale)。"""
+    DTYPE_MAX = torch.finfo(dtype_quant).max
+    M, N = x.shape
+
+    pad = (group_size - (N % group_size)) % group_size
+    if pad:
+        x_reshape = F.pad(x, (0, pad, 0, 0), "constant", 0)
+    else:
+        x_reshape = x
+
+    G = (N + group_size - 1) // group_size
+    x_reshape = x_reshape.view(M, G, group_size).to(torch.float32)
+    x_max = torch.max(torch.abs(x_reshape), dim=-1, keepdim=True)[0].clamp_min_(1e-10)
+    x_scale = x_max / DTYPE_MAX
+    inv = 1.0 / x_scale
+
+    x_q = torch.clamp(x_reshape * inv, -DTYPE_MAX, DTYPE_MAX).to(dtype_quant)
+    x_q = x_q.view(M, G * group_size)
+    if pad:
+        x_q = x_q[:, :N]
+    x_scale = x_scale.squeeze(-1)  # [M, G]
+    return x_q, x_scale
+
+
+def _upcast_fp8_group(x_q, x_s, out_dtype=torch.float32, group_size=128):
+    """unqaunt"""
+    M, N = x_q.shape
+    G = (N + group_size - 1) // group_size
+    pad = (group_size - (N % group_size)) % group_size
+
+    if pad:
+        x_q = F.pad(x_q, (0, pad, 0, 0), "constant", 0)
+
+    x_q = x_q.view(M, G, group_size).to(torch.float32)
+    x = x_q * x_s.view(M, G, 1)
+    x = x.view(M, G * group_size)[:, :N]
+    return x.to(out_dtype)
+
+
+class TestFusedRMSFP8GroupQuant(CustomTestCase):
+    #
+    DTYPES = [torch.bfloat16, torch.float16]
+    # (M, N1, N2)
+    SHAPES = [(32, 128, 7168), (128, 7168, 7168)]
+    GROUP_SIZE = [128]
+    SEEDS = [0]
+
+    @classmethod
+    def setUpClass(cls):
+        if not _fp8_available():
+            raise unittest.SkipTest("Skip: ROCm/FP8 is not available")
+        torch.set_default_device("cuda")
+
+    def _run_ref(self, x1, w1, eps1, x2, w2, eps2, res1, dtype_quant, group_size):
+        s = x1 + (res1 if res1 is not None else 0)
+        y1 = _rmsnorm(s, w1, eps1)
+        y2 = _rmsnorm(x2, w2, eps2) if x2 is not None else None
+        y1_q, y1_s = _per_token_fp8_group_quant(y1, dtype_quant, group_size)
+        return (
+            (y1_q, y1_s),
+            y1.to(x1.dtype),
+            (y2.to(x1.dtype) if y2 is not None else None),
+            (s.to(x1.dtype) if res1 is not None else None),
+        )
+
+    def _case(self, M, N1, N2, group_size, dtype, seed):
+        torch.manual_seed(seed)
+        fp8 = torch.float8_e4m3fn
+        device = "cuda"
+
+        x1 = torch.randn(M, N1, dtype=dtype, device=device) / 10
+        x2 = torch.randn(M, N2, dtype=dtype, device=device) / 10
+        w1 = torch.ones(N1, dtype=torch.float32, device=device)
+        w2 = torch.ones(N2, dtype=torch.float32, device=device)
+        res1 = torch.randn(M, N1, dtype=dtype, device=device) / 10
+
+        # ref
+        (y1_q_ref, y1_s_ref), y1_ref, y2_ref, s_ref = self._run_ref(
+            x1, w1, 1e-6, x2, w2, 1e-6, res1, fp8, group_size
+        )
+
+        # be tested：aiter fused op
+        from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
+
+        (y1_q, y1_s), y1, y2, s = fused_rms_fp8_group_quant(
+            x1,
+            w1,
+            1e-6,
+            inp2=x2,
+            inp2_weight=w2,
+            inp2_epsilon=1e-6,
+            group_size=group_size,
+            dtype_quant=fp8,
+            res1=res1,
+            output_unquantized_inp1=True,  # get unqaunt y1
+        )
+
+        torch.testing.assert_close(y1, y1_ref, atol=0.1, rtol=0.1)
+        torch.testing.assert_close(y2, y2_ref, atol=0.1, rtol=0.1)
+        torch.testing.assert_close(s, s_ref, atol=0.1, rtol=0.1)
+
+        # check unquant
+        y1_up_ref = _upcast_fp8_group(
+            y1_q_ref, y1_s_ref, out_dtype=torch.float32, group_size=group_size
+        )
+        y1_up = _upcast_fp8_group(
+            y1_q, y1_s, out_dtype=torch.float32, group_size=group_size
+        )
+        torch.testing.assert_close(y1_up, y1_up_ref, atol=0.1, rtol=0.1)
+
+    def test_fused_rms_fp8_group_quant(self):
+        for params in itertools.product(
+            self.SHAPES, self.GROUP_SIZE, self.DTYPES, self.SEEDS
+        ):
+            (M, N1, N2), g, dtype, seed = params
+            with self.subTest(M=M, N1=N1, N2=N2, group_size=g, dtype=dtype, seed=seed):
+                self._case(M, N1, N2, g, dtype, seed)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

test/srt/run_suite.py

@@ -421,6 +421,7 @@ suite_amd = {
         TestFile("openai_server/validation/test_request_length_validation.py", 31),
         TestFile("quant/test_awq_dequant.py", 2),
         TestFile("quant/test_block_int8.py", 22),
+        TestFile("quant/test_fused_rms_fp8_group_quant.py", 10),
         TestFile("rl/test_update_weights_from_disk.py", 210),
         TestFile("test_abort.py", 51),
         TestFile("test_chunked_prefill.py", 410),