#023

mrope

bf16 vllm · both · vllm.model_executor.layers.rotary_embedding.mrope · importance 1.0%

Reference Implementation

reference.py

import torch
import torch.nn as nn

class Model(nn.Module):

    def __init__(self, head_size: int, rotary_dim: int, mrope_section: list[int], mrope_interleaved: bool=False, is_neox_style: bool=True) -> None:
        super().__init__()
        self.head_size = int(head_size)
        self.rotary_dim = int(rotary_dim)
        self.mrope_section = [int(x) for x in mrope_section]
        self.mrope_interleaved = bool(mrope_interleaved)
        self.is_neox_style = bool(is_neox_style)

    def forward(self, q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> dict[str, torch.Tensor]:
        return {'q': self._rotate(q, cos, sin), 'k': self._rotate(k, cos, sin)}

    def _rotate(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        tokens = x.shape[0]
        heads = x.shape[1] // self.head_size
        x_view = x.reshape(tokens, heads, self.head_size).clone()
        half = self.rotary_dim // 2
        axis = self._axis_map(half, x.device)
        pos = torch.arange(tokens, device=x.device)[:, None]
        freq = torch.arange(half, device=x.device)[None, :]
        cos_row = cos[axis[None, :], pos, freq]
        sin_row = sin[axis[None, :], pos, freq]
        if self.is_neox_style:
            x1 = x_view[:, :, :half].float()
            x2 = x_view[:, :, half:self.rotary_dim].float()
            c = cos_row[:, None, :].float()
            s = sin_row[:, None, :].float()
            x_view[:, :, :half] = (x1 * c - x2 * s).to(x.dtype)
            x_view[:, :, half:self.rotary_dim] = (x2 * c + x1 * s).to(x.dtype)
        else:
            even = x_view[:, :, :self.rotary_dim:2].float()
            odd = x_view[:, :, 1:self.rotary_dim:2].float()
            c = cos_row[:, None, :].float()
            s = sin_row[:, None, :].float()
            x_view[:, :, :self.rotary_dim:2] = (even * c - odd * s).to(x.dtype)
            x_view[:, :, 1:self.rotary_dim:2] = (odd * c + even * s).to(x.dtype)
        return x_view.reshape_as(x)

    def _axis_map(self, half: int, device: torch.device) -> torch.Tensor:
        offsets = torch.arange(half, device=device)
        if self.mrope_interleaved:
            h = (offsets % 3 == 1) & (offsets <= 3 * self.mrope_section[1])
            w = (offsets % 3 == 2) & (offsets <= 3 * self.mrope_section[2])
            return torch.where(h, 1, torch.where(w, 2, 0)).long()
        t_end = self.mrope_section[0]
        h_end = t_end + self.mrope_section[1]
        return torch.where(offsets < t_end, 0, torch.where(offsets < h_end, 1, 2)).long()

Shapes

T_SOL hardware:

#	token_count	num_query_heads	num_kv_heads	head_size	rotary_dim	T_SOL(XPU-A)	T_Prod	S
0	3592	16	1	128	128	6.94 us	46.00 us	15.1%
1	8373	16	1	128	128	16.18 us	78.30 us	20.7%
2	16530	16	1	128	128	31.94 us	139.00 us	23.0%
3	838	16	4	128	128	1.86 us	17.70 us	10.5%
4	1673	16	4	128	128	3.72 us	28.40 us	13.1%

Input Generation

input.py

import torch

def _make_inputs(token_count: int, num_query_heads: int, num_kv_heads: int, head_size: int, rotary_dim: int) -> dict[str, torch.Tensor]:
    q = torch.randn(token_count, num_query_heads * head_size, dtype=torch.bfloat16, device='cuda') * 0.02
    k = torch.randn(token_count, num_kv_heads * head_size, dtype=torch.bfloat16, device='cuda') * 0.02
    half = rotary_dim // 2
    angles = torch.randn(3, token_count, half, dtype=torch.float32, device='cuda') * 0.1
    return {'q': q, 'k': k, 'cos': torch.cos(angles), 'sin': torch.sin(angles)}