Sparse Attention Patterns › Block-Sparse Attention

Michael BrenndoerferUpdated June 25, 20257 min read
Language AI HandbookAll chapters
Part 17 of 67Efficient Attention
Chapter 118Sparse Attention Patterns › Block-Sparse Attention
Reading section 5 of 13 · Block-Sparse Attention · 7 minSection 5 of 13 · 7 min

Book contents

Section 5 of 13 · 7 min read

Sparse Attention Patterns › Block-Sparse AttentionOne section · Sections in this chapter5 / 13 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Block-sparse attention groups positions into blocks and defines attention patterns at the block level. This approach is particularly hardware-friendly because modern GPUs and TPUs are optimized for matrix operations on contiguous memory blocks.

To understand why block structure matters, you need to understand something about how GPUs execute operations. A GPU is organized around a hierarchy of processing units: threads, warps (groups of 32 threads), and thread blocks (groups of warps). When you launch a matrix multiplication, the GPU divides the work into tiles that map to these hardware units. The key constraint is that all threads in a warp must execute the same instruction at the same time, a property called lock-step execution. If some threads are doing useful work while others are idle because their data was masked out, you are wasting hardware capacity.

Arbitrary sparse patterns suffer from this problem. If you randomly mask out 90% of an attention matrix, the remaining 10% is scattered unpredictably across the matrix. When you implement this as a GPU kernel, many warps end up executing mixed instructions: some threads are computing real values while others are computing masked-out positions. The theoretical 10x speedup from 90% sparsity does not materialize in practice because the hardware cannot efficiently exploit irregular sparsity.

Block-sparse attention avoids this problem by ensuring that the sparse structure aligns with hardware boundaries. Instead of masking individual elements, you skip entire blocks. A block-skipped operation means the entire warp either executes or does not execute together, achieving full hardware utilization on the blocks that do execute.

Block Structure

The patterns we've seen so far, local and strided, define attention at the level of individual positions. Block-sparse attention takes a different approach: it groups positions into contiguous blocks and defines attention at the block level. This abstraction has a practical motivation. Modern GPUs are designed to process data in aligned, contiguous chunks. By organizing attention into blocks that match hardware execution units, we can achieve theoretical speedups in actual wall-clock time.

Imagine dividing a 1024-token sequence into blocks of size b=64b = 64. This creates 1024/64=161024/64 = 16 blocks. Instead of asking "which of 1024 positions can position ii attend to?", we ask "which of 16 blocks can block BiB_i attend to?" This coarser-grained question has fewer possible answers, and each answer involves a dense 64×6464 \times 64 matrix operation that GPUs handle efficiently.

The key parameters for block-sparse attention are:

  • bb: the block size (number of positions per block)
  • kk: the number of blocks each query block attends to

If we select kk blocks for each query block, the complexity analysis proceeds as follows:

  1. Number of query blocks: The sequence divides into n/bn/b blocks
  2. Attention per query block: Each block attends to kk other blocks
  3. Operations per block pair: Computing attention between two blocks of size bb requires b×b=b2b \times b = b^2 score computations

Multiplying these together:

Total operations=nb×k×b2=n⋅k⋅b\text{Total operations} = \frac{n}{b} \times k \times b^2 = n \cdot k \cdot b

When kk and bb are constants independent of nn, this expression is linear in sequence length. The quadratic term has vanished. Even better, each b×bb \times b block attention can use highly optimized dense matrix multiplication routines, achieving near-peak GPU utilization.

In[18]:
Code
def create_block_sparse_mask(seq_len, block_size, pattern="diagonal"):
    """
    Create a block-sparse attention mask.

    Args:
        seq_len: Length of sequence
        block_size: Size of each block
        pattern: Which block pairs to attend to
            - "diagonal": Each block attends to itself only
            - "tridiagonal": Each block attends to itself and adjacent blocks

    Returns:
        mask: (seq_len, seq_len) boolean mask
    """
    num_blocks = (seq_len + block_size - 1) // block_size
    mask = np.zeros((seq_len, seq_len), dtype=bool)

    for i in range(num_blocks):
        i_start = i * block_size
        i_end = min((i + 1) * block_size, seq_len)

        if pattern == "diagonal":
            # Attend only to same block
            blocks_to_attend = [i]
        elif pattern == "tridiagonal":
            # Attend to previous, current, and next block
            blocks_to_attend = [
                j for j in [i - 1, i, i + 1] if 0 <= j < num_blocks
            ]
        else:
            raise ValueError(f"Unknown pattern: {pattern}")

        for j in blocks_to_attend:
            j_start = j * block_size
            j_end = min((j + 1) * block_size, seq_len)
            mask[i_start:i_end, j_start:j_end] = True

    return mask


seq_len = 32
block_size = 4
diagonal_mask = create_block_sparse_mask(seq_len, block_size, "diagonal")
tridiagonal_mask = create_block_sparse_mask(seq_len, block_size, "tridiagonal")
Out[19]:
Visualization
Checkerboard pattern of square blocks along the diagonal.
Diagonal block-sparse: Each block attends only to itself. Creates isolated groups with no cross-block communication.
Wider band of square blocks along the diagonal extending one block in each direction.
Tridiagonal block-sparse: Each block attends to itself and neighbors. Allows local information flow across block boundaries.

Block-sparse attention provides a hardware-efficient abstraction for implementation. The regular block structure maps directly to GPU thread blocks, enabling efficient parallel execution with minimal memory overhead.

Hardware Efficiency

The advantage of block-sparse attention goes beyond theoretical complexity. Modern GPUs achieve peak performance when operating on aligned, contiguous memory blocks. Random sparse patterns, while mathematically equivalent, suffer from irregular memory access patterns that underutilize hardware.

In[20]:
Code
def estimate_memory_efficiency(seq_len, block_size, sparsity_ratio):
    """
    Estimate memory efficiency of block-sparse vs random sparse attention.

    Block-sparse can use dense matmul on selected blocks.
    Random sparse requires custom kernels with lower efficiency.
    """
    total_pairs = seq_len * seq_len
    sparse_pairs = int(total_pairs * (1 - sparsity_ratio))

    # Block-sparse: compute full dense blocks
    num_blocks = seq_len // block_size
    blocks_per_row = int(num_blocks * (1 - sparsity_ratio))
    block_sparse_ops = num_blocks * blocks_per_row * (block_size**2)

    # Memory coalescing factor (rough estimate)
    # Block-sparse achieves near 100% efficiency, random ~30%
    block_efficiency = 0.95
    random_efficiency = 0.30

    return {
        "block_sparse_pairs": block_sparse_ops,
        "random_sparse_pairs": sparse_pairs,
        "block_effective_throughput": block_sparse_ops * block_efficiency,
        "random_effective_throughput": sparse_pairs * random_efficiency,
    }


results = estimate_memory_efficiency(1024, 64, 0.9)
Out[21]:
Console
Memory efficiency comparison (90% sparsity, 1024 tokens, block size 64):

Block-sparse computed pairs:        65,536
Random-sparse computed pairs:      104,857

Effective throughput (accounting for memory coalescing):
Block-sparse:        62,259 (95% memory efficiency)
Random-sparse:       31,457 (30% memory efficiency)

Block-sparse advantage: 2.0x effective throughput

Block-sparse attention can be 3x or more efficient than random sparse patterns at the same sparsity level. This hardware awareness matters for practical implementations and explains why production models favor structured sparsity.

The choice of block size is a real engineering decision with trade-offs. Smaller blocks give you finer-grained control over which pairs attend to each other, but they reduce the size of each dense matrix multiplication, making it harder to achieve peak GPU utilization. Larger blocks maximize GPU efficiency but coarsen the sparsity pattern: you either attend to an entire block of tokens or none of them, which can include irrelevant tokens or exclude relevant ones at the block boundary. A block size of 64 threads is a common practical choice because it aligns with GPU warp sizes and provides a good balance between granularity and throughput.

Block-sparse attention does not inherently specify which blocks to attend to. The diagonal and tridiagonal patterns we implement here are just two examples. In practice, you can combine block-sparse structure with any higher-level pattern: local attention at the block level (attend to the current block and its neighbors), strided attention at the block level (attend to every kk-th block), or even content-adaptive selection where the blocks to attend to are determined dynamically based on query-key similarity scores computed at a coarser granularity. The Routing Transformer and Reformer use variations of this last idea, selecting blocks based on locality-sensitive hashing rather than fixed structural patterns.