@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 165 lines 5.61 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for CPU unquantized GEMM dispatch behavior."""
4
5 import pytest
6 import torch
7
8 from vllm.model_executor.layers import utils
9 from vllm.platforms import CpuArchEnum, current_platform
10 from vllm.utils.torch_utils import set_default_torch_dtype
11
12
13 def test_dispatch_prepacks_arm_bf16_causal_conv(monkeypatch):
14 monkeypatch.setattr(
15 current_platform, "get_cpu_architecture", lambda: CpuArchEnum.ARM
16 )
17 monkeypatch.setattr(torch.cpu, "get_capabilities", lambda: {"bf16": True})
18 monkeypatch.setattr(torch.cpu, "_is_avx512_bf16_supported", lambda: False)
19
20 packed = torch.randn(32, 4, dtype=torch.bfloat16)
21 pack_inputs = []
22
23 def pack(weight):
24 pack_inputs.append(weight.clone())
25 return packed
26
27 monkeypatch.setattr(utils.ops, "causal_conv1d_weight_pack", pack)
28 original = torch.randn(32, 1, 4, dtype=torch.bfloat16)
29 layer = torch.nn.Module()
30 layer.weight = torch.nn.Parameter(original.clone(), requires_grad=False)
31
32 utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
33
34 assert len(pack_inputs) == 1
35 # Check causal_conv1d_weight_pack received the expected weights
36 torch.testing.assert_close(pack_inputs[0], original.view(32, 4))
37 # Check we have correctly stashed the unpacked weights
38 torch.testing.assert_close(layer._cpu_unpacked_conv_weight, original.view(32, 4))
39 # Check we have correctly stored the packed weights
40 assert layer.weight.data_ptr() == packed.data_ptr()
41
42
43 def test_dispatch_does_not_pack_3d_expert_weight(monkeypatch):
44 monkeypatch.setattr(
45 current_platform, "get_cpu_architecture", lambda: CpuArchEnum.ARM
46 )
47 monkeypatch.setattr(torch.cpu, "get_capabilities", lambda: {"bf16": True})
48 monkeypatch.setattr(torch.cpu, "_is_avx512_bf16_supported", lambda: False)
49 pack_calls = []
50 monkeypatch.setattr(
51 utils.ops,
52 "causal_conv1d_weight_pack",
53 lambda weight: pack_calls.append(weight),
54 )
55
56 layer = torch.nn.Module()
57 layer.weight = torch.nn.Parameter(
58 torch.randn(8, 32, 4, dtype=torch.bfloat16), requires_grad=False
59 )
60 utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
61 assert pack_calls == []
62
63
64 @pytest.fixture(scope="module")
65 def _mock_zentorch_linear_unary():
66 """Register a mock zentorch_linear_unary op when zentorch is not installed.
67
68 Allows the dispatch tests to run in CI without a real zentorch build.
69 Skips registration when zentorch is already available.
70 """
71 if hasattr(torch.ops.zentorch, "zentorch_linear_unary"):
72 yield
73 return
74
75 lib_def = torch.library.Library("zentorch", "DEF")
76 lib_def.define(
77 "zentorch_linear_unary("
78 "Tensor input, "
79 "Tensor weight, "
80 "Tensor? bias, "
81 "bool is_weight_prepacked=False"
82 ") -> Tensor"
83 )
84
85 lib_impl = torch.library.Library("zentorch", "IMPL", "CPU")
86 lib_impl.impl(
87 "zentorch_linear_unary",
88 lambda input, weight, bias, is_weight_prepacked=False: (
89 torch.nn.functional.linear(input, weight, bias)
90 ),
91 )
92
93 yield
94
95 lib_impl._destroy()
96 lib_def._destroy()
97
98
99 @pytest.mark.usefixtures("_mock_zentorch_linear_unary")
100 def test_dispatch_cpu_unquantized_gemm_uses_zentorch_on_zen(monkeypatch):
101 monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
102
103 layer = torch.nn.Linear(16, 8, bias=True)
104 x = torch.randn(4, 16)
105 expected = torch.nn.functional.linear(x, layer.weight, layer.bias)
106
107 utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
108 output = layer.cpu_linear(x, layer.weight, layer.bias)
109
110 torch.testing.assert_close(output, expected)
111
112
113 @pytest.mark.usefixtures("_mock_zentorch_linear_unary")
114 def test_dispatch_cpu_unquantized_gemm_zen_remove_weight(monkeypatch):
115 monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
116
117 layer = torch.nn.Linear(16, 8, bias=True)
118 utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
119
120 assert layer.weight.numel() == 0
121
122
123 @pytest.mark.usefixtures("_mock_zentorch_linear_unary")
124 def test_dispatch_cpu_unquantized_gemm_logs_zentorch_dispatch(monkeypatch):
125 monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
126 expected_prepacked = bool(utils.envs.VLLM_ZENTORCH_WEIGHT_PREPACK) and hasattr(
127 torch.ops.zentorch, "zentorch_weight_prepack_for_linear"
128 )
129
130 log_calls = []
131 monkeypatch.setattr(
132 utils.logger, "debug_once", lambda *args: log_calls.append(args)
133 )
134
135 layer = torch.nn.Linear(16, 8, bias=True)
136 utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
137
138 assert log_calls == [
139 (
140 "CPU unquantized GEMM dispatch: using zentorch_linear_unary (prepacked=%s)",
141 expected_prepacked,
142 )
143 ]
144
145
146 @pytest.mark.usefixtures("_mock_zentorch_linear_unary")
147 @pytest.mark.parametrize("weight_dtype", [torch.bfloat16, torch.float16, torch.float32])
148 def test_dispatch_cpu_unquantized_gemm_remove_weight_keeps_dtype(
149 monkeypatch, weight_dtype
150 ):
151 monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
152
153 layer = torch.nn.Linear(16, 8, bias=False, dtype=weight_dtype)
154 loading_dtype = (
155 torch.float32 if weight_dtype is not torch.float32 else torch.bfloat16
156 )
157 with set_default_torch_dtype(loading_dtype):
158 utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
159
160 assert layer.weight.numel() == 0
161 assert layer.weight.dtype is weight_dtype
162
163 x = torch.randn(4, 16, dtype=weight_dtype)
164 output = layer.cpu_linear(x, layer.weight, None)
165 assert not output.isnan().any()