be consistent with origin vllm

Signed-off-by: kewang-xlnx <kewang@xilinx.com>

Author: kewang-xlnx

tests/quantization/test_quark.py

@@ -6,7 +6,7 @@
 import torch
 
 from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501
- QuarkLinearMethod, QuarkW8A8Fp8, QuarkW8A8Int8)
+ QuarkLinearMethod, QuarkW8A8Fp8)
 
 
 def test_quark_fp8(vllm_runner):
@@ -28,21 +28,3 @@ def test_quark_fp8(vllm_runner):
 
  output = llm.generate_greedy("Hello my name is", max_tokens=20)
  assert output
-
-
-def test_quark_int8(vllm_runner):
- model_path = "amd/Llama-3.1-8B-Instruct-w-int8-a-int8-sym-test"
- with vllm_runner(model_path) as llm:
- model = llm.model.llm_engine.model_executor.driver_worker.model_runner.model # noqa: E501
- layer = model.model.layers[0]
-
- qkv_proj = layer.self_attn.qkv_proj
-
- assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
- assert isinstance(qkv_proj.scheme, QuarkW8A8Int8)
-
- if isinstance(qkv_proj.scheme, QuarkW8A8Int8):
- assert qkv_proj.weight.dtype is torch.int8
-
- output = llm.generate_greedy("Hello my name is", max_tokens=20)
- assert output

vllm/model_executor/layers/quantization/quark/quark_moe.py

@@ -4,13 +4,15 @@
 
 import vllm.model_executor.layers.fused_moe # noqa
 from vllm import _custom_ops as ops
+from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEMethodBase,
  FusedMoeWeightScaleSupported)
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
  all_close_1d, normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize)
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
-from vllm.utils import print_warning_once
+
+logger = init_logger(__name__)
 
 __all__ = ["QuarkMoEMethod", "QuarkW8A8Fp8MoEMethod"]
 
@@ -127,7 +129,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
  "activation scales are None.")
  if (not all_close_1d(layer.w13_input_scale)
  or not all_close_1d(layer.w2_input_scale)):
- print_warning_once(
+ logger.warning_once(
  "Found input_scales that are not equal for "
  "fp8 MoE layer. Using the maximum across experts "
  "for each layer. ")

vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py

@@ -1,12 +1,11 @@
-from typing import Callable, List, Optional
+from typing import Callable, List, Optional, Set
 
 import torch
-from torch.nn import Parameter
 
 from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
+ ScaledMMLinearLayerConfig, choose_scaled_mm_linear_kernel)
 from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
- apply_int8_linear, convert_to_channelwise)
 from vllm.model_executor.parameter import (BasevLLMParameter,
  ChannelQuantScaleParameter,
  ModelWeightParameter,
@@ -16,6 +15,7 @@
 
 
 class QuarkW8A8Int8(QuarkScheme):
+ _kernel_backends_being_used: Set[str] = set()
 
  def __init__(self, qscheme: str, is_static_input_scheme: Optional[bool],
  input_symmetric: Optional[bool]):
@@ -28,77 +28,25 @@ def get_min_capability(cls) -> int:
  # turing and up
  return 75
 
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
- # WEIGHT
- # Cutlass kernels need transposed weight.
- weight = layer.weight
- layer.weight = Parameter(weight.t(), requires_grad=False)
-
- # WEIGHT SCALE
- # Cutlass kernels support only per-tensor and per-channel.
- # If we have a fused module (QKV, MLP) with per tensor scales (thus N
- # scales being passed to the kernel), convert to the per-channel case.
- is_fused_module = len(self.logical_widths) > 1
- if is_fused_module and self.qscheme == "per_tensor":
- ws_channelwise = convert_to_channelwise(layer.weight_scale,
- self.logical_widths)
- layer.weight_scale = Parameter(ws_channelwise, requires_grad=False)
- else:
- layer.weight_scale = Parameter(layer.weight_scale.data,
- requires_grad=False)
- layer.weight_zero_point = None
-
- # INPUT SCALE
- if self.is_static_input_scheme:
- if self.input_symmetric:
- layer.input_scale = Parameter(layer.input_scale.max(),
- requires_grad=False)
- layer.input_zero_point = None
- else:
- # reconstruct the ranges
- int8_traits = torch.iinfo(torch.int8)
- azps = layer.input_zero_point.to(dtype=torch.int32)
- range_max = (layer.input_scale *
- (int8_traits.max - azps)).max()
- range_min = (layer.input_scale *
- (int8_traits.min - azps)).min()
-
- scale = (range_max - range_min) / (int8_traits.max -
- int8_traits.min)
- layer.input_scale = Parameter(scale, requires_grad=False)
-
- # AZP loaded as int8 but used as int32
- azp = (int8_traits.min -
- range_min / scale).to(dtype=torch.int32)
- layer.input_zero_point = Parameter(azp, requires_grad=False)
-
- else:
- layer.input_scale = None
- layer.input_zero_point = None
-
- # azp_adj is the AZP adjustment term, used to account for weights.
- # It does not depend on scales or azp, so it is the same for
- # static and dynamic quantization.
- # For more details, see csrc/quantization/cutlass_w8a8/Epilogues.md
- # https://github.com/vllm-project/vllm/blob/8d59dbb00044a588cab96bcdc028006ed922eb06/csrc/quantization/cutlass_w8a8/Epilogues.md
- if not self.input_symmetric:
- azp_adj = layer.weight.sum(dim=0, keepdim=True, dtype=torch.int32)
- if self.is_static_input_scheme:
- # cutlass_w8a8 requires azp to be folded into azp_adj
- # in the per-tensor case
- azp_adj = layer.input_zero_point * azp_adj
-
- layer.azp_adj = azp_adj
- else:
- layer.azp_adj = None
-
  def create_weights(self, layer: torch.nn.Module,
  output_partition_sizes: List[int],
  input_size_per_partition: int,
  params_dtype: torch.dtype, weight_loader: Callable,
  **kwargs):
  self.logical_widths = output_partition_sizes
 
+ scaled_mm_linear_kernel_config = ScaledMMLinearLayerConfig(
+ is_channelwise=(self.qscheme == "per_channel"),
+ is_static_input_scheme=(self.is_static_input_scheme is True),
+ input_symmetric=(self.input_symmetric is True))
+
+ kernel_type = choose_scaled_mm_linear_kernel(
+ scaled_mm_linear_kernel_config)
+
+ if kernel_type.__name__ not in self._kernel_backends_being_used:
+ logger.info("Using %s for QuarkW8A8Int8", kernel_type.__name__)
+ self._kernel_backends_being_used.add(kernel_type.__name__)
+
  # WEIGHT
  weight = ModelWeightParameter(data=torch.empty(
  sum(output_partition_sizes),
@@ -117,22 +65,12 @@ def create_weights(self, layer: torch.nn.Module,
  dtype=torch.float32),
  output_dim=0,
  weight_loader=weight_loader)
- weight_zero_point = ChannelQuantScaleParameter(
- data=torch.zeros((sum(output_partition_sizes), 1),
- dtype=torch.int8),
- output_dim=0,
- weight_loader=weight_loader)
  else:
  assert self.qscheme == "per_tensor"
  weight_scale = PerTensorScaleParameter(data=torch.empty(
  len(output_partition_sizes), dtype=torch.float32),
  weight_loader=weight_loader)
- weight_zero_point = PerTensorScaleParameter(
- data=torch.zeros(len(output_partition_sizes),
- dtype=torch.int8),
- weight_loader=weight_loader)
  layer.register_parameter("weight_scale", weight_scale)
- layer.register_parameter("weight_zero_point", weight_zero_point)
 
  # INPUT SCALE
  if self.is_static_input_scheme:
@@ -142,24 +80,26 @@ def create_weights(self, layer: torch.nn.Module,
  layer.register_parameter("input_scale", input_scale)
 
  if not self.input_symmetric:
- # Note: compressed-tensors stores the zp using the same dtype
+ # Note: quark stores the zp using the same dtype
  # as the weights
  # AZP loaded as int8 but used as int32
  input_zero_point = BasevLLMParameter(
  data=torch.empty(1, dtype=torch.int8),
  weight_loader=weight_loader)
- else:
- input_zero_point = BasevLLMParameter(
- data=torch.zeros(1, dtype=torch.int8),
- weight_loader=weight_loader)
- layer.register_parameter("input_zero_point", input_zero_point)
+ layer.register_parameter("input_zero_point", input_zero_point)
+
+ self.kernel = kernel_type(c=scaled_mm_linear_kernel_config,
+ w_q_param_name="weight",
+ w_s_param_name="weight_scale",
+ i_s_param_name="input_scale",
+ i_zp_param_name="input_zero_point",
+ azp_adj_param_name="azp_adj")
+
+ # Checkpoints are serialized in quark format, which is
+ # different from the format the kernel may want. Handle repacking here.
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+ self.kernel.process_weights_after_loading(layer)
 
  def apply_weights(self, layer: torch.nn.Module, x: torch.Tensor,
  bias: Optional[torch.Tensor]) -> torch.Tensor:
- return apply_int8_linear(input=x,
- weight=layer.weight,
- weight_scale=layer.weight_scale,
- input_scale=layer.input_scale,
- input_zero_point=layer.input_zero_point,
- azp_adj=layer.azp_adj,
- bias=bias)
+ return self.kernel.apply_weights(layer, x, bias)