| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from .mrope import MRotaryEmbedding |
| 8 | |
| 9 | |
| 10 | class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding): |
| 11 | """3D rotary positional embedding. 3D is t:time h:height w:width.""" |
| 12 | |
| 13 | def forward_native( # type: ignore[override] |
| 14 | self, |
| 15 | positions: torch.Tensor, |
| 16 | query: torch.Tensor, |
| 17 | key: torch.Tensor | None = None, |
| 18 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 19 | assert positions.ndim == 1 or positions.ndim == 2 |
| 20 | assert key is not None |
| 21 | |
| 22 | num_tokens = positions.shape[-1] |
| 23 | cos_sin = self.cos_sin_cache[positions] |
| 24 | cos, sin = cos_sin.chunk(2, dim=-1) |
| 25 | if positions.ndim == 2: |
| 26 | assert self.mrope_section |
| 27 | |
| 28 | section_h = self.mrope_section[0] # 22 |
| 29 | section_w = self.mrope_section[1] # 22 |
| 30 | section_t = self.mrope_section[2] # 20 |
| 31 | assert section_h == section_w |
| 32 | # Split according to [h w h w h w h w... t t t...] |
| 33 | section_cos_t = cos[..., -section_t:] |
| 34 | section_cos_h = cos[..., : section_h + section_w : 2] |
| 35 | section_cos_w = cos[..., 1 : section_h + section_w : 2] |
| 36 | |
| 37 | cos_t, cos_h, cos_w = section_cos_t[0], section_cos_h[1], section_cos_w[2] |
| 38 | cos_hw = torch.stack([cos_h, cos_w], dim=-1).reshape( |
| 39 | cos_h.shape[:-1] + (cos_h.shape[-1] * 2,) |
| 40 | ) |
| 41 | cos = torch.cat([cos_hw, cos_t], dim=-1) |
| 42 | |
| 43 | section_sin_t = sin[..., -section_t:] |
| 44 | section_sin_h = sin[..., : section_h + section_w : 2] |
| 45 | section_sin_w = sin[..., 1 : section_h + section_w : 2] |
| 46 | |
| 47 | sin_t, sin_h, sin_w = section_sin_t[0], section_sin_h[1], section_sin_w[2] |
| 48 | sin_hw = torch.stack([sin_h, sin_w], dim=-1).reshape( |
| 49 | sin_h.shape[:-1] + (sin_h.shape[-1] * 2,) |
| 50 | ) |
| 51 | sin = torch.cat([sin_hw, sin_t], dim=-1) |
| 52 | |
| 53 | query_shape = query.shape |
| 54 | query = query.view(num_tokens, -1, self.head_size) |
| 55 | query_rot = query[..., : self.rotary_dim] |
| 56 | query_pass = query[..., self.rotary_dim :] |
| 57 | query_rot = self.apply_rotary_emb.forward_native( |
| 58 | query_rot, |
| 59 | cos, |
| 60 | sin, |
| 61 | ) |
| 62 | query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape) |
| 63 | |
| 64 | key_shape = key.shape |
| 65 | key = key.view(num_tokens, -1, self.head_size) |
| 66 | key_rot = key[..., : self.rotary_dim] |
| 67 | key_pass = key[..., self.rotary_dim :] |
| 68 | key_rot = self.apply_rotary_emb.forward_native( |
| 69 | key_rot, |
| 70 | cos, |
| 71 | sin, |
| 72 | ) |
| 73 | key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) |
| 74 | return query, key |
| 75 | |
| 76 | def forward_cuda( # type: ignore[override] |
| 77 | self, |
| 78 | positions: torch.Tensor, |
| 79 | query: torch.Tensor, |
| 80 | key: torch.Tensor | None = None, |
| 81 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 82 | return self.forward_native(positions, query, key) |
| 83 | |
| 84 | def forward_xpu( # type: ignore[override] |
| 85 | self, |
| 86 | positions: torch.Tensor, |
| 87 | query: torch.Tensor, |
| 88 | key: torch.Tensor | None = None, |
| 89 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 90 | # No fused XPU kernel for this 3D t/h/w rope; base |
| 91 | # MRotaryEmbedding.forward_xpu forwards an extra `offsets` arg that |
| 92 | # this class's forward_cuda override doesn't accept. Use native path. |
| 93 | return self.forward_native(positions, query, key) |