Sparse Attention Patterns › Implementing Sparse Attention

def sparse_attention(query, key, value, mask, scale=None):
    """
    Compute sparse attention given a boolean mask.

    Args:
        query: (seq_len, d_k) query vectors
        key: (seq_len, d_k) key vectors
        value: (seq_len, d_v) value vectors
        mask: (seq_len, seq_len) boolean mask (True = attend, False = block)
        scale: Optional scaling factor (default: 1/sqrt(d_k))

    Returns:
        output: (seq_len, d_v) attention output
        weights: (seq_len, seq_len) attention weights (masked)
    """
    seq_len, d_k = query.shape

    # Step 0: Determine scaling factor
    # The 1/sqrt(d_k) scaling prevents attention scores from growing
    # too large as dimension increases, which would push softmax into
    # saturation regions with near-zero gradients
    if scale is None:
        scale = 1.0 / np.sqrt(d_k)

    # Step 1: Compute raw attention scores
    # Each score[i,j] measures similarity between query[i] and key[j]
    scores = query @ key.T * scale  # (seq_len, seq_len)

    # Step 2: Apply the sparse mask
    # Where mask is True, keep the score unchanged
    # Where mask is False, replace with a very large negative number
    # This ensures exp(-1e9) ≈ 0 after softmax
    masked_scores = np.where(mask, scores, -1e9)

    # Step 3: Softmax normalization
    # We subtract the max for numerical stability (prevents overflow)
    # This doesn't change the result since softmax is shift-invariant
    exp_scores = np.exp(
        masked_scores - masked_scores.max(axis=1, keepdims=True)
    )
    weights = exp_scores / exp_scores.sum(axis=1, keepdims=True)

    # Explicitly zero blocked positions
    # They're already near-zero from softmax, but this ensures exactness
    weights = np.where(mask, weights, 0.0)

    # Step 4: Weighted sum of values
    # Each output[i] is a weighted combination of all value vectors,
    # with weights determined by attention
    output = weights @ value

    return output, weights

# Test with random inputs
seq_len = 16
d_model = 32

query = np.random.randn(seq_len, d_model)
key = np.random.randn(seq_len, d_model)
value = np.random.randn(seq_len, d_model)

# Compare full vs sparse attention
full_mask = np.ones((seq_len, seq_len), dtype=bool)
sparse_mask = create_combined_sparse_mask(seq_len, local_window=2, stride=4)[0]

full_output, full_weights = sparse_attention(query, key, value, full_mask)
sparse_output, sparse_weights = sparse_attention(query, key, value, sparse_mask)

import time


def benchmark_attention(seq_lens, patterns, num_trials=5):
    """
    Benchmark different attention patterns across sequence lengths.
    """
    results = []
    d_model = 64

    for n in seq_lens:
        query = np.random.randn(n, d_model)
        key = np.random.randn(n, d_model)
        value = np.random.randn(n, d_model)

        for pattern_name, create_mask_fn in patterns.items():
            mask = create_mask_fn(n)

            # Warm up
            _ = sparse_attention(query, key, value, mask)

            # Time multiple trials
            times = []
            for _ in range(num_trials):
                start = time.perf_counter()
                _ = sparse_attention(query, key, value, mask)
                times.append(time.perf_counter() - start)

            avg_time = np.mean(times)
            pairs = mask.sum()

            results.append(
                {
                    "seq_len": n,
                    "pattern": pattern_name,
                    "time_ms": avg_time * 1000,
                    "pairs": pairs,
                    "sparsity": 1 - pairs / (n * n),
                }
            )

    return results


# Define patterns
patterns = {
    "full": lambda n: np.ones((n, n), dtype=bool),
    "local_16": lambda n: create_local_attention_mask(n, 8),
    "sparse": lambda n: create_combined_sparse_mask(n, 4, 8)[0],
}

# Benchmark
seq_lens = [64, 128, 256, 512]
benchmark_results = benchmark_attention(seq_lens, patterns)