| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Rotary Positional Embeddings Base Class.""" |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm._aiter_ops import rocm_aiter_ops |
| 8 | from vllm.model_executor.custom_op import CustomOp |
| 9 | |
| 10 | from .common import ApplyRotaryEmb |
| 11 | |
| 12 | |
| 13 | # --8<-- [start:rotary_embedding] |
| 14 | @CustomOp.register("rotary_embedding") |
| 15 | class RotaryEmbeddingBase(CustomOp): |
| 16 | """Original rotary positional embedding.""" |
| 17 | |
| 18 | # --8<-- [end:rotary_embedding] |
| 19 | |
| 20 | def __init__( |
| 21 | self, |
| 22 | head_size: int, |
| 23 | rotary_dim: int, |
| 24 | max_position_embeddings: int, |
| 25 | base: float, |
| 26 | is_neox_style: bool, |
| 27 | dtype: torch.dtype, |
| 28 | init_cache: bool = True, |
| 29 | ) -> None: |
| 30 | super().__init__() |
| 31 | self.head_size = head_size |
| 32 | self.rotary_dim = rotary_dim |
| 33 | self.max_position_embeddings = max_position_embeddings |
| 34 | self.base = base |
| 35 | self.is_neox_style = is_neox_style |
| 36 | self.dtype = dtype |
| 37 | # TODO(mgoin): disabled for now due to failures |
| 38 | # Flashinfer only supports head_size=64, 128, 256, 512. |
| 39 | # https://github.com/flashinfer-ai/flashinfer/blob/ebfd655efe830048dba5d582aaa61d61d1cf9a87/include/flashinfer/utils.cuh#L174-L202 |
| 40 | # self.use_flashinfer = (self.enabled() |
| 41 | # and dtype in (torch.float16, torch.bfloat16) |
| 42 | # and current_platform.is_cuda() |
| 43 | # and has_flashinfer() |
| 44 | # and self.head_size in [64, 128, 256, 512]) |
| 45 | |
| 46 | # Check if use_flashinfer is already set |
| 47 | if not hasattr(self, "use_flashinfer"): |
| 48 | self.use_flashinfer = False |
| 49 | |
| 50 | self.use_aiter = ( |
| 51 | self.enabled() and rocm_aiter_ops.is_triton_rotary_embed_enabled() |
| 52 | ) |
| 53 | if self.use_aiter: |
| 54 | self.rocm_aiter_triton_rotary_embedding = ( |
| 55 | rocm_aiter_ops.get_triton_rotary_embedding_op() |
| 56 | ) |
| 57 | |
| 58 | if init_cache: |
| 59 | cache = self._compute_cos_sin_cache() |
| 60 | if not self.use_flashinfer: |
| 61 | cache = cache.to(dtype) |
| 62 | self.cos_sin_cache: torch.Tensor |
| 63 | self.register_buffer("cos_sin_cache", cache, persistent=False) |
| 64 | |
| 65 | # Reuse a precomputed bf16 cache for the AITER compile path. |
| 66 | if self.use_aiter and cache.dtype != torch.bfloat16: |
| 67 | self.cos_sin_cache_bf16: torch.Tensor | None |
| 68 | self.register_buffer( |
| 69 | "cos_sin_cache_bf16", |
| 70 | cache.to(torch.bfloat16), |
| 71 | persistent=False, |
| 72 | ) |
| 73 | else: |
| 74 | self.cos_sin_cache_bf16 = None |
| 75 | |
| 76 | self.apply_rotary_emb = ApplyRotaryEmb( |
| 77 | is_neox_style=self.is_neox_style, |
| 78 | ) |
| 79 | |
| 80 | def _compute_inv_freq(self, base: float) -> torch.Tensor: |
| 81 | """Compute the inverse frequency.""" |
| 82 | # NOTE(woosuk): To exactly match the HF implementation, we need to |
| 83 | # use CPU to compute the cache and then move it to GPU. However, we |
| 84 | # create the cache on GPU for faster initialization. This may cause |
| 85 | # a slight numerical difference between the HF implementation and ours. |
| 86 | inv_freq = 1.0 / ( |
| 87 | base |
| 88 | ** ( |
| 89 | torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim |
| 90 | ) |
| 91 | ) |
| 92 | return inv_freq |
| 93 | |
| 94 | def _compute_cos_sin_cache(self) -> torch.Tensor: |
| 95 | """Compute the cos and sin cache.""" |
| 96 | inv_freq = self._compute_inv_freq(self.base) |
| 97 | t = torch.arange(self.max_position_embeddings, dtype=torch.float) |
| 98 | |
| 99 | freqs = torch.einsum("i,j -> ij", t, inv_freq) |
| 100 | cos = freqs.cos() |
| 101 | sin = freqs.sin() |
| 102 | cache = torch.cat((cos, sin), dim=-1) |
| 103 | return cache |
| 104 | |
| 105 | def _match_cos_sin_cache_dtype(self, query: torch.Tensor) -> torch.Tensor: |
| 106 | # __setattr__ in nn.Module (called by `self.cos_sin_cache = ...`) |
| 107 | # is expensive, so avoid calling it if possible |
| 108 | cos_sin_cache = self.cos_sin_cache |
| 109 | if ( |
| 110 | cos_sin_cache.device == query.device |
| 111 | and self.cos_sin_cache.dtype == query.dtype |
| 112 | ): |
| 113 | return cos_sin_cache |
| 114 | |
| 115 | # Reuse precomputed bf16 cache in the AITER compile path. |
| 116 | if ( |
| 117 | self.use_aiter |
| 118 | and torch.compiler.is_compiling() |
| 119 | and query.dtype == torch.bfloat16 |
| 120 | ): |
| 121 | cache_bf16 = getattr(self, "cos_sin_cache_bf16", None) |
| 122 | if cache_bf16 is not None and cache_bf16.device == query.device: |
| 123 | return cache_bf16 |
| 124 | |
| 125 | cos_sin_cache = cos_sin_cache.to(query.device, dtype=query.dtype) |
| 126 | # Avoid mutating buffers during torch.compile (cudagraph) tracing. |
| 127 | if torch.compiler.is_compiling(): |
| 128 | return cos_sin_cache |
| 129 | |
| 130 | self.cos_sin_cache = cos_sin_cache |
| 131 | return cos_sin_cache |
| 132 | |
| 133 | def get_cos_sin(self, seqlen: int) -> tuple[torch.Tensor, torch.Tensor]: |
| 134 | cos_sin = self.cos_sin_cache[:seqlen] |
| 135 | cos, sin = cos_sin.chunk(2, dim=-1) |
| 136 | return cos, sin |
| 137 | |
| 138 | |
| 139 | class RotaryEmbedding(RotaryEmbeddingBase): |
| 140 | def __init__( |
| 141 | self, |
| 142 | head_size: int, |
| 143 | rotary_dim: int, |
| 144 | max_position_embeddings: int, |
| 145 | base: float, |
| 146 | is_neox_style: bool, |
| 147 | dtype: torch.dtype, |
| 148 | init_cache: bool = True, |
| 149 | ) -> None: |
| 150 | super().__init__( |
| 151 | head_size=head_size, |
| 152 | rotary_dim=rotary_dim, |
| 153 | max_position_embeddings=max_position_embeddings, |
| 154 | base=base, |
| 155 | is_neox_style=is_neox_style, |
| 156 | dtype=dtype, |
| 157 | init_cache=init_cache, |
| 158 | ) |
| 159 | |
| 160 | @staticmethod |
| 161 | def forward_static( |
| 162 | positions: torch.Tensor, |
| 163 | query: torch.Tensor, |
| 164 | key: torch.Tensor | None, |
| 165 | head_size: int, |
| 166 | rotary_dim: int, |
| 167 | cos_sin_cache: torch.Tensor, |
| 168 | is_neox_style: bool, |
| 169 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 170 | """A PyTorch-native implementation of forward().""" |
| 171 | positions = positions.flatten() |
| 172 | num_tokens = positions.shape[0] |
| 173 | cos_sin = cos_sin_cache.index_select(0, positions) |
| 174 | cos, sin = cos_sin.chunk(2, dim=-1) |
| 175 | |
| 176 | query_shape = query.shape |
| 177 | query = query.view(num_tokens, -1, head_size) |
| 178 | query_rot = query[..., :rotary_dim] |
| 179 | query_pass = query[..., rotary_dim:] |
| 180 | query_rot = ApplyRotaryEmb.forward_static( |
| 181 | query_rot, |
| 182 | cos, |
| 183 | sin, |
| 184 | is_neox_style, |
| 185 | ) |
| 186 | query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape) |
| 187 | |
| 188 | # key may be None in some cases, e.g. cross-layer KV sharing |
| 189 | if key is not None: |
| 190 | key_shape = key.shape |
| 191 | key = key.view(num_tokens, -1, head_size) |
| 192 | key_rot = key[..., :rotary_dim] |
| 193 | key_pass = key[..., rotary_dim:] |
| 194 | key_rot = ApplyRotaryEmb.forward_static( |
| 195 | key_rot, |
| 196 | cos, |
| 197 | sin, |
| 198 | is_neox_style, |
| 199 | ) |
| 200 | key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) |
| 201 | return query, key |
| 202 | |
| 203 | def forward_native( |
| 204 | self, |
| 205 | positions: torch.Tensor, |
| 206 | query: torch.Tensor, |
| 207 | key: torch.Tensor | None = None, |
| 208 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 209 | """A PyTorch-native implementation of forward().""" |
| 210 | cos_sin_cache = self._match_cos_sin_cache_dtype(query) |
| 211 | return self.forward_static( |
| 212 | positions, |
| 213 | query, |
| 214 | key, |
| 215 | self.head_size, |
| 216 | self.rotary_dim, |
| 217 | cos_sin_cache, |
| 218 | self.is_neox_style, |
| 219 | ) |
| 220 | |
| 221 | def forward_cuda( |
| 222 | self, |
| 223 | positions: torch.Tensor, |
| 224 | query: torch.Tensor, |
| 225 | key: torch.Tensor | None = None, |
| 226 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 227 | if self.use_flashinfer: |
| 228 | torch.ops.vllm.flashinfer_rotary_embedding( |
| 229 | positions, |
| 230 | query, |
| 231 | key, |
| 232 | self.head_size, |
| 233 | self.cos_sin_cache, |
| 234 | self.is_neox_style, |
| 235 | ) |
| 236 | return query, key |
| 237 | |
| 238 | from vllm import _custom_ops as ops |
| 239 | |
| 240 | cos_sin_cache = self._match_cos_sin_cache_dtype(query) |
| 241 | |
| 242 | # ops.rotary_embedding() is an in-place operation |
| 243 | # that updates the query and key tensors. |
| 244 | ops.rotary_embedding( |
| 245 | positions, |
| 246 | query, |
| 247 | key, |
| 248 | self.head_size, |
| 249 | cos_sin_cache, |
| 250 | self.is_neox_style, |
| 251 | ) |
| 252 | return query, key |
| 253 | |
| 254 | def forward_hip( |
| 255 | self, |
| 256 | positions: torch.Tensor, |
| 257 | query: torch.Tensor, |
| 258 | key: torch.Tensor | None = None, |
| 259 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 260 | if self.use_aiter: |
| 261 | cos_sin_cache = self._match_cos_sin_cache_dtype(query) |
| 262 | self.rocm_aiter_triton_rotary_embedding( |
| 263 | positions, |
| 264 | query, |
| 265 | key, |
| 266 | self.head_size, |
| 267 | cos_sin_cache, |
| 268 | self.is_neox_style, |
| 269 | ) |
| 270 | return query, key |
| 271 | return self.forward_cuda(positions, query, key) |
| 272 | |
| 273 | def forward_xpu( |
| 274 | self, |
| 275 | positions: torch.Tensor, |
| 276 | query: torch.Tensor, |
| 277 | key: torch.Tensor | None = None, |
| 278 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 279 | self._match_cos_sin_cache_dtype(query) |
| 280 | # ops.rotary_embedding() is an in-place operation |
| 281 | # that updates the query and key tensors. |
| 282 | if key is None: |
| 283 | return self.forward_native(positions, query, key) |
| 284 | else: |
| 285 | from vllm import _custom_ops as ops |
| 286 | |
| 287 | cos_sin_cache = self._match_cos_sin_cache_dtype(query) |
| 288 | ops.rotary_embedding( |
| 289 | positions, |
| 290 | query, |
| 291 | key, |
| 292 | self.head_size, |
| 293 | cos_sin_cache, |
| 294 | self.is_neox_style, |
| 295 | ) |
| 296 | return query, key |
| 297 | |
| 298 | def forward_cpu( |
| 299 | self, |
| 300 | positions: torch.Tensor, |
| 301 | query: torch.Tensor, |
| 302 | key: torch.Tensor | None = None, |
| 303 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 304 | from vllm import _custom_ops as ops |
| 305 | |
| 306 | cos_sin_cache = self._match_cos_sin_cache_dtype(query) |
| 307 | |
| 308 | # ops.rotary_embedding() is an in-place operation |
| 309 | # that updates the query and key tensors. |
| 310 | ops.rotary_embedding( |
| 311 | positions, |
| 312 | query, |
| 313 | key, |
| 314 | self.head_size, |
| 315 | cos_sin_cache, |
| 316 | self.is_neox_style, |
| 317 | ) |
| 318 | return query, key |
| 319 | |
| 320 | def extra_repr(self) -> str: |
| 321 | s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}" |
| 322 | s += f", max_position_embeddings={self.max_position_embeddings}" |
| 323 | s += f", base={self.base}, is_neox_style={self.is_neox_style}" |
| 324 | return s |