BlockMask: Efficient Sparse Attention

• BlockMask is a mask-aware algorithm that processes only non-masked block tiles in transformer attention, optimizing sparse computations.
• It uses a preprocessing phase to create a Binary Block Matrix and auxiliary arrays, reducing the effective computational complexity from O(N²) to near O(N·s).
• Empirical tests show up to a 9× speedup in scenarios like ALPACA and Longformer masks, highlighting its efficiency in structured sparse attention.

BlockMask, also known as Binary Block Masking, is a mask-aware algorithmic modification to Flash Attention (FA) that reduces computation for sparse or partially filled attention masks common in modern transformer workloads. BlockMask introduces a lightweight dispatcher that selectively computes attention for only those block tiles containing non-masked entries, thus providing significant speed and efficiency advantages in scenarios where standard Flash Attention expends resources on masked-out regions.

1. Motivation and Problem Formulation

Flash Attention computes $N \times N$ attention matrices in block-wise tiles of size $B_I \times B_J$, regardless of the sparsity of the attention mask $M \in \{0,1\}^{N \times N}$. This results in unnecessary memory bandwidth utilization and arithmetic operations, retaining the standard $\mathcal{O}(N^2)$ complexity even when large portions of the matrix are masked out. A naïve approach of reading every block and applying zeroes introduces additional memory accesses but does not reduce the total number of blocks that require computation. BlockMask addresses this inefficiency by creating a dispatch mechanism that only processes blocks with at least one nonzero (unmasked) entry [(Sharma et al., 2024), Sec. 3, Fig. 1d].

2. Methodology and Data Structures

BlockMask employs a preprocessing phase producing a Binary Block Matrix, denoted $\mathrm{BinBlkMat} \in \{0,1\}^{n_I \times n_J}$, where $n_I = \lceil N / B_I \rceil$ and $n_J = \lceil N / B_J \rceil$. For block $(i,j)$,

$$\mathrm{BinBlkMat}_{i,j} = \mathop{\mathrm{OR}}\bigl(M_{iB_I:(i+1)B_I,\, jB_J:(j+1)B_J}\bigr)$$

Thus, only those tiles where $\mathrm{BinBlkMat}_{i,j} = 1$ are passed to the computation kernel.

For masks with contiguous all-one patterns, two auxiliary arrays of length $B_I \times B_J$0 are used:

• $B_I \times B_J$1 (start index of contiguous all-ones run in row $B_I \times B_J$2)
• $B_I \times B_J$3 (length of that run)

The preprocessing step, which is parallelized over blocks, establishes these for both basic and contiguous-optimized masking. Empirical runtime of this preprocessing kernel is approximately 0.2 ms for $B_I \times B_J$4 on an RTX 3060 [(Sharma et al., 2024), Table A.1].

Dispatch at runtime replaces the traditional dense $B_I \times B_J$5 block loop with an iteration over $B_I \times B_J$6, skipping any with $B_I \times B_J$7. Contiguous optimizations skip mask reads within known all-ones regions, further reducing bandwidth consumption and control divergence.

3. Algorithmic Complexity and Optimization Modes

The computational complexities are compared as follows:

• Baseline FA (dense):

$B_I \times B_J$8

• BlockMask: If $B_I \times B_J$9 denotes the number of nonzero blocks,

$M \in \{0,1\}^{N \times N}$0

If each query attends to at most $M \in \{0,1\}^{N \times N}$1 keys, BlockMask complexity approaches $M \in \{0,1\}^{N \times N}$2.

BlockMask supports two major optimization strategies:

• Contiguous Non-Zero Pattern (Dense BinBlkMsk): For masks with long runs of contiguous all-ones blocks (as in sequence packing), $M \in \{0,1\}^{N \times N}$3 and $M \in \{0,1\}^{N \times N}$4 arrays allow mask reads to be skipped entirely for such runs.
• Extremely Sparse Mask (RCM Preprocessing): When the nonzero blocks are highly localized but occupy many $M \in \{0,1\}^{N \times N}$5 tiles, Reverse Cuthill–McKee (RCM) graph reordering is applied to $M \in \{0,1\}^{N \times N}$6 to minimize the "bandwidth" of the mask. This reordering can turn many previously nonzero blocks into all-zero ones, reducing the total number of blocks processed. Empirical results report up to a 90% reduction in active blocks on synthetic masks [Fig. 2b].

4. Empirical Performance

BlockMask was evaluated across three real-world mask scenarios on an RTX 3060 (bfloat16, batch 4, 32 heads, blocksize $M \in \{0,1\}^{N \times N}$7):

• MEDUSA Tree Masks: $M \in \{0,1\}^{N \times N}$8 shape; BlockMask yields 2–3× speedup over PyTorch-native attention at small sizes, with scalability as decoding heads $M \in \{0,1\}^{N \times N}$9 or tree candidates $\mathcal{O}(N^2)$0 increase. Flash Attention alone does not show gains in this regime [Sec. 5.2, Fig. 3].
• ALPACA Packing Masks: Both sequential and input-bidirectional patterns exhibit long contiguous nonzero runs. BlockMask achieves up to 9× end-to-end speedup in forward and backward passes compared to standard Flash Attention [Sec. 5.3, Fig. 4].
• Longformer Masks: Windowed pattern (contiguous runs) achieves up to 5× gains; dilated and global patterns (less contiguous) yield 2–3× speedups utilizing basic BinBlkMat [Sec. 5.4, Fig. 5].

Memory and bandwidth efficiency is also improved: $\mathcal{O}(N^2)$1 requires $\mathcal{O}(N^2)$2 bytes plus $\mathcal{O}(N^2)$3 32-bit auxiliary entries. Mask-read bandwidth can be reduced by up to 90% in optimal cases compared to dense approaches [Fig. 2].

5. Integration and Implementation Guidance

BlockMask exposes two principal Triton API entry points:

1. preprocess_mask(mask, B_I, B_J) → (BinBlkMat, offset, total_ones)
2. flash_attn_blockmasked(Q, K, V, BinBlkMat, offset, total_ones)

Integration into existing Flash Attention pipelines consists of:

• Replacing the mask-application kernel with a call to preprocess_mask (once per configuration).
• Modifying the standard block computation loop to the BlockMask-aware version, dispatching only relevant thread-blocks.
• In CUDA, thread-block dispatch is limited to $\mathcal{O}(N^2)$4 blocks, reading only the relevant $\mathcal{O}(N^2)$5 chunks.

Because $\mathcal{O}(N^2)$6 and its auxiliary arrays are shared across all heads and layers, preprocessing cost is amortized even in large-scale models. A plausible implication is that the runtime benefit increases with overall model size and mask sparsity.

6. Applications and Context within Transformer Architectures

BlockMask directly benefits use-cases where attention patterns are highly structured or sparse:

• Sequence packing (e.g., ALPACA patterns)
• Tree masking for structured output models such as MEDUSA
• Sparse long-range attention (e.g., Longformer, windowed/dilated/global patterns)

In each domain, BlockMask leverages block-level sparsity for runtime reduction and scalability benefits, enabling the practical use of more efficient masking strategies in large transformer deployments [(Sharma et al., 2024), Sec. 5, Figs. 3–5].

7. Limitations and Further Considerations

BlockMask maintains effectiveness only when mask-induced block-level sparsity is substantial. Highly irregular but dense active block distributions may yield less benefit, though RCM reordering can mitigate this. Memory overhead for $\mathcal{O}(N^2)$7 and auxiliary arrays is small relative to overall activation size, but not negligible for extremely large $\mathcal{O}(N^2)$8. Mask dispatch and preprocessing overheads are minimal but remain as a fixed cost irrespective of mask structure [(Sharma et al., 2024), Sec. 4.4, App. A].