rope.py

"""
P0-4: Unified RoPE (Rotary Positional Embedding) Implementation

This module provides a single, consistent RoPE implementation for both
pd_engine_full.py and hybrid_inference.py.

Key Fix: rope_theta should be 1,000,000 for Qwen3, not 10,000
"""

import mlx.core as mx
import torch


# Default RoPE theta for Qwen3 models (from config.json)
DEFAULT_ROPE_THETA = 1000000.0


def mlx_apply_rope(q: mx.array, positions: mx.array, theta: float = None) -> mx.array:
    """Apply Rotary Positional Embedding using MLX.
    
    Args:
        q: Query tensor of shape (batch, num_heads, seq_len, head_dim)
        positions: Position indices, shape (batch, seq_len) or (1, seq_len)
        theta: RoPE base frequency (default: 1000000 for Qwen3)
    
    Returns:
        RoPE-applied tensor of same shape
    """
    if theta is None:
        theta = DEFAULT_ROPE_THETA
    
    batch, num_heads, seq_len, head_dim = q.shape
    
    # Compute frequency matrix
    freqs = 1.0 / (theta ** (mx.arange(0, head_dim, 2)[: head_dim // 2] / head_dim))
    
    # positions: shape (batch, seq_len) or (1, seq_len)
    # Ensure positions is 2D
    if positions.ndim == 1:
        positions = positions.reshape(1, -1)
    
    pos_batch, pos_seq_len = positions.shape
    
    # angles: (pos_batch, pos_seq_len, head_dim//2)
    angles = positions[:, :, None] * freqs[None, None, :]
    
    cos = mx.cos(angles)  # (batch, seq_len, head_dim//2)
    sin = mx.sin(angles)
    
    # Split and rotate
    x1 = q[..., : head_dim // 2]  # (batch, num_heads, seq_len, head_dim//2)
    x2 = q[..., head_dim // 2:]
    
    # Broadcast: cos/sin from (batch, seq, 64) to match (batch, heads, seq, 64)
    # Add head dimension: (batch, seq, 64) -> (batch, 1, seq, 64)
    cos = cos[:, None, :, :]  # (batch, 1, seq_len, head_dim//2)
    sin = sin[:, None, :, :]
    
    result = mx.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1)
    
    return result


def torch_apply_rope(x: torch.Tensor, positions: torch.Tensor, theta: float = None) -> torch.Tensor:
    """Apply Rotary Positional Embedding using PyTorch.
    
    This version is used in convert_to_ane.py for CoreML model conversion.
    
    Args:
        x: Tensor of shape (batch, seq_len, num_kv_heads, head_dim)
        positions: Position indices, shape (batch, seq_len)
        theta: RoPE base frequency (default: 1000000 for Qwen3)
    
    Returns:
        RoPE-applied tensor of same shape
    """
    if theta is None:
        theta = DEFAULT_ROPE_THETA
    
    batch, seq_len, num_kv_heads, head_dim = x.shape
    
    # Compute frequency matrix
    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=x.device).float() / head_dim))
    
    # Create angle matrix: (seq_len, head_dim//2)
    positions = positions.float()
    angles = positions.unsqueeze(-1) * freqs  # (batch, seq_len, head_dim//2)
    
    # Compute cos and sin
    cos = torch.cos(angles)
    sin = torch.sin(angles)
    
    # Split x into two halves
    x1 = x[..., :head_dim // 2]  # (batch, seq_len, num_kv_heads, head_dim//2)
    x2 = x[..., head_dim // 2:]  # (batch, seq_len, num_kv_heads, head_dim//2)
    
    # Apply rotation: x1*cos - x2*sin, x2*cos + x1*sin
    # Need to broadcast cos/sin properly
    cos = cos.unsqueeze(2)  # (batch, seq_len, 1, head_dim//2)
    sin = sin.unsqueeze(2)  # (batch, seq_len, 1, head_dim//2)
    
    new_x1 = x1 * cos - x2 * sin
    new_x2 = x2 * cos + x1 * sin
    
    return torch.cat([new_x1, new_x2], dim=-1)


def get_rope_theta_from_config(model_path: str) -> float:
    """Load rope_theta from model's config.json.
    
    Args:
        model_path: Path to model directory containing config.json
    
    Returns:
        rope_theta value from config (default 1000000 if not found)
    """
    import os
    import json
    
    config_path = os.path.join(model_path, "config.json")
    if os.path.exists(config_path):
        try:
            with open(config_path) as f:
                config = json.load(f)
            return config.get("rope_theta", DEFAULT_ROPE_THETA)
        except Exception:
            pass
    return DEFAULT_ROPE_THETA