Mã Hóa Vị Trí Xoay (Rotary Position Embedding - RoPE) LLaMA từ Gốc

Mục tiêu kiến thức

• RoPE (Su et al., 2021) là chuẩn mã hóa vị trí thống trị trong các LLM hiện đại (LLaMA-3, Mistral, Qwen, DeepSeek).
• Thay vì cộng thêm positional embedding vào input token, RoPE nhân vector Query và Key với một ma trận quay trực giao:

Với vector 2D [x₁, x₂] ở vị trí thứ m và tần số góc θ: x₁'
x₂' = cos(mθ) & -sin(mθ)
sin(mθ) & cos(mθ) x₁
x₂ = x₁ cos(mθ) - x₂ sin(mθ)
x₁ sin(mθ) + x₂ cos(mθ)

• Tần số góc: θ_i = base^{-2i / d} với i ∈ {0, 1, …, d/2 - 1}, thường base = 10000.0.
• Tích vô hướng giữa Q ở vị trí m và K ở vị trí n:

langle R_{Theta, m} Q, R_{Theta, n} K rangle = Q^T R_{Theta, n - m} K Tích vô hướng tự nhiên chỉ phụ thuộc vào khoảng cách tương đối (m - n), giúp mô hình có khả năng ngoại suy chiều dài ngữ cảnh (length extrapolation) vượt trội.

Yêu cầu

Xây dựng lớp RotaryEmbedding(torch.nn.Module):

class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int, max_seq_len: int = 2048, base: float = 10000.0):
        pass

    def forward(self, x: torch.Tensor, seq_len: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
        # Returns cos, sin cached tensors of shape (seq_len, dim)
        pass

def apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    # x shape: (B, num_heads, seq_len, head_dim)
    pass

Ràng buộc

• Thời gian chạy tối đa: 3000ms.
• Giới hạn bộ nhớ: 512MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.