@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 193 lines 5.99 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Callable
5 from itertools import product
6
7 import pytest
8 import torch
9
10 from tests.kernels.allclose_default import get_default_atol, get_default_rtol
11 from vllm.model_executor.layers.rotary_embedding import get_rope
12 from vllm.utils.torch_utils import set_random_seed
13
14 IS_NEOX_STYLE = [True, False]
15 DTYPES = [torch.bfloat16, torch.float]
16 HEAD_SIZES = [64, 80, 120, 256]
17 ROTARY_DIMS = [None, 32] # None means rotary dim == head size
18 NUM_HEADS = [17] # Arbitrary values for testing
19 BATCH_SIZES = [5] # Arbitrary values for testing
20 SEQ_LENS = [11, 8192] # Arbitrary values for testing
21 SEEDS = [0]
22 CUDA_DEVICES = [
23 f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
24 ]
25 USE_KEY = [True, False]
26
27
28 def _get_flat_tensor_shape(
29 batch_size: int, seq_len: int, num_heads: int, head_size: int
30 ) -> tuple[int, ...]:
31 return (batch_size, seq_len, num_heads * head_size)
32
33
34 # For testing sliced tensors
35 def _get_padded_tensor_shape(
36 batch_size: int, seq_len: int, num_heads: int, head_size: int
37 ) -> tuple[int, ...]:
38 return (batch_size, seq_len, num_heads, head_size + 64)
39
40
41 def _get_batch_tensor_shape(
42 batch_size: int, seq_len: int, num_heads: int, head_size: int
43 ) -> tuple[int, ...]:
44 return (batch_size, seq_len, num_heads, head_size)
45
46
47 TENSORS_SHAPES_FN = [
48 _get_batch_tensor_shape,
49 _get_flat_tensor_shape,
50 _get_padded_tensor_shape,
51 ]
52
53
54 @pytest.mark.parametrize("is_neox_style", IS_NEOX_STYLE)
55 @pytest.mark.parametrize("tensor_shape_fn", TENSORS_SHAPES_FN)
56 @pytest.mark.parametrize("batch_size", BATCH_SIZES)
57 @pytest.mark.parametrize("seq_len", SEQ_LENS)
58 @pytest.mark.parametrize("num_heads", NUM_HEADS)
59 @pytest.mark.parametrize("head_size", HEAD_SIZES)
60 @pytest.mark.parametrize("rotary_dim", ROTARY_DIMS)
61 @pytest.mark.parametrize("dtype", DTYPES)
62 @pytest.mark.parametrize("seed", SEEDS)
63 @pytest.mark.parametrize("device", CUDA_DEVICES)
64 @pytest.mark.parametrize("use_key", USE_KEY)
65 @torch.inference_mode()
66 def test_rotary_embedding(
67 default_vllm_config,
68 is_neox_style: bool,
69 tensor_shape_fn: Callable[[int, int, int, int], tuple[int, ...]],
70 batch_size: int,
71 seq_len: int,
72 num_heads: int,
73 head_size: int,
74 rotary_dim: int | None,
75 dtype: torch.dtype,
76 seed: int,
77 device: str,
78 use_key: bool,
79 max_position: int = 8192,
80 rope_theta: float = 10000,
81 ) -> None:
82 if rotary_dim is None:
83 rotary_dim = head_size
84
85 set_random_seed(seed)
86 torch.set_default_device(device)
87 if rotary_dim is None:
88 rotary_dim = head_size
89 rope_parameters = {
90 "rope_type": "default",
91 "rope_theta": rope_theta,
92 "partial_rotary_factor": rotary_dim / head_size,
93 }
94 rope = get_rope(head_size, max_position, is_neox_style, rope_parameters)
95 rope = rope.to(dtype=dtype, device=torch.get_default_device())
96
97 positions = torch.randint(0, max_position, (batch_size, seq_len))
98 query_shape = tensor_shape_fn(batch_size, seq_len, num_heads, head_size)
99 # slice tensor if required, noop otherwise
100 query = torch.randn(query_shape, dtype=dtype)[..., :head_size]
101 key = torch.randn_like(query)[..., :head_size] if use_key else None
102
103 # NOTE(woosuk): The reference implementation should be executed first
104 # because the custom kernel is in-place.
105 ref_query, ref_key = rope.forward_native(positions, query, key)
106 out_query, out_key = rope.forward(positions, query, key)
107 # Compare the results.
108 torch.testing.assert_close(
109 out_query,
110 ref_query,
111 atol=get_default_atol(out_query),
112 rtol=get_default_rtol(out_query),
113 )
114 if use_key:
115 torch.testing.assert_close(
116 out_key,
117 ref_key,
118 atol=get_default_atol(out_key),
119 rtol=get_default_rtol(out_key),
120 )
121 else:
122 assert ref_key is None and out_key is None, "expected returned key to be None"
123
124
125 @torch.inference_mode()
126 def test_rope_module_cache(default_vllm_config):
127 MAX_POSITIONS = [123, 1234]
128 ROPE_THETAS = [10000, 1000000]
129 ROPE_PARAMETERS = (
130 {"rope_type": "default"},
131 {"rope_type": "linear", "factor": (1,)},
132 {"rope_type": "dynamic", "factor": 1},
133 )
134 settings = (
135 HEAD_SIZES,
136 ROTARY_DIMS,
137 MAX_POSITIONS,
138 ROPE_THETAS,
139 IS_NEOX_STYLE,
140 ROPE_PARAMETERS,
141 DTYPES,
142 )
143 rope_setting_id_map: dict[str, int] = {}
144 for setting in product(*settings):
145 (
146 head_size,
147 rotary_dim,
148 max_position,
149 rope_theta,
150 is_neox_style,
151 rope_parameters,
152 dtype,
153 ) = setting
154 if rotary_dim is None:
155 rotary_dim = head_size
156 rope_parameters["rope_theta"] = rope_theta
157 rope_parameters["partial_rotary_factor"] = rotary_dim / head_size
158 rope = get_rope(
159 head_size,
160 max_position,
161 is_neox_style,
162 rope_parameters,
163 dtype,
164 )
165 # different settings cannot share the same rope module
166 assert id(rope) not in rope_setting_id_map.values()
167 assert all(x.dtype == dtype for x in rope.buffers())
168 assert all(x.dtype == dtype for x in rope.parameters())
169 rope_setting_id_map[str(setting)] = id(rope)
170
171 for setting in product(*settings):
172 (
173 head_size,
174 rotary_dim,
175 max_position,
176 rope_theta,
177 is_neox_style,
178 rope_parameters,
179 dtype,
180 ) = setting
181 if rotary_dim is None:
182 rotary_dim = head_size
183 rope_parameters["rope_theta"] = rope_theta
184 rope_parameters["partial_rotary_factor"] = rotary_dim / head_size
185 rope = get_rope(
186 head_size,
187 max_position,
188 is_neox_style,
189 rope_parameters,
190 dtype,
191 )
192 # check if cache take effect
193 assert id(rope) == rope_setting_id_map[str(setting)]