| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.model_executor.custom_op import CustomOp |
| 8 | |
| 9 | from .common import rotate_gptj, rotate_neox |
| 10 | |
| 11 | |
| 12 | # --8<-- [start:dual_chunk_rotary_embedding] |
| 13 | @CustomOp.register("dual_chunk_rotary_embedding") |
| 14 | class DualChunkRotaryEmbedding(CustomOp): |
| 15 | """Rotary positional embedding for Dual Chunk Attention.""" |
| 16 | |
| 17 | # --8<-- [end:dual_chunk_rotary_embedding] |
| 18 | |
| 19 | def __init__( |
| 20 | self, |
| 21 | head_size: int, |
| 22 | rotary_dim: int, |
| 23 | max_position_embeddings: int, |
| 24 | base: float, |
| 25 | is_neox_style: bool, |
| 26 | dtype: torch.dtype, |
| 27 | chunk_size: int, |
| 28 | local_size: int, |
| 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.chunk_size = chunk_size |
| 37 | self.local_size = local_size |
| 38 | self.dtype = dtype |
| 39 | device_idx = torch.accelerator.current_device_index() |
| 40 | self.device = torch.device(f"cuda:{device_idx}") |
| 41 | (q_cache, qc_cache, k_cache, qc_no_clamp_cache, q_inter_cache) = ( |
| 42 | self._compute_cos_sin_cache() |
| 43 | ) |
| 44 | |
| 45 | self.register_buffer("cos_sin_q_cache", q_cache, persistent=False) |
| 46 | self.register_buffer("cos_sin_qc_cache", qc_cache, persistent=False) |
| 47 | self.register_buffer("cos_sin_k_cache", k_cache, persistent=False) |
| 48 | self.register_buffer( |
| 49 | "cos_sin_qc_no_clamp_cache", qc_no_clamp_cache, persistent=False |
| 50 | ) |
| 51 | self.register_buffer("cos_sin_q_inter_cache", q_inter_cache, persistent=False) |
| 52 | |
| 53 | def _compute_inv_freq(self, base: float) -> torch.Tensor: |
| 54 | """Compute the inverse frequency.""" |
| 55 | # NOTE(woosuk): The HF implementation uses `torch.arange(...).float()`. |
| 56 | # However, we use `torch.arange(..., dtype=torch.float)` instead to |
| 57 | # avoid numerical issues with large base values (e.g., 10000000). |
| 58 | # This may cause a slight numerical difference between the HF |
| 59 | # implementation and ours. |
| 60 | # NOTE(woosuk): To exactly match the HF implementation, we need to |
| 61 | # use CPU to compute the cache and then move it to GPU. However, we |
| 62 | # create the cache on GPU for faster initialization. This may cause |
| 63 | # a slight numerical difference between the HF implementation and ours. |
| 64 | inv_freq = 1.0 / ( |
| 65 | base |
| 66 | ** ( |
| 67 | torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim |
| 68 | ) |
| 69 | ) |
| 70 | return inv_freq |
| 71 | |
| 72 | def _compute_cos_sin_cache(self) -> torch.Tensor: |
| 73 | """Compute the cos and sin cache.""" |
| 74 | inv_freq = self._compute_inv_freq(self.base) |
| 75 | chunk_len = self.chunk_size - self.local_size |
| 76 | q_t = torch.arange(chunk_len, dtype=torch.float) |
| 77 | qc_t = (torch.arange(chunk_len, dtype=torch.float) + chunk_len).clamp( |
| 78 | max=self.chunk_size |
| 79 | ) |
| 80 | k_t = torch.arange(self.max_position_embeddings, dtype=torch.float) % chunk_len |
| 81 | |
| 82 | # count from chunk_len, no clamp(self.chunk_size) restriction |
| 83 | qc_no_clamp_t = torch.arange(chunk_len, dtype=torch.float) + chunk_len |
| 84 | # count from self.chunk_size for q_inter's rope |
| 85 | q_inter_t = torch.arange(chunk_len, dtype=torch.float) + self.chunk_size |
| 86 | |
| 87 | q_freqs = torch.outer(q_t, inv_freq) |
| 88 | qc_freqs = torch.outer(qc_t, inv_freq) |
| 89 | k_freqs = torch.outer(k_t, inv_freq) |
| 90 | qc_no_clamp_freqs = torch.outer(qc_no_clamp_t, inv_freq) |
| 91 | q_inter_freqs = torch.outer(q_inter_t, inv_freq) |
| 92 | |
| 93 | q_cos = q_freqs.cos() |
| 94 | q_sin = q_freqs.sin() |
| 95 | qc_cos = qc_freqs.cos() |
| 96 | qc_sin = qc_freqs.sin() |
| 97 | k_cos = k_freqs.cos() |
| 98 | k_sin = k_freqs.sin() |
| 99 | |
| 100 | qc_no_clamp_cos = qc_no_clamp_freqs.cos() |
| 101 | qc_no_clamp_sin = qc_no_clamp_freqs.sin() |
| 102 | q_inter_cos = q_inter_freqs.cos() |
| 103 | q_inter_sin = q_inter_freqs.sin() |
| 104 | |
| 105 | q_cache = torch.cat((q_cos, q_sin), dim=-1).to( |
| 106 | dtype=self.dtype, device=self.device |
| 107 | ) |
| 108 | qc_cache = torch.cat((qc_cos, qc_sin), dim=-1).to( |
| 109 | dtype=self.dtype, device=self.device |
| 110 | ) |
| 111 | k_cache = torch.cat((k_cos, k_sin), dim=-1).to( |
| 112 | dtype=self.dtype, device=self.device |
| 113 | ) |
| 114 | qc_no_clamp_cache = torch.cat((qc_no_clamp_cos, qc_no_clamp_sin), dim=-1).to( |
| 115 | dtype=self.dtype, device=self.device |
| 116 | ) |
| 117 | q_inter_cache = torch.cat((q_inter_cos, q_inter_sin), dim=-1).to( |
| 118 | dtype=self.dtype, device=self.device |
| 119 | ) |
| 120 | return q_cache, qc_cache, k_cache, qc_no_clamp_cache, q_inter_cache |
| 121 | |
| 122 | def forward_native( |
| 123 | self, |
| 124 | positions: torch.Tensor, |
| 125 | query: torch.Tensor, |
| 126 | key: torch.Tensor, |
| 127 | offsets: torch.Tensor | None = None, |
| 128 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 129 | query = query.view(*query.shape[:-1], -1, self.head_size) |
| 130 | key = key.view(*key.shape[:-1], -1, self.head_size) |
| 131 | query_rot = query[..., : self.rotary_dim] |
| 132 | key_rot = key[..., : self.rotary_dim] |
| 133 | if self.rotary_dim < self.head_size: |
| 134 | query_pass = query[..., self.rotary_dim :] |
| 135 | key_pass = key[..., self.rotary_dim :] |
| 136 | else: |
| 137 | query_pass = None |
| 138 | key_pass = None |
| 139 | |
| 140 | positions_with_offsets = ( |
| 141 | torch.add(positions, offsets) if offsets is not None else positions |
| 142 | ) |
| 143 | key = self._apply_rotary_embedding( |
| 144 | self.cos_sin_k_cache[positions_with_offsets], key_rot, key_pass |
| 145 | ) |
| 146 | chunk_len = self.chunk_size - self.local_size |
| 147 | query = self._apply_rotary_embedding( |
| 148 | self.cos_sin_q_cache[positions_with_offsets % chunk_len], |
| 149 | query_rot, |
| 150 | query_pass, |
| 151 | ) |
| 152 | query_succ = self._apply_rotary_embedding( |
| 153 | self.cos_sin_qc_cache[positions_with_offsets % chunk_len], |
| 154 | query_rot, |
| 155 | query_pass, |
| 156 | ) |
| 157 | query_inter = self._apply_rotary_embedding( |
| 158 | self.cos_sin_qc_cache[chunk_len - 1].repeat(positions.shape[0], 1), |
| 159 | query_rot, |
| 160 | query_pass, |
| 161 | ) |
| 162 | query_succ_critical = self._apply_rotary_embedding( |
| 163 | self.cos_sin_qc_no_clamp_cache[positions_with_offsets % chunk_len], |
| 164 | query_rot, |
| 165 | query_pass, |
| 166 | ) |
| 167 | query_inter_critical = self._apply_rotary_embedding( |
| 168 | self.cos_sin_q_inter_cache[positions_with_offsets % chunk_len], |
| 169 | query_rot, |
| 170 | query_pass, |
| 171 | ) |
| 172 | |
| 173 | # merge query into one tensor to simplify the interfaces |
| 174 | query = torch.cat( |
| 175 | ( |
| 176 | query, |
| 177 | query_succ, |
| 178 | query_inter, |
| 179 | query_succ_critical, |
| 180 | query_inter_critical, |
| 181 | ), |
| 182 | dim=-1, |
| 183 | ) |
| 184 | return query, key |
| 185 | |
| 186 | def forward_cuda( |
| 187 | self, |
| 188 | positions: torch.Tensor, |
| 189 | query: torch.Tensor, |
| 190 | key: torch.Tensor, |
| 191 | offsets: torch.Tensor | None = None, |
| 192 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 193 | return self.forward_native(positions, query, key, offsets) |
| 194 | |
| 195 | def _apply_rotary_embedding(self, cos_sin, hidden_rot, hidden_pass): |
| 196 | cos, sin = cos_sin.chunk(2, dim=-1) |
| 197 | if self.is_neox_style: |
| 198 | cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2) |
| 199 | sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2) |
| 200 | else: |
| 201 | cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2) |
| 202 | sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2) |
| 203 | rotate_fn = rotate_neox if self.is_neox_style else rotate_gptj |
| 204 | hidden_rot = hidden_rot * cos + rotate_fn(hidden_rot) * sin |
| 205 | |
| 206 | if self.rotary_dim < self.head_size: |
| 207 | hidden = torch.cat((hidden_rot, hidden_pass), dim=-1) |
| 208 | else: |
| 209 | hidden = hidden_rot |
| 210 | return hidden.flatten(-2).squeeze(0) |
| 211 | |
| 212 | def extra_repr(self) -> str: |
| 213 | s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}" |
| 214 | s += f", max_position_embeddings={self.max_position_embeddings}" |
| 215 | s += f", base={self.base}, is_neox_style={self.is_neox_style}" |
| 216 | s += f", chunk_size={self.chunk_size}, local_size={self.local_size}" |
| 217 | return s |