@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 162 lines 5.54 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the triton_scaled_mm kernel.
4
5 Run `pytest tests/kernels/quantization/test_triton_scaled_mm.py`.
6 """
7
8 import importlib
9
10 import pytest
11 import torch
12
13 from vllm.platforms import current_platform
14 from vllm.utils.torch_utils import set_random_seed
15
16 device = current_platform.device_type
17
18 triton_scaled_mm_module = importlib.import_module(
19 "vllm.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm"
20 )
21 triton_scaled_mm = triton_scaled_mm_module.triton_scaled_mm
22
23
24 def torch_scaled_mm(
25 a: torch.Tensor,
26 b: torch.Tensor,
27 scale_a: torch.Tensor,
28 scale_b: torch.Tensor,
29 out_dtype: type[torch.dtype],
30 bias: torch.Tensor | None = None,
31 ) -> torch.Tensor:
32 out = torch.mm(a.to(torch.float32), b.to(torch.float32))
33 out = scale_a * out
34 out = scale_b.T * out
35 out = out.to(out_dtype)
36 if bias is not None:
37 out = out + bias
38
39 return out
40
41
42 def get_8bit_types():
43 types = [torch.int8]
44 if current_platform.supports_fp8():
45 types.append(current_platform.fp8_dtype())
46 return types
47
48
49 # This test is to check regressions for int8 support on ROCm.
50 @pytest.mark.parametrize(
51 "model_path",
52 [
53 "neuralmagic/Llama-3.2-1B-quantized.w8a8",
54 ],
55 )
56 @pytest.mark.parametrize("max_tokens", [32])
57 @pytest.mark.parametrize("num_logprobs", [10])
58 @pytest.mark.skipif(not current_platform.is_rocm(), reason="Should only run on ROCm")
59 def test_rocm_compressed_tensors_w8a8(
60 vllm_runner, example_prompts, model_path, max_tokens, num_logprobs
61 ):
62 dtype = "bfloat16"
63
64 with vllm_runner(model_path, dtype=dtype) as vllm_model:
65 vllm_model.generate_greedy_logprobs(example_prompts, max_tokens, num_logprobs)
66
67
68 MNK_FACTORS = [
69 (1, 256, 128),
70 (33, 256, 496),
71 (64, 971, 1024),
72 (64, 20486, 128),
73 (512, 256, 496),
74 (512, 20486, 1024),
75 ]
76
77
78 @pytest.mark.parametrize("M,N,K", MNK_FACTORS)
79 @pytest.mark.parametrize("out_dtype", [torch.bfloat16])
80 @pytest.mark.parametrize("in_dtype", get_8bit_types())
81 @pytest.mark.parametrize("use_scalar_scale_a", [True, False])
82 @pytest.mark.parametrize("use_scalar_scale_b", [True, False])
83 @pytest.mark.parametrize("use_bias", [True, False])
84 def test_scaled_mm(
85 M, N, K, in_dtype, out_dtype, use_scalar_scale_a, use_scalar_scale_b, use_bias
86 ):
87 is_floating_point_type = lambda t: torch.tensor([1, 1], dtype=t).is_floating_point()
88
89 set_random_seed(0)
90
91 # NOTE: There are cases, where if the matrix is large enough, an output
92 # like 65504.4 can be produced, and can easily turn into inf when
93 # multiplied when using float16/bfloat16. This means one function, e.g.,
94 # testing function, and another function, e.g. golden function, can
95 # produce a non-inf value while the other produces an inf value, and
96 # will cause assert_close/allclose to fail, even though if overflow
97 # wouldn't have occurred, the values would have been "close."
98 #
99 # So, the values here are kept small enough to avoid this situation.
100 if is_floating_point_type(in_dtype):
101 a = (0.25 * torch.rand((M, K), dtype=torch.float32, device=device)).to(in_dtype)
102 b = (0.25 * torch.rand((K, N), dtype=torch.float32, device=device)).to(in_dtype)
103 else:
104 a = torch.randint(-32, 32, (M, K), dtype=in_dtype, device=device)
105 b = torch.randint(-32, 32, (K, N), dtype=in_dtype, device=device)
106
107 if use_scalar_scale_a:
108 scale_a = torch.rand((1, 1), device=device)
109 else:
110 scale_a = 0.25 * torch.rand((M, 1), device=device)
111
112 if use_scalar_scale_b:
113 scale_b = torch.rand((1, 1), device=device)
114 else:
115 scale_b = 0.25 * torch.rand((N, 1), device=device)
116
117 bias = None
118 if use_bias:
119 bias = torch.rand((N,), device=device, dtype=out_dtype)
120
121 c_check = triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
122
123 c_actual = torch_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
124
125 torch.testing.assert_close(c_check, c_actual, rtol=1e-1, atol=1e-1)
126
127
128 # TD operand loads must be bit-exact vs the plain masked-load path.
129 @pytest.mark.skipif(
130 not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
131 reason="Triton scaled_mm runs on CUDA-alike or XPU.",
132 )
133 @pytest.mark.parametrize(
134 "M,N,K", [(1, 4096, 4096), (64, 4096, 4096), (256, 2048, 4096)]
135 )
136 @pytest.mark.parametrize("in_dtype", get_8bit_types())
137 @pytest.mark.parametrize("use_scalar_scale_a", [True, False])
138 @pytest.mark.parametrize("use_bias", [True, False])
139 def test_scaled_mm_td_matches_plain(M, N, K, in_dtype, use_scalar_scale_a, use_bias):
140 dev = current_platform.device_type
141 out_dtype = torch.bfloat16
142 set_random_seed(0)
143
144 is_fp = torch.tensor([1, 1], dtype=in_dtype).is_floating_point()
145 if is_fp:
146 a = (0.25 * torch.rand((M, K), dtype=torch.float32, device=dev)).to(in_dtype)
147 b = (0.25 * torch.rand((K, N), dtype=torch.float32, device=dev)).to(in_dtype)
148 else:
149 a = torch.randint(-32, 32, (M, K), dtype=in_dtype, device=dev)
150 b = torch.randint(-32, 32, (K, N), dtype=in_dtype, device=dev)
151
152 scale_a = (
153 torch.rand((1, 1), device=dev)
154 if use_scalar_scale_a
155 else 0.25 * torch.rand((M, 1), device=dev)
156 )
157 scale_b = 0.25 * torch.rand((N, 1), device=dev)
158 bias = torch.rand((N,), device=dev, dtype=out_dtype) if use_bias else None
159
160 out_plain = triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias, use_td=False)
161 out_td = triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias, use_td=True)
162 torch.testing.assert_close(out_td, out_plain, rtol=0, atol=0)