add fuse rms quant

Author: k50112113

vllm/model_executor/layers/linear.py

@@ -387,11 +387,17 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
         param.data.copy_(loaded_weight)
 
     def forward(
-        self, x: torch.Tensor
+        self, x: torch.Tensor, x_quant_scales: torch.Tensor = None,
     ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
         bias = self.bias if not self.skip_bias_add else None
         assert self.quant_method is not None
-        output = self.quant_method.apply(self, x, bias)
+        # output = self.quant_method.apply(self, x, bias)
+        if isinstance(self.quant_method, UnquantizedLinearMethod):
+            assert x_quant_scales is None, "UnquantizedLinearMethod should not have quantized input"
+            output = self.quant_method.apply(self, x, bias)
+        else:
+            output = self.quant_method.apply(self, x, bias, x_quant_scales=x_quant_scales)
+        
         output_bias = self.bias if self.skip_bias_add else None
         if not self.return_bias:
             return output

vllm/model_executor/layers/quantization/fp8.py

@@ -404,7 +404,8 @@ def process_weights_after_loading(self, layer: Module) -> None:
     def apply(self,
               layer: torch.nn.Module,
               x: torch.Tensor,
-              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+              bias: Optional[torch.Tensor] = None,
+              x_quant_scales: torch.Tensor = None) -> torch.Tensor:
 
         if self.use_marlin:
             return apply_fp8_marlin_linear(
@@ -427,6 +428,7 @@ def apply(self,
                 bias=bias,
                 cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
                 use_aiter_and_is_supported=self.use_aiter_and_is_supported,
+                input_quant_scale=x_quant_scales,
             )
 
         return self.fp8_linear.apply(input=x,

vllm/model_executor/layers/quantization/utils/fp8_utils.py

@@ -115,6 +115,7 @@ def apply_w8a8_block_fp8_linear(
     bias: Optional[torch.Tensor] = None,
     cutlass_block_fp8_supported: bool = CUTLASS_BLOCK_FP8_SUPPORTED,
     use_aiter_and_is_supported: bool = False,
+    input_quant_scale: torch.Tensor = None,
 ) -> torch.Tensor:
     assert input_scale is None
     # View input as 2D matrix for fp8 methods
@@ -135,7 +136,12 @@ def apply_w8a8_block_fp8_linear(
                                        scale_a=x_scale,
                                        scale_b=weight_scale.T)
     else:
-        if use_aiter_and_is_supported:
+        output_dtype = input.dtype
+        if input_quant_scale is not None:
+            q_input = input
+            x_scale = input_quant_scale
+            output_dtype = torch.bfloat16
+        elif use_aiter_and_is_supported:
             # print(f"input_2d.shape: {input_2d.shape}, input_2d.stride: {input_2d.stride()}")
             q_input, x_scale = aiter_per1x128_quant(input_2d.contiguous(), quant_dtype=rocm_aiter.dtypes.fp8)
         else:
@@ -150,11 +156,11 @@ def apply_w8a8_block_fp8_linear(
                                       x_scale,
                                       weight_scale,
                                       block_size,
-                                      output_dtype=input.dtype)
+                                      output_dtype=output_dtype)
 
     if bias is not None:
         output = output + bias
-    return output.to(dtype=input.dtype).view(*output_shape)
+    return output.to(dtype=output_dtype).view(*output_shape)
 
 
 def apply_w8a8_block_fp8_linear_fake(
@@ -166,6 +172,7 @@ def apply_w8a8_block_fp8_linear_fake(
     bias: Optional[torch.Tensor] = None,
     cutlass_block_fp8_supported: bool = CUTLASS_BLOCK_FP8_SUPPORTED,
     use_aiter_and_is_supported: bool = False,
+    input_quant_scale: torch.Tensor = None,
 ) -> torch.Tensor:
     output_shape = [*input.shape[:-1], weight.shape[0]]
     return torch.empty(output_shape, dtype=input.dtype, device=input.device)

vllm/model_executor/models/deepseek_v2.py

@@ -63,6 +63,10 @@
     is_rocm_aiter_fuse_routed_scaling_factor,
 )
 
+from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
+import aiter as rocm_aiter
+rocm_aiter_fp8_dtype = rocm_aiter.dtypes.fp8
+rocm_aiter_fp8_quant_group_size = 128
 
 class DeepseekV2MLP(nn.Module):
 
@@ -484,10 +488,15 @@ def forward(
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
+        hidden_states_quant = None
+        if isinstance(hidden_states, tuple):
+            hidden_states, hidden_states_quant = hidden_states
+
         if self.q_lora_rank is not None:
             # q_c = self.q_a_proj(hidden_states)[0]
             # kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]
-            qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
+            qkv_lora = self.fused_qkv_a_proj(hidden_states, x_quant_scales = hidden_states_quant)[0]
+            # qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
             q_c, kv_lora = qkv_lora.split(
                 [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                 dim=-1,
@@ -576,12 +585,28 @@ def forward(
         residual: Optional[torch.Tensor],
     ) -> torch.Tensor:
         # Self Attention
+        weight = self.input_layernorm.weight
+        eps = self.input_layernorm.variance_epsilon
         if residual is None:
             residual = hidden_states
-            hidden_states = self.input_layernorm(hidden_states)
+            hidden_states, hidden_states_quant = fused_rms_fp8_group_quant(hidden_states, weight, eps, 
+                                                        None, None, eps, 
+                                                        group_size = rocm_aiter_fp8_quant_group_size,
+                                                        dtype_quant=rocm_aiter_fp8_dtype, 
+                                                        res1=None)
         else:
-            hidden_states, residual = self.input_layernorm(
-                hidden_states, residual)
+            (hidden_states, hidden_states_quant), residual = fused_rms_fp8_group_quant(hidden_states, weight, eps, 
+                                                        None, None, eps, 
+                                                        group_size = rocm_aiter_fp8_quant_group_size,
+                                                        dtype_quant=rocm_aiter_fp8_dtype, 
+                                                        res1=residual)
+        hidden_states = (hidden_states, hidden_states_quant)
+        # if residual is None:
+        #     residual = hidden_states
+        #     hidden_states = self.input_layernorm(hidden_states)
+        # else:
+        #     hidden_states, residual = self.input_layernorm(
+        #         hidden_states, residual)
         hidden_states = self.self_attn(
             positions=positions,
             hidden_states=hidden_states,