# Naive attention (what PyTorch does under the hood)
def standard_attention(Q, K, V):
    """
    Q, K, V: [batch, heads, seq_len, head_dim]
    Returns: [batch, heads, seq_len, head_dim]
    """
    d_k = Q.size(-1)
    
    # Step 1: Compute attention scores
    # O(n²d) operations, stores n² matrix in HBM
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    # scores: [batch, heads, seq_len, seq_len] ← MEMORY BOTTLENECK
    
    # Step 2: Softmax normalization
    attn_weights = F.softmax(scores, dim=-1)
    # Still [batch, heads, seq_len, seq_len] in memory
    
    # Step 3: Weighted sum of values
    output = torch.matmul(attn_weights, V)
    
    return output

Standard Attention:
  [Compute all n² scores] → [Softmax] → [Multiply by V]
  Memory: O(n²)
 
FlashAttention:
  For each block of Q:
    For each block of K, V:
      Compute partial attention in SRAM
      Update running statistics
  Memory: O(n) ← only stores final output!

def flash_attention_concept(Q, K, V, block_size=256):
    """
    Conceptual implementation (real version uses CUDA kernels)
    Key idea: Process attention in blocks, never materialize full n×n matrix
    """
    seq_len = Q.size(2)
    output = torch.zeros_like(Q)
    
    # Running max for numerical stability (prevents overflow)
    row_max = torch.full((Q.size(0), Q.size(1), seq_len), float('-inf'))
    row_sum = torch.zeros((Q.size(0), Q.size(1), seq_len))
    
    # Tile over queries
    for q_start in range(0, seq_len, block_size):
        q_end = min(q_start + block_size, seq_len)
        Q_block = Q[:, :, q_start:q_end, :]  # Load to SRAM
        
        # Tile over keys/values
        for kv_start in range(0, seq_len, block_size):
            kv_end = min(kv_start + block_size, kv_start)
            K_block = K[:, :, kv_start:kv_end, :]  # Load to SRAM
            V_block = V[:, :, kv_start:kv_end, :]  # Load to SRAM
            
            # Compute attention scores for this block (in SRAM!)
            scores = torch.matmul(Q_block, K_block.transpose(-2, -1))
            scores = scores / math.sqrt(Q.size(-1))
            
            # Online softmax (avoid storing full scores)
            # Update running max and sum for numerical stability
            block_max = scores.max(dim=-1, keepdim=True).values
            new_max = torch.maximum(row_max[:, :, q_start:q_end], block_max)
            
            # Rescale previous output and add new contribution
            exp_scores = torch.exp(scores - new_max)
            block_sum = exp_scores.sum(dim=-1, keepdim=True)
            
            # Update output with this block's contribution
            output[:, :, q_start:q_end] = (
                output[:, :, q_start:q_end] * torch.exp(row_max[:, :, q_start:q_end] - new_max) +
                torch.matmul(exp_scores, V_block)
            ) / (row_sum[:, :, q_start:q_end] * torch.exp(row_max[:, :, q_start:q_end] - new_max) + block_sum)
            
            row_max[:, :, q_start:q_end] = new_max
            row_sum[:, :, q_start:q_end] = row_sum[:, :, q_start:q_end] * torch.exp(row_max[:, :, q_start:q_end] - new_max) + block_sum
    
    return output

import torch
from torch.nn.functional import scaled_dot_product_attention
# Modern PyTorch (2.0+) has Flash Attention built-in!
def efficient_attention(query, key, value, mask=None):
    """
    Uses FlashAttention automatically if:
    1. CUDA is available
    2. Data types are FP16 or BF16
    3. No custom attention mask (or simple causal mask)
    """
    # PyTorch automatically dispatches to FlashAttention
    output = scaled_dot_product_attention(
        query, key, value,
        attn_mask=mask,
        dropout_p=0.0,
        is_causal=True  # Enables causal masking efficiently
    )
    return output
# Example usage
batch, heads, seq_len, head_dim = 4, 12, 4096, 64
Q = torch.randn(batch, heads, seq_len, head_dim, device='cuda', dtype=torch.float16)
K = torch.randn(batch, heads, seq_len, head_dim, device='cuda', dtype=torch.float16)
V = torch.randn(batch, heads, seq_len, head_dim, device='cuda', dtype=torch.float16)
# This uses FlashAttention under the hood
output = efficient_attention(Q, K, V)
print(f"Output shape: {output.shape}")  # [4, 12, 4096, 64]

For dimension pair (d₁, d₂), position m:
[x_d₁]   [cos(mθ) -sin(mθ)] [x_d₁]
[x_d₂] = [sin(mθ)  cos(mθ)] [x_d₂]

Q_m · K_n = f(Q, K, m - n)

import torch
import torch.nn as nn
class RotaryPositionEmbedding(nn.Module):
    """
    Rotary Position Embedding (RoPE) implementation
    Used in LLaMA, Mistral, and most modern LLMs
    """
    def __init__(self, dim, max_seq_len=8192, base=10000):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.base = base
        
        # Precompute theta values for each dimension pair
        # θᵢ = base^(-2i/d) where i ∈ [0, d/2)
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)
        
        # Precompute rotation matrices for all positions
        self._precompute_freqs(max_seq_len)
    
    def _precompute_freqs(self, seq_len):
        """Precompute cos and sin for all positions"""
        # Positions: [0, 1, 2, ..., seq_len-1]
        t = torch.arange(seq_len, dtype=self.inv_freq.dtype)
        
        # Outer product: [seq_len, dim/2]
        freqs = torch.outer(t, self.inv_freq)
        
        # Combine to get [seq_len, dim]
        emb = torch.cat([freqs, freqs], dim=-1)
        
        self.register_buffer('cos_cached', emb.cos())
        self.register_buffer('sin_cached', emb.sin())
    
    def rotate_half(self, x):
        """
        Rotate half the dimensions
        [x1, x2, x3, x4] → [-x3, -x4, x1, x2]
        """
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat([-x2, x1], dim=-1)
    
    def forward(self, q, k, seq_len=None):
        """
        Apply rotary embeddings to queries and keys
        Args:
            q: [batch, heads, seq_len, head_dim]
            k: [batch, heads, seq_len, head_dim]
        Returns:
            q_rotated, k_rotated (same shape)
        """
        if seq_len is None:
            seq_len = q.size(2)
        
        # Get cached cos/sin for this sequence length
        cos = self.cos_cached[:seq_len, :].unsqueeze(0).unsqueeze(0)
        sin = self.sin_cached[:seq_len, :].unsqueeze(0).unsqueeze(0)
        
        # Apply rotation
        # q_rotated = q * cos + rotate_half(q) * sin
        q_rotated = q * cos + self.rotate_half(q) * sin
        k_rotated = k * cos + self.rotate_half(k) * sin
        
        return q_rotated, k_rotated
# Usage example
rope = RotaryPositionEmbedding(dim=64, max_seq_len=8192)
# Sample query and key tensors
batch, heads, seq_len, head_dim = 2, 8, 512, 64
q = torch.randn(batch, heads, seq_len, head_dim)
k = torch.randn(batch, heads, seq_len, head_dim)
# Apply RoPE
q_rotated, k_rotated = rope(q, k)
print(f"Rotated Q shape: {q_rotated.shape}")  # [2, 8, 512, 64]

Attention_scores[i, j] = Q[i] · K[j] - m × |i - j|

def get_alibi_slopes(num_heads):
    """
    Compute ALiBi slopes for each attention head
    Uses geometric sequence: 2^(-8/n), 2^(-16/n), ..., 2^(-8)
    """
    def get_slopes_power_of_2(n):
        start = 2 ** (-8)
        ratio = start
        return [start * (ratio ** i) for i in range(n)]
    
    if (num_heads & (num_heads - 1)) == 0:  # Power of 2
        return get_slopes_power_of_2(num_heads)
    else:  # Not power of 2: interpolate
        closest_power = 2 ** math.floor(math.log2(num_heads))
        return (
            get_slopes_power_of_2(closest_power) +
            get_alibi_slopes(2 * closest_power)[0::2][:num_heads - closest_power]
        )
def apply_alibi_bias(attention_scores, num_heads):
    """
    Apply ALiBi bias to attention scores
    Args:
        attention_scores: [batch, heads, seq_len, seq_len]
    Returns:
        Biased scores (same shape)
    """
    seq_len = attention_scores.size(-1)
    
    # Create position distance matrix
    # distances[i, j] = |i - j|
    positions = torch.arange(seq_len, device=attention_scores.device)
    distances = (positions.unsqueeze(0) - positions.unsqueeze(1)).abs()
    
    # Get slopes for each head
    slopes = torch.tensor(
        get_alibi_slopes(num_heads),
        device=attention_scores.device
    ).view(1, num_heads, 1, 1)
    
    # Compute bias: -slope × distance
    bias = -slopes * distances.unsqueeze(0).unsqueeze(0)
    
    # Add bias to scores
    return attention_scores + bias
# Example usage
batch, heads, seq_len = 2, 12, 512
scores = torch.randn(batch, heads, seq_len, seq_len)
# Apply ALiBi
biased_scores = apply_alibi_bias(scores, num_heads=heads)
print(f"Biased scores shape: {biased_scores.shape}")  # [2, 12, 512, 512]

# Inefficient autoregressive generation
def generate_naive(model, prompt_ids, max_new_tokens=100):
    """
    Recomputes Q, K, V for ALL tokens at EVERY step
    For 100 tokens: ~5000 redundant computations!
    """
    input_ids = prompt_ids.clone()
    
    for _ in range(max_new_tokens):
        # Process entire sequence every time (wasteful!)
        logits = model(input_ids)  # Recomputes K, V for all past tokens
        next_token = logits[:, -1, :].argmax(dim=-1)
        input_ids = torch.cat([input_ids, next_token.unsqueeze(-1)], dim=-1)
    
    return input_ids

def generate_with_kv_cache(model, prompt_ids, max_new_tokens=100):
    """
    Efficient generation with KV-cache
    Reduces computation by ~30-50× for long sequences
    """
    input_ids = prompt_ids.clone()
    kv_cache = None  # Will store past K, V tensors
    
    for i in range(max_new_tokens):
        if i == 0:
            # First step: process full prompt
            logits, kv_cache = model(input_ids, use_cache=True, past_kv=None)
        else:
            # Subsequent steps: only process new token
            logits, kv_cache = model(
                input_ids[:, -1:],  # Only last token!
                use_cache=True,
                past_kv=kv_cache  # Reuse cached K, V
            )
        
        next_token = logits[:, -1, :].argmax(dim=-1)
        input_ids = torch.cat([input_ids, next_token.unsqueeze(-1)], dim=-1)
    
    return input_ids

class TransformerLayerWithCache(nn.Module):
    """
    Transformer layer with KV-caching support
    """
    def __init__(self, d_model=768, num_heads=12):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)
    
    def forward(self, x, past_kv=None):
        """
        Args:
            x: [batch, seq_len, d_model] (seq_len=1 when using cache)
            past_kv: Tuple of (past_key, past_value) or None
                     past_key: [batch, num_heads, past_seq_len, head_dim]
        Returns:
            output: [batch, seq_len, d_model]
            new_kv: Updated (key, value) cache
        """
        batch, seq_len, d_model = x.shape
        
        # Compute Q, K, V for current token(s)
        Q = self.q_proj(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.k_proj(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.v_proj(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        
        # If we have past KV, concatenate with current
        if past_kv is not None:
            past_key, past_value = past_kv
            K = torch.cat([past_key, K], dim=2)  # Concatenate along sequence dimension
            V = torch.cat([past_value, V], dim=2)
        
        # Store new KV cache (includes past + current)
        new_kv = (K, V)
        
        # Compute attention with full key/value sequence
        # Q: [batch, heads, 1, head_dim] (current token)
        # K, V: [batch, heads, total_seq_len, head_dim] (all past + current)
        attn_output = F.scaled_dot_product_attention(Q, K, V, is_causal=False)
        
        # Reshape and project output
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch, seq_len, d_model)
        output = self.out_proj(attn_output)
        
        return output, new_kv
# Example: Generate with KV-cache
layer = TransformerLayerWithCache(d_model=768, num_heads=12)
# Step 1: Process prompt (seq_len=10)
prompt = torch.randn(1, 10, 768)
output1, kv_cache = layer(prompt, past_kv=None)
print(f"Step 1 - Output: {output1.shape}, Cache K shape: {kv_cache[0].shape}")
# Output: [1, 10, 768], Cache K: [1, 12, 10, 64]
# Step 2: Generate next token (seq_len=1, reuse cache)
new_token = torch.randn(1, 1, 768)
output2, kv_cache = layer(new_token, past_kv=kv_cache)
print(f"Step 2 - Output: {output2.shape}, Cache K shape: {kv_cache[0].shape}")
# Output: [1, 1, 768], Cache K: [1, 12, 11, 64] ← cache grew by 1

Memory = 2 × L × b × h × n × d_h × 2 bytes
         ↑   ↑   ↑   ↑   ↑   ↑     ↑
         K+V layers batch heads seq head_dim FP16

def extend_rope_context(rope_module, original_max_len, new_max_len):
    """
    Extend RoPE context by interpolating position indices
    Used in Llama 2 Long, Code Llama 100K
    """
    scale = original_max_len / new_max_len  # e.g., 2048 / 8192 = 0.25
    
    # Adjust frequency computation
    # Original: θᵢ = base^(-2i/d)
    # Extended: θᵢ = base^(-2i/d) × scale
    rope_module.inv_freq = rope_module.inv_freq * scale
    
    # Recompute cached cos/sin with new frequencies
    rope_module._precompute_freqs(new_max_len)
    
    return rope_module
# Example: Extend from 2K to 8K context
rope = RotaryPositionEmbedding(dim=64, max_seq_len=2048)
rope_extended = extend_rope_context(rope, original_max_len=2048, new_max_len=8192)
# Now can handle 8K sequences without retraining!
q = torch.randn(1, 8, 8192, 64)  # 8K sequence
k = torch.randn(1, 8, 8192, 64)
q_rot, k_rot = rope_extended(q, k)
print(f"Extended to {q_rot.size(2)} tokens")  # 8192

def yarn_scaling(original_length, new_length, dim, base=10000):
    """
    YaRN scaling: Non-uniform interpolation
    - Low frequencies (long-range): Interpolate aggressively
    - High frequencies (short-range): Minimal interpolation
    """
    scale = new_length / original_length
    
    # Compute per-dimension scaling factors
    dim_range = torch.arange(0, dim, 2).float()
    freqs = base ** (-dim_range / dim)
    
    # Lower frequencies get more scaling (they capture long-range)
    # Higher frequencies get less scaling (they capture local patterns)
    yarn_scale = torch.where(
        freqs < 0.1,  # Low frequency threshold
        scale,  # Full interpolation
        1.0 + (scale - 1.0) * (freqs / 0.1)  # Gradual interpolation
    )
    
    return yarn_scale

# Experiment: Where does the model look?
def test_context_recall(model, tokenizer, needle_position='middle'):
    """
    Test if model can find info at different positions
    "Needle in haystack" benchmark
    """
    context = "Random text... " * 1000  # Long distractor
    needle = "The secret password is BANANA."
    
    if needle_position == 'start':
        full_text = needle + context
    elif needle_position == 'middle':
        mid = len(context) // 2
        full_text = context[:mid] + needle + context[mid:]
    else:  # end
        full_text = context + needle
    
    # Ask model to recall the password
    prompt = full_text + "
What is the secret password?"
    response = model.generate(tokenizer.encode(prompt))
    
    return "BANANA" in response
# Results (GPT-3.5 on 16K context):
# Start: 95% recall ✅
# Middle: 62% recall ⚠️
# End: 98% recall ✅

def diagnose_attention_collapse(attention_weights):
    """
    Check if attention has collapsed to uniform distribution
    Args:
        attention_weights: [batch, heads, seq_len, seq_len]
    Returns:
        entropy: Higher = more diffuse attention
    """
    # Compute entropy of attention distribution
    # Uniform distribution has max entropy
    eps = 1e-10
    entropy = -(attention_weights * torch.log(attention_weights + eps)).sum(dim=-1)
    max_entropy = math.log(attention_weights.size(-1))
    
    # Ratio: 1.0 = completely uniform (collapsed)
    entropy_ratio = entropy / max_entropy
    
    return entropy_ratio.mean().item()
# Example
attn = torch.softmax(torch.randn(1, 12, 4096, 4096), dim=-1)
collapse_score = diagnose_attention_collapse(attn)
print(f"Attention collapse score: {collapse_score:.3f}")
# > 0.9 indicates potential collapse

def streaming_kv_cache(kv_cache, max_cache_len=2048, keep_first=128):
    """
    Maintain fixed-size KV-cache for infinite generation
    Keep first N tokens (attention sinks) + recent tokens
    """
    key_cache, value_cache = kv_cache
    current_len = key_cache.size(2)
    
    if current_len <= max_cache_len:
        return kv_cache  # No eviction needed
    
    # Keep first keep_first tokens + most recent
    recent_tokens = max_cache_len - keep_first
    
    key_cache = torch.cat([
        key_cache[:, :, :keep_first, :],  # First N tokens
        key_cache[:, :, -recent_tokens:, :]  # Recent tokens
    ], dim=2)
    
    value_cache = torch.cat([
        value_cache[:, :, :keep_first, :],
        value_cache[:, :, -recent_tokens:, :]
    ], dim=2)
    
    return (key_cache, value_cache)