wehub-resource-sync
443 lines
17 KiB
Python
443 lines
17 KiB
Python
"""Utilities for Multimodal Rotary Position Embeddings (MRoPE)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
from typing import List, Optional, Tuple # noqa: UP035
|
|
|
|
import numpy as np
|
|
from tvm import te, tirx
|
|
from tvm.relax.frontend import nn
|
|
from tvm.relax.frontend.nn import Tensor, op
|
|
|
|
|
|
def _rotate_half(x: Tensor) -> Tensor:
|
|
"""Rotate the last dimension of ``x`` by swapping pairs."""
|
|
|
|
x1, x2 = op.split(x, 2, axis=-1)
|
|
return op.concat([op.negative(x2), x1], dim=-1)
|
|
|
|
|
|
def _repeat_mrope_section(section: Sequence[int]) -> Tuple[int, ...]: # noqa: UP006
|
|
if not section:
|
|
raise ValueError("mrope_section must not be empty.")
|
|
if any(s <= 0 for s in section):
|
|
raise ValueError(f"All mrope_section entries must be positive, got {section}.")
|
|
return tuple(section) * 2
|
|
|
|
|
|
def _split_indices_from_sizes(sizes: Sequence[int]) -> List[int]: # noqa: UP006
|
|
indices: List[int] = [] # noqa: UP006
|
|
running = 0
|
|
# Drop the final cumulative sum so split() keeps the last chunk.
|
|
for size in sizes[:-1]:
|
|
running += size
|
|
indices.append(running)
|
|
return indices
|
|
|
|
|
|
def _reorder_cos_sin(
|
|
tensor: Tensor,
|
|
split_sizes: Sequence[int],
|
|
) -> Tensor:
|
|
"""Reorder cos/sin tensors so the head dimension follows T/H/W repeating sections."""
|
|
|
|
if not split_sizes:
|
|
raise ValueError("split_sizes must not be empty.")
|
|
split_points = _split_indices_from_sizes(split_sizes)
|
|
# relax.op.split returns a Python tuple, so we can iterate directly.
|
|
sections = op.split(tensor, indices_or_sections=split_points, axis=-1)
|
|
reordered = []
|
|
for idx, chunk in enumerate(sections):
|
|
axis_selector = nn.Tensor.from_const(np.array([idx % 3], dtype="int32"))
|
|
axis_slice = op.take(chunk, axis_selector, axis=0)
|
|
reordered.append(nn.op.squeeze(axis_slice, 0))
|
|
return op.concat(reordered, dim=-1)
|
|
|
|
|
|
class MultimodalRotaryEmbedding(nn.Module):
|
|
"""Generate cosine/sine tables for multimodal rotary embeddings."""
|
|
|
|
def __init__(
|
|
self,
|
|
head_dim: int,
|
|
theta: float,
|
|
mrope_section: Sequence[int],
|
|
attention_scaling: float = 1.0,
|
|
) -> None:
|
|
if head_dim % 2 != 0:
|
|
raise ValueError(f"head_dim must be even for RoPE, got {head_dim}.")
|
|
self.head_dim = head_dim
|
|
self.theta = theta
|
|
self.attention_scaling = attention_scaling
|
|
self.mrope_section = tuple(mrope_section)
|
|
self._inv_freq = 1.0 / (
|
|
theta ** (np.arange(0, head_dim, 2, dtype="float32") / np.float32(head_dim))
|
|
)
|
|
|
|
def forward(self, reference: Tensor, position_ids: Tensor) -> Tuple[Tensor, Tensor]: # noqa: UP006
|
|
"""Return ``(cos, sin)`` with shape ``(3, batch, seq, head_dim)``."""
|
|
if len(position_ids.shape) != 3:
|
|
raise ValueError(
|
|
"position_ids must be rank-3 with either "
|
|
"(batch, seq, 3) or (3, batch, seq) layout, "
|
|
f"got shape {position_ids.shape}."
|
|
)
|
|
if isinstance(position_ids.shape[0], int) and position_ids.shape[0] == 3:
|
|
batch_size, seq_len = position_ids.shape[1], position_ids.shape[2]
|
|
pos_tensor = op.reshape(position_ids, (3, batch_size, 1, seq_len))
|
|
elif isinstance(position_ids.shape[-1], int) and position_ids.shape[-1] == 3:
|
|
batch_size, seq_len = position_ids.shape[0], position_ids.shape[1]
|
|
permuted_pos = op.permute_dims(position_ids, axes=[2, 0, 1])
|
|
pos_tensor = op.reshape(permuted_pos, (3, batch_size, 1, seq_len))
|
|
else:
|
|
raise ValueError(
|
|
"position_ids must have exactly one static dimension of size 3, "
|
|
f"got shape {position_ids.shape}."
|
|
)
|
|
|
|
dtype = reference.dtype
|
|
inv_freq_tensor = nn.Tensor.from_const(self._inv_freq.reshape(1, 1, -1, 1))
|
|
inv_freq_tensor = op.broadcast_to(inv_freq_tensor, (3, batch_size, self._inv_freq.size, 1))
|
|
|
|
freqs = op.matmul(inv_freq_tensor.astype("float32"), pos_tensor.astype("float32"))
|
|
freqs = op.permute_dims(freqs, axes=[0, 1, 3, 2])
|
|
emb = op.concat([freqs, freqs], dim=-1)
|
|
|
|
def _apply_trig(func_name: str) -> Tensor:
|
|
def compute(x: te.Tensor):
|
|
return te.compute(
|
|
x.shape,
|
|
lambda *indices: getattr(tirx, func_name)(x[indices]),
|
|
name=f"mrope_{func_name}",
|
|
)
|
|
|
|
return op.tensor_expr_op(compute, f"mrope_{func_name}", [emb])
|
|
|
|
cos = _apply_trig("cos") * self.attention_scaling
|
|
sin = _apply_trig("sin") * self.attention_scaling
|
|
return cos.astype(dtype), sin.astype(dtype)
|
|
|
|
|
|
def apply_multimodal_rotary_pos_emb(
|
|
q: Tensor,
|
|
k: Tensor,
|
|
cos: Tensor,
|
|
sin: Tensor,
|
|
mrope_section: Sequence[int],
|
|
unsqueeze_dim: int = 2,
|
|
) -> Tuple[Tensor, Tensor]: # noqa: UP006
|
|
"""Apply multimodal rotary embedding to query and key tensors."""
|
|
|
|
split_sizes = _repeat_mrope_section(mrope_section)
|
|
reordered_cos = _reorder_cos_sin(cos, split_sizes)
|
|
reordered_sin = _reorder_cos_sin(sin, split_sizes)
|
|
cos_term = op.unsqueeze(reordered_cos, dim=unsqueeze_dim)
|
|
sin_term = op.unsqueeze(reordered_sin, dim=unsqueeze_dim)
|
|
cos_term = cos_term.astype(q.dtype)
|
|
sin_term = sin_term.astype(q.dtype)
|
|
q_embed = op.add(op.multiply(q, cos_term), op.multiply(_rotate_half(q), sin_term))
|
|
k_embed = op.add(op.multiply(k, cos_term), op.multiply(_rotate_half(k), sin_term))
|
|
return q_embed, k_embed
|
|
|
|
|
|
@dataclass
|
|
class VisionPositionMetadata:
|
|
"""Metadata required to build multimodal position IDs."""
|
|
|
|
vision_start_token_id: int
|
|
image_token_id: int
|
|
video_token_id: int
|
|
spatial_merge_size: int
|
|
tokens_per_second: float
|
|
|
|
def merged_hw(self, height: int, width: int) -> Tuple[int, int]: # noqa: UP006
|
|
"""Return merged height/width after applying ``spatial_merge_size``."""
|
|
|
|
if height % self.spatial_merge_size != 0 or width % self.spatial_merge_size != 0:
|
|
raise ValueError(
|
|
"Image or video grid is not divisible by spatial_merge_size "
|
|
f"(got h={height}, w={width}, merge={self.spatial_merge_size})."
|
|
)
|
|
return height // self.spatial_merge_size, width // self.spatial_merge_size
|
|
|
|
|
|
def _text_chunk(length: int, offset: int) -> np.ndarray:
|
|
"""Create a text-position chunk with a shared scalar offset for T/H/W axes."""
|
|
|
|
if length <= 0:
|
|
return np.zeros((3, 0), dtype=np.int64)
|
|
seq: np.ndarray = np.arange(length, dtype=np.int64)
|
|
chunk = np.broadcast_to(seq.reshape(1, -1), (3, length))
|
|
return chunk + offset
|
|
|
|
|
|
def _grid_chunk(
|
|
grid_t: int,
|
|
grid_h: int,
|
|
grid_w: int,
|
|
offset: int,
|
|
tokens_per_second: float,
|
|
second_per_grid: float,
|
|
) -> np.ndarray:
|
|
if grid_t <= 0 or grid_h <= 0 or grid_w <= 0:
|
|
raise ValueError(
|
|
f"Invalid grid shape t={grid_t}, h={grid_h}, w={grid_w} for multimodal positions."
|
|
)
|
|
time_axis = (np.arange(grid_t, dtype=np.float32) * second_per_grid * tokens_per_second).astype(
|
|
np.int64
|
|
)
|
|
t_index = np.repeat(time_axis, grid_h * grid_w)
|
|
h_index = np.tile(np.repeat(np.arange(grid_h, dtype=np.int64), grid_w), grid_t)
|
|
w_index = np.tile(np.tile(np.arange(grid_w, dtype=np.int64), grid_h), grid_t)
|
|
stacked = np.stack([t_index, h_index, w_index])
|
|
return stacked + offset
|
|
|
|
|
|
def _find_token_index(tokens: Sequence[int], token_id: int, start: int) -> int:
|
|
for idx in range(start, len(tokens)):
|
|
if tokens[idx] == token_id:
|
|
return idx
|
|
return len(tokens)
|
|
|
|
|
|
def _next_chunk_offset(chunks: Sequence[np.ndarray]) -> int:
|
|
if not chunks:
|
|
return 0
|
|
return int(chunks[-1].max()) + 1
|
|
|
|
|
|
def _count_vision_items(
|
|
token_array: np.ndarray,
|
|
vision_start_token_id: int,
|
|
image_token_id: int,
|
|
video_token_id: int,
|
|
) -> Tuple[int, int]: # noqa: UP006
|
|
vision_starts = np.where(token_array == vision_start_token_id)[0]
|
|
valid_starts = vision_starts[vision_starts + 1 < token_array.shape[0]]
|
|
following_tokens = token_array[valid_starts + 1]
|
|
image_count = int(np.sum(following_tokens == image_token_id))
|
|
video_count = int(np.sum(following_tokens == video_token_id))
|
|
return image_count, video_count
|
|
|
|
|
|
def _next_vision_block(
|
|
tokens: Sequence[int],
|
|
start: int,
|
|
meta: VisionPositionMetadata,
|
|
has_images: bool,
|
|
has_videos: bool,
|
|
) -> Tuple[str, int]: # noqa: UP006
|
|
sentinel = len(tokens) + 1
|
|
image_end = _find_token_index(tokens, meta.image_token_id, start) if has_images else sentinel
|
|
video_end = _find_token_index(tokens, meta.video_token_id, start) if has_videos else sentinel
|
|
if image_end < video_end:
|
|
return "image", image_end
|
|
return "video", video_end
|
|
|
|
|
|
def _load_grid_for_block(
|
|
block_kind: str,
|
|
image_grid_thw: Optional[np.ndarray], # noqa: UP045
|
|
video_grid_thw: Optional[np.ndarray], # noqa: UP045
|
|
second_per_grid_ts: Optional[np.ndarray], # noqa: UP045
|
|
image_index: int,
|
|
video_index: int,
|
|
) -> Tuple[int, int, int, float, int, int]: # noqa: UP006
|
|
if block_kind == "image":
|
|
if image_grid_thw is None:
|
|
raise ValueError("Image grids are required for sequences with image tokens.")
|
|
grid_t, grid_h, grid_w = image_grid_thw[image_index]
|
|
return int(grid_t), int(grid_h), int(grid_w), 0.0, image_index + 1, video_index
|
|
|
|
if video_grid_thw is None:
|
|
raise ValueError("Video grids are required for sequences with video tokens.")
|
|
grid_t, grid_h, grid_w = video_grid_thw[video_index]
|
|
second_per_grid = (
|
|
float(second_per_grid_ts[video_index]) if second_per_grid_ts is not None else 1.0
|
|
)
|
|
return int(grid_t), int(grid_h), int(grid_w), second_per_grid, image_index, video_index + 1
|
|
|
|
|
|
def _build_sequence_position_ids(
|
|
input_tokens: Sequence[int],
|
|
meta: VisionPositionMetadata,
|
|
image_grid_thw: Optional[np.ndarray], # noqa: UP045
|
|
video_grid_thw: Optional[np.ndarray], # noqa: UP045
|
|
second_per_grid_ts: Optional[np.ndarray], # noqa: UP045
|
|
image_index: int,
|
|
video_index: int,
|
|
) -> Tuple[np.ndarray, int, int, int]: # noqa: UP006
|
|
token_array = np.asarray(input_tokens, dtype=np.int64)
|
|
image_count, video_count = _count_vision_items(
|
|
token_array,
|
|
vision_start_token_id=meta.vision_start_token_id,
|
|
image_token_id=meta.image_token_id,
|
|
video_token_id=meta.video_token_id,
|
|
)
|
|
if image_count > 0 and image_grid_thw is None:
|
|
raise ValueError("Image grids are required for sequences with image tokens.")
|
|
if video_count > 0 and video_grid_thw is None:
|
|
raise ValueError("Video grids are required for sequences with video tokens.")
|
|
|
|
llm_pos_ids_list: List[np.ndarray] = [] # noqa: UP006
|
|
start = 0
|
|
remain_images = image_count
|
|
remain_videos = video_count
|
|
for _ in range(image_count + video_count):
|
|
block_kind, block_end = _next_vision_block(
|
|
tokens=input_tokens,
|
|
start=start,
|
|
meta=meta,
|
|
has_images=remain_images > 0,
|
|
has_videos=remain_videos > 0,
|
|
)
|
|
(
|
|
grid_t,
|
|
grid_h,
|
|
grid_w,
|
|
second_per_grid,
|
|
image_index,
|
|
video_index,
|
|
) = _load_grid_for_block(
|
|
block_kind=block_kind,
|
|
image_grid_thw=image_grid_thw,
|
|
video_grid_thw=video_grid_thw,
|
|
second_per_grid_ts=second_per_grid_ts,
|
|
image_index=image_index,
|
|
video_index=video_index,
|
|
)
|
|
if block_kind == "image":
|
|
remain_images -= 1
|
|
else:
|
|
remain_videos -= 1
|
|
|
|
llm_grid_h, llm_grid_w = meta.merged_hw(grid_h, grid_w)
|
|
text_len = block_end - start
|
|
text_offset = _next_chunk_offset(llm_pos_ids_list)
|
|
llm_pos_ids_list.append(_text_chunk(text_len, text_offset))
|
|
grid_offset = text_offset + text_len
|
|
llm_pos_ids_list.append(
|
|
_grid_chunk(
|
|
grid_t=grid_t,
|
|
grid_h=llm_grid_h,
|
|
grid_w=llm_grid_w,
|
|
offset=grid_offset,
|
|
tokens_per_second=meta.tokens_per_second,
|
|
second_per_grid=second_per_grid,
|
|
)
|
|
)
|
|
start = block_end + grid_t * llm_grid_h * llm_grid_w
|
|
|
|
if start < len(input_tokens):
|
|
tail_len = len(input_tokens) - start
|
|
tail_offset = _next_chunk_offset(llm_pos_ids_list)
|
|
llm_pos_ids_list.append(_text_chunk(tail_len, tail_offset))
|
|
|
|
if not llm_pos_ids_list:
|
|
empty_positions: np.ndarray = np.zeros((3, 0), dtype=np.int64)
|
|
return empty_positions, 0, image_index, video_index
|
|
llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1)
|
|
delta = int(llm_positions.max()) + 1 - len(input_tokens)
|
|
return llm_positions, delta, image_index, video_index
|
|
|
|
|
|
def _text_only_position_ids(
|
|
input_ids: np.ndarray,
|
|
attention_mask: Optional[np.ndarray], # noqa: UP045
|
|
) -> Tuple[np.ndarray, np.ndarray]: # noqa: UP006
|
|
batch, seq_len = input_ids.shape
|
|
if attention_mask is None:
|
|
base: np.ndarray = np.arange(seq_len, dtype=np.int64).reshape(1, 1, -1)
|
|
tiled = np.broadcast_to(base, (3, batch, seq_len))
|
|
return tiled, np.zeros((batch, 1), dtype=np.int64)
|
|
|
|
position = attention_mask.cumsum(axis=-1) - 1
|
|
position = np.where(attention_mask == 0, 1, position)
|
|
position = np.expand_dims(position, axis=0).repeat(3, axis=0)
|
|
max_pos = position.max(axis=0, keepdims=False).max(axis=-1, keepdims=True)
|
|
delta = (max_pos + 1 - seq_len).astype(np.int64)
|
|
return position.astype(np.int64), delta
|
|
|
|
|
|
def get_mrope_position_ids(
|
|
input_ids: np.ndarray,
|
|
meta: VisionPositionMetadata,
|
|
attention_mask: Optional[np.ndarray] = None, # noqa: UP045
|
|
image_grid_thw: Optional[np.ndarray] = None, # noqa: UP045
|
|
video_grid_thw: Optional[np.ndarray] = None, # noqa: UP045
|
|
second_per_grid_ts: Optional[np.ndarray] = None, # noqa: UP045
|
|
) -> Tuple[np.ndarray, np.ndarray]: # noqa: UP006
|
|
"""Generate 3D position IDs and deltas following Hugging Face Qwen2.5-VL."""
|
|
|
|
input_ids = np.asarray(input_ids, dtype=np.int64)
|
|
batch, seq_len = input_ids.shape
|
|
position_ids = np.ones((3, batch, seq_len), dtype=np.int64)
|
|
|
|
attention = None
|
|
if attention_mask is not None:
|
|
attention_mask = np.asarray(attention_mask, dtype=np.int64)
|
|
if attention_mask.shape != input_ids.shape:
|
|
raise ValueError(
|
|
"attention_mask shape must match input_ids shape: "
|
|
f"{attention_mask.shape} vs {input_ids.shape}"
|
|
)
|
|
attention = attention_mask.astype(bool)
|
|
|
|
image_grid_thw = None if image_grid_thw is None else np.asarray(image_grid_thw, dtype=np.int64)
|
|
video_grid_thw = None if video_grid_thw is None else np.asarray(video_grid_thw, dtype=np.int64)
|
|
if second_per_grid_ts is not None:
|
|
second_per_grid_ts = np.asarray(second_per_grid_ts, dtype=np.float32)
|
|
|
|
contains_image_tokens = bool(np.any(input_ids == meta.image_token_id))
|
|
contains_video_tokens = bool(np.any(input_ids == meta.video_token_id))
|
|
if contains_image_tokens and image_grid_thw is None:
|
|
raise ValueError("image_grid_thw must be provided when image tokens exist in input_ids.")
|
|
if contains_video_tokens and video_grid_thw is None:
|
|
raise ValueError("video_grid_thw must be provided when video tokens exist in input_ids.")
|
|
if (
|
|
second_per_grid_ts is not None
|
|
and video_grid_thw is not None
|
|
and second_per_grid_ts.shape[0] != video_grid_thw.shape[0]
|
|
):
|
|
raise ValueError(
|
|
"second_per_grid_ts length must match number of video grids "
|
|
f"({second_per_grid_ts.shape[0]} vs {video_grid_thw.shape[0]})."
|
|
)
|
|
|
|
if not (contains_image_tokens or contains_video_tokens):
|
|
return _text_only_position_ids(input_ids, attention_mask)
|
|
|
|
image_index = 0
|
|
video_index = 0
|
|
deltas: List[int] = [] # noqa: UP006
|
|
|
|
for batch_idx in range(batch):
|
|
tokens = input_ids[batch_idx]
|
|
if attention is not None:
|
|
tokens = tokens[attention[batch_idx]]
|
|
token_values = np.asarray(tokens, dtype=np.int64).tolist()
|
|
input_tokens: List[int] = [int(token) for token in token_values] # noqa: UP006
|
|
if not input_tokens:
|
|
deltas.append(0)
|
|
continue
|
|
|
|
llm_positions, delta, image_index, video_index = _build_sequence_position_ids(
|
|
input_tokens=input_tokens,
|
|
meta=meta,
|
|
image_grid_thw=image_grid_thw,
|
|
video_grid_thw=video_grid_thw,
|
|
second_per_grid_ts=second_per_grid_ts,
|
|
image_index=image_index,
|
|
video_index=video_index,
|
|
)
|
|
if attention is not None:
|
|
position_ids[:, batch_idx, attention[batch_idx]] = llm_positions
|
|
else:
|
|
position_ids[:, batch_idx, :] = llm_positions
|
|
deltas.append(delta)
|
|
|
|
delta_array = np.asarray(deltas, dtype=np.int64).reshape(batch, 1)
|
|
return position_ids, delta_array
|