[Diffusion] Restruct and clean Diffusion rotary embedding (#19064)
This commit is contained in:
Xiaoyu Zhang
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,48 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/rotary_embedding.py
|
||||
|
||||
# Adapted from
|
||||
# https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/llama/modeling_llama.py
|
||||
# Copyright 2023 The vLLM team.
|
||||
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
||||
# and OPT implementations in this library. It has been modified from its
|
||||
# original forms to accommodate minor architectural differences compared
|
||||
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Rotary Positional Embeddings — unified public API (drop-in replacement)."""
|
||||
|
||||
from ._base import RotaryEmbedding
|
||||
from ._factory import get_rope, get_rotary_pos_embed
|
||||
from ._mrope import NDRotaryEmbedding
|
||||
from ._utils import (
|
||||
_apply_rotary_emb,
|
||||
apply_flashinfer_rope_qk_inplace,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# _utils
|
||||
"_apply_rotary_emb",
|
||||
"apply_flashinfer_rope_qk_inplace",
|
||||
# _base
|
||||
"RotaryEmbedding",
|
||||
# _mrope
|
||||
"NDRotaryEmbedding",
|
||||
# _factory
|
||||
"get_rope",
|
||||
"get_rotary_pos_embed",
|
||||
]
|
||||
@@ -0,0 +1,130 @@
|
||||
"""RotaryEmbedding base class and LinearScalingRotaryEmbedding variant."""
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
|
||||
|
||||
from ._utils import _apply_rotary_emb
|
||||
|
||||
|
||||
@CustomOp.register("rotary_embedding")
|
||||
class RotaryEmbedding(CustomOp):
|
||||
"""Original rotary positional embedding."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
head_size: int,
|
||||
rotary_dim: int,
|
||||
max_position_embeddings: int,
|
||||
base: int | float,
|
||||
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()
|
||||
cache = cache.to(dtype)
|
||||
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."""
|
||||
# NOTE(woosuk): 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_cuda(self, *args, **kwargs):
|
||||
return self.forward_native(*args, **kwargs)
|
||||
|
||||
def forward_native(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
offsets: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""A PyTorch-native implementation of forward()."""
|
||||
if offsets is not None:
|
||||
positions = positions + offsets
|
||||
positions = positions.flatten()
|
||||
num_tokens = positions.shape[0]
|
||||
cos_sin = self.cos_sin_cache.index_select(0, positions)
|
||||
cos, sin = cos_sin.chunk(2, 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
|
||||
|
||||
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):
|
||||
def __init__(
|
||||
self,
|
||||
head_size: int,
|
||||
rotary_dim: int,
|
||||
max_position_embeddings: int,
|
||||
base: int | float,
|
||||
is_neox_style: bool,
|
||||
dtype: torch.dtype,
|
||||
scaling_factor: float,
|
||||
) -> None:
|
||||
self.scaling_factor = float(scaling_factor)
|
||||
super().__init__(
|
||||
head_size=head_size,
|
||||
rotary_dim=rotary_dim,
|
||||
max_position_embeddings=max_position_embeddings,
|
||||
base=base,
|
||||
is_neox_style=is_neox_style,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def _compute_cos_sin_cache(self) -> torch.Tensor:
|
||||
inv_freq = self._compute_inv_freq(self.base)
|
||||
t = torch.arange(self.max_position_embeddings, dtype=torch.float)
|
||||
t = t / self.scaling_factor
|
||||
freqs = torch.einsum("i,j -> ij", t, inv_freq)
|
||||
cos = freqs.cos()
|
||||
sin = freqs.sin()
|
||||
cache = torch.cat((cos, sin), dim=-1)
|
||||
return cache
|
||||
@@ -0,0 +1,171 @@
|
||||
"""get_rope / get_rotary_pos_embed factory functions and module-level caches."""
|
||||
|
||||
from collections import OrderedDict
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from ._base import LinearScalingRotaryEmbedding, RotaryEmbedding
|
||||
from ._mrope import NDRotaryEmbedding, _to_tuple
|
||||
|
||||
_ROPE_DICT: dict[tuple, RotaryEmbedding] = {}
|
||||
_ND_ROPE_CACHE: "OrderedDict[tuple, NDRotaryEmbedding]" = OrderedDict()
|
||||
_ROPE_3D_CACHE: "OrderedDict[tuple, tuple[torch.Tensor, torch.Tensor]]" = OrderedDict()
|
||||
|
||||
|
||||
def get_rope(
|
||||
head_size: int,
|
||||
rotary_dim: int,
|
||||
max_position: int,
|
||||
base: int | float,
|
||||
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)
|
||||
max_position_embeddings = max_position
|
||||
rope_type = None
|
||||
if rope_scaling is not None:
|
||||
rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None))
|
||||
if rope_type in (None, "default"):
|
||||
rope_scaling = None
|
||||
elif rope_type == "linear":
|
||||
factor = float(rope_scaling.get("factor", 1.0))
|
||||
original_max = rope_scaling.get("original_max_position_embeddings", None)
|
||||
if original_max is not None:
|
||||
max_position_embeddings = max(
|
||||
max_position_embeddings, int(float(original_max) * factor)
|
||||
)
|
||||
key = (
|
||||
head_size,
|
||||
rotary_dim,
|
||||
max_position_embeddings,
|
||||
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_embeddings, base, is_neox_style, dtype
|
||||
)
|
||||
else:
|
||||
if rope_type == "linear":
|
||||
factor = float(rope_scaling.get("factor", 1.0))
|
||||
rotary_emb = LinearScalingRotaryEmbedding(
|
||||
head_size=head_size,
|
||||
rotary_dim=rotary_dim,
|
||||
max_position_embeddings=max_position_embeddings,
|
||||
base=base,
|
||||
is_neox_style=is_neox_style,
|
||||
dtype=dtype,
|
||||
scaling_factor=factor,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown RoPE scaling {rope_scaling}")
|
||||
_ROPE_DICT[key] = rotary_emb
|
||||
return rotary_emb
|
||||
|
||||
|
||||
def get_rotary_pos_embed(
|
||||
rope_sizes,
|
||||
hidden_size,
|
||||
heads_num,
|
||||
rope_dim_list,
|
||||
rope_theta,
|
||||
theta_rescale_factor=1.0,
|
||||
interpolation_factor=1.0,
|
||||
shard_dim: int = 0,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
start_frame: int = 0,
|
||||
device: torch.device | str | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Generate rotary positional embeddings for the given sizes.
|
||||
|
||||
Args:
|
||||
rope_sizes: Tuple of dimensions (t, h, w)
|
||||
hidden_size: Hidden dimension size
|
||||
heads_num: Number of attention heads
|
||||
rope_dim_list: List of dimensions for each axis, or None
|
||||
rope_theta: Base for frequency calculations
|
||||
theta_rescale_factor: Rescale factor for theta. Defaults to 1.0
|
||||
interpolation_factor: Factor to scale positions. Defaults to 1.0
|
||||
shard_dim: Which dimension to shard for sequence parallelism. Defaults to 0.
|
||||
|
||||
Returns:
|
||||
Tuple of (cos, sin) tensors for rotary embeddings
|
||||
"""
|
||||
|
||||
target_ndim = 3
|
||||
head_dim = hidden_size // heads_num
|
||||
|
||||
if rope_dim_list is None:
|
||||
rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)]
|
||||
|
||||
assert (
|
||||
sum(rope_dim_list) == head_dim
|
||||
), "sum(rope_dim_list) should equal to head_dim of attention layer"
|
||||
|
||||
# Get SP info - now handled within NDRotaryEmbedding
|
||||
# sp_group = get_sp_group()
|
||||
# sp_rank = sp_group.rank_in_group
|
||||
# sp_world_size = sp_group.world_size
|
||||
|
||||
# Simple LRU cache keyed by parameters
|
||||
global _ND_ROPE_CACHE
|
||||
key = (
|
||||
tuple(rope_dim_list),
|
||||
float(rope_theta),
|
||||
(
|
||||
tuple(theta_rescale_factor)
|
||||
if isinstance(theta_rescale_factor, list)
|
||||
else float(theta_rescale_factor)
|
||||
),
|
||||
(
|
||||
tuple(interpolation_factor)
|
||||
if isinstance(interpolation_factor, list)
|
||||
else float(interpolation_factor)
|
||||
),
|
||||
dtype,
|
||||
)
|
||||
|
||||
cache_hit = key in _ND_ROPE_CACHE
|
||||
if cache_hit:
|
||||
rope_emb = _ND_ROPE_CACHE.pop(key)
|
||||
_ND_ROPE_CACHE[key] = rope_emb # move to end (most-recent)
|
||||
else:
|
||||
rope_emb = NDRotaryEmbedding(
|
||||
rope_dim_list=rope_dim_list,
|
||||
rope_theta=rope_theta,
|
||||
theta_rescale_factor=theta_rescale_factor,
|
||||
interpolation_factor=interpolation_factor,
|
||||
dtype=dtype,
|
||||
)
|
||||
_ND_ROPE_CACHE[key] = rope_emb
|
||||
if len(_ND_ROPE_CACHE) > 16:
|
||||
# pop least-recently-used
|
||||
_ND_ROPE_CACHE.pop(next(iter(_ND_ROPE_CACHE)))
|
||||
|
||||
freqs_cos, freqs_sin = rope_emb.forward_from_grid(
|
||||
grid_size=_to_tuple(rope_sizes, dim=3),
|
||||
shard_dim=shard_dim,
|
||||
start_frame=start_frame,
|
||||
device=device,
|
||||
)
|
||||
return freqs_cos, freqs_sin
|
||||
@@ -0,0 +1,392 @@
|
||||
"""MRotaryEmbedding, YaRNScalingMRotaryEmbedding, NDRotaryEmbedding, OneDRotaryEmbedding."""
|
||||
|
||||
import functools
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_group
|
||||
|
||||
|
||||
def _to_tuple(x: int | tuple[int, ...], dim: int = 2) -> tuple[int, ...]:
|
||||
if isinstance(x, int):
|
||||
return (x,) * dim
|
||||
elif len(x) == dim:
|
||||
return x
|
||||
else:
|
||||
raise ValueError(f"Expected length {dim} or int, but got {x}")
|
||||
|
||||
|
||||
def get_1d_rotary_pos_embed(
|
||||
dim: int,
|
||||
pos: torch.FloatTensor | int,
|
||||
theta: float = 10000.0,
|
||||
theta_rescale_factor: float = 1.0,
|
||||
interpolation_factor: float = 1.0,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
device: torch.device | str | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Precompute the frequency tensor for complex exponential (cis) with given dimensions.
|
||||
(Note: `cis` means `cos + i * sin`, where i is the imaginary unit.)
|
||||
|
||||
This function calculates a frequency tensor with complex exponential using the given dimension 'dim'
|
||||
and the end index 'end'. The 'theta' parameter scales the frequencies.
|
||||
|
||||
Args:
|
||||
dim (int): Dimension of the frequency tensor.
|
||||
pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar
|
||||
theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.
|
||||
theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0.
|
||||
interpolation_factor (float, optional): Factor to scale positions. Defaults to 1.0.
|
||||
|
||||
Returns:
|
||||
freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D]
|
||||
"""
|
||||
if isinstance(pos, int):
|
||||
pos = torch.arange(pos, dtype=dtype, device=device)
|
||||
elif (
|
||||
isinstance(pos, torch.Tensor)
|
||||
and device is not None
|
||||
and pos.device != torch.device(device)
|
||||
):
|
||||
# Ensure positions are on the requested device to avoid implicit CPU ops.
|
||||
pos = pos.to(device)
|
||||
|
||||
# proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
|
||||
# has some connection to NTK literature
|
||||
if theta_rescale_factor != 1.0:
|
||||
theta *= theta_rescale_factor ** (dim / (dim - 2))
|
||||
|
||||
freqs = 1.0 / (
|
||||
theta
|
||||
** (torch.arange(0, dim, 2, device=device)[: (dim // 2)].to(dtype) / dim).to(
|
||||
device=device
|
||||
)
|
||||
) # [D/2]
|
||||
freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2]
|
||||
freqs_cos = freqs.cos() # [S, D/2]
|
||||
freqs_sin = freqs.sin() # [S, D/2]
|
||||
return freqs_cos, freqs_sin
|
||||
|
||||
|
||||
class OneDRotaryEmbedding(torch.nn.Module):
|
||||
"""1D rotary positional embedding with caching."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
theta: float = 10000.0,
|
||||
theta_rescale_factor: float = 1.0,
|
||||
interpolation_factor: float = 1.0,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
use_real: bool = False,
|
||||
repeat_interleave_real: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
assert dim % 2 == 0
|
||||
self.dim = dim
|
||||
self.theta = theta
|
||||
self.theta_rescale_factor = theta_rescale_factor
|
||||
self.interpolation_factor = interpolation_factor
|
||||
# dtype of freqs
|
||||
self.dtype = dtype
|
||||
self.use_real = use_real
|
||||
self.repeat_interleave_real = repeat_interleave_real
|
||||
|
||||
def build_freqs(self, device):
|
||||
freqs = 1.0 / (
|
||||
self.theta
|
||||
** (
|
||||
torch.arange(0, self.dim, 2, dtype=self.dtype, device=device)[
|
||||
: (self.dim // 2)
|
||||
]
|
||||
/ self.dim
|
||||
).to(device=device)
|
||||
)
|
||||
return freqs
|
||||
|
||||
def build_freqs_outer(self, pos: torch.Tensor, device):
|
||||
theta = self.theta
|
||||
# proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
|
||||
# has some connection to NTK literature
|
||||
if self.theta_rescale_factor != 1.0:
|
||||
theta *= self.theta_rescale_factor ** (self.dim / (self.dim - 2))
|
||||
|
||||
freqs = self.build_freqs(device)
|
||||
|
||||
freqs = torch.outer(pos * self.interpolation_factor, freqs)
|
||||
freqs_cos = freqs.cos()
|
||||
freqs_sin = freqs.sin()
|
||||
|
||||
if self.use_real and self.repeat_interleave_real:
|
||||
freqs_cos = freqs_cos.repeat_interleave(2, dim=1)
|
||||
freqs_sin = freqs_sin.repeat_interleave(2, dim=1)
|
||||
|
||||
return freqs_cos.float(), freqs_sin.float()
|
||||
|
||||
@functools.lru_cache(maxsize=16)
|
||||
def forward_from_grid(
|
||||
self, seq_len: int, start_pos: int, device_str: str
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
device = torch.device(device_str)
|
||||
pos = torch.arange(
|
||||
start_pos, start_pos + seq_len, dtype=self.dtype, device=device
|
||||
)
|
||||
|
||||
freqs_cos, freqs_sin = self.build_freqs_outer(pos, device)
|
||||
return freqs_cos, freqs_sin
|
||||
|
||||
def forward(self, pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Calculates 1D rotary embeddings for the given positions.
|
||||
|
||||
This method converts the input tensor to a hashable representation
|
||||
and calls a cached helper method to perform the computation.
|
||||
"""
|
||||
pos_tuple = tuple(pos.tolist())
|
||||
device_str = str(pos.device)
|
||||
return self._forward_cached(pos_tuple, device_str)
|
||||
|
||||
@functools.lru_cache(maxsize=16)
|
||||
def _forward_cached(
|
||||
self, pos_tuple: tuple, device_str: str
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
The core implementation that computes 1D rotary embeddings.
|
||||
This method is wrapped by an LRU cache.
|
||||
"""
|
||||
device = torch.device(device_str)
|
||||
pos = torch.as_tensor(pos_tuple, dtype=self.dtype, device=device)
|
||||
freqs_cos, freqs_sin = self.build_freqs_outer(pos, device)
|
||||
return freqs_cos, freqs_sin
|
||||
|
||||
|
||||
class NDRotaryEmbedding(torch.nn.Module):
|
||||
"""N-dimensional rotary positional embedding."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rope_dim_list: list[int],
|
||||
rope_theta: float,
|
||||
theta_rescale_factor: float | list[float] = 1.0,
|
||||
interpolation_factor: float | list[float] = 1.0,
|
||||
use_real: bool = False,
|
||||
repeat_interleave_real: bool = False,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
):
|
||||
super().__init__()
|
||||
self.rope_dim_list = rope_dim_list
|
||||
self.ndim = len(rope_dim_list)
|
||||
self.rope_theta = rope_theta
|
||||
# dtype of freqs
|
||||
# does not control the output dtype
|
||||
self.dtype = dtype
|
||||
|
||||
if isinstance(theta_rescale_factor, (int, float)):
|
||||
self.theta_rescale_factor = [theta_rescale_factor] * self.ndim
|
||||
elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1:
|
||||
self.theta_rescale_factor = [theta_rescale_factor[0]] * self.ndim
|
||||
else:
|
||||
self.theta_rescale_factor = theta_rescale_factor
|
||||
assert (
|
||||
len(self.theta_rescale_factor) == self.ndim
|
||||
), "len(theta_rescale_factor) should equal to len(rope_dim_list)"
|
||||
|
||||
if isinstance(interpolation_factor, (int, float)):
|
||||
self.interpolation_factor = [interpolation_factor] * self.ndim
|
||||
elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1:
|
||||
self.interpolation_factor = [interpolation_factor[0]] * self.ndim
|
||||
else:
|
||||
self.interpolation_factor = interpolation_factor
|
||||
assert (
|
||||
len(self.interpolation_factor) == self.ndim
|
||||
), "len(interpolation_factor) should equal to len(rope_dim_list)"
|
||||
|
||||
self.rope_generators: list[OneDRotaryEmbedding] = torch.nn.ModuleList()
|
||||
_config_to_gen_idx: dict[tuple, int] = {}
|
||||
self.dim_idx_to_gen_idx: list[int] = []
|
||||
|
||||
for i in range(self.ndim):
|
||||
dim = self.rope_dim_list[i]
|
||||
rescale = self.theta_rescale_factor[i]
|
||||
interp = self.interpolation_factor[i]
|
||||
|
||||
config_key = (dim, rescale, interp, use_real, repeat_interleave_real)
|
||||
if config_key not in _config_to_gen_idx:
|
||||
generator = OneDRotaryEmbedding(
|
||||
dim=dim,
|
||||
theta=self.rope_theta,
|
||||
theta_rescale_factor=rescale,
|
||||
interpolation_factor=interp,
|
||||
dtype=self.dtype,
|
||||
use_real=use_real,
|
||||
repeat_interleave_real=repeat_interleave_real,
|
||||
)
|
||||
_config_to_gen_idx[config_key] = len(self.rope_generators)
|
||||
self.rope_generators.append(generator)
|
||||
|
||||
gen_idx = _config_to_gen_idx[config_key]
|
||||
self.dim_idx_to_gen_idx.append(gen_idx)
|
||||
|
||||
def forward(self, positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Calculates n-d rotary embeddings for given absolute positions.
|
||||
|
||||
Args:
|
||||
positions (torch.Tensor): A tensor of shape `[num_tokens, ndim]`
|
||||
containing the integer coordinates for each token.
|
||||
|
||||
Returns:
|
||||
A tuple of (cos, sin) tensors.
|
||||
"""
|
||||
# Caching wrapper: convert tensor to a hashable tuple of tuples.
|
||||
pos_tuple = tuple(map(tuple, positions.tolist()))
|
||||
device_str = str(positions.device)
|
||||
return self._forward_cached(pos_tuple, device_str)
|
||||
|
||||
@functools.lru_cache(maxsize=16)
|
||||
def _forward_cached(
|
||||
self, pos_tuple: tuple[tuple[int, ...], ...], device_str: str
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
The core implementation that computes embeddings from a position tensor.
|
||||
This method is wrapped by an LRU cache.
|
||||
"""
|
||||
device = torch.device(device_str)
|
||||
positions = torch.tensor(pos_tuple, dtype=torch.long, device=device)
|
||||
return self.forward_uncached(pos=positions)
|
||||
|
||||
def forward_uncached(self, pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
The core implementation that computes embeddings from a position tensor.
|
||||
This method is wrapped by an LRU cache.
|
||||
"""
|
||||
device = pos.device
|
||||
|
||||
# Pre-allocate the final tensors for efficiency.
|
||||
num_tokens = pos.shape[0]
|
||||
first_generator = self.rope_generators[0]
|
||||
if first_generator.use_real and first_generator.repeat_interleave_real:
|
||||
head_dim = sum(self.rope_dim_list)
|
||||
else:
|
||||
head_dim = sum(self.rope_dim_list) // 2
|
||||
|
||||
cos = torch.empty((num_tokens, head_dim), device=device, dtype=self.dtype)
|
||||
sin = torch.empty((num_tokens, head_dim), device=device, dtype=self.dtype)
|
||||
|
||||
col_offset = 0
|
||||
for i in range(self.ndim):
|
||||
# Extract position coordinates for the current dimension for all tokens.
|
||||
pos_i = pos[:, i].to(self.dtype)
|
||||
|
||||
# Get the appropriate 1D generator.
|
||||
gen_idx = self.dim_idx_to_gen_idx[i]
|
||||
generator = self.rope_generators[gen_idx]
|
||||
|
||||
# Calculate 1D embeddings.
|
||||
cos_1d, sin_1d = generator(pos_i)
|
||||
|
||||
slice_width = cos_1d.shape[1]
|
||||
cos[:, col_offset : col_offset + slice_width] = cos_1d
|
||||
sin[:, col_offset : col_offset + slice_width] = sin_1d
|
||||
col_offset += slice_width
|
||||
|
||||
return cos.float(), sin.float()
|
||||
|
||||
def forward_from_grid(
|
||||
self,
|
||||
grid_size: tuple[int, ...],
|
||||
shard_dim: int = 0,
|
||||
start_frame: int = 0,
|
||||
device: torch.device | str | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Handles sp internally
|
||||
"""
|
||||
# Caching wrapper: use grid parameters directly as the key.
|
||||
# grid_tuple = _to_tuple(grid_size, dim=self.ndim)
|
||||
device_str = str(device) if device is not None else "cpu"
|
||||
return self._forward_cached_from_grid(
|
||||
grid_size, shard_dim, start_frame, device_str
|
||||
)
|
||||
|
||||
@functools.lru_cache(maxsize=16)
|
||||
def _forward_cached_from_grid(
|
||||
self,
|
||||
grid_size: tuple[int, ...],
|
||||
shard_dim: int,
|
||||
start_frame: int,
|
||||
device_str: str,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Computes embeddings for a structured grid, using a highly efficient
|
||||
implementation that avoids materializing the full position tensor.
|
||||
This method is wrapped by an LRU cache.
|
||||
"""
|
||||
device = torch.device(device_str)
|
||||
sp_group = get_sp_group()
|
||||
sp_rank = sp_group.rank_in_group
|
||||
sp_world_size = sp_group.world_size
|
||||
|
||||
sizes = _to_tuple(grid_size, dim=self.ndim)
|
||||
starts = (0,) * self.ndim
|
||||
|
||||
# Apply sequence parallel sharding to the sizes and compute shard offset
|
||||
shard_sizes = list(sizes)
|
||||
shard_offsets = [0] * self.ndim
|
||||
if sp_world_size > 1:
|
||||
assert sizes[shard_dim] % sp_world_size == 0, (
|
||||
f"Dimension {shard_dim} with size {sizes[shard_dim]} is not divisible "
|
||||
f"by sequence parallel world size {sp_world_size}"
|
||||
)
|
||||
shard_size = sizes[shard_dim] // sp_world_size
|
||||
shard_offsets[shard_dim] = sp_rank * shard_size
|
||||
shard_sizes[shard_dim] = shard_size
|
||||
|
||||
# Pre-allocate outputs on the requested device to avoid CPU ops and extra cats
|
||||
num_tokens = 1
|
||||
for s in shard_sizes:
|
||||
num_tokens *= int(s)
|
||||
head_dim_half = sum(self.rope_dim_list) // 2
|
||||
cos = torch.empty((num_tokens, head_dim_half), device=device, dtype=self.dtype)
|
||||
sin = torch.empty((num_tokens, head_dim_half), device=device, dtype=self.dtype)
|
||||
|
||||
# Compute per-axis 1D embeddings once and expand via repeats to [N, d_i/2]
|
||||
col_offset = 0
|
||||
for i in range(self.ndim):
|
||||
dim_i = self.rope_dim_list[i]
|
||||
dim_i_half = dim_i // 2
|
||||
size_i = int(shard_sizes[i])
|
||||
|
||||
# Starting position for this axis, with optional frame offset for time axis (i==0)
|
||||
base_offset = starts[i]
|
||||
if i == 0 and start_frame > 0:
|
||||
base_offset += start_frame
|
||||
if sp_world_size > 1 and i == shard_dim:
|
||||
base_offset += shard_offsets[i]
|
||||
|
||||
gen_idx = self.dim_idx_to_gen_idx[i]
|
||||
generator = self.rope_generators[gen_idx]
|
||||
cos_1d, sin_1d = generator.forward_from_grid(
|
||||
size_i, base_offset, device_str
|
||||
)
|
||||
|
||||
# Expand to [num_tokens, dim_i/2] matching flatten order (last dims vary fastest)
|
||||
repeats_per_entry = 1
|
||||
for j in range(i + 1, self.ndim):
|
||||
repeats_per_entry *= int(shard_sizes[j])
|
||||
tile_count = 1
|
||||
for j in range(0, i):
|
||||
tile_count *= int(shard_sizes[j])
|
||||
|
||||
cos_expanded = cos_1d.repeat_interleave(repeats_per_entry, dim=0)
|
||||
sin_expanded = sin_1d.repeat_interleave(repeats_per_entry, dim=0)
|
||||
if tile_count > 1:
|
||||
cos_expanded = cos_expanded.repeat(tile_count, 1)
|
||||
sin_expanded = sin_expanded.repeat(tile_count, 1)
|
||||
|
||||
cos[:, col_offset : col_offset + dim_i_half] = cos_expanded
|
||||
sin[:, col_offset : col_offset + dim_i_half] = sin_expanded
|
||||
col_offset += dim_i_half
|
||||
|
||||
return cos.float(), sin.float()
|
||||
@@ -0,0 +1,121 @@
|
||||
"""Primitive RoPE ops: rotate helpers and apply_rotary_emb utilities."""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.diffusion.triton.rotary import apply_rotary_embedding
|
||||
|
||||
|
||||
def _apply_rotary_emb(
|
||||
x: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
is_neox_style: bool,
|
||||
interleaved: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: [num_tokens, num_heads, head_size] or [num_tokens, 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:
|
||||
cos = cos.unsqueeze(-2)
|
||||
sin = sin.unsqueeze(-2)
|
||||
if is_neox_style:
|
||||
x1, x2 = torch.chunk(x, 2, dim=-1)
|
||||
else:
|
||||
x1 = x[..., ::2]
|
||||
x2 = x[..., 1::2]
|
||||
o1 = (x1.float() * cos - x2.float() * sin).type_as(x)
|
||||
o2 = (x2.float() * cos + x1.float() * sin).type_as(x)
|
||||
return torch.cat((o1, o2), dim=-1)
|
||||
else:
|
||||
return apply_rotary_embedding(x, cos, sin, interleaved)
|
||||
|
||||
|
||||
def apply_flashinfer_rope_qk_inplace(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
*,
|
||||
head_size: Optional[int] = None,
|
||||
is_neox: bool = False,
|
||||
positions: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
if q.dim() != 4 or k.dim() != 4:
|
||||
raise ValueError(
|
||||
f"Expected q/k to be 4D [bsz, seqlen, nheads, head_size], "
|
||||
f"got q:{tuple(q.shape)} k:{tuple(k.shape)}"
|
||||
)
|
||||
if q.shape != k.shape:
|
||||
raise ValueError(
|
||||
f"q and k must have the same shape, got {q.shape} vs {k.shape}"
|
||||
)
|
||||
|
||||
if not (isinstance(cos_sin_cache, torch.Tensor) and cos_sin_cache.dim() == 2):
|
||||
raise ValueError("cos_sin_cache must be a 2D torch.Tensor")
|
||||
|
||||
bsz, seqlen, nheads, d = q.shape
|
||||
if head_size is None:
|
||||
head_size = d
|
||||
if head_size != d:
|
||||
raise ValueError(f"head_size mismatch: inferred {d}, but head_size={head_size}")
|
||||
|
||||
try:
|
||||
from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
|
||||
except ImportError:
|
||||
# Triton fallback for AMD/ROCm where FlashInfer is not available
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"FlashInfer not available, using Triton fallback for RoPE",
|
||||
stacklevel=2,
|
||||
)
|
||||
half_size = cos_sin_cache.shape[-1] // 2
|
||||
if positions is None:
|
||||
cos = cos_sin_cache[:seqlen, :half_size].to(q.dtype)
|
||||
sin = cos_sin_cache[:seqlen, half_size:].to(q.dtype)
|
||||
cos = cos.unsqueeze(0).expand(bsz, -1, -1).reshape(bsz * seqlen, -1)
|
||||
sin = sin.unsqueeze(0).expand(bsz, -1, -1).reshape(bsz * seqlen, -1)
|
||||
else:
|
||||
positions = positions.to(cos_sin_cache.device).view(-1)
|
||||
cos = cos_sin_cache[positions, :half_size].to(q.dtype)
|
||||
sin = cos_sin_cache[positions, half_size:].to(q.dtype)
|
||||
q_flat = q.reshape(bsz * seqlen, nheads, d)
|
||||
k_flat = k.reshape(bsz * seqlen, nheads, d)
|
||||
q_rot = apply_rotary_embedding(q_flat, cos, sin, interleaved=not is_neox)
|
||||
k_rot = apply_rotary_embedding(k_flat, cos, sin, interleaved=not is_neox)
|
||||
return q_rot.view(bsz, seqlen, nheads, d), k_rot.view(bsz, seqlen, nheads, d)
|
||||
|
||||
if positions is None:
|
||||
pos_1d = torch.arange(seqlen, device=q.device, dtype=torch.long)
|
||||
positions = pos_1d if bsz == 1 else pos_1d.repeat(bsz)
|
||||
else:
|
||||
if not (
|
||||
isinstance(positions, torch.Tensor)
|
||||
and positions.dtype == torch.long
|
||||
and positions.dim() == 1
|
||||
):
|
||||
raise ValueError("positions must be a 1D torch.long Tensor")
|
||||
if positions.numel() != bsz * seqlen:
|
||||
raise ValueError(
|
||||
f"positions length must be bsz*seqlen={bsz*seqlen}, got {positions.numel()}"
|
||||
)
|
||||
|
||||
q_flat = q.reshape(bsz * seqlen, nheads * d).contiguous()
|
||||
k_flat = k.reshape(bsz * seqlen, nheads * d).contiguous()
|
||||
apply_rope_with_cos_sin_cache_inplace(
|
||||
positions=positions,
|
||||
query=q_flat,
|
||||
key=k_flat,
|
||||
head_size=d,
|
||||
cos_sin_cache=cos_sin_cache,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
return q_flat.view(bsz, seqlen, nheads, d), k_flat.view(bsz, seqlen, nheads, d)
|
||||
Reference in New Issue
Block a user