| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Gemma4-specific Rotary Positional Embeddings (proportional scaling). |
| 4 | |
| 5 | Gemma4 uses "proportional" RoPE which computes inv_freq frequencies scaled |
| 6 | by head_dim (not rotary_dim), and zero-pads for non-rotated dimensions when |
| 7 | partial_rotary_factor < 1. The actual rotation uses standard neox-style |
| 8 | rotate_half, matching HF transformers' apply_rotary_pos_emb. |
| 9 | """ |
| 10 | |
| 11 | import torch |
| 12 | |
| 13 | from .base import RotaryEmbedding |
| 14 | |
| 15 | |
| 16 | class Gemma4RotaryEmbedding(RotaryEmbedding): |
| 17 | """Gemma4 proportional RoPE. |
| 18 | |
| 19 | Extends RotaryEmbedding (which provides standard neox-style rotation |
| 20 | via ops.rotary_embedding CUDA kernel) but overrides the inv_freq |
| 21 | computation to match HF's _compute_proportional_rope_parameters: |
| 22 | - Frequency exponents use head_dim (not rotary_dim) as denominator |
| 23 | - Non-rotated dims are zero-padded (cos=1, sin=0 = identity rotation) |
| 24 | |
| 25 | When partial_rotary_factor=1.0 (the default for some variants), ALL dims are |
| 26 | rotated and this is equivalent to standard RotaryEmbedding with |
| 27 | head_dim-scaled frequencies. |
| 28 | """ |
| 29 | |
| 30 | def __init__( |
| 31 | self, |
| 32 | head_size: int, |
| 33 | rotary_dim: int, |
| 34 | max_position_embeddings: int, |
| 35 | base: float, |
| 36 | is_neox_style: bool, |
| 37 | dtype: torch.dtype, |
| 38 | ) -> None: |
| 39 | # Number of rotation angle pairs (from partial_rotary_factor) |
| 40 | self.rope_angles = rotary_dim // 2 |
| 41 | # Non-rotated angle pairs per half |
| 42 | self.nope_angles = (head_size // 2) - self.rope_angles |
| 43 | |
| 44 | # Important: set rotary_dim = head_size so the base class's |
| 45 | # forward_static applies rotation to ALL dims of the cos/sin cache. |
| 46 | # The non-rotated dims will have cos=1, sin=0 (identity) thanks |
| 47 | # to our _compute_inv_freq zero-padding. |
| 48 | super().__init__( |
| 49 | head_size, |
| 50 | head_size, # rotary_dim = head_size (full application) |
| 51 | max_position_embeddings, |
| 52 | base, |
| 53 | is_neox_style, |
| 54 | dtype, |
| 55 | ) |
| 56 | |
| 57 | def _compute_inv_freq(self, base: float) -> torch.Tensor: |
| 58 | """Compute frequencies matching HF proportional RoPE. |
| 59 | |
| 60 | Key difference from base: exponent denominator is head_size (not |
| 61 | rotary_dim), and non-rotated dims are zero-padded. |
| 62 | """ |
| 63 | # HF formula: base ** (arange(0, 2*rope_angles, 2) / head_dim) |
| 64 | freq_exponents = ( |
| 65 | torch.arange(0, 2 * self.rope_angles, 2, dtype=torch.float) / self.head_size |
| 66 | ) |
| 67 | inv_freq = 1.0 / (base**freq_exponents) |
| 68 | |
| 69 | # Zero-pad for non-rotated dims (identity rotation: cos=1, sin=0) |
| 70 | if self.nope_angles > 0: |
| 71 | inv_freq = torch.cat( |
| 72 | [ |
| 73 | inv_freq, |
| 74 | torch.zeros(self.nope_angles, dtype=torch.float), |
| 75 | ] |
| 76 | ) |
| 77 | return inv_freq |
| 78 | |
| 79 | def extra_repr(self) -> str: |
| 80 | s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}" |
| 81 | s += f", rope_angles={self.rope_angles}, nope_angles={self.nope_angles}" |
| 82 | s += f", max_position_embeddings={self.max_position_embeddings}" |
| 83 | s += f", base={self.base}, is_neox_style={self.is_neox_style}" |
| 84 | return s |