| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | """Kernel-selection / gating tests for the ROCm RDNA3 W4A16 GPTQ kernel. |
| 5 | |
| 6 | Verifies that ``choose_mp_linear_kernel`` resolves a supported W4A16 GPTQ |
| 7 | config to ``RDNA3W4A16LinearKernel`` on gfx1100 (it is registered ahead of |
| 8 | ``TritonW4A16LinearKernel`` in the ROCm priority list), and that |
| 9 | ``RDNA3W4A16LinearKernel.can_implement`` rejects the configs it does not |
| 10 | support so selection falls through to the next kernel. |
| 11 | |
| 12 | Run `pytest tests/kernels/quantization/test_rdna3_w4a16_selection.py`. |
| 13 | """ |
| 14 | |
| 15 | import pytest |
| 16 | import torch |
| 17 | |
| 18 | from vllm.platforms import current_platform |
| 19 | |
| 20 | if not current_platform.is_rocm(): |
| 21 | pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True) |
| 22 | |
| 23 | from vllm.model_executor.kernels.linear import ( # noqa: E402 |
| 24 | choose_mp_linear_kernel, |
| 25 | ) |
| 26 | from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E402 |
| 27 | MPLinearLayerConfig, |
| 28 | ) |
| 29 | from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E402 |
| 30 | RDNA3W4A16LinearKernel, |
| 31 | ) |
| 32 | from vllm.platforms.rocm import on_gfx1100 # noqa: E402 |
| 33 | from vllm.scalar_type import scalar_types # noqa: E402 |
| 34 | |
| 35 | WEIGHT_TYPE = scalar_types.uint4b8 # symmetric int4, bias = 8 |
| 36 | |
| 37 | # The kernel is only selectable when running on gfx1100 with the custom op |
| 38 | # compiled in; otherwise can_implement rejects and selection falls through. |
| 39 | gfx1100_only = pytest.mark.skipif( |
| 40 | not ( |
| 41 | on_gfx1100() |
| 42 | and hasattr(torch.ops, "_rocm_C") |
| 43 | and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3") |
| 44 | ), |
| 45 | reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in", |
| 46 | ) |
| 47 | |
| 48 | |
| 49 | @gfx1100_only |
| 50 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 51 | def test_selection_prefers_rdna3(dtype): |
| 52 | """A supported W4A16 GPTQ config resolves to the RDNA3 kernel on gfx1100.""" |
| 53 | config = MPLinearLayerConfig( |
| 54 | full_weight_shape=(1024, 256), |
| 55 | partition_weight_shape=(1024, 256), |
| 56 | weight_type=WEIGHT_TYPE, |
| 57 | act_type=dtype, |
| 58 | group_size=128, |
| 59 | zero_points=False, |
| 60 | ) |
| 61 | assert choose_mp_linear_kernel(config).__name__ == "RDNA3W4A16LinearKernel" |
| 62 | |
| 63 | |
| 64 | @gfx1100_only |
| 65 | @pytest.mark.parametrize( |
| 66 | "weight_type,group_size,N,full_k,expected_ok", |
| 67 | [ |
| 68 | (scalar_types.uint4b8, 128, 256, 1024, True), # nominal: supported |
| 69 | (scalar_types.uint4b8, -1, 256, 1024, False), # channelwise unsupported |
| 70 | (scalar_types.uint4b8, 128, 252, 1024, False), # N not a multiple of 8 |
| 71 | (scalar_types.uint4b8, 96, 256, 1024, False), # group does not divide K |
| 72 | (scalar_types.uint8b128, 128, 256, 1024, False), # wrong quant type |
| 73 | ], |
| 74 | ids=["ok", "channelwise", "bad_n", "group_ndiv_k", "wrong_qtype"], |
| 75 | ) |
| 76 | def test_can_implement(weight_type, group_size, N, full_k, expected_ok): |
| 77 | """can_implement gates on quant type, group size, and N divisibility.""" |
| 78 | config = MPLinearLayerConfig( |
| 79 | full_weight_shape=(full_k, N), |
| 80 | partition_weight_shape=(full_k, N), |
| 81 | weight_type=weight_type, |
| 82 | act_type=torch.float16, |
| 83 | group_size=group_size, |
| 84 | zero_points=False, |
| 85 | ) |
| 86 | ok, reason = RDNA3W4A16LinearKernel.can_implement(config) |
| 87 | assert ok is expected_ok, reason |