@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 220 lines 9.52 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from unittest.mock import MagicMock
5
6 import pytest
7 import torch
8
9 from vllm.platforms import current_platform
10
11 if current_platform.is_cuda():
12 pytest.skip(
13 "ROCm skinny GEMM tests are not supported on CUDA.",
14 allow_module_level=True,
15 )
16
17 from vllm.model_executor.layers import utils
18
19
20 @pytest.mark.parametrize("m", [1, 128])
21 def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch, m):
22 x = torch.randn(1, 64, dtype=torch.float16)
23 weight = torch.randn(m, 64, dtype=torch.float16)
24
25 monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
26 monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
27 monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
28 monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
29 monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
30 monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False)
31 monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
32
33 wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
34 monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
35 llmm1_mock = MagicMock(side_effect=lambda w, x_view, _: x_view @ w.t())
36 monkeypatch.setattr(utils.ops, "LLMM1", llmm1_mock)
37
38 out = utils.rocm_unquantized_gemm_impl(x, weight, None)
39 ref = torch.nn.functional.linear(x, weight, None)
40
41 wvsplitk_mock.assert_called_once()
42 llmm1_mock.assert_not_called()
43 assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
44
45
46 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only kernel test")
47 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
48 @pytest.mark.parametrize("n", range(1, 6))
49 def test_rocm_unquantized_gemm_single_output_real_kernel(dtype, n):
50 torch.manual_seed(0)
51 x = torch.randn(n, 2048, device="cuda", dtype=dtype)
52 weight = torch.randn(1, 2048, device="cuda", dtype=dtype)
53
54 out = utils.rocm_unquantized_gemm_impl(x, weight, None)
55 ref = torch.nn.functional.linear(x, weight, None)
56
57 torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
58
59
60 def test_rocm_unquantized_gemm_makes_skinny_activation_contiguous(monkeypatch):
61 x = torch.randn(64, 4, dtype=torch.float16).t()
62 weight = torch.randn(128, 64, dtype=torch.float16)
63 assert x.shape == (4, 64)
64 assert x.stride() == (1, 4)
65
66 monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
67 monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
68 monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
69 monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
70 monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
71 monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False)
72 monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
73
74 wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
75 monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
76
77 out = utils.rocm_unquantized_gemm_impl(x, weight, None)
78 ref = torch.nn.functional.linear(x, weight, None)
79
80 wvsplitk_mock.assert_called_once()
81 x_view = wvsplitk_mock.call_args.args[1]
82 assert x_view.is_contiguous()
83 assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
84
85
86 def test_rocm_unquantized_gemm_makes_llmm1_activation_contiguous(monkeypatch):
87 x = torch.randn(1, 128, dtype=torch.float16)[:, ::2]
88 weight = torch.randn(4, 64, dtype=torch.float16)
89 assert x.shape == (1, 64)
90 assert x.stride() == (128, 2)
91
92 monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
93 monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
94 monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
95 monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
96 monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
97 monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False)
98 monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
99
100 llmm1_mock = MagicMock(side_effect=lambda w, x_view, _: x_view @ w.t())
101 monkeypatch.setattr(utils.ops, "LLMM1", llmm1_mock)
102
103 out = utils.rocm_unquantized_gemm_impl(x, weight, None)
104 ref = torch.nn.functional.linear(x, weight, None)
105
106 llmm1_mock.assert_called_once()
107 x_view = llmm1_mock.call_args.args[1]
108 assert x_view.is_contiguous()
109 assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
110
111
112 @pytest.mark.parametrize("noncontiguous_operand", ["weight", "bias"])
113 def test_rocm_unquantized_gemm_rejects_unsupported_skinny_layouts(
114 monkeypatch, noncontiguous_operand
115 ):
116 x = torch.randn(4, 64, dtype=torch.float16)
117 weight = torch.randn(128, 64, dtype=torch.float16)
118 bias = torch.randn(128, dtype=torch.float16)
119 if noncontiguous_operand == "weight":
120 weight = torch.randn(64, 128, dtype=torch.float16).t()
121 assert not weight.is_contiguous()
122 else:
123 bias = torch.randn(256, dtype=torch.float16)[::2]
124 assert not bias.is_contiguous()
125
126 monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
127 monkeypatch.setattr(utils.rocm_aiter_ops, "is_tgemm_enabled", lambda: False)
128 monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
129 monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
130 monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
131 monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
132 monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False)
133 monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
134
135 wvsplitk_mock = MagicMock()
136 monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
137 llmm1_mock = MagicMock()
138 monkeypatch.setattr(utils.ops, "LLMM1", llmm1_mock)
139
140 out = utils.rocm_unquantized_gemm_impl(x, weight, bias)
141 ref = torch.nn.functional.linear(x, weight, bias)
142
143 wvsplitk_mock.assert_not_called()
144 llmm1_mock.assert_not_called()
145 assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
146
147
148 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only kernel test")
149 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
150 def test_rocm_unquantized_gemm_noncontiguous_activation_real_kernel(monkeypatch, dtype):
151 x = torch.randn(64, 4, device="cuda", dtype=dtype).t()
152 weight = torch.randn(128, 64, device="cuda", dtype=dtype)
153 assert x.stride() == (1, 4)
154
155 monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
156 original_wvsplitk = utils.ops.wvSplitK
157 wvsplitk_mock = MagicMock(side_effect=original_wvsplitk)
158 monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
159
160 out = utils.rocm_unquantized_gemm_impl(x, weight, None)
161 ref = torch.nn.functional.linear(x, weight, None)
162
163 wvsplitk_mock.assert_called_once()
164 torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
165
166
167 @pytest.mark.parametrize("m", [1, 128])
168 def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch, m):
169 # wvSplitK skinny GEMM handles n in [1, 5] (see PR #40687); n > 5 must
170 # fall back to torch.nn.functional.linear.
171 x = torch.randn(6, 64, dtype=torch.float16)
172 weight = torch.randn(m, 64, dtype=torch.float16)
173
174 monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
175 monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
176 monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
177 monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
178 monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
179 monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False)
180 monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
181
182 wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
183 monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
184 llmm1_mock = MagicMock(side_effect=lambda w, x_view, _: x_view @ w.t())
185 monkeypatch.setattr(utils.ops, "LLMM1", llmm1_mock)
186
187 out = utils.rocm_unquantized_gemm_impl(x, weight, None)
188 ref = torch.nn.functional.linear(x, weight, None)
189
190 wvsplitk_mock.assert_not_called()
191 llmm1_mock.assert_not_called()
192 assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
193
194
195 def test_rocm_unquantized_gemm_gfx950_wvsplitkrc_path(monkeypatch):
196 x = torch.randn(1024, 16, dtype=torch.float16).t()
197 weight = torch.randn(256, 1024, dtype=torch.float16)
198 assert x.stride() == (1, 16)
199
200 monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
201 monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
202 monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: False)
203 monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
204 monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True)
205 monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: True)
206 monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
207
208 wvsplitkrc_mock = MagicMock(side_effect=lambda x_view, w, _, __: x_view @ w.t())
209 monkeypatch.setattr(utils.ops, "wvSplitKrc", wvsplitkrc_mock)
210 wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
211 monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
212
213 out = utils.rocm_unquantized_gemm_impl(x, weight, None)
214 ref = torch.nn.functional.linear(x, weight, None)
215
216 wvsplitkrc_mock.assert_called_once()
217 wvsplitk_mock.assert_not_called()
218 x_view = wvsplitkrc_mock.call_args.args[0]
219 assert x_view.is_contiguous()
220 assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)