[Misc][Quark] Upstream Quark format to VLLM (#10765)

Signed-off-by: kewang-xlnx <kewang@xilinx.com> Signed-off-by: kewang2 <kewang2@amd.com> Co-authored-by: kewang2 <kewang2@amd.com> Co-authored-by: Michael Goin <michael@neuralmagic.com>

Author: 3 people

tests/quantization/test_quark.py

@@ -0,0 +1,30 @@
+"""Test model set-up and weight loading for quark-quantized models.
+
+Run `pytest tests/quantization/test_quark.py`.
+"""
+
+import torch
+
+from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501
+ QuarkLinearMethod, QuarkW8A8Fp8)
+
+
+def test_quark_fp8(vllm_runner):
+ model_path = "amd/Llama-3.1-8B-Instruct-FP8-KV-Quark-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, QuarkW8A8Fp8)
+
+ if isinstance(qkv_proj.scheme, QuarkW8A8Fp8):
+ assert len(qkv_proj.input_scale.shape) == 0
+ assert qkv_proj.weight.dtype is torch.float8_e4m3fn
+ #assert qkv_proj.weight.dtype is torch.float8_e4m3fnuz
+ assert len(qkv_proj.weight_scale.shape) == 0
+
+ output = llm.generate_greedy("Hello my name is", max_tokens=20)
+ assert output

vllm/config.py

@@ -553,7 +553,7 @@ def _verify_quantization(self) -> None:
  optimized_quantization_methods = [
  "fp8", "marlin", "modelopt", "gptq_marlin_24", "gptq_marlin",
  "awq_marlin", "fbgemm_fp8", "compressed_tensors",
- "compressed-tensors", "experts_int8"
+ "compressed-tensors", "experts_int8", "quark"
  ]
  if self.quantization is not None:
  self.quantization = self.quantization.lower()

vllm/model_executor/layers/linear.py

@@ -32,7 +32,7 @@
  "MarlinLinearMethod", "QQQLinearMethod", "GPTQMarlin24LinearMethod",
  "TPUInt8LinearMethod", "GPTQLinearMethod", "FBGEMMFp8LinearMethod",
  "ModelOptFp8LinearMethod", "IPEXAWQLinearMethod", "IPEXGPTQLinearMethod",
- "HQQMarlinMethod"
+ "HQQMarlinMethod", "QuarkLinearMethod"
 ]
 
 

vllm/model_executor/layers/quantization/__init__.py

@@ -26,6 +26,7 @@
  "experts_int8",
  "neuron_quant",
  "ipex",
+ "quark"
 ]
 
 
@@ -34,6 +35,8 @@ def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:
  raise ValueError(f"Invalid quantization method: {quantization}")
 
  # lazy import to avoid triggering `torch.compile` too early
+ from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
+
  from .aqlm import AQLMConfig
  from .awq import AWQConfig
  from .awq_marlin import AWQMarlinConfig
@@ -79,6 +82,7 @@ def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:
  "experts_int8": ExpertsInt8Config,
  "neuron_quant": NeuronQuantConfig,
  "ipex": IPEXConfig,
+ "quark": QuarkConfig
  }
 
  return method_to_config[quantization]

vllm/model_executor/layers/quantization/base_config.py

@@ -133,3 +133,6 @@ def get_quant_method(self, layer: torch.nn.Module,
  method.
  """
  raise NotImplementedError
+
+ def get_cache_scale(self, name: str) -> Optional[str]:
+ return None

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py

@@ -412,6 +412,22 @@ def get_scheme(
  self._check_scheme_supported(scheme.get_min_capability())
  return scheme
 
+ def get_cache_scale(self, name: str) -> Optional[str]:
+ """
+ Check whether the param name matches the format for k/v cache scales
+ in compressed-tensors. If this is the case, return its equivalent
+ param name expected by vLLM
+
+ :param name: param name
+ :return: matching param name for KV cache scale in vLLM
+ """
+ if name.endswith(".output_scale") and ".k_proj" in name:
+ return name.replace(".k_proj.output_scale", ".attn.k_scale")
+ if name.endswith(".output_scale") and ".v_proj" in name:
+ return name.replace(".v_proj.output_scale", ".attn.v_scale")
+ # If no matches, return None
+ return None
+
  @staticmethod
  def supports_cutlass_24(
  weight_quant: Optional[QuantizationArgs],

vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py

@@ -136,6 +136,10 @@ def triton_scaled_mm(input: torch.Tensor,
  assert N > 0 and K > 0 and M > 0
  assert weight.shape[0] == K
  assert input.dtype == weight.dtype
+
+ scale_a = scale_a.reshape(-1, 1) if scale_a.dim() <= 1 else scale_a
+ scale_b = scale_b.reshape(-1, 1) if scale_b.dim() <= 1 else scale_b
+
  assert scale_a.dtype == scale_b.dtype and scale_a.is_floating_point()
  assert scale_a.shape == torch.Size([1, 1]) or scale_a.shape == torch.Size(
  [M, 1])

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

@@ -133,23 +133,6 @@ def _find_first_match(value: str,
  return None
 
 
-def get_compressed_tensors_cache_scale(name: str) -> Optional[str]:
- """
- Check whether the param name matches the format for k/v cache scales
- in compressed-tensors. If this is the case, return its equivalent
- param name expected by vLLM
-
- :param name: param name
- :return: matching param name for KV cache scale in vLLM
- """
- if name.endswith(".output_scale") and ".k_proj" in name:
- return name.replace(".k_proj.output_scale", ".attn.k_scale")
- if name.endswith(".output_scale") and ".v_proj" in name:
- return name.replace(".v_proj.output_scale", ".attn.v_scale")
- # If no matches, return None
- return None
-
-
 def _is_equal_or_regex_match(value: str,
  target: str,
  check_contains: bool = False) -> bool: