Mặc dù các API cấp cao giúp che giấu đi sự phức tạp khi thực thi Mô hình Ngôn ngữ Lớn (LLM), việc tối ưu hóa độ trễ suy luận và thông lượng VRAM GPU đòi hỏi sự hiểu biết sâu sắc về các thuật toán tối ưu hóa theo phần cứng.

Trong bài viết kỹ nghệ chuyên sâu này, chúng ta phân tích ba đột phá công nghệ cốt lõi đứng sau các mô hình tiên tiến hàng đầu hiện nay: Rotary Position Embeddings (RoPE), FlashAttention-2 Kernel Logic, và Quản lý Bộ nhớ KV-Cache.

1. Công thức Toán học của Rotary Position Embeddings (RoPE)

Các phương pháp nhúng vị trí tuyệt đối (vectơ hình sin cộng dồn) không thể tổng quát hóa khi độ dài chuỗi vượt quá cửa sổ huấn luyện. RoPE mã hóa thông tin vị trí tương đối trực tiếp vào vectơ Query ($Q$) và Key ($K$) thông qua phép quay trong mặt phẳng phức.

Đối với một cặp đặc trưng 2D $(x_1, x_2)$ tại vị trí $m$ với tần số góc $\theta_i = 10000^{-2(i-1)/d}$, phép quay $\mathbf{R}_{\Theta,m}^d$ được xác định bởi:

$$\mathbf{R}_{\Theta, m}^d \mathbf{x}_m = \begin{pmatrix} \cos m\theta_1 & -\sin m\theta_1 & 0 & 0 & \dots \\ \sin m\theta_1 & \cos m\theta_1 & 0 & 0 & \dots \\ 0 & 0 & \cos m\theta_2 & -\sin m\theta_2 & \dots \\ \vdots & \vdots & \vdots & \vdots & \ddots \end{pmatrix} \begin{pmatrix} x_1 \\ x_2 \\ x_3 \\ \vdots \end{pmatrix}$$

Tích vô hướng giữa query và key sau khi quay bảo toàn chính xác khoảng cách tương đối $(m - n)$:

$$\langle \mathbf{R}_{\Theta, m}^d \mathbf{q}_m, \mathbf{R}_{\Theta, n}^d \mathbf{k}_n \rangle = \mathbf{q}_m^T \mathbf{R}_{\Theta, n-m}^d \mathbf{k}_n$$

1.1 Cài đặt PyTorch RoPE Tối ưu

import torch

def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
    """
    Áp dụng Rotary Position Embeddings cho tensor query và key.
    Hình dạng x: [batch, seq_len, num_heads, head_dim]
    """
    x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
    freqs_cis = freqs_cis.unsqueeze(0).unsqueeze(2)
    x_rotated = x_complex * freqs_cis
    return torch.view_as_real(x_rotated).flatten(3).type_as(x)