| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from .base import RotaryEmbeddingBase |
| 9 | |
| 10 | |
| 11 | class Llama4VisionRotaryEmbedding(RotaryEmbeddingBase): |
| 12 | def __init__( |
| 13 | self, |
| 14 | head_size: int, |
| 15 | rotary_dim: int, |
| 16 | max_position_embeddings: int, |
| 17 | base: float, |
| 18 | is_neox_style: bool, |
| 19 | dtype: torch.dtype, |
| 20 | ): |
| 21 | super().__init__( |
| 22 | head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype |
| 23 | ) |
| 24 | |
| 25 | def _compute_inv_freq(self, base: float) -> torch.Tensor: |
| 26 | inv_freqs = super()._compute_inv_freq(base) |
| 27 | inv_freqs = inv_freqs[: (self.rotary_dim // 2)] |
| 28 | return inv_freqs |
| 29 | |
| 30 | def _compute_cos_sin_cache(self) -> torch.Tensor: |
| 31 | inv_freq = self._compute_inv_freq(self.base) |
| 32 | |
| 33 | # self.max_position_embeddings here is number of image patches |
| 34 | # i.e. (image_size // patch_size) ** 2 |
| 35 | num_patches = self.max_position_embeddings |
| 36 | img_idx = torch.arange(num_patches, dtype=torch.int32).reshape(num_patches, 1) |
| 37 | img_idx = torch.cat([img_idx, img_idx[:1]], dim=0) |
| 38 | img_idx[-1, -1] = -2 # set to ID_CLS_TOKEN |
| 39 | num_patches_single_dim = int(math.sqrt(num_patches)) |
| 40 | frequencies_x = img_idx % num_patches_single_dim |
| 41 | frequencies_y = img_idx // num_patches_single_dim |
| 42 | freqs_x = ( |
| 43 | (frequencies_x + 1)[..., None] * inv_freq[None, None, :] |
| 44 | ).repeat_interleave(2, dim=-1) |
| 45 | freqs_y = ( |
| 46 | (frequencies_y + 1)[..., None] * inv_freq[None, None, :] |
| 47 | ).repeat_interleave(2, dim=-1) |
| 48 | freqs = torch.cat([freqs_x, freqs_y], dim=-1).float().contiguous()[..., ::2] |
| 49 | freqs = freqs.masked_fill(img_idx.reshape(-1, 1, 1) < 0, 0) |
| 50 | cache = torch.view_as_complex( |
| 51 | torch.stack([torch.cos(freqs), torch.sin(freqs)], dim=-1) |
| 52 | ) |
| 53 | return cache |
| 54 | |
| 55 | def forward_native( # type: ignore[override] |
| 56 | self, |
| 57 | query: torch.Tensor, |
| 58 | key: torch.Tensor | None = None, |
| 59 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 60 | assert key is not None |
| 61 | # self.cos_sin_cache here is complex tensor so we cannot cast into |
| 62 | # query's dtype directly with self._match_cos_sin_cache_dtype |
| 63 | |
| 64 | # NOTE: by not storing cos_sin_cache in self, we can avoid |
| 65 | # memory buffer update which is costly to runtime |
| 66 | cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(query.device) |
| 67 | query_ = torch.view_as_complex(query.float().reshape(*query.shape[:-1], -1, 2)) |
| 68 | key_ = torch.view_as_complex(key.float().reshape(*key.shape[:-1], -1, 2)) |
| 69 | broadcast_shape = [ |
| 70 | d if i == 1 or i == (query_.ndim - 1) else 1 |
| 71 | for i, d in enumerate(query_.shape) |
| 72 | ] |
| 73 | freqs_ci = cos_sin_cache.view(*broadcast_shape) |
| 74 | query_out = torch.view_as_real(query_ * freqs_ci).flatten(3) |
| 75 | key_out = torch.view_as_real(key_ * freqs_ci).flatten(3) |
| 76 | return query_out.type_as(query), key_out.type_as(key) |
| 77 | |
| 78 | def forward_cuda( # type: ignore[override] |
| 79 | self, |
| 80 | query: torch.Tensor, |
| 81 | key: torch.Tensor | None = None, |
| 82 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 83 | return self.forward_native(query, key) |