wehub-resource-sync
59a0a3844c
PR Test AMD / cancel-on-close (push) Has been skipped
PR Test NVIDIA ARM / scan (push) Has been skipped
PR Test NVIDIA / cancel-on-close (push) Has been skipped
PR Test AMD / scan (push) Has been skipped
PR Test NVIDIA ARM / cancel-on-close (push) Has been skipped
PR Test NVIDIA / scan (push) Has been skipped
Release Docker Images / build (cu129-torch-2.11.0) (push) Has been skipped
Release Docker Images / build (cu130-torch-2.11.0) (push) Has been skipped
Release PyPI / publish (push) Has been skipped
Scheduler Python Test / test (push) Successful in 27m19s
Docs / build (push) Successful in 28m8s
Scheduler C++ Test / test (push) Successful in 28m19s
Scheduler C++ Test / test-flat (push) Successful in 28m18s
Docs / deploy (push) Has been cancelled
PR Test AMD / finish (push) Has been cancelled
PR Test NVIDIA / finish (push) Has been cancelled
PR Test NVIDIA ARM / finish (push) Has been cancelled
PR Test NVIDIA ARM / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test AMD / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test NVIDIA / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
1469 lines
53 KiB
Python
Executable File
1469 lines
53 KiB
Python
Executable File
# Copyright (c) 2026 LightSeek Foundation
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
# copies of the Software, and to permit persons to whom the Software is
|
|
# furnished to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in
|
|
# all copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
# SOFTWARE.
|
|
|
|
|
|
"""Rotary Positional Embeddings."""
|
|
|
|
import itertools
|
|
import logging
|
|
import math
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from tokenspeed_kernel.ops.embedding import FusedSetKVBufferArg, apply_rope
|
|
from tokenspeed_kernel.platform import current_platform
|
|
from tokenspeed_kernel.torch_compile import get_compiler_backend
|
|
|
|
_is_nvidia = current_platform().is_nvidia
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
|
|
x1 = x[..., : x.shape[-1] // 2]
|
|
x2 = x[..., x.shape[-1] // 2 :]
|
|
return torch.cat((-x2, x1), dim=-1)
|
|
|
|
|
|
def _rotate_gptj(x: torch.Tensor) -> torch.Tensor:
|
|
x1 = x[..., ::2]
|
|
x2 = x[..., 1::2]
|
|
x = torch.stack((-x2, x1), dim=-1)
|
|
return x.flatten(-2)
|
|
|
|
|
|
def _apply_rotary_emb(
|
|
x: torch.Tensor,
|
|
cos: torch.Tensor,
|
|
sin: torch.Tensor,
|
|
is_neox_style: bool,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
x: [num_tokens, num_heads, head_size]
|
|
cos: [num_tokens, head_size // 2]
|
|
sin: [num_tokens, head_size // 2]
|
|
is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
|
|
positional embeddings.
|
|
"""
|
|
cos = cos.unsqueeze(-2).to(x.dtype)
|
|
sin = sin.unsqueeze(-2).to(x.dtype)
|
|
if is_neox_style:
|
|
x1, x2 = torch.chunk(x, 2, dim=-1)
|
|
else:
|
|
x1 = x[..., ::2]
|
|
x2 = x[..., 1::2]
|
|
o1 = x1 * cos - x2 * sin
|
|
o2 = x2 * cos + x1 * sin
|
|
if is_neox_style:
|
|
return torch.cat((o1, o2), dim=-1)
|
|
else:
|
|
return torch.stack((o1, o2), dim=-1).flatten(-2)
|
|
|
|
|
|
# Copied from transformers
|
|
def rotate_half(x):
|
|
"""Rotates half the hidden dims of the input."""
|
|
x1 = x[..., : x.shape[-1] // 2]
|
|
x2 = x[..., x.shape[-1] // 2 :]
|
|
return torch.cat((-x2, x1), dim=-1)
|
|
|
|
|
|
@torch.compile(dynamic=True, backend=get_compiler_backend())
|
|
def apply_rotary_pos_emb_native(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
cos: torch.Tensor,
|
|
sin: torch.Tensor,
|
|
unsqueeze_dim=1,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
orig_q_dtype = q.dtype
|
|
orig_k_dtype = k.dtype
|
|
q, k = q.float(), k.float()
|
|
|
|
# embedding is performed in float
|
|
cos = cos.unsqueeze(unsqueeze_dim).float()
|
|
sin = sin.unsqueeze(unsqueeze_dim).float()
|
|
q_embed = (q * cos) + (rotate_half(q) * sin)
|
|
k_embed = (k * cos) + (rotate_half(k) * sin)
|
|
|
|
q_embed = q_embed.to(orig_q_dtype)
|
|
k_embed = k_embed.to(orig_k_dtype)
|
|
|
|
return q_embed, k_embed
|
|
|
|
|
|
def apply_interleaved_rope(x: torch.Tensor, mrope_section: list) -> torch.Tensor:
|
|
x_t = x[0].clone()
|
|
x_t[..., 1 : mrope_section[1] * 3 : 3] = x[1, ..., 1 : mrope_section[1] * 3 : 3]
|
|
x_t[..., 2 : mrope_section[2] * 3 : 3] = x[2, ..., 2 : mrope_section[2] * 3 : 3]
|
|
return x_t
|
|
|
|
|
|
class RotaryEmbedding(torch.nn.Module):
|
|
"""Original rotary positional embedding."""
|
|
|
|
def __init__(
|
|
self,
|
|
head_size: int,
|
|
rotary_dim: int,
|
|
max_position_embeddings: int,
|
|
base: int,
|
|
is_neox_style: bool,
|
|
dtype: torch.dtype,
|
|
) -> None:
|
|
super().__init__()
|
|
self.head_size = head_size
|
|
self.rotary_dim = rotary_dim
|
|
self.max_position_embeddings = max_position_embeddings
|
|
self.base = base
|
|
self.is_neox_style = is_neox_style
|
|
self.dtype = dtype
|
|
|
|
cache = self._compute_cos_sin_cache()
|
|
self.cos_sin_cache: torch.Tensor
|
|
self.register_buffer("cos_sin_cache", cache, persistent=False)
|
|
|
|
def _compute_inv_freq(self, base: int | float) -> torch.Tensor:
|
|
"""Compute the inverse frequency."""
|
|
# To exactly match the HF implementation, we need to
|
|
# use CPU to compute the cache and then move it to GPU. However, we
|
|
# create the cache on GPU for faster initialization. This may cause
|
|
# a slight numerical difference between the HF implementation and ours.
|
|
inv_freq = 1.0 / (
|
|
base
|
|
** (
|
|
torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim
|
|
)
|
|
)
|
|
return inv_freq
|
|
|
|
def _compute_cos_sin_cache(self) -> torch.Tensor:
|
|
"""Compute the cos and sin cache."""
|
|
inv_freq = self._compute_inv_freq(self.base)
|
|
t = torch.arange(self.max_position_embeddings, dtype=torch.float)
|
|
|
|
freqs = torch.einsum("i,j -> ij", t, inv_freq)
|
|
cos = freqs.cos()
|
|
sin = freqs.sin()
|
|
cache = torch.cat((cos, sin), dim=-1)
|
|
return cache
|
|
|
|
def forward(
|
|
self,
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
offsets: torch.Tensor | None = None,
|
|
fused_set_kv_buffer_arg: FusedSetKVBufferArg | None = None,
|
|
output_q_rope: torch.Tensor | None = None,
|
|
output_k_rope: torch.Tensor | None = None,
|
|
enable_pdl: bool = False,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if offsets is not None:
|
|
raise ValueError("embedding.rope does not support offsets")
|
|
return apply_rope(
|
|
positions=positions,
|
|
q=query,
|
|
k=key,
|
|
head_size=self.head_size,
|
|
cos_sin_cache=self.cos_sin_cache,
|
|
is_neox=self.is_neox_style,
|
|
fused_set_kv_buffer_arg=fused_set_kv_buffer_arg,
|
|
q_rope_out=output_q_rope,
|
|
k_rope_out=output_k_rope,
|
|
enable_pdl=enable_pdl,
|
|
)
|
|
|
|
def extra_repr(self) -> str:
|
|
s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
|
|
s += f", max_position_embeddings={self.max_position_embeddings}"
|
|
s += f", base={self.base}, is_neox_style={self.is_neox_style}"
|
|
return s
|
|
|
|
|
|
class LinearScalingRotaryEmbedding(RotaryEmbedding):
|
|
"""RotaryEmbedding extended with linear scaling.
|
|
|
|
It supports multiple scaling factors. Since multiple LoRA adapters may have
|
|
different scaling factors, we need multiple cos/sin caches. In this way,
|
|
instead of running rotary embedding kernel per lora, we can run multiple
|
|
lora in a batched way.
|
|
|
|
In addition to that, we also keep the cos/sin cache for the scaling factor
|
|
of 1 (default) at all times.
|
|
|
|
Exemplary for two scaling factors x=1, y and z with embeddings
|
|
[[x11, x12, ... x1m], ..., [xn1, xn2, ..., xnm]] and
|
|
[[y11, y12, ... y1o], ..., [yn1, yn2, ..., yno]], and
|
|
[[z11, z12, ... z1p], ..., [zn1, zn2, ..., znp]],
|
|
|
|
we construct the cos/sin cache as follows:
|
|
[[x11, x12, ... x1m, y11, y12, ... y1o, z11, z12, ... z1p],
|
|
...
|
|
[xn1, xn2, ... xnm, yn1, yn2, ... yno, zn1, zn2, ... znp]]
|
|
|
|
We then use offsets to index into the cos/sin cache for
|
|
the respective scaling factors.
|
|
|
|
The offset to cache can be accessed via `scaling_factor_to_offset` API.
|
|
|
|
Credits to the Reddit user /u/kaiokendev
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
head_size: int,
|
|
rotary_dim: int,
|
|
max_position_embeddings: int,
|
|
base: int,
|
|
is_neox_style: bool,
|
|
scaling_factors: list[float] | float,
|
|
dtype: torch.dtype,
|
|
) -> None:
|
|
if isinstance(scaling_factors, float):
|
|
scaling_factors = [scaling_factors]
|
|
self.scaling_factors: list[float] = scaling_factors
|
|
super().__init__(
|
|
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
|
|
)
|
|
# Lazy initialized.
|
|
self._scaling_factor_to_offset: dict[float, int]
|
|
|
|
def _compute_cos_sin_cache(self) -> torch.Tensor:
|
|
inv_freq = self._compute_inv_freq(self.base)
|
|
cache_list: list[torch.Tensor] = []
|
|
# offsets to the next cache in a tensor.
|
|
# Each offset corresponds to the same index in scaling_factors.
|
|
offsets: list[int] = []
|
|
for scaling_factor in self.scaling_factors:
|
|
# self.max_position_embeddings is the original
|
|
# maximum length before applying the rope scaling.
|
|
# Thus, the maximum length after applying the rope scaling is
|
|
# self.max_position_embeddings * self.scaling_factor.
|
|
max_len = self.max_position_embeddings * scaling_factor
|
|
t = torch.arange(max_len, dtype=torch.float)
|
|
t = t / scaling_factor
|
|
|
|
freqs = torch.einsum("i,j -> ij", t, inv_freq)
|
|
cos = freqs.cos()
|
|
sin = freqs.sin()
|
|
cache = torch.cat((cos, sin), dim=-1)
|
|
if not cache_list:
|
|
offset = 0
|
|
else:
|
|
last_offset = offsets[-1]
|
|
next_max_len = cache_list[-1].shape[0]
|
|
offset = last_offset + next_max_len
|
|
offsets.append(offset)
|
|
cache_list.append(cache)
|
|
self._scaling_factor_to_offset = {
|
|
float(scaling_factor): offsets[i]
|
|
for i, scaling_factor in enumerate(self.scaling_factors)
|
|
}
|
|
if len(self.scaling_factors) != len(offsets):
|
|
raise RuntimeError("scaling factor offsets were not initialized correctly.")
|
|
return torch.cat(cache_list, dim=0)
|
|
|
|
@property
|
|
def scaling_factor_to_offset(self) -> dict[float, int]:
|
|
return self._scaling_factor_to_offset
|
|
|
|
|
|
class DynamicNTKScalingRotaryEmbedding(RotaryEmbedding):
|
|
"""RotaryEmbedding extended with Dynamic NTK scaling.
|
|
|
|
Credits to the Reddit users /u/bloc97 and /u/emozilla
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
head_size: int,
|
|
rotary_dim: int,
|
|
max_position_embeddings: int,
|
|
base: int,
|
|
is_neox_style: bool,
|
|
scaling_factor: float,
|
|
dtype: torch.dtype,
|
|
) -> None:
|
|
self.scaling_factor = scaling_factor
|
|
super().__init__(
|
|
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
|
|
)
|
|
|
|
def _compute_cos_sin_cache(self) -> torch.Tensor:
|
|
# self.max_position_embeddings is the original
|
|
# maximum length before applying the rope scaling.
|
|
# Thus, the maximum length after applying the rope scaling is
|
|
# self.max_position_embeddings * self.scaling_factor.
|
|
max_len = self.max_position_embeddings * self.scaling_factor
|
|
base = self.base * (
|
|
(self.scaling_factor * max_len / self.max_position_embeddings)
|
|
- (self.scaling_factor - 1)
|
|
) ** (self.rotary_dim / (self.rotary_dim - 2))
|
|
inv_freq = self._compute_inv_freq(base)
|
|
t = torch.arange(max_len, dtype=torch.float)
|
|
|
|
freqs = torch.einsum("i,j -> ij", t, inv_freq)
|
|
cos = freqs.cos()
|
|
sin = freqs.sin()
|
|
cache = torch.cat((cos, sin), dim=-1)
|
|
return cache
|
|
|
|
|
|
# Inverse dim formula to find dim based on number of rotations
|
|
def _yarn_find_correction_dim(
|
|
num_rotations: int,
|
|
dim: int,
|
|
base: float = 10000,
|
|
max_position_embeddings: int = 2048,
|
|
) -> float:
|
|
return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (
|
|
2 * math.log(base)
|
|
)
|
|
|
|
|
|
# Find dim range bounds based on rotations
|
|
def _yarn_find_correction_range(
|
|
low_rot: int,
|
|
high_rot: int,
|
|
dim: int,
|
|
base: float = 10000,
|
|
max_position_embeddings: int = 2048,
|
|
) -> tuple[int, int]:
|
|
low = math.floor(
|
|
_yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
|
|
)
|
|
high = math.ceil(
|
|
_yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
|
|
)
|
|
return max(low, 0), min(high, dim - 1) # Clamp values just in case
|
|
|
|
|
|
def _yarn_linear_ramp_mask(
|
|
low: float, high: float, dim: int, dtype: torch.dtype, device: torch.device = None
|
|
) -> torch.Tensor:
|
|
if low == high:
|
|
high += 0.001 # Prevent singularity
|
|
|
|
linear_func = (torch.arange(dim, dtype=dtype, device=device) - low) / (high - low)
|
|
ramp_func = torch.clamp(linear_func, 0, 1)
|
|
return ramp_func
|
|
|
|
|
|
def _yarn_get_mscale(scale: float = 1) -> float:
|
|
if scale <= 1:
|
|
return 1.0
|
|
return 0.1 * math.log(scale) + 1.0
|
|
|
|
|
|
class YaRNScalingRotaryEmbedding(RotaryEmbedding):
|
|
"""RotaryEmbedding extended with YaRN method.
|
|
|
|
Credits to Peng et al. github.com/jquesnelle/yarn
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
head_size: int,
|
|
rotary_dim: int,
|
|
max_position_embeddings: int,
|
|
base: int,
|
|
is_neox_style: bool,
|
|
scaling_factor: float,
|
|
dtype: torch.dtype,
|
|
*,
|
|
extrapolation_factor: float = 1,
|
|
attn_factor: float = 1,
|
|
beta_fast: int = 32,
|
|
beta_slow: int = 1,
|
|
) -> None:
|
|
self.scaling_factor = scaling_factor
|
|
self.extrapolation_factor = extrapolation_factor
|
|
self.attn_factor = attn_factor
|
|
self.beta_fast = beta_fast
|
|
self.beta_slow = beta_slow
|
|
# Get n-d magnitude scaling corrected for interpolation
|
|
self.mscale = float(_yarn_get_mscale(self.scaling_factor) * attn_factor)
|
|
super().__init__(
|
|
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
|
|
)
|
|
|
|
def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
|
|
pos_freqs = self.base ** (
|
|
torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim
|
|
)
|
|
inv_freq_extrapolation = 1.0 / pos_freqs
|
|
inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
|
|
|
|
low, high = _yarn_find_correction_range(
|
|
self.beta_fast,
|
|
self.beta_slow,
|
|
self.rotary_dim,
|
|
self.base,
|
|
self.max_position_embeddings,
|
|
)
|
|
# Get n-d rotational scaling corrected for extrapolation
|
|
inv_freq_mask = (
|
|
1
|
|
- _yarn_linear_ramp_mask(low, high, self.rotary_dim // 2, dtype=torch.float)
|
|
) * self.extrapolation_factor
|
|
inv_freq = (
|
|
inv_freq_interpolation * (1 - inv_freq_mask)
|
|
+ inv_freq_extrapolation * inv_freq_mask
|
|
)
|
|
return inv_freq
|
|
|
|
def _compute_cos_sin_cache(self) -> torch.Tensor:
|
|
inv_freq = self._compute_inv_freq(self.scaling_factor)
|
|
t = torch.arange(
|
|
self.max_position_embeddings * self.scaling_factor, dtype=torch.float32
|
|
)
|
|
freqs = torch.einsum("i,j -> ij", t, inv_freq)
|
|
cos = freqs.cos() * self.mscale
|
|
sin = freqs.sin() * self.mscale
|
|
cache = torch.cat((cos, sin), dim=-1)
|
|
return cache
|
|
|
|
|
|
class Phi3LongRoPEScaledRotaryEmbedding(nn.Module):
|
|
"""Phi3 family of models scaled rotary embedding."""
|
|
|
|
def __init__(
|
|
self,
|
|
head_size: int,
|
|
rotary_dim: int,
|
|
max_position_embeddings: int,
|
|
original_max_position_embeddings: int,
|
|
base: int,
|
|
is_neox_style: bool,
|
|
dtype: torch.dtype,
|
|
short_factor: list[float],
|
|
long_factor: list[float],
|
|
short_mscale: float | None = None,
|
|
long_mscale: float | None = None,
|
|
):
|
|
super().__init__()
|
|
|
|
if is_neox_style is False:
|
|
raise ValueError(
|
|
"`Phi3LongRoPEScaledRotaryEmbedding` only supports neox_style."
|
|
)
|
|
|
|
self.rotary_dim = rotary_dim
|
|
self.head_size = head_size
|
|
self.max_position_embeddings = max_position_embeddings
|
|
self.original_max_position_embeddings = original_max_position_embeddings
|
|
self.base = base
|
|
self.short_factor = short_factor
|
|
self.long_factor = long_factor
|
|
|
|
scale = self.max_position_embeddings / self.original_max_position_embeddings
|
|
if scale <= 1.0:
|
|
scaling_factor = 1.0
|
|
else:
|
|
scaling_factor = math.sqrt(
|
|
1 + math.log(scale) / math.log(self.original_max_position_embeddings)
|
|
)
|
|
if short_mscale is None:
|
|
short_mscale = scaling_factor
|
|
if long_mscale is None:
|
|
long_mscale = scaling_factor
|
|
|
|
self.short_mscale = short_mscale
|
|
self.long_mscale = long_mscale
|
|
|
|
short_cache = self._compute_cos_sin_cache(
|
|
original_max_position_embeddings, short_factor, short_mscale
|
|
)
|
|
short_cache = short_cache.to(dtype)
|
|
self.register_buffer("short_cos_sin_cache", short_cache, persistent=False)
|
|
|
|
long_cache = self._compute_cos_sin_cache(
|
|
max_position_embeddings, long_factor, long_mscale
|
|
)
|
|
long_cache = long_cache.to(dtype)
|
|
self.register_buffer("long_cos_sin_cache", long_cache, persistent=False)
|
|
|
|
long_short_cache = torch.cat(
|
|
[self.short_cos_sin_cache, self.long_cos_sin_cache], dim=0
|
|
)
|
|
self.register_buffer(
|
|
"long_short_cos_sin_cache", long_short_cache, persistent=False
|
|
)
|
|
|
|
def _compute_inv_freq(self, rescale_factors: list[float]) -> torch.Tensor:
|
|
rescale_factors = torch.tensor(rescale_factors, dtype=torch.float32)
|
|
inv_freq = 1.0 / (
|
|
rescale_factors
|
|
* (
|
|
self.base
|
|
** (
|
|
torch.arange(0, self.rotary_dim, 2, dtype=torch.float)
|
|
/ self.rotary_dim
|
|
)
|
|
)
|
|
)
|
|
return inv_freq
|
|
|
|
def _compute_cos_sin_cache(
|
|
self,
|
|
max_position_embeddings: int,
|
|
rescale_factors: list[float],
|
|
mscale: float,
|
|
) -> torch.Tensor:
|
|
inv_freq = self._compute_inv_freq(rescale_factors)
|
|
t = torch.arange(max_position_embeddings, dtype=torch.float)
|
|
freqs = torch.einsum("i,j -> ij", t, inv_freq)
|
|
cos = freqs.cos() * mscale
|
|
sin = freqs.sin() * mscale
|
|
cache = torch.cat((cos, sin), dim=-1)
|
|
return cache
|
|
|
|
def forward(
|
|
self,
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
offsets: torch.Tensor | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
query = query.view(*query.shape[:-1], -1, self.head_size)
|
|
key = key.view(*key.shape[:-1], -1, self.head_size)
|
|
|
|
k = self.original_max_position_embeddings
|
|
long_prompt_offset = (
|
|
torch.any(positions > k).float() * torch.full_like(positions, k)
|
|
).long()
|
|
idx = (
|
|
torch.add(positions, long_prompt_offset)
|
|
if long_prompt_offset is not None
|
|
else positions
|
|
)
|
|
self.long_short_cos_sin_cache: torch.Tensor = self.long_short_cos_sin_cache.to(
|
|
idx.device
|
|
)
|
|
idx = torch.add(idx, offsets) if offsets is not None else idx
|
|
cos_sin = torch.index_select(self.long_short_cos_sin_cache, 0, idx)
|
|
|
|
cos, sin = cos_sin.chunk(2, dim=-1)
|
|
cos = cos.repeat(1, 2).unsqueeze(-2)
|
|
sin = sin.repeat(1, 2).unsqueeze(-2)
|
|
|
|
query_rot = query[..., : self.rotary_dim]
|
|
query_pass = query[..., self.rotary_dim :]
|
|
query_rot = query_rot * cos + _rotate_neox(query_rot) * sin
|
|
query = torch.cat((query_rot, query_pass), dim=-1)
|
|
|
|
key_rot = key[..., : self.rotary_dim]
|
|
key_pass = key[..., self.rotary_dim :]
|
|
key_rot = key_rot * cos + _rotate_neox(key_rot) * sin
|
|
key = torch.cat((key_rot, key_pass), dim=-1)
|
|
|
|
return query.flatten(-2), key.flatten(-2)
|
|
|
|
|
|
def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
|
|
if scale <= 1:
|
|
return 1.0
|
|
return 0.1 * mscale * math.log(scale) + 1.0
|
|
|
|
|
|
class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
|
|
"""RotaryEmbedding extended with YaRN method.
|
|
|
|
Credits to Peng et al. github.com/jquesnelle/yarn
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
head_size: int,
|
|
rotary_dim: int,
|
|
max_position_embeddings: int,
|
|
base: int,
|
|
is_neox_style: bool,
|
|
scaling_factor: float,
|
|
dtype: torch.dtype,
|
|
*,
|
|
extrapolation_factor: float = 1,
|
|
attn_factor: float = 1,
|
|
beta_fast: int = 32,
|
|
beta_slow: int = 1,
|
|
mscale: float = 1,
|
|
mscale_all_dim: float = 0,
|
|
device: str | None = "cuda",
|
|
) -> None:
|
|
self.scaling_factor = scaling_factor
|
|
self.extrapolation_factor = extrapolation_factor
|
|
self.attn_factor = attn_factor
|
|
self.beta_fast = beta_fast
|
|
self.beta_slow = beta_slow
|
|
# Get n-d magnitude scaling corrected for interpolation.
|
|
self.mscale = float(
|
|
yarn_get_mscale(self.scaling_factor, float(mscale))
|
|
/ yarn_get_mscale(self.scaling_factor, float(mscale_all_dim))
|
|
* attn_factor
|
|
)
|
|
self.device = device
|
|
super().__init__(
|
|
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
|
|
)
|
|
|
|
def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
|
|
pos_freqs = self.base ** (
|
|
torch.arange(0, self.rotary_dim, 2, dtype=torch.float, device=self.device)
|
|
/ self.rotary_dim
|
|
)
|
|
inv_freq_extrapolation = 1.0 / pos_freqs
|
|
inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
|
|
|
|
low, high = _yarn_find_correction_range(
|
|
self.beta_fast,
|
|
self.beta_slow,
|
|
self.rotary_dim,
|
|
self.base,
|
|
self.max_position_embeddings,
|
|
)
|
|
# Get n-d rotational scaling corrected for extrapolation
|
|
inv_freq_mask = (
|
|
1
|
|
- _yarn_linear_ramp_mask(
|
|
low, high, self.rotary_dim // 2, dtype=torch.float, device=self.device
|
|
)
|
|
) * self.extrapolation_factor
|
|
inv_freq = (
|
|
inv_freq_interpolation * (1 - inv_freq_mask)
|
|
+ inv_freq_extrapolation * inv_freq_mask
|
|
)
|
|
return inv_freq
|
|
|
|
def _compute_cos_sin_cache(self) -> torch.Tensor:
|
|
inv_freq = self._compute_inv_freq(self.scaling_factor)
|
|
t = torch.arange(
|
|
self.max_position_embeddings * self.scaling_factor,
|
|
device=self.device,
|
|
dtype=torch.float32,
|
|
)
|
|
freqs = torch.einsum("i,j -> ij", t, inv_freq)
|
|
cos = freqs.cos() * self.mscale
|
|
sin = freqs.sin() * self.mscale
|
|
cache = torch.cat((cos, sin), dim=-1)
|
|
return cache
|
|
|
|
def forward(
|
|
self,
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
fused_set_kv_buffer_arg=None,
|
|
output_q_rope=None,
|
|
offsets: torch.Tensor | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if _is_nvidia:
|
|
return super().forward(
|
|
positions=positions,
|
|
query=query,
|
|
key=key,
|
|
fused_set_kv_buffer_arg=fused_set_kv_buffer_arg,
|
|
output_q_rope=output_q_rope,
|
|
offsets=offsets,
|
|
)
|
|
|
|
dtype = query.dtype
|
|
query_rot = query[..., : self.rotary_dim]
|
|
key_rot = key[..., : self.rotary_dim]
|
|
if self.rotary_dim < self.head_size:
|
|
query_pass = query[..., self.rotary_dim :]
|
|
key_pass = key[..., self.rotary_dim :]
|
|
|
|
self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(positions.device)
|
|
cos_sin = self.cos_sin_cache[
|
|
torch.add(positions, offsets) if offsets is not None else positions
|
|
]
|
|
cos, sin = cos_sin.chunk(2, dim=-1)
|
|
if self.is_neox_style:
|
|
# Here we assume that the positions tensor has the
|
|
# shape [batch_size, seq_len].
|
|
cos = cos.repeat(1, 1, 2).unsqueeze(-2)
|
|
sin = sin.repeat(1, 1, 2).unsqueeze(-2)
|
|
else:
|
|
cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
|
|
sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)
|
|
|
|
rotate_fn = _rotate_neox if self.is_neox_style else _rotate_gptj
|
|
query_rot = query_rot * cos + rotate_fn(query_rot) * sin
|
|
key_rot = key_rot * cos + rotate_fn(key_rot) * sin
|
|
|
|
if self.rotary_dim < self.head_size:
|
|
query = torch.cat((query_rot, query_pass), dim=-1)
|
|
key = torch.cat((key_rot, key_pass), dim=-1)
|
|
else:
|
|
query = query_rot
|
|
key = key_rot
|
|
return query.to(dtype), key.to(dtype)
|
|
|
|
|
|
class Llama3RotaryEmbedding(RotaryEmbedding):
|
|
|
|
def __init__(
|
|
self,
|
|
head_size: int,
|
|
rotary_dim: int,
|
|
max_position_embeddings: int,
|
|
base: int,
|
|
is_neox_style: bool,
|
|
dtype: torch.dtype,
|
|
scaling_factor: float,
|
|
low_freq_factor: float,
|
|
high_freq_factor: float,
|
|
orig_max_position: int,
|
|
) -> None:
|
|
self.scaling_factor = scaling_factor
|
|
self.low_freq_factor = low_freq_factor
|
|
self.high_freq_factor = high_freq_factor
|
|
self.orig_max_position = orig_max_position
|
|
super().__init__(
|
|
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
|
|
)
|
|
|
|
def _compute_inv_freq(self, base: int | float) -> torch.Tensor:
|
|
inv_freqs = super()._compute_inv_freq(base)
|
|
low_freq_wavelen = self.orig_max_position / self.low_freq_factor
|
|
high_freq_wavelen = self.orig_max_position / self.high_freq_factor
|
|
|
|
wave_len = 2 * math.pi / inv_freqs
|
|
if self.low_freq_factor != self.high_freq_factor:
|
|
smooth = (self.orig_max_position / wave_len - self.low_freq_factor) / (
|
|
self.high_freq_factor - self.low_freq_factor
|
|
)
|
|
else:
|
|
smooth = 0
|
|
new_freqs = torch.where(
|
|
wave_len < high_freq_wavelen,
|
|
inv_freqs,
|
|
torch.where(
|
|
wave_len > low_freq_wavelen,
|
|
inv_freqs / self.scaling_factor,
|
|
(1 - smooth) * inv_freqs / self.scaling_factor + smooth * inv_freqs,
|
|
),
|
|
)
|
|
return new_freqs
|
|
|
|
|
|
class DynamicNTKAlphaRotaryEmbedding(RotaryEmbedding):
|
|
"""RotaryEmbedding extended with Dynamic NTK scaling.
|
|
|
|
Credits to the Reddit users /u/bloc97 and /u/emozilla
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
head_size: int,
|
|
rotary_dim: int,
|
|
max_position_embeddings: int,
|
|
base: int,
|
|
is_neox_style: bool,
|
|
scaling_alpha: float,
|
|
dtype: torch.dtype,
|
|
) -> None:
|
|
self.scaling_alpha = scaling_alpha
|
|
super().__init__(
|
|
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
|
|
)
|
|
|
|
def _compute_cos_sin_cache(self) -> torch.Tensor:
|
|
max_len = self.max_position_embeddings
|
|
base = self.base * self.scaling_alpha ** (
|
|
self.rotary_dim / (self.rotary_dim - 2)
|
|
)
|
|
|
|
inv_freq = self._compute_inv_freq(base)
|
|
t = torch.arange(max_len, dtype=torch.float)
|
|
|
|
freqs = torch.einsum("i,j -> ij", t, inv_freq)
|
|
cos = freqs.cos()
|
|
sin = freqs.sin()
|
|
cache = torch.cat((cos, sin), dim=-1)
|
|
return cache
|
|
|
|
|
|
class MRotaryEmbedding(RotaryEmbedding):
|
|
"""Rotary Embedding with Multimodal Sections."""
|
|
|
|
def __init__(
|
|
self,
|
|
head_size: int,
|
|
rotary_dim: int,
|
|
max_position_embeddings: int,
|
|
base: int,
|
|
is_neox_style: bool,
|
|
dtype: torch.dtype,
|
|
mrope_section: list[int] | None = None,
|
|
mrope_interleaved: bool = False,
|
|
) -> None:
|
|
super().__init__(
|
|
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
|
|
)
|
|
|
|
self.mrope_section = mrope_section
|
|
self.mrope_interleaved = mrope_interleaved
|
|
if self.mrope_section:
|
|
expected_sum = rotary_dim // 2
|
|
actual_sum = sum(self.mrope_section)
|
|
if actual_sum != expected_sum:
|
|
logger.warning(
|
|
f"MRoPE section sum mismatch: expected {expected_sum}, got {actual_sum}. "
|
|
f"Adjusting mrope_section to match rotary_dim // 2 = {expected_sum}"
|
|
)
|
|
# Auto-correct by scaling the mrope_section proportionally
|
|
if actual_sum > 0:
|
|
scale_factor = expected_sum / actual_sum
|
|
self.mrope_section = [
|
|
max(1, int(section * scale_factor))
|
|
for section in self.mrope_section
|
|
]
|
|
# Ensure the sum exactly matches by adjusting the last element
|
|
current_sum = sum(self.mrope_section)
|
|
if current_sum != expected_sum:
|
|
self.mrope_section[-1] += expected_sum - current_sum
|
|
else:
|
|
# If all sections are 0, create a default distribution
|
|
self.mrope_section = [
|
|
expected_sum // len(self.mrope_section)
|
|
] * len(self.mrope_section)
|
|
# Handle remainder
|
|
remainder = expected_sum % len(self.mrope_section)
|
|
for i in range(remainder):
|
|
self.mrope_section[i] += 1
|
|
|
|
logger.warning(
|
|
f"Corrected mrope_section: {self.mrope_section} (sum={sum(self.mrope_section)})"
|
|
)
|
|
|
|
@torch.compile(dynamic=True, backend=get_compiler_backend())
|
|
def forward(
|
|
self,
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
fused_set_kv_buffer_arg: FusedSetKVBufferArg | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""PyTorch-native implementation equivalent to forward().
|
|
|
|
Args:
|
|
positions:
|
|
[num_tokens,] (text only) or
|
|
[3, num_tokens] (T/H/W positions with multimodal inputs)
|
|
query: [num_tokens, num_heads * head_size]
|
|
key: [num_tokens, num_kv_heads * head_size]
|
|
"""
|
|
if fused_set_kv_buffer_arg is not None:
|
|
raise ValueError("save kv cache is not supported for MRotaryEmbedding.")
|
|
if positions.ndim not in (1, 2):
|
|
raise ValueError(f"positions must be 1D or 2D, got ndim={positions.ndim}.")
|
|
|
|
num_tokens = positions.shape[-1]
|
|
cos_sin = self.cos_sin_cache[positions]
|
|
cos, sin = cos_sin.chunk(2, dim=-1)
|
|
if positions.ndim == 2:
|
|
if not self.mrope_section:
|
|
raise RuntimeError("mrope_section must be set for 2D M-RoPE.")
|
|
|
|
if self.mrope_interleaved:
|
|
cos = apply_interleaved_rope(cos, self.mrope_section)
|
|
sin = apply_interleaved_rope(sin, self.mrope_section)
|
|
else:
|
|
cos = torch.cat(
|
|
[m[i] for i, m in enumerate(cos.split(self.mrope_section, dim=-1))],
|
|
dim=-1,
|
|
)
|
|
sin = torch.cat(
|
|
[m[i] for i, m in enumerate(sin.split(self.mrope_section, dim=-1))],
|
|
dim=-1,
|
|
)
|
|
|
|
query_shape = query.shape
|
|
query = query.view(num_tokens, -1, self.head_size)
|
|
query_rot = query[..., : self.rotary_dim]
|
|
query_pass = query[..., self.rotary_dim :]
|
|
query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
|
|
query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
|
|
|
|
key_shape = key.shape
|
|
key = key.view(num_tokens, -1, self.head_size)
|
|
key_rot = key[..., : self.rotary_dim]
|
|
key_pass = key[..., self.rotary_dim :]
|
|
key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
|
|
key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
|
|
return query, key
|
|
|
|
@staticmethod
|
|
def get_rope_index(
|
|
spatial_merge_size: int,
|
|
image_token_id: int,
|
|
video_token_id: int,
|
|
vision_start_token_id: int,
|
|
model_type: str,
|
|
tokens_per_second: int | None = None,
|
|
input_ids: torch.LongTensor | None = None,
|
|
image_grid_thw: torch.LongTensor | None = None,
|
|
video_grid_thw: torch.LongTensor | None = None,
|
|
second_per_grid_ts: torch.Tensor | None = None,
|
|
**kwargs,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
mrope_position_deltas = []
|
|
if input_ids is not None and (
|
|
image_grid_thw is not None or video_grid_thw is not None
|
|
):
|
|
total_input_ids = input_ids
|
|
position_ids = torch.ones(
|
|
3,
|
|
input_ids.shape[0],
|
|
input_ids.shape[1],
|
|
dtype=input_ids.dtype,
|
|
device=input_ids.device,
|
|
)
|
|
image_index, video_index = 0, 0
|
|
for i, input_ids in enumerate(total_input_ids):
|
|
image_nums, video_nums = 0, 0
|
|
vision_start_indices = torch.argwhere(
|
|
input_ids == vision_start_token_id
|
|
).squeeze(1)
|
|
vision_tokens = input_ids[vision_start_indices + 1]
|
|
image_nums = (vision_tokens == image_token_id).sum()
|
|
video_nums = (vision_tokens == video_token_id).sum()
|
|
input_tokens = input_ids.tolist()
|
|
llm_pos_ids_list: list = []
|
|
st = 0
|
|
remain_images, remain_videos = image_nums, video_nums
|
|
for _ in range(image_nums + video_nums):
|
|
if image_token_id in input_tokens and remain_images > 0:
|
|
ed_image = input_tokens.index(image_token_id, st)
|
|
else:
|
|
ed_image = len(input_tokens) + 1
|
|
if video_token_id in input_tokens and remain_videos > 0:
|
|
ed_video = input_tokens.index(video_token_id, st)
|
|
else:
|
|
ed_video = len(input_tokens) + 1
|
|
if ed_image < ed_video:
|
|
t, h, w = (
|
|
image_grid_thw[image_index][0],
|
|
image_grid_thw[image_index][1],
|
|
image_grid_thw[image_index][2],
|
|
)
|
|
second_per_grid_t = 0
|
|
image_index += 1
|
|
remain_images -= 1
|
|
ed = ed_image
|
|
else:
|
|
t, h, w = (
|
|
video_grid_thw[video_index][0],
|
|
video_grid_thw[video_index][1],
|
|
video_grid_thw[video_index][2],
|
|
)
|
|
if second_per_grid_ts is not None:
|
|
second_per_grid_t = second_per_grid_ts[video_index]
|
|
else:
|
|
second_per_grid_t = 1.0
|
|
video_index += 1
|
|
remain_videos -= 1
|
|
ed = ed_video
|
|
llm_grid_t, llm_grid_h, llm_grid_w = (
|
|
t.item(),
|
|
h.item() // spatial_merge_size,
|
|
w.item() // spatial_merge_size,
|
|
)
|
|
text_len = ed - st
|
|
|
|
st_idx = (
|
|
llm_pos_ids_list[-1].max() + 1
|
|
if len(llm_pos_ids_list) > 0
|
|
else 0
|
|
)
|
|
llm_pos_ids_list.append(
|
|
torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
|
|
)
|
|
|
|
if model_type == "qwen2_5_vl":
|
|
range_tensor = torch.arange(llm_grid_t).view(-1, 1)
|
|
expanded_range = range_tensor.expand(
|
|
-1, llm_grid_h * llm_grid_w
|
|
)
|
|
|
|
time_tensor = (
|
|
expanded_range * second_per_grid_t * tokens_per_second
|
|
)
|
|
|
|
time_tensor_long = time_tensor.long()
|
|
t_index = time_tensor_long.flatten()
|
|
elif model_type in (
|
|
"qwen2_vl",
|
|
"qwen3_vl",
|
|
"qwen3_vl_moe",
|
|
"qwen3_5",
|
|
"qwen3_5_moe",
|
|
):
|
|
t_index = (
|
|
torch.arange(llm_grid_t)
|
|
.view(-1, 1)
|
|
.expand(-1, llm_grid_h * llm_grid_w)
|
|
.flatten()
|
|
)
|
|
else:
|
|
raise RuntimeError("Unimplemented")
|
|
h_index = (
|
|
torch.arange(llm_grid_h)
|
|
.view(1, -1, 1)
|
|
.expand(llm_grid_t, -1, llm_grid_w)
|
|
.flatten()
|
|
)
|
|
w_index = (
|
|
torch.arange(llm_grid_w)
|
|
.view(1, 1, -1)
|
|
.expand(llm_grid_t, llm_grid_h, -1)
|
|
.flatten()
|
|
)
|
|
llm_pos_ids_list.append(
|
|
torch.stack([t_index, h_index, w_index]) + text_len + st_idx
|
|
)
|
|
st = ed + llm_grid_t * llm_grid_h * llm_grid_w
|
|
|
|
if st < len(input_tokens):
|
|
st_idx = (
|
|
llm_pos_ids_list[-1].max() + 1
|
|
if len(llm_pos_ids_list) > 0
|
|
else 0
|
|
)
|
|
text_len = len(input_tokens) - st
|
|
llm_pos_ids_list.append(
|
|
torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
|
|
)
|
|
|
|
llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
|
|
position_ids[..., i, :] = llm_positions.to(position_ids.device)
|
|
mrope_position_deltas.append(
|
|
llm_positions.max() + 1 - len(total_input_ids[i])
|
|
)
|
|
mrope_position_deltas = torch.tensor(
|
|
mrope_position_deltas, device=input_ids.device
|
|
).unsqueeze(1)
|
|
return position_ids, mrope_position_deltas
|
|
else:
|
|
s = input_ids.shape[1]
|
|
position_ids = torch.arange(s)
|
|
position_ids = (
|
|
position_ids.unsqueeze(0).expand(3, -1, -1).to(input_ids.device)
|
|
)
|
|
max_position_ids = position_ids.max(0, keepdim=False)[0].max(
|
|
-1, keepdim=True
|
|
)[0]
|
|
mrope_position_deltas = max_position_ids + 1 - s
|
|
return position_ids, mrope_position_deltas
|
|
|
|
@staticmethod
|
|
def get_rope_index_glm4v(
|
|
input_ids: torch.Tensor,
|
|
hf_config: Any,
|
|
image_grid_thw: list[list[int]] | torch.Tensor,
|
|
video_grid_thw: list[list[int]] | torch.Tensor,
|
|
attention_mask: torch.Tensor,
|
|
**kwargs,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Get mrope input positions and delta value for GLM4V."""
|
|
image_token_id = hf_config.image_token_id
|
|
video_start_token_id = hf_config.video_start_token_id
|
|
video_end_token_id = hf_config.video_end_token_id
|
|
spatial_merge_size = hf_config.vision_config.spatial_merge_size
|
|
|
|
mrope_position_deltas = []
|
|
if input_ids is not None and (
|
|
image_grid_thw is not None or video_grid_thw is not None
|
|
):
|
|
total_input_ids = input_ids
|
|
if attention_mask is None:
|
|
attention_mask = torch.ones_like(total_input_ids)
|
|
position_ids = torch.ones(
|
|
3,
|
|
input_ids.shape[0],
|
|
input_ids.shape[1],
|
|
dtype=input_ids.dtype,
|
|
device=input_ids.device,
|
|
)
|
|
image_index, video_index = 0, 0
|
|
video_group_index = 0
|
|
attention_mask = attention_mask.to(total_input_ids.device)
|
|
for i, input_ids in enumerate(total_input_ids):
|
|
input_ids = input_ids[attention_mask[i] == 1]
|
|
input_tokens = input_ids.tolist()
|
|
|
|
input_token_type = []
|
|
video_check_flg = False
|
|
for token in input_tokens:
|
|
if token == video_start_token_id:
|
|
video_check_flg = True
|
|
elif token == video_end_token_id:
|
|
video_check_flg = False
|
|
|
|
if token == image_token_id and not video_check_flg:
|
|
input_token_type.append("image")
|
|
elif token == image_token_id and video_check_flg:
|
|
input_token_type.append("video")
|
|
else:
|
|
input_token_type.append("text")
|
|
|
|
input_type_group = []
|
|
for key, group in itertools.groupby(
|
|
enumerate(input_token_type), lambda x: x[1]
|
|
):
|
|
group = list(group)
|
|
start_index = group[0][0]
|
|
end_index = group[-1][0] + 1
|
|
input_type_group.append((key, start_index, end_index))
|
|
|
|
llm_pos_ids_list = []
|
|
video_frame_num = 1
|
|
for modality_type, start_idx, end_idx in input_type_group:
|
|
st_idx = (
|
|
llm_pos_ids_list[-1].max() + 1
|
|
if len(llm_pos_ids_list) > 0
|
|
else 0
|
|
)
|
|
|
|
if modality_type == "image":
|
|
t, h, w = (
|
|
image_grid_thw[image_index][0],
|
|
image_grid_thw[image_index][1],
|
|
image_grid_thw[image_index][2],
|
|
)
|
|
llm_grid_t, llm_grid_h, llm_grid_w = (
|
|
t.item(),
|
|
h.item() // spatial_merge_size,
|
|
w.item() // spatial_merge_size,
|
|
)
|
|
|
|
t_index = (
|
|
torch.arange(llm_grid_t)
|
|
.view(-1, 1)
|
|
.expand(-1, llm_grid_h * llm_grid_w)
|
|
.flatten()
|
|
)
|
|
h_index = (
|
|
torch.arange(llm_grid_h)
|
|
.view(1, -1, 1)
|
|
.expand(llm_grid_t, -1, llm_grid_w)
|
|
.flatten()
|
|
)
|
|
w_index = (
|
|
torch.arange(llm_grid_w)
|
|
.view(1, 1, -1)
|
|
.expand(llm_grid_t, llm_grid_h, -1)
|
|
.flatten()
|
|
)
|
|
llm_pos_ids_list.append(
|
|
torch.stack([t_index, h_index, w_index]) + st_idx
|
|
)
|
|
|
|
image_index += 1
|
|
video_frame_num = 1
|
|
|
|
elif modality_type == "video":
|
|
t, h, w = (
|
|
video_frame_num,
|
|
video_grid_thw[video_index][1],
|
|
video_grid_thw[video_index][2],
|
|
)
|
|
|
|
llm_grid_t, llm_grid_h, llm_grid_w = (
|
|
t,
|
|
h.item() // spatial_merge_size,
|
|
w.item() // spatial_merge_size,
|
|
)
|
|
|
|
for t_idx in range(llm_grid_t):
|
|
t_index = (
|
|
torch.tensor(t_idx)
|
|
.view(-1, 1)
|
|
.expand(-1, llm_grid_h * llm_grid_w)
|
|
.flatten()
|
|
)
|
|
|
|
h_index = (
|
|
torch.arange(llm_grid_h)
|
|
.view(1, -1, 1)
|
|
.expand(1, -1, llm_grid_w)
|
|
.flatten()
|
|
)
|
|
w_index = (
|
|
torch.arange(llm_grid_w)
|
|
.view(1, 1, -1)
|
|
.expand(1, llm_grid_h, -1)
|
|
.flatten()
|
|
)
|
|
llm_pos_ids_list.append(
|
|
torch.stack([t_index, h_index, w_index]) + st_idx
|
|
)
|
|
|
|
video_group_index += 1
|
|
|
|
if video_group_index >= video_grid_thw[video_index][0]:
|
|
video_index += 1
|
|
video_group_index = 0
|
|
|
|
video_frame_num += 1
|
|
|
|
else:
|
|
text_len = end_idx - start_idx
|
|
llm_pos_ids_list.append(
|
|
torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
|
|
)
|
|
|
|
video_frame_num = 1
|
|
|
|
llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
|
|
position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(
|
|
position_ids.device
|
|
)
|
|
mrope_position_deltas.append(
|
|
llm_positions.max() + 1 - len(total_input_ids[i])
|
|
)
|
|
mrope_position_deltas = torch.tensor(
|
|
mrope_position_deltas, device=input_ids.device
|
|
).unsqueeze(1)
|
|
return position_ids, mrope_position_deltas
|
|
else:
|
|
if attention_mask is not None:
|
|
position_ids = attention_mask.long().cumsum(-1) - 1
|
|
position_ids.masked_fill_(attention_mask == 0, 1)
|
|
position_ids = (
|
|
position_ids.unsqueeze(0)
|
|
.expand(3, -1, -1)
|
|
.to(attention_mask.device)
|
|
)
|
|
max_position_ids = position_ids.max(0, keepdim=False)[0].max(
|
|
-1, keepdim=True
|
|
)[0]
|
|
mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
|
|
else:
|
|
position_ids = (
|
|
torch.arange(input_ids.shape[1], device=input_ids.device)
|
|
.view(1, 1, -1)
|
|
.expand(3, input_ids.shape[0], -1)
|
|
)
|
|
mrope_position_deltas = torch.zeros(
|
|
[input_ids.shape[0], 1],
|
|
device=input_ids.device,
|
|
dtype=input_ids.dtype,
|
|
)
|
|
|
|
return position_ids, mrope_position_deltas
|
|
|
|
|
|
_ROPE_DICT: dict[tuple, RotaryEmbedding] = {}
|
|
|
|
|
|
def get_rope(
|
|
head_size: int,
|
|
rotary_dim: int,
|
|
max_position: int,
|
|
base: int,
|
|
is_neox_style: bool = True,
|
|
rope_scaling: dict[str, Any] | None = None,
|
|
dtype: torch.dtype | None = None,
|
|
partial_rotary_factor: float = 1.0,
|
|
) -> RotaryEmbedding:
|
|
if dtype is None:
|
|
dtype = torch.get_default_dtype()
|
|
if rope_scaling is not None:
|
|
# Transforms every value that is a list into a tuple for caching calls
|
|
rope_scaling_tuple = {
|
|
k: tuple(v) if isinstance(v, list) else v for k, v in rope_scaling.items()
|
|
}
|
|
rope_scaling_args = tuple(rope_scaling_tuple.items())
|
|
else:
|
|
rope_scaling_args = None
|
|
|
|
if partial_rotary_factor < 1.0:
|
|
rotary_dim = int(rotary_dim * partial_rotary_factor)
|
|
|
|
key = (
|
|
head_size,
|
|
rotary_dim,
|
|
max_position,
|
|
base,
|
|
is_neox_style,
|
|
rope_scaling_args,
|
|
dtype,
|
|
)
|
|
if key in _ROPE_DICT:
|
|
return _ROPE_DICT[key]
|
|
|
|
if rope_scaling is None:
|
|
rotary_emb = RotaryEmbedding(
|
|
head_size, rotary_dim, max_position, base, is_neox_style, dtype
|
|
)
|
|
else:
|
|
if "rope_type" in rope_scaling:
|
|
scaling_type = rope_scaling["rope_type"]
|
|
elif "type" in rope_scaling:
|
|
scaling_type = rope_scaling["type"]
|
|
else:
|
|
raise ValueError("Unknown RoPE scaling type")
|
|
|
|
if scaling_type == "llama3":
|
|
scaling_factor = rope_scaling["factor"]
|
|
low_freq_factor = rope_scaling["low_freq_factor"]
|
|
high_freq_factor = rope_scaling["high_freq_factor"]
|
|
original_max_position = rope_scaling["original_max_position_embeddings"]
|
|
rotary_emb = Llama3RotaryEmbedding(
|
|
head_size,
|
|
rotary_dim,
|
|
max_position,
|
|
base,
|
|
is_neox_style,
|
|
dtype,
|
|
scaling_factor,
|
|
low_freq_factor,
|
|
high_freq_factor,
|
|
original_max_position,
|
|
)
|
|
elif scaling_type == "default":
|
|
if "mrope_section" in rope_scaling:
|
|
rotary_emb = MRotaryEmbedding(
|
|
head_size,
|
|
rotary_dim,
|
|
max_position,
|
|
base,
|
|
is_neox_style,
|
|
dtype,
|
|
mrope_section=rope_scaling["mrope_section"],
|
|
mrope_interleaved=rope_scaling.get("mrope_interleaved", False),
|
|
)
|
|
else:
|
|
rotary_emb = RotaryEmbedding(
|
|
head_size,
|
|
rotary_dim,
|
|
max_position,
|
|
base,
|
|
is_neox_style,
|
|
dtype,
|
|
)
|
|
elif scaling_type == "linear":
|
|
scaling_factor = rope_scaling["factor"]
|
|
rotary_emb = LinearScalingRotaryEmbedding(
|
|
head_size,
|
|
rotary_dim,
|
|
max_position,
|
|
base,
|
|
is_neox_style,
|
|
scaling_factor,
|
|
dtype,
|
|
)
|
|
elif scaling_type == "dynamic":
|
|
scaling_factor = rope_scaling["factor"]
|
|
if "alpha" in rope_scaling:
|
|
rotary_emb = DynamicNTKAlphaRotaryEmbedding(
|
|
head_size,
|
|
rotary_dim,
|
|
max_position,
|
|
base,
|
|
is_neox_style,
|
|
rope_scaling["alpha"],
|
|
dtype,
|
|
)
|
|
else:
|
|
rotary_emb = DynamicNTKScalingRotaryEmbedding(
|
|
head_size,
|
|
rotary_dim,
|
|
max_position,
|
|
base,
|
|
is_neox_style,
|
|
scaling_factor,
|
|
dtype,
|
|
)
|
|
elif scaling_type == "yarn":
|
|
scaling_factor = rope_scaling["factor"]
|
|
original_max_position = rope_scaling["original_max_position_embeddings"]
|
|
extra_kwargs = {
|
|
k: v
|
|
for k, v in rope_scaling.items()
|
|
if k
|
|
in ("extrapolation_factor", "attn_factor", "beta_fast", "beta_slow")
|
|
}
|
|
rotary_emb = YaRNScalingRotaryEmbedding(
|
|
head_size,
|
|
rotary_dim,
|
|
original_max_position,
|
|
base,
|
|
is_neox_style,
|
|
scaling_factor,
|
|
dtype,
|
|
**extra_kwargs,
|
|
)
|
|
elif scaling_type == "deepseek_yarn":
|
|
scaling_factor = rope_scaling["factor"]
|
|
original_max_position = rope_scaling["original_max_position_embeddings"]
|
|
extra_kwargs = {
|
|
k: v
|
|
for k, v in rope_scaling.items()
|
|
if k
|
|
in (
|
|
"extrapolation_factor",
|
|
"attn_factor",
|
|
"beta_fast",
|
|
"beta_slow",
|
|
"mscale",
|
|
"mscale_all_dim",
|
|
)
|
|
}
|
|
rotary_emb = DeepseekScalingRotaryEmbedding(
|
|
head_size,
|
|
rotary_dim,
|
|
original_max_position,
|
|
base,
|
|
is_neox_style,
|
|
scaling_factor,
|
|
dtype,
|
|
**extra_kwargs,
|
|
)
|
|
elif scaling_type == "longrope":
|
|
short_factor = rope_scaling["short_factor"]
|
|
long_factor = rope_scaling["long_factor"]
|
|
original_max_position = rope_scaling["original_max_position_embeddings"]
|
|
extra_kwargs = {
|
|
k: v
|
|
for k, v in rope_scaling.items()
|
|
if k in ("short_mscale", "long_mscale")
|
|
}
|
|
rotary_emb = Phi3LongRoPEScaledRotaryEmbedding(
|
|
head_size,
|
|
rotary_dim,
|
|
max_position,
|
|
original_max_position,
|
|
base,
|
|
is_neox_style,
|
|
dtype,
|
|
short_factor,
|
|
long_factor,
|
|
**extra_kwargs,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
|
|
_ROPE_DICT[key] = rotary_emb
|
|
return rotary_emb
|