Split View: LLM 컨텍스트 윈도우 확장 기술 완벽 가이드: RoPE, ALiBi, YaRN부터 Ring Attention까지

import torch
import torch.nn as nn

class AbsolutePositionalEncoding(nn.Module):
    """Learnable absolute position embeddings"""
    def __init__(self, max_len, d_model):
        super().__init__()
        # Learnable parameter - max_len is fixed
        self.pe = nn.Embedding(max_len, d_model)

    def forward(self, x):
        positions = torch.arange(x.size(1), device=x.device)
        return x + self.pe(positions)

import math

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, max_len, d_model):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

RoPE(q, m) = q * e^{im*theta_i}

import torch

def precompute_freqs_cis(dim, max_seq_len, base=10000.0):
    """Precompute RoPE frequencies"""
    # theta_i = base^(-2i/d)
    freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    # Angle per position: m * theta_i
    t = torch.arange(max_seq_len)
    freqs = torch.outer(t, freqs)  # (max_seq_len, dim/2)
    # Convert to complex: e^{i*m*theta}
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis

def apply_rotary_emb(xq, xk, freqs_cis):
    """Apply RoPE to queries and keys"""
    # Convert real tensors to complex
    xq_complex = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_complex = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

    # Apply rotation
    freqs_cis = freqs_cis[:xq.shape[1]]
    xq_out = torch.view_as_real(xq_complex * freqs_cis).flatten(-2)
    xk_out = torch.view_as_real(xk_complex * freqs_cis).flatten(-2)

    return xq_out.type_as(xq), xk_out.type_as(xk)

# Proof:
# <RoPE(q,m), RoPE(k,n)>
# = <q*e^{im*theta}, k*e^{in*theta}>
# = <q, k> * e^{i(m-n)*theta}  <- depends only on relative position (m-n)!

def position_interpolation(freqs_cis, original_max_len, target_max_len):
    """Extend by linearly interpolating positions"""
    scale = original_max_len / target_max_len
    t = torch.arange(target_max_len) * scale  # Compress positions
    # The rest is the same
    freqs = 1.0 / (10000.0 ** (torch.arange(0, dim, 2).float() / dim))
    freqs = torch.outer(t, freqs)
    return torch.polar(torch.ones_like(freqs), freqs)

# Example: 4K -> 32K extension
# Map positions 0-31999 to the 0-3999 range

def ntk_aware_scaling(dim, max_seq_len, base=10000.0, scale_factor=4):
    """NTK-aware RoPE: Adjust base to preserve high frequencies"""
    # Increase base proportionally to scale_factor
    new_base = base * (scale_factor ** (dim / (dim - 2)))

    freqs = 1.0 / (new_base ** (torch.arange(0, dim, 2).float() / dim))
    t = torch.arange(max_seq_len)
    freqs = torch.outer(t, freqs)
    return torch.polar(torch.ones_like(freqs), freqs)

# 4K -> 16K (scale_factor=4)
# base: 10000 -> ~161489

def ntk_by_parts(dim, max_seq_len, base=10000.0, original_max_len=4096,
                 target_max_len=32768, beta_fast=32, beta_slow=1):
    """Per-dimension differential scaling"""
    scale = target_max_len / original_max_len

    # Original frequencies
    freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))

    # Wavelength of each frequency
    wavelengths = 2 * math.pi / freqs

    # Calculate interpolation ratio (ramp function)
    low_threshold = original_max_len / beta_fast
    high_threshold = original_max_len / beta_slow

    ramp = (wavelengths - low_threshold) / (high_threshold - low_threshold)
    ramp = ramp.clamp(0, 1)

    # High frequency (ramp=0): keep original, Low frequency (ramp=1): full interpolation
    scaled_freqs = freqs / scale  # PI method
    new_freqs = (1 - ramp) * freqs + ramp * scaled_freqs

    t = torch.arange(target_max_len)
    freqs_matrix = torch.outer(t, new_freqs)
    return torch.polar(torch.ones_like(freqs_matrix), freqs_matrix)

def yarn_rope(dim, max_seq_len, base=10000.0, original_max_len=4096,
              target_max_len=131072, beta_fast=32, beta_slow=1):
    """YaRN: NTK-by-Parts + Temperature Scaling"""
    scale = target_max_len / original_max_len

    freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))

    wavelengths = 2 * math.pi / freqs
    low_threshold = original_max_len / beta_fast
    high_threshold = original_max_len / beta_slow

    ramp = (wavelengths - low_threshold) / (high_threshold - low_threshold)
    ramp = ramp.clamp(0, 1)

    scaled_freqs = freqs / scale
    new_freqs = (1 - ramp) * freqs + ramp * scaled_freqs

    # Key: Attention Temperature
    # sqrt(s) correction to preserve attention score distribution
    temperature = 0.1 * math.log(scale) + 1.0

    t = torch.arange(target_max_len)
    freqs_matrix = torch.outer(t, new_freqs)
    freqs_cis = torch.polar(torch.ones_like(freqs_matrix), freqs_matrix)

    return freqs_cis, temperature

# During attention computation:
# attn_weights = (Q @ K.T) / (sqrt(d) * temperature)

def alibi_bias(num_heads, max_seq_len):
    """ALiBi: linear distance-based penalty on attention"""
    # Per-head slopes
    slopes = torch.pow(2, -torch.arange(1, num_heads + 1) * 8.0 / num_heads)

    # Distance matrix
    positions = torch.arange(max_seq_len)
    distances = positions.unsqueeze(0) - positions.unsqueeze(1)  # (L, L)
    distances = distances.abs().neg()  # Negative distances

    # Per-head biases
    biases = slopes.unsqueeze(1).unsqueeze(2) * distances.unsqueeze(0)
    return biases  # (num_heads, L, L)

# Attention computation:
# attn = softmax(Q @ K.T / sqrt(d) + alibi_bias)

# Ring Attention pseudocode
def ring_attention(Q, K, V, num_devices):
    """
    Split the sequence into num_devices chunks for distributed processing.
    Each device computes attention with its own Q chunk and circulating KV chunks.
    """
    seq_len = Q.shape[1]
    chunk_size = seq_len // num_devices

    # Q chunk for each device
    Q_local = Q[:, rank * chunk_size:(rank + 1) * chunk_size]

    # KV circulates in a ring pattern
    K_local = K[:, rank * chunk_size:(rank + 1) * chunk_size]
    V_local = V[:, rank * chunk_size:(rank + 1) * chunk_size]

    output = torch.zeros_like(Q_local)
    log_sum_exp = torch.full_like(Q_local[:, :, :1], float('-inf'))

    for step in range(num_devices):
        # Compute attention with current KV chunk
        attn = Q_local @ K_local.transpose(-2, -1) / math.sqrt(d)
        # Accumulate with online softmax
        output, log_sum_exp = online_softmax_update(
            output, log_sum_exp, attn, V_local
        )
        # Receive KV from the next device (ring communication)
        K_local = ring_send_recv(K_local)
        V_local = ring_send_recv(V_local)

    return output

from transformers import AutoModelForCausalLM, AutoTokenizer

# Qwen2.5-7B (default 32K -> 131K extension)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B",
    torch_dtype=torch.bfloat16,
    device_map="auto",
    rope_scaling={
        "type": "yarn",
        "factor": 4.0,
        "original_max_position_embeddings": 32768,
    }
)

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B")

# Process long documents
long_text = open("long_document.txt").read()
inputs = tokenizer(long_text, return_tensors="pt", truncation=True, max_length=131072)
outputs = model.generate(**inputs, max_new_tokens=512)
print(tokenizer.decode(outputs[0]))