| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for AWQ INT4 W4A8 GEMM pipeline (SGLang kernel migration). |
| 4 | |
| 5 | Part 1: Weight packing tests |
| 6 | - convert_weight_packed_scale_zp correctness |
| 7 | |
| 8 | Part 2: INT4 W4A8 GEMM tests |
| 9 | - int4_scaled_mm_cpu correctness w.r.t. float reference |
| 10 | - Bias, 3D input, various shapes |
| 11 | |
| 12 | Part 3: create_weights shapes |
| 13 | |
| 14 | cmd: |
| 15 | VLLM_CPU_INT4_W4A8=1 python -m pytest tests/kernels/test_awq_int4_to_int8.py -v -s |
| 16 | """ |
| 17 | |
| 18 | import numpy as np |
| 19 | import pytest |
| 20 | import torch |
| 21 | |
| 22 | from vllm._custom_ops import ( |
| 23 | CPUQuantAlgo, |
| 24 | convert_weight_packed_scale_zp, |
| 25 | int4_scaled_mm_cpu, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 28 | pack_cols, |
| 29 | ) |
| 30 | from vllm.platforms import current_platform |
| 31 | |
| 32 | if not current_platform.is_cpu(): |
| 33 | pytest.skip("skipping CPU-only tests", allow_module_level=True) |
| 34 | |
| 35 | |
| 36 | def make_awq_checkpoint_data(K, N, group_size, seed=42): |
| 37 | """Create synthetic AWQ checkpoint data in packed int32 format. |
| 38 | |
| 39 | Returns: |
| 40 | packed_qweight: [K, N//8] int32 (AWQ interleaved + packed) |
| 41 | packed_qzeros: [num_groups, N//8] int32 (AWQ interleaved + packed) |
| 42 | scales: [num_groups, N] float32 |
| 43 | float_ref: [K, N] float32, reference dequantized weights |
| 44 | weight_int4_orig: [K, N] int32, original int4 values (0-15) |
| 45 | zeros_int4_orig: [num_groups, N] int32, original zero points (0-15) |
| 46 | |
| 47 | """ |
| 48 | rng = np.random.RandomState(seed) |
| 49 | num_groups = K // group_size |
| 50 | |
| 51 | weight_int4_orig = torch.from_numpy( |
| 52 | rng.randint(0, 16, size=(K, N)).astype(np.int32) |
| 53 | ) |
| 54 | zeros_int4_orig = torch.from_numpy( |
| 55 | rng.randint(0, 16, size=(num_groups, N)).astype(np.int32) |
| 56 | ) |
| 57 | scales = torch.from_numpy((rng.randn(num_groups, N) * 0.05).astype(np.float32)) |
| 58 | |
| 59 | scales_exp = scales.repeat_interleave(group_size, dim=0) |
| 60 | zeros_exp = zeros_int4_orig.repeat_interleave(group_size, dim=0) |
| 61 | float_ref = (weight_int4_orig.float() - zeros_exp.float()) * scales_exp |
| 62 | |
| 63 | awq_interleave = [0, 2, 4, 6, 1, 3, 5, 7] |
| 64 | weight_interleaved = ( |
| 65 | weight_int4_orig.reshape(-1, 8)[:, awq_interleave].reshape(K, N).contiguous() |
| 66 | ) |
| 67 | packed_qweight = pack_cols(weight_interleaved, 4, K, N) |
| 68 | |
| 69 | zeros_interleaved = ( |
| 70 | zeros_int4_orig.reshape(-1, 8)[:, awq_interleave] |
| 71 | .reshape(num_groups, N) |
| 72 | .contiguous() |
| 73 | ) |
| 74 | packed_qzeros = pack_cols(zeros_interleaved, 4, num_groups, N) |
| 75 | |
| 76 | return ( |
| 77 | packed_qweight, |
| 78 | packed_qzeros, |
| 79 | scales, |
| 80 | float_ref, |
| 81 | weight_int4_orig, |
| 82 | zeros_int4_orig, |
| 83 | ) |
| 84 | |
| 85 | |
| 86 | class TestConvertWeightPackedScaleZp: |
| 87 | """Tests for convert_weight_packed_scale_zp weightpacking.""" |
| 88 | |
| 89 | @pytest.mark.parametrize( |
| 90 | "K,N,group_size", |
| 91 | [ |
| 92 | (128, 128, 128), |
| 93 | (256, 256, 128), |
| 94 | (512, 256, 64), |
| 95 | ], |
| 96 | ) |
| 97 | def test_packing_output_shapes(self, K, N, group_size): |
| 98 | """Packed outputs should have expected shapes.""" |
| 99 | (packed_qweight, packed_qzeros, scales, _, _, _) = make_awq_checkpoint_data( |
| 100 | K, N, group_size |
| 101 | ) |
| 102 | |
| 103 | blocked_w, blocked_zp, blocked_s = convert_weight_packed_scale_zp( |
| 104 | packed_qweight, |
| 105 | packed_qzeros, |
| 106 | scales, |
| 107 | CPUQuantAlgo.AWQ, |
| 108 | ) |
| 109 | |
| 110 | block_n = 32 |
| 111 | Nc = N // block_n |
| 112 | |
| 113 | assert blocked_w.dim() >= 2, ( |
| 114 | f"blocked_w should have >= 2 dims, got {blocked_w.dim()}" |
| 115 | ) |
| 116 | assert blocked_s.size(0) == Nc, ( |
| 117 | f"Expected Nc={Nc} scale blocks, got {blocked_s.size(0)}" |
| 118 | ) |
| 119 | assert blocked_zp.size(0) == Nc, ( |
| 120 | f"Expected Nc={Nc} qzeros blocks, got {blocked_zp.size(0)}" |
| 121 | ) |
| 122 | |
| 123 | print( |
| 124 | f" [PASS] packing shapes K={K}, N={N}, gs={group_size}: " |
| 125 | f"blocked_w={list(blocked_w.shape)}, " |
| 126 | f"blocked_s={list(blocked_s.shape)}, blocked_zp={list(blocked_zp.shape)}" |
| 127 | ) |
| 128 | |
| 129 | |
| 130 | class TestInt4ScaledMmCpu: |
| 131 | """Tests for int4_scaled_mm_cpu GEMM kernel.""" |
| 132 | |
| 133 | @pytest.mark.parametrize( |
| 134 | "M,K,N,group_size", |
| 135 | [ |
| 136 | (1, 128, 128, 128), |
| 137 | (4, 256, 256, 128), |
| 138 | (16, 512, 256, 64), |
| 139 | (32, 256, 512, 128), |
| 140 | (64, 512, 512, 128), |
| 141 | ], |
| 142 | ) |
| 143 | def test_gemm_vs_float_reference(self, M, K, N, group_size): |
| 144 | """INT4 W4A8 GEMM should approximate float matmul.""" |
| 145 | (packed_qweight, packed_qzeros, scales, float_ref, _, _) = ( |
| 146 | make_awq_checkpoint_data(K, N, group_size) |
| 147 | ) |
| 148 | |
| 149 | blocked_w, blocked_zp, blocked_s = convert_weight_packed_scale_zp( |
| 150 | packed_qweight, |
| 151 | packed_qzeros, |
| 152 | scales, |
| 153 | CPUQuantAlgo.AWQ, |
| 154 | ) |
| 155 | |
| 156 | x = torch.randn(M, K, dtype=torch.bfloat16) |
| 157 | out = int4_scaled_mm_cpu(x, blocked_w, blocked_zp, blocked_s, None) |
| 158 | |
| 159 | ref_out = torch.mm(x.float(), float_ref) |
| 160 | |
| 161 | abs_diff = (out.float() - ref_out).abs() |
| 162 | mean_abs = abs_diff.mean().item() |
| 163 | pct95 = torch.quantile(abs_diff, 0.95).item() |
| 164 | ref_mag = ref_out.abs().mean().item() + 1e-6 |
| 165 | mean_rel = mean_abs / ref_mag |
| 166 | |
| 167 | assert mean_rel < 0.05, ( |
| 168 | f"Mean relative error {mean_rel:.4f} exceeds 5% threshold" |
| 169 | ) |
| 170 | assert pct95 < ref_mag * 0.15, ( |
| 171 | f"95th-pctile abs_diff {pct95:.4f} exceeds 15% of ref magnitude" |
| 172 | ) |
| 173 | print(f" [PASS] INT4 GEMM correct: M={M}, K={K}, N={N}") |
| 174 | |
| 175 | @pytest.mark.parametrize("M", [1, 8, 32]) |
| 176 | def test_gemm_with_bias(self, M): |
| 177 | """INT4 W4A8 GEMM with bias should match reference.""" |
| 178 | K, N, group_size = 256, 128, 128 |
| 179 | (packed_qweight, packed_qzeros, scales, float_ref, _, _) = ( |
| 180 | make_awq_checkpoint_data(K, N, group_size) |
| 181 | ) |
| 182 | |
| 183 | blocked_w, blocked_zp, blocked_s = convert_weight_packed_scale_zp( |
| 184 | packed_qweight, |
| 185 | packed_qzeros, |
| 186 | scales, |
| 187 | CPUQuantAlgo.AWQ, |
| 188 | ) |
| 189 | |
| 190 | bias = torch.randn(N, dtype=torch.float32) |
| 191 | x = torch.randn(M, K, dtype=torch.bfloat16) |
| 192 | |
| 193 | out = int4_scaled_mm_cpu(x, blocked_w, blocked_zp, blocked_s, bias) |
| 194 | |
| 195 | ref_out = torch.mm(x.float(), float_ref) + bias |
| 196 | abs_diff = (out.float() - ref_out).abs() |
| 197 | mean_abs = abs_diff.mean().item() |
| 198 | ref_mag = ref_out.abs().mean().item() + 1e-6 |
| 199 | mean_rel = mean_abs / ref_mag |
| 200 | assert mean_rel < 0.05, ( |
| 201 | f"Mean relative error {mean_rel:.4f} with bias exceeds 5%" |
| 202 | ) |
| 203 | print(f" [PASS] INT4 GEMM with bias: M={M}") |
| 204 | |
| 205 | def test_gemm_3d_input(self): |
| 206 | """int4_scaled_mm_cpu flattens 3D [B, S, K] to 2D and reshapes back.""" |
| 207 | K, N, group_size = 256, 128, 128 |
| 208 | (packed_qweight, packed_qzeros, scales, float_ref, _, _) = ( |
| 209 | make_awq_checkpoint_data(K, N, group_size) |
| 210 | ) |
| 211 | |
| 212 | blocked_w, blocked_zp, blocked_s = convert_weight_packed_scale_zp( |
| 213 | packed_qweight, |
| 214 | packed_qzeros, |
| 215 | scales, |
| 216 | CPUQuantAlgo.AWQ, |
| 217 | ) |
| 218 | |
| 219 | B, S = 2, 8 |
| 220 | x_3d = torch.randn(B, S, K, dtype=torch.bfloat16) |
| 221 | x_2d = x_3d.reshape(-1, K) |
| 222 | |
| 223 | out_3d = int4_scaled_mm_cpu(x_3d, blocked_w, blocked_zp, blocked_s, None) |
| 224 | out_2d = int4_scaled_mm_cpu(x_2d, blocked_w, blocked_zp, blocked_s, None) |
| 225 | |
| 226 | assert out_3d.shape == (B, S, N) |
| 227 | torch.testing.assert_close(out_3d.reshape(-1, N), out_2d) |
| 228 | |
| 229 | ref_out = torch.mm(x_2d.float(), float_ref).reshape(B, S, N) |
| 230 | abs_diff = (out_3d.float() - ref_out).abs() |
| 231 | mean_abs = abs_diff.mean().item() |
| 232 | ref_mag = ref_out.abs().mean().item() + 1e-6 |
| 233 | mean_rel = mean_abs / ref_mag |
| 234 | |
| 235 | assert mean_rel < 0.05, f"Mean relative error {mean_rel:.4f} for 3D exceeds 5%" |
| 236 | print(f" [PASS] 3D input [{B},{S},{K}] -> output [{B},{S},{N}]") |
| 237 | |
| 238 | def test_gemm_fp16_input(self): |
| 239 | """INT4 GEMM should also work with fp16 input.""" |
| 240 | K, N, group_size, M = 256, 256, 128, 8 |
| 241 | (packed_qweight, packed_qzeros, scales, float_ref, _, _) = ( |
| 242 | make_awq_checkpoint_data(K, N, group_size) |
| 243 | ) |
| 244 | |
| 245 | blocked_w, blocked_zp, blocked_s = convert_weight_packed_scale_zp( |
| 246 | packed_qweight, |
| 247 | packed_qzeros, |
| 248 | scales, |
| 249 | CPUQuantAlgo.AWQ, |
| 250 | ) |
| 251 | |
| 252 | x = torch.randn(M, K, dtype=torch.float16) |
| 253 | out = int4_scaled_mm_cpu(x, blocked_w, blocked_zp, blocked_s, None) |
| 254 | |
| 255 | ref_out = torch.mm(x.float(), float_ref) |
| 256 | abs_diff = (out.float() - ref_out).abs() |
| 257 | ref_mag = ref_out.abs().mean().item() + 1e-6 |
| 258 | mean_rel = abs_diff.mean().item() / ref_mag |
| 259 | |
| 260 | assert mean_rel < 0.05, ( |
| 261 | f"Mean relative error {mean_rel:.4f} for fp16 exceeds 5%" |
| 262 | ) |
| 263 | print(f" [PASS] fp16 input M={M}, K={K}, N={N}") |
| 264 | |
| 265 | |
| 266 | class TestCreateWeightsUnchanged: |
| 267 | """Create_weights should still produce correct int4 placeholder shapes.""" |
| 268 | |
| 269 | @pytest.mark.parametrize( |
| 270 | "K,N,group_size", |
| 271 | [ |
| 272 | (128, 128, 128), |
| 273 | (256, 256, 128), |
| 274 | (512, 256, 64), |
| 275 | ], |
| 276 | ) |
| 277 | def test_int4_placeholder_shapes(self, K, N, group_size): |
| 278 | """Verify qweight, qzeros, scales shapes.""" |
| 279 | pack_factor = 8 |
| 280 | num_groups = K // group_size |
| 281 | |
| 282 | qweight = torch.empty(K, N // pack_factor, dtype=torch.int32) |
| 283 | qzeros = torch.empty(num_groups, N // pack_factor, dtype=torch.int32) |
| 284 | scales = torch.empty(num_groups, N, dtype=torch.bfloat16) |
| 285 | |
| 286 | assert qweight.shape == (K, N // pack_factor) |
| 287 | assert qzeros.shape == (num_groups, N // pack_factor) |
| 288 | assert scales.shape == (num_groups, N) |
| 289 | print(f" [PASS] create_weights shapes: K={K}, N={N}, gs={group_size}") |