@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 240 lines 7.76 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4 import torch
5 import torch.nn.functional as F
6
7 import vllm._custom_ops as ops
8 from tests.kernels.utils import opcheck
9 from vllm.model_executor.layers.activation import SiluAndMul
10 from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant
11 from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
12 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
13 per_token_group_quant_fp8,
14 )
15 from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 kFp8Dynamic128Sym,
17 kFp8StaticTensorSym,
18 )
19 from vllm.platforms import current_platform
20 from vllm.utils.torch_utils import set_random_seed
21
22 DTYPES = [torch.bfloat16, torch.float16]
23 QUANT_DTYPES = [current_platform.fp8_dtype()]
24 NUM_TOKENS = [1, 17, 86, 1234, 3045] # Arbitrary values for testing
25 HIDDEN_SIZES = [16, 48, 128, 1562, 4096] # Arbitrary values for testing
26 SEEDS = [0]
27 CUDA_DEVICES = [
28 f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
29 ]
30
31
32 def ref_impl(
33 silu_and_mul: SiluAndMul, x: torch.Tensor, scale: torch.Tensor
34 ) -> torch.Tensor:
35 silu_and_mul_out = silu_and_mul.forward_native(x)
36 out, scales = ops.scaled_fp8_quant(silu_and_mul_out, scale)
37 return out
38
39
40 def ops_impl(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
41 out_shape = (x.shape[0], x.shape[1] // 2)
42 out = torch.empty(out_shape, dtype=current_platform.fp8_dtype(), device=x.device)
43 torch.ops._C.silu_and_mul_quant(out, x, scale)
44 return out
45
46
47 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
48 @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
49 @pytest.mark.parametrize("dtype", DTYPES)
50 @pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
51 @pytest.mark.parametrize("seed", SEEDS)
52 @pytest.mark.parametrize("device", CUDA_DEVICES)
53 @torch.inference_mode()
54 def test_silu_and_mul(
55 default_vllm_config,
56 num_tokens: int,
57 hidden_size: int,
58 dtype: torch.dtype,
59 quant_dtype: torch.dtype,
60 seed: int,
61 device: str,
62 ) -> None:
63 set_random_seed(seed)
64 torch.set_default_device(device)
65
66 layer = SiluAndMul()
67
68 # Make inputs
69 scale = torch.randn((1), device=device, dtype=torch.float32)
70 x = torch.randn(num_tokens, hidden_size, dtype=dtype)
71
72 ref_out = ref_impl(layer, x, scale)
73 ops_out = ops_impl(x, scale)
74
75 assert ref_out.dtype == quant_dtype
76 assert ops_out.dtype == quant_dtype
77 assert ref_out.shape == ops_out.shape
78 assert torch.allclose(
79 ref_out.to(dtype=torch.float32), ops_out.to(dtype=torch.float32)
80 )
81 opcheck(torch.ops._C.silu_and_mul_quant, (ops_out, x, scale))
82
83
84 # ---------------------------------------------------------------------------
85 # Tests for maybe_fused_act_quant interface
86 # ---------------------------------------------------------------------------
87
88
89 class MockLinearFp8Static(torch.nn.Module):
90 """Mock linear layer advertising kFp8StaticTensorSym."""
91
92 def __init__(self, input_scale: torch.Tensor):
93 super().__init__()
94 self._input_quant_key = kFp8StaticTensorSym
95 self.input_scale = input_scale
96
97
98 class MockLinearFp8Dynamic128(torch.nn.Module):
99 """Mock linear layer advertising kFp8Dynamic128Sym."""
100
101 def __init__(self):
102 super().__init__()
103 self._input_quant_key = kFp8Dynamic128Sym
104
105
106 class MockLinearNoQuant(torch.nn.Module):
107 """Mock linear layer with no input quantization key."""
108
109 pass
110
111
112 class MockLinearRequiresUnquantized(MockLinearFp8Static):
113 """Mock consumer that needs the original activation for another branch."""
114
115 requires_unquantized_input = True
116
117
118 @pytest.mark.parametrize("num_tokens", [1, 16, 128])
119 @pytest.mark.parametrize("hidden_size", [128, 512, 1024])
120 @pytest.mark.parametrize("dtype", DTYPES)
121 @torch.inference_mode()
122 def test_maybe_fused_act_quant_fp8_static(
123 default_vllm_config,
124 num_tokens: int,
125 hidden_size: int,
126 dtype: torch.dtype,
127 ) -> None:
128 """Test maybe_fused_act_quant with FP8 static per-tensor quantization."""
129 device = "cuda:0"
130 torch.set_default_device(device)
131
132 act_fn = SiluAndMul()
133 scale = torch.tensor([0.5], device=device, dtype=torch.float32)
134 linear = MockLinearFp8Static(scale)
135
136 x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device)
137 result = maybe_fused_act_quant(act_fn, x, linear)
138
139 assert isinstance(result, QuantizedActivation)
140 assert result.quant_key == kFp8StaticTensorSym
141 assert result.data.dtype == current_platform.fp8_dtype()
142 assert result.orig_dtype == dtype
143 assert result.orig_shape == (num_tokens, hidden_size)
144
145 ref_out = ref_impl(act_fn, x, scale)
146 torch.testing.assert_close(result.data.to(torch.float32), ref_out.to(torch.float32))
147
148
149 @pytest.mark.parametrize("num_tokens", [1, 16, 128])
150 @pytest.mark.parametrize("hidden_size", [128, 512, 1024])
151 @pytest.mark.parametrize("dtype", DTYPES)
152 @torch.inference_mode()
153 def test_maybe_fused_act_quant_fp8_dynamic_block(
154 default_vllm_config,
155 num_tokens: int,
156 hidden_size: int,
157 dtype: torch.dtype,
158 ) -> None:
159 """Test maybe_fused_act_quant with FP8 dynamic per-block quantization."""
160 group_size = 128 # We only support 128 for now
161
162 device = "cuda:0"
163 torch.set_default_device(device)
164
165 act_fn = SiluAndMul()
166 linear = MockLinearFp8Dynamic128()
167
168 scale = 1 / hidden_size
169 x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) * scale
170 result = maybe_fused_act_quant(act_fn, x, linear)
171
172 assert isinstance(result, QuantizedActivation)
173 assert result.quant_key == kFp8Dynamic128Sym
174 assert result.data.dtype == current_platform.fp8_dtype()
175 assert result.orig_dtype == dtype
176 assert result.orig_shape == (num_tokens, hidden_size)
177
178 num_groups = hidden_size // group_size
179 assert result.scale.shape == (num_tokens, num_groups)
180
181 gate, up = x.split(hidden_size, dim=-1)
182 silu_out = F.silu(gate) * up
183 ref_out, ref_scales = per_token_group_quant_fp8(
184 silu_out, group_size=group_size, use_ue8m0=False
185 )
186
187 torch.testing.assert_close(result.scale, ref_scales, rtol=1e-5, atol=1e-5)
188
189 ref_deq = ref_out.to(torch.float32) * ref_scales.repeat_interleave(
190 group_size, dim=1
191 )
192 result_deq = result.data.to(torch.float32) * result.scale.repeat_interleave(
193 group_size, dim=1
194 )
195 torch.testing.assert_close(ref_deq, result_deq, atol=5e-2, rtol=5e-2)
196
197
198 @pytest.mark.parametrize("num_tokens", [1, 16, 128])
199 @pytest.mark.parametrize("hidden_size", [128, 512])
200 @pytest.mark.parametrize("dtype", DTYPES)
201 @torch.inference_mode()
202 def test_maybe_fused_act_quant_fallback(
203 default_vllm_config,
204 num_tokens: int,
205 hidden_size: int,
206 dtype: torch.dtype,
207 ) -> None:
208 """Test maybe_fused_act_quant falls back without an input quantization key."""
209 device = "cuda:0"
210 torch.set_default_device(device)
211
212 act_fn = SiluAndMul()
213 linear = MockLinearNoQuant()
214 x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device)
215
216 result = maybe_fused_act_quant(act_fn, x, linear)
217
218 assert isinstance(result, torch.Tensor)
219 assert not isinstance(result, QuantizedActivation)
220 assert result.dtype == dtype
221 assert result.shape == (num_tokens, hidden_size)
222
223 ref_out = act_fn(x)
224 torch.testing.assert_close(result, ref_out)
225
226
227 @torch.inference_mode()
228 def test_maybe_fused_act_quant_preserves_required_unquantized_input(
229 default_vllm_config,
230 ) -> None:
231 device = "cuda:0"
232 act_fn = SiluAndMul()
233 scale = torch.tensor([0.5], device=device, dtype=torch.float32)
234 linear = MockLinearRequiresUnquantized(scale)
235 x = torch.randn(17, 256, dtype=torch.bfloat16, device=device)
236
237 result = maybe_fused_act_quant(act_fn, x, linear)
238
239 assert isinstance(result, torch.Tensor)
240 torch.testing.assert_close(result, act_fn(x))