update linear interface

Author: k50112113

vllm/model_executor/layers/linear.py

@@ -391,6 +391,12 @@ def forward(
     ) -> 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
+        from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
+        if isinstance(self.quant_method, Fp8LinearMethod):
+            output = self.quant_method.apply(self, x, bias, x_quant_scales=x_quant_scales)
+        else:
+            assert x_quant_scales is None, f"x_quant_scales input is not supported for {self.quant_method.__class__}"
+            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"
@@ -611,13 +617,18 @@ def weight_loader_v2(self, param: Parameter, loaded_weight: torch.Tensor):
         param.load_column_parallel_weight(loaded_weight=loaded_weight)
 
     def forward(
-        self, input_
+        self, input_, 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
 
         # Matrix multiply.
         assert self.quant_method is not None
-        output_parallel = self.quant_method.apply(self, input_, bias)
+        from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
+        if isinstance(self.quant_method, Fp8LinearMethod):
+            output_parallel = self.quant_method.apply(self, input_, bias, x_quant_scales=x_quant_scales)
+        else:
+            assert x_quant_scales is None, f"x_quant_scales input is not supported for {self.quant_method.__class__}"
+            output_parallel = self.quant_method.apply(self, input_, bias)
         if self.gather_output:
             # All-gather across the partitions.
             output = tensor_model_parallel_all_gather(output_parallel)
@@ -1386,7 +1397,8 @@ def weight_loader_v2(self, param: BasevLLMParameter,
         param.load_row_parallel_weight(loaded_weight=loaded_weight)
 
     def forward(
-        self, input_
+        self, input_,
+        x_quant_scales = None
     ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
         if self.input_is_parallel:
             input_parallel = input_
@@ -1401,9 +1413,12 @@ def forward(
         # Only fuse bias add into GEMM for rank 0 (this ensures that
         # bias will not get added more than once in TP>1 case)
         bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
-        output_parallel = self.quant_method.apply(self,
-                                                  input_parallel,
-                                                  bias=bias_)
+        from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
+        if isinstance(self.quant_method, Fp8LinearMethod):
+            output_parallel = self.quant_method.apply(self, input_parallel, bias_, x_quant_scales=x_quant_scales)
+        else:
+            assert x_quant_scales is None, f"x_quant_scales input is not supported for {self.quant_method.__class__}"
+            output_parallel = self.quant_method.apply(self, input_parallel, bias_)
         if self.reduce_results and self.tp_size > 1:
             output = tensor_model_parallel_all_reduce(output_parallel)
         else: