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 14, 05:30 UTC
[quantization] use channel scales for w4a8 + misc fixes (#23570)
Signed-off-by: czhu-cohere <conway.zhu@cohere.com>
czhu-cohere committed
Aug 26, 2025 at 21:23 UTC
2c2b140ae8c60dc0c38e4d37274fc7106a72c21b
4 files changed
+63
-14
tests/quantization/test_compressed_tensors.py
+40
-4
@@ -14,10 +14,10 @@ from compressed_tensors.quantization import QuantizationType
14
from tests.models.utils import check_logprobs_close
15
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
16
CompressedTensors24, CompressedTensorsLinearMethod,
17
- CompressedTensorsW4A4Fp4, CompressedTensorsW4A16Fp4,
18
- CompressedTensorsW4A16Sparse24, CompressedTensorsW8A8Fp8,
19
- CompressedTensorsW8A8Int8, CompressedTensorsW8A16Fp8,
20
- CompressedTensorsWNA16)
17
+ CompressedTensorsW4A4Fp4, CompressedTensorsW4A8Fp8,
18
+ CompressedTensorsW4A16Fp4, CompressedTensorsW4A16Sparse24,
19
+ CompressedTensorsW8A8Fp8, CompressedTensorsW8A8Int8,
20
+ CompressedTensorsW8A16Fp8, CompressedTensorsWNA16)
21
from vllm.model_executor.layers.quantization.utils.quant_utils import (
22
cutlass_fp4_supported)
23
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
@@ -683,3 +683,39 @@ def test_compressed_tensors_nvfp4(vllm_runner, args):
683
output = llm.generate_greedy("Hello my name is", max_tokens=20)
684
print(output)
685
assert output
686
+
687
+
688
+@pytest.mark.skipif(
689
+ not current_platform.is_cuda()
690
+ or not current_platform.has_device_capability(90),
691
+ reason="W4A8 FP8 is not yet supported on this GPU type.",
692
+)
693
+@pytest.mark.parametrize("args", [
694
+ ("czhu-cohere/TinyLlama-1.1B-Chat-v1.0-W4A8-e2e", CompressedTensorsW4A8Fp8)
695
+])
696
+def test_compressed_tensors_w4a8_fp8(vllm_runner, args):
697
+ model, scheme = args
698
+ with vllm_runner(model, enforce_eager=True) as llm:
699
+
700
+ def check_model(model):
701
+ layer = model.model.layers[0]
702
+
703
+ qkv_proj = layer.self_attn.qkv_proj
704
+ o_proj = layer.self_attn.o_proj
705
+ gate_up_proj = layer.mlp.gate_up_proj
706
+ down_proj = layer.mlp.down_proj
707
+
708
+ for proj in (qkv_proj, o_proj, gate_up_proj, down_proj):
709
+ assert isinstance(proj.quant_method,
710
+ CompressedTensorsLinearMethod)
711
+ assert isinstance(proj.scheme, scheme)
712
+
713
+ assert proj.weight_packed.dtype is torch.int32
714
+ assert proj.weight_scale.dtype is torch.float8_e4m3fn
715
+ assert proj.weight_chan_scale.dtype is torch.float32
716
+ assert proj.scheme.group_size == 128
717
+
718
+ llm.apply_model(check_model)
719
+ output = llm.generate_greedy("Hello my name is", max_tokens=20)
720
+ print(output)
721
+ assert output
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py
+11
-2
@@ -79,7 +79,8 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):
79
act_type=torch.float8_e4m3fn, # always use fp8(e4m3)
80
group_size=self.group_size,
81
zero_points=not self.symmetric,
82
- has_g_idx=self.has_g_idx
82
+ has_g_idx=self.has_g_idx,
83
+ out_type=params_dtype
84
)
85
86
kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
@@ -122,7 +123,7 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):
123
torch.empty(
124
output_size_per_partition,
125
scales_and_zp_size,
125
- dtype=params_dtype,
126
+ dtype=torch.float8_e4m3fn,
127
)
128
}
129
@@ -140,9 +141,17 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):
141
dtype=torch.int64),
142
weight_loader=weight_loader)
143
144
+ # per-channel scales
145
+ weight_chan_scale = ChannelQuantScaleParameter(
146
+ data=torch.empty((output_size_per_partition, 1),
147
+ dtype=torch.float32),
148
+ output_dim=0,
149
+ weight_loader=weight_loader)
150
+
151
layer.register_parameter("weight_packed", weight)
152
layer.register_parameter("weight_scale", weight_scale)
153
layer.register_parameter("weight_shape", weight_shape)
154
+ layer.register_parameter("weight_chan_scale", weight_chan_scale)
155
156
self.kernel = kernel_type(mp_linear_kernel_config,
157
w_q_param_name="weight_packed",
vllm/model_executor/layers/quantization/kernels/mixed_precision/MPLinearKernel.py
+1
@@ -20,6 +20,7 @@ class MPLinearLayerConfig:
20
group_size: int
21
zero_points: bool
22
has_g_idx: bool
23
+ out_type: Optional[torch.dtype] = None
24
25
26
class MPLinearKernel(ABC):
vllm/model_executor/layers/quantization/kernels/mixed_precision/cutlass.py
+11
-8
@@ -60,13 +60,17 @@ class CutlassW4A8LinearKernel(MPLinearKernel):
60
if in_features % 128 or out_features % 128:
61
return False, "K and N must be divisible by 128, got "\
62
f"{c.partition_weight_shape}"
63
+
64
+ if c.out_type != torch.bfloat16:
65
+ return False, "Only bfloat16 output type currently supported"\
66
+ f"got {c.out_type=}"
67
+
68
return True, None
69
70
# note assumes that
71
# `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
72
# `weight_scale` is: {input_dim = 0, output_dim = 1}
73
def process_weights_after_loading(self, layer: torch.nn.Module):
69
- c = self.config
74
75
# TODO(czhu): optimize speed/mem usage
76
def transform_w_q(x):
@@ -86,19 +90,15 @@ class CutlassW4A8LinearKernel(MPLinearKernel):
90
# Encode/reorder weights and pack scales
91
self._transform_param(layer, self.w_q_name, transform_w_q)
92
self._transform_param(layer, self.w_s_name, transform_w_s)
89
-
90
- # TODO(czhu): support loading channel scales
91
- self.w_ch_s = torch.ones((c.partition_weight_shape[1], ),
92
- dtype=torch.float32,
93
- device='cuda')
93
+ self._transform_param(layer, "weight_chan_scale", lambda x: x)
94
95
def apply_weights(self,
96
layer: torch.nn.Module,
97
x: torch.Tensor,
98
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
99
- assert bias is None, "bias not supported by CUTLASS W4A8"
99
c = self.config
100
w_q, w_s, _, _ = self._get_weight_params(layer)
101
+ w_ch_s = layer.weight_chan_scale
102
103
x_2d = x.reshape(-1, x.shape[-1])
104
out_shape = x.shape[:-1] + (c.partition_weight_shape[1], )
@@ -109,6 +109,9 @@ class CutlassW4A8LinearKernel(MPLinearKernel):
109
b_group_scales=w_s,
110
b_group_size=c.group_size,
111
a_token_scales=act_scales,
112
- b_channel_scales=self.w_ch_s)
112
+ b_channel_scales=w_ch_s)
113
+
114
+ if bias is not None:
115
+ output.add_(bias) # In-place add
116
117
return output.reshape(out_shape)