"""MRotaryEmbedding, YaRNScalingMRotaryEmbedding, Ernie4_5_VLRotaryEmbedding, apply_interleaved_rope for multimodal RoPE.""" from __future__ import annotations from typing import List, Optional, Tuple import torch from sglang.srt.layers.rotary_embedding.base import RotaryEmbedding from sglang.srt.layers.rotary_embedding.triton_kernels import ( triton_ernie45_rope_fused_inplace, triton_mrope_fused, ) from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb from sglang.srt.layers.rotary_embedding.yarn import ( yarn_find_correction_range, yarn_get_mscale_simple, yarn_linear_ramp_mask, ) from sglang.srt.runtime_context import get_server_args from sglang.srt.utils import ( cpu_has_amx_support, is_cuda, is_npu, is_xpu, support_triton, ) _is_cuda = is_cuda() _is_npu = is_npu() _is_xpu = is_xpu() _is_cpu_amx_available = cpu_has_amx_support() if _is_cuda: from sglang.jit_kernel.rope import apply_rope_with_cos_sin_cache_inplace if _is_npu: import torch_npu if _is_xpu: from sgl_kernel import multimodal_rotary_embedding import triton import triton.language as tl from sglang.srt.runtime_context import get_server_args @triton.jit def apply_interleaved_rope_kernel( x_ptr, out_ptr, S: tl.constexpr, D: tl.constexpr, stride_x_m, stride_x_s, stride_out_s, section_1_end, section_2_end, BLOCK_S: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): start_s = tl.program_id(0) * BLOCK_S s_offsets = start_s + tl.arange(0, BLOCK_S) dim_offset = tl.program_id(1) * BLOCK_SIZE dim_indices = dim_offset + tl.arange(0, BLOCK_SIZE) mask_s = s_offsets < S mask_d = dim_indices < D mask = mask_s[:, None] & mask_d[None, :] val_ptr = ( x_ptr + 0 * stride_x_m + s_offsets[:, None] * stride_x_s + dim_indices[None, :] ) val = tl.load(val_ptr, mask=mask, other=0.0) cond_a = (dim_indices[None, :] % 3 == 1) & ( dim_indices[None, :] < section_1_end * 3 ) val_a_ptr = ( x_ptr + 1 * stride_x_m + s_offsets[:, None] * stride_x_s + dim_indices[None, :] ) val_a = tl.load(val_a_ptr, mask=mask & cond_a, other=0.0) cond_b = (dim_indices[None, :] % 3 == 2) & ( dim_indices[None, :] < section_2_end * 3 ) val_b_ptr = ( x_ptr + 2 * stride_x_m + s_offsets[:, None] * stride_x_s + dim_indices[None, :] ) val_b = tl.load(val_b_ptr, mask=mask & cond_b, other=0.0) val = tl.where(cond_a, val_a, val) val = tl.where(cond_b, val_b, val) out_ptr = out_ptr + s_offsets[:, None] * stride_out_s + dim_indices[None, :] tl.store(out_ptr, val, mask=mask) def apply_interleaved_rope_triton(x: torch.Tensor, mrope_section: list) -> torch.Tensor: x = x.contiguous() M, S, D = x.shape out = torch.empty((S, D), dtype=x.dtype, device=x.device) BLOCK_S = 64 BLOCK_SIZE = 128 grid = (triton.cdiv(S, BLOCK_S), triton.cdiv(D, BLOCK_SIZE)) section_1_end = mrope_section[1] section_2_end = mrope_section[2] apply_interleaved_rope_kernel[grid]( x, out, S, D, x.stride(0), x.stride(1), out.stride(0), section_1_end, section_2_end, BLOCK_S=BLOCK_S, BLOCK_SIZE=BLOCK_SIZE, ) return out 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 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: Optional[List[int]] = None, mrope_interleaved: bool = False, mrope_interleaved_glm: 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 self.mrope_interleaved_glm = mrope_interleaved_glm if self.mrope_section: expected_sum = rotary_dim // 2 actual_sum = sum(self.mrope_section) if actual_sum != expected_sum: print( f"MRoPE section sum mismatch: expected {expected_sum}, got {actual_sum}. " f"Adjusting mrope_section to match rotary_dim // 2 = {expected_sum}" ) 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 ] current_sum = sum(self.mrope_section) if current_sum != expected_sum: self.mrope_section[-1] += expected_sum - current_sum else: self.mrope_section = [ expected_sum // len(self.mrope_section) ] * len(self.mrope_section) remainder = expected_sum % len(self.mrope_section) for i in range(remainder): self.mrope_section[i] += 1 print( f"Corrected mrope_section: {self.mrope_section} (sum={sum(self.mrope_section)})" ) # MRoPE axis_map interleaving pattern depends on mrope_section sizes. # The algorithm cycles through axes [0(T), 1(H), 2(W)] round-robin, # skipping any axis that has exhausted its allocated pairs. # # For GLM-V (mrope_section=[8,12,12]): # T(8) < H(12) = W(12), so T exhausts first at pair 24. # Result: [0,1,2, 0,1,2, 0,1,2, 0,1,2, 0,1,2, 0,1,2, 0,1,2, 0,1,2, 1,1,2, 1,1,2, 2,2] # After T runs out, only H and W fill the remaining slots. # # For Qwen3-VL (mrope_section=[24,20,20]): # T(24) > H(20) = W(20), so H and W exhaust first near the tail. # Result: [0,1,2, 0,1,2, ...repeated evenly..., 0,1, 0,1, 0,0] # After H/W run out, T fills the remaining slots. if self.mrope_interleaved_glm: num_pairs = rotary_dim // 2 axis_map = torch.empty(num_pairs, dtype=torch.long) assert sum(self.mrope_section) == num_pairs counts = [0, 0, 0] current_ax = 0 for i in range(num_pairs): current_ax = i % 3 while counts[current_ax] >= self.mrope_section[current_ax]: current_ax = (current_ax + 1) % 3 axis_map[i] = current_ax counts[current_ax] += 1 self.register_buffer("axis_map", axis_map, persistent=False) else: self.axis_map = None if get_server_args().rl_on_policy_target is not None: self._forward_method = self.forward_native def get_cos_sin_with_position(self, positions): if positions.ndim == 1: return super().get_cos_sin_with_position(positions) assert positions.ndim == 2 assert self.mrope_section cos_sin = self.cos_sin_cache[positions] last_dim = cos_sin.size()[-1] cos, sin = cos_sin.chunk(2, dim=-1) if self.mrope_interleaved: if support_triton(get_server_args().attention_backend): cos = apply_interleaved_rope_triton(cos, self.mrope_section) sin = apply_interleaved_rope_triton(sin, self.mrope_section) else: 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, ) self.position_cos = cos.repeat(1, 2).view(-1, 1, 1, last_dim).contiguous() self.position_sin = sin.repeat(1, 2).view(-1, 1, 1, last_dim).contiguous() def _match_cos_sin_cache_dtype(self, query: torch.Tensor) -> None: if ( self.cos_sin_cache.device != query.device or self.cos_sin_cache.dtype != query.dtype ): self.cos_sin_cache = self.cos_sin_cache.to(query.device, dtype=query.dtype) def forward_native( self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor, fused_set_kv_buffer_arg=None, ) -> Tuple[torch.Tensor, torch.Tensor]: assert ( fused_set_kv_buffer_arg is None ), "save kv cache is not supported for MRotaryEmbedding." assert positions.ndim == 1 or positions.ndim == 2 cos_sin = self.cos_sin_cache[positions] cos, sin = cos_sin.chunk(2, dim=-1) if positions.ndim == 2: assert self.mrope_section 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, ) seq_len_q = query.shape[0] query_shape = query.shape query = query.view(seq_len_q, -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) seq_len_k = key.shape[0] key_shape = key.shape key = key.view(seq_len_k, -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 forward_cpu( self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor, fused_set_kv_buffer_arg=None, ) -> Tuple[torch.Tensor, torch.Tensor]: if _is_cpu_amx_available: return torch.ops.sgl_kernel.multimodal_rotary_embedding_cpu( positions, query, key, self.head_size, self.cos_sin_cache, self.mrope_section if self.mrope_section else None, self.mrope_interleaved, self.is_neox_style, ) return self.forward_native(positions, query, key, fused_set_kv_buffer_arg) def forward_cuda( self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor, fused_set_kv_buffer_arg=None, ) -> Tuple[torch.Tensor, torch.Tensor]: assert positions.ndim == 1 or positions.ndim == 2 if positions.ndim == 2 and self.mrope_section: return self.forward_triton(positions, query, key) return self.forward_native(positions, query, key, fused_set_kv_buffer_arg) def forward_triton( self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: assert self.mrope_section self._match_cos_sin_cache_dtype(query) triton_mrope_fused( query, key, self.cos_sin_cache, positions, self.mrope_section, self.head_size, self.rotary_dim, self.mrope_interleaved, self.mrope_interleaved_glm, self.is_neox_style, self.axis_map, ) return query, key def forward_npu( self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor, fused_set_kv_buffer_arg=None, ) -> Tuple[torch.Tensor, torch.Tensor]: assert ( fused_set_kv_buffer_arg is None ), "fused_set_kv_buffer_arg is not supported for npu implementation" if query.shape[1] > 4096: return self.forward_native(positions, query, key, fused_set_kv_buffer_arg) rotary_mode = "half" if self.is_neox_style else "interleave" mrope_section = [0, 0, 0] query_out, key_out = torch_npu.npu_mrope( positions, query, key, self.cos_sin_cache, self.head_size, mrope_section=mrope_section, rotary_mode=rotary_mode, ) return query_out, key_out def forward_xpu( self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor, fused_set_kv_buffer_arg=None, ) -> Tuple[torch.Tensor, torch.Tensor]: assert positions.ndim in (1, 2) if positions.ndim == 2 and self.mrope_section: multimodal_rotary_embedding( query, key, self.cos_sin_cache, positions, self.mrope_section, self.head_size, self.rotary_dim, self.mrope_interleaved, self.mrope_interleaved_glm, self.is_neox_style, self.axis_map, ) return query, key return self.forward_native(positions, query, key, fused_set_kv_buffer_arg) @staticmethod def get_rope_index( spatial_merge_size, image_token_id, video_token_id, vision_start_token_id, model_type, tokens_per_second=None, input_ids=None, image_grid_thw=None, video_grid_thw=None, second_per_grid_ts=None, **kwargs, ): from sglang.srt.layers.rotary_embedding.mrope_rope_index import get_rope_index return get_rope_index( spatial_merge_size, image_token_id, video_token_id, vision_start_token_id, model_type, tokens_per_second, input_ids, image_grid_thw, video_grid_thw, second_per_grid_ts, **kwargs, ) @staticmethod def get_rope_index_qwen3_omni( spatial_merge_size, image_token_id, video_token_id, vision_start_token_id, tokens_per_second=None, input_ids=None, image_grid_thw=None, video_grid_thw=None, second_per_grid_ts=None, **kwargs, ): from sglang.srt.layers.rotary_embedding.mrope_rope_index import ( get_rope_index_qwen3_omni, ) return get_rope_index_qwen3_omni( spatial_merge_size, image_token_id, video_token_id, vision_start_token_id, tokens_per_second, input_ids, image_grid_thw, video_grid_thw, second_per_grid_ts, **kwargs, ) @staticmethod def get_rope_index_glm4v( input_ids, hf_config, image_grid_thw, video_grid_thw, attention_mask, **kwargs ): from sglang.srt.layers.rotary_embedding.mrope_rope_index import ( get_rope_index_glm4v, ) return get_rope_index_glm4v( input_ids, hf_config, image_grid_thw, video_grid_thw, attention_mask, **kwargs, ) @staticmethod def get_rope_index_ernie45( input_ids, hf_config, image_grid_thw, video_grid_thw, **kwargs ): from sglang.srt.layers.rotary_embedding.mrope_rope_index import ( get_rope_index_ernie45, ) return get_rope_index_ernie45( input_ids, hf_config, image_grid_thw, video_grid_thw, **kwargs ) class YaRNScalingMRotaryEmbedding(MRotaryEmbedding): """MRoPE-enabled rotary embedding with YaRN context scaling.""" 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, *, mrope_section: Optional[List[int]] = None, mrope_interleaved: bool = False, extrapolation_factor: float = 1, attn_factor: float = 1, beta_fast: int = 32, beta_slow: int = 1, truncate: bool = True, ) -> 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 self.truncate = truncate self.mscale = float(yarn_get_mscale_simple(self.scaling_factor) * attn_factor) super().__init__( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, mrope_section=mrope_section, mrope_interleaved=mrope_interleaved, ) 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, self.truncate, ) 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 Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding): """3D rotary positional embedding. [h w h w h w h w... t t t...]""" def __init__( self, head_size: int, rotary_dim: int, max_position_embeddings: int, base: int, is_neox_style: bool, dtype: torch.dtype, mrope_section: Optional[List[int]] = None, mrope_interleaved: bool = False, ) -> None: super().__init__( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, mrope_section=mrope_section, mrope_interleaved=mrope_interleaved, ) self._apply_rotary_emb_wrapped = torch.compile(dynamic=True)(apply_rotary_emb) def forward_native( self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor = None, ): assert positions.ndim == 1 or positions.ndim == 2 assert key is not None num_tokens = positions.shape[-1] cos_sin = self.cos_sin_cache[positions] cos, sin = cos_sin.chunk(2, dim=-1) if positions.ndim == 2: assert self.mrope_section section_h = self.mrope_section[0] section_w = self.mrope_section[1] section_t = self.mrope_section[2] assert section_h == section_w section_cos_t = cos[..., -section_t:] section_cos_h = cos[..., : section_h + section_w : 2] section_cos_w = cos[..., 1 : section_h + section_w : 2] cos_t, cos_h, cos_w = section_cos_t[0], section_cos_h[1], section_cos_w[2] cos_hw = torch.stack([cos_h, cos_w], dim=-1).reshape( cos_h.shape[:-1] + (cos_h.shape[-1] * 2,) ) cos = torch.cat([cos_hw, cos_t], dim=-1) section_sin_t = sin[..., -section_t:] section_sin_h = sin[..., : section_h + section_w : 2] section_sin_w = sin[..., 1 : section_h + section_w : 2] sin_t, sin_h, sin_w = section_sin_t[0], section_sin_h[1], section_sin_w[2] sin_hw = torch.stack([sin_h, sin_w], dim=-1).reshape( sin_h.shape[:-1] + (sin_h.shape[-1] * 2,) ) sin = torch.cat([sin_hw, sin_t], 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 = self._apply_rotary_emb_wrapped( 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 = self._apply_rotary_emb_wrapped(key_rot, cos, sin, self.is_neox_style) key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) return query, key def forward_cuda( self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor = None, ): assert key is not None assert positions.ndim in (1, 2) self._match_cos_sin_cache_dtype(query) if positions.ndim == 2: assert self.mrope_section is not None triton_ernie45_rope_fused_inplace( q=query, k=key, cos_sin_cache=self.cos_sin_cache, positions=positions, mrope_section=self.mrope_section, head_size=self.head_size, rotary_dim=self.rotary_dim, is_neox_style=self.is_neox_style, ) return query, key if _is_cuda and (apply_rope_with_cos_sin_cache_inplace is not None): apply_rope_with_cos_sin_cache_inplace( positions=positions, query=query, key=key, head_size=self.head_size, cos_sin_cache=self.cos_sin_cache, is_neox=self.is_neox_style, ) return query, key return self.forward_native(positions, query, key) def forward( self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor, fused_set_kv_buffer_arg=None, ) -> Tuple[torch.Tensor, torch.Tensor]: assert positions.ndim == 1 or positions.ndim == 2 return self.forward_cuda(positions, query, key)