LLM·

FlashMask-3 sends each GPU its mask without sending any mask — Four-vector sparse attention masks — What does it mean?

The news. On October 13, 2026, Guoxia Wang and eight co-authors submitted FlashMask-3: Efficient and Expressive Mask-Aware Distributed Attention. On one H100 across 12 masks at 4K–128K context, the paper reports 45–141% more achieved TFLOPs than FlashMask, 34–71% over FlexAttention and 2–44% over MagiAttention. With context parallelism at degrees 4–32, it reports 1.80×–3.71× the throughput of Megatron-Core's ring CP. For a 118B-A14B mixture-of-experts model at 128K on 128 H100 GPUs, it reports 7.3–11.0% higher end-to-end throughput. Read the paper →

Picture a city where every avenue has road closures. A mask says, for each key (an avenue), which query rows (street numbers) are closed to it. The obvious way to store that is a full map with one open-or-closed square for every street-and-avenue pair, which grows as n² with the sequence length. FlashMask, the method this paper builds on, found a shorter way: for structured masks, each avenue has at most two closed stretches, one below the diagonal and one above it, so four numbers per avenue are enough. Those numbers are the start and end of each stretch, named LTS, LTE, UTS and UTE in the paper. The map shrinks from n² squares to four vectors of n numbers, and the family it covers includes document, prefix-LM, share-question, sliding-window and block-sparse masks.

The trouble starts when the street is too long for one GPU. In context parallelism the query sequence is split among ranks, so each rank, a delivery district, owns some blocks of street numbers and needs the closure list written in its own numbering. Building each district's slice and sending it costs exactly what a long-context run is short of: communication. FlashMask had no distributed support at all, and MagiAttention, a system that accepts arbitrary masks, materializes per-rank mask structure by cutting slices.

The paper's observation is about how the four lists are stored. They are indexed by key position but hold query-row numbers. A district owns rows, not avenues, so every closure mark is already written in the unit the district is cutting by. Each rank can crop the global marks to its own chunks and renumber them from its local zero using only the mask it already has and its own chunk offsets. The paper names the operation clip–shift–max: clip each mark to a chunk, shift it into the rank's concatenated numbering, and take the maximum across the rank's chunks so that only the chunk that actually contains the mark contributes. Its Theorem 1 states that the rank-local intervals match the global ones at the corresponding positions, so the crop loses nothing, and the output has the same shape the single-GPU kernel already reads.

RouteHow the mask is storedHow a rank gets its sliceSource
FlashMask (previous version)Column-wise sparse: four query-axis boundary vectorsNot supported; Ampere-only, single deviceFlashMask-3, §2.2
FlexAttentionA user-written mask function, compiled just in time with TritonSeparate code path; a compile per maskFlashMask-3, §2.2
MagiAttentionAttention slices with a dedicated kernel, any mask typeSlice-cutting operations build per-rank mask structureFlashMask-3, §2.2
FlashMask-3Same four boundary vectorsClip–shift–max on the vectors, no cross-rank communicationFlashMask-3, §3.2

Cropping the marks is only the start of the savings. Inside the GPU, the same four vectors let the kernel sort every tile of the score matrix as fully masked, partly masked or unmasked, using the maximum and minimum along the key axis. Fully masked tiles are never loaded, the same IO-aware idea that makes FlashAttention fast, and only partly masked tiles apply the row-interval test. On Hopper, FlashAttention-3's producer warp group leaves 96 of its 128 threads idle each iteration, and the paper puts them to work as a third role, a KV Index Generator that scans the marks and queues the valid blocks.

Across ranks the same classification does two more jobs. Under expressive masks the load no longer grows smoothly with position as it does for a causal mask, so the paper adds a sparsity-aware load balancer, HCC-IPO, to avoid stragglers. And the forward all-gather skips key-value chunks that are fully masked, with a two-phase route that follows the cluster topology. The authors report 31–75% compute–communication overlap efficiency at context-parallel degrees 4–32 with 8K tokens per device.

Hold the context at 131,072 tokens (128K), and assume one byte per entry for a dense mask and four bytes per number for the vectors (illustrative; the paper states O(n²) → O(n), not these byte counts). The dense mask is 131,072² = about 17.2 billion bytes. The four vectors are 4 × 131,072 × 4 = 2,097,152 bytes. That is about 17.2 GB (16 GiB) → 2 MiB, a reduction of 8,192×, and the gap widens linearly with the sequence length. None of those bytes has to cross the network, because each rank crops them locally.

The method has stated limits. It covers masks whose blocked rows per key fit two intervals; other masks need an extended representation. The Hopper kernel assumes TMA and WGMMA, and the overlap design assumes nodes of eight GPUs joined by NVLink with a matching InfiniBand layout. Hierarchical routing is applied to the forward all-gather only, and the authors note that smaller contexts may not repay the scheduler overhead. Treat the headline numbers as results for production-scale long-context training on Hopper, measured by the authors on their own benchmarks.

Goes deeper in: LLM Internals → Self-Attention → Causal Masking

Related explainers

Frequently Asked Questions

Check what you knowMap your AI & GPU knowledge across every track — free, role-based