| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the Triton dequant-gather kernel used by |
| 4 | ``CompressedTensorsEmbeddingWNA16Int`` (quantized embedding lookup).""" |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | from compressed_tensors.compressors.pack_quantized.helpers import unpack_from_int32 |
| 9 | |
| 10 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_embedding import ( # noqa: E501 |
| 11 | _dequant_gather_triton, |
| 12 | ) |
| 13 | from vllm.platforms import current_platform |
| 14 | |
| 15 | |
| 16 | def _dequant_gather_torch( |
| 17 | ids: torch.Tensor, |
| 18 | weight_packed: torch.Tensor, |
| 19 | weight_scale: torch.Tensor, |
| 20 | hidden: int, |
| 21 | num_bits: int, |
| 22 | ) -> torch.Tensor: |
| 23 | """Reference: gather packed rows by id, unpack int32-packed INT, dequant.""" |
| 24 | n = ids.shape[0] |
| 25 | int8 = unpack_from_int32(weight_packed[ids], num_bits, torch.Size([n, hidden])) |
| 26 | scale_rows = weight_scale[ids] |
| 27 | w = int8.to(scale_rows.dtype) |
| 28 | if scale_rows.shape[1] == 1: |
| 29 | return w * scale_rows |
| 30 | ng = scale_rows.shape[1] |
| 31 | return (w.view(n, ng, hidden // ng) * scale_rows.unsqueeze(-1)).view(n, hidden) |
| 32 | |
| 33 | |
| 34 | @pytest.mark.skipif( |
| 35 | not current_platform.is_cuda(), reason="Triton dequant kernel requires CUDA" |
| 36 | ) |
| 37 | @pytest.mark.parametrize("num_bits", [2, 4, 8]) |
| 38 | @pytest.mark.parametrize("group_size", [0, 256]) # 0 -> channel |
| 39 | @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) |
| 40 | @pytest.mark.parametrize("num_ids", [1, 17, 4096]) |
| 41 | def test_dequant_gather(num_bits, group_size, dtype, num_ids): |
| 42 | torch.manual_seed(0) |
| 43 | device = "cuda" |
| 44 | vocab, hidden = 1000, 2048 |
| 45 | pack_factor = 32 // num_bits |
| 46 | |
| 47 | # Random full-range int32 packed weights (covers the sign bit -> exercises the |
| 48 | # arithmetic-shift + mask unpack path). |
| 49 | weight_packed = torch.randint( |
| 50 | -(2**31), |
| 51 | 2**31, |
| 52 | (vocab, hidden // pack_factor), |
| 53 | dtype=torch.int32, |
| 54 | device=device, |
| 55 | ) |
| 56 | |
| 57 | num_groups = 1 if group_size == 0 else hidden // group_size |
| 58 | weight_scale = torch.rand(vocab, num_groups, dtype=dtype, device=device) + 0.01 |
| 59 | |
| 60 | ids = torch.randint(0, vocab, (num_ids,), dtype=torch.long, device=device) |
| 61 | |
| 62 | out = _dequant_gather_triton(ids, weight_packed, weight_scale, hidden, num_bits) |
| 63 | ref = _dequant_gather_torch(ids, weight_packed, weight_scale, hidden, num_bits) |
| 64 | |
| 65 | assert out.shape == (num_ids, hidden) |
| 66 | assert out.dtype == dtype |
| 67 | torch.testing.assert_close(out, ref) |