A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[FEAT] Refactor ROPE into module (#22192)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
Aug 4, 2025 at 19:12 UTC
6ad6b8e115b8b46ad918284d862bdadded3af447
15 files changed
+2111
-1967
vllm/model_executor/layers/rotary_embedding.py
deleted
-1967
@@ -1,1967 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-# Adapted from
5
-# https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/llama/modeling_llama.py
6
-# Copyright 2023 The vLLM team.
7
-# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
8
-#
9
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
10
-# and OPT implementations in this library. It has been modified from its
11
-# original forms to accommodate minor architectural differences compared
12
-# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13
-#
14
-# Licensed under the Apache License, Version 2.0 (the "License");
15
-# you may not use this file except in compliance with the License.
16
-# You may obtain a copy of the License at
17
-#
18
-# http://www.apache.org/licenses/LICENSE-2.0
19
-#
20
-# Unless required by applicable law or agreed to in writing, software
21
-# distributed under the License is distributed on an "AS IS" BASIS,
22
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23
-# See the License for the specific language governing permissions and
24
-# limitations under the License.
25
-"""Rotary Positional Embeddings."""
26
-import itertools
27
-import math
28
-from typing import Any, Optional, Union
29
-
30
-import numpy as np
31
-import torch
32
-import torch.nn as nn
33
-from transformers import PretrainedConfig
34
-
35
-from vllm.model_executor.custom_op import CustomOp
36
-from vllm.platforms import current_platform
37
-
38
-if current_platform.is_cuda():
39
- from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
40
-
41
-
42
-def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
43
- x1 = x[..., :x.shape[-1] // 2]
44
- x2 = x[..., x.shape[-1] // 2:]
45
- return torch.cat((-x2, x1), dim=-1)
46
-
47
-
48
-def _rotate_gptj(x: torch.Tensor) -> torch.Tensor:
49
- x1 = x[..., ::2]
50
- x2 = x[..., 1::2]
51
- x = torch.stack((-x2, x1), dim=-1)
52
- return x.flatten(-2)
53
-
54
-
55
-def _apply_rotary_emb_torch(
56
- x: torch.Tensor,
57
- cos: torch.Tensor,
58
- sin: torch.Tensor,
59
- is_neox_style: bool,
60
-) -> torch.Tensor:
61
- cos = cos.unsqueeze(-2).to(x.dtype)
62
- sin = sin.unsqueeze(-2).to(x.dtype)
63
- if is_neox_style:
64
- x1, x2 = torch.chunk(x, 2, dim=-1)
65
- else:
66
- x1 = x[..., ::2]
67
- x2 = x[..., 1::2]
68
- o1 = x1 * cos - x2 * sin
69
- o2 = x2 * cos + x1 * sin
70
- if is_neox_style:
71
- return torch.cat((o1, o2), dim=-1)
72
- else:
73
- return torch.stack((o1, o2), dim=-1).flatten(-2)
74
-
75
-
76
-def _apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
77
- is_neox_style: bool) -> torch.Tensor:
78
- """
79
- Args:
80
- x: [num_tokens, num_heads, head_size]
81
- cos: [num_tokens, head_size // 2]
82
- sin: [num_tokens, head_size // 2]
83
- is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
84
- positional embeddings.
85
- """
86
- if current_platform.is_cuda():
87
- return apply_rotary_emb(x.unsqueeze(0), cos, sin,
88
- not is_neox_style).squeeze(0)
89
- else:
90
- return _apply_rotary_emb_torch(x, cos, sin, is_neox_style)
91
-
92
-
93
-@CustomOp.register("rotary_embedding")
94
-class RotaryEmbedding(CustomOp):
95
- """Original rotary positional embedding."""
96
-
97
- def __init__(
98
- self,
99
- head_size: int,
100
- rotary_dim: int,
101
- max_position_embeddings: int,
102
- base: float,
103
- is_neox_style: bool,
104
- dtype: torch.dtype,
105
- ) -> None:
106
- super().__init__()
107
- self.head_size = head_size
108
- self.rotary_dim = rotary_dim
109
- self.max_position_embeddings = max_position_embeddings
110
- self.base = base
111
- self.is_neox_style = is_neox_style
112
- self.dtype = dtype
113
-
114
- cache = self._compute_cos_sin_cache()
115
- cache = cache.to(dtype)
116
- self.cos_sin_cache: torch.Tensor
117
- self.register_buffer("cos_sin_cache", cache, persistent=False)
118
-
119
- def _compute_inv_freq(self, base: float) -> torch.Tensor:
120
- """Compute the inverse frequency."""
121
- # NOTE(woosuk): To exactly match the HF implementation, we need to
122
- # use CPU to compute the cache and then move it to GPU. However, we
123
- # create the cache on GPU for faster initialization. This may cause
124
- # a slight numerical difference between the HF implementation and ours.
125
- inv_freq = 1.0 / (base**(torch.arange(
126
- 0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim))
127
- return inv_freq
128
-
129
- def _compute_cos_sin_cache(self) -> torch.Tensor:
130
- """Compute the cos and sin cache."""
131
- inv_freq = self._compute_inv_freq(self.base)
132
- t = torch.arange(self.max_position_embeddings, dtype=torch.float)
133
-
134
- freqs = torch.einsum("i,j -> ij", t, inv_freq)
135
- cos = freqs.cos()
136
- sin = freqs.sin()
137
- cache = torch.cat((cos, sin), dim=-1)
138
- return cache
139
-
140
- def forward_native(
141
- self,
142
- positions: torch.Tensor,
143
- query: torch.Tensor,
144
- key: Optional[torch.Tensor] = None,
145
- offsets: Optional[torch.Tensor] = None,
146
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
147
- """A PyTorch-native implementation of forward()."""
148
- if offsets is not None:
149
- positions = positions + offsets
150
- positions = positions.flatten()
151
- num_tokens = positions.shape[0]
152
- cos_sin = self.cos_sin_cache.index_select(0, positions)
153
- cos, sin = cos_sin.chunk(2, dim=-1)
154
-
155
- query_shape = query.shape
156
- query = query.view(num_tokens, -1, self.head_size)
157
- query_rot = query[..., :self.rotary_dim]
158
- query_pass = query[..., self.rotary_dim:]
159
- query_rot = _apply_rotary_emb_torch(query_rot, cos, sin,
160
- self.is_neox_style)
161
- query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
162
-
163
- # key may be None in some cases, e.g. cross-layer KV sharing
164
- if key is not None:
165
- key_shape = key.shape
166
- key = key.view(num_tokens, -1, self.head_size)
167
- key_rot = key[..., :self.rotary_dim]
168
- key_pass = key[..., self.rotary_dim:]
169
- key_rot = _apply_rotary_emb_torch(key_rot, cos, sin,
170
- self.is_neox_style)
171
- key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
172
- return query, key
173
-
174
- def forward_cuda(
175
- self,
176
- positions: torch.Tensor,
177
- query: torch.Tensor,
178
- key: Optional[torch.Tensor] = None,
179
- offsets: Optional[torch.Tensor] = None,
180
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
181
- from vllm import _custom_ops as ops
182
-
183
- # __setattr__ in nn.Module (called by `self.cos_sin_cache = ...`)
184
- # is expensive, so avoid calling it if possible
185
- if self.cos_sin_cache.device != query.device or \
186
- self.cos_sin_cache.dtype != query.dtype:
187
- self.cos_sin_cache = self.cos_sin_cache.to(query.device,
188
- dtype=query.dtype)
189
-
190
- # ops.rotary_embedding()/batched_rotary_embedding()
191
- # are in-place operations that update the query and key tensors.
192
- if offsets is not None:
193
- ops.batched_rotary_embedding(positions, query, key, self.head_size,
194
- self.cos_sin_cache,
195
- self.is_neox_style, self.rotary_dim,
196
- offsets)
197
- else:
198
- ops.rotary_embedding(positions, query, key, self.head_size,
199
- self.cos_sin_cache, self.is_neox_style)
200
- return query, key
201
-
202
- def forward_xpu(
203
- self,
204
- positions: torch.Tensor,
205
- query: torch.Tensor,
206
- key: Optional[torch.Tensor] = None,
207
- offsets: Optional[torch.Tensor] = None,
208
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
209
- from vllm._ipex_ops import ipex_ops as ops
210
-
211
- self.cos_sin_cache = self.cos_sin_cache.to(positions.device,
212
- dtype=query.dtype)
213
- # ops.rotary_embedding()/batched_rotary_embedding()
214
- # are in-place operations that update the query and key tensors.
215
- if key is None:
216
- # XPU kernel doesn't support key=None so fall back to native impl
217
- # TODO(sarckk): add support for optional key in
218
- # ipex.llm.functional.rotary_embedding_batched
219
- return self.forward_native(positions, query, key, offsets)
220
- else:
221
- if offsets is not None:
222
- ops.batched_rotary_embedding(positions, query, key,
223
- self.head_size,
224
- self.cos_sin_cache,
225
- self.is_neox_style,
226
- self.rotary_dim, offsets)
227
- else:
228
- ops.rotary_embedding(positions, query, key, self.head_size,
229
- self.cos_sin_cache, self.is_neox_style)
230
- return query, key
231
-
232
- def forward_neuron(
233
- self,
234
- positions: torch.Tensor,
235
- query: torch.Tensor,
236
- key: Optional[torch.Tensor] = None,
237
- offsets: Optional[torch.Tensor] = None,
238
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
239
-
240
- def _apply_rotary_emb_neuron(
241
- x: torch.Tensor,
242
- cos: torch.Tensor,
243
- sin: torch.Tensor,
244
- is_neox_style: bool,
245
- ) -> torch.Tensor:
246
- cos = cos.unsqueeze(-2).to(x.dtype)
247
- sin = sin.unsqueeze(-2).to(x.dtype)
248
- if is_neox_style:
249
- x1, x2 = torch.chunk(x, 2, dim=-1)
250
- else:
251
- # x1 = x[..., ::2]
252
-
253
- # x2 = x[..., 1::2]
254
- d = x.shape[-1] // 2
255
- x_reshaped = x.view(-1, x.shape[-1])
256
- x1 = x_reshaped[:, ::2].view(*x.shape[:-1], d)
257
- x2 = x_reshaped[:, 1::2].view(*x.shape[:-1], d)
258
- o1 = x1 * cos - x2 * sin
259
- o2 = x2 * cos + x1 * sin
260
- if is_neox_style:
261
- return torch.cat((o1, o2), dim=-1)
262
- else:
263
- return torch.stack((o1, o2), dim=-1).flatten(-2)
264
-
265
- if offsets is not None:
266
- positions = positions + offsets
267
-
268
- self.cos_sin_cache = self.cos_sin_cache.to(query.device,
269
- dtype=query.dtype)
270
-
271
- positions = positions.flatten()
272
- num_tokens = positions.shape[0]
273
- cos_sin = self.cos_sin_cache.index_select(0, positions)
274
- cos, sin = cos_sin.chunk(2, dim=-1)
275
-
276
- query_shape = query.shape
277
- query = query.view(num_tokens, -1, self.head_size)
278
- if key is not None:
279
- key_shape = key.shape
280
- key = key.view(num_tokens, -1, self.head_size)
281
-
282
- if self.rotary_dim == self.head_size:
283
- query = _apply_rotary_emb(query, cos, sin, self.is_neox_style)
284
- query = query.reshape(query_shape)
285
- if key is not None:
286
- key = _apply_rotary_emb(key, cos, sin, self.is_neox_style)
287
- key = key.reshape(key_shape)
288
- else:
289
- head_size = query.shape[-1]
290
- query_reshaped = query.view(-1, head_size)
291
- query_pass = query_reshaped[:, self.rotary_dim:].view(
292
- *query.shape[:-1], head_size - self.rotary_dim)
293
- query_rot = query_reshaped[:, :self.rotary_dim].view(
294
- *query.shape[:-1], self.rotary_dim)
295
- query_rot = _apply_rotary_emb_neuron(query_rot, cos, sin,
296
- self.is_neox_style)
297
- query = torch.cat((query_rot, query_pass),
298
- dim=-1).reshape(query_shape)
299
-
300
- if key is not None:
301
- key_reshaped = key.view(-1, head_size)
302
- key_pass = key_reshaped[:, self.rotary_dim:].view(
303
- *key.shape[:-1], head_size - self.rotary_dim)
304
- key_rot = key_reshaped[:, :self.rotary_dim].view(
305
- *key.shape[:-1], self.rotary_dim)
306
- key_rot = _apply_rotary_emb_neuron(key_rot, cos, sin,
307
- self.is_neox_style)
308
- key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
309
- return query, key
310
-
311
- def extra_repr(self) -> str:
312
- s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
313
- s += f", max_position_embeddings={self.max_position_embeddings}"
314
- s += f", base={self.base}, is_neox_style={self.is_neox_style}"
315
- return s
316
-
317
-
318
-class LinearScalingRotaryEmbedding(RotaryEmbedding):
319
- """RotaryEmbedding extended with linear scaling.
320
-
321
- It supports multiple scaling factors. Since multiple LoRA adapters may have
322
- different scaling factors, we need multiple cos/sin caches. In this way,
323
- instead of running rotary embedding kernel per lora, we can run multiple
324
- lora in a batched way.
325
-
326
- In addition to that, we also keep the cos/sin cache for the scaling factor
327
- of 1 (default) at all times.
328
-
329
- Exemplary for two scaling factors x=1, y and z with embeddings
330
- [[x11, x12, ... x1m], ..., [xn1, xn2, ..., xnm]] and
331
- [[y11, y12, ... y1o], ..., [yn1, yn2, ..., yno]], and
332
- [[z11, z12, ... z1p], ..., [zn1, zn2, ..., znp]],
333
-
334
- we construct the cos/sin cache as follows:
335
- [[x11, x12, ... x1m, y11, y12, ... y1o, z11, z12, ... z1p],
336
- ...
337
- [xn1, xn2, ... xnm, yn1, yn2, ... yno, zn1, zn2, ... znp]]
338
-
339
- We then use offsets to index into the cos/sin cache for
340
- the respective scaling factors.
341
-
342
- The offset to cache can be accessed via `scaling_factor_to_offset` API.
343
-
344
- Credits to the Reddit user /u/kaiokendev
345
- """
346
-
347
- def __init__(
348
- self,
349
- head_size: int,
350
- rotary_dim: int,
351
- max_position_embeddings: int,
352
- base: float,
353
- is_neox_style: bool,
354
- scaling_factors: Union[list[float], float],
355
- dtype: torch.dtype,
356
- ) -> None:
357
- if isinstance(scaling_factors, float):
358
- scaling_factors = [scaling_factors]
359
- self.scaling_factors: list[float] = scaling_factors # noqa
360
- super().__init__(head_size, rotary_dim, max_position_embeddings, base,
361
- is_neox_style, dtype)
362
- # Lazy initialized.
363
- self._scaling_factor_to_offset: dict[float, int]
364
-
365
- def _compute_cos_sin_cache(self) -> torch.Tensor:
366
- inv_freq = self._compute_inv_freq(self.base)
367
- cache_list: list[torch.Tensor] = []
368
- # offsets to the next cache in a tensor.
369
- # Each offset corresponds to the same index in scaling_factors.
370
- offsets: list[int] = []
371
- for scaling_factor in self.scaling_factors:
372
- # NOTE(woosuk): self.max_position_embeddings is the original
373
- # maximum length before applying the rope scaling.
374
- # Thus, the maximum length after applying the rope scaling is
375
- # self.max_position_embeddings * self.scaling_factor.
376
- max_len = self.max_position_embeddings * scaling_factor
377
- t = torch.arange(max_len, dtype=torch.float)
378
- t = t / scaling_factor
379
-
380
- freqs = torch.einsum("i,j -> ij", t, inv_freq)
381
- cos = freqs.cos()
382
- sin = freqs.sin()
383
- cache = torch.cat((cos, sin), dim=-1)
384
- if not cache_list:
385
- offset = 0
386
- else:
387
- last_offset = offsets[-1]
388
- next_max_len = cache_list[-1].shape[0]
389
- offset = last_offset + next_max_len
390
- offsets.append(offset)
391
- cache_list.append(cache)
392
- self._scaling_factor_to_offset = {
393
- float(scaling_factor): offsets[i]
394
- for i, scaling_factor in enumerate(self.scaling_factors)
395
- }
396
- assert len(self.scaling_factors) == len(offsets)
397
- return torch.cat(cache_list, dim=0)
398
-
399
- @property
400
- def scaling_factor_to_offset(self) -> dict[float, int]:
401
- return self._scaling_factor_to_offset
402
-
403
-
404
-class NTKScalingRotaryEmbedding(RotaryEmbedding):
405
- """RotaryEmbedding extended with fixed and mixed NTK scaling.
406
- https://kexue.fm/archives/9706 """
407
-
408
- def __init__(self,
409
- head_size: int,
410
- rotary_dim: int,
411
- max_position_embeddings: int,
412
- base: float,
413
- is_neox_style: bool,
414
- scaling_factor: float,
415
- dtype: torch.dtype,
416
- mixed_b: Optional[float] = None) -> None:
417
- self.scaling_factor = scaling_factor
418
- self.mixed_b = mixed_b
419
- super().__init__(head_size, rotary_dim, max_position_embeddings, base,
420
- is_neox_style, dtype)
421
-
422
- def _compute_inv_freq(self, base: float) -> torch.Tensor:
423
- base = self.base * (self.scaling_factor if self.mixed_b is None else 1)
424
- inv_freq = super()._compute_inv_freq(base)
425
-
426
- if self.mixed_b is None:
427
- inv_freq = inv_freq / self.scaling_factor**(2 / self.rotary_dim)
428
- else:
429
- a = torch.tensor(self.scaling_factor).log() / (self.rotary_dim /
430
- 2)**self.mixed_b
431
- lambda_1_m = (a * torch.arange(
432
- 1, self.rotary_dim // 2 + 1).float()**self.mixed_b).exp()
433
- inv_freq = inv_freq / lambda_1_m
434
-
435
- return inv_freq
436
-
437
-
438
-class DynamicNTKScalingRotaryEmbedding(RotaryEmbedding):
439
- """RotaryEmbedding extended with Dynamic NTK scaling.
440
-
441
- Credits to the Reddit users /u/bloc97 and /u/emozilla
442
- """
443
-
444
- def __init__(
445
- self,
446
- head_size: int,
447
- rotary_dim: int,
448
- max_position_embeddings: int,
449
- base: float,
450
- is_neox_style: bool,
451
- scaling_factor: float,
452
- dtype: torch.dtype,
453
- ) -> None:
454
- self.scaling_factor = scaling_factor
455
- super().__init__(head_size, rotary_dim, max_position_embeddings, base,
456
- is_neox_style, dtype)
457
-
458
- def _compute_cos_sin_cache(self) -> torch.Tensor:
459
- # NOTE(woosuk): self.max_position_embeddings is the original
460
- # maximum length before applying the rope scaling.
461
- # Thus, the maximum length after applying the rope scaling is
462
- # self.max_position_embeddings * self.scaling_factor.
463
- max_len = self.max_position_embeddings * self.scaling_factor
464
- base = self.base * (
465
- (self.scaling_factor * max_len / self.max_position_embeddings) -
466
- (self.scaling_factor - 1))**(self.rotary_dim /
467
- (self.rotary_dim - 2))
468
- inv_freq = self._compute_inv_freq(base)
469
- t = torch.arange(max_len, dtype=torch.float)
470
-
471
- freqs = torch.einsum("i,j -> ij", t, inv_freq)
472
- cos = freqs.cos()
473
- sin = freqs.sin()
474
- cache = torch.cat((cos, sin), dim=-1)
475
- return cache
476
-
477
-
478
-class DynamicNTKAlphaRotaryEmbedding(RotaryEmbedding):
479
- """RotaryEmbedding extended with Dynamic NTK alpha.
480
-
481
- Based on the original RotaryEmbedding implementation.
482
- """
483
-
484
- def __init__(
485
- self,
486
- head_size: int,
487
- rotary_dim: int,
488
- max_position_embeddings: int,
489
- base: float,
490
- is_neox_style: bool,
491
- scaling_alpha: float,
492
- dtype: torch.dtype,
493
- ) -> None:
494
- self.scaling_alpha = scaling_alpha
495
- super().__init__(head_size, rotary_dim, max_position_embeddings, base,
496
- is_neox_style, dtype)
497
-
498
- def _compute_cos_sin_cache(self) -> torch.Tensor:
499
- # For Hunyuan DynamicNTKAlphaRotaryEmbedding
500
- max_len = self.max_position_embeddings
501
- base = self.base * self.scaling_alpha**(self.rotary_dim /
502
- (self.rotary_dim - 2))
503
- inv_freq = self._compute_inv_freq(base)
504
- t = torch.arange(max_len, dtype=torch.float)
505
-
506
- freqs = torch.einsum("i,j -> ij", t, inv_freq)
507
- cos = freqs.cos()
508
- sin = freqs.sin()
509
- cache = torch.cat((cos, sin), dim=-1)
510
- return cache
511
-
512
-
513
-# Inverse dim formula to find dim based on number of rotations
514
-def _yarn_find_correction_dim(num_rotations: int,
515
- dim: int,
516
- base: float = 10000,
517
- max_position_embeddings: int = 2048) -> float:
518
- return (dim * math.log(max_position_embeddings /
519
- (num_rotations * 2 * math.pi))) / (2 *
520
- math.log(base))
521
-
522
-
523
-# Find dim range bounds based on rotations
524
-def _yarn_find_correction_range(
525
- low_rot: int,
526
- high_rot: int,
527
- dim: int,
528
- base: float = 10000,
529
- max_position_embeddings: int = 2048) -> tuple[int, int]:
530
- low = math.floor(
531
- _yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings))
532
- high = math.ceil(
533
- _yarn_find_correction_dim(high_rot, dim, base,
534
- max_position_embeddings))
535
- return max(low, 0), min(high, dim - 1) # Clamp values just in case
536
-
537
-
538
-def _yarn_linear_ramp_mask(low: float, high: float, dim: int,
539
- dtype: torch.dtype) -> torch.Tensor:
540
- if low == high:
541
- high += 0.001 # Prevent singularity
542
-
543
- linear_func = (torch.arange(dim, dtype=dtype) - low) / (high - low)
544
- ramp_func = torch.clamp(linear_func, 0, 1)
545
- return ramp_func
546
-
547
-
548
-def _yarn_get_mscale(scale: float = 1) -> float:
549
- if scale <= 1:
550
- return 1.0
551
- return 0.1 * math.log(scale) + 1.0
552
-
553
-
554
-class YaRNScalingRotaryEmbedding(RotaryEmbedding):
555
- """RotaryEmbedding extended with YaRN method.
556
-
557
- Credits to Peng et al. github.com/jquesnelle/yarn
558
- """
559
-
560
- def __init__(
561
- self,
562
- head_size: int,
563
- rotary_dim: int,
564
- max_position_embeddings: int,
565
- base: float,
566
- is_neox_style: bool,
567
- scaling_factor: float,
568
- dtype: torch.dtype,
569
- *,
570
- extrapolation_factor: float = 1,
571
- attn_factor: float = 1,
572
- beta_fast: int = 32,
573
- beta_slow: int = 1,
574
- ) -> None:
575
- self.scaling_factor = scaling_factor
576
- self.extrapolation_factor = extrapolation_factor
577
- self.attn_factor = attn_factor
578
- self.beta_fast = beta_fast
579
- self.beta_slow = beta_slow
580
- # Get n-d magnitude scaling corrected for interpolation
581
- self.mscale = float(
582
- _yarn_get_mscale(self.scaling_factor) * attn_factor)
583
- super().__init__(head_size, rotary_dim, max_position_embeddings, base,
584
- is_neox_style, dtype)
585
-
586
- def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
587
- pos_freqs = self.base**(
588
- torch.arange(0, self.rotary_dim, 2, dtype=torch.float) /
589
- self.rotary_dim)
590
- inv_freq_extrapolation = 1.0 / pos_freqs
591
- inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
592
-
593
- low, high = _yarn_find_correction_range(self.beta_fast, self.beta_slow,
594
- self.rotary_dim, self.base,
595
- self.max_position_embeddings)
596
- # Get n-d rotational scaling corrected for extrapolation
597
- inv_freq_mask = (1 - _yarn_linear_ramp_mask(
598
- low, high, self.rotary_dim // 2,
599
- dtype=torch.float)) * self.extrapolation_factor
600
- inv_freq = inv_freq_interpolation * (
601
- 1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask
602
- return inv_freq
603
-
604
- def _compute_cos_sin_cache(self) -> torch.Tensor:
605
- inv_freq = self._compute_inv_freq(self.scaling_factor)
606
- t = torch.arange(self.max_position_embeddings * self.scaling_factor,
607
- dtype=torch.float32)
608
- freqs = torch.einsum("i,j -> ij", t, inv_freq)
609
- cos = (freqs.cos() * self.mscale)
610
- sin = (freqs.sin() * self.mscale)
611
- cache = torch.cat((cos, sin), dim=-1)
612
- return cache
613
-
614
-
615
-class Phi3LongRoPEScaledRotaryEmbedding(nn.Module):
616
- """Phi3 family of models scaled rotary embedding.
617
-
618
- Based on the original RotaryEmbedding implementation.
619
- """
620
-
621
- def __init__(
622
- self,
623
- head_size: int,
624
- rotary_dim: int,
625
- max_position_embeddings: int,
626
- original_max_position_embeddings: int,
627
- base: float,
628
- is_neox_style: bool,
629
- dtype: torch.dtype,
630
- short_factor: list[float],
631
- long_factor: list[float],
632
- short_mscale: Optional[float] = None,
633
- long_mscale: Optional[float] = None,
634
- ):
635
- super().__init__()
636
-
637
- if is_neox_style is False:
638
- raise ValueError(
639
- "`Phi3LongRoPEScaledRotaryEmbedding` only supports neox_style."
640
- )
641
-
642
- self.rotary_dim = rotary_dim
643
- self.head_size = head_size
644
- self.max_position_embeddings = max_position_embeddings
645
- self.original_max_position_embeddings = original_max_position_embeddings
646
- self.base = base
647
- self.short_factor = short_factor
648
- self.long_factor = long_factor
649
-
650
- scale = self.max_position_embeddings / \
651
- self.original_max_position_embeddings
652
- if scale <= 1.0:
653
- scaling_factor = 1.0
654
- else:
655
- scaling_factor = math.sqrt(
656
- 1 + math.log(scale) /
657
- math.log(self.original_max_position_embeddings))
658
- if short_mscale is None:
659
- short_mscale = scaling_factor
660
- if long_mscale is None:
661
- long_mscale = scaling_factor
662
-
663
- self.short_mscale = short_mscale
664
- self.long_mscale = long_mscale
665
-
666
- short_cache = self._compute_cos_sin_cache(
667
- original_max_position_embeddings, short_factor, short_mscale)
668
- short_cache = short_cache.to(dtype)
669
-
670
- long_cache = self._compute_cos_sin_cache(max_position_embeddings,
671
- long_factor, long_mscale)
672
- long_cache = long_cache.to(dtype)
673
-
674
- long_short_cache = torch.cat([short_cache, long_cache], dim=0)
675
- self.register_buffer("long_short_cos_sin_cache",
676
- long_short_cache,
677
- persistent=False)
678
-
679
- def _compute_inv_freq(self, rescale_factors: list[float]) -> torch.Tensor:
680
- rescale_factors = torch.tensor(rescale_factors, dtype=torch.float32)
681
- inv_freq = 1.0 / (rescale_factors * (self.base**(torch.arange(
682
- 0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim)))
683
- return inv_freq
684
-
685
- def _compute_cos_sin_cache(
686
- self,
687
- max_position_embeddings: int,
688
- rescale_factors: list[float],
689
- mscale: float,
690
- ) -> torch.Tensor:
691
- inv_freq = self._compute_inv_freq(rescale_factors)
692
- t = torch.arange(max_position_embeddings, dtype=torch.float)
693
- freqs = torch.einsum("i,j -> ij", t, inv_freq)
694
- cos = freqs.cos() * mscale
695
- sin = freqs.sin() * mscale
696
- cache = torch.cat((cos, sin), dim=-1)
697
- return cache
698
-
699
- def forward(
700
- self,
701
- positions: torch.Tensor,
702
- query: torch.Tensor,
703
- key: Optional[torch.Tensor] = None,
704
- offsets: Optional[torch.Tensor] = None,
705
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
706
- assert key is not None
707
- query = query.view(*query.shape[:-1], -1, self.head_size)
708
- key = key.view(*key.shape[:-1], -1, self.head_size)
709
-
710
- k = self.original_max_position_embeddings
711
- long_prompt_offset = (torch.any(positions > k).float() *
712
- torch.full_like(positions, k)).long()
713
- idx = (torch.add(positions, long_prompt_offset)
714
- if long_prompt_offset is not None else positions)
715
- idx = torch.add(idx, offsets) if offsets is not None else idx
716
- cos_sin = torch.index_select(self.long_short_cos_sin_cache, 0, idx)
717
-
718
- cos, sin = cos_sin.chunk(2, dim=-1)
719
- cos = cos.repeat(1, 2).unsqueeze(-2)
720
- sin = sin.repeat(1, 2).unsqueeze(-2)
721
-
722
- query_rot = query[..., :self.rotary_dim]
723
- query_pass = query[..., self.rotary_dim:]
724
- query_rot = query_rot * cos + _rotate_neox(query_rot) * sin
725
- query = torch.cat((query_rot, query_pass), dim=-1)
726
-
727
- key_rot = key[..., :self.rotary_dim]
728
- key_pass = key[..., self.rotary_dim:]
729
- key_rot = key_rot * cos + _rotate_neox(key_rot) * sin
730
- key = torch.cat((key_rot, key_pass), dim=-1)
731
-
732
- return query.flatten(-2), key.flatten(-2)
733
-
734
-
735
-def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
736
- if scale <= 1:
737
- return 1.0
738
- return 0.1 * mscale * math.log(scale) + 1.0
739
-
740
-
741
-class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
742
- """RotaryEmbedding extended with YaRN method.
743
-
744
- Credits to Peng et al. github.com/jquesnelle/yarn
745
- """
746
-
747
- def __init__(
748
- self,
749
- head_size: int,
750
- rotary_dim: int,
751
- max_position_embeddings: int,
752
- base: float,
753
- is_neox_style: bool,
754
- scaling_factor: float,
755
- dtype: torch.dtype,
756
- *,
757
- extrapolation_factor: float = 1,
758
- attn_factor: float = 1,
759
- beta_fast: int = 32,
760
- beta_slow: int = 1,
761
- mscale: float = 1,
762
- mscale_all_dim: float = 0,
763
- ) -> None:
764
- self.scaling_factor = scaling_factor
765
- self.extrapolation_factor = extrapolation_factor
766
- self.attn_factor = attn_factor
767
- self.beta_fast = beta_fast
768
- self.beta_slow = beta_slow
769
- # Get n-d magnitude scaling corrected for interpolation.
770
- self.mscale = float(
771
- yarn_get_mscale(self.scaling_factor, float(mscale)) /
772
- yarn_get_mscale(self.scaling_factor, float(mscale_all_dim)) *
773
- attn_factor)
774
- super().__init__(head_size, rotary_dim, max_position_embeddings, base,
775
- is_neox_style, dtype)
776
-
777
- def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
778
- pos_freqs = self.base**(
779
- torch.arange(0,
780
- self.rotary_dim,
781
- 2,
782
- dtype=torch.float,
783
- device=current_platform.device_type) /
784
- self.rotary_dim)
785
- inv_freq_extrapolation = 1.0 / pos_freqs
786
- inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
787
-
788
- low, high = _yarn_find_correction_range(self.beta_fast, self.beta_slow,
789
- self.rotary_dim, self.base,
790
- self.max_position_embeddings)
791
- # Get n-d rotational scaling corrected for extrapolation
792
- inv_freq_mask = (1 - _yarn_linear_ramp_mask(
793
- low, high, self.rotary_dim // 2,
794
- dtype=torch.float)) * self.extrapolation_factor
795
- inv_freq = inv_freq_interpolation * (
796
- 1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask
797
- return inv_freq
798
-
799
- def _compute_cos_sin_cache(self) -> torch.Tensor:
800
- inv_freq = self._compute_inv_freq(self.scaling_factor)
801
- t = torch.arange(self.max_position_embeddings * self.scaling_factor,
802
- device=current_platform.device_type,
803
- dtype=torch.float32)
804
- freqs = torch.einsum("i,j -> ij", t, inv_freq)
805
- cos = (freqs.cos() * self.mscale)
806
- sin = (freqs.sin() * self.mscale)
807
- cache = torch.cat((cos, sin), dim=-1)
808
- return cache
809
-
810
- def forward(
811
- self,
812
- positions: torch.Tensor,
813
- query: torch.Tensor,
814
- key: Optional[torch.Tensor] = None,
815
- offsets: Optional[torch.Tensor] = None,
816
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
817
- """PyTorch-native implementation equivalent to forward()."""
818
- assert key is not None
819
- query_rot = query[..., :self.rotary_dim]
820
- key_rot = key[..., :self.rotary_dim]
821
- if self.rotary_dim < self.head_size:
822
- query_pass = query[..., self.rotary_dim:]
823
- key_pass = key[..., self.rotary_dim:]
824
-
825
- if self.cos_sin_cache.device != positions.device:
826
- self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(
827
- positions.device)
828
- cos_sin = self.cos_sin_cache[torch.add(positions, offsets)
829
- if offsets is not None else positions]
830
- cos, sin = cos_sin.chunk(2, dim=-1)
831
- if self.is_neox_style:
832
- # NOTE(woosuk): Here we assume that the positions tensor has the
833
- # shape [batch_size, seq_len].
834
- cos = cos.repeat(1, 1, 2).unsqueeze(-2)
835
- sin = sin.repeat(1, 1, 2).unsqueeze(-2)
836
- else:
837
- cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
838
- sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)
839
-
840
- rotate_fn = _rotate_neox if self.is_neox_style else _rotate_gptj
841
- query_rot = query_rot * cos + rotate_fn(query_rot) * sin
842
- key_rot = key_rot * cos + rotate_fn(key_rot) * sin
843
-
844
- if self.rotary_dim < self.head_size:
845
- query = torch.cat((query_rot, query_pass), dim=-1)
846
- key = torch.cat((key_rot, key_pass), dim=-1)
847
- else:
848
- query = query_rot
849
- key = key_rot
850
- return query, key
851
-
852
-
853
-class Llama3RotaryEmbedding(RotaryEmbedding):
854
-
855
- def __init__(
856
- self,
857
- head_size: int,
858
- rotary_dim: int,
859
- max_position_embeddings: int,
860
- base: float,
861
- is_neox_style: bool,
862
- dtype: torch.dtype,
863
- scaling_factor: float,
864
- low_freq_factor: float,
865
- high_freq_factor: float,
866
- orig_max_position: int,
867
- ) -> None:
868
- self.scaling_factor = scaling_factor
869
- self.low_freq_factor = low_freq_factor
870
- self.high_freq_factor = high_freq_factor
871
- self.orig_max_position = orig_max_position
872
- super().__init__(head_size, rotary_dim, max_position_embeddings, base,
873
- is_neox_style, dtype)
874
-
875
- def _compute_inv_freq(self, base: float) -> torch.Tensor:
876
- inv_freqs = super()._compute_inv_freq(base)
877
- low_freq_wavelen = self.orig_max_position / self.low_freq_factor
878
- high_freq_wavelen = self.orig_max_position / self.high_freq_factor
879
-
880
- wave_len = 2 * math.pi / inv_freqs
881
- if self.low_freq_factor != self.high_freq_factor:
882
- smooth = (self.orig_max_position / wave_len - self.low_freq_factor
883
- ) / (self.high_freq_factor - self.low_freq_factor)
884
- else:
885
- smooth = 0
886
- new_freqs = torch.where(
887
- wave_len < high_freq_wavelen,
888
- inv_freqs,
889
- torch.where(
890
- wave_len > low_freq_wavelen,
891
- inv_freqs / self.scaling_factor,
892
- (1 - smooth) * inv_freqs / self.scaling_factor +
893
- smooth * inv_freqs,
894
- ),
895
- )
896
- return new_freqs
897
-
898
-
899
-class Llama4VisionRotaryEmbedding(RotaryEmbedding):
900
-
901
- def __init__(
902
- self,
903
- head_size: int,
904
- rotary_dim: int,
905
- max_position_embeddings: int,
906
- base: float,
907
- is_neox_style: bool,
908
- dtype: torch.dtype,
909
- ):
910
- super().__init__(head_size, rotary_dim, max_position_embeddings, base,
911
- is_neox_style, dtype)
912
-
913
- def _compute_inv_freq(self, base: float) -> torch.Tensor:
914
- inv_freqs = super()._compute_inv_freq(base)
915
- inv_freqs = inv_freqs[:(self.rotary_dim // 2)]
916
- return inv_freqs
917
-
918
- def _compute_cos_sin_cache(self) -> torch.Tensor:
919
- inv_freq = self._compute_inv_freq(self.base)
920
-
921
- # self.max_position_embeddings here is number of image patches
922
- # i.e. (image_size // patch_size) ** 2
923
- num_patches = self.max_position_embeddings
924
- img_idx = torch.arange(num_patches,
925
- dtype=torch.int32) \
926
- .reshape(num_patches, 1)
927
- img_idx = torch.cat([img_idx, img_idx[:1]], dim=0)
928
- img_idx[-1, -1] = -2 # set to ID_CLS_TOKEN
929
- num_patches_single_dim = int(math.sqrt(num_patches))
930
- frequencies_x = img_idx % num_patches_single_dim
931
- frequencies_y = img_idx // num_patches_single_dim
932
- freqs_x = ((frequencies_x + 1)[..., None] *
933
- inv_freq[None, None, :]).repeat_interleave(2, dim=-1)
934
- freqs_y = ((frequencies_y + 1)[..., None] *
935
- inv_freq[None, None, :]).repeat_interleave(2, dim=-1)
936
- freqs = torch.cat([freqs_x, freqs_y],
937
- dim=-1).float().contiguous()[..., ::2]
938
- freqs = freqs.masked_fill(img_idx.reshape(-1, 1, 1) < 0, 0)
939
- cache = torch.view_as_complex(
940
- torch.stack([torch.cos(freqs), torch.sin(freqs)], dim=-1))
941
- return cache
942
-
943
- def forward(
944
- self,
945
- query: torch.Tensor,
946
- key: Optional[torch.Tensor] = None,
947
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
948
- assert key is not None
949
- self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(query.device)
950
- query_ = torch.view_as_complex(query.float().reshape(
951
- *query.shape[:-1], -1, 2))
952
- key_ = torch.view_as_complex(key.float().reshape(
953
- *key.shape[:-1], -1, 2))
954
- broadcast_shape = [
955
- d if i == 1 or i == (query_.ndim - 1) else 1
956
- for i, d in enumerate(query_.shape)
957
- ]
958
- freqs_ci = self.cos_sin_cache.view(*broadcast_shape)
959
- query_out = torch.view_as_real(query_ * freqs_ci).flatten(3)
960
- key_out = torch.view_as_real(key_ * freqs_ci).flatten(3)
961
- return query_out.type_as(query), key_out.type_as(key)
962
-
963
-
964
-class MRotaryEmbedding(RotaryEmbedding):
965
- """Rotary Embedding with Multimodal Sections."""
966
-
967
- def __init__(
968
- self,
969
- head_size: int,
970
- rotary_dim: int,
971
- max_position_embeddings: int,
972
- base: float,
973
- is_neox_style: bool,
974
- dtype: torch.dtype,
975
- mrope_section: Optional[list[int]] = None,
976
- ) -> None:
977
- # In Qwen2.5-VL, the maximum index value is related to the duration of
978
- # the input video. We enlarge max_position_embeddings to 4 times to get
979
- # a larger the cos and sin cache.
980
- self.cache_max_position_num = max_position_embeddings * 4
981
- super().__init__(head_size, rotary_dim, self.cache_max_position_num,
982
- base, is_neox_style, dtype)
983
-
984
- self.mrope_section = mrope_section
985
- if self.mrope_section:
986
- assert sum(self.mrope_section) == rotary_dim // 2
987
-
988
- def forward(
989
- self,
990
- positions: torch.Tensor,
991
- query: torch.Tensor,
992
- key: Optional[torch.Tensor] = None,
993
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
994
- """PyTorch-native implementation equivalent to forward().
995
-
996
- Args:
997
- positions:
998
- [num_tokens,] (text only) or
999
- [3, num_tokens] (T/H/W positions with multimodal inputs)
1000
- query: [num_tokens, num_heads * head_size]
1001
- key: [num_tokens, num_kv_heads * head_size]
1002
- """
1003
- assert positions.ndim == 1 or positions.ndim == 2
1004
- assert key is not None
1005
-
1006
- num_tokens = positions.shape[-1]
1007
- cos_sin = self.cos_sin_cache[positions]
1008
- cos, sin = cos_sin.chunk(2, dim=-1)
1009
- if positions.ndim == 2:
1010
- assert self.mrope_section
1011
-
1012
- cos = torch.cat([
1013
- m[i]
1014
- for i, m in enumerate(cos.split(self.mrope_section, dim=-1))
1015
- ],
1016
- dim=-1)
1017
- sin = torch.cat([
1018
- m[i]
1019
- for i, m in enumerate(sin.split(self.mrope_section, dim=-1))
1020
- ],
1021
- dim=-1)
1022
-
1023
- query_shape = query.shape
1024
- query = query.view(num_tokens, -1, self.head_size)
1025
- query_rot = query[..., :self.rotary_dim]
1026
- query_pass = query[..., self.rotary_dim:]
1027
- query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
1028
- query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
1029
-
1030
- key_shape = key.shape
1031
- key = key.view(num_tokens, -1, self.head_size)
1032
- key_rot = key[..., :self.rotary_dim]
1033
- key_pass = key[..., self.rotary_dim:]
1034
- key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
1035
- key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
1036
- return query, key
1037
-
1038
- @classmethod
1039
- def get_input_positions(
1040
- cls,
1041
- input_tokens: list[int],
1042
- hf_config: PretrainedConfig,
1043
- image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
1044
- video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
1045
- second_per_grid_ts: Optional[list[float]],
1046
- context_len: int = 0,
1047
- seq_len: Optional[int] = None,
1048
- audio_feature_lengths: Optional[torch.Tensor] = None,
1049
- use_audio_in_video: bool = False,
1050
- ) -> tuple[list[list[int]], int]:
1051
- """Get mrope input positions and delta value."""
1052
-
1053
- image_grid_thw = [] if image_grid_thw is None else image_grid_thw
1054
- video_grid_thw = [] if video_grid_thw is None else video_grid_thw
1055
- second_per_grid_ts = [] if second_per_grid_ts is None else \
1056
- second_per_grid_ts
1057
-
1058
- llm_positions, mrope_position_delta = \
1059
- cls.get_input_positions_tensor(
1060
- input_tokens=input_tokens,
1061
- hf_config=hf_config,
1062
- image_grid_thw=image_grid_thw,
1063
- video_grid_thw=video_grid_thw,
1064
- second_per_grid_ts=second_per_grid_ts,
1065
- context_len=context_len,
1066
- seq_len=seq_len,
1067
- audio_feature_lengths=audio_feature_lengths,
1068
- use_audio_in_video=use_audio_in_video,
1069
- )
1070
-
1071
- return llm_positions.tolist(), mrope_position_delta
1072
-
1073
- @classmethod
1074
- def get_input_positions_tensor(
1075
- cls,
1076
- input_tokens: list[int],
1077
- hf_config: PretrainedConfig,
1078
- image_grid_thw: Union[list[list[int]], torch.Tensor],
1079
- video_grid_thw: Union[list[list[int]], torch.Tensor],
1080
- second_per_grid_ts: list[float],
1081
- context_len: int = 0,
1082
- seq_len: Optional[int] = None,
1083
- audio_feature_lengths: Optional[torch.Tensor] = None,
1084
- use_audio_in_video: bool = False,
1085
- ) -> tuple[torch.Tensor, int]:
1086
- from vllm.transformers_utils.config import thinker_uses_mrope
1087
- if thinker_uses_mrope(hf_config):
1088
- return cls._omni_get_input_positions_tensor(
1089
- input_tokens=input_tokens,
1090
- hf_config=hf_config,
1091
- image_grid_thw=image_grid_thw,
1092
- video_grid_thw=video_grid_thw,
1093
- second_per_grid_ts=second_per_grid_ts,
1094
- context_len=context_len,
1095
- seq_len=seq_len,
1096
- audio_feature_lengths=audio_feature_lengths,
1097
- use_audio_in_video=use_audio_in_video,
1098
- )
1099
- elif hf_config.model_type in ["glm4v", "glm4v_moe"]:
1100
- return cls._glm4v_get_input_positions_tensor(
1101
- input_tokens=input_tokens,
1102
- hf_config=hf_config,
1103
- image_grid_thw=image_grid_thw,
1104
- video_grid_thw=video_grid_thw,
1105
- context_len=context_len,
1106
- seq_len=seq_len,
1107
- )
1108
- else:
1109
- return cls._vl_get_input_positions_tensor(
1110
- input_tokens=input_tokens,
1111
- hf_config=hf_config,
1112
- image_grid_thw=image_grid_thw,
1113
- video_grid_thw=video_grid_thw,
1114
- second_per_grid_ts=second_per_grid_ts,
1115
- context_len=context_len,
1116
- seq_len=seq_len,
1117
- )
1118
-
1119
- @classmethod
1120
- def _glm4v_get_input_positions_tensor(
1121
- cls,
1122
- input_tokens: list[int],
1123
- hf_config: PretrainedConfig,
1124
- image_grid_thw: Union[list[list[int]], torch.Tensor],
1125
- video_grid_thw: Union[list[list[int]], torch.Tensor],
1126
- context_len: int = 0,
1127
- seq_len: Optional[int] = None,
1128
- ) -> tuple[torch.Tensor, int]:
1129
- """Get mrope input positions and delta value for GLM4V."""
1130
-
1131
- image_token_id = hf_config.image_token_id
1132
- video_start_token_id = hf_config.video_start_token_id
1133
- video_end_token_id = hf_config.video_end_token_id
1134
- spatial_merge_size = hf_config.vision_config.spatial_merge_size
1135
- llm_pos_ids_list: list = []
1136
-
1137
- if not (image_grid_thw is None and video_grid_thw is None):
1138
- if isinstance(image_grid_thw, torch.Tensor):
1139
- image_grid_thw = image_grid_thw.tolist()
1140
-
1141
- input_token_type: list[str] = []
1142
- video_check_flg = False
1143
- for token in input_tokens:
1144
- if token == video_start_token_id:
1145
- video_check_flg = True
1146
- elif token == video_end_token_id:
1147
- video_check_flg = False
1148
-
1149
- if (token == image_token_id) and (video_check_flg is False):
1150
- input_token_type.append("image")
1151
- elif (token == image_token_id) and (video_check_flg is True):
1152
- input_token_type.append("video")
1153
- else:
1154
- input_token_type.append("text")
1155
-
1156
- input_type_group: list[tuple[str, int, int]] = []
1157
- for key, group_iter in itertools.groupby(
1158
- enumerate(input_token_type), lambda x: x[1]):
1159
- group_list = list(group_iter)
1160
- start_index = group_list[0][0]
1161
- end_index = group_list[-1][0] + 1
1162
- input_type_group.append((key, start_index, end_index))
1163
-
1164
- video_frame_num = 1
1165
- mm_data_idx = 0
1166
- for modality_type, start_idx, end_idx in input_type_group:
1167
- st_idx = llm_pos_ids_list[-1].max() + 1 if len(
1168
- llm_pos_ids_list) > 0 else 0
1169
- if modality_type == "image":
1170
- t, h, w = (
1171
- image_grid_thw[mm_data_idx][0],
1172
- image_grid_thw[mm_data_idx][1],
1173
- image_grid_thw[mm_data_idx][2],
1174
- )
1175
- llm_grid_t, llm_grid_h, llm_grid_w = \
1176
- t, h // spatial_merge_size, w // spatial_merge_size
1177
-
1178
- t_index = torch.arange(llm_grid_t).view(-1, 1).expand(
1179
- -1, llm_grid_h * llm_grid_w).flatten()
1180
- h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(
1181
- llm_grid_t, -1, llm_grid_w).flatten()
1182
- w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(
1183
- llm_grid_t, llm_grid_h, -1).flatten()
1184
- llm_pos_ids_list.append(
1185
- torch.stack([t_index, h_index, w_index]) + st_idx)
1186
- mm_data_idx += 1
1187
-
1188
- elif modality_type == "video":
1189
- t, h, w = (
1190
- video_frame_num,
1191
- image_grid_thw[mm_data_idx][1],
1192
- image_grid_thw[mm_data_idx][2],
1193
- )
1194
- llm_grid_t, llm_grid_h, llm_grid_w = \
1195
- t, h // spatial_merge_size, w // spatial_merge_size
1196
-
1197
- for t_idx in range(llm_grid_t):
1198
- t_index = torch.tensor(t_idx).view(-1, 1).expand(
1199
- -1, llm_grid_h * llm_grid_w).flatten()
1200
- h_index = torch.arange(llm_grid_h).view(
1201
- 1, -1, 1).expand(1, -1, llm_grid_w).flatten()
1202
- w_index = torch.arange(llm_grid_w).view(
1203
- 1, 1, -1).expand(1, llm_grid_h, -1).flatten()
1204
- llm_pos_ids_list.append(
1205
- torch.stack([t_index, h_index, w_index]) + st_idx)
1206
-
1207
- mm_data_idx += 1
1208
- video_frame_num += 1
1209
-
1210
- else:
1211
- text_len = end_idx - start_idx
1212
- llm_pos_ids_list.append(
1213
- torch.arange(text_len).view(1, -1).expand(3, -1) +
1214
- st_idx)
1215
- video_frame_num = 1
1216
-
1217
- else:
1218
- text_len = len(input_tokens)
1219
- llm_pos_ids_list.append(
1220
- torch.arange(text_len).view(1, -1).expand(3, -1))
1221
-
1222
- llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
1223
- llm_positions = llm_positions[:, context_len:seq_len]
1224
- mrope_position_delta = (llm_positions.max() + 1 -
1225
- len(input_tokens)).item()
1226
- return llm_positions, mrope_position_delta
1227
-
1228
- @classmethod
1229
- def _vl_get_input_positions_tensor(
1230
- cls,
1231
- input_tokens: list[int],
1232
- hf_config: PretrainedConfig,
1233
- image_grid_thw: Union[list[list[int]], torch.Tensor],
1234
- video_grid_thw: Union[list[list[int]], torch.Tensor],
1235
- second_per_grid_ts: list[float],
1236
- context_len: int = 0,
1237
- seq_len: Optional[int] = None,
1238
- ) -> tuple[torch.Tensor, int]:
1239
- """Get mrope input positions and delta value."""
1240
-
1241
- image_token_id = hf_config.image_token_id
1242
- video_token_id = hf_config.video_token_id
1243
- vision_start_token_id = hf_config.vision_start_token_id
1244
- spatial_merge_size = hf_config.vision_config.spatial_merge_size
1245
- tokens_per_second = getattr(hf_config.vision_config,
1246
- "tokens_per_second", 1.0)
1247
-
1248
- input_tokens_tensor = torch.tensor(input_tokens)
1249
- vision_start_indices = torch.argwhere(
1250
- input_tokens_tensor == vision_start_token_id).squeeze(1)
1251
- vision_tokens = input_tokens_tensor[vision_start_indices + 1]
1252
- image_nums = (vision_tokens == image_token_id).sum()
1253
- video_nums = (vision_tokens == video_token_id).sum()
1254
- llm_pos_ids_list: list = []
1255
-
1256
- st = 0
1257
- remain_images, remain_videos = image_nums, video_nums
1258
-
1259
- image_index, video_index = 0, 0
1260
- for _ in range(image_nums + video_nums):
1261
- video_second_per_grid_t = 0.0
1262
- if image_token_id in input_tokens and remain_images > 0:
1263
- ed_image = input_tokens.index(image_token_id, st)
1264
- else:
1265
- ed_image = len(input_tokens) + 1
1266
- if video_token_id in input_tokens and remain_videos > 0:
1267
- ed_video = input_tokens.index(video_token_id, st)
1268
- else:
1269
- ed_video = len(input_tokens) + 1
1270
- if ed_image < ed_video:
1271
- t, h, w = (
1272
- image_grid_thw[image_index][0],
1273
- image_grid_thw[image_index][1],
1274
- image_grid_thw[image_index][2],
1275
- )
1276
- image_index += 1
1277
- remain_images -= 1
1278
- ed = ed_image
1279
- else:
1280
- t, h, w = (
1281
- video_grid_thw[video_index][0],
1282
- video_grid_thw[video_index][1],
1283
- video_grid_thw[video_index][2],
1284
- )
1285
- video_second_per_grid_t = 1.0
1286
- if second_per_grid_ts:
1287
- video_second_per_grid_t = second_per_grid_ts[video_index]
1288
- video_index += 1
1289
- remain_videos -= 1
1290
- ed = ed_video
1291
-
1292
- llm_grid_t, llm_grid_h, llm_grid_w = \
1293
- t, h // spatial_merge_size, w // spatial_merge_size
1294
- text_len = ed - st
1295
-
1296
- st_idx = llm_pos_ids_list[-1].max() + 1 if len(
1297
- llm_pos_ids_list) > 0 else 0
1298
- llm_pos_ids_list.append(
1299
- torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
1300
-
1301
- t_index = (torch.arange(llm_grid_t).view(-1, 1).expand(
1302
- -1, llm_grid_h * llm_grid_w) * video_second_per_grid_t *
1303
- tokens_per_second).long().flatten()
1304
-
1305
- h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(
1306
- llm_grid_t, -1, llm_grid_w).flatten()
1307
- w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(
1308
- llm_grid_t, llm_grid_h, -1).flatten()
1309
- llm_pos_ids_list.append(
1310
- torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
1311
- st = ed + llm_grid_t * llm_grid_h * llm_grid_w
1312
-
1313
- if st < len(input_tokens):
1314
- st_idx = llm_pos_ids_list[-1].max() + 1 if len(
1315
- llm_pos_ids_list) > 0 else 0
1316
- text_len = len(input_tokens) - st
1317
- llm_pos_ids_list.append(
1318
- torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
1319
-
1320
- llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
1321
- mrope_position_delta = (llm_positions.max() + 1 -
1322
- len(input_tokens)).item()
1323
- llm_positions = llm_positions[:, context_len:seq_len]
1324
-
1325
- return llm_positions, mrope_position_delta
1326
-
1327
- @classmethod
1328
- def _omni_get_input_positions_tensor(
1329
- cls,
1330
- input_tokens: list[int],
1331
- hf_config: PretrainedConfig,
1332
- image_grid_thw: Union[list[list[int]], torch.Tensor],
1333
- video_grid_thw: Union[list[list[int]], torch.Tensor],
1334
- second_per_grid_ts: Optional[list[float]] = None,
1335
- context_len: int = 0,
1336
- seq_len: Optional[int] = None,
1337
- audio_feature_lengths: Optional[torch.Tensor] = None,
1338
- use_audio_in_video: bool = False,
1339
- ) -> tuple[torch.Tensor, int]:
1340
- """Get mrope input positions and delta value (Qwen2.5-Omni version).
1341
-
1342
- Differences from MRotaryEmbedding:
1343
- 1. Add audio support (and related `audio_feature_lengths`).
1344
- 2. Add `use_audio_in_video` option to read audio from video inputs.
1345
- In this case, audio and vision position ids will be split into
1346
- chunks and interleaved.
1347
-
1348
- Example:
1349
-
1350
- (V_i are vision position ids, A_i are audio position ids)
1351
-
1352
- |V_1 ... V_n|A_1 ... A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|...
1353
- |vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |...
1354
- """
1355
-
1356
- # TODO(fyabc): refactor and share more code with
1357
- # _vl_get_input_positions_tensor.
1358
-
1359
- thinker_config = hf_config.thinker_config
1360
- audio_token_id = thinker_config.audio_token_index
1361
- image_token_id = thinker_config.image_token_index
1362
- video_token_id = thinker_config.video_token_index
1363
- audio_start_token_id = thinker_config.audio_start_token_id
1364
- audio_end_token_id = thinker_config.audio_end_token_id
1365
- vision_start_token_id = thinker_config.vision_start_token_id
1366
- vision_end_token_id = thinker_config.vision_end_token_id
1367
- seconds_per_chunk = thinker_config.seconds_per_chunk
1368
- spatial_merge_size = thinker_config.vision_config.spatial_merge_size
1369
- tokens_per_second = getattr(thinker_config.vision_config,
1370
- "tokens_per_second", 25)
1371
-
1372
- if isinstance(image_grid_thw, list):
1373
- image_grid_thw = torch.tensor(image_grid_thw)
1374
- if isinstance(video_grid_thw, list):
1375
- video_grid_thw = torch.tensor(video_grid_thw)
1376
-
1377
- src_item = input_tokens
1378
- audio_seqlens = audio_feature_lengths
1379
- if not second_per_grid_ts:
1380
- second_per_grid_ts = [1] * video_grid_thw.shape[0]
1381
- audio_idx = 0
1382
- video_idx = 0
1383
- image_idx = 0
1384
- new_src_item: list[int] = []
1385
- llm_pos_ids_list: list[torch.Tensor] = []
1386
-
1387
- idx = 0
1388
- while idx < len(src_item):
1389
- new_src_item_len = len(new_src_item)
1390
- start_idx = llm_pos_ids_list[-1].max() + 1 if len(
1391
- llm_pos_ids_list) > 0 else 0
1392
- if src_item[idx] not in [
1393
- audio_token_id, video_token_id, image_token_id
1394
- ]:
1395
- if use_audio_in_video and idx > 0:
1396
- if src_item[idx] == vision_end_token_id and \
1397
- src_item[idx - 1] == audio_end_token_id:
1398
- # processing the <|audio_eos|> before <|vision_eos|>
1399
- start_idx -= 1
1400
- elif src_item[idx] == audio_start_token_id and \
1401
- src_item[idx - 1] == vision_start_token_id:
1402
- # processing the <|audio_bos|> after <|vision_eos|>
1403
- start_idx -= 1
1404
- new_src_item.append(src_item[idx])
1405
- llm_pos_ids = torch.tensor([start_idx],
1406
- dtype=torch.long).expand(3, -1)
1407
- llm_pos_ids_list.append(llm_pos_ids)
1408
- elif src_item[idx] == audio_token_id:
1409
- assert audio_seqlens is not None
1410
- audio_seqlen = audio_seqlens[audio_idx]
1411
- place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1)
1412
- new_src_item.extend([audio_token_id] * place_num)
1413
- llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx
1414
- llm_pos_ids_list.append(llm_pos_ids)
1415
- audio_idx += 1
1416
- elif src_item[idx] == image_token_id:
1417
- grid_t = image_grid_thw[image_idx][0]
1418
- grid_hs = image_grid_thw[:, 1]
1419
- grid_ws = image_grid_thw[:, 2]
1420
- t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long()
1421
- llm_pos_ids = cls._get_llm_pos_ids_for_vision(
1422
- start_idx, image_idx, spatial_merge_size, t_index, grid_hs,
1423
- grid_ws)
1424
- llm_pos_ids_list.append(llm_pos_ids)
1425
- vision_seqlen = image_grid_thw[image_idx].prod() // (
1426
- spatial_merge_size**2)
1427
- new_src_item.extend([image_token_id] * vision_seqlen)
1428
- image_idx += 1
1429
- elif src_item[idx] == video_token_id and not use_audio_in_video:
1430
- grid_t = video_grid_thw[video_idx][0]
1431
- grid_hs = video_grid_thw[:, 1]
1432
- grid_ws = video_grid_thw[:, 2]
1433
- t_index = (torch.arange(grid_t) *
1434
- second_per_grid_ts[video_idx] *
1435
- tokens_per_second).long()
1436
- llm_pos_ids = cls._get_llm_pos_ids_for_vision(
1437
- start_idx, video_idx, spatial_merge_size, t_index, grid_hs,
1438
- grid_ws)
1439
- llm_pos_ids_list.append(llm_pos_ids)
1440
- vision_seqlen = video_grid_thw[video_idx].prod() // (
1441
- spatial_merge_size**2)
1442
- new_src_item.extend([video_token_id] * vision_seqlen)
1443
- video_idx += 1
1444
- else:
1445
- # read audio from video
1446
- assert audio_seqlens is not None
1447
- audio_seqlen = audio_seqlens[audio_idx]
1448
- vision_seqlen = video_grid_thw[video_idx].prod() // (
1449
- spatial_merge_size**2)
1450
- grid_t = video_grid_thw[video_idx][0]
1451
- grid_h = video_grid_thw[video_idx][1]
1452
- grid_w = video_grid_thw[video_idx][2]
1453
- grid_hs = video_grid_thw[:, 1]
1454
- grid_ws = video_grid_thw[:, 2]
1455
- t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
1456
- t_index = (torch.arange(grid_t) *
1457
- second_per_grid_ts[video_idx] *
1458
- tokens_per_second).long()
1459
- t_index_split_chunk = cls._split_list_into_ranges(
1460
- t_index, t_ntoken_per_chunk)
1461
- place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2
1462
- pure_audio_len = place_num - 2
1463
- added_audio_len = 0
1464
- audio_llm_pos_ids_list: list[torch.Tensor] = []
1465
- for t_chunk in t_index_split_chunk:
1466
- vision_ntoken_per_chunk = len(
1467
- t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
1468
- new_src_item.extend([video_token_id] *
1469
- vision_ntoken_per_chunk)
1470
- vision_llm_pos_ids_list = cls._get_llm_pos_ids_for_vision(
1471
- start_idx, video_idx, spatial_merge_size, t_chunk,
1472
- grid_hs, grid_ws).split(1, dim=1)
1473
- llm_pos_ids_list.extend(vision_llm_pos_ids_list)
1474
- new_src_item.extend(
1475
- min(t_ntoken_per_chunk, pure_audio_len -
1476
- added_audio_len) * [audio_token_id])
1477
- audio_start_idx = start_idx if len(
1478
- audio_llm_pos_ids_list
1479
- ) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1
1480
- if min(t_ntoken_per_chunk,
1481
- pure_audio_len - added_audio_len) > 0:
1482
- audio_llm_pos_ids_list = (torch.arange(
1483
- min(t_ntoken_per_chunk, pure_audio_len -
1484
- added_audio_len)).expand(3, -1) +
1485
- audio_start_idx).split(1,
1486
- dim=1)
1487
- else:
1488
- audio_llm_pos_ids_list = []
1489
- added_audio_len += min(t_ntoken_per_chunk,
1490
- pure_audio_len - added_audio_len)
1491
- llm_pos_ids_list.extend(audio_llm_pos_ids_list)
1492
- if added_audio_len < pure_audio_len:
1493
- new_src_item.extend(
1494
- (pure_audio_len - added_audio_len) * [audio_token_id])
1495
- audio_llm_pos_ids_list = (
1496
- torch.arange(pure_audio_len - added_audio_len).expand(
1497
- 3, -1) + llm_pos_ids_list[-1].max() + 1).split(
1498
- 1, dim=1)
1499
- llm_pos_ids_list.extend(audio_llm_pos_ids_list)
1500
- audio_idx += 1
1501
- video_idx += 1
1502
- # move to the next token
1503
- idx += len(new_src_item) - new_src_item_len
1504
-
1505
- llm_positions = torch.cat(llm_pos_ids_list, dim=1)
1506
- mrope_position_delta = torch.cat(llm_pos_ids_list,
1507
- dim=1).max() + 1 - len(src_item)
1508
- llm_positions = llm_positions[:, context_len:seq_len]
1509
-
1510
- return llm_positions, mrope_position_delta
1511
-
1512
- @staticmethod
1513
- def _get_llm_pos_ids_for_vision(
1514
- start_idx: int,
1515
- vision_idx: int,
1516
- spatial_merge_size: int,
1517
- t_index: list[int],
1518
- grid_hs: torch.Tensor,
1519
- grid_ws: torch.Tensor,
1520
- ) -> torch.Tensor:
1521
- llm_pos_ids_list = []
1522
- llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
1523
- llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
1524
- h_index = (torch.arange(llm_grid_h).view(1, -1, 1).expand(
1525
- len(t_index), -1, llm_grid_w).flatten())
1526
- w_index = (torch.arange(llm_grid_w).view(1, 1, -1).expand(
1527
- len(t_index), llm_grid_h, -1).flatten())
1528
- t_index_tensor = torch.Tensor(t_index).to(llm_grid_h.device).view(
1529
- -1, 1).expand(-1, llm_grid_h * llm_grid_w).long().flatten()
1530
- _llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
1531
- llm_pos_ids_list.append(_llm_pos_ids + start_idx)
1532
- llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
1533
- return llm_pos_ids
1534
-
1535
- @staticmethod
1536
- def _split_list_into_ranges(lst: torch.Tensor,
1537
- interval: int) -> list[list[int]]:
1538
- ranges: list[list[int]] = [[]
1539
- for _ in range((max(lst) // interval) + 1)]
1540
- for num in lst:
1541
- index = num // interval
1542
- ranges[index].append(num)
1543
- return ranges
1544
-
1545
- @staticmethod
1546
- def get_next_input_positions(
1547
- mrope_position_delta: int,
1548
- context_len: int,
1549
- seq_len: int,
1550
- ) -> list[list[int]]:
1551
- return [
1552
- list(
1553
- range(context_len + mrope_position_delta,
1554
- seq_len + mrope_position_delta)) for _ in range(3)
1555
- ]
1556
-
1557
- @staticmethod
1558
- def get_next_input_positions_tensor(out: np.ndarray, out_offset: int,
1559
- mrope_position_delta: int,
1560
- context_len: int, num_new_tokens: int):
1561
-
1562
- values = np.arange(mrope_position_delta + context_len,
1563
- mrope_position_delta + context_len + num_new_tokens,
1564
- dtype=out.dtype)
1565
- out[:, out_offset:out_offset + num_new_tokens] = values
1566
-
1567
- @classmethod
1568
- def omni_get_updates_use_audio_in_video(
1569
- cls,
1570
- thinker_config: PretrainedConfig,
1571
- audio_len: int,
1572
- video_grid_thw: Union[list[int], torch.Tensor],
1573
- video_second_per_grid_t: float,
1574
- ) -> list[int]:
1575
- """Get video prompt updates when `use_audio_in_video` is True.
1576
-
1577
- In this case, audio and vision update ids will be split into
1578
- chunks and interleaved (details in `_omni_get_input_positions_tensor`).
1579
-
1580
- <|video_bos|><|VIDEO|><|video_eos|> =>
1581
- <|video_bos|><|audio_bos|>(... chunks ...)<|audio_eos|><|video_eos|>
1582
- """
1583
-
1584
- audio_token_id = thinker_config.audio_token_index
1585
- video_token_id = thinker_config.video_token_index
1586
- audio_start_token_id = thinker_config.audio_start_token_id
1587
- audio_end_token_id = thinker_config.audio_end_token_id
1588
- seconds_per_chunk = thinker_config.seconds_per_chunk
1589
- spatial_merge_size = thinker_config.vision_config.spatial_merge_size
1590
- tokens_per_second = getattr(thinker_config.vision_config,
1591
- "tokens_per_second", 25)
1592
-
1593
- grid_t = video_grid_thw[0]
1594
- grid_h = video_grid_thw[1]
1595
- grid_w = video_grid_thw[2]
1596
- t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
1597
- t_index = (torch.arange(grid_t) * video_second_per_grid_t *
1598
- tokens_per_second).long()
1599
- t_index_split_chunk = cls._split_list_into_ranges(
1600
- t_index, t_ntoken_per_chunk)
1601
-
1602
- updates = [audio_start_token_id]
1603
- added_audio_len = 0
1604
- for t_chunk in t_index_split_chunk:
1605
- vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (
1606
- spatial_merge_size**2)
1607
- updates.extend([video_token_id] * vision_ntoken_per_chunk)
1608
-
1609
- audio_chunk_size = min(t_ntoken_per_chunk,
1610
- audio_len - added_audio_len)
1611
- updates.extend(audio_chunk_size * [audio_token_id])
1612
- added_audio_len += audio_chunk_size
1613
- if added_audio_len < audio_len:
1614
- updates.extend((audio_len - added_audio_len) * [audio_token_id])
1615
- updates.extend([audio_end_token_id])
1616
-
1617
- return updates
1618
-
1619
-
1620
-@CustomOp.register("dual_chunk_rotary_embedding")
1621
-class DualChunkRotaryEmbedding(CustomOp):
1622
- """Rotary positional embedding for Dual Chunk Attention."""
1623
-
1624
- def __init__(
1625
- self,
1626
- head_size: int,
1627
- rotary_dim: int,
1628
- max_position_embeddings: int,
1629
- base: float,
1630
- is_neox_style: bool,
1631
- dtype: torch.dtype,
1632
- chunk_size: int,
1633
- local_size: int,
1634
- ) -> None:
1635
- super().__init__()
1636
- self.head_size = head_size
1637
- self.rotary_dim = rotary_dim
1638
- self.max_position_embeddings = max_position_embeddings
1639
- self.base = base
1640
- self.is_neox_style = is_neox_style
1641
- self.chunk_size = chunk_size
1642
- self.local_size = local_size
1643
- self.dtype = dtype
1644
- self.device = torch.device(f"cuda:{torch.cuda.current_device()}")
1645
- (q_cache, qc_cache, k_cache, qc_no_clamp_cache,
1646
- q_inter_cache) = self._compute_cos_sin_cache()
1647
-
1648
- self.register_buffer("cos_sin_q_cache", q_cache, persistent=False)
1649
- self.register_buffer("cos_sin_qc_cache", qc_cache, persistent=False)
1650
- self.register_buffer("cos_sin_k_cache", k_cache, persistent=False)
1651
- self.register_buffer("cos_sin_qc_no_clamp_cache",
1652
- qc_no_clamp_cache,
1653
- persistent=False)
1654
- self.register_buffer("cos_sin_q_inter_cache",
1655
- q_inter_cache,
1656
- persistent=False)
1657
-
1658
- def _compute_inv_freq(self, base: float) -> torch.Tensor:
1659
- """Compute the inverse frequency."""
1660
- # NOTE(woosuk): The HF implementation uses `torch.arange(...).float()`.
1661
- # However, we use `torch.arange(..., dtype=torch.float)` instead to
1662
- # avoid numerical issues with large base values (e.g., 10000000).
1663
- # This may cause a slight numerical difference between the HF
1664
- # implementation and ours.
1665
- # NOTE(woosuk): To exactly match the HF implementation, we need to
1666
- # use CPU to compute the cache and then move it to GPU. However, we
1667
- # create the cache on GPU for faster initialization. This may cause
1668
- # a slight numerical difference between the HF implementation and ours.
1669
- inv_freq = 1.0 / (base**(torch.arange(
1670
- 0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim))
1671
- return inv_freq
1672
-
1673
- def _compute_cos_sin_cache(self) -> torch.Tensor:
1674
- """Compute the cos and sin cache."""
1675
- inv_freq = self._compute_inv_freq(self.base)
1676
- chunk_len = self.chunk_size - self.local_size
1677
- q_t = torch.arange(chunk_len, dtype=torch.float)
1678
- qc_t = (torch.arange(chunk_len, dtype=torch.float) +
1679
- chunk_len).clamp(max=self.chunk_size)
1680
- k_t = torch.arange(self.max_position_embeddings,
1681
- dtype=torch.float) % chunk_len
1682
-
1683
- # count from chunk_len, no clamp(self.chunk_size) restriction
1684
- qc_no_clamp_t = torch.arange(chunk_len, dtype=torch.float) + chunk_len
1685
- # count from self.chunk_size for q_inter's rope
1686
- q_inter_t = torch.arange(chunk_len,
1687
- dtype=torch.float) + self.chunk_size
1688
-
1689
- q_freqs = torch.outer(q_t, inv_freq)
1690
- qc_freqs = torch.outer(qc_t, inv_freq)
1691
- k_freqs = torch.outer(k_t, inv_freq)
1692
- qc_no_clamp_freqs = torch.outer(qc_no_clamp_t, inv_freq)
1693
- q_inter_freqs = torch.outer(q_inter_t, inv_freq)
1694
-
1695
- q_cos = q_freqs.cos()
1696
- q_sin = q_freqs.sin()
1697
- qc_cos = qc_freqs.cos()
1698
- qc_sin = qc_freqs.sin()
1699
- k_cos = k_freqs.cos()
1700
- k_sin = k_freqs.sin()
1701
-
1702
- qc_no_clamp_cos = qc_no_clamp_freqs.cos()
1703
- qc_no_clamp_sin = qc_no_clamp_freqs.sin()
1704
- q_inter_cos = q_inter_freqs.cos()
1705
- q_inter_sin = q_inter_freqs.sin()
1706
-
1707
- q_cache = torch.cat((q_cos, q_sin), dim=-1).to(dtype=self.dtype,
1708
- device=self.device)
1709
- qc_cache = torch.cat((qc_cos, qc_sin), dim=-1).to(dtype=self.dtype,
1710
- device=self.device)
1711
- k_cache = torch.cat((k_cos, k_sin), dim=-1).to(dtype=self.dtype,
1712
- device=self.device)
1713
- qc_no_clamp_cache = torch.cat((qc_no_clamp_cos, qc_no_clamp_sin),
1714
- dim=-1).to(dtype=self.dtype,
1715
- device=self.device)
1716
- q_inter_cache = torch.cat((q_inter_cos, q_inter_sin),
1717
- dim=-1).to(dtype=self.dtype,
1718
- device=self.device)
1719
- return q_cache, qc_cache, k_cache, qc_no_clamp_cache, q_inter_cache
1720
-
1721
- def forward(
1722
- self,
1723
- positions: torch.Tensor,
1724
- query: torch.Tensor,
1725
- key: torch.Tensor,
1726
- offsets: Optional[torch.Tensor] = None,
1727
- ) -> tuple[torch.Tensor, torch.Tensor]:
1728
- query = query.view(*query.shape[:-1], -1, self.head_size)
1729
- key = key.view(*key.shape[:-1], -1, self.head_size)
1730
- query_rot = query[..., :self.rotary_dim]
1731
- key_rot = key[..., :self.rotary_dim]
1732
- if self.rotary_dim < self.head_size:
1733
- query_pass = query[..., self.rotary_dim:]
1734
- key_pass = key[..., self.rotary_dim:]
1735
- else:
1736
- query_pass = None
1737
- key_pass = None
1738
-
1739
- positions_with_offsets = (torch.add(positions, offsets)
1740
- if offsets is not None else positions)
1741
- key = self._apply_rotary_embedding(
1742
- self.cos_sin_k_cache[positions_with_offsets], key_rot, key_pass)
1743
- chunk_len = self.chunk_size - self.local_size
1744
- query = self._apply_rotary_embedding(
1745
- self.cos_sin_q_cache[positions_with_offsets % chunk_len],
1746
- query_rot, query_pass)
1747
- query_succ = self._apply_rotary_embedding(
1748
- self.cos_sin_qc_cache[positions_with_offsets % chunk_len],
1749
- query_rot, query_pass)
1750
- query_inter = self._apply_rotary_embedding(
1751
- self.cos_sin_qc_cache[chunk_len - 1].repeat(positions.shape[0], 1),
1752
- query_rot, query_pass)
1753
- query_succ_critical = self._apply_rotary_embedding(
1754
- self.cos_sin_qc_no_clamp_cache[positions_with_offsets % chunk_len],
1755
- query_rot, query_pass)
1756
- query_inter_critical = self._apply_rotary_embedding(
1757
- self.cos_sin_q_inter_cache[positions_with_offsets % chunk_len],
1758
- query_rot, query_pass)
1759
-
1760
- # merge query into one tensor to simplify the interfaces
1761
- query = torch.cat((
1762
- query,
1763
- query_succ,
1764
- query_inter,
1765
- query_succ_critical,
1766
- query_inter_critical,
1767
- ),
1768
- dim=-1)
1769
- return query, key
1770
-
1771
- def _apply_rotary_embedding(self, cos_sin, hidden_rot, hidden_pass):
1772
- cos, sin = cos_sin.chunk(2, dim=-1)
1773
- if self.is_neox_style:
1774
- # NOTE(woosuk): Here we assume that the positions tensor has the
1775
- # shape [batch_size, seq_len].
1776
- cos = cos.repeat(1, 1, 2).unsqueeze(-2)
1777
- sin = sin.repeat(1, 1, 2).unsqueeze(-2)
1778
- else:
1779
- cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
1780
- sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)
1781
- rotate_fn = _rotate_neox if self.is_neox_style else _rotate_gptj
1782
- hidden_rot = hidden_rot * cos + rotate_fn(hidden_rot) * sin
1783
-
1784
- if self.rotary_dim < self.head_size:
1785
- hidden = torch.cat((hidden_rot, hidden_pass), dim=-1)
1786
- else:
1787
- hidden = hidden_rot
1788
- return hidden.flatten(-2).squeeze(0)
1789
-
1790
- def extra_repr(self) -> str:
1791
- s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
1792
- s += f", max_position_embeddings={self.max_position_embeddings}"
1793
- s += f", base={self.base}, is_neox_style={self.is_neox_style}"
1794
- s += f", chunk_size={self.chunk_size}, local_size={self.local_size}"
1795
- return s
1796
-
1797
-
1798
-_ROPE_DICT: dict[tuple, RotaryEmbedding] = {}
1799
-
1800
-
1801
-def get_rope(
1802
- head_size: int,
1803
- rotary_dim: int,
1804
- max_position: int,
1805
- base: float,
1806
- is_neox_style: bool = True,
1807
- rope_scaling: Optional[dict[str, Any]] = None,
1808
- dtype: Optional[torch.dtype] = None,
1809
- partial_rotary_factor: float = 1.0,
1810
- dual_chunk_attention_config: Optional[dict[str, Any]] = None,
1811
-) -> RotaryEmbedding:
1812
- if dtype is None:
1813
- dtype = torch.get_default_dtype()
1814
- if rope_scaling is not None:
1815
- # Transforms every value that is a list into a tuple for caching calls
1816
- rope_scaling_tuple = {
1817
- k: tuple(v) if isinstance(v, list) else v
1818
- for k, v in rope_scaling.items()
1819
- }
1820
- rope_scaling_args = tuple(rope_scaling_tuple.items())
1821
- else:
1822
- rope_scaling_args = None
1823
-
1824
- if dual_chunk_attention_config is not None:
1825
- dual_chunk_attention_tuple = {
1826
- k: tuple(v) if isinstance(v, list) else v
1827
- for k, v in dual_chunk_attention_config.items()
1828
- if k != "sparse_attention_config"
1829
- }
1830
- dual_chunk_attention_args = tuple(dual_chunk_attention_tuple.items())
1831
- else:
1832
- dual_chunk_attention_args = None
1833
-
1834
- if partial_rotary_factor < 1.0:
1835
- rotary_dim = int(rotary_dim * partial_rotary_factor)
1836
- key = (head_size, rotary_dim, max_position, base, is_neox_style,
1837
- rope_scaling_args, dual_chunk_attention_args, dtype)
1838
- if key in _ROPE_DICT:
1839
- return _ROPE_DICT[key]
1840
-
1841
- if dual_chunk_attention_config is not None:
1842
- extra_kwargs = {
1843
- k: v
1844
- for k, v in dual_chunk_attention_config.items()
1845
- if k in ("chunk_size", "local_size")
1846
- }
1847
- rotary_emb = DualChunkRotaryEmbedding(head_size, rotary_dim,
1848
- max_position, base,
1849
- is_neox_style, dtype,
1850
- **extra_kwargs)
1851
- elif not rope_scaling:
1852
- rotary_emb = RotaryEmbedding(head_size, rotary_dim, max_position, base,
1853
- is_neox_style, dtype)
1854
- else:
1855
- scaling_type = rope_scaling["rope_type"]
1856
-
1857
- if scaling_type == "llama3":
1858
- scaling_factor = rope_scaling["factor"]
1859
- low_freq_factor = rope_scaling["low_freq_factor"]
1860
- high_freq_factor = rope_scaling["high_freq_factor"]
1861
- original_max_position = rope_scaling[
1862
- "original_max_position_embeddings"]
1863
- rotary_emb = Llama3RotaryEmbedding(head_size, rotary_dim,
1864
- max_position, base,
1865
- is_neox_style, dtype,
1866
- scaling_factor, low_freq_factor,
1867
- high_freq_factor,
1868
- original_max_position)
1869
- elif scaling_type == "mllama4":
1870
- rotary_emb = Llama4VisionRotaryEmbedding(head_size, rotary_dim,
1871
- max_position, base,
1872
- is_neox_style, dtype)
1873
- elif scaling_type == "default":
1874
- if "mrope_section" in rope_scaling:
1875
- rotary_emb = MRotaryEmbedding(
1876
- head_size,
1877
- rotary_dim,
1878
- max_position,
1879
- base,
1880
- is_neox_style,
1881
- dtype,
1882
- mrope_section=rope_scaling["mrope_section"],
1883
- )
1884
- else:
1885
- rotary_emb = RotaryEmbedding(
1886
- head_size,
1887
- rotary_dim,
1888
- max_position,
1889
- base,
1890
- is_neox_style,
1891
- dtype,
1892
- )
1893
- elif scaling_type == "linear":
1894
- scaling_factor = rope_scaling["factor"]
1895
- rotary_emb = LinearScalingRotaryEmbedding(head_size, rotary_dim,
1896
- max_position, base,
1897
- is_neox_style,
1898
- scaling_factor, dtype)
1899
- elif scaling_type == "ntk":
1900
- scaling_factor = rope_scaling["factor"]
1901
- mixed_b = rope_scaling.get('mixed_b', None)
1902
- rotary_emb = NTKScalingRotaryEmbedding(head_size, rotary_dim,
1903
- max_position, base,
1904
- is_neox_style,
1905
- scaling_factor, dtype,
1906
- mixed_b)
1907
- elif scaling_type == "dynamic":
1908
- if "alpha" in rope_scaling:
1909
- scaling_alpha = rope_scaling["alpha"]
1910
- rotary_emb = DynamicNTKAlphaRotaryEmbedding(
1911
- head_size, rotary_dim, max_position, base, is_neox_style,
1912
- scaling_alpha, dtype)
1913
- elif "factor" in rope_scaling:
1914
- scaling_factor = rope_scaling["factor"]
1915
- rotary_emb = DynamicNTKScalingRotaryEmbedding(
1916
- head_size, rotary_dim, max_position, base, is_neox_style,
1917
- scaling_factor, dtype)
1918
- else:
1919
- raise ValueError("Dynamic rope scaling must contain either "
1920
- "'alpha' or 'factor' field")
1921
- elif scaling_type == "yarn":
1922
- scaling_factor = rope_scaling["factor"]
1923
- original_max_position = rope_scaling[
1924
- "original_max_position_embeddings"]
1925
- extra_kwargs = {
1926
- k: v
1927
- for k, v in rope_scaling.items()
1928
- if k in ("extrapolation_factor", "attn_factor", "beta_fast",
1929
- "beta_slow")
1930
- }
1931
- rotary_emb = YaRNScalingRotaryEmbedding(head_size, rotary_dim,
1932
- original_max_position,
1933
- base, is_neox_style,
1934
- scaling_factor, dtype,
1935
- **extra_kwargs)
1936
- elif scaling_type == "deepseek_yarn":
1937
- scaling_factor = rope_scaling["factor"]
1938
- original_max_position = rope_scaling[
1939
- "original_max_position_embeddings"]
1940
- # assert max_position == original_max_position * scaling_factor
1941
- extra_kwargs = {
1942
- k: v
1943
- for k, v in rope_scaling.items()
1944
- if k in ("extrapolation_factor", "attn_factor", "beta_fast",
1945
- "beta_slow", "mscale", "mscale_all_dim")
1946
- }
1947
- rotary_emb = DeepseekScalingRotaryEmbedding(
1948
- head_size, rotary_dim, original_max_position, base,
1949
- is_neox_style, scaling_factor, dtype, **extra_kwargs)
1950
- elif scaling_type == "longrope":
1951
- short_factor = rope_scaling["short_factor"]
1952
- long_factor = rope_scaling["long_factor"]
1953
- original_max_position = rope_scaling[
1954
- "original_max_position_embeddings"]
1955
- extra_kwargs = {
1956
- k: v
1957
- for k, v in rope_scaling.items()
1958
- if k in ("short_mscale", "long_mscale")
1959
- }
1960
- rotary_emb = Phi3LongRoPEScaledRotaryEmbedding(
1961
- head_size, rotary_dim, max_position, original_max_position,
1962
- base, is_neox_style, dtype, short_factor, long_factor,
1963
- **extra_kwargs)
1964
- else:
1965
- raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
1966
- _ROPE_DICT[key] = rotary_emb
1967
- return rotary_emb
vllm/model_executor/layers/rotary_embedding/__init__.py
new
+190
@@ -0,0 +1,190 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Rotary Positional Embeddings."""
4
+from typing import Any, Optional
5
+
6
+import torch
7
+
8
+from .base import RotaryEmbedding
9
+from .deepseek_scaling_rope import DeepseekScalingRotaryEmbedding
10
+from .dual_chunk_rope import DualChunkRotaryEmbedding
11
+from .dynamic_ntk_alpha_rope import DynamicNTKAlphaRotaryEmbedding
12
+from .dynamic_ntk_scaling_rope import DynamicNTKScalingRotaryEmbedding
13
+from .linear_scaling_rope import LinearScalingRotaryEmbedding
14
+from .llama3_rope import Llama3RotaryEmbedding
15
+from .llama4_vision_rope import Llama4VisionRotaryEmbedding
16
+from .mrope import MRotaryEmbedding
17
+from .ntk_scaling_rope import NTKScalingRotaryEmbedding
18
+from .phi3_long_rope_scaled_rope import Phi3LongRoPEScaledRotaryEmbedding
19
+from .yarn_scaling_rope import YaRNScalingRotaryEmbedding
20
+
21
+_ROPE_DICT: dict[tuple, RotaryEmbedding] = {}
22
+
23
+
24
+def get_rope(
25
+ head_size: int,
26
+ rotary_dim: int,
27
+ max_position: int,
28
+ base: float,
29
+ is_neox_style: bool = True,
30
+ rope_scaling: Optional[dict[str, Any]] = None,
31
+ dtype: Optional[torch.dtype] = None,
32
+ partial_rotary_factor: float = 1.0,
33
+ dual_chunk_attention_config: Optional[dict[str, Any]] = None,
34
+) -> RotaryEmbedding:
35
+ if dtype is None:
36
+ dtype = torch.get_default_dtype()
37
+ if rope_scaling is not None:
38
+ # Transforms every value that is a list into a tuple for caching calls
39
+ rope_scaling_tuple = {
40
+ k: tuple(v) if isinstance(v, list) else v
41
+ for k, v in rope_scaling.items()
42
+ }
43
+ rope_scaling_args = tuple(rope_scaling_tuple.items())
44
+ else:
45
+ rope_scaling_args = None
46
+
47
+ if dual_chunk_attention_config is not None:
48
+ dual_chunk_attention_tuple = {
49
+ k: tuple(v) if isinstance(v, list) else v
50
+ for k, v in dual_chunk_attention_config.items()
51
+ if k != "sparse_attention_config"
52
+ }
53
+ dual_chunk_attention_args = tuple(dual_chunk_attention_tuple.items())
54
+ else:
55
+ dual_chunk_attention_args = None
56
+
57
+ if partial_rotary_factor < 1.0:
58
+ rotary_dim = int(rotary_dim * partial_rotary_factor)
59
+ key = (head_size, rotary_dim, max_position, base, is_neox_style,
60
+ rope_scaling_args, dual_chunk_attention_args, dtype)
61
+ if key in _ROPE_DICT:
62
+ return _ROPE_DICT[key]
63
+
64
+ if dual_chunk_attention_config is not None:
65
+ extra_kwargs = {
66
+ k: v
67
+ for k, v in dual_chunk_attention_config.items()
68
+ if k in ("chunk_size", "local_size")
69
+ }
70
+ rotary_emb = DualChunkRotaryEmbedding(head_size, rotary_dim,
71
+ max_position, base,
72
+ is_neox_style, dtype,
73
+ **extra_kwargs)
74
+ elif not rope_scaling:
75
+ rotary_emb = RotaryEmbedding(head_size, rotary_dim, max_position, base,
76
+ is_neox_style, dtype)
77
+ else:
78
+ scaling_type = rope_scaling["rope_type"]
79
+
80
+ if scaling_type == "llama3":
81
+ scaling_factor = rope_scaling["factor"]
82
+ low_freq_factor = rope_scaling["low_freq_factor"]
83
+ high_freq_factor = rope_scaling["high_freq_factor"]
84
+ original_max_position = rope_scaling[
85
+ "original_max_position_embeddings"]
86
+ rotary_emb = Llama3RotaryEmbedding(head_size, rotary_dim,
87
+ max_position, base,
88
+ is_neox_style, dtype,
89
+ scaling_factor, low_freq_factor,
90
+ high_freq_factor,
91
+ original_max_position)
92
+ elif scaling_type == "mllama4":
93
+ rotary_emb = Llama4VisionRotaryEmbedding(head_size, rotary_dim,
94
+ max_position, base,
95
+ is_neox_style, dtype)
96
+ elif scaling_type == "default":
97
+ if "mrope_section" in rope_scaling:
98
+ rotary_emb = MRotaryEmbedding(
99
+ head_size,
100
+ rotary_dim,
101
+ max_position,
102
+ base,
103
+ is_neox_style,
104
+ dtype,
105
+ mrope_section=rope_scaling["mrope_section"],
106
+ )
107
+ else:
108
+ rotary_emb = RotaryEmbedding(
109
+ head_size,
110
+ rotary_dim,
111
+ max_position,
112
+ base,
113
+ is_neox_style,
114
+ dtype,
115
+ )
116
+ elif scaling_type == "linear":
117
+ scaling_factor = rope_scaling["factor"]
118
+ rotary_emb = LinearScalingRotaryEmbedding(head_size, rotary_dim,
119
+ max_position, base,
120
+ is_neox_style,
121
+ scaling_factor, dtype)
122
+ elif scaling_type == "ntk":
123
+ scaling_factor = rope_scaling["factor"]
124
+ mixed_b = rope_scaling.get('mixed_b', None)
125
+ rotary_emb = NTKScalingRotaryEmbedding(head_size, rotary_dim,
126
+ max_position, base,
127
+ is_neox_style,
128
+ scaling_factor, dtype,
129
+ mixed_b)
130
+ elif scaling_type == "dynamic":
131
+ if "alpha" in rope_scaling:
132
+ scaling_alpha = rope_scaling["alpha"]
133
+ rotary_emb = DynamicNTKAlphaRotaryEmbedding(
134
+ head_size, rotary_dim, max_position, base, is_neox_style,
135
+ scaling_alpha, dtype)
136
+ elif "factor" in rope_scaling:
137
+ scaling_factor = rope_scaling["factor"]
138
+ rotary_emb = DynamicNTKScalingRotaryEmbedding(
139
+ head_size, rotary_dim, max_position, base, is_neox_style,
140
+ scaling_factor, dtype)
141
+ else:
142
+ raise ValueError("Dynamic rope scaling must contain either "
143
+ "'alpha' or 'factor' field")
144
+ elif scaling_type == "yarn":
145
+ scaling_factor = rope_scaling["factor"]
146
+ original_max_position = rope_scaling[
147
+ "original_max_position_embeddings"]
148
+ extra_kwargs = {
149
+ k: v
150
+ for k, v in rope_scaling.items()
151
+ if k in ("extrapolation_factor", "attn_factor", "beta_fast",
152
+ "beta_slow")
153
+ }
154
+ rotary_emb = YaRNScalingRotaryEmbedding(head_size, rotary_dim,
155
+ original_max_position,
156
+ base, is_neox_style,
157
+ scaling_factor, dtype,
158
+ **extra_kwargs)
159
+ elif scaling_type == "deepseek_yarn":
160
+ scaling_factor = rope_scaling["factor"]
161
+ original_max_position = rope_scaling[
162
+ "original_max_position_embeddings"]
163
+ # assert max_position == original_max_position * scaling_factor
164
+ extra_kwargs = {
165
+ k: v
166
+ for k, v in rope_scaling.items()
167
+ if k in ("extrapolation_factor", "attn_factor", "beta_fast",
168
+ "beta_slow", "mscale", "mscale_all_dim")
169
+ }
170
+ rotary_emb = DeepseekScalingRotaryEmbedding(
171
+ head_size, rotary_dim, original_max_position, base,
172
+ is_neox_style, scaling_factor, dtype, **extra_kwargs)
173
+ elif scaling_type == "longrope":
174
+ short_factor = rope_scaling["short_factor"]
175
+ long_factor = rope_scaling["long_factor"]
176
+ original_max_position = rope_scaling[
177
+ "original_max_position_embeddings"]
178
+ extra_kwargs = {
179
+ k: v
180
+ for k, v in rope_scaling.items()
181
+ if k in ("short_mscale", "long_mscale")
182
+ }
183
+ rotary_emb = Phi3LongRoPEScaledRotaryEmbedding(
184
+ head_size, rotary_dim, max_position, original_max_position,
185
+ base, is_neox_style, dtype, short_factor, long_factor,
186
+ **extra_kwargs)
187
+ else:
188
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
189
+ _ROPE_DICT[key] = rotary_emb
190
+ return rotary_emb
vllm/model_executor/layers/rotary_embedding/base.py
new
+237
@@ -0,0 +1,237 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Rotary Positional Embeddings Base Class."""
4
+from typing import Optional
5
+
6
+import torch
7
+
8
+from vllm.model_executor.custom_op import CustomOp
9
+
10
+from .common import apply_rotary_emb_dispatch, apply_rotary_emb_torch
11
+
12
+
13
+@CustomOp.register("rotary_embedding")
14
+class RotaryEmbedding(CustomOp):
15
+ """Original rotary positional embedding."""
16
+
17
+ def __init__(
18
+ self,
19
+ head_size: int,
20
+ rotary_dim: int,
21
+ max_position_embeddings: int,
22
+ base: float,
23
+ is_neox_style: bool,
24
+ dtype: torch.dtype,
25
+ ) -> None:
26
+ super().__init__()
27
+ self.head_size = head_size
28
+ self.rotary_dim = rotary_dim
29
+ self.max_position_embeddings = max_position_embeddings
30
+ self.base = base
31
+ self.is_neox_style = is_neox_style
32
+ self.dtype = dtype
33
+
34
+ cache = self._compute_cos_sin_cache()
35
+ cache = cache.to(dtype)
36
+ self.cos_sin_cache: torch.Tensor
37
+ self.register_buffer("cos_sin_cache", cache, persistent=False)
38
+
39
+ def _compute_inv_freq(self, base: float) -> torch.Tensor:
40
+ """Compute the inverse frequency."""
41
+ # NOTE(woosuk): To exactly match the HF implementation, we need to
42
+ # use CPU to compute the cache and then move it to GPU. However, we
43
+ # create the cache on GPU for faster initialization. This may cause
44
+ # a slight numerical difference between the HF implementation and ours.
45
+ inv_freq = 1.0 / (base**(torch.arange(
46
+ 0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim))
47
+ return inv_freq
48
+
49
+ def _compute_cos_sin_cache(self) -> torch.Tensor:
50
+ """Compute the cos and sin cache."""
51
+ inv_freq = self._compute_inv_freq(self.base)
52
+ t = torch.arange(self.max_position_embeddings, dtype=torch.float)
53
+
54
+ freqs = torch.einsum("i,j -> ij", t, inv_freq)
55
+ cos = freqs.cos()
56
+ sin = freqs.sin()
57
+ cache = torch.cat((cos, sin), dim=-1)
58
+ return cache
59
+
60
+ def forward_native(
61
+ self,
62
+ positions: torch.Tensor,
63
+ query: torch.Tensor,
64
+ key: Optional[torch.Tensor] = None,
65
+ offsets: Optional[torch.Tensor] = None,
66
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
67
+ """A PyTorch-native implementation of forward()."""
68
+ if offsets is not None:
69
+ positions = positions + offsets
70
+ positions = positions.flatten()
71
+ num_tokens = positions.shape[0]
72
+ cos_sin = self.cos_sin_cache.index_select(0, positions)
73
+ cos, sin = cos_sin.chunk(2, dim=-1)
74
+
75
+ query_shape = query.shape
76
+ query = query.view(num_tokens, -1, self.head_size)
77
+ query_rot = query[..., :self.rotary_dim]
78
+ query_pass = query[..., self.rotary_dim:]
79
+ query_rot = apply_rotary_emb_torch(query_rot, cos, sin,
80
+ self.is_neox_style)
81
+ query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
82
+
83
+ # key may be None in some cases, e.g. cross-layer KV sharing
84
+ if key is not None:
85
+ key_shape = key.shape
86
+ key = key.view(num_tokens, -1, self.head_size)
87
+ key_rot = key[..., :self.rotary_dim]
88
+ key_pass = key[..., self.rotary_dim:]
89
+ key_rot = apply_rotary_emb_torch(key_rot, cos, sin,
90
+ self.is_neox_style)
91
+ key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
92
+ return query, key
93
+
94
+ def forward_cuda(
95
+ self,
96
+ positions: torch.Tensor,
97
+ query: torch.Tensor,
98
+ key: Optional[torch.Tensor] = None,
99
+ offsets: Optional[torch.Tensor] = None,
100
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
101
+ from vllm import _custom_ops as ops
102
+
103
+ # __setattr__ in nn.Module (called by `self.cos_sin_cache = ...`)
104
+ # is expensive, so avoid calling it if possible
105
+ if self.cos_sin_cache.device != query.device or \
106
+ self.cos_sin_cache.dtype != query.dtype:
107
+ self.cos_sin_cache = self.cos_sin_cache.to(query.device,
108
+ dtype=query.dtype)
109
+
110
+ # ops.rotary_embedding()/batched_rotary_embedding()
111
+ # are in-place operations that update the query and key tensors.
112
+ if offsets is not None:
113
+ ops.batched_rotary_embedding(positions, query, key, self.head_size,
114
+ self.cos_sin_cache,
115
+ self.is_neox_style, self.rotary_dim,
116
+ offsets)
117
+ else:
118
+ ops.rotary_embedding(positions, query, key, self.head_size,
119
+ self.cos_sin_cache, self.is_neox_style)
120
+ return query, key
121
+
122
+ def forward_xpu(
123
+ self,
124
+ positions: torch.Tensor,
125
+ query: torch.Tensor,
126
+ key: Optional[torch.Tensor] = None,
127
+ offsets: Optional[torch.Tensor] = None,
128
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
129
+ from vllm._ipex_ops import ipex_ops as ops
130
+
131
+ self.cos_sin_cache = self.cos_sin_cache.to(positions.device,
132
+ dtype=query.dtype)
133
+ # ops.rotary_embedding()/batched_rotary_embedding()
134
+ # are in-place operations that update the query and key tensors.
135
+ if key is None:
136
+ # XPU kernel doesn't support key=None so fall back to native impl
137
+ # TODO(sarckk): add support for optional key in
138
+ # ipex.llm.functional.rotary_embedding_batched
139
+ return self.forward_native(positions, query, key, offsets)
140
+ else:
141
+ if offsets is not None:
142
+ ops.batched_rotary_embedding(positions, query, key,
143
+ self.head_size,
144
+ self.cos_sin_cache,
145
+ self.is_neox_style,
146
+ self.rotary_dim, offsets)
147
+ else:
148
+ ops.rotary_embedding(positions, query, key, self.head_size,
149
+ self.cos_sin_cache, self.is_neox_style)
150
+ return query, key
151
+
152
+ def forward_neuron(
153
+ self,
154
+ positions: torch.Tensor,
155
+ query: torch.Tensor,
156
+ key: Optional[torch.Tensor] = None,
157
+ offsets: Optional[torch.Tensor] = None,
158
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
159
+
160
+ def _apply_rotary_emb_neuron(
161
+ x: torch.Tensor,
162
+ cos: torch.Tensor,
163
+ sin: torch.Tensor,
164
+ is_neox_style: bool,
165
+ ) -> torch.Tensor:
166
+ cos = cos.unsqueeze(-2).to(x.dtype)
167
+ sin = sin.unsqueeze(-2).to(x.dtype)
168
+ if is_neox_style:
169
+ x1, x2 = torch.chunk(x, 2, dim=-1)
170
+ else:
171
+ # x1 = x[..., ::2]
172
+
173
+ # x2 = x[..., 1::2]
174
+ d = x.shape[-1] // 2
175
+ x_reshaped = x.view(-1, x.shape[-1])
176
+ x1 = x_reshaped[:, ::2].view(*x.shape[:-1], d)
177
+ x2 = x_reshaped[:, 1::2].view(*x.shape[:-1], d)
178
+ o1 = x1 * cos - x2 * sin
179
+ o2 = x2 * cos + x1 * sin
180
+ if is_neox_style:
181
+ return torch.cat((o1, o2), dim=-1)
182
+ else:
183
+ return torch.stack((o1, o2), dim=-1).flatten(-2)
184
+
185
+ if offsets is not None:
186
+ positions = positions + offsets
187
+
188
+ self.cos_sin_cache = self.cos_sin_cache.to(query.device,
189
+ dtype=query.dtype)
190
+
191
+ positions = positions.flatten()
192
+ num_tokens = positions.shape[0]
193
+ cos_sin = self.cos_sin_cache.index_select(0, positions)
194
+ cos, sin = cos_sin.chunk(2, dim=-1)
195
+
196
+ query_shape = query.shape
197
+ query = query.view(num_tokens, -1, self.head_size)
198
+ if key is not None:
199
+ key_shape = key.shape
200
+ key = key.view(num_tokens, -1, self.head_size)
201
+
202
+ if self.rotary_dim == self.head_size:
203
+ query = apply_rotary_emb_dispatch(query, cos, sin,
204
+ self.is_neox_style)
205
+ query = query.reshape(query_shape)
206
+ if key is not None:
207
+ key = apply_rotary_emb_dispatch(key, cos, sin,
208
+ self.is_neox_style)
209
+ key = key.reshape(key_shape)
210
+ else:
211
+ head_size = query.shape[-1]
212
+ query_reshaped = query.view(-1, head_size)
213
+ query_pass = query_reshaped[:, self.rotary_dim:].view(
214
+ *query.shape[:-1], head_size - self.rotary_dim)
215
+ query_rot = query_reshaped[:, :self.rotary_dim].view(
216
+ *query.shape[:-1], self.rotary_dim)
217
+ query_rot = _apply_rotary_emb_neuron(query_rot, cos, sin,
218
+ self.is_neox_style)
219
+ query = torch.cat((query_rot, query_pass),
220
+ dim=-1).reshape(query_shape)
221
+
222
+ if key is not None:
223
+ key_reshaped = key.view(-1, head_size)
224
+ key_pass = key_reshaped[:, self.rotary_dim:].view(
225
+ *key.shape[:-1], head_size - self.rotary_dim)
226
+ key_rot = key_reshaped[:, :self.rotary_dim].view(
227
+ *key.shape[:-1], self.rotary_dim)
228
+ key_rot = _apply_rotary_emb_neuron(key_rot, cos, sin,
229
+ self.is_neox_style)
230
+ key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
231
+ return query, key
232
+
233
+ def extra_repr(self) -> str:
234
+ s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
235
+ s += f", max_position_embeddings={self.max_position_embeddings}"
236
+ s += f", base={self.base}, is_neox_style={self.is_neox_style}"
237
+ return s
vllm/model_executor/layers/rotary_embedding/common.py
new
+105
@@ -0,0 +1,105 @@
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 vllm.platforms import current_platform
9
+
10
+if current_platform.is_cuda():
11
+ from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
12
+
13
+
14
+# common functions
15
+def rotate_neox(x: torch.Tensor) -> torch.Tensor:
16
+ x1 = x[..., :x.shape[-1] // 2]
17
+ x2 = x[..., x.shape[-1] // 2:]
18
+ return torch.cat((-x2, x1), dim=-1)
19
+
20
+
21
+def rotate_gptj(x: torch.Tensor) -> torch.Tensor:
22
+ x1 = x[..., ::2]
23
+ x2 = x[..., 1::2]
24
+ x = torch.stack((-x2, x1), dim=-1)
25
+ return x.flatten(-2)
26
+
27
+
28
+def apply_rotary_emb_torch(
29
+ x: torch.Tensor,
30
+ cos: torch.Tensor,
31
+ sin: torch.Tensor,
32
+ is_neox_style: bool,
33
+) -> torch.Tensor:
34
+ cos = cos.unsqueeze(-2).to(x.dtype)
35
+ sin = sin.unsqueeze(-2).to(x.dtype)
36
+ if is_neox_style:
37
+ x1, x2 = torch.chunk(x, 2, dim=-1)
38
+ else:
39
+ x1 = x[..., ::2]
40
+ x2 = x[..., 1::2]
41
+ o1 = x1 * cos - x2 * sin
42
+ o2 = x2 * cos + x1 * sin
43
+ if is_neox_style:
44
+ return torch.cat((o1, o2), dim=-1)
45
+ else:
46
+ return torch.stack((o1, o2), dim=-1).flatten(-2)
47
+
48
+
49
+def apply_rotary_emb_dispatch(x: torch.Tensor, cos: torch.Tensor,
50
+ sin: torch.Tensor,
51
+ is_neox_style: bool) -> torch.Tensor:
52
+ """
53
+ Args:
54
+ x: [num_tokens, num_heads, head_size]
55
+ cos: [num_tokens, head_size // 2]
56
+ sin: [num_tokens, head_size // 2]
57
+ is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
58
+ positional embeddings.
59
+ """
60
+ if current_platform.is_cuda():
61
+ return apply_rotary_emb(x.unsqueeze(0), cos, sin,
62
+ not is_neox_style).squeeze(0)
63
+ else:
64
+ return apply_rotary_emb_torch(x, cos, sin, is_neox_style)
65
+
66
+
67
+# yarn functions
68
+# Inverse dim formula to find dim based on number of rotations
69
+def yarn_find_correction_dim(num_rotations: int,
70
+ dim: int,
71
+ base: float = 10000,
72
+ max_position_embeddings: int = 2048) -> float:
73
+ return (dim * math.log(max_position_embeddings /
74
+ (num_rotations * 2 * math.pi))) / (2 *
75
+ math.log(base))
76
+
77
+
78
+# Find dim range bounds based on rotations
79
+def yarn_find_correction_range(
80
+ low_rot: int,
81
+ high_rot: int,
82
+ dim: int,
83
+ base: float = 10000,
84
+ max_position_embeddings: int = 2048) -> tuple[int, int]:
85
+ low = math.floor(
86
+ yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings))
87
+ high = math.ceil(
88
+ yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings))
89
+ return max(low, 0), min(high, dim - 1) # Clamp values just in case
90
+
91
+
92
+def yarn_linear_ramp_mask(low: float, high: float, dim: int,
93
+ dtype: torch.dtype) -> torch.Tensor:
94
+ if low == high:
95
+ high += 0.001 # Prevent singularity
96
+
97
+ linear_func = (torch.arange(dim, dtype=dtype) - low) / (high - low)
98
+ ramp_func = torch.clamp(linear_func, 0, 1)
99
+ return ramp_func
100
+
101
+
102
+def yarn_get_mscale(scale: float = 1) -> float:
103
+ if scale <= 1:
104
+ return 1.0
105
+ return 0.1 * math.log(scale) + 1.0
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
new
+131
@@ -0,0 +1,131 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import math
5
+from typing import Optional
6
+
7
+import torch
8
+
9
+from vllm.platforms import current_platform
10
+
11
+from .base import RotaryEmbedding
12
+from .common import (rotate_gptj, rotate_neox, yarn_find_correction_range,
13
+ yarn_linear_ramp_mask)
14
+
15
+
16
+def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
17
+ if scale <= 1:
18
+ return 1.0
19
+ return 0.1 * mscale * math.log(scale) + 1.0
20
+
21
+
22
+class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
23
+ """RotaryEmbedding extended with YaRN method.
24
+
25
+ Credits to Peng et al. github.com/jquesnelle/yarn
26
+ """
27
+
28
+ def __init__(
29
+ self,
30
+ head_size: int,
31
+ rotary_dim: int,
32
+ max_position_embeddings: int,
33
+ base: float,
34
+ is_neox_style: bool,
35
+ scaling_factor: float,
36
+ dtype: torch.dtype,
37
+ *,
38
+ extrapolation_factor: float = 1,
39
+ attn_factor: float = 1,
40
+ beta_fast: int = 32,
41
+ beta_slow: int = 1,
42
+ mscale: float = 1,
43
+ mscale_all_dim: float = 0,
44
+ ) -> None:
45
+ self.scaling_factor = scaling_factor
46
+ self.extrapolation_factor = extrapolation_factor
47
+ self.attn_factor = attn_factor
48
+ self.beta_fast = beta_fast
49
+ self.beta_slow = beta_slow
50
+ # Get n-d magnitude scaling corrected for interpolation.
51
+ self.mscale = float(
52
+ yarn_get_mscale(self.scaling_factor, float(mscale)) /
53
+ yarn_get_mscale(self.scaling_factor, float(mscale_all_dim)) *
54
+ attn_factor)
55
+ super().__init__(head_size, rotary_dim, max_position_embeddings, base,
56
+ is_neox_style, dtype)
57
+
58
+ def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
59
+ pos_freqs = self.base**(
60
+ torch.arange(0,
61
+ self.rotary_dim,
62
+ 2,
63
+ dtype=torch.float,
64
+ device=current_platform.device_type) /
65
+ self.rotary_dim)
66
+ inv_freq_extrapolation = 1.0 / pos_freqs
67
+ inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
68
+
69
+ low, high = yarn_find_correction_range(self.beta_fast, self.beta_slow,
70
+ self.rotary_dim, self.base,
71
+ self.max_position_embeddings)
72
+ # Get n-d rotational scaling corrected for extrapolation
73
+ inv_freq_mask = (1 - yarn_linear_ramp_mask(
74
+ low, high, self.rotary_dim // 2,
75
+ dtype=torch.float)) * self.extrapolation_factor
76
+ inv_freq = inv_freq_interpolation * (
77
+ 1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask
78
+ return inv_freq
79
+
80
+ def _compute_cos_sin_cache(self) -> torch.Tensor:
81
+ inv_freq = self._compute_inv_freq(self.scaling_factor)
82
+ t = torch.arange(self.max_position_embeddings * self.scaling_factor,
83
+ device=current_platform.device_type,
84
+ dtype=torch.float32)
85
+ freqs = torch.einsum("i,j -> ij", t, inv_freq)
86
+ cos = (freqs.cos() * self.mscale)
87
+ sin = (freqs.sin() * self.mscale)
88
+ cache = torch.cat((cos, sin), dim=-1)
89
+ return cache
90
+
91
+ def forward(
92
+ self,
93
+ positions: torch.Tensor,
94
+ query: torch.Tensor,
95
+ key: Optional[torch.Tensor] = None,
96
+ offsets: Optional[torch.Tensor] = None,
97
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
98
+ """PyTorch-native implementation equivalent to forward()."""
99
+ assert key is not None
100
+ query_rot = query[..., :self.rotary_dim]
101
+ key_rot = key[..., :self.rotary_dim]
102
+ if self.rotary_dim < self.head_size:
103
+ query_pass = query[..., self.rotary_dim:]
104
+ key_pass = key[..., self.rotary_dim:]
105
+
106
+ if self.cos_sin_cache.device != positions.device:
107
+ self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(
108
+ positions.device)
109
+ cos_sin = self.cos_sin_cache[torch.add(positions, offsets)
110
+ if offsets is not None else positions]
111
+ cos, sin = cos_sin.chunk(2, dim=-1)
112
+ if self.is_neox_style:
113
+ # NOTE(woosuk): Here we assume that the positions tensor has the
114
+ # shape [batch_size, seq_len].
115
+ cos = cos.repeat(1, 1, 2).unsqueeze(-2)
116
+ sin = sin.repeat(1, 1, 2).unsqueeze(-2)
117
+ else:
118
+ cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
119
+ sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)
120
+
121
+ rotate_fn = rotate_neox if self.is_neox_style else rotate_gptj
122
+ query_rot = query_rot * cos + rotate_fn(query_rot) * sin
123
+ key_rot = key_rot * cos + rotate_fn(key_rot) * sin
124
+
125
+ if self.rotary_dim < self.head_size:
126
+ query = torch.cat((query_rot, query_pass), dim=-1)
127
+ key = torch.cat((key_rot, key_pass), dim=-1)
128
+ else:
129
+ query = query_rot
130
+ key = key_rot
131
+ return query, key
vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py
new
+188
@@ -0,0 +1,188 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Optional
5
+
6
+import torch
7
+
8
+from vllm.model_executor.custom_op import CustomOp
9
+
10
+from .common import rotate_gptj, rotate_neox
11
+
12
+
13
+@CustomOp.register("dual_chunk_rotary_embedding")
14
+class DualChunkRotaryEmbedding(CustomOp):
15
+ """Rotary positional embedding for Dual Chunk Attention."""
16
+
17
+ def __init__(
18
+ self,
19
+ head_size: int,
20
+ rotary_dim: int,
21
+ max_position_embeddings: int,
22
+ base: float,
23
+ is_neox_style: bool,
24
+ dtype: torch.dtype,
25
+ chunk_size: int,
26
+ local_size: int,
27
+ ) -> None:
28
+ super().__init__()
29
+ self.head_size = head_size
30
+ self.rotary_dim = rotary_dim
31
+ self.max_position_embeddings = max_position_embeddings
32
+ self.base = base
33
+ self.is_neox_style = is_neox_style
34
+ self.chunk_size = chunk_size
35
+ self.local_size = local_size
36
+ self.dtype = dtype
37
+ self.device = torch.device(f"cuda:{torch.cuda.current_device()}")
38
+ (q_cache, qc_cache, k_cache, qc_no_clamp_cache,
39
+ q_inter_cache) = self._compute_cos_sin_cache()
40
+
41
+ self.register_buffer("cos_sin_q_cache", q_cache, persistent=False)
42
+ self.register_buffer("cos_sin_qc_cache", qc_cache, persistent=False)
43
+ self.register_buffer("cos_sin_k_cache", k_cache, persistent=False)
44
+ self.register_buffer("cos_sin_qc_no_clamp_cache",
45
+ qc_no_clamp_cache,
46
+ persistent=False)
47
+ self.register_buffer("cos_sin_q_inter_cache",
48
+ q_inter_cache,
49
+ persistent=False)
50
+
51
+ def _compute_inv_freq(self, base: float) -> torch.Tensor:
52
+ """Compute the inverse frequency."""
53
+ # NOTE(woosuk): The HF implementation uses `torch.arange(...).float()`.
54
+ # However, we use `torch.arange(..., dtype=torch.float)` instead to
55
+ # avoid numerical issues with large base values (e.g., 10000000).
56
+ # This may cause a slight numerical difference between the HF
57
+ # implementation and ours.
58
+ # NOTE(woosuk): To exactly match the HF implementation, we need to
59
+ # use CPU to compute the cache and then move it to GPU. However, we
60
+ # create the cache on GPU for faster initialization. This may cause
61
+ # a slight numerical difference between the HF implementation and ours.
62
+ inv_freq = 1.0 / (base**(torch.arange(
63
+ 0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim))
64
+ return inv_freq
65
+
66
+ def _compute_cos_sin_cache(self) -> torch.Tensor:
67
+ """Compute the cos and sin cache."""
68
+ inv_freq = self._compute_inv_freq(self.base)
69
+ chunk_len = self.chunk_size - self.local_size
70
+ q_t = torch.arange(chunk_len, dtype=torch.float)
71
+ qc_t = (torch.arange(chunk_len, dtype=torch.float) +
72
+ chunk_len).clamp(max=self.chunk_size)
73
+ k_t = torch.arange(self.max_position_embeddings,
74
+ dtype=torch.float) % chunk_len
75
+
76
+ # count from chunk_len, no clamp(self.chunk_size) restriction
77
+ qc_no_clamp_t = torch.arange(chunk_len, dtype=torch.float) + chunk_len
78
+ # count from self.chunk_size for q_inter's rope
79
+ q_inter_t = torch.arange(chunk_len,
80
+ dtype=torch.float) + self.chunk_size
81
+
82
+ q_freqs = torch.outer(q_t, inv_freq)
83
+ qc_freqs = torch.outer(qc_t, inv_freq)
84
+ k_freqs = torch.outer(k_t, inv_freq)
85
+ qc_no_clamp_freqs = torch.outer(qc_no_clamp_t, inv_freq)
86
+ q_inter_freqs = torch.outer(q_inter_t, inv_freq)
87
+
88
+ q_cos = q_freqs.cos()
89
+ q_sin = q_freqs.sin()
90
+ qc_cos = qc_freqs.cos()
91
+ qc_sin = qc_freqs.sin()
92
+ k_cos = k_freqs.cos()
93
+ k_sin = k_freqs.sin()
94
+
95
+ qc_no_clamp_cos = qc_no_clamp_freqs.cos()
96
+ qc_no_clamp_sin = qc_no_clamp_freqs.sin()
97
+ q_inter_cos = q_inter_freqs.cos()
98
+ q_inter_sin = q_inter_freqs.sin()
99
+
100
+ q_cache = torch.cat((q_cos, q_sin), dim=-1).to(dtype=self.dtype,
101
+ device=self.device)
102
+ qc_cache = torch.cat((qc_cos, qc_sin), dim=-1).to(dtype=self.dtype,
103
+ device=self.device)
104
+ k_cache = torch.cat((k_cos, k_sin), dim=-1).to(dtype=self.dtype,
105
+ device=self.device)
106
+ qc_no_clamp_cache = torch.cat((qc_no_clamp_cos, qc_no_clamp_sin),
107
+ dim=-1).to(dtype=self.dtype,
108
+ device=self.device)
109
+ q_inter_cache = torch.cat((q_inter_cos, q_inter_sin),
110
+ dim=-1).to(dtype=self.dtype,
111
+ device=self.device)
112
+ return q_cache, qc_cache, k_cache, qc_no_clamp_cache, q_inter_cache
113
+
114
+ def forward(
115
+ self,
116
+ positions: torch.Tensor,
117
+ query: torch.Tensor,
118
+ key: torch.Tensor,
119
+ offsets: Optional[torch.Tensor] = None,
120
+ ) -> tuple[torch.Tensor, torch.Tensor]:
121
+ query = query.view(*query.shape[:-1], -1, self.head_size)
122
+ key = key.view(*key.shape[:-1], -1, self.head_size)
123
+ query_rot = query[..., :self.rotary_dim]
124
+ key_rot = key[..., :self.rotary_dim]
125
+ if self.rotary_dim < self.head_size:
126
+ query_pass = query[..., self.rotary_dim:]
127
+ key_pass = key[..., self.rotary_dim:]
128
+ else:
129
+ query_pass = None
130
+ key_pass = None
131
+
132
+ positions_with_offsets = (torch.add(positions, offsets)
133
+ if offsets is not None else positions)
134
+ key = self._apply_rotary_embedding(
135
+ self.cos_sin_k_cache[positions_with_offsets], key_rot, key_pass)
136
+ chunk_len = self.chunk_size - self.local_size
137
+ query = self._apply_rotary_embedding(
138
+ self.cos_sin_q_cache[positions_with_offsets % chunk_len],
139
+ query_rot, query_pass)
140
+ query_succ = self._apply_rotary_embedding(
141
+ self.cos_sin_qc_cache[positions_with_offsets % chunk_len],
142
+ query_rot, query_pass)
143
+ query_inter = self._apply_rotary_embedding(
144
+ self.cos_sin_qc_cache[chunk_len - 1].repeat(positions.shape[0], 1),
145
+ query_rot, query_pass)
146
+ query_succ_critical = self._apply_rotary_embedding(
147
+ self.cos_sin_qc_no_clamp_cache[positions_with_offsets % chunk_len],
148
+ query_rot, query_pass)
149
+ query_inter_critical = self._apply_rotary_embedding(
150
+ self.cos_sin_q_inter_cache[positions_with_offsets % chunk_len],
151
+ query_rot, query_pass)
152
+
153
+ # merge query into one tensor to simplify the interfaces
154
+ query = torch.cat((
155
+ query,
156
+ query_succ,
157
+ query_inter,
158
+ query_succ_critical,
159
+ query_inter_critical,
160
+ ),
161
+ dim=-1)
162
+ return query, key
163
+
164
+ def _apply_rotary_embedding(self, cos_sin, hidden_rot, hidden_pass):
165
+ cos, sin = cos_sin.chunk(2, dim=-1)
166
+ if self.is_neox_style:
167
+ # NOTE(woosuk): Here we assume that the positions tensor has the
168
+ # shape [batch_size, seq_len].
169
+ cos = cos.repeat(1, 1, 2).unsqueeze(-2)
170
+ sin = sin.repeat(1, 1, 2).unsqueeze(-2)
171
+ else:
172
+ cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
173
+ sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)
174
+ rotate_fn = rotate_neox if self.is_neox_style else rotate_gptj
175
+ hidden_rot = hidden_rot * cos + rotate_fn(hidden_rot) * sin
176
+
177
+ if self.rotary_dim < self.head_size:
178
+ hidden = torch.cat((hidden_rot, hidden_pass), dim=-1)
179
+ else:
180
+ hidden = hidden_rot
181
+ return hidden.flatten(-2).squeeze(0)
182
+
183
+ def extra_repr(self) -> str:
184
+ s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
185
+ s += f", max_position_embeddings={self.max_position_embeddings}"
186
+ s += f", base={self.base}, is_neox_style={self.is_neox_style}"
187
+ s += f", chunk_size={self.chunk_size}, local_size={self.local_size}"
188
+ return s
vllm/model_executor/layers/rotary_embedding/dynamic_ntk_alpha_rope.py
new
+41
@@ -0,0 +1,41 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from .base import RotaryEmbedding
7
+
8
+
9
+class DynamicNTKAlphaRotaryEmbedding(RotaryEmbedding):
10
+ """RotaryEmbedding extended with Dynamic NTK alpha.
11
+
12
+ Based on the original RotaryEmbedding implementation.
13
+ """
14
+
15
+ def __init__(
16
+ self,
17
+ head_size: int,
18
+ rotary_dim: int,
19
+ max_position_embeddings: int,
20
+ base: float,
21
+ is_neox_style: bool,
22
+ scaling_alpha: float,
23
+ dtype: torch.dtype,
24
+ ) -> None:
25
+ self.scaling_alpha = scaling_alpha
26
+ super().__init__(head_size, rotary_dim, max_position_embeddings, base,
27
+ is_neox_style, dtype)
28
+
29
+ def _compute_cos_sin_cache(self) -> torch.Tensor:
30
+ # For Hunyuan DynamicNTKAlphaRotaryEmbedding
31
+ max_len = self.max_position_embeddings
32
+ base = self.base * self.scaling_alpha**(self.rotary_dim /
33
+ (self.rotary_dim - 2))
34
+ inv_freq = self._compute_inv_freq(base)
35
+ t = torch.arange(max_len, dtype=torch.float)
36
+
37
+ freqs = torch.einsum("i,j -> ij", t, inv_freq)
38
+ cos = freqs.cos()
39
+ sin = freqs.sin()
40
+ cache = torch.cat((cos, sin), dim=-1)
41
+ return cache
vllm/model_executor/layers/rotary_embedding/dynamic_ntk_scaling_rope.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# Adapted from
5
+# https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/llama/modeling_llama.py
6
+# Copyright 2023 The vLLM team.
7
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
8
+#
9
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
10
+# and OPT implementations in this library. It has been modified from its
11
+# original forms to accommodate minor architectural differences compared
12
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13
+#
14
+# Licensed under the Apache License, Version 2.0 (the "License");
15
+# you may not use this file except in compliance with the License.
16
+# You may obtain a copy of the License at
17
+#
18
+# http://www.apache.org/licenses/LICENSE-2.0
19
+#
20
+# Unless required by applicable law or agreed to in writing, software
21
+# distributed under the License is distributed on an "AS IS" BASIS,
22
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23
+# See the License for the specific language governing permissions and
24
+# limitations under the License.
25
+import torch
26
+
27
+from .base import RotaryEmbedding
28
+
29
+
30
+class DynamicNTKScalingRotaryEmbedding(RotaryEmbedding):
31
+ """RotaryEmbedding extended with Dynamic NTK scaling.
32
+
33
+ Credits to the Reddit users /u/bloc97 and /u/emozilla
34
+ """
35
+
36
+ def __init__(
37
+ self,
38
+ head_size: int,
39
+ rotary_dim: int,
40
+ max_position_embeddings: int,
41
+ base: float,
42
+ is_neox_style: bool,
43
+ scaling_factor: float,
44
+ dtype: torch.dtype,
45
+ ) -> None:
46
+ self.scaling_factor = scaling_factor
47
+ super().__init__(head_size, rotary_dim, max_position_embeddings, base,
48
+ is_neox_style, dtype)
49
+
50
+ def _compute_cos_sin_cache(self) -> torch.Tensor:
51
+ # NOTE(woosuk): self.max_position_embeddings is the original
52
+ # maximum length before applying the rope scaling.
53
+ # Thus, the maximum length after applying the rope scaling is
54
+ # self.max_position_embeddings * self.scaling_factor.
55
+ max_len = self.max_position_embeddings * self.scaling_factor
56
+ base = self.base * (
57
+ (self.scaling_factor * max_len / self.max_position_embeddings) -
58
+ (self.scaling_factor - 1))**(self.rotary_dim /
59
+ (self.rotary_dim - 2))
60
+ inv_freq = self._compute_inv_freq(base)
61
+ t = torch.arange(max_len, dtype=torch.float)
62
+
63
+ freqs = torch.einsum("i,j -> ij", t, inv_freq)
64
+ cos = freqs.cos()
65
+ sin = freqs.sin()
66
+ cache = torch.cat((cos, sin), dim=-1)
67
+ return cache
vllm/model_executor/layers/rotary_embedding/linear_scaling_rope.py
new
+115
@@ -0,0 +1,115 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Union
5
+
6
+# Adapted from
7
+# https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/llama/modeling_llama.py
8
+# Copyright 2023 The vLLM team.
9
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
10
+#
11
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
12
+# and OPT implementations in this library. It has been modified from its
13
+# original forms to accommodate minor architectural differences compared
14
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
15
+#
16
+# Licensed under the Apache License, Version 2.0 (the "License");
17
+# you may not use this file except in compliance with the License.
18
+# You may obtain a copy of the License at
19
+#
20
+# http://www.apache.org/licenses/LICENSE-2.0
21
+#
22
+# Unless required by applicable law or agreed to in writing, software
23
+# distributed under the License is distributed on an "AS IS" BASIS,
24
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
25
+# See the License for the specific language governing permissions and
26
+# limitations under the License.
27
+import torch
28
+
29
+from .base import RotaryEmbedding
30
+
31
+
32
+class LinearScalingRotaryEmbedding(RotaryEmbedding):
33
+ """RotaryEmbedding extended with linear scaling.
34
+
35
+ It supports multiple scaling factors. Since multiple LoRA adapters may have
36
+ different scaling factors, we need multiple cos/sin caches. In this way,
37
+ instead of running rotary embedding kernel per lora, we can run multiple
38
+ lora in a batched way.
39
+
40
+ In addition to that, we also keep the cos/sin cache for the scaling factor
41
+ of 1 (default) at all times.
42
+
43
+ Exemplary for two scaling factors x=1, y and z with embeddings
44
+ [[x11, x12, ... x1m], ..., [xn1, xn2, ..., xnm]] and
45
+ [[y11, y12, ... y1o], ..., [yn1, yn2, ..., yno]], and
46
+ [[z11, z12, ... z1p], ..., [zn1, zn2, ..., znp]],
47
+
48
+ we construct the cos/sin cache as follows:
49
+ [[x11, x12, ... x1m, y11, y12, ... y1o, z11, z12, ... z1p],
50
+ ...
51
+ [xn1, xn2, ... xnm, yn1, yn2, ... yno, zn1, zn2, ... znp]]
52
+
53
+ We then use offsets to index into the cos/sin cache for
54
+ the respective scaling factors.
55
+
56
+ The offset to cache can be accessed via `scaling_factor_to_offset` API.
57
+
58
+ Credits to the Reddit user /u/kaiokendev
59
+ """
60
+
61
+ def __init__(
62
+ self,
63
+ head_size: int,
64
+ rotary_dim: int,
65
+ max_position_embeddings: int,
66
+ base: float,
67
+ is_neox_style: bool,
68
+ scaling_factors: Union[list[float], float],
69
+ dtype: torch.dtype,
70
+ ) -> None:
71
+ if isinstance(scaling_factors, float):
72
+ scaling_factors = [scaling_factors]
73
+ self.scaling_factors: list[float] = scaling_factors # noqa
74
+ super().__init__(head_size, rotary_dim, max_position_embeddings, base,
75
+ is_neox_style, dtype)
76
+ # Lazy initialized.
77
+ self._scaling_factor_to_offset: dict[float, int]
78
+
79
+ def _compute_cos_sin_cache(self) -> torch.Tensor:
80
+ inv_freq = self._compute_inv_freq(self.base)
81
+ cache_list: list[torch.Tensor] = []
82
+ # offsets to the next cache in a tensor.
83
+ # Each offset corresponds to the same index in scaling_factors.
84
+ offsets: list[int] = []
85
+ for scaling_factor in self.scaling_factors:
86
+ # NOTE(woosuk): self.max_position_embeddings is the original
87
+ # maximum length before applying the rope scaling.
88
+ # Thus, the maximum length after applying the rope scaling is
89
+ # self.max_position_embeddings * self.scaling_factor.
90
+ max_len = self.max_position_embeddings * scaling_factor
91
+ t = torch.arange(max_len, dtype=torch.float)
92
+ t = t / scaling_factor
93
+
94
+ freqs = torch.einsum("i,j -> ij", t, inv_freq)
95
+ cos = freqs.cos()
96
+ sin = freqs.sin()
97
+ cache = torch.cat((cos, sin), dim=-1)
98
+ if not cache_list:
99
+ offset = 0
100
+ else:
101
+ last_offset = offsets[-1]
102
+ next_max_len = cache_list[-1].shape[0]
103
+ offset = last_offset + next_max_len
104
+ offsets.append(offset)
105
+ cache_list.append(cache)
106
+ self._scaling_factor_to_offset = {
107
+ float(scaling_factor): offsets[i]
108
+ for i, scaling_factor in enumerate(self.scaling_factors)
109
+ }
110
+ assert len(self.scaling_factors) == len(offsets)
111
+ return torch.cat(cache_list, dim=0)
112
+
113
+ @property
114
+ def scaling_factor_to_offset(self) -> dict[float, int]:
115
+ return self._scaling_factor_to_offset
vllm/model_executor/layers/rotary_embedding/llama3_rope.py
new
+54
@@ -0,0 +1,54 @@
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 RotaryEmbedding
9
+
10
+
11
+class Llama3RotaryEmbedding(RotaryEmbedding):
12
+
13
+ def __init__(
14
+ self,
15
+ head_size: int,
16
+ rotary_dim: int,
17
+ max_position_embeddings: int,
18
+ base: float,
19
+ is_neox_style: bool,
20
+ dtype: torch.dtype,
21
+ scaling_factor: float,
22
+ low_freq_factor: float,
23
+ high_freq_factor: float,
24
+ orig_max_position: int,
25
+ ) -> None:
26
+ self.scaling_factor = scaling_factor
27
+ self.low_freq_factor = low_freq_factor
28
+ self.high_freq_factor = high_freq_factor
29
+ self.orig_max_position = orig_max_position
30
+ super().__init__(head_size, rotary_dim, max_position_embeddings, base,
31
+ is_neox_style, dtype)
32
+
33
+ def _compute_inv_freq(self, base: float) -> torch.Tensor:
34
+ inv_freqs = super()._compute_inv_freq(base)
35
+ low_freq_wavelen = self.orig_max_position / self.low_freq_factor
36
+ high_freq_wavelen = self.orig_max_position / self.high_freq_factor
37
+
38
+ wave_len = 2 * math.pi / inv_freqs
39
+ if self.low_freq_factor != self.high_freq_factor:
40
+ smooth = (self.orig_max_position / wave_len - self.low_freq_factor
41
+ ) / (self.high_freq_factor - self.low_freq_factor)
42
+ else:
43
+ smooth = 0
44
+ new_freqs = torch.where(
45
+ wave_len < high_freq_wavelen,
46
+ inv_freqs,
47
+ torch.where(
48
+ wave_len > low_freq_wavelen,
49
+ inv_freqs / self.scaling_factor,
50
+ (1 - smooth) * inv_freqs / self.scaling_factor +
51
+ smooth * inv_freqs,
52
+ ),
53
+ )
54
+ return new_freqs
vllm/model_executor/layers/rotary_embedding/llama4_vision_rope.py
new
+74
@@ -0,0 +1,74 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import math
5
+from typing import Optional
6
+
7
+import torch
8
+
9
+from .base import RotaryEmbedding
10
+
11
+
12
+class Llama4VisionRotaryEmbedding(RotaryEmbedding):
13
+
14
+ def __init__(
15
+ self,
16
+ head_size: int,
17
+ rotary_dim: int,
18
+ max_position_embeddings: int,
19
+ base: float,
20
+ is_neox_style: bool,
21
+ dtype: torch.dtype,
22
+ ):
23
+ super().__init__(head_size, rotary_dim, max_position_embeddings, base,
24
+ is_neox_style, dtype)
25
+
26
+ def _compute_inv_freq(self, base: float) -> torch.Tensor:
27
+ inv_freqs = super()._compute_inv_freq(base)
28
+ inv_freqs = inv_freqs[:(self.rotary_dim // 2)]
29
+ return inv_freqs
30
+
31
+ def _compute_cos_sin_cache(self) -> torch.Tensor:
32
+ inv_freq = self._compute_inv_freq(self.base)
33
+
34
+ # self.max_position_embeddings here is number of image patches
35
+ # i.e. (image_size // patch_size) ** 2
36
+ num_patches = self.max_position_embeddings
37
+ img_idx = torch.arange(num_patches,
38
+ dtype=torch.int32) \
39
+ .reshape(num_patches, 1)
40
+ img_idx = torch.cat([img_idx, img_idx[:1]], dim=0)
41
+ img_idx[-1, -1] = -2 # set to ID_CLS_TOKEN
42
+ num_patches_single_dim = int(math.sqrt(num_patches))
43
+ frequencies_x = img_idx % num_patches_single_dim
44
+ frequencies_y = img_idx // num_patches_single_dim
45
+ freqs_x = ((frequencies_x + 1)[..., None] *
46
+ inv_freq[None, None, :]).repeat_interleave(2, dim=-1)
47
+ freqs_y = ((frequencies_y + 1)[..., None] *
48
+ inv_freq[None, None, :]).repeat_interleave(2, dim=-1)
49
+ freqs = torch.cat([freqs_x, freqs_y],
50
+ dim=-1).float().contiguous()[..., ::2]
51
+ freqs = freqs.masked_fill(img_idx.reshape(-1, 1, 1) < 0, 0)
52
+ cache = torch.view_as_complex(
53
+ torch.stack([torch.cos(freqs), torch.sin(freqs)], dim=-1))
54
+ return cache
55
+
56
+ def forward(
57
+ self,
58
+ query: torch.Tensor,
59
+ key: Optional[torch.Tensor] = None,
60
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
61
+ assert key is not None
62
+ self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(query.device)
63
+ query_ = torch.view_as_complex(query.float().reshape(
64
+ *query.shape[:-1], -1, 2))
65
+ key_ = torch.view_as_complex(key.float().reshape(
66
+ *key.shape[:-1], -1, 2))
67
+ broadcast_shape = [
68
+ d if i == 1 or i == (query_.ndim - 1) else 1
69
+ for i, d in enumerate(query_.shape)
70
+ ]
71
+ freqs_ci = self.cos_sin_cache.view(*broadcast_shape)
72
+ query_out = torch.view_as_real(query_ * freqs_ci).flatten(3)
73
+ key_out = torch.view_as_real(key_ * freqs_ci).flatten(3)
74
+ return query_out.type_as(query), key_out.type_as(key)
vllm/model_executor/layers/rotary_embedding/mrope.py
new
+670
@@ -0,0 +1,670 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import itertools
5
+from typing import Optional, Union
6
+
7
+import numpy as np
8
+import torch
9
+from transformers import PretrainedConfig
10
+
11
+from .base import RotaryEmbedding
12
+from .common import apply_rotary_emb_dispatch
13
+
14
+
15
+class MRotaryEmbedding(RotaryEmbedding):
16
+ """Rotary Embedding with Multimodal Sections."""
17
+
18
+ def __init__(
19
+ self,
20
+ head_size: int,
21
+ rotary_dim: int,
22
+ max_position_embeddings: int,
23
+ base: float,
24
+ is_neox_style: bool,
25
+ dtype: torch.dtype,
26
+ mrope_section: Optional[list[int]] = None,
27
+ ) -> None:
28
+ # In Qwen2.5-VL, the maximum index value is related to the duration of
29
+ # the input video. We enlarge max_position_embeddings to 4 times to get
30
+ # a larger the cos and sin cache.
31
+ self.cache_max_position_num = max_position_embeddings * 4
32
+ super().__init__(head_size, rotary_dim, self.cache_max_position_num,
33
+ base, is_neox_style, dtype)
34
+
35
+ self.mrope_section = mrope_section
36
+ if self.mrope_section:
37
+ assert sum(self.mrope_section) == rotary_dim // 2
38
+
39
+ def forward(
40
+ self,
41
+ positions: torch.Tensor,
42
+ query: torch.Tensor,
43
+ key: Optional[torch.Tensor] = None,
44
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
45
+ """PyTorch-native implementation equivalent to forward().
46
+
47
+ Args:
48
+ positions:
49
+ [num_tokens,] (text only) or
50
+ [3, num_tokens] (T/H/W positions with multimodal inputs)
51
+ query: [num_tokens, num_heads * head_size]
52
+ key: [num_tokens, num_kv_heads * head_size]
53
+ """
54
+ assert positions.ndim == 1 or positions.ndim == 2
55
+ assert key is not None
56
+
57
+ num_tokens = positions.shape[-1]
58
+ cos_sin = self.cos_sin_cache[positions]
59
+ cos, sin = cos_sin.chunk(2, dim=-1)
60
+ if positions.ndim == 2:
61
+ assert self.mrope_section
62
+
63
+ cos = torch.cat([
64
+ m[i]
65
+ for i, m in enumerate(cos.split(self.mrope_section, dim=-1))
66
+ ],
67
+ dim=-1)
68
+ sin = torch.cat([
69
+ m[i]
70
+ for i, m in enumerate(sin.split(self.mrope_section, dim=-1))
71
+ ],
72
+ dim=-1)
73
+
74
+ query_shape = query.shape
75
+ query = query.view(num_tokens, -1, self.head_size)
76
+ query_rot = query[..., :self.rotary_dim]
77
+ query_pass = query[..., self.rotary_dim:]
78
+ query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin,
79
+ self.is_neox_style)
80
+ query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
81
+
82
+ key_shape = key.shape
83
+ key = key.view(num_tokens, -1, self.head_size)
84
+ key_rot = key[..., :self.rotary_dim]
85
+ key_pass = key[..., self.rotary_dim:]
86
+ key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin,
87
+ self.is_neox_style)
88
+ key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
89
+ return query, key
90
+
91
+ @classmethod
92
+ def get_input_positions(
93
+ cls,
94
+ input_tokens: list[int],
95
+ hf_config: PretrainedConfig,
96
+ image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
97
+ video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
98
+ second_per_grid_ts: Optional[list[float]],
99
+ context_len: int = 0,
100
+ seq_len: Optional[int] = None,
101
+ audio_feature_lengths: Optional[torch.Tensor] = None,
102
+ use_audio_in_video: bool = False,
103
+ ) -> tuple[list[list[int]], int]:
104
+ """Get mrope input positions and delta value."""
105
+
106
+ image_grid_thw = [] if image_grid_thw is None else image_grid_thw
107
+ video_grid_thw = [] if video_grid_thw is None else video_grid_thw
108
+ second_per_grid_ts = [] if second_per_grid_ts is None else \
109
+ second_per_grid_ts
110
+
111
+ llm_positions, mrope_position_delta = \
112
+ cls.get_input_positions_tensor(
113
+ input_tokens=input_tokens,
114
+ hf_config=hf_config,
115
+ image_grid_thw=image_grid_thw,
116
+ video_grid_thw=video_grid_thw,
117
+ second_per_grid_ts=second_per_grid_ts,
118
+ context_len=context_len,
119
+ seq_len=seq_len,
120
+ audio_feature_lengths=audio_feature_lengths,
121
+ use_audio_in_video=use_audio_in_video,
122
+ )
123
+
124
+ return llm_positions.tolist(), mrope_position_delta
125
+
126
+ @classmethod
127
+ def get_input_positions_tensor(
128
+ cls,
129
+ input_tokens: list[int],
130
+ hf_config: PretrainedConfig,
131
+ image_grid_thw: Union[list[list[int]], torch.Tensor],
132
+ video_grid_thw: Union[list[list[int]], torch.Tensor],
133
+ second_per_grid_ts: list[float],
134
+ context_len: int = 0,
135
+ seq_len: Optional[int] = None,
136
+ audio_feature_lengths: Optional[torch.Tensor] = None,
137
+ use_audio_in_video: bool = False,
138
+ ) -> tuple[torch.Tensor, int]:
139
+ from vllm.transformers_utils.config import thinker_uses_mrope
140
+ if thinker_uses_mrope(hf_config):
141
+ return cls._omni_get_input_positions_tensor(
142
+ input_tokens=input_tokens,
143
+ hf_config=hf_config,
144
+ image_grid_thw=image_grid_thw,
145
+ video_grid_thw=video_grid_thw,
146
+ second_per_grid_ts=second_per_grid_ts,
147
+ context_len=context_len,
148
+ seq_len=seq_len,
149
+ audio_feature_lengths=audio_feature_lengths,
150
+ use_audio_in_video=use_audio_in_video,
151
+ )
152
+ elif hf_config.model_type in ["glm4v", "glm4v_moe"]:
153
+ return cls._glm4v_get_input_positions_tensor(
154
+ input_tokens=input_tokens,
155
+ hf_config=hf_config,
156
+ image_grid_thw=image_grid_thw,
157
+ video_grid_thw=video_grid_thw,
158
+ context_len=context_len,
159
+ seq_len=seq_len,
160
+ )
161
+ else:
162
+ return cls._vl_get_input_positions_tensor(
163
+ input_tokens=input_tokens,
164
+ hf_config=hf_config,
165
+ image_grid_thw=image_grid_thw,
166
+ video_grid_thw=video_grid_thw,
167
+ second_per_grid_ts=second_per_grid_ts,
168
+ context_len=context_len,
169
+ seq_len=seq_len,
170
+ )
171
+
172
+ @classmethod
173
+ def _glm4v_get_input_positions_tensor(
174
+ cls,
175
+ input_tokens: list[int],
176
+ hf_config: PretrainedConfig,
177
+ image_grid_thw: Union[list[list[int]], torch.Tensor],
178
+ video_grid_thw: Union[list[list[int]], torch.Tensor],
179
+ context_len: int = 0,
180
+ seq_len: Optional[int] = None,
181
+ ) -> tuple[torch.Tensor, int]:
182
+ """Get mrope input positions and delta value for GLM4V."""
183
+
184
+ image_token_id = hf_config.image_token_id
185
+ video_start_token_id = hf_config.video_start_token_id
186
+ video_end_token_id = hf_config.video_end_token_id
187
+ spatial_merge_size = hf_config.vision_config.spatial_merge_size
188
+ llm_pos_ids_list: list = []
189
+
190
+ if not (image_grid_thw is None and video_grid_thw is None):
191
+ if isinstance(image_grid_thw, torch.Tensor):
192
+ image_grid_thw = image_grid_thw.tolist()
193
+
194
+ input_token_type: list[str] = []
195
+ video_check_flg = False
196
+ for token in input_tokens:
197
+ if token == video_start_token_id:
198
+ video_check_flg = True
199
+ elif token == video_end_token_id:
200
+ video_check_flg = False
201
+
202
+ if (token == image_token_id) and (video_check_flg is False):
203
+ input_token_type.append("image")
204
+ elif (token == image_token_id) and (video_check_flg is True):
205
+ input_token_type.append("video")
206
+ else:
207
+ input_token_type.append("text")
208
+
209
+ input_type_group: list[tuple[str, int, int]] = []
210
+ for key, group_iter in itertools.groupby(
211
+ enumerate(input_token_type), lambda x: x[1]):
212
+ group_list = list(group_iter)
213
+ start_index = group_list[0][0]
214
+ end_index = group_list[-1][0] + 1
215
+ input_type_group.append((key, start_index, end_index))
216
+
217
+ video_frame_num = 1
218
+ mm_data_idx = 0
219
+ for modality_type, start_idx, end_idx in input_type_group:
220
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(
221
+ llm_pos_ids_list) > 0 else 0
222
+ if modality_type == "image":
223
+ t, h, w = (
224
+ image_grid_thw[mm_data_idx][0],
225
+ image_grid_thw[mm_data_idx][1],
226
+ image_grid_thw[mm_data_idx][2],
227
+ )
228
+ llm_grid_t, llm_grid_h, llm_grid_w = \
229
+ t, h // spatial_merge_size, w // spatial_merge_size
230
+
231
+ t_index = torch.arange(llm_grid_t).view(-1, 1).expand(
232
+ -1, llm_grid_h * llm_grid_w).flatten()
233
+ h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(
234
+ llm_grid_t, -1, llm_grid_w).flatten()
235
+ w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(
236
+ llm_grid_t, llm_grid_h, -1).flatten()
237
+ llm_pos_ids_list.append(
238
+ torch.stack([t_index, h_index, w_index]) + st_idx)
239
+ mm_data_idx += 1
240
+
241
+ elif modality_type == "video":
242
+ t, h, w = (
243
+ video_frame_num,
244
+ image_grid_thw[mm_data_idx][1],
245
+ image_grid_thw[mm_data_idx][2],
246
+ )
247
+ llm_grid_t, llm_grid_h, llm_grid_w = \
248
+ t, h // spatial_merge_size, w // spatial_merge_size
249
+
250
+ for t_idx in range(llm_grid_t):
251
+ t_index = torch.tensor(t_idx).view(-1, 1).expand(
252
+ -1, llm_grid_h * llm_grid_w).flatten()
253
+ h_index = torch.arange(llm_grid_h).view(
254
+ 1, -1, 1).expand(1, -1, llm_grid_w).flatten()
255
+ w_index = torch.arange(llm_grid_w).view(
256
+ 1, 1, -1).expand(1, llm_grid_h, -1).flatten()
257
+ llm_pos_ids_list.append(
258
+ torch.stack([t_index, h_index, w_index]) + st_idx)
259
+
260
+ mm_data_idx += 1
261
+ video_frame_num += 1
262
+
263
+ else:
264
+ text_len = end_idx - start_idx
265
+ llm_pos_ids_list.append(
266
+ torch.arange(text_len).view(1, -1).expand(3, -1) +
267
+ st_idx)
268
+ video_frame_num = 1
269
+
270
+ else:
271
+ text_len = len(input_tokens)
272
+ llm_pos_ids_list.append(
273
+ torch.arange(text_len).view(1, -1).expand(3, -1))
274
+
275
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
276
+ llm_positions = llm_positions[:, context_len:seq_len]
277
+ mrope_position_delta = (llm_positions.max() + 1 -
278
+ len(input_tokens)).item()
279
+ return llm_positions, mrope_position_delta
280
+
281
+ @classmethod
282
+ def _vl_get_input_positions_tensor(
283
+ cls,
284
+ input_tokens: list[int],
285
+ hf_config: PretrainedConfig,
286
+ image_grid_thw: Union[list[list[int]], torch.Tensor],
287
+ video_grid_thw: Union[list[list[int]], torch.Tensor],
288
+ second_per_grid_ts: list[float],
289
+ context_len: int = 0,
290
+ seq_len: Optional[int] = None,
291
+ ) -> tuple[torch.Tensor, int]:
292
+ """Get mrope input positions and delta value."""
293
+
294
+ image_token_id = hf_config.image_token_id
295
+ video_token_id = hf_config.video_token_id
296
+ vision_start_token_id = hf_config.vision_start_token_id
297
+ spatial_merge_size = hf_config.vision_config.spatial_merge_size
298
+ tokens_per_second = getattr(hf_config.vision_config,
299
+ "tokens_per_second", 1.0)
300
+
301
+ input_tokens_tensor = torch.tensor(input_tokens)
302
+ vision_start_indices = torch.argwhere(
303
+ input_tokens_tensor == vision_start_token_id).squeeze(1)
304
+ vision_tokens = input_tokens_tensor[vision_start_indices + 1]
305
+ image_nums = (vision_tokens == image_token_id).sum()
306
+ video_nums = (vision_tokens == video_token_id).sum()
307
+ llm_pos_ids_list: list = []
308
+
309
+ st = 0
310
+ remain_images, remain_videos = image_nums, video_nums
311
+
312
+ image_index, video_index = 0, 0
313
+ for _ in range(image_nums + video_nums):
314
+ video_second_per_grid_t = 0.0
315
+ if image_token_id in input_tokens and remain_images > 0:
316
+ ed_image = input_tokens.index(image_token_id, st)
317
+ else:
318
+ ed_image = len(input_tokens) + 1
319
+ if video_token_id in input_tokens and remain_videos > 0:
320
+ ed_video = input_tokens.index(video_token_id, st)
321
+ else:
322
+ ed_video = len(input_tokens) + 1
323
+ if ed_image < ed_video:
324
+ t, h, w = (
325
+ image_grid_thw[image_index][0],
326
+ image_grid_thw[image_index][1],
327
+ image_grid_thw[image_index][2],
328
+ )
329
+ image_index += 1
330
+ remain_images -= 1
331
+ ed = ed_image
332
+ else:
333
+ t, h, w = (
334
+ video_grid_thw[video_index][0],
335
+ video_grid_thw[video_index][1],
336
+ video_grid_thw[video_index][2],
337
+ )
338
+ video_second_per_grid_t = 1.0
339
+ if second_per_grid_ts:
340
+ video_second_per_grid_t = second_per_grid_ts[video_index]
341
+ video_index += 1
342
+ remain_videos -= 1
343
+ ed = ed_video
344
+
345
+ llm_grid_t, llm_grid_h, llm_grid_w = \
346
+ t, h // spatial_merge_size, w // spatial_merge_size
347
+ text_len = ed - st
348
+
349
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(
350
+ llm_pos_ids_list) > 0 else 0
351
+ llm_pos_ids_list.append(
352
+ torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
353
+
354
+ t_index = (torch.arange(llm_grid_t).view(-1, 1).expand(
355
+ -1, llm_grid_h * llm_grid_w) * video_second_per_grid_t *
356
+ tokens_per_second).long().flatten()
357
+
358
+ h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(
359
+ llm_grid_t, -1, llm_grid_w).flatten()
360
+ w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(
361
+ llm_grid_t, llm_grid_h, -1).flatten()
362
+ llm_pos_ids_list.append(
363
+ torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
364
+ st = ed + llm_grid_t * llm_grid_h * llm_grid_w
365
+
366
+ if st < len(input_tokens):
367
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(
368
+ llm_pos_ids_list) > 0 else 0
369
+ text_len = len(input_tokens) - st
370
+ llm_pos_ids_list.append(
371
+ torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
372
+
373
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
374
+ mrope_position_delta = (llm_positions.max() + 1 -
375
+ len(input_tokens)).item()
376
+ llm_positions = llm_positions[:, context_len:seq_len]
377
+
378
+ return llm_positions, mrope_position_delta
379
+
380
+ @classmethod
381
+ def _omni_get_input_positions_tensor(
382
+ cls,
383
+ input_tokens: list[int],
384
+ hf_config: PretrainedConfig,
385
+ image_grid_thw: Union[list[list[int]], torch.Tensor],
386
+ video_grid_thw: Union[list[list[int]], torch.Tensor],
387
+ second_per_grid_ts: Optional[list[float]] = None,
388
+ context_len: int = 0,
389
+ seq_len: Optional[int] = None,
390
+ audio_feature_lengths: Optional[torch.Tensor] = None,
391
+ use_audio_in_video: bool = False,
392
+ ) -> tuple[torch.Tensor, int]:
393
+ """Get mrope input positions and delta value (Qwen2.5-Omni version).
394
+
395
+ Differences from MRotaryEmbedding:
396
+ 1. Add audio support (and related `audio_feature_lengths`).
397
+ 2. Add `use_audio_in_video` option to read audio from video inputs.
398
+ In this case, audio and vision position ids will be split into
399
+ chunks and interleaved.
400
+
401
+ Example:
402
+
403
+ (V_i are vision position ids, A_i are audio position ids)
404
+
405
+ |V_1 ... V_n|A_1 ... A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|...
406
+ |vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |...
407
+ """
408
+
409
+ # TODO(fyabc): refactor and share more code with
410
+ # _vl_get_input_positions_tensor.
411
+
412
+ thinker_config = hf_config.thinker_config
413
+ audio_token_id = thinker_config.audio_token_index
414
+ image_token_id = thinker_config.image_token_index
415
+ video_token_id = thinker_config.video_token_index
416
+ audio_start_token_id = thinker_config.audio_start_token_id
417
+ audio_end_token_id = thinker_config.audio_end_token_id
418
+ vision_start_token_id = thinker_config.vision_start_token_id
419
+ vision_end_token_id = thinker_config.vision_end_token_id
420
+ seconds_per_chunk = thinker_config.seconds_per_chunk
421
+ spatial_merge_size = thinker_config.vision_config.spatial_merge_size
422
+ tokens_per_second = getattr(thinker_config.vision_config,
423
+ "tokens_per_second", 25)
424
+
425
+ if isinstance(image_grid_thw, list):
426
+ image_grid_thw = torch.tensor(image_grid_thw)
427
+ if isinstance(video_grid_thw, list):
428
+ video_grid_thw = torch.tensor(video_grid_thw)
429
+
430
+ src_item = input_tokens
431
+ audio_seqlens = audio_feature_lengths
432
+ if not second_per_grid_ts:
433
+ second_per_grid_ts = [1] * video_grid_thw.shape[0]
434
+ audio_idx = 0
435
+ video_idx = 0
436
+ image_idx = 0
437
+ new_src_item: list[int] = []
438
+ llm_pos_ids_list: list[torch.Tensor] = []
439
+
440
+ idx = 0
441
+ while idx < len(src_item):
442
+ new_src_item_len = len(new_src_item)
443
+ start_idx = llm_pos_ids_list[-1].max() + 1 if len(
444
+ llm_pos_ids_list) > 0 else 0
445
+ if src_item[idx] not in [
446
+ audio_token_id, video_token_id, image_token_id
447
+ ]:
448
+ if use_audio_in_video and idx > 0:
449
+ if src_item[idx] == vision_end_token_id and \
450
+ src_item[idx - 1] == audio_end_token_id:
451
+ # processing the <|audio_eos|> before <|vision_eos|>
452
+ start_idx -= 1
453
+ elif src_item[idx] == audio_start_token_id and \
454
+ src_item[idx - 1] == vision_start_token_id:
455
+ # processing the <|audio_bos|> after <|vision_eos|>
456
+ start_idx -= 1
457
+ new_src_item.append(src_item[idx])
458
+ llm_pos_ids = torch.tensor([start_idx],
459
+ dtype=torch.long).expand(3, -1)
460
+ llm_pos_ids_list.append(llm_pos_ids)
461
+ elif src_item[idx] == audio_token_id:
462
+ assert audio_seqlens is not None
463
+ audio_seqlen = audio_seqlens[audio_idx]
464
+ place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1)
465
+ new_src_item.extend([audio_token_id] * place_num)
466
+ llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx
467
+ llm_pos_ids_list.append(llm_pos_ids)
468
+ audio_idx += 1
469
+ elif src_item[idx] == image_token_id:
470
+ grid_t = image_grid_thw[image_idx][0]
471
+ grid_hs = image_grid_thw[:, 1]
472
+ grid_ws = image_grid_thw[:, 2]
473
+ t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long()
474
+ llm_pos_ids = cls._get_llm_pos_ids_for_vision(
475
+ start_idx, image_idx, spatial_merge_size, t_index, grid_hs,
476
+ grid_ws)
477
+ llm_pos_ids_list.append(llm_pos_ids)
478
+ vision_seqlen = image_grid_thw[image_idx].prod() // (
479
+ spatial_merge_size**2)
480
+ new_src_item.extend([image_token_id] * vision_seqlen)
481
+ image_idx += 1
482
+ elif src_item[idx] == video_token_id and not use_audio_in_video:
483
+ grid_t = video_grid_thw[video_idx][0]
484
+ grid_hs = video_grid_thw[:, 1]
485
+ grid_ws = video_grid_thw[:, 2]
486
+ t_index = (torch.arange(grid_t) *
487
+ second_per_grid_ts[video_idx] *
488
+ tokens_per_second).long()
489
+ llm_pos_ids = cls._get_llm_pos_ids_for_vision(
490
+ start_idx, video_idx, spatial_merge_size, t_index, grid_hs,
491
+ grid_ws)
492
+ llm_pos_ids_list.append(llm_pos_ids)
493
+ vision_seqlen = video_grid_thw[video_idx].prod() // (
494
+ spatial_merge_size**2)
495
+ new_src_item.extend([video_token_id] * vision_seqlen)
496
+ video_idx += 1
497
+ else:
498
+ # read audio from video
499
+ assert audio_seqlens is not None
500
+ audio_seqlen = audio_seqlens[audio_idx]
501
+ vision_seqlen = video_grid_thw[video_idx].prod() // (
502
+ spatial_merge_size**2)
503
+ grid_t = video_grid_thw[video_idx][0]
504
+ grid_h = video_grid_thw[video_idx][1]
505
+ grid_w = video_grid_thw[video_idx][2]
506
+ grid_hs = video_grid_thw[:, 1]
507
+ grid_ws = video_grid_thw[:, 2]
508
+ t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
509
+ t_index = (torch.arange(grid_t) *
510
+ second_per_grid_ts[video_idx] *
511
+ tokens_per_second).long()
512
+ t_index_split_chunk = cls._split_list_into_ranges(
513
+ t_index, t_ntoken_per_chunk)
514
+ place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2
515
+ pure_audio_len = place_num - 2
516
+ added_audio_len = 0
517
+ audio_llm_pos_ids_list: list[torch.Tensor] = []
518
+ for t_chunk in t_index_split_chunk:
519
+ vision_ntoken_per_chunk = len(
520
+ t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
521
+ new_src_item.extend([video_token_id] *
522
+ vision_ntoken_per_chunk)
523
+ vision_llm_pos_ids_list = cls._get_llm_pos_ids_for_vision(
524
+ start_idx, video_idx, spatial_merge_size, t_chunk,
525
+ grid_hs, grid_ws).split(1, dim=1)
526
+ llm_pos_ids_list.extend(vision_llm_pos_ids_list)
527
+ new_src_item.extend(
528
+ min(t_ntoken_per_chunk, pure_audio_len -
529
+ added_audio_len) * [audio_token_id])
530
+ audio_start_idx = start_idx if len(
531
+ audio_llm_pos_ids_list
532
+ ) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1
533
+ if min(t_ntoken_per_chunk,
534
+ pure_audio_len - added_audio_len) > 0:
535
+ audio_llm_pos_ids_list = (torch.arange(
536
+ min(t_ntoken_per_chunk, pure_audio_len -
537
+ added_audio_len)).expand(3, -1) +
538
+ audio_start_idx).split(1,
539
+ dim=1)
540
+ else:
541
+ audio_llm_pos_ids_list = []
542
+ added_audio_len += min(t_ntoken_per_chunk,
543
+ pure_audio_len - added_audio_len)
544
+ llm_pos_ids_list.extend(audio_llm_pos_ids_list)
545
+ if added_audio_len < pure_audio_len:
546
+ new_src_item.extend(
547
+ (pure_audio_len - added_audio_len) * [audio_token_id])
548
+ audio_llm_pos_ids_list = (
549
+ torch.arange(pure_audio_len - added_audio_len).expand(
550
+ 3, -1) + llm_pos_ids_list[-1].max() + 1).split(
551
+ 1, dim=1)
552
+ llm_pos_ids_list.extend(audio_llm_pos_ids_list)
553
+ audio_idx += 1
554
+ video_idx += 1
555
+ # move to the next token
556
+ idx += len(new_src_item) - new_src_item_len
557
+
558
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1)
559
+ mrope_position_delta = torch.cat(llm_pos_ids_list,
560
+ dim=1).max() + 1 - len(src_item)
561
+ llm_positions = llm_positions[:, context_len:seq_len]
562
+
563
+ return llm_positions, mrope_position_delta
564
+
565
+ @staticmethod
566
+ def _get_llm_pos_ids_for_vision(
567
+ start_idx: int,
568
+ vision_idx: int,
569
+ spatial_merge_size: int,
570
+ t_index: list[int],
571
+ grid_hs: torch.Tensor,
572
+ grid_ws: torch.Tensor,
573
+ ) -> torch.Tensor:
574
+ llm_pos_ids_list = []
575
+ llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
576
+ llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
577
+ h_index = (torch.arange(llm_grid_h).view(1, -1, 1).expand(
578
+ len(t_index), -1, llm_grid_w).flatten())
579
+ w_index = (torch.arange(llm_grid_w).view(1, 1, -1).expand(
580
+ len(t_index), llm_grid_h, -1).flatten())
581
+ t_index_tensor = torch.Tensor(t_index).to(llm_grid_h.device).view(
582
+ -1, 1).expand(-1, llm_grid_h * llm_grid_w).long().flatten()
583
+ _llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
584
+ llm_pos_ids_list.append(_llm_pos_ids + start_idx)
585
+ llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
586
+ return llm_pos_ids
587
+
588
+ @staticmethod
589
+ def _split_list_into_ranges(lst: torch.Tensor,
590
+ interval: int) -> list[list[int]]:
591
+ ranges: list[list[int]] = [[]
592
+ for _ in range((max(lst) // interval) + 1)]
593
+ for num in lst:
594
+ index = num // interval
595
+ ranges[index].append(num)
596
+ return ranges
597
+
598
+ @staticmethod
599
+ def get_next_input_positions(
600
+ mrope_position_delta: int,
601
+ context_len: int,
602
+ seq_len: int,
603
+ ) -> list[list[int]]:
604
+ return [
605
+ list(
606
+ range(context_len + mrope_position_delta,
607
+ seq_len + mrope_position_delta)) for _ in range(3)
608
+ ]
609
+
610
+ @staticmethod
611
+ def get_next_input_positions_tensor(out: np.ndarray, out_offset: int,
612
+ mrope_position_delta: int,
613
+ context_len: int, num_new_tokens: int):
614
+
615
+ values = np.arange(mrope_position_delta + context_len,
616
+ mrope_position_delta + context_len + num_new_tokens,
617
+ dtype=out.dtype)
618
+ out[:, out_offset:out_offset + num_new_tokens] = values
619
+
620
+ @classmethod
621
+ def omni_get_updates_use_audio_in_video(
622
+ cls,
623
+ thinker_config: PretrainedConfig,
624
+ audio_len: int,
625
+ video_grid_thw: Union[list[int], torch.Tensor],
626
+ video_second_per_grid_t: float,
627
+ ) -> list[int]:
628
+ """Get video prompt updates when `use_audio_in_video` is True.
629
+
630
+ In this case, audio and vision update ids will be split into
631
+ chunks and interleaved (details in `_omni_get_input_positions_tensor`).
632
+
633
+ <|video_bos|><|VIDEO|><|video_eos|> =>
634
+ <|video_bos|><|audio_bos|>(... chunks ...)<|audio_eos|><|video_eos|>
635
+ """
636
+
637
+ audio_token_id = thinker_config.audio_token_index
638
+ video_token_id = thinker_config.video_token_index
639
+ audio_start_token_id = thinker_config.audio_start_token_id
640
+ audio_end_token_id = thinker_config.audio_end_token_id
641
+ seconds_per_chunk = thinker_config.seconds_per_chunk
642
+ spatial_merge_size = thinker_config.vision_config.spatial_merge_size
643
+ tokens_per_second = getattr(thinker_config.vision_config,
644
+ "tokens_per_second", 25)
645
+
646
+ grid_t = video_grid_thw[0]
647
+ grid_h = video_grid_thw[1]
648
+ grid_w = video_grid_thw[2]
649
+ t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
650
+ t_index = (torch.arange(grid_t) * video_second_per_grid_t *
651
+ tokens_per_second).long()
652
+ t_index_split_chunk = cls._split_list_into_ranges(
653
+ t_index, t_ntoken_per_chunk)
654
+
655
+ updates = [audio_start_token_id]
656
+ added_audio_len = 0
657
+ for t_chunk in t_index_split_chunk:
658
+ vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (
659
+ spatial_merge_size**2)
660
+ updates.extend([video_token_id] * vision_ntoken_per_chunk)
661
+
662
+ audio_chunk_size = min(t_ntoken_per_chunk,
663
+ audio_len - added_audio_len)
664
+ updates.extend(audio_chunk_size * [audio_token_id])
665
+ added_audio_len += audio_chunk_size
666
+ if added_audio_len < audio_len:
667
+ updates.extend((audio_len - added_audio_len) * [audio_token_id])
668
+ updates.extend([audio_end_token_id])
669
+
670
+ return updates
vllm/model_executor/layers/rotary_embedding/ntk_scaling_rope.py
new
+42
@@ -0,0 +1,42 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Optional
5
+
6
+import torch
7
+
8
+from .base import RotaryEmbedding
9
+
10
+
11
+class NTKScalingRotaryEmbedding(RotaryEmbedding):
12
+ """RotaryEmbedding extended with fixed and mixed NTK scaling.
13
+ https://kexue.fm/archives/9706 """
14
+
15
+ def __init__(self,
16
+ head_size: int,
17
+ rotary_dim: int,
18
+ max_position_embeddings: int,
19
+ base: float,
20
+ is_neox_style: bool,
21
+ scaling_factor: float,
22
+ dtype: torch.dtype,
23
+ mixed_b: Optional[float] = None) -> None:
24
+ self.scaling_factor = scaling_factor
25
+ self.mixed_b = mixed_b
26
+ super().__init__(head_size, rotary_dim, max_position_embeddings, base,
27
+ is_neox_style, dtype)
28
+
29
+ def _compute_inv_freq(self, base: float) -> torch.Tensor:
30
+ base = self.base * (self.scaling_factor if self.mixed_b is None else 1)
31
+ inv_freq = super()._compute_inv_freq(base)
32
+
33
+ if self.mixed_b is None:
34
+ inv_freq = inv_freq / self.scaling_factor**(2 / self.rotary_dim)
35
+ else:
36
+ a = torch.tensor(self.scaling_factor).log() / (self.rotary_dim /
37
+ 2)**self.mixed_b
38
+ lambda_1_m = (a * torch.arange(
39
+ 1, self.rotary_dim // 2 + 1).float()**self.mixed_b).exp()
40
+ inv_freq = inv_freq / lambda_1_m
41
+
42
+ return inv_freq
vllm/model_executor/layers/rotary_embedding/phi3_long_rope_scaled_rope.py
new
+129
@@ -0,0 +1,129 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import math
4
+from typing import Optional
5
+
6
+import torch
7
+import torch.nn as nn
8
+
9
+from .common import rotate_neox
10
+
11
+
12
+class Phi3LongRoPEScaledRotaryEmbedding(nn.Module):
13
+ """Phi3 family of models scaled rotary embedding.
14
+
15
+ Based on the original RotaryEmbedding implementation.
16
+ """
17
+
18
+ def __init__(
19
+ self,
20
+ head_size: int,
21
+ rotary_dim: int,
22
+ max_position_embeddings: int,
23
+ original_max_position_embeddings: int,
24
+ base: float,
25
+ is_neox_style: bool,
26
+ dtype: torch.dtype,
27
+ short_factor: list[float],
28
+ long_factor: list[float],
29
+ short_mscale: Optional[float] = None,
30
+ long_mscale: Optional[float] = None,
31
+ ):
32
+ super().__init__()
33
+
34
+ if is_neox_style is False:
35
+ raise ValueError(
36
+ "`Phi3LongRoPEScaledRotaryEmbedding` only supports neox_style."
37
+ )
38
+
39
+ self.rotary_dim = rotary_dim
40
+ self.head_size = head_size
41
+ self.max_position_embeddings = max_position_embeddings
42
+ self.original_max_position_embeddings = original_max_position_embeddings
43
+ self.base = base
44
+ self.short_factor = short_factor
45
+ self.long_factor = long_factor
46
+
47
+ scale = self.max_position_embeddings / \
48
+ self.original_max_position_embeddings
49
+ if scale <= 1.0:
50
+ scaling_factor = 1.0
51
+ else:
52
+ scaling_factor = math.sqrt(
53
+ 1 + math.log(scale) /
54
+ math.log(self.original_max_position_embeddings))
55
+ if short_mscale is None:
56
+ short_mscale = scaling_factor
57
+ if long_mscale is None:
58
+ long_mscale = scaling_factor
59
+
60
+ self.short_mscale = short_mscale
61
+ self.long_mscale = long_mscale
62
+
63
+ short_cache = self._compute_cos_sin_cache(
64
+ original_max_position_embeddings, short_factor, short_mscale)
65
+ short_cache = short_cache.to(dtype)
66
+
67
+ long_cache = self._compute_cos_sin_cache(max_position_embeddings,
68
+ long_factor, long_mscale)
69
+ long_cache = long_cache.to(dtype)
70
+
71
+ long_short_cache = torch.cat([short_cache, long_cache], dim=0)
72
+ self.register_buffer("long_short_cos_sin_cache",
73
+ long_short_cache,
74
+ persistent=False)
75
+
76
+ def _compute_inv_freq(self, rescale_factors: list[float]) -> torch.Tensor:
77
+ rescale_factors = torch.tensor(rescale_factors, dtype=torch.float32)
78
+ inv_freq = 1.0 / (rescale_factors * (self.base**(torch.arange(
79
+ 0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim)))
80
+ return inv_freq
81
+
82
+ def _compute_cos_sin_cache(
83
+ self,
84
+ max_position_embeddings: int,
85
+ rescale_factors: list[float],
86
+ mscale: float,
87
+ ) -> torch.Tensor:
88
+ inv_freq = self._compute_inv_freq(rescale_factors)
89
+ t = torch.arange(max_position_embeddings, dtype=torch.float)
90
+ freqs = torch.einsum("i,j -> ij", t, inv_freq)
91
+ cos = freqs.cos() * mscale
92
+ sin = freqs.sin() * mscale
93
+ cache = torch.cat((cos, sin), dim=-1)
94
+ return cache
95
+
96
+ def forward(
97
+ self,
98
+ positions: torch.Tensor,
99
+ query: torch.Tensor,
100
+ key: Optional[torch.Tensor] = None,
101
+ offsets: Optional[torch.Tensor] = None,
102
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
103
+ assert key is not None
104
+ query = query.view(*query.shape[:-1], -1, self.head_size)
105
+ key = key.view(*key.shape[:-1], -1, self.head_size)
106
+
107
+ k = self.original_max_position_embeddings
108
+ long_prompt_offset = (torch.any(positions > k).float() *
109
+ torch.full_like(positions, k)).long()
110
+ idx = (torch.add(positions, long_prompt_offset)
111
+ if long_prompt_offset is not None else positions)
112
+ idx = torch.add(idx, offsets) if offsets is not None else idx
113
+ cos_sin = torch.index_select(self.long_short_cos_sin_cache, 0, idx)
114
+
115
+ cos, sin = cos_sin.chunk(2, dim=-1)
116
+ cos = cos.repeat(1, 2).unsqueeze(-2)
117
+ sin = sin.repeat(1, 2).unsqueeze(-2)
118
+
119
+ query_rot = query[..., :self.rotary_dim]
120
+ query_pass = query[..., self.rotary_dim:]
121
+ query_rot = query_rot * cos + rotate_neox(query_rot) * sin
122
+ query = torch.cat((query_rot, query_pass), dim=-1)
123
+
124
+ key_rot = key[..., :self.rotary_dim]
125
+ key_pass = key[..., self.rotary_dim:]
126
+ key_rot = key_rot * cos + rotate_neox(key_rot) * sin
127
+ key = torch.cat((key_rot, key_pass), dim=-1)
128
+
129
+ return query.flatten(-2), key.flatten(-2)
vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py
new
+68
@@ -0,0 +1,68 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from .base import RotaryEmbedding
7
+from .common import (yarn_find_correction_range, yarn_get_mscale,
8
+ yarn_linear_ramp_mask)
9
+
10
+
11
+class YaRNScalingRotaryEmbedding(RotaryEmbedding):
12
+ """RotaryEmbedding extended with YaRN method.
13
+
14
+ Credits to Peng et al. github.com/jquesnelle/yarn
15
+ """
16
+
17
+ def __init__(
18
+ self,
19
+ head_size: int,
20
+ rotary_dim: int,
21
+ max_position_embeddings: int,
22
+ base: float,
23
+ is_neox_style: bool,
24
+ scaling_factor: float,
25
+ dtype: torch.dtype,
26
+ *,
27
+ extrapolation_factor: float = 1,
28
+ attn_factor: float = 1,
29
+ beta_fast: int = 32,
30
+ beta_slow: int = 1,
31
+ ) -> None:
32
+ self.scaling_factor = scaling_factor
33
+ self.extrapolation_factor = extrapolation_factor
34
+ self.attn_factor = attn_factor
35
+ self.beta_fast = beta_fast
36
+ self.beta_slow = beta_slow
37
+ # Get n-d magnitude scaling corrected for interpolation
38
+ self.mscale = float(yarn_get_mscale(self.scaling_factor) * attn_factor)
39
+ super().__init__(head_size, rotary_dim, max_position_embeddings, base,
40
+ is_neox_style, dtype)
41
+
42
+ def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
43
+ pos_freqs = self.base**(
44
+ torch.arange(0, self.rotary_dim, 2, dtype=torch.float) /
45
+ self.rotary_dim)
46
+ inv_freq_extrapolation = 1.0 / pos_freqs
47
+ inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
48
+
49
+ low, high = yarn_find_correction_range(self.beta_fast, self.beta_slow,
50
+ self.rotary_dim, self.base,
51
+ self.max_position_embeddings)
52
+ # Get n-d rotational scaling corrected for extrapolation
53
+ inv_freq_mask = (1 - yarn_linear_ramp_mask(
54
+ low, high, self.rotary_dim // 2,
55
+ dtype=torch.float)) * self.extrapolation_factor
56
+ inv_freq = inv_freq_interpolation * (
57
+ 1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask
58
+ return inv_freq
59
+
60
+ def _compute_cos_sin_cache(self) -> torch.Tensor:
61
+ inv_freq = self._compute_inv_freq(self.scaling_factor)
62
+ t = torch.arange(self.max_position_embeddings * self.scaling_factor,
63
+ dtype=torch.float32)
64
+ freqs = torch.einsum("i,j -> ij", t, inv_freq)
65
+ cos = (freqs.cos() * self.mscale)
66
+ sin = (freqs.sin() * self.mscale)
67
+ cache = torch.cat((cos, sin), dim=-1)
68
+ return cache