| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | from contextlib import suppress |
| 6 | from importlib import import_module |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.model_executor.custom_op import CustomOp |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.utils.torch_utils import direct_register_custom_op |
| 14 | |
| 15 | logger = init_logger(__name__) |
| 16 | |
| 17 | |
| 18 | # common functions |
| 19 | def rotate_neox(x: torch.Tensor) -> torch.Tensor: |
| 20 | x1 = x[..., : x.shape[-1] // 2] |
| 21 | x2 = x[..., x.shape[-1] // 2 :] |
| 22 | return torch.cat((-x2, x1), dim=-1) |
| 23 | |
| 24 | |
| 25 | def rotate_gptj(x: torch.Tensor) -> torch.Tensor: |
| 26 | x1 = x[..., ::2] |
| 27 | x2 = x[..., 1::2] |
| 28 | x = torch.stack((-x2, x1), dim=-1) |
| 29 | return x.flatten(-2) |
| 30 | |
| 31 | |
| 32 | # yarn functions |
| 33 | # Inverse dim formula to find dim based on number of rotations |
| 34 | def yarn_find_correction_dim( |
| 35 | num_rotations: int, |
| 36 | dim: int, |
| 37 | base: float = 10000, |
| 38 | max_position_embeddings: int = 2048, |
| 39 | ) -> float: |
| 40 | return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / ( |
| 41 | 2 * math.log(base) |
| 42 | ) |
| 43 | |
| 44 | |
| 45 | # Find dim range bounds based on rotations |
| 46 | def yarn_find_correction_range( |
| 47 | low_rot: int, |
| 48 | high_rot: int, |
| 49 | dim: int, |
| 50 | base: float = 10000, |
| 51 | max_position_embeddings: int = 2048, |
| 52 | truncate: bool = True, |
| 53 | ) -> tuple[float | int, float | int]: |
| 54 | low = yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings) |
| 55 | high = yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings) |
| 56 | if truncate: |
| 57 | low = math.floor(low) |
| 58 | high = math.ceil(high) |
| 59 | return max(low, 0), min(high, dim - 1) # Clamp values just in case |
| 60 | |
| 61 | |
| 62 | def yarn_linear_ramp_mask( |
| 63 | low: float, high: float, dim: int, dtype: torch.dtype |
| 64 | ) -> torch.Tensor: |
| 65 | if low == high: |
| 66 | high += 0.001 # Prevent singularity |
| 67 | |
| 68 | linear_func = (torch.arange(dim, dtype=dtype) - low) / (high - low) |
| 69 | ramp_func = torch.clamp(linear_func, 0, 1) |
| 70 | return ramp_func |
| 71 | |
| 72 | |
| 73 | def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: |
| 74 | if scale <= 1: |
| 75 | return 1.0 |
| 76 | return 0.1 * mscale * math.log(scale) + 1.0 |
| 77 | |
| 78 | |
| 79 | def _flashinfer_rotary_embedding( |
| 80 | positions: torch.Tensor, |
| 81 | query: torch.Tensor, |
| 82 | key: torch.Tensor, |
| 83 | head_size: int, |
| 84 | cos_sin_cache: torch.Tensor, |
| 85 | is_neox: bool, |
| 86 | ) -> None: |
| 87 | """Custom op wrapper for flashinfer's rotary embedding. |
| 88 | |
| 89 | This is an in-place operation that modifies query and key tensors directly. |
| 90 | """ |
| 91 | from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace |
| 92 | |
| 93 | apply_rope_with_cos_sin_cache_inplace( |
| 94 | positions=positions, |
| 95 | query=query, |
| 96 | key=key, |
| 97 | head_size=head_size, |
| 98 | cos_sin_cache=cos_sin_cache, |
| 99 | is_neox=is_neox, |
| 100 | ) |
| 101 | |
| 102 | |
| 103 | # Register flashinfer rotary embedding custom op |
| 104 | direct_register_custom_op( |
| 105 | op_name="flashinfer_rotary_embedding", |
| 106 | op_func=_flashinfer_rotary_embedding, |
| 107 | mutates_args=["query", "key"], # These tensors are modified in-place |
| 108 | ) |
| 109 | |
| 110 | |
| 111 | # --8<-- [start:apply_rotary_emb] |
| 112 | @CustomOp.register("apply_rotary_emb") |
| 113 | class ApplyRotaryEmb(CustomOp): |
| 114 | # --8<-- [end:apply_rotary_emb] |
| 115 | |
| 116 | def __init__( |
| 117 | self, |
| 118 | enforce_enable: bool = False, |
| 119 | is_neox_style: bool = True, |
| 120 | enable_fp32_compute: bool = False, |
| 121 | ) -> None: |
| 122 | super().__init__(enforce_enable=enforce_enable) |
| 123 | self.is_neox_style = is_neox_style |
| 124 | self.enable_fp32_compute = enable_fp32_compute |
| 125 | |
| 126 | self.apply_rotary_emb_flash_attn = None |
| 127 | if not current_platform.is_cpu(): |
| 128 | with suppress(ModuleNotFoundError): |
| 129 | self.apply_rotary_emb_flash_attn = import_module( |
| 130 | "flash_attn.ops.triton.rotary" |
| 131 | ).apply_rotary |
| 132 | |
| 133 | @staticmethod |
| 134 | def forward_static( |
| 135 | x: torch.Tensor, |
| 136 | cos: torch.Tensor, |
| 137 | sin: torch.Tensor, |
| 138 | is_neox_style: bool = True, |
| 139 | enable_fp32_compute: bool = False, |
| 140 | ) -> torch.Tensor: |
| 141 | """Args: |
| 142 | x: [batch_size (optional), seq_len, num_heads, head_size] |
| 143 | cos: [seq_len, head_size // 2] |
| 144 | sin: [seq_len, head_size // 2] |
| 145 | is_neox_style: Whether to use the Neox-style or GPT-J-style. |
| 146 | enable_fp32_compute: Temporarily convert x, cos, sin to FP32 dtype |
| 147 | for higher accuracy. |
| 148 | |
| 149 | """ |
| 150 | origin_dtype = x.dtype |
| 151 | if enable_fp32_compute: |
| 152 | x = x.float() |
| 153 | |
| 154 | cos = cos.unsqueeze(-2).to(x.dtype) |
| 155 | sin = sin.unsqueeze(-2).to(x.dtype) |
| 156 | |
| 157 | if is_neox_style: |
| 158 | x1, x2 = torch.chunk(x, 2, dim=-1) |
| 159 | else: |
| 160 | x1 = x[..., ::2] |
| 161 | x2 = x[..., 1::2] |
| 162 | |
| 163 | o1 = x1 * cos - x2 * sin |
| 164 | o2 = x2 * cos + x1 * sin |
| 165 | |
| 166 | if is_neox_style: |
| 167 | output = torch.cat((o1, o2), dim=-1) |
| 168 | else: |
| 169 | output = torch.stack((o1, o2), dim=-1).flatten(-2) |
| 170 | |
| 171 | if enable_fp32_compute: |
| 172 | output = output.to(origin_dtype) |
| 173 | return output |
| 174 | |
| 175 | def _pre_process( |
| 176 | self, |
| 177 | x: torch.Tensor, |
| 178 | cos: torch.Tensor, |
| 179 | sin: torch.Tensor, |
| 180 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Size, torch.dtype]: |
| 181 | origin_shape = x.shape |
| 182 | if len(origin_shape) == 3: |
| 183 | # x: [seq_len, num_heads, head_size] |
| 184 | x = x.unsqueeze(0) |
| 185 | |
| 186 | origin_dtype = x.dtype |
| 187 | if self.enable_fp32_compute: |
| 188 | x = x.float() |
| 189 | cos = cos.float() |
| 190 | sin = sin.float() |
| 191 | |
| 192 | return x, cos, sin, origin_shape, origin_dtype |
| 193 | |
| 194 | def _post_process( |
| 195 | self, |
| 196 | output: torch.Tensor, |
| 197 | origin_shape: torch.Size, |
| 198 | origin_dtype: torch.dtype, |
| 199 | ) -> torch.Tensor: |
| 200 | if len(origin_shape) == 3: |
| 201 | output = output.squeeze(0) |
| 202 | if self.enable_fp32_compute: |
| 203 | output = output.to(origin_dtype) |
| 204 | return output |
| 205 | |
| 206 | def forward_native( |
| 207 | self, |
| 208 | x: torch.Tensor, |
| 209 | cos: torch.Tensor, |
| 210 | sin: torch.Tensor, |
| 211 | ) -> torch.Tensor: |
| 212 | output = self.forward_static( |
| 213 | x, cos, sin, self.is_neox_style, self.enable_fp32_compute |
| 214 | ) |
| 215 | return output |
| 216 | |
| 217 | def forward_cuda( |
| 218 | self, |
| 219 | x: torch.Tensor, |
| 220 | cos: torch.Tensor, |
| 221 | sin: torch.Tensor, |
| 222 | ) -> torch.Tensor: |
| 223 | from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb |
| 224 | |
| 225 | x, cos, sin, origin_shape, origin_dtype = self._pre_process(x, cos, sin) |
| 226 | |
| 227 | """ |
| 228 | Arguments of apply_rotary_emb() in vllm_flash_attn: |
| 229 | x: [batch_size, seq_len, nheads, headdim] |
| 230 | cos, sin: [seqlen_rotary, rotary_dim / 2] |
| 231 | interleaved: default as False (Neox-style). |
| 232 | ... |
| 233 | """ |
| 234 | interleaved = not self.is_neox_style |
| 235 | output = apply_rotary_emb(x, cos, sin, interleaved) |
| 236 | |
| 237 | output = self._post_process(output, origin_shape, origin_dtype) |
| 238 | return output |
| 239 | |
| 240 | def forward_xpu( |
| 241 | self, |
| 242 | x: torch.Tensor, |
| 243 | cos: torch.Tensor, |
| 244 | sin: torch.Tensor, |
| 245 | ) -> torch.Tensor: |
| 246 | import vllm._xpu_ops # noqa: F401 registers torch.ops.vllm.xpu_apply_rotary_emb |
| 247 | |
| 248 | x, cos, sin, origin_shape, origin_dtype = self._pre_process(x, cos, sin) |
| 249 | output = torch.ops.vllm.xpu_apply_rotary_emb(x, cos, sin, self.is_neox_style) |
| 250 | return self._post_process(output, origin_shape, origin_dtype) |
| 251 | |
| 252 | def forward_hip( |
| 253 | self, |
| 254 | x: torch.Tensor, |
| 255 | cos: torch.Tensor, |
| 256 | sin: torch.Tensor, |
| 257 | ) -> torch.Tensor: |
| 258 | _HIP_MAX_GRID_DIM = 65535 |
| 259 | """ |
| 260 | HIP/ROCm has a per-dim grid limit of 65535 on gridY/gridZ. The |
| 261 | flash_attn triton rotary kernel uses |
| 262 | grid = (cdiv(nheads, BLOCK_H), cdiv(seq_len, BLOCK_M), batch) |
| 263 | with BLOCK_M=8 (rotary_dim<=128) or BLOCK_M=4 (otherwise) and |
| 264 | BLOCK_H=2. When the visual encoder packs many image patches into one |
| 265 | batch (e.g. vLLM profile_run with max_num_seqs images), gridY can |
| 266 | exceed 65535 and hipModuleLaunchKernel returns |
| 267 | `Triton Error [HIP]: Code: 1, invalid argument`. Fall back to the |
| 268 | native PyTorch implementation in that case. |
| 269 | """ |
| 270 | if self.apply_rotary_emb_flash_attn is not None: |
| 271 | x, cos, sin, origin_shape, origin_dtype = self._pre_process(x, cos, sin) |
| 272 | |
| 273 | seq_len = x.shape[-3] |
| 274 | batch = x.shape[0] |
| 275 | rotary_dim = cos.shape[-1] * 2 |
| 276 | block_m = 8 if rotary_dim <= 128 else 4 |
| 277 | grid_y = (seq_len + block_m - 1) // block_m |
| 278 | if grid_y > _HIP_MAX_GRID_DIM or batch > _HIP_MAX_GRID_DIM: |
| 279 | output = self.forward_static( |
| 280 | x, cos, sin, self.is_neox_style, self.enable_fp32_compute |
| 281 | ) |
| 282 | return self._post_process(output, origin_shape, origin_dtype) |
| 283 | |
| 284 | """ |
| 285 | Arguments of apply_rotary() in flash_attn: |
| 286 | x: [batch_size, seq_len, nheads, headdim] |
| 287 | cos, sin: [seqlen_rotary, rotary_dim / 2] |
| 288 | interleaved: default as False (Neox-style). |
| 289 | ... |
| 290 | """ |
| 291 | interleaved = not self.is_neox_style |
| 292 | output = self.apply_rotary_emb_flash_attn( |
| 293 | x, cos, sin, interleaved=interleaved |
| 294 | ).type_as(x) |
| 295 | |
| 296 | output = self._post_process(output, origin_shape, origin_dtype) |
| 297 | else: |
| 298 | # Falling back to PyTorch native implementation. |
| 299 | output = self.forward_native(x, cos, sin) |
| 300 | |
| 301 | return output |
| 302 | |
| 303 | def forward_cpu( |
| 304 | self, |
| 305 | x: torch.Tensor, |
| 306 | cos: torch.Tensor, |
| 307 | sin: torch.Tensor, |
| 308 | ) -> torch.Tensor: |
| 309 | # TODO (bigPYJ1151): need to enable fused CPU ROPE here |
| 310 | return self.forward_native(x, cos, sin) |
| 311 | |
| 312 | def extra_repr(self) -> str: |
| 313 | s = f"is_neox_style={self.is_neox_style}" |
| 314 | s += f", enable_fp32_compute={self.enable_fp32_compute}" |
| 315 | return s |