FlexPrefill A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference

import torch

def query_aware_index_search(Q, K, gamma, block_size):
    """
    Implements Algorithm 4: Query Aware Index Search.
    
    Args:
        Q: Full resolution Query matrix
        K: Full resolution Key matrix
        gamma: Cumulative attention threshold (e.g., 0.95)
        block_size: Size for pooling (e.g., 128)
    
    Returns:
        S: The set of sparse indices (block coordinates) to compute.
    """
    
    # 1. Compute estimated attention scores using pooled vectors
    # Average pool Q and K to reduce dimensions (block-level representation)
    Q_pooled = avg_pool(Q, kernel_size=block_size) 
    K_pooled = avg_pool(K, kernel_size=block_size)
    
    # Calculate coarse-grained attention map (Equation: softmax(Q_pooled * K_pooled.T / sqrt(d)))
    # This represents the "thumbnail" or low-res attention map
    A_coarse = torch.softmax(
        torch.matmul(Q_pooled, K_pooled.transpose(-2, -1)) / math.sqrt(d), 
        dim=-1
    )
    
    # 2. Flatten and normalize the attention map
    # We treat all (query_block, key_block) pairs as global candidates
    A_flat = A_coarse.flatten()
    
    # Re-normalize so the sum equals 1 (to handle the probability mass correctly)
    A_flat_normalized = A_flat / A_flat.sum()
    
    # 3. Sort attention scores in descending order
    # We want to pick the most important blocks first
    sorted_scores, sorted_indices = torch.sort(A_flat_normalized, descending=True)
    
    # 4. Obtain the minimum computational budget
    # Calculate cumulative sum of the sorted scores
    cumulative_scores = torch.cumsum(sorted_scores, dim=0)
    
    # Find the cut-off index where the cumulative sum first exceeds gamma
    # This dynamically determines how many blocks we need to compute
    cutoff_k = torch.searchsorted(cumulative_scores, gamma).item() + 1
    
    # 5. Get final index set S
    # Select the top-k block indices
    selected_flat_indices = sorted_indices[:cutoff_k]
    
    # Convert flat indices back to 2D (query_block_idx, key_block_idx) coordinates
    # These are the blocks where we will perform full dense attention later
    S = unflatten_indices(selected_flat_indices, shape=A_coarse.shape)
    
    return S

import torch

def vertical_slash_index_search(Q, K, gamma, block_size):
    """
    Implements Algorithm 3: Vertical Slash Index Search.
    
    Args:
        Q: Full resolution Query matrix (only the last block is used effectively)
        K: Full resolution Key matrix
        gamma: Cumulative attention threshold (e.g., 0.95)
        block_size: Size of the representative query block
    
    Returns:
        S: The sparse index set combining Vertical and Slash patterns.
    """
    
    # 1. Probe: Compute exact attention for the representative query subset
    # We select the last block of queries to estimate the pattern
    Q_probe = Q[-block_size:] 
    
    # Compute full attention for this small block
    # A_probe shape: [block_size, seq_len]
    A_probe = torch.softmax(torch.matmul(Q_probe, K.transpose(-2, -1)) / math.sqrt(d), dim=-1)
    
    # 2. Projection: Decouple into Vertical and Slash components
    total_mass = A_probe.sum()
    
    # Vertical: Sum over the query dimension (rows) to get column importance
    # a_vertical shape: [seq_len]
    a_vertical = A_probe.sum(dim=0) / total_mass
    
    # Slash: Sum over diagonals to get relative distance importance
    # We map (i, j) to offset = i - j. 
    # Implementation usually involves skewing the matrix or gathering by offset.
    # a_slash shape: [seq_len] (representing offsets 0 to seq_len)
    a_slash = sum_by_diagonals(A_probe) / total_mass
    
    # 3. Selection: Find top indices/offsets to meet the budget (gamma)
    
    def get_top_k_indices(scores, threshold):
        # Sort scores descending
        sorted_scores, sorted_indices = torch.sort(scores, descending=True)
        # Cumulative sum to find cutoff
        cumsum = torch.cumsum(sorted_scores, dim=0)
        # Find how many elements are needed to cross the threshold
        k = torch.searchsorted(cumsum, threshold).item() + 1
        return sorted_indices[:k]
    
    # Get the critical columns (e.g., [0, 1, 5])
    selected_vertical_indices = get_top_k_indices(a_vertical, gamma)
    
    # Get the critical offsets (e.g., [0, -1, -2])
    selected_slash_offsets = get_top_k_indices(a_slash, gamma)
    
    # 4. Extension: Construct the final sparse mask S
    # S is the union of:
    #   - The selected vertical columns for ALL queries
    #   - The selected diagonal offsets for ALL queries
    S = construct_mask_from_indices_and_offsets(
        selected_vertical_indices, 
        selected_slash_offsets, 
        total_rows=Q.shape[0], 
        total_cols=K.shape[0]
    )
    
    return S