Fuse rmsnorm and quant

python/sglang/srt/layers/communicator.py

@@ -34,6 +34,7 @@
     is_nsa_enable_prefill_cp,
     nsa_use_prefill_cp,
 )
+from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
 from sglang.srt.layers.dp_attention import (
     attn_tp_all_gather_into_tensor,
     attn_tp_reduce_scatter_tensor,
@@ -77,6 +78,10 @@
 _is_gfx95_supported = is_gfx95_supported()
 _is_npu = is_npu()
 
+if _use_aiter:
+    from aiter import rmsnorm2d_fwd_with_add_dynamicquant as _fused_add_rms_norm_quant
+    from aiter import rmsnorm2d_fwd_with_dynamicquant as _fused_rms_norm_quant
+
 if _use_aiter and _is_gfx95_supported:
     from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
 
@@ -495,6 +500,26 @@ def prepare_attn(
                             res1=None,
                             output_unquantized_inp1=False,
                         )
+                    elif _use_aiter and quant_format == "fp8_per_token":
+                        out_fp8 = torch.empty(
+                            hidden_states.shape,
+                            dtype=fp8_dtype,
+                            device=hidden_states.device,
+                        )
+                        out_scale = torch.empty(
+                            hidden_states.shape[0],
+                            1,
+                            dtype=torch.float32,
+                            device=hidden_states.device,
+                        )
+                        _fused_rms_norm_quant(
+                            out_fp8,
+                            hidden_states,
+                            out_scale,
+                            self.input_layernorm.weight,
+                            self.input_layernorm.variance_epsilon,
+                        )
+                        hidden_states = (hidden_states, out_fp8, out_scale)
 
                     else:
                         hidden_states = self.input_layernorm(hidden_states)
@@ -527,6 +552,32 @@ def prepare_attn(
                             res1=residual,
                             output_unquantized_inp1=False,
                         )
+                    elif _use_aiter and quant_format == "fp8_per_token":
+                        out_fp8 = torch.empty(
+                            hidden_states.shape,
+                            dtype=fp8_dtype,
+                            device=hidden_states.device,
+                        )
+                        residual_out = torch.empty_like(hidden_states)
+                        out_scale = torch.empty(
+                            hidden_states.shape[0],
+                            1,
+                            dtype=torch.float32,
+                            device=hidden_states.device,
+                        )
+                        if post_residual_addition is not None:
+                            residual = residual + post_residual_addition
+                        _fused_add_rms_norm_quant(
+                            out_fp8,
+                            hidden_states,
+                            residual,
+                            residual_out,
+                            out_scale,
+                            self.input_layernorm.weight,
+                            self.input_layernorm.variance_epsilon,
+                        )
+                        residual = residual_out
+                        hidden_states = (residual_out, out_fp8, out_scale)
                     else:
                         hidden_states, residual = self.input_layernorm(
                             hidden_states,

python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py

@@ -230,6 +230,18 @@ def apply_weights(
         x: torch.Tensor,
         bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
+        pre_quantized = None
+        if (
+            isinstance(x, tuple)
+            and len(x) == 3
+            and isinstance(x[0], torch.Tensor)
+            and isinstance(x[1], torch.Tensor)
+            and isinstance(x[2], torch.Tensor)
+        ):
+            bf16_view, fp8_view, scale_view = x
+            pre_quantized = (fp8_view, scale_view)
+            x = bf16_view
+
         if self.weight_block_size is not None:
             return self.w8a8_block_fp8_linear(
                 input=x,
@@ -249,6 +261,7 @@ def apply_weights(
                 bias=bias,
                 use_per_token_if_dynamic=True,
                 compressed_tensor_quant=True,
+                pre_quantized=pre_quantized,
             )
         else:
             return apply_fp8_linear(
@@ -259,4 +272,5 @@ def apply_weights(
                 bias=bias,
                 use_per_token_if_dynamic=True,
                 compressed_tensor_quant=True,
+                pre_quantized=pre_quantized,
             )

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

@@ -1305,11 +1305,15 @@ def apply_fp8_ptpc_linear(
     use_per_token_if_dynamic: bool = False,
     pad_output: Optional[bool] = None,
     compressed_tensor_quant: bool = False,
+    pre_quantized: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
 ) -> torch.Tensor:
     # View input as 2D matrix for fp8 methods
     input_2d = input.view(-1, input.shape[-1])
 
-    q_input, x_scale = aiter.per_token_quant_hip(input_2d, quant_dtype=aiter.dtypes.fp8)
+    if pre_quantized is not None:
+        q_input, x_scale = pre_quantized
+    else:
+        q_input, x_scale = aiter.per_token_quant_hip(input_2d, quant_dtype=aiter.dtypes.fp8)
 
     # per_tensor_weights = (weight_scale.numel() == 1) and weight_scale.dim() < 2
     # per_tensor_activations = (x_scale.numel() == 1) and x_scale.dim() < 2

python/sglang/srt/models/qwen3_5.py

@@ -893,6 +893,14 @@ def __init__(
 
         self.alt_stream = alt_stream
 
+        self._fuse_rmsnorm_quant = (
+            _use_aiter
+            and not is_dp_attention_enabled()
+            and attn_quant_config is not None
+            and hasattr(attn_quant_config, "get_name")
+            and "compressed" in attn_quant_config.get_name()
+        )
+
     def _apply_qk_norm(
         self, q: torch.Tensor, k: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -961,8 +969,15 @@ def forward(
         forward_batch: ForwardBatch,
         **kwargs,
     ):
+        _fuse_quant = (
+            self._fuse_rmsnorm_quant
+            and not forward_batch.forward_mode.is_idle()
+        )
         hidden_states, residual = self.layer_communicator.prepare_attn(
-            hidden_states, residual, forward_batch
+            hidden_states,
+            residual,
+            forward_batch,
+            quant_format="fp8_per_token" if _fuse_quant else "",
         )
 
         if not forward_batch.forward_mode.is_idle():