diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index bf65192a8f3b..e727843fbeee 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/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, diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py index 1032271e73de..d443528c33ff 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +++ b/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, ) diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 2339cd636eef..ae91f04d87ad 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/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 diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 234b4e827e03..e2eceb4b9020 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/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():