| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | """End-to-end smoke test for CT W4A16 models on ROCm. |
| 5 | |
| 6 | This validates that a real compressed-tensors W4A16 model can run inference |
| 7 | end-to-end (which will exercise the Triton W4A16 kernel when selected). |
| 8 | |
| 9 | Run `pytest tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py`. |
| 10 | """ |
| 11 | |
| 12 | import pytest |
| 13 | |
| 14 | from vllm.platforms import current_platform |
| 15 | |
| 16 | |
| 17 | @pytest.mark.parametrize( |
| 18 | "model_path", |
| 19 | [ |
| 20 | "nm-testing/tinyllama-oneshot-w4a16-group128-v2", |
| 21 | ], |
| 22 | ) |
| 23 | @pytest.mark.parametrize("max_tokens", [32]) |
| 24 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="Should only run on ROCm") |
| 25 | def test_rocm_compressed_tensors_w4a16_e2e( |
| 26 | vllm_runner, example_prompts, model_path, max_tokens |
| 27 | ): |
| 28 | # Use fp16 activations for maximum compatibility. |
| 29 | # gpu_memory_utilization lowered to work on shared nodes. |
| 30 | with vllm_runner( |
| 31 | model_path, dtype="float16", gpu_memory_utilization=0.3 |
| 32 | ) as vllm_model: |
| 33 | # If the W4A16 kernel is broken, this will typically throw. |
| 34 | vllm_model.generate_greedy(example_prompts, max_tokens=max_tokens) |