import math from abc import ABC, abstractmethod import torch class CustomRoPE: @staticmethod def pad_position_ids(position_ids: list, padding_mode="right", padding_value=0): """ position_ids shape: [seq_len, dim] """ max_len = max(pos.shape[0] for pos in position_ids) padded_position_ids = [] for pos in position_ids: padding_tensor = torch.zeros((max_len - pos.shape[0], pos.shape[1]), dtype=pos.dtype, device=pos.device) + padding_value if padding_mode == "right": pad_list = [pos, padding_tensor] else: pad_list = [padding_tensor, pos] padded_position_ids.append(torch.cat(pad_list, dim=0)) return torch.stack(padded_position_ids, dim=0) class DataShape(ABC): @abstractmethod def __call__(self, rope_obj: CustomRoPE) -> torch.Tensor: pass class VRoPEVideoShape(DataShape): def __init__(self, t: int, h: int, w: int): self.t = t self.h = h self.w = w def __call__(self, rope_obj: "VRoPE") -> torch.Tensor: return rope_obj.generate_nd_positions((self.t, self.h, self.w)) class VRoPEImageShape(VRoPEVideoShape): def __init__(self, h: int, w: int): super().__init__(1, h, w) class VRoPETextShape(DataShape): def __init__(self, seq_len: int): self.seq_len = seq_len def __call__(self, rope_obj: "VRoPE") -> torch.Tensor: return rope_obj.generate_1d_positions(self.seq_len) class VRoPESharedTextShape(DataShape): # NOTE: VRoPE shared text shape allows generating grouped suffix input def __init__(self, seq_len: int): self.seq_len = seq_len def __call__(self, rope_obj: "VRoPE") -> torch.Tensor: # NOTE: do not increase ``current_index`` position_ids = torch.arange(rope_obj.current_index, rope_obj.current_index + self.seq_len, dtype=torch.long).unsqueeze(-1).expand(-1, rope_obj.half_head_dim) rope_obj.total_token_len += self.seq_len return position_ids class VRoPE(CustomRoPE): """ VRoPE implementation for any-dimensional input. """ def __init__( self, half_head_dim: int, half_time_dim: int, ): assert half_time_dim < half_head_dim, "``half_time_dim`` must be smaller than ``half_head_dim``" # channel allocation config self.half_head_dim = half_head_dim # self.half_time_dim = half_time_dim self.half_time_dim = 0 # total length of the tokens self.total_token_len = 0 # current position index self.current_index = 0 # interleaved data shapes self.data_shapes: list[DataShape] = [] def __call__(self): if not self.data_shapes: # Assume that it's autoregressive mode return self.generate_1d_positions(1) position_ids = torch.cat([data_shape(self) for data_shape in self.data_shapes], dim=0) # Clear after data shapes are computed self.data_shapes.clear() return position_ids def generate_1d_positions(self, seq_len: int): """ Generate 1d positions given seq_len. """ position_ids = torch.arange(self.current_index, self.current_index + seq_len, dtype=torch.long).unsqueeze(-1).expand(-1, self.half_head_dim) self.increase_index_and_total_len(seq_len, seq_len) return position_ids # seq_len, half_head_dim def generate_nd_positions(self, data_shape: tuple): """ Generate nd positions given (time, d1, d2, d3, ..., dn). TODO: Add mixed nd positions (e.g., cls + patch) """ time_shape = data_shape[0] nd_shape = data_shape[1:] n = len(nd_shape) nd_len = math.prod(nd_shape) # number of tokens in nd nd_index_len = sum(nd_shape) - n + 1 # number of indexes used in nd def generate_xd_positions(x): """ Generate positions at xd (x >= 0). """ shape = nd_shape[x] shape_before = math.prod(nd_shape[0:x]) shape_after = math.prod(nd_shape[x + 1 :]) positions = torch.arange(shape, dtype=torch.long).unsqueeze(0).unsqueeze(-1).expand(shape_before, -1, shape_after) # set prod_shape[x] = 2 to apply broadcasting in the following sum operation. prod_shape = [1] * n prod_shape[x] = 2 return torch.stack( [ positions.flatten(), positions.flip(dims=(1,)).flatten(), ], dim=0, ).view(*(prod_shape + [nd_len])) nd_positions = list(map(generate_xd_positions, range(len(nd_shape)))) # Apply symmetric operations. two_pow_nd_positions = sum(nd_positions).reshape(-1, nd_len).unsqueeze(0).expand(time_shape, -1, -1) # Add time shift. # TODO: add time delta time_positions = (torch.arange(time_shape, dtype=torch.long) * nd_index_len).unsqueeze(-1).expand(-1, nd_len) two_pow_nd_positions = two_pow_nd_positions + time_positions.unsqueeze(1) # Make time index center time_positions = time_positions + nd_index_len // 2 # flatten to 1d sequence and add ``current_index`` offset. time_positions = time_positions.flatten() + self.current_index # seq_len two_pow_nd_positions = two_pow_nd_positions.permute(1, 0, 2).flatten(start_dim=1) + self.current_index # 2^n, seq_len # channel allocation # TODO: channel allocation strategies two_pow_nd_position_ids = torch.stack([two_pow_nd_positions[i % two_pow_nd_positions.shape[0]] for i in range(self.half_head_dim - self.half_time_dim)], dim=-1) # seq_len, half_head_dim - half_time_dim time_position_ids = time_positions.unsqueeze(-1).expand(-1, self.half_time_dim) # seq_len, half_time_dim position_ids = torch.cat([two_pow_nd_position_ids, time_position_ids], dim=-1) # seq_len, half_head_dim self.increase_index_and_total_len(nd_index_len * time_shape, nd_len * time_shape) return position_ids # seq_len, half_head_dim def add_shape(self, shape: DataShape) -> None: """ Add a data shape to ``self.data_shapes``. """ self.data_shapes.append(shape) def increase_index_and_total_len(self, index_len: int, token_len: int) -> None: """ The ``current_index`` and ``total_token_len`` should always be consistent, so we need to increase them together. """ self.current_index += int(index_len) self.total_token_len += int(token_len) class GroupedRoPE(CustomRoPE): def __call__(self, group_info: list, padding_mode: str): """ Apply 1d RoPE based on group_info """ position_ids = [] for group in group_info: prefix_len = group[0] grouped_position_ids = [torch.arange(prefix_len)] for g in group[1:]: grouped_position_ids.append(torch.arange(prefix_len, prefix_len + g)) position_ids.append(torch.cat(grouped_position_ids, dim=0).unsqueeze(-1)) return self.pad_position_ids(position_ids=position_ids, padding_mode=padding_mode).squeeze(-1)