test_compressed_tensors.py 23 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
"""Test model set-up and weight loading for llmcompressor-quantized models.
3
4
5

Run `pytest tests/quantization/test_compressed_tensors.py`.
"""
6

7
from typing import Optional
8

9
import pytest
10
import torch
11
import os
12

13
from compressed_tensors.quantization import QuantizationType
14

15
from tests.models.utils import check_logprobs_close
16
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (  # noqa: E501
17
18
19
20
    CompressedTensors24, CompressedTensorsLinearMethod,
    CompressedTensorsW4A16Sparse24, CompressedTensorsW8A8Fp8,
    CompressedTensorsW8A8Int8, CompressedTensorsW8A16Fp8,
    CompressedTensorsWNA16)
21
22
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
    sparse_cutlass_supported)
zhuwenwen's avatar
zhuwenwen committed
23
from vllm.platforms import current_platform
24
from ..utils import models_path_prefix
25

26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
# AITER only supports per-channel-per-channel INT8 gemm
# and per-tensor-per-tensor INT8 GEMM.
# It does not support mix precision MM and mix quantization scheme.
ROCM_AITER_SUPPORTED_INT8_MODEL = [
    "neuralmagic/Llama-3.2-1B-quantized.w8a8",
    "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2"
]

# TritonScaledMMLinearKernel only supports symmetric quantization.
ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL = [
    "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change",
    "nm-testing/tinyllama-oneshot-w8-channel-a8-tensor",
    "neuralmagic/Llama-3.2-1B-quantized.w8a8",
    "nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2",
    "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2",
]

43

44
45
46
47
48
49
50
51
@pytest.fixture(scope="function", autouse=True)
def use_v0_only(monkeypatch):
    """
    This module relies on V0 internals, so set VLLM_USE_V1=0.
    """
    monkeypatch.setenv('VLLM_USE_V1', '0')


52
53
@pytest.mark.parametrize(
    "model_args",
54
55
    [
        (
zhuwenwen's avatar
zhuwenwen committed
56
            os.path.join(models_path_prefix, "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change"),
57
58
59
60
61
62
            "tensor",
            QuantizationType.INT,
            2560,
            True,
        ),
        (
zhuwenwen's avatar
zhuwenwen committed
63
            os.path.join(models_path_prefix, "nm-testing/tinyllama-oneshot-w8-channel-a8-tensor"),
64
65
66
67
68
69
            "channel",
            QuantizationType.INT,
            2560,
            True,
        ),
        (
zhuwenwen's avatar
zhuwenwen committed
70
            os.path.join(models_path_prefix, "nm-testing/asym-w8w8-int8-static-per-tensor-tiny-llama"),
71
72
73
74
75
76
77
            "tensor",
            QuantizationType.INT,
            2560,
            False,
        ),
    ],
)
78
def test_compressed_tensors_w8a8_static_setup(vllm_runner, model_args):
79
    model_path, strategy, quant_type, shape_0, is_symmetric = model_args
80
81
82
83
84

    if current_platform.is_rocm(
    ) and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL:
        pytest.skip(f"Skip model {model_path} as it is not support on ROCm.")

85
    with vllm_runner(model_path, enforce_eager=True) as llm:
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

        def check_model(model):
            layer = model.model.layers[0]

            qkv_proj = layer.self_attn.qkv_proj
            o_proj = layer.self_attn.o_proj
            gate_up_proj = layer.mlp.gate_up_proj
            down_proj = layer.mlp.down_proj

            # assert zp for symmetric and asymmetric cases
            def zp_valid(zp: Optional[torch.Tensor]):
                if is_symmetric:
                    return zp is None

                return zp is not None and zp.dtype is torch.int32

            assert zp_valid(qkv_proj.input_zero_point)
            assert zp_valid(o_proj.input_zero_point)
            assert zp_valid(gate_up_proj.input_zero_point)
            assert zp_valid(down_proj.input_zero_point)

            assert isinstance(qkv_proj.quant_method,
                              CompressedTensorsLinearMethod)
            assert isinstance(o_proj.quant_method,
                              CompressedTensorsLinearMethod)
            assert isinstance(gate_up_proj.quant_method,
                              CompressedTensorsLinearMethod)
            assert isinstance(down_proj.quant_method,
                              CompressedTensorsLinearMethod)
            assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8)

            assert qkv_proj.scheme.strategy == strategy
            assert qkv_proj.scheme.is_static_input_scheme
            expected_type = torch.int8

            assert qkv_proj.weight.dtype is expected_type
            assert o_proj.weight.dtype is expected_type
            assert gate_up_proj.weight.dtype is expected_type

            if qkv_proj.scheme.strategy == "tensor":
                # Make sure it is a channelwise buffer
                # After running process_weights_after_loading
                assert len(qkv_proj.weight_scale.shape) == 2
                assert qkv_proj.weight_scale.shape[0] == shape_0
                assert qkv_proj.weight_scale.shape[1] == 1
            assert qkv_proj.weight_scale.dtype is torch.float32
            assert qkv_proj.input_scale.dtype is torch.float32

        llm.apply_model(check_model)
135

136
        output = llm.generate_greedy(["Hello my name is"], max_tokens=20)
137
138
        assert output

139

140
141
142
@pytest.mark.parametrize(
    "model_path",
    [
zhuwenwen's avatar
zhuwenwen committed
143
144
145
146
        os.path.join(models_path_prefix, "neuralmagic/Llama-3.2-1B-quantized.w8a8"),
        os.path.join(models_path_prefix, "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Dynamic-Asym"),
        os.path.join(models_path_prefix, "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Sym"),
        os.path.join(models_path_prefix, "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Asym"),
147
148
    ],
)
149
150
@pytest.mark.parametrize("max_tokens", [32])
@pytest.mark.parametrize("num_logprobs", [10])
151
152
@pytest.mark.parametrize(
    "use_aiter", [True, False] if current_platform.is_rocm() else [False])
153
154
155
156
157
158
159
def test_compressed_tensors_w8a8_logprobs(
    hf_runner,
    vllm_runner,
    example_prompts,
    model_path,
    max_tokens,
    num_logprobs,
160
161
    use_aiter,
    monkeypatch,
162
):
163
164
165
166
167
168
169
170
171
172
173
174

    if current_platform.is_rocm(
    ) and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL:
        pytest.skip(f"Skip model {model_path} as it is not support on ROCm.")

    if use_aiter:
        if model_path not in ROCM_AITER_SUPPORTED_INT8_MODEL:
            pytest.skip(
                f"Skip model {model_path} as it is not support by aiter.")
        # this will enable VLLM_ROCM_USE_AITER_LINEAR
        monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")

175
176
    dtype = "bfloat16"

177
    # skip language translation prompt for the static per tensor asym model
178
    if (model_path ==
zhuwenwen's avatar
zhuwenwen committed
179
            os.path.join(models_path_prefix, "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Asym")
180
        ):  # noqa: E501
181
182
        example_prompts = example_prompts[0:-1]

183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
    with hf_runner(model_path, dtype=dtype) as hf_model:
        hf_outputs = hf_model.generate_greedy_logprobs_limit(
            example_prompts, max_tokens, num_logprobs)

    with vllm_runner(model_path, dtype=dtype) as vllm_model:
        vllm_outputs = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)

    check_logprobs_close(
        outputs_0_lst=hf_outputs,
        outputs_1_lst=vllm_outputs,
        name_0="hf",
        name_1="vllm",
    )

198
199
200
    if current_platform.is_rocm():
        torch.cuda.synchronize()

201

202
def test_compressed_tensors_no_enforce_eager(vllm_runner):
203
    model_path = os.path.join(models_path_prefix, "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change")
204
    with vllm_runner(model_path) as llm:
205
        output = llm.generate_greedy("Hello my name is", max_tokens=20)
206
207
208
        assert output


209
210
211
@pytest.mark.parametrize(
    "model_args",
    [
zhuwenwen's avatar
zhuwenwen committed
212
213
        (os.path.join(models_path_prefix, "nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2"), "tensor"),
        (os.path.join(models_path_prefix, "nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2-asym"), "tensor"),
214
        (
zhuwenwen's avatar
zhuwenwen committed
215
            os.path.join(models_path_prefix, "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2"),
216
217
218
            "channel",
        ),
        (
zhuwenwen's avatar
zhuwenwen committed
219
            os.path.join(models_path_prefix, "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2-asym"),
220
221
222
223
            "channel",
        ),
    ],
)
224
225
226
227
228
229
230
231
@pytest.mark.parametrize(
    "use_aiter", [True, False] if current_platform.is_rocm() else [False])
def test_compressed_tensors_w8a8_dynamic_per_token(
    vllm_runner,
    model_args,
    use_aiter,
    monkeypatch,
):
232
    model_path, strategy = model_args
233
234
235
236
237
238
239
240
241
242
243
244

    if current_platform.is_rocm(
    ) and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL:
        pytest.skip(f"Skip model {model_path} as it is not support on ROCm.")

    if use_aiter:
        if model_path not in ROCM_AITER_SUPPORTED_INT8_MODEL:
            pytest.skip(
                f"Skip model {model_path} as it is not support by aiter.")
        # this will enable VLLM_ROCM_USE_AITER_LINEAR
        monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")

245
    with vllm_runner(model_path, dtype=torch.float16) as llm:
246

247
248
        def check_model(model):
            layer = model.model.layers[0]
249

250
251
252
253
254
255
256
257
            qkv_proj = layer.self_attn.qkv_proj

            assert isinstance(qkv_proj.quant_method,
                              CompressedTensorsLinearMethod)
            assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8)
            assert not qkv_proj.scheme.is_static_input_scheme
            assert qkv_proj.scheme.strategy == strategy
            assert qkv_proj.weight.dtype is torch.int8
258

259
        llm.apply_model(check_model)
260

261
        output = llm.generate_greedy(["Hello my name is"], max_tokens=20)
262
263
        assert output

264

zhuwenwen's avatar
zhuwenwen committed
265
@pytest.mark.skipif(current_platform(),
266
                    reason="WNA16 is not supported on ROCm.")
267
268
@pytest.mark.parametrize(
    "wNa16_args",
zhuwenwen's avatar
zhuwenwen committed
269
    [(os.path.join(models_path_prefix, "nm-testing/tinyllama-oneshot-w4a16-channel-v2"), "channel", None, 8,
270
      True, False),
zhuwenwen's avatar
zhuwenwen committed
271
     (os.path.join(models_path_prefix, "nm-testing/tinyllama-oneshot-w4a16-group128-v2"), "group", 128, 8, True,
272
      False),
zhuwenwen's avatar
zhuwenwen committed
273
     (os.path.join(models_path_prefix, "nm-testing/tinyllama-oneshot-w8a16-per-channel"), "channel", None, 4,
274
      True, False),
zhuwenwen's avatar
zhuwenwen committed
275
     (os.path.join(models_path_prefix, "nm-testing/TinyLlama-1.1B-Chat-v1.0-awq-group128-asym256"), "group", 128,
276
      8, False, False),
zhuwenwen's avatar
zhuwenwen committed
277
     (os.path.join(models_path_prefix, "nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-Channel"),
278
      "channel", None, 8, False, False),
zhuwenwen's avatar
zhuwenwen committed
279
     (os.path.join(models_path_prefix, "nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-ActOrder"),
280
      "group", 128, 8, False, True)],
281
)
282
283
@pytest.mark.skipif(not current_platform.is_cuda(),
                    reason="The tests are skipped on non-CUDA platform.")
284
def test_compressed_tensors_wNa16(vllm_runner, wNa16_args):
285
    model, strategy, group, pack_factor, symmetric, has_g_idx = wNa16_args
286
287
    with vllm_runner(model) as llm:

288
289
        def check_model(model):
            layer = model.model.layers[0]
290

291
292
293
294
            qkv_proj = layer.self_attn.qkv_proj
            assert isinstance(qkv_proj.quant_method,
                              CompressedTensorsLinearMethod)
            assert isinstance(qkv_proj.scheme, CompressedTensorsWNA16)
295

296
297
298
            assert qkv_proj.scheme.strategy == strategy
            assert qkv_proj.scheme.group_size == (-1
                                                  if group is None else group)
299

300
            assert qkv_proj.scheme.pack_factor == pack_factor
301
302
            assert qkv_proj.scheme.symmetric == symmetric
            assert qkv_proj.scheme.has_g_idx == has_g_idx
303
304

        llm.apply_model(check_model)
305

306
307
308
        output = llm.generate_greedy("Hello my name is", max_tokens=20)
        assert output

309

310
311
@pytest.mark.skipif(not current_platform.is_cuda(),
                    reason="This test is skipped on non-CUDA platform.")
312
def test_compressed_tensors_w4a16_marlin24(vllm_runner):
313
    model_path = os.path.join(models_path_prefix,"nm-testing/llama7b-one-shot-2_4-w4a16-marlin24-t")
314
315
    with vllm_runner(model_path) as llm:

316
317
318
319
        def check_model(model):
            layer = model.model.layers[0]

            qkv_proj = layer.self_attn.qkv_proj
320

321
322
323
324
            assert isinstance(qkv_proj.quant_method,
                              CompressedTensorsLinearMethod)
            assert isinstance(qkv_proj.scheme, CompressedTensorsW4A16Sparse24)
            assert qkv_proj.weight_packed.dtype is torch.int32
325

326
        llm.apply_model(check_model)
327

328
        output = llm.generate_greedy("Hello my name is", max_tokens=20)
329
        assert output
330
331


zhuwenwen's avatar
zhuwenwen committed
332
@pytest.mark.skipif(current_platform(),
333
                    reason="FP8 is not supported on ROCm.")
334
def test_compressed_tensors_fp8(vllm_runner):
335
    model_path = os.path.join(models_path_prefix,"nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test")
336
337
    with vllm_runner(model_path) as llm:

338
339
        def check_model(model):
            layer = model.model.layers[0]
340

341
            qkv_proj = layer.self_attn.qkv_proj
342

343
344
345
346
            assert isinstance(qkv_proj.quant_method,
                              CompressedTensorsLinearMethod)
            assert isinstance(
                qkv_proj.scheme,
347
348
                (CompressedTensorsW8A8Fp8, CompressedTensorsW8A16Fp8),
            )
349

350
351
352
353
            assert qkv_proj.input_scale.dtype is torch.float32

            if isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8):
                assert len(qkv_proj.input_scale.shape) == 0
354
                assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
355
356
357
358
                assert qkv_proj.weight_scale.dtype is torch.float32
                assert len(qkv_proj.weight_scale.shape) == 0

        llm.apply_model(check_model)
359

360
        output = llm.generate_greedy("Hello my name is", max_tokens=20)
361
        assert output
362
363


364
365
@pytest.mark.skipif(not current_platform.is_cuda(),
                    reason="This test is skipped on non-CUDA platform.")
366
def test_compressed_tensors_kv_cache(vllm_runner):
367
    model_path = os.path.join(models_path_prefix,"nm-testing/TinyLlama-1.1B-compressed-tensors-kv-cache-scheme")
368
369
    with vllm_runner(model_path, kv_cache_dtype="fp8") as llm:
        output = llm.generate_greedy("Hello world!", max_tokens=20)
370
        assert output
371
372


373
374
375
376
377
378
379
380
@pytest.mark.skipif(
    not sparse_cutlass_supported(),
    reason="Sparse FP8 is not yet supported on this GPU type.",
)
def _test_2of4_quant_models(qkv_proj,
                            weight_strategy,
                            input_strategy,
                            format="dense"):
381
382
383
384
385
386
387
388
    assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
    assert isinstance(qkv_proj.scheme, CompressedTensors24)

    assert qkv_proj.scheme.weight_quant.strategy == weight_strategy
    assert qkv_proj.scheme.input_quant.strategy == input_strategy
    assert qkv_proj.scheme.quantized
    assert qkv_proj.quant_method.quantization_config.sparsity_scheme_map
    sparsity_map = qkv_proj.quant_method.quantization_config.sparsity_scheme_map  # noqa: E501
389
    assert sparsity_map.get("Linear").format == format
390
391
392
    assert sparsity_map.get("Linear").sparsity_structure == "2:4"


393
@pytest.mark.skipif(
394
395
    not current_platform.is_cuda()
    or not current_platform.has_device_capability(90),
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
    reason="Sparse FP8 is not yet supported on this GPU type.",
)
@pytest.mark.parametrize(
    "args_2of4",
    [
        (
            "nm-testing/Meta-Llama-3-8B-Instruct-FP8-Dynamic-2of4-testing",
            "channel",
            "token",
        ),
        (
            "nm-testing/Meta-Llama-3-8B-Instruct-FP8-Static-Per-Tensor-testing",
            "channel",
            "tensor",
        ),
        (
            "nm-testing/Meta-Llama-3-8B-Instruct-FP8-Static-testing",
            "tensor",
            "tensor",
        ),
        (
            "nm-testing/Meta-Llama-3-8B-Instruct-FP8-Dynamic-IA-Per-Tensor-Weight-testing",
            "tensor",
            "token",
        ),
    ],
)
423
424
425
426
def test_compressed_tensors_2of4_quant_fp8(vllm_runner, args_2of4):
    model, weight_strategy, input_strategy = args_2of4
    with vllm_runner(model) as llm:

427
428
429
430
431
432
433
434
        def check_model(model):
            layer = model.model.layers[0]

            qkv_proj = layer.self_attn.qkv_proj
            assert qkv_proj.scheme.weights_dtype == torch.float8_e4m3fn
            _test_2of4_quant_models(qkv_proj, weight_strategy, input_strategy)

        llm.apply_model(check_model)
435
436
437
438
439
440

        output = llm.generate_greedy("Hello my name is", max_tokens=20)
        print(output)
        assert output


441
@pytest.mark.skipif(
442
443
    not current_platform.is_cuda()
    or not current_platform.has_device_capability(90),
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
    reason="Sparse FP8 is not yet supported on this GPU type.",
)
@pytest.mark.parametrize(
    "args_2of4",
    [
        (
            "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_fp8-BitM",
            "channel",
            "token",
        ),
        (
            "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_tensor_act_fp8-BitM",
            "channel",
            "tensor",
        ),
        (
            "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_per_tok_dyn_act_fp8-BitM",
            "tensor",
            "token",
        ),
        (
            "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_tensor_act_fp8-BitM",
            "tensor",
            "tensor",
        ),
    ],
)
def test_compressed_tensors_2of4_quant_fp8_compressed(vllm_runner, args_2of4):
    model, weight_strategy, input_strategy = args_2of4
    with vllm_runner(model) as llm:

        def check_model(model):
            layer = model.model.layers[0]

            qkv_proj = layer.self_attn.qkv_proj
            assert qkv_proj.scheme.weights_dtype == torch.float8_e4m3fn
            _test_2of4_quant_models(
                qkv_proj,
                weight_strategy,
                input_strategy,
                format="sparse-24-bitmask",
            )

        llm.apply_model(check_model)

        output = llm.generate_greedy("Hello my name is", max_tokens=20)
        print(output)
        assert output


@pytest.mark.skipif(
    not sparse_cutlass_supported(),
    reason="cutlass is not yet supported on this GPU type.",
)
@pytest.mark.parametrize(
    "args_2of4",
    [
        (
            "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_int8-BitM",
            "channel",
            "token",
        ),
        (
            "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_tensor_act_int8-BitM",
            "channel",
            "tensor",
        ),
        (
            "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_per_tok_dyn_act_int8-BitM",
            "tensor",
            "token",
        ),
        (
            "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_tensor_act_int8-BitM",
            "tensor",
            "tensor",
        ),
    ],
)
def test_compressed_tensors_2of4_quant_int8_compressed(vllm_runner, args_2of4):
    model, weight_strategy, input_strategy = args_2of4
    with vllm_runner(model) as llm:

        def check_model(model):
            layer = model.model.layers[0]

            qkv_proj = layer.self_attn.qkv_proj
            assert qkv_proj.scheme.weights_dtype == torch.int8
            _test_2of4_quant_models(
                qkv_proj,
                weight_strategy,
                input_strategy,
                format="sparse-24-bitmask",
            )

        llm.apply_model(check_model)

        output = llm.generate_greedy("Hello my name is", max_tokens=20)
        print(output)
        assert output


@pytest.mark.skipif(
    not sparse_cutlass_supported(),
    reason="Sparse FP8 is not yet supported on this GPU type.",
)
@pytest.mark.parametrize(
    "args_2of4",
    [
        (
            "nm-testing/TinyLlama-1.1B-Chat-v1.0-INT8-Dynamic-IA-Per-Channel-Weight-testing",
            "channel",
            "token",
        ),
        (
            "nm-testing/TinyLlama-1.1B-Chat-v1.0-INT8-Static-testing",
            "tensor",
            "tensor",
        ),
        (
            "nm-testing/TinyLlama-1.1B-Chat-v1.0-INT8-Dynamic-IA-Per-Tensor-Weight-testing",
            "tensor",
            "token",
        ),
    ],
)
570
571
572
573
def test_compressed_tensors_2of4_quant_int8(vllm_runner, args_2of4):
    model, weight_strategy, input_strategy = args_2of4
    with vllm_runner(model) as llm:

574
575
576
577
578
579
580
581
        def check_model(model):
            layer = model.model.layers[0]

            qkv_proj = layer.self_attn.qkv_proj
            assert qkv_proj.scheme.weights_dtype == torch.int8
            _test_2of4_quant_models(qkv_proj, weight_strategy, input_strategy)

        llm.apply_model(check_model)
582
583
584
585
586
587

        output = llm.generate_greedy("Hello my name is", max_tokens=20)
        print(output)
        assert output


588
589
@pytest.mark.skipif(
    not sparse_cutlass_supported(),
590
591
    reason="2of4 Sparse is not yet supported on this GPU type.",
)
592
593
@pytest.mark.parametrize(
    "args_2of4",
594
595
    [("nm-testing/TinyLlama-1.1B-Chat-v1.0-2of4-Sparse-Dense-Compressor")],
)
596
597
598
def test_compressed_tensors_2of4_sparse(vllm_runner, args_2of4):
    model = args_2of4
    with vllm_runner(model) as llm:
599
600
601
602
603
604
605
606
607
608
609
610
611

        def check_model(model):
            layer = model.model.layers[0]

            qkv_proj = layer.self_attn.qkv_proj
            assert isinstance(qkv_proj.quant_method,
                              CompressedTensorsLinearMethod)
            assert isinstance(qkv_proj.scheme, CompressedTensors24)

            assert qkv_proj.scheme.weight_quant is None
            assert qkv_proj.scheme.input_quant is None
            assert not qkv_proj.scheme.quantized
            assert qkv_proj.quant_method.quantization_config.sparsity_scheme_map
612
613
614
            sparsity_map = (
                qkv_proj.quant_method.quantization_config.sparsity_scheme_map
            )  # noqa: E501
615
616
617
618
            assert sparsity_map.get("Linear").format == "dense"
            assert sparsity_map.get("Linear").sparsity_structure == "2:4"

        llm.apply_model(check_model)
619
620
621
622

        output = llm.generate_greedy("Hello my name is", max_tokens=20)
        print(output)
        assert output
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657


@pytest.mark.skipif(
    not sparse_cutlass_supported(),
    reason="Cutlass is not yet supported on this GPU type.",
)
@pytest.mark.parametrize(
    "args_2of4", [("nm-testing/llama2.c-stories42M-pruned2.4-compressed")])
def test_compressed_tensors_2of4_sparse_compressed(vllm_runner, args_2of4):
    model = args_2of4
    with vllm_runner(model) as llm:

        def check_model(model):
            layer = model.model.layers[0]

            qkv_proj = layer.self_attn.qkv_proj
            assert isinstance(qkv_proj.quant_method,
                              CompressedTensorsLinearMethod)
            assert isinstance(qkv_proj.scheme, CompressedTensors24)

            assert qkv_proj.scheme.weight_quant is None
            assert qkv_proj.scheme.input_quant is None
            assert not qkv_proj.scheme.quantized
            assert qkv_proj.quant_method.quantization_config.sparsity_scheme_map
            sparsity_map = (
                qkv_proj.quant_method.quantization_config.sparsity_scheme_map
            )  # noqa: E501
            assert sparsity_map.get("Linear").format == "sparse-24-bitmask"
            assert sparsity_map.get("Linear").sparsity_structure == "2:4"

        llm.apply_model(check_model)

        output = llm.generate_greedy("Hello my name is", max_tokens=20)
        print(output)
        assert output