CompileMask: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking

• Runtime-dependent masks arise in workloads such as preference optimization, multimodal generation, and recommendation, while tile-based kernels natively support few patterns.Mask boundaries can depend on sequence lengths and segment boundaries.Source: §1 · p. 1, Figure 1 · p. 2
• Predicate-based masking is expressive but costly, whereas interval-based masking is efficient but restricted by a fixed interval count.These limitations reflect a trade-off between user-facing expressiveness and kernel-facing compactness.Source: §1 · p. 1
• The Execution Descriptor uses precomputed interval endpoints for visibility checks and Tile-State Bitmaps for tile scheduling.The compiler derives the required interval count from each Region’s visibility rule instead of imposing a predefined count.Source: §1 · p. 2
• COMPILEMASK’s symbolic compiler analyzes a user’s MaskSpec offline and generates a mask-specific preprocess kernel for runtime descriptor construction.Users express segment-structured visibility rules without hand-writing kernels.Source: §1 · pp. 1–2
• Up to 3.80× the single-operator training throughput of FlexAttention with the FA4 backend is achieved by COMPILEMASK.The reported result is from the paper’s experiments.Source: §1 · p. 2

2.1 Background

• Tiled attention skips fully masked tiles, processes fully visible tiles without masking, and checks visibility element-wise in partially masked tiles.Custom-mask support therefore requires both tile classification and visibility checks within partial tiles.Source: §2.1 · p. 2
• Predicate-based masking supports arbitrary rules but can incur costly blockmask construction and per-element predicate evaluation during attention.Metadata changes can require rebuilding the BlockMask.Source: §2.1 · pp. 2–3
• FlashMask’s four boundary vectors represent at most two separate masked intervals per column, so they cannot directly represent CLEAR’s four intervals in column 9.Source: §2.1 · p. 3

2.2 Motivation

• For targeted masks, visibility structure is fixed by the rule while boundary values depend on batch metadata, allowing structure analysis offline and boundary computation online.The HSTU example varies history length without changing the visibility rule.Source: §2.2 · p. 3

3.1 The COMPILEMASK Pipeline

• CompileMask compiles each MaskSpec once into a mask-specific preprocess kernel, which reads batch metadata online to build descriptors consumed by the attention kernel.The compiler derives interval counts from the visibility rules rather than imposing a predefined count.Source: §3.1 · p. 3

3.2 Offline Mask Specification and Compilation

• MaskSpec groups query rows into Regions with shared visibility rules, letting users describe segment-based patterns without manually deriving interval endpoints.Users can reuse a MaskSpec across batches while the masking rules remain unchanged, even when sequence lengths and segment boundaries vary.Source: §3.2 · p. 4, §3.2.1 · p. 4

3.2.1 Mask Specification

• Separating row guards from visibility rules lets users compose different patterns across Regions and reuse the same specification as batch boundaries vary.Source: §3.2.1 · p. 4

3.2.2 Symbolic Lowering to Visible Intervals

• The symbolic compiler translates visibility rules through AST construction, DNF normalization, interval extraction, and canonicalization; Region C's rule yields two interval slots.For Region C, the intervals are [0, b) and [s(i), i + 1); the compiler retains both because they are not mergeable across the Region as a whole.Source: §3.2.2 · pp. 4–5

3.3 Online Efficient Mask Attention Execution

• The online preprocess kernel evaluates compiled endpoints using batch metadata and builds descriptors for forward and backward attention execution.Forward and backward kernels use separate MaskSpecs to produce intervals aligned with their respective traversal directions.Source: §3.2.2 · p. 5, §3.3 · p. 6

3.3.1 Execution Descriptor

• The Execution Descriptor combines an Interval Table of visible-key intervals with two-bit-per-tile bitmaps that encode tile states for scheduling.Rows with fewer intervals than the compiler-determined maximum are padded with zero-length intervals; the bitmaps distinguish empty, full, and partial tiles.Source: §3.3.1 · p. 6

3.3.2 Descriptor-Guided Attention Execution

• Tile-State Bitmaps let the kernel skip empty tiles, process full tiles without masking, and apply interval-based visibility checks within partial tiles.Unsigned endpoint comparisons determine visibility in partial tiles, and invisible scores are replaced with −∞ before online softmax.Source: §3.3.2 · pp. 6–7

4.1 Experimental Setup

• The evaluation covers HSTU, shared-prompt reward model, CLEAR, and causal masks, with synthetic lengths from 1K to 32K and variable-length HSTU inputs constructed from KuaiRand-1K user-day interaction logs.For CLEAR, lengths refer to image tokens; for the other three mask patterns, they refer to total sequence length.Source: §4.1 · p. 7

4.2 Single-Operator Performance

• At 32K, COMPILEMASK achieves 1.44×, 2.13×, and 2.87× Flex-FA4 throughput on HSTU, RM, and CLEAR, respectively, and retains 85%–89% of native FA4 throughput on causal attention.It supports CLEAR at every evaluated length, unlike FlashMask, whose fixed representation cannot cover columns with three or more masked intervals.Source: Figure 5 · p. 8, §4.2 · p. 8, Table 1 · p. 8

4.3 End-to-End Performance

• At 32K, COMPILEMASK reaches 10.19 kTok/s on Qwen3-8B HSTU, 17.82 kTok/s on Qwen3-8B CLEAR, and 10.64 kTok/s on Qwen3-14B CLEAR, improving over Flex-FA4.On Qwen3-14B Reward Model, its 32K speedup over Flex-FA4 is 1.21×; causal throughput stays close to native FA4.Source: Figure 6 · p. 9, §4.3 · p. 9
• On variable-length KuaiRand HSTU traces, COMPILEMASK reaches 15.04 kTok/s, 3.83× Flex-Triton and 1.33× Flex-FA4, with all methods using packed inputs.The comparison uses Qwen3-8B end-to-end training.Source: Table 2 · p. 9, §4.3 · p. 9

4.4 Efficiency of Visibility Checks

• Visible intervals with unsigned arithmetic reduce combined HSTU forward and backward latency by 6.5%–8.2% versus masked intervals with ordinary comparisons.The implementation uses fewer masking instructions for visible intervals, and unsigned arithmetic replaces two endpoint comparisons and their conjunction with one subtraction and one comparison.Source: §4.4 · p. 9, Table 3 · p. 10

Efficient Attention

• FlashAttention-family methods improve efficiency through IO-aware tiling, work partitioning, and hardware-specific pipelines, while sparse methods reduce computation through content-dependent token or block selection.FlashInfer extends execution to diverse KV-cache layouts and dynamic inference workloads; the paper focuses on efficiently executing user-specified visibility rules.Source: §5 · p. 10

Attention Customization and Mask Representations

• Prior customization systems include FlexAttention, AttentionEngine, Flashlight, FlashMask, SPLAT, SAS, and SplashAttention, with different approaches to programmable masks, compressed representations, code generation, or preprocessing.FlashMask’s fixed four-vector representation limits expressible patterns; CompileMask instead derives interval counts from visibility rules and moves endpoint calculations out of attention execution.Source: §5 · p. 10

GPU Programming and Compilation

• Triton, ThunderKittens, and TileLang support tiled GPU programming, while SparseTIR and Mirage explore sparse-format transformations and joint algebraic and scheduling optimizations.These systems provide complementary implementation and optimization mechanisms; CompileMask derives interval representations from attention visibility rules.Source: §5 · p. 10

6 CONCLUSION

• COMPILEMASK compiles region-based MaskSpec rules into preprocess kernels that build an Execution Descriptor from runtime batch metadata, enabling efficient visibility checks with precomputed interval endpoints.The design targets expressive mask specifications while improving single-operator and end-to-end performance over evaluated baselines.Source: §6 · p. 10