@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 260 lines 9.08 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5 import torch
6
7 from tests.kernels.quant_utils import FP8_DTYPE
8 from tests.kernels.utils import fp8_allclose, fp8_ulp_distance, opcheck
9 from vllm import ir
10 from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
11 from vllm.platforms import current_platform
12 from vllm.utils.torch_utils import set_random_seed
13
14 if current_platform.is_rocm():
15 from vllm.platforms.rocm import on_gfx90a
16
17 on_mi250 = on_gfx90a()
18 else:
19 on_mi250 = False
20
21 DTYPES = [torch.half, torch.bfloat16, torch.float]
22 NUM_TOKENS = [7, 83, 4096] # Arbitrary values for testing
23 HIDDEN_SIZES = [8, 768, 769, 5120, 5125, 8192] # Arbitrary values for testing
24 ADD_RESIDUAL = [False, True] if not on_mi250 else [True]
25 SEEDS = [0]
26 CUDA_DEVICES = [
27 f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
28 ]
29
30
31 def _rms_norm_tolerance(dtype: torch.dtype) -> dict[str, float]:
32 return ir.ops.rms_norm.get_tolerance(dtype)
33
34
35 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
36 @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
37 @pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
38 @pytest.mark.parametrize("dtype", DTYPES)
39 @pytest.mark.parametrize("seed", SEEDS)
40 @pytest.mark.parametrize("device", CUDA_DEVICES)
41 @pytest.mark.parametrize("strided_input", [False, True])
42 @torch.inference_mode()
43 def test_rms_norm(
44 default_vllm_config,
45 num_tokens: int,
46 hidden_size: int,
47 add_residual: bool,
48 dtype: torch.dtype,
49 seed: int,
50 device: str,
51 strided_input: bool,
52 ) -> None:
53 set_random_seed(seed)
54 torch.set_default_device(device)
55 layer = RMSNorm(hidden_size).to(dtype=dtype)
56 layer.weight.data.normal_(mean=1.0, std=0.1)
57 scale = 1 / (2 * hidden_size)
58 last_dim = 2 * hidden_size if strided_input else hidden_size
59 x = torch.randn(num_tokens, last_dim, dtype=dtype)
60 x = x[..., :hidden_size]
61 assert x.is_contiguous() != strided_input
62 x *= scale
63 residual = x.new_empty_strided(x.size(), x.stride()) if add_residual else None
64 if residual is not None:
65 residual.normal_(std=scale)
66
67 # NOTE(woosuk): The reference implementation should be executed first
68 # because the custom kernel is in-place.
69 ref_out = layer.forward_native(x, residual)
70 out = layer(x, residual)
71 # NOTE(woosuk): LayerNorm operators (including RMS) typically have larger
72 # numerical errors than other operators because they involve reductions.
73 # Therefore, we use a larger tolerance.
74 if add_residual:
75 torch.testing.assert_close(out[0], ref_out[0], atol=1e-2, rtol=1e-2)
76 torch.testing.assert_close(out[1], ref_out[1], atol=1e-2, rtol=1e-2)
77 else:
78 torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
79
80 if residual is not None:
81 opcheck(
82 torch.ops._C.fused_add_rms_norm,
83 (x, residual, layer.weight.data, layer.variance_epsilon),
84 )
85 else:
86 opcheck(
87 torch.ops._C.rms_norm, (out, x, layer.weight.data, layer.variance_epsilon)
88 )
89
90
91 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
92 @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
93 @pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
94 @pytest.mark.parametrize("dtype", DTYPES)
95 @pytest.mark.parametrize("seed", SEEDS)
96 @pytest.mark.parametrize("device", CUDA_DEVICES)
97 @torch.inference_mode()
98 def test_rms_norm_weightless(
99 default_vllm_config,
100 num_tokens: int,
101 hidden_size: int,
102 add_residual: bool,
103 dtype: torch.dtype,
104 seed: int,
105 device: str,
106 ) -> None:
107 set_random_seed(seed)
108 torch.set_default_device(device)
109 layer = RMSNorm(hidden_size, has_weight=False).to(dtype=dtype)
110 x = torch.randn(num_tokens, hidden_size, dtype=dtype)
111 residual = torch.randn_like(x) if add_residual else None
112
113 ref_out = layer.forward_native(x, residual)
114 out = layer(x, residual)
115 tol = _rms_norm_tolerance(dtype)
116 if add_residual:
117 torch.testing.assert_close(out[0], ref_out[0], **tol)
118 torch.testing.assert_close(out[1], ref_out[1], **tol)
119 else:
120 torch.testing.assert_close(out, ref_out, **tol)
121
122 if residual is not None:
123 opcheck(
124 torch.ops._C.fused_add_rms_norm,
125 (x, residual, None, layer.variance_epsilon),
126 )
127 else:
128 opcheck(
129 torch.ops._C.rms_norm,
130 (out, x, None, layer.variance_epsilon),
131 )
132
133
134 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
135 @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
136 @pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
137 @pytest.mark.parametrize("dtype", DTYPES)
138 @pytest.mark.parametrize("quant_scale", [0.01, 1.0, 10.0])
139 @pytest.mark.parametrize("seed", SEEDS)
140 @pytest.mark.parametrize("device", CUDA_DEVICES)
141 @pytest.mark.parametrize("strided_input", [False, True])
142 def test_fused_rms_norm_quant(
143 num_tokens: int,
144 hidden_size: int,
145 add_residual: bool,
146 dtype: torch.dtype,
147 quant_scale: float,
148 seed: int,
149 device: str,
150 strided_input: bool,
151 ) -> None:
152 set_random_seed(seed)
153 torch.set_default_device(device)
154
155 weight = torch.empty(hidden_size, dtype=dtype).normal_(mean=1.0, std=0.1)
156 scale = 1 / (2 * hidden_size)
157 last_dim = 2 * hidden_size if strided_input else hidden_size
158 x_base = torch.randn(num_tokens, last_dim, dtype=dtype)
159 x = x_base[..., :hidden_size]
160 assert x.is_contiguous() != strided_input
161
162 x *= scale
163 if add_residual:
164 residual = torch.randn_like(x) * scale
165 residual_fused = residual.clone()
166 else:
167 residual = residual_fused = None
168
169 out_norm = torch.empty_like(x)
170 out_quant = torch.empty_like(x, dtype=FP8_DTYPE)
171 out_quant_fused = torch.empty_like(out_quant)
172
173 quant_scale_t = torch.tensor(quant_scale, dtype=torch.float32)
174
175 if add_residual:
176 torch.ops._C.fused_add_rms_norm_static_fp8_quant(
177 out_quant_fused, x, residual_fused, weight, quant_scale_t, 1e-6
178 )
179
180 # Unfused kernel is in-place so it goes second
181 # Also use a separate clone of x to avoid modifying the input
182 x_unfused_base = x_base.clone()
183 x_unfused = x_unfused_base[..., :hidden_size]
184 assert x_unfused.is_contiguous() != strided_input
185 torch.ops._C.fused_add_rms_norm(x_unfused, residual, weight, 1e-6)
186 torch.ops._C.static_scaled_fp8_quant(
187 out_quant, x_unfused.contiguous(), quant_scale_t
188 )
189
190 torch.accelerator.synchronize()
191 torch.testing.assert_close(residual_fused, residual, atol=1e-2, rtol=1e-2)
192 opcheck(
193 torch.ops._C.fused_add_rms_norm_static_fp8_quant,
194 (out_quant_fused, x, residual_fused, weight, quant_scale_t, 1e-6),
195 )
196 else:
197 torch.ops._C.rms_norm_static_fp8_quant(
198 out_quant_fused, x, weight, quant_scale_t, 1e-6
199 )
200
201 torch.ops._C.rms_norm(out_norm, x, weight, 1e-6)
202 torch.ops._C.static_scaled_fp8_quant(out_quant, out_norm, quant_scale_t)
203
204 opcheck(
205 torch.ops._C.rms_norm_static_fp8_quant,
206 (out_quant_fused, x, weight, quant_scale_t, 1e-6),
207 )
208
209 if current_platform.is_rocm():
210 from vllm.platforms.rocm import on_gfx950
211
212 if on_gfx950() and dtype == torch.float16 and not add_residual:
213 # Fusion may round normalized FP16 across an E4M3 boundary on gfx950.
214 assert fp8_allclose(out_quant_fused, out_quant, rtol=0.125, atol=2e-3)
215 assert int(fp8_ulp_distance(out_quant_fused, out_quant).max()) <= 1
216 else:
217 # Fused and unfused FP8 paths can land on opposite sides of an E4M3
218 # tie; tolerate a tiny number of isolated fp8 outliers on ROCm.
219 ulp = fp8_ulp_distance(out_quant, out_quant_fused)
220 max_outliers = ulp.numel() // 100_000 + 8
221 num_outliers = int((ulp > 0).sum().item())
222 assert num_outliers <= max_outliers, (
223 f"FP8 quant mismatch: {num_outliers} fp8 outliers "
224 f"(allowed {max_outliers})"
225 )
226 else:
227 torch.testing.assert_close(
228 out_quant.to(dtype=torch.float32),
229 out_quant_fused.to(dtype=torch.float32),
230 atol=1e-3,
231 rtol=1e-3,
232 )
233
234
235 @torch.inference_mode()
236 def test_gemma_rms_norm_mixed_input_weight_dtype(default_vllm_config) -> None:
237 if not torch.cuda.is_available():
238 pytest.skip("CUDA required")
239
240 device = CUDA_DEVICES[0]
241 torch.set_default_device(device)
242
243 num_tokens, hidden_size = 32, 1024
244 x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
245 layer = GemmaRMSNorm(hidden_size, eps=1e-6).to(device=device)
246 layer.weight.data.normal_(mean=0.0, std=0.1)
247
248 # Gemma uses fp32 weight parameter while activations can be bf16.
249 assert layer.weight.dtype == torch.float32
250 out = layer(x)
251
252 x_fp32 = x.float()
253 weight_fp32 = layer.weight.data.float() + 1.0
254 variance = x_fp32.pow(2).mean(dim=-1, keepdim=True)
255 ref = (x_fp32 * torch.rsqrt(variance + layer.variance_epsilon) * weight_fp32).to(
256 x.dtype
257 )
258
259 assert out.dtype == x.dtype
260 torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)