A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 03:30 UTC
[Bugfix] fix quark ptpc (#20251)
Signed-off-by: Haoyang Li <Haoyang.Li@amd.com> Co-authored-by: Haoyang Li <307790822@qq.com>
li haoyang committed
Jun 30, 2025 at 21:24 UTC
1c50e100a9c5dc439aceb9c4437b262d564baa53
2 files changed
+23
-16
vllm/model_executor/layers/quantization/quark/quark.py
+1
-5
@@ -312,11 +312,7 @@ class QuarkConfig(QuantizationConfig):
312
is_fp8_w8a8_supported = self._check_scheme_supported(
313
QuarkW8A8Fp8.get_min_capability(), error=False)
314
if is_fp8_w8a8_supported:
315
- weight_qscheme = cast(str, weight_config.get("qscheme"))
316
- input_static = (input_config is not None and
317
- not cast(bool, input_config.get("is_dynamic")))
318
- return QuarkW8A8Fp8(qscheme=weight_qscheme,
319
- is_static_input_scheme=input_static)
315
+ return QuarkW8A8Fp8(weight_config, input_config)
316
elif self._is_static_tensor_w8a8(weight_config, input_config):
317
weight_qscheme = cast(str, weight_config.get("qscheme"))
318
return QuarkW8A8Int8(qscheme=weight_qscheme,
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py
+22
-11
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import Callable, Optional
4
+from typing import Any, Callable, Optional, cast
5
6
import torch
7
from torch.nn import Parameter
@@ -19,10 +19,19 @@ __all__ = ["QuarkW8A8Fp8"]
19
20
class QuarkW8A8Fp8(QuarkScheme):
21
22
- def __init__(self, qscheme: str, is_static_input_scheme: Optional[bool]):
23
- self.qscheme = qscheme
24
- self.is_static_input_scheme = is_static_input_scheme
25
- self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=False)
22
+ def __init__(self, weight_config: dict[str, Any],
23
+ input_config: Optional[dict[str, Any]]):
24
+ self.weight_qscheme = cast(str, weight_config.get("qscheme"))
25
+ self.is_static_input_scheme: bool = False
26
+ self.input_qscheme: Optional[str] = None
27
+ if input_config is not None:
28
+ self.is_static_input_scheme = not cast(
29
+ bool, input_config.get("is_dynamic"))
30
+ self.input_qscheme = cast(str, input_config.get("qscheme"))
31
+ self.use_per_token_if_dynamic = (not self.is_static_input_scheme \
32
+ and self.input_qscheme == "per_channel")
33
+ self.fp8_linear = Fp8LinearOp(
34
+ use_per_token_if_dynamic=self.use_per_token_if_dynamic)
35
self.out_dtype = torch.get_default_dtype()
36
37
@classmethod
@@ -34,7 +43,7 @@ class QuarkW8A8Fp8(QuarkScheme):
43
# If per tensor, when we have a fused module (e.g. QKV) with per
44
# tensor scales (thus N scales being passed to the kernel),
45
# requantize so we can always run per tensor
37
- if self.qscheme == "per_tensor":
46
+ if self.weight_qscheme == "per_tensor":
47
if current_platform.is_rocm():
48
input_scale = getattr(layer, 'input_scale', None)
49
weight, max_w_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
@@ -58,7 +67,7 @@ class QuarkW8A8Fp8(QuarkScheme):
67
layer.weight_scale = Parameter(max_w_scale, requires_grad=False)
68
69
# If channelwise, scales are already lined up, so just transpose.
61
- elif self.qscheme == "per_channel":
70
+ elif self.weight_qscheme == "per_channel":
71
weight = layer.weight
72
73
if current_platform.is_fp8_fnuz():
@@ -73,13 +82,15 @@ class QuarkW8A8Fp8(QuarkScheme):
82
requires_grad=False)
83
else:
84
weight_scale = layer.weight_scale.data
76
-
85
+ if self.use_per_token_if_dynamic:
86
+ weight_scale = weight_scale.view(-1, 1)
87
layer.weight = Parameter(weight.t(), requires_grad=False)
88
# required by torch.compile to be torch.nn.Parameter
89
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
90
91
else:
82
- raise ValueError(f"Unknown quantization scheme {self.qscheme}")
92
+ raise ValueError(
93
+ f"Unknown quantization scheme {self.weight_qscheme}")
94
95
# INPUT SCALE
96
if self.is_static_input_scheme:
@@ -109,14 +120,14 @@ class QuarkW8A8Fp8(QuarkScheme):
120
# WEIGHT SCALE
121
# TODO: update create_xxx_parameter functions to return
122
# the newly added parameters
112
- if self.qscheme == "per_channel":
123
+ if self.weight_qscheme == "per_channel":
124
weight_scale = ChannelQuantScaleParameter(
125
data=torch.empty((sum(output_partition_sizes)),
126
dtype=torch.float32),
127
output_dim=0,
128
weight_loader=weight_loader)
129
else:
119
- assert self.qscheme == "per_tensor"
130
+ assert self.weight_qscheme == "per_tensor"
131
weight_scale = PerTensorScaleParameter(data=torch.empty(
132
len(output_partition_sizes), dtype=torch.float32),
133
weight_loader=weight_loader)