| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the marlin kernel. |
| 4 | |
| 5 | Run `pytest tests/kernels/quantization/test_marlin_gemm.py`. |
| 6 | """ |
| 7 | |
| 8 | import itertools |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | |
| 13 | from tests.kernels.utils import opcheck |
| 14 | from tests.quantization.utils import is_quant_method_supported |
| 15 | from vllm import _custom_ops as ops |
| 16 | from vllm.model_executor.layers.quantization.utils.int8_utils import ( |
| 17 | per_token_quant_int8, |
| 18 | ) |
| 19 | from vllm.model_executor.layers.quantization.utils.marlin_utils import ( |
| 20 | get_marlin_workspace, |
| 21 | marlin_make_empty, |
| 22 | marlin_make_workspace_new, |
| 23 | marlin_permute_bias, |
| 24 | query_marlin_supported_quant_types, |
| 25 | ) |
| 26 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( |
| 27 | rand_marlin_weight_mxfp4_like, |
| 28 | rand_marlin_weight_nvfp4_like, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( |
| 31 | marlin_quant_fp8_torch, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.quantization.utils.marlin_utils_test import ( |
| 34 | awq_marlin_quantize, |
| 35 | get_weight_perm, |
| 36 | marlin_quantize, |
| 37 | marlin_weights, |
| 38 | ) |
| 39 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 40 | awq_pack, |
| 41 | gptq_pack, |
| 42 | gptq_quantize_weights, |
| 43 | quantize_weights, |
| 44 | ) |
| 45 | from vllm.platforms import current_platform |
| 46 | from vllm.scalar_type import scalar_types |
| 47 | from vllm.v1.worker.workspace import current_workspace_manager |
| 48 | |
| 49 | if current_platform.is_rocm(): |
| 50 | pytest.skip( |
| 51 | "These tests require marlin, which is not supported on ROCm.", |
| 52 | allow_module_level=True, |
| 53 | ) |
| 54 | |
| 55 | USE_ATOMIC_ADD_OPTS = [False, True] |
| 56 | USE_FP32_REDUCE_OPTS = [True] |
| 57 | |
| 58 | MARLIN_K_CHUNKS = [128] |
| 59 | MARLIN_N_CHUNKS = [64, 256] |
| 60 | |
| 61 | MARLIN_REPACK_NK_FACTORS = [ |
| 62 | (4, 8), |
| 63 | (7, 5), |
| 64 | (13, 11), |
| 65 | ] |
| 66 | |
| 67 | MNK_FACTORS = [ |
| 68 | (1, 1, 1), |
| 69 | (1, 4, 8), |
| 70 | (26, 37, 13), |
| 71 | (257, 13, 11), |
| 72 | ] |
| 73 | |
| 74 | DTYPES = [torch.float16, torch.bfloat16] |
| 75 | |
| 76 | DENSE_MARLIN_QUANT_TEST_CONFIGS = [ |
| 77 | # AWQ-INT4 |
| 78 | {"b_type": scalar_types.uint4, "group_blocks": [-1, 2, 4, 8]}, |
| 79 | # GPTQ-INT4 |
| 80 | { |
| 81 | "b_type": scalar_types.uint4b8, |
| 82 | "group_blocks": [-1, 2, 4, 8], |
| 83 | }, |
| 84 | # GPTQ-INT8 |
| 85 | { |
| 86 | "b_type": scalar_types.uint8b128, |
| 87 | "group_blocks": [-1, 2, 4, 8], |
| 88 | }, |
| 89 | # FP8 |
| 90 | {"b_type": scalar_types.float8_e4m3fn, "group_blocks": [-1, 8]}, |
| 91 | # NVFP4 |
| 92 | {"b_type": scalar_types.float4_e2m1f, "group_blocks": [1]}, |
| 93 | # MXFP4 |
| 94 | { |
| 95 | "a_type": [scalar_types.bfloat16], |
| 96 | "b_type": scalar_types.float4_e2m1f, |
| 97 | "group_blocks": [2], |
| 98 | }, |
| 99 | # AWQ-INT4 with INT8 activation |
| 100 | { |
| 101 | "a_type": [scalar_types.int8], |
| 102 | "b_type": scalar_types.uint4, |
| 103 | "group_blocks": [-1, 2, 4, 8], |
| 104 | }, |
| 105 | # GPTQ-INT4 with INT8 activation |
| 106 | { |
| 107 | "a_type": [scalar_types.int8], |
| 108 | "b_type": scalar_types.uint4b8, |
| 109 | "group_blocks": [-1, 2, 4, 8], |
| 110 | }, |
| 111 | # GPTQ-INT4 with FP8 activation |
| 112 | { |
| 113 | "a_type": [scalar_types.float8_e4m3fn], |
| 114 | "b_type": scalar_types.uint4b8, |
| 115 | "group_blocks": [-1, 2, 4, 8], |
| 116 | }, |
| 117 | # AWQ-INT4 with FP8 activation |
| 118 | { |
| 119 | "a_type": [scalar_types.float8_e4m3fn], |
| 120 | "b_type": scalar_types.uint4, |
| 121 | "group_blocks": [-1, 2, 4, 8], |
| 122 | }, |
| 123 | # MXFP4 with FP8 activation |
| 124 | { |
| 125 | "a_type": [scalar_types.float8_e4m3fn], |
| 126 | "b_type": scalar_types.float4_e2m1f, |
| 127 | "c_type": [scalar_types.bfloat16], |
| 128 | "group_blocks": [2], |
| 129 | }, |
| 130 | ] |
| 131 | |
| 132 | |
| 133 | def compute_max_diff(output, output_ref): |
| 134 | return torch.mean(torch.abs(output - output_ref)) / torch.mean( |
| 135 | torch.abs(output_ref) |
| 136 | ) |
| 137 | |
| 138 | |
| 139 | def rand_data(shape, dtype=torch.float16): |
| 140 | return torch.randn(shape, dtype=dtype, device="cuda") |
| 141 | |
| 142 | |
| 143 | @pytest.mark.skipif( |
| 144 | not is_quant_method_supported("gptq_marlin"), |
| 145 | reason="Marlin is not supported on this GPU type.", |
| 146 | ) |
| 147 | def test_marlin_int4_fp8_preprocess_without_zp(): |
| 148 | qweight_unpacked = torch.randint( |
| 149 | 0, 16, size=(2048, 2048), dtype=torch.int32, device="cuda" |
| 150 | ) |
| 151 | qweight_packed = qweight_unpacked[:, ::2] * 16 + qweight_unpacked[:, 1::2] |
| 152 | qweight_packed = qweight_packed.to(torch.int8).view(torch.int32) |
| 153 | |
| 154 | cuda_res = ops.marlin_int4_fp8_preprocess(qweight_packed) |
| 155 | |
| 156 | torch_res = torch.where( |
| 157 | qweight_unpacked >= 8, qweight_unpacked - 8, 15 - qweight_unpacked |
| 158 | ) |
| 159 | torch_res = torch_res[:, ::2] * 16 + torch_res[:, 1::2] |
| 160 | torch_res = torch_res.to(torch.int8).view(torch.int32) |
| 161 | |
| 162 | assert (cuda_res == torch_res).all() |
| 163 | |
| 164 | |
| 165 | @pytest.mark.skipif( |
| 166 | not is_quant_method_supported("gptq_marlin"), |
| 167 | reason="Marlin is not supported on this GPU type.", |
| 168 | ) |
| 169 | def test_marlin_int4_fp8_preprocess_awq(): |
| 170 | group_size = 128 |
| 171 | |
| 172 | qweight_unpacked = torch.randint( |
| 173 | 0, 16, size=(2048, 2048), dtype=torch.int32, device="cuda" |
| 174 | ) |
| 175 | qzeros_unpacked = torch.randint( |
| 176 | 0, 16, size=(2048 // group_size, 2048), dtype=torch.int32, device="cuda" |
| 177 | ) |
| 178 | |
| 179 | qweight_packed = qweight_unpacked[:, ::2] * 16 + qweight_unpacked[:, 1::2] |
| 180 | qweight_packed = qweight_packed.to(torch.int8).view(torch.int32) |
| 181 | qzeros_packed = qzeros_unpacked[:, ::2] * 16 + qzeros_unpacked[:, 1::2] |
| 182 | qzeros_packed = qzeros_packed.to(torch.int8).view(torch.int32) |
| 183 | |
| 184 | cuda_res = ops.marlin_int4_fp8_preprocess(qweight_packed, qzeros_packed) |
| 185 | |
| 186 | repeated_zp = qzeros_unpacked.repeat_interleave(group_size, 0) |
| 187 | torch_res = qweight_unpacked - repeated_zp |
| 188 | torch_res[torch_res < 0] = 15 - qweight_unpacked[torch_res < 0] |
| 189 | torch_res = torch_res[:, ::2] * 16 + torch_res[:, 1::2] |
| 190 | torch_res = torch_res.to(torch.int8).view(torch.int32) |
| 191 | |
| 192 | assert (cuda_res == torch_res).all() |
| 193 | |
| 194 | |
| 195 | @pytest.mark.skipif( |
| 196 | not is_quant_method_supported("gptq_marlin"), |
| 197 | reason="Marlin is not supported on this GPU type.", |
| 198 | ) |
| 199 | @pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS) |
| 200 | @pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS) |
| 201 | @pytest.mark.parametrize("quant_type", query_marlin_supported_quant_types(False, False)) |
| 202 | @pytest.mark.parametrize("is_a_8bit", [True, False]) |
| 203 | @pytest.mark.parametrize("nk_factors", MARLIN_REPACK_NK_FACTORS) |
| 204 | def test_gptq_marlin_repack(k_chunk, n_chunk, quant_type, is_a_8bit, nk_factors): |
| 205 | n_factor, k_factor = nk_factors |
| 206 | |
| 207 | size_k = k_chunk * k_factor |
| 208 | size_n = n_chunk * n_factor |
| 209 | group_size = 128 |
| 210 | |
| 211 | # Normalize group_size |
| 212 | if group_size == -1: |
| 213 | group_size = size_k |
| 214 | assert group_size <= size_k |
| 215 | |
| 216 | # Create input |
| 217 | b_weight = rand_data((size_k, size_n)) |
| 218 | |
| 219 | # Quantize |
| 220 | _, q_w, _ = gptq_quantize_weights(b_weight, quant_type, group_size) |
| 221 | |
| 222 | # Pack to GPTQ format |
| 223 | q_w_gptq = gptq_pack(q_w, quant_type.size_bits, size_k, size_n) |
| 224 | |
| 225 | # Pack to Marlin format |
| 226 | weight_perm = get_weight_perm(quant_type.size_bits, is_a_8bit) |
| 227 | marlin_q_w_1 = marlin_weights( |
| 228 | q_w, size_k, size_n, quant_type.size_bits, weight_perm, is_a_8bit |
| 229 | ) |
| 230 | |
| 231 | opcheck( |
| 232 | torch.ops._C.gptq_marlin_repack, |
| 233 | (q_w_gptq, size_k, size_n, quant_type.size_bits, is_a_8bit), |
| 234 | ) |
| 235 | |
| 236 | # Run Marlin repack GPU kernel |
| 237 | marlin_q_w_2 = ops.gptq_marlin_repack( |
| 238 | q_w_gptq, size_k, size_n, quant_type.size_bits, is_a_8bit |
| 239 | ) |
| 240 | torch.accelerator.synchronize() |
| 241 | |
| 242 | torch.testing.assert_close(marlin_q_w_1, marlin_q_w_2) |
| 243 | |
| 244 | |
| 245 | @pytest.mark.skipif( |
| 246 | not is_quant_method_supported("gptq_marlin"), |
| 247 | reason="Marlin is not supported on this GPU type.", |
| 248 | ) |
| 249 | @pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS) |
| 250 | @pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS) |
| 251 | @pytest.mark.parametrize("quant_type", query_marlin_supported_quant_types(True)) |
| 252 | @pytest.mark.parametrize("is_a_8bit", [True, False]) |
| 253 | @pytest.mark.parametrize("nk_factors", MARLIN_REPACK_NK_FACTORS) |
| 254 | def test_awq_marlin_repack(k_chunk, n_chunk, quant_type, is_a_8bit, nk_factors): |
| 255 | n_factor, k_factor = nk_factors |
| 256 | |
| 257 | size_k = k_chunk * k_factor |
| 258 | size_n = n_chunk * n_factor |
| 259 | |
| 260 | group_size = 128 |
| 261 | |
| 262 | # Create input |
| 263 | b_weight = rand_data((size_k, size_n)) |
| 264 | |
| 265 | # Quantize |
| 266 | w_ref, q_w, s, zp = quantize_weights( |
| 267 | b_weight, quant_type, group_size, zero_points=True |
| 268 | ) |
| 269 | |
| 270 | # Pack to AWQ format |
| 271 | q_w_awq = awq_pack(q_w, quant_type.size_bits, size_k, size_n) |
| 272 | |
| 273 | # Pack to Marlin format |
| 274 | weight_perm = get_weight_perm(quant_type.size_bits, is_a_8bit) |
| 275 | marlin_q_w_1 = marlin_weights( |
| 276 | q_w, size_k, size_n, quant_type.size_bits, weight_perm, is_a_8bit |
| 277 | ) |
| 278 | |
| 279 | opcheck( |
| 280 | torch.ops._C.awq_marlin_repack, |
| 281 | (q_w_awq, size_k, size_n, quant_type.size_bits, is_a_8bit), |
| 282 | ) |
| 283 | |
| 284 | # Run Marlin repack GPU kernel |
| 285 | marlin_q_w_2 = ops.awq_marlin_repack( |
| 286 | q_w_awq, size_k, size_n, quant_type.size_bits, is_a_8bit |
| 287 | ) |
| 288 | torch.accelerator.synchronize() |
| 289 | |
| 290 | torch.testing.assert_close(marlin_q_w_1, marlin_q_w_2) |
| 291 | |
| 292 | |
| 293 | def marlin_generate_valid_test_cases(): |
| 294 | all_combinations = itertools.product( |
| 295 | DENSE_MARLIN_QUANT_TEST_CONFIGS, |
| 296 | MNK_FACTORS, |
| 297 | MARLIN_N_CHUNKS, |
| 298 | MARLIN_K_CHUNKS, |
| 299 | USE_ATOMIC_ADD_OPTS, |
| 300 | USE_FP32_REDUCE_OPTS, |
| 301 | ) |
| 302 | |
| 303 | def is_invalid( |
| 304 | a_type, |
| 305 | b_type, |
| 306 | c_type, |
| 307 | group_blocks, |
| 308 | size_m, |
| 309 | size_n, |
| 310 | size_k, |
| 311 | use_atomic_add, |
| 312 | use_fp32_reduce, |
| 313 | ): |
| 314 | if use_atomic_add: |
| 315 | if use_fp32_reduce: |
| 316 | return False |
| 317 | if ( |
| 318 | c_type == scalar_types.bfloat16 |
| 319 | and torch.cuda.get_device_capability()[0] < 9 |
| 320 | ): |
| 321 | return False |
| 322 | |
| 323 | group_size = group_blocks if group_blocks <= 0 else group_blocks * 16 |
| 324 | if group_size > 0 and size_k % group_size != 0: |
| 325 | return False |
| 326 | |
| 327 | if group_size == size_k: |
| 328 | return False |
| 329 | |
| 330 | return a_type.size_bits < 16 or a_type is c_type |
| 331 | |
| 332 | cases = [] |
| 333 | for case in all_combinations: |
| 334 | quant_test_config, mnk_factors, n_chunk, k_chunk, *_ = case |
| 335 | size_m = mnk_factors[0] |
| 336 | size_n = mnk_factors[1] * n_chunk |
| 337 | size_k = mnk_factors[2] * k_chunk |
| 338 | |
| 339 | f16_types = [scalar_types.float16, scalar_types.bfloat16] |
| 340 | inner_combinations = itertools.product( |
| 341 | quant_test_config.get("a_type", f16_types), |
| 342 | [quant_test_config["b_type"]], |
| 343 | quant_test_config.get("c_type", f16_types), |
| 344 | quant_test_config["group_blocks"], |
| 345 | ) # type: ignore[call-overload] |
| 346 | |
| 347 | for sub_case in inner_combinations: |
| 348 | if ( |
| 349 | sub_case[0] == scalar_types.float8_e4m3fn |
| 350 | and not current_platform.is_device_capability(89) |
| 351 | and not current_platform.is_device_capability_family(120) |
| 352 | ): |
| 353 | continue |
| 354 | args = sub_case + (size_m, size_n, size_k) + case[4:] |
| 355 | if is_invalid(*args): |
| 356 | cases.append(args) |
| 357 | return cases |
| 358 | |
| 359 | |
| 360 | @pytest.mark.skipif( |
| 361 | not is_quant_method_supported("gptq_marlin"), |
| 362 | reason="Marlin is not supported on this GPU type.", |
| 363 | ) |
| 364 | @pytest.mark.parametrize( |
| 365 | ( |
| 366 | "a_type, b_type, c_type, group_blocks," |
| 367 | "size_m, size_n, size_k," |
| 368 | "use_atomic_add, use_fp32_reduce" |
| 369 | ), |
| 370 | marlin_generate_valid_test_cases(), |
| 371 | ) |
| 372 | def test_marlin_gemm( |
| 373 | a_type, |
| 374 | b_type, |
| 375 | c_type, |
| 376 | group_blocks, |
| 377 | size_m, |
| 378 | size_n, |
| 379 | size_k, |
| 380 | use_atomic_add, |
| 381 | use_fp32_reduce, |
| 382 | ): |
| 383 | has_zp = b_type in [scalar_types.uint4, scalar_types.uint8] |
| 384 | |
| 385 | group_size = group_blocks if group_blocks <= 0 else group_blocks * 16 |
| 386 | |
| 387 | if c_type == scalar_types.float16: |
| 388 | dtype = torch.float16 |
| 389 | elif c_type == scalar_types.bfloat16: |
| 390 | dtype = torch.bfloat16 |
| 391 | else: |
| 392 | raise RuntimeError("unsupported c_type") |
| 393 | |
| 394 | if a_type == scalar_types.int8: |
| 395 | a_dtype = torch.int8 |
| 396 | elif a_type == scalar_types.float8_e4m3fn: |
| 397 | a_dtype = torch.float8_e4m3fn |
| 398 | else: |
| 399 | a_dtype = dtype |
| 400 | |
| 401 | a_input = rand_data((size_m, size_k), dtype=dtype) |
| 402 | b_weight = rand_data((size_k, size_n), dtype=dtype) |
| 403 | |
| 404 | if b_type == scalar_types.float4_e2m1f: |
| 405 | if group_size == 16: |
| 406 | w_ref, marlin_q_w, marlin_s, marlin_s2 = rand_marlin_weight_nvfp4_like( |
| 407 | b_weight.T, group_size, input_dtype=a_dtype |
| 408 | ) |
| 409 | else: |
| 410 | w_ref, marlin_q_w, marlin_s = rand_marlin_weight_mxfp4_like( |
| 411 | b_weight.T, group_size, input_dtype=a_dtype |
| 412 | ) |
| 413 | marlin_s2 = None |
| 414 | |
| 415 | marlin_zp = None |
| 416 | elif b_type == scalar_types.float8_e4m3fn: |
| 417 | w_ref, marlin_q_w, marlin_s = marlin_quant_fp8_torch( |
| 418 | b_weight.T, group_size, input_dtype=a_dtype |
| 419 | ) |
| 420 | marlin_zp = None |
| 421 | marlin_s2 = None |
| 422 | elif has_zp: |
| 423 | w_ref, marlin_q_w, marlin_s, marlin_zp = awq_marlin_quantize( |
| 424 | b_weight, b_type, group_size, input_dtype=a_dtype |
| 425 | ) |
| 426 | marlin_s2 = None |
| 427 | else: |
| 428 | w_ref, marlin_q_w, marlin_s = marlin_quantize( |
| 429 | b_weight, b_type, group_size, input_dtype=a_dtype |
| 430 | ) |
| 431 | |
| 432 | marlin_zp = None |
| 433 | marlin_s2 = None |
| 434 | |
| 435 | workspace = marlin_make_workspace_new(w_ref.device) |
| 436 | |
| 437 | if a_type == scalar_types.int8: |
| 438 | a_input, a_scales = per_token_quant_int8(a_input) |
| 439 | a_input_ref = a_input.to(a_scales.dtype) * a_scales.view(-1, 1) |
| 440 | a_input_ref = a_input_ref.to(dtype) |
| 441 | |
| 442 | if group_size != -1: |
| 443 | a_scales = a_scales / 4096 * marlin_s.max() |
| 444 | a_scales = a_scales.float() |
| 445 | marlin_s = marlin_s / marlin_s.max() * 4096 |
| 446 | marlin_s = marlin_s.round().to(torch.int16).view(dtype) |
| 447 | elif a_type == scalar_types.float8_e4m3fn: |
| 448 | a_input, a_scales = ops.scaled_fp8_quant(a_input, use_per_token_if_dynamic=True) |
| 449 | a_input_ref = a_input.to(a_scales.dtype) * a_scales.view(-1, 1) |
| 450 | a_input_ref = a_input_ref.to(dtype) |
| 451 | else: |
| 452 | assert a_type.size_bits == 16 |
| 453 | a_input_ref = a_input |
| 454 | a_scales = None |
| 455 | |
| 456 | output = torch.empty((size_m, size_n), dtype=dtype, device=a_input.device) |
| 457 | |
| 458 | output = ops.marlin_gemm( |
| 459 | a_input, |
| 460 | output, |
| 461 | marlin_q_w, |
| 462 | None, |
| 463 | marlin_s, |
| 464 | a_scales, |
| 465 | marlin_s2, |
| 466 | marlin_zp, |
| 467 | workspace, |
| 468 | b_type, |
| 469 | a_input.shape[0], |
| 470 | b_weight.shape[1], |
| 471 | a_input.shape[1], |
| 472 | use_atomic_add=use_atomic_add, |
| 473 | use_fp32_reduce=use_fp32_reduce, |
| 474 | is_zp_float=False, |
| 475 | ) |
| 476 | output_ref = torch.matmul(a_input_ref, w_ref) |
| 477 | |
| 478 | max_diff = compute_max_diff(output, output_ref) |
| 479 | assert max_diff < 0.04 |
| 480 | |
| 481 | |
| 482 | def test_marlin_persistent_locks_cuda_graph_streams(workspace_init): |
| 483 | """Compiled GEMMs select stream-local locks at runtime and survive replay.""" |
| 484 | size_m, size_k, size_n = 32, 1024, 2048 |
| 485 | inputs = [rand_data((size_m, size_k)) for _ in range(2)] |
| 486 | weight = rand_data((size_k, size_n)) |
| 487 | quant_type = scalar_types.uint4b8 |
| 488 | w_ref, qweight, scales = marlin_quantize(weight, quant_type, 128) |
| 489 | zeros = marlin_make_empty(scales.device) |
| 490 | streams = [torch.cuda.Stream() for _ in inputs] |
| 491 | graphs = [torch.cuda.CUDAGraph() for _ in inputs] |
| 492 | locks = [] |
| 493 | outputs = [] |
| 494 | |
| 495 | @torch.compile(backend="eager", fullgraph=True) |
| 496 | def run(x): |
| 497 | return torch.ops.vllm.marlin_gemm( |
| 498 | x, |
| 499 | qweight, |
| 500 | None, |
| 501 | scales, |
| 502 | None, |
| 503 | None, |
| 504 | zeros, |
| 505 | None, |
| 506 | quant_type.id, |
| 507 | size_m, |
| 508 | size_n, |
| 509 | size_k, |
| 510 | use_atomic_add=False, |
| 511 | use_fp32_reduce=True, |
| 512 | ) |
| 513 | |
| 514 | for stream, graph, x in zip(streams, graphs, inputs): |
| 515 | stream.wait_stream(torch.cuda.current_stream()) |
| 516 | with torch.cuda.stream(stream): |
| 517 | locks.append(get_marlin_workspace(x.device)) |
| 518 | run(x) |
| 519 | with torch.cuda.graph(graph, stream=stream): |
| 520 | outputs.append(run(x)) |
| 521 | torch.cuda.current_stream().wait_stream(stream) |
| 522 | |
| 523 | assert locks[0].data_ptr() != locks[1].data_ptr() |
| 524 | current_workspace_manager().lock() |
| 525 | for _ in range(3): |
| 526 | for stream, graph, x, lock in zip(streams, graphs, inputs, locks): |
| 527 | with torch.cuda.stream(stream): |
| 528 | assert get_marlin_workspace(x.device) is lock |
| 529 | x.add_(0.1) |
| 530 | graph.replay() |
| 531 | for stream in streams: |
| 532 | torch.cuda.current_stream().wait_stream(stream) |
| 533 | for x, output, lock in zip(inputs, outputs, locks): |
| 534 | assert compute_max_diff(output, x @ w_ref) < 0.04 |
| 535 | assert torch.count_nonzero(lock) == 0 |
| 536 | |
| 537 | |
| 538 | def test_marlin_gemm_subset_input(): |
| 539 | quant_type = scalar_types.uint4b8 |
| 540 | group_size = 128 |
| 541 | |
| 542 | size_m, size_k, size_n = 32, 1024, 2048 |
| 543 | big_m = size_m * 2 |
| 544 | big_k = size_k * 2 |
| 545 | |
| 546 | a_input = rand_data((big_m, big_k))[8 : size_m + 8, 8 : size_k + 8] |
| 547 | b_weight = rand_data((size_k, size_n)) |
| 548 | |
| 549 | w_ref, marlin_q_w, marlin_s = marlin_quantize(b_weight, quant_type, group_size) |
| 550 | |
| 551 | marlin_zp = marlin_make_empty(marlin_s.device) |
| 552 | workspace = marlin_make_workspace_new(a_input.device) |
| 553 | |
| 554 | output = ops.marlin_gemm( |
| 555 | a_input, |
| 556 | None, |
| 557 | marlin_q_w, |
| 558 | None, |
| 559 | marlin_s, |
| 560 | None, |
| 561 | None, |
| 562 | marlin_zp, |
| 563 | workspace, |
| 564 | quant_type, |
| 565 | a_input.shape[0], |
| 566 | b_weight.shape[1], |
| 567 | a_input.shape[1], |
| 568 | use_atomic_add=False, |
| 569 | use_fp32_reduce=True, |
| 570 | is_zp_float=False, |
| 571 | ) |
| 572 | output_ref = torch.matmul(a_input, w_ref) |
| 573 | |
| 574 | torch.accelerator.synchronize() |
| 575 | |
| 576 | max_diff = compute_max_diff(output, output_ref) |
| 577 | |
| 578 | assert max_diff < 0.04 |
| 579 | |
| 580 | |
| 581 | @pytest.mark.parametrize("size_m", [1, 256]) |
| 582 | def test_marlin_gemm_with_bias(size_m): |
| 583 | quant_type = scalar_types.uint4b8 |
| 584 | group_size = 128 |
| 585 | |
| 586 | size_k, size_n = 1024, 2048 |
| 587 | a_input = rand_data((size_m, size_k)) |
| 588 | b_weight = rand_data((size_k, size_n)) |
| 589 | b_bias = rand_data((size_n,)) * 10 |
| 590 | |
| 591 | marlin_bias = marlin_permute_bias(b_bias) |
| 592 | |
| 593 | w_ref, marlin_q_w, marlin_s = marlin_quantize(b_weight, quant_type, group_size) |
| 594 | |
| 595 | marlin_zp = marlin_make_empty(marlin_s.device) |
| 596 | workspace = marlin_make_workspace_new(a_input.device) |
| 597 | |
| 598 | output = ops.marlin_gemm( |
| 599 | a_input, |
| 600 | None, |
| 601 | marlin_q_w, |
| 602 | marlin_bias, |
| 603 | marlin_s, |
| 604 | None, |
| 605 | None, |
| 606 | marlin_zp, |
| 607 | workspace, |
| 608 | quant_type, |
| 609 | a_input.shape[0], |
| 610 | b_weight.shape[1], |
| 611 | a_input.shape[1], |
| 612 | use_atomic_add=False, |
| 613 | use_fp32_reduce=True, |
| 614 | is_zp_float=False, |
| 615 | ) |
| 616 | output_ref = torch.matmul(a_input, w_ref) + b_bias.view(1, -1) |
| 617 | |
| 618 | torch.accelerator.synchronize() |
| 619 | |
| 620 | max_diff = compute_max_diff(output, output_ref) |
| 621 | |
| 622 | assert max_diff < 0.04 |