arXiv CS open access

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

Hongzheng Li et al. · arxiv_cs

arXiv CS · Papers

cloud distributed-computing parallel-computing

Modern attention workloads require diverse masking patterns whose boundaries depend on runtime metadata, yet efficient attention kernels natively support only a small set of patterns. Existing methods for supporting custom masks in these kernels face an expressiveness-efficiency trade-off. Predicate

Conceptio Open Archive

Analyze & explore related research in Conceptio

Connect this record directly into your research workflow or AI agent. Explore citations across open repositories, query full text, and export structured proof bundles.

Explore in Search →Connect via MCP / API ↗

C OMPILE M ASK : B RIDGING THE G AP B ETWEEN E XPRESSIVENESS AND E FFICIENCY IN ATTENTION M ASKING

arXiv:2610.13474v1 [cs.DC] 9 Oct 2026

Hongzheng Li * † 1 Jiacheng Wan * 2 Zhaolong Xing 2 Zhida Jiang 2 Yuanhang Yang 2 Junhua Shu 2 Jiaxing Wang 2 Xupeng Miao 3 Yingxia Shao ‡ 1 Zhen Chen ‡ 2 Ke Zhang 2

A BSTRACT Modern attention workloads require diverse masking patterns whose boundaries depend on runtime metadata, yet efficient attention kernels natively support only a small set of patterns. Existing methods for supporting custom masks in these kernels face an expressiveness-efficiency trade-off. Predicate-based methods are expressive but incur substantial preprocessing and execution overhead, whereas interval-based methods execute efficiently but restrict expressiveness through a predefined interval count. We present C OMPILE M ASK, which addresses this trade-off by connecting expressive mask specifications with efficient attention execution. We design a compiler that analyzes visibility rules during the offline stage and generates a mask-specific preprocess kernel. We further propose MaskSpec, a region-based DSL that lets users express these rules. During the online stage, the generated kernel reads batch metadata and produces a compact Execution Descriptor that lets the attention kernel efficiently check visibility by comparing element indices against precomputed endpoints. Experiments show that C OMPILE M ASK achieves up to 3.8× speedup over state-of-the-art methods. 1

I NTRODUCTION

The predicate-based method (Dong et al., 2025) supports arbitrary masking rules through element-wise Boolean predicates. However, its efficiency is limited by expensive blockmask construction and per-element predicate evaluation during attention execution. The interval-based method (Wang et al., 2025), in contrast, enables efficient masking through compact interval representations. However, its expressiveness is limited because it cannot represent masks whose columns contain more intervals than its fixed format supports.

The attention mechanism is the central computational primitive of modern Transformers (Vaswani et al., 2017). Tasks such as preference optimization (Rafailov et al., 2023) with prefix sharing (Wang & Hegde, 2024), unified multimodal generation (Zhou et al., 2025), and generative recommendation (Zhai et al., 2024) each call for their own attentionmasking pattern. The attention mask realizes these patterns by specifying which keys and values each query may attend to. This visibility often depends on runtime metadata such as sequence lengths and segment boundaries, so the mask’s shape is not fixed ahead of time but instantiated per batch from these parameters.

We find that the limitations of both approaches stem from their reliance on a single representation to satisfy two conflicting demands. The user defining the mask needs a representation that is expressive and easy to write, while the kernel executing it needs one that is compact and of fixed layout. These approaches prioritize either user-facing expressiveness or kernel efficiency within a single representation. In this paper, we call this the expressiveness-efficiency trade-off.

Tile-based attention kernels (Dao et al., 2022) have become the de facto standard for their efficiency. However, these hardware-efficient kernels natively support only a small, fixed set of patterns such as causal or sliding-window. Realizing a mask outside this set forces a practitioner either to hand-modify the kernel or to materialize a dense O(N 2 ) mask, leaving custom masking a second-class citizen.

To bridge this gap, we present C OMPILE M ASK, which connects expressive user-facing mask specifications with an efficient kernel-facing representation through an offline compiler. C OMPILE M ASK works in two stages. During the offline stage, a user writes the masking logic once in a highlevel, user-facing form, and the compiler analyzes it and emits a mask-specific preprocess kernel that encodes how the visible structure of every row is composed. During the online stage, for each batch, this preprocess kernel reads

Recently, two lines of work have been proposed to inject custom attention patterns into tile-based attention kernels. * † Equal contribution. Work done during internship at JD.com. 1 Beijing University of Posts and Telecommunications 2 JD.com 3 Peking University. Correspondence to: ‡ Yingxia Shao <shaoyx@bupt.edu.cn>, Zhen Chen <chenzhen48@jd.com>.

C OMPILE M ASK: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking

Context History 2

Targets 4

Prompt 7

Res 1

Res 2

Image 8

9 10 11 12 13 14 15 16 0 1 2 3 4

Res 1

Targets

Text

7 0

Prompt

5 6 7 8 9 10 11

(a) HSTU mask

Res 2

The question that follows is how to design the user-facing representation for offline compilation and the kernel-facing representation for online execution. On the user side, C OM PILE M ASK provides a MaskSpec composed of Regions, each of which specifies a range of query rows and a visibility predicate for those rows. These Regions let users flexibly specify visibility relationships between token segments. On the kernel side, C OMPILE M ASK uses an Execution Descriptor comprising an Interval Table and Tile-State Bitmaps. The Interval Table stores precomputed interval endpoints for fast element-wise visibility checks, while the Tile-State Bitmaps compactly encode tile scheduling information to reduce its memory footprint. The offline compiler connects these two representations to preserve both expressiveness and efficiency. For expressiveness, it derives the required interval count from each Region’s visibility predicate rather than imposing a predefined count, allowing C OM PILE M ASK to support a wider range of masks compared with the existing interval-based method. For efficiency, the compiled interval representation lets the attention kernel check visibility by comparing element indices against precomputed endpoints, keeping online masking lightweight. Beyond these benefits, C OMPILE M ASK simplifies development by requiring users to write only a MaskSpec, from which the compiler generates the preprocess kernels. In our production workloads, this reduces user-written code from approximately 1,300 lines to 80 and development time from days to hours.

(b) Shared-prefix reward-model mask

14 16

Figure 1. Segment-structured attention masks. (a) HSTU mask with context, history, and target segments. (b) Shared-prefix reward-model mask, where each response attends to the prompt and causally within itself. (c) CLEAR multimodal mask with global text interactions and local image-to-image attention.

BACKGROUND AND M OTIVATION

2.1

Background

Tiled attention. Given query, key, and value matrices Q, K, V ∈ RN ×d , where N is the sequence length and d is the head dimension, attention forms the score matrix S and output O as √ S = QK ⊤ / d, O = softmax(S + M )V, ( (1) 0 if key j is visible to query i, Mi,j = −∞ otherwise. The mask determines which keys each query attends to and may depend on runtime metadata such as sequence lengths and segment boundaries. Examples include masks for preference optimization with prefix sharing (Rafailov et al., 2023; Wang & Hegde, 2024), unified multimodal generation (Zhou et al., 2025), and generative recommendation (Zhai et al., 2024).

Our contributions are summarized as follows:

FlashAttention (Dao et al., 2022) computes attention tile by tile without materializing the full score matrix. Fully masked tiles can be skipped, fully visible tiles require no masking, and partially masked tiles require element-wise visibility checks. Supporting a custom mask therefore involves two tasks: classifying tiles and resolving visibility within partial tiles.

• We design a compiler that symbolically analyzes boolean visibility rules and generates a preprocess kernel that constructs the Execution Descriptor at runtime. • We propose MaskSpec, a region-based DSL that lets users declare segment-structured attention masks without hand-writing kernels. • We design an attention kernel that uses the Execution Descriptor’s Tile-State Bitmaps to schedule tiles and its Interval Table to determine element-wise visibility within partial tiles. • We conduct comprehensive experiments, and the results show that C OMPILE M ASK achieves up to 3.80× the single-operator training throughput of FlexAttention with the FA4 backend.

History Context

Masked

Text

Visible

Image

the batch metadata and produces a compact Execution Descriptor that the attention kernel consumes directly. This descriptor is built once per batch and shared by attention layers that use the same mask. This separation allows users to express masking rules independently of the representation consumed by the attention kernel, while moving structural analysis offline and leaving only batch-dependent values to be computed online.

Predicate-based masking. FlexAttention (Dong et al., 2025) accepts an element-wise Boolean function mask mod(i, j) that determines whether key j is visible to query i, providing an expressively complete interface. Its flexibility comes with overhead in both preprocessing and attention execution. During preprocessing, create block mask evaluates the predicate across query–key pairs to identify tiles that can be skipped or

C OMPILE M ASK: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking

processed without masking. This construction can be costly and must be repeated when changes in runtime metadata invalidate the existing BlockMask. During attention execution, the kernel still evaluates the user-defined predicate for each element within partially masked tiles. Complex predicates therefore introduce additional arithmetic and logical operations into the attention kernel, making online visibility checks expensive even after the BlockMask has been constructed.

MaskSpec

Batch Metadata

Online Stage (Per Batch)

Interval-based masking. FlashMask (Wang et al., 2025) represents a mask using four per-column boundary vectors (LTS, LTE, UTS, UTE), where LTS and LTE specify the start and end of a masked interval in the lower triangular part of each column and UTS and UTE specify those in the upper triangular part. During preprocessing, it computes MinMax summaries by taking the minimum and maximum of each of these four vectors within each column block. The MinMax summaries are compared against a query tile’s row range to classify the tile as fully masked, fully visible, or partially masked. Within partially masked tiles, the attention kernel checks the column intervals to determine element-wise visibility. This representation supports efficient execution, but its fixed number of boundary vectors limits expressiveness. Each pair of vectors describes only one contiguous masked interval, so the four-vector format cannot directly represent a column containing more than two separate masked intervals. For example, as shown in Figure 1(c), the CLEAR (Liu et al., 2025) multimodal mask contains four separate masked intervals in column 9. This column exceeds the two-interval limit and cannot be represented directly by FlashMask’s four-vector format. 2.2

Symbolic Compiler

Sequence Length Region Length, …

Preprocess Kernel

Offline Stage (Per Mask)

Preprocess Kernel

Batch Tensor Data

Execution Descriptor

Attention Kernel

Attention Result

Figure 2. The C OMPILE M ASK pipeline. Pink blocks denote the two mask representations, and green blocks denote the computational components.

M ETHODOLOGY

3.1

The C OMPILE M ASK Pipeline

C OMPILE M ASK addresses the expressiveness-efficiency trade-off by using an offline compiler to connect expressive mask specifications with efficient online kernel execution. As shown in Figure 2, the pipeline operates in two stages. During the offline stage, the symbolic compiler analyzes the user-provided MaskSpec and generates a mask-specific preprocess kernel. The compiler generates this kernel once per MaskSpec, and it is reused across all batches even when their metadata differ. During the online stage, this kernel reads each batch’s mask metadata and instantiates an Execution Descriptor. Our custom attention kernel, built on FlashAttention (Shah et al., 2024; Zadouri et al., 2026), then consumes this descriptor together with the batch’s Q, K, and V tensors to compute the attention output. The two representations are designed around distinct goals of expressiveness and execution efficiency. Our observation is that query rows can share a visibility rule even when their visible keys differ with row indices and batch metadata. A Region therefore provides a natural unit for describing these rows together, motivating the region-based MaskSpec detailed in Section 3.2. For kernel execution, comparing indices against precomputed interval endpoints provides a lightweight way to determine visibility. This motivates the interval-based Execution Descriptor detailed in Section 3.3.

Motivation

Both approaches use a single representation to serve the user and the kernel at once, which forces the expressivenessefficiency trade-off. For the masks we target, however, a mask’s structure need not be discovered at runtime. The structure of each row’s visible keys is fixed by the masking rule, while only the boundary values that the rule refers to depend on the batch. For example, in the HSTU mask shown in Figure 1(a), the history length can vary across sequences, shifting the boundary between history and target tokens. The visibility rule remains unchanged, with each history query attending to all context tokens and causally to history tokens. This separation makes it possible to determine the mask structure offline and compute only the batch-dependent boundary values online.

The offline compiler connects these representations to address the expressiveness-efficiency trade-off. It derives the required interval count from each Region’s visibility rule rather than imposing a predefined count on all patterns. It also moves rule analysis offline and generates a preprocess kernel that computes endpoint values once per batch for reuse during attention execution. By deriving the interval representation from user-defined rules and precomputing its endpoints, C OMPILE M ASK preserves expressive mask specification while keeping attention-time visibility checks lightweight. 3

C OMPILE M ASK: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking

3.2

Offline Mask Specification and Compilation Listing 1. The HSTU mask expressed as a three-Region

To support offline compilation, we design a user-facing MaskSpec. mask representation that makes the structure of masking guards rules explicit and a symbolic compiler that translates these 12 #defRowwhen_A(i): rules into visible intervals. The key observation is that at- 3 return i < region_a_end tention masks in many practical workloads arise from input 45 def when_B(i): return region_a_end <= i and i < region_b_end sequences composed of token segments with different roles 6 def when_C(i): return region_b_end <= i and i < sequence_end and visibility rules. A natural interface therefore lets users 78 # Visibility rules specify the visibility rule for each segment directly. For 9 def visible_A(i, j): return j < region_b_end example, the HSTU mask in Figure 1(a) assigns different 10 11 def visible_B(i, j): rules to context, history, and target segments, while the 12 return j <= i def visible_C(i, j): CLEAR mask in Figure 1(c) combines global text interac- 13 14 group = (i - region_b_end) // group_size tions with local image-to-image attention. Rows within a 15 group_start = region_b_end + group * group_size return (j < region_b_end) or ( segment share a rule, although their visible keys may vary 16 17 group_start <= j and j < i + 1) with the query index and batch metadata. We therefore de- 18 # MaskSpec mask_spec = MaskSpec(regions=[ sign MaskSpec around Regions, each pairing a contiguous 19 20 RegionMask(name="Context", range of query rows with a user-defined visibility rule. This 21 when=when_A, visible=visible_A), organization makes the row structure explicit rather than 22 23 RegionMask(name="History", leaving it implicit in an element-wise predicate. The sym- 24 when=when_B, visible=visible_B), bolic compiler then derives visible intervals from these rules, 25 26 RegionMask(name="Target", allowing users to specify visibility without manually deriv- 27 when=when_C, visible=visible_C)]) ing interval endpoints. In the following, we introduce the 28 MaskSpec interface and the symbolic compiler in detail in Section 3.2.1 and 3.2.2, respectively. Descriptor, which will be detailed in Section 3.3.1. 3.2.1 Mask Specification By separating row guards from visibility rules, MaskSpec A MaskSpec formalizes this organization as a set of dislets users compose different visibility patterns across Rejoint Regions. Each Region Rg contains a row guard gions without manually deriving interval endpoints. Users wheng (i) and a visibility rule visibleg (i, j). The row guards write a MaskSpec once for each new masking pattern are mutually exclusive and jointly cover the valid query and reuse it across batches without rewriting it as long as rows, so each row i selects one Region ρ(i). Let visible(i, j) the masking rules remain unchanged, even when sequence denote the Boolean visibility predicate for element (i, j), lengths and segment boundaries vary. where a true value means that the element is visible. It is defined as 3.2.2 Symbolic Lowering to Visible Intervals ρ(i) = g

s.t.

wheng (i),

visible(i, j) = visibleρ(i) (i, j).

The compiler then lowers each Region’s visibility rule into symbolic intervals through four offline transformations and generates a preprocess kernel in TileLang (Wang et al., 2026) to evaluate their endpoints at runtime.

(2)

For example, Listing 1 specifies the HSTU mask (Zhai et al., 2024) shown in Figure 1(a). Regions A, B, and C correspond to the context, history, and target row segments, respectively, and pair runtime-dependent row ranges with their respective visibility rules.

The symbolic lowering process consists of the following 1 Abstract Syntax Tree (AST) Construction four steps. ⃝ translates the Python when and visible functions into symbolic trees that preserve the rules’ structure while leav2 Disjunctive Noring batch-dependent values symbolic. ⃝ mal Form (DNF) Normalization rewrites the visibility expression into DNF, expressed as a disjunction of conjunc3 Interval Extraction converts each supported contions. ⃝ junction of lower- and upper-bound constraints on the key index into one symbolic half-open interval. The compiler takes the maximum of the lower bounds and the minimum of the upper bounds to obtain the interval endpoints, so the

The guards and visibility rules in a MaskSpec may depend on batch metadata that varies across input sequences. In Listing 1, region a end and region b end mark the ends of Regions A and B, sequence end bounds the valid index range, and group size determines the grouped visibility pattern in Region C. A complete MaskSpec also declares where these values are read from in the input batch. During offline compilation, the compiler preserves them as symbolic inputs and emits a preprocess kernel that loads their concrete values at runtime to instantiate the Execution 4

C OMPILE M ASK: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking def when(i): return b <= i and i < N

Benefits of endpoint precomputation. Separating endpoint computation from visibility checks offers two benefits. First, the preprocess kernel computes endpoint values once per batch for reuse during attention execution, including across heads and layers that share the same mask. This avoids recomputing endpoint expressions whenever the attention kernel checks visibility. Second, endpoint precomputation reduces non-matrix computation inside the attention kernel. On Hopper and Blackwell GPUs, matrix multiplication runs on high-throughput Tensor Cores, while non-matrix arithmetic and softmax exponentials rely on CUDA cores and special function units with substantially lower throughput (Shah et al., 2024; Zadouri et al., 2026). Additional endpoint computation can therefore extend the non-matrix portion of attention execution and delay dependent matrix multiplications. Moving this computation into preprocessing leaves lightweight endpoint comparisons in the attention kernel, reducing the risk that masking overhead stalls the Tensor Core pipeline.

MaskSpec

def visible(i, j): s = b + ((i - b) // g) * g return (j < b) or (s <= j and j < i + 1)

① AST Construction OR

Region Guard

Disjunctive Normal Form

Symbolic AST ② DNF Norm

AND

③ Interval Extraction Symbolic Intervals

④ Interval Canonicalization

Preprocess Code Generation

Figure 3. Offline compilation of a MaskSpec into preprocess code, shown for Region C of Listing 1. Here, b, g, and N denote region b end, group size, and sequence end, respectively, and s denotes group start.

Traversal-aligned layouts. The forward and backward kernels traverse the attention matrix in different directions. The forward kernel holds a query tile and streams key tiles, requiring visible-key intervals for each query row. The backward kernel holds a key tile and streams query tiles to accumulate key and value gradients, requiring visiblequery intervals for each key column. Users provide separate forward and backward MaskSpecs to describe these two views of the same mask. In the backward specification, Regions select key columns and visibility rules describe the query rows that can attend to them. The same compilation pipeline generates a preprocess kernel for each specification, producing descriptors aligned with the corresponding attention kernel’s traversal. Each kernel therefore consumes its required interval representation directly, without transposing mask metadata at runtime.

outer disjunction becomes an interval union.   Pg _ ^  visibleg (i, j) = a(i, j) , p=1

a∈Ag,p

{j | visibleg (i, j)} =⇒

[

(3)

[Lg,p (i), Ug,p (i)).

p=1

Here, Pg is the number of conjunctive terms in the DNF of Region g, and Ag,p is the set of atomic comparisons in 4 Interval Canonicalization exploits the Region term p. ⃝ guard to establish endpoint relationships and merge provably overlapping or adjacent intervals, preserving the visible set. For example, Figure 3 illustrates this process using Region C in Listing 1. The visibility rule for a row in this Region is visibleC (i, j) = (j < b) ∨ (s(i) ≤ j ∧ j < i + 1) ,

Discussion. To compile a visibility rule into intervals, the compiler must determine an upper bound on the interval count that does not grow with sequence length and derive endpoint expressions that can be evaluated from row indices and batch metadata. When either requirement cannot be met, C OMPILE M ASK retains the original visibility predicate as a callable function and evaluates it online within partially masked tiles. This fallback preserves the expressiveness of the callable interface, supporting masking rules even when they cannot be compiled into intervals.

(4)

where b is the boundary between Regions B and C and s(i) is the beginning of the group containing row i. Interval Extraction maps the two DNF conjunctions to [0, b) and [s(i), i + 1), whose union is the visible set before Interval Canonicalization. Interval Canonicalization then consults the Region guard b ≤ i < N : since s(i) ≥ b for every row in Region C, the two intervals are provably disjoint and adjacent only for rows of the first group, so no merge is valid for the Region as a whole and Region C requires two interval slots. The canonical intervals are finally emitted as the preprocess kernel, whose runtime behavior is described in Section 3.3.1.

3.3

Online Efficient Mask Attention Execution

The Execution Descriptor represents visibility as intervals so that attention kernels can determine visibility through lightweight endpoint comparisons rather than evaluating user-defined masking predicates. At runtime, the preprocess 5

C OMPILE M ASK: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking Tile-State Bitmaps 00XX

Comp.

Mask

State

1100

01XX

Empty

1110

001X

Full

00X1 Compute Mask Bitmap Bitmap

Partial

1100

Context Region

History Region

e Ch

T il

1101

Partial Tile

1 3 5

2 3

6 7

Start Len Start Len

2 3 E le C h e m e ck nt S ta te

Masked Attention Score

Interval 1 Interval 2

Target Region

t S ta

𝑎 -inf

Compute Interval Overlap Apply 1

Visibility Predicate

Interval Table

𝑐

𝑑

𝑎

𝑏

𝑐

𝑑

Attention Score

Figure 4. Descriptor-guided attention execution. The Tile-State Bitmaps determine the tile-level execution path. For a partial tile, the Interval Table produces a per-element visibility predicate that is applied directly to the attention score tile.

kernels generated in Section 3.2 evaluate the compiled endpoint expressions using batch metadata and construct the Execution Descriptors for the forward and backward passes. The attention kernels consume these descriptors to schedule tiles and determine element-wise visibility. We first describe the descriptor’s structure and runtime construction, then explain how it enables efficient attention execution. 3.3.1

encode this information in the Tile-State Bitmaps, using only two bits per tile. The preprocess kernel reuses the computed interval endpoints to check visibility within each tile and reduces the results across query rows. The compute bit indicates whether any element is visible, while the mask bit indicates whether any element is masked. Together, they distinguish empty, full, and partial tiles, as shown in Figure 4. The resulting bitmaps and Interval Table form the Execution Descriptor, which is reused across attention layers that share the same mask.

Execution Descriptor

The Execution Descriptor consists of an Interval Table and Tile-State Bitmaps. The Interval Table stores precomputed interval endpoints for element-wise visibility checks within partial tiles, while the Tile-State Bitmaps compactly encode tile states to guide scheduling.

3.3.2

As shown in Figure 4, the attention kernel uses the TileState Bitmaps to select each tile’s execution path. Empty tiles are skipped without loading their key and value tiles or performing matrix multiplication. Full tiles follow the standard attention computation without masking, while partial tiles require element-wise visibility checks using the Interval Table.

The Interval Table stores K visible-key intervals for each query row i, encoding each interval by its start Li,m and length ℓi,m , K

Ii = {(Li,m , ℓi,m )}m=1 ,

ℓi,m = Ui,m − Li,m . (5)

For a partial tile, the kernel loads the interval starts and lengths for its query rows. For query row i and key position j, visibility is determined by

Each entry represents the half-open interval [Li,m , Li,m + ℓi,m ). This encoding supports the unsigned interval test described in Section 3.3.2. The compiler determines K as the maximum interval count required across the Regions in the MaskSpec. Rows requiring fewer than K intervals are padded with zero-length intervals so that every row stores K entries. For backward execution, the table stores visiblequery intervals for each key column, with Kbwd determined from the backward MaskSpec. To skip empty tiles and avoid masking in full tiles, the attention kernel needs tile-level visibility information. We

Descriptor-Guided Attention Execution

visible(i, j) =

K _

[Uint(j − Li,m ) < Uint(ℓi,m )] . (6)

m=1

When j < Li,m , unsigned subtraction underflows and the comparison fails. Otherwise, the comparison succeeds exactly when j lies within the interval. The kernel replaces invisible scores with −∞ before applying the online softmax.

C OMPILE M ASK: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking

• Causal (Vaswani et al., 2017) allows each query to attend to itself and preceding tokens.

Three design choices make element-wise masking lightweight. First, unsigned arithmetic combines the lowerand upper-bound checks into one subtraction and one comparison per interval, avoiding two endpoint comparisons and their logical conjunction. Second, preprocessing resolves Region-specific rules into per-row interval data, so the attention kernel no longer needs to identify which Region each row belongs to. Padding all rows to K intervals further enables a fully unrolled sequence of checks, even when a tile spans multiple Regions. Zero-length intervals naturally yield false comparisons and require no special handling. Third, the Interval Table directly represents visible intervals rather than masked intervals. In our implementation, this choice generates fewer masking instructions and improves performance, as evaluated in Section 4.4.

Our synthetic workloads cover these patterns in both singleoperator benchmarks and end-to-end training with Qwen38B and Qwen3-14B (Yang et al., 2025). We evaluate lengths from 1K to 32K, doubling at each step, where K denotes 1,024. These lengths refer to the number of image tokens for CLEAR and the total sequence length for the other three mask patterns. All methods receive identical inputs and mask metadata within each comparison. For real-world workloads, we construct variable-length HSTU inputs from KuaiRand-1K interaction logs (Gao et al., 2022), treating each user-day as a sequence. Baselines. We compare against dense softmax attention compiled by TorchInductor, FlexAttention (Dong et al., 2025), and FlashMask (Wang et al., 2025). We abbreviate FlashAttention versions 2, 3, and 4 as FA2, FA3, and FA4, respectively. FlexAttention uses either the Triton or FA4 backend, denoted as Flex-Triton and Flex-FA4, respectively. FlashMask uses PaddlePaddle 3.3.0 and executes its FA2 fallback1 on Blackwell for the evaluated four-vector representation. For causal attention, we include native FA4 (Zadouri et al., 2026) as a reference for the practical performance ceiling. We validate attention outputs and gradients with respect to queries, keys, and values against dense masked attention.

E XPERIMENTS

We evaluate whether C OMPILE M ASK preserves mask expressiveness while delivering efficient attention execution. 4.1

Experimental Setup

Environment Settings. Unless otherwise specified, all experiments are conducted on NVIDIA Blackwell GPUs. We use Python 3.12.9, PyTorch 2.11.0 with CUDA 13.0, and BF16 precision. C OMPILE M ASK uses TileLang 0.1.9 for preprocessing and CUTLASS DSL 4.5.2 for attention kernels.

4.2

Test Workloads. We evaluate four attention-masking patterns using synthetic inputs and real-world sequence lengths. For the synthetic inputs, let N denote the total sequence length.

Single-Operator Performance

We first study single-operator performance by measuring an attention operator together with its mask preprocessing, forward pass, and backward pass. The single-operator benchmarks use eight attention heads with a head dimension of 128. We keep 128K image tokens per batch for CLEAR and 128K total tokens for the other masks, with the batch size decreasing from 128 to 4 as the evaluated length increases. All sequences within each batch have the same length. Results average repeated measurements across four input instances after warmup.

• HSTU (Zhai et al., 2024) combines context, history, and grouped target visibility. We uniformly sample the context length between ⌊N/4⌋ and ⌊N/2⌋, and independently sample the target group size and number of groups from 3 to 8. The history occupies the remaining positions. • Shared-prompt reward model (RM) uses prefix sharing (Wang & Hegde, 2024), with six responses per prompt. Each response attends to the common prompt and causally within itself. We independently sample each response length uniformly between ⌊N/16⌋ and ⌊N/11⌋, with the prompt occupying the remaining positions.

Overall Performance. Figure 5 reports effective training throughput across the four mask patterns. We measure the total time for mask preprocessing and the forward and backward passes. Effective throughput is computed from the attention FLOPs over visible query–key pairs divided by the measured latency. Following FlashAttention (Dao, 2024), we count the backward pass as 2.5 times the forward-pass FLOPs to include score recomputation.

• CLEAR (Liu et al., 2025) combines global text attention with local image attention, using a spatial radius of R = 2. We uniformly sample the text length from the integers 128 through 384.

For HSTU, causal, and RM, C OMPILE M ASK outperforms 1

https://github.com/PaddlePaddle/Paddle/ blob/v3.3.0/python/paddle/nn/functional/ flash_attention.py

C OMPILE M ASK: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking Dense compile

Flex-Triton

FlashMask FA2

Reward Model

Native FA4

CompileMask

CLEAR

Causal

FlashMask: Unsupported

1,000

8K 16K 32K

OOM

500 OOM

Eﬀective TFLOP/s

HSTU

Flex-FA4

8K 16K 32K

Figure 5. Effective training throughput of a single operator at lengths from 1K to 32K. C OMPILE M ASK and Flex-FA4 use FlashAttention4-based kernels. OOM denotes out of memory. Higher is better.

workloads to measure GPU time for mask preprocessing, forward execution, and backward execution with Nsight Systems. This separate GPU-phase study uses a different timing protocol, HSTU metadata generator, and C OMPILE M ASK layout from Figure 5. For C OMPILE M ASK, preprocessing constructs the Interval Tables and Tile-State Bitmaps for both passes from batch metadata. For Flex-Triton and Flex-FA4, it constructs the BlockMask. For FlashMask, preprocessing generates the four boundary vectors and computes the MinMax summaries required by the forward and backward passes. Forward and backward timings include the attention computation and its associated auxiliary GPU operations, including gradient preparation and postprocessing in the backward pass.

Table 1. Time breakdown of single-operator execution, in milliseconds. Pre, Fwd, and Bwd denote preprocessing, forward, and backward, respectively. Lower is better. Method

Phase

16K

32K

Flex-Triton

Pre 0.292 0.568 1.125 2.238 5.643 Fwd 1.483 2.284 3.886 7.010 13.273 Bwd 7.688 9.170 12.628 20.014 34.947

11.617 25.294 64.904

Flex-FA4

Pre 0.280 0.543 Fwd 2.404 2.623 Bwd 3.435 4.394

11.977 12.647 34.977

FlashMask-FA2

Pre 0.105 0.105 0.105 0.105 0.105 0.106 Fwd 1.624 2.856 5.275 10.030 19.540 38.194 Bwd 4.237 7.438 13.801 26.406 51.546 100.532

Pre 0.011 0.013 C OMPILE M ASK Fwd 0.540 0.833 Bwd 1.465 2.286

1.063 2.164 5.773 3.077 4.409 7.433 6.254 10.099 18.354

0.019 1.507 4.107

0.033 0.060 2.723 5.217 8.191 16.047

0.118 10.075 31.322

Table 1 shows that C OMPILE M ASK reduces both mask preprocessing time and attention execution time compared with FlexAttention. As sequence length increases from 1K to 32K, preprocessing time grows from 0.011 to 0.118 ms for C OMPILE M ASK, compared with 0.280 to 11.977 ms for Flex-FA4. This advantage comes from computing interval endpoints and tile states directly from batch metadata, avoiding the element-wise predicate evaluation used to construct the BlockMask. Over the same range of sequence length, C OMPILE M ASK achieves forward speedups of 1.26–4.45× and backward speedups of 1.12–2.34× over Flex-FA4. Endpoint precomputation also contributes to these execution speedups by moving boundary calculations out of the attention kernel, leaving lightweight endpoint comparisons for visibility checks and reducing the non-matrix computation performed during attention execution. FlashMask FA2 has comparable preprocessing time at 32K, but its forward and backward execution remains slower.

the evaluated general-mask baselines at all six sequence lengths. At 32K, it achieves 1.44×, 2.13×, and 2.87× the throughput of Flex-FA4 on HSTU, RM, and CLEAR, respectively. For causal attention, it retains 85%–89% of native FA4 throughput, showing that support for custom masks introduces limited overhead relative to the specialized implementation. Here dense compile runs out of memory at long sequence lengths because it materializes quadratic-size attention intermediates. We emphasize that FlashMask does not support the evaluated CLEAR mask because some columns contain three or more masked intervals, exceeding the two-interval limit of its fixed four-vector representation. C OMPILE M ASK supports this mask at all evaluated sequence lengths, demonstrating its ability to handle visibility patterns beyond FlashMask’s fixed interval format. Performance Breakdown. We use fixed-length HSTU

4.3

End-to-End Performance

We evaluate end-to-end training throughput using Qwen38B and Qwen3-14B in TorchTitan (Liang et al., 2024) with 8

C OMPILE M ASK: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking Flex-Triton

HSTU Qwen3-8B Throughput (kTok/s)

Native FA4

Reward Model

HSTU

Reward Model

CLEAR

Causal

2K 4K 8K 16K Sequence length

32K

Causal 15

CompileMask

CLEAR

Qwen3-14B Throughput (kTok/s)

Flex-FA4

2K 4K 8K 16K Sequence length

32K

4 2 1K

4K 8K 16K Image tokens

32K

2K 4K 8K 16K Sequence length

32K

Figure 6. End-to-end training throughput. Higher is better.

BF16, AdamW, torch.compile, and activation checkpointing. We use a batch size of one. For CLEAR, text lengths vary across training steps according to the same presampled schedule for all methods, and throughput includes both image and text tokens. Mask preprocessing runs once per training step, and its results are reused across attention layers and activation recomputation. We measure complete training steps and report average throughput over three independent runs. FlashMask is excluded because its PaddlePaddle implementation is not integrated into this PyTorch training stack.

Table 2. End-to-end Qwen3-8B training on KuaiRand length traces. Method Flex-Triton Flex-FA4 C OMPILE M ASK

Step latency (ms) ↓

Throughput (kTok/s) ↑

962.79 ± 1.07 334.04 ± 1.49 251.58 ± 3.06

3.929 ± 0.004 11.325 ± 0.051 15.038 ± 0.184

8,192 tokens and crops its metadata to the actual packed length before attention execution. Table 2 shows that C OM PILE M ASK achieves 15.04 kTok/s, providing 3.83× and 1.33× speedups over Flex-Triton and Flex-FA4, respectively. These results demonstrate that C OMPILE M ASK also improves end-to-end training throughput on variable-length workloads when all methods use packed inputs.

Figure 6 shows that C OMPILE M ASK improves end-to-end training throughput over FlexAttention on HSTU, Reward Model, and CLEAR workloads. Throughput is reported in kTok/s, denoting thousands of tokens per second. For Qwen3-8B with HSTU at 32K, C OMPILE M ASK achieves 10.19 kTok/s, compared with 9.67 kTok/s for Flex-FA4 and 5.74 kTok/s for Flex-Triton, yielding speedups of 1.05× and 1.78×, respectively. For Qwen3-14B with Reward Model at 32K, the speedup over Flex-FA4 reaches 1.21×. With 32K image tokens on CLEAR, C OMPILE M ASK achieves 17.82 and 10.64 kTok/s for Qwen3-8B and Qwen3-14B, respectively, corresponding to speedups of 1.31× and 1.16× over Flex-FA4. For causal attention, C OMPILE M ASK maintains throughput close to native FA4. These results show that C OMPILE M ASK retains its performance benefits in complete training steps, even when mask preprocessing is amortized across attention layers and activation recomputation.

4.4

Efficiency of Visibility Checks

We further evaluate the visibility-check designs in Section 3.3.2 by comparing visible and masked intervals with unsigned arithmetic or ordinary endpoint comparisons. Table 3 reports combined forward and backward latency on HSTU with precomputed descriptors. Visible intervals with unsigned arithmetic reduce latency by 6.5%–8.2% compared with masked intervals with ordinary comparisons. The performance gains result from fewer masking instructions in our implementation, as visible intervals generate fewer instructions than masked intervals and unsigned arithmetic replaces two endpoint comparisons and their logical conjunction with one subtraction and one comparison.

We further evaluate Qwen3-8B training with HSTU masks using real-world sequence lengths from KuaiRand user-day traces. All methods use packed inputs. Our FlexAttention integration constructs a global BlockMask at a capacity of 9

C OMPILE M ASK: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking

GPU Programming and Compilation. GPU programming systems such as Triton and ThunderKittens (Tillet et al., 2019; Spector et al., 2025) simplify kernel development through tile-oriented abstractions and support efficient mapping onto GPU hardware. TileLang (Wang et al., 2026) provides a composable tiled programming model and serves as the target for our generated preprocess kernels. Complementary compiler research explores composable sparse formats and transformations in SparseTIR (Ye et al., 2023), and joint algebraic and scheduling optimizations in Mirage (Wu et al., 2025). These systems provide complementary mechanisms for kernel implementation and optimization, while C OMPILE M ASK derives interval representations from attention visibility rules.

Table 3. Combined forward and backward latency of different visibility-check designs on HSTU, in milliseconds. Lower is better. Descriptor

Test

16K

64K

Visible interval Unsigned 2.227 5.267 17.642 64.578 Masked interval Unsigned 2.249 5.305 17.684 64.687 Visible interval Ordinary 2.278 5.442 17.649 64.546 Masked interval Ordinary 2.381 5.672 18.905 70.379

R ELATED W ORK

Efficient Attention. The FlashAttention family (Dao et al., 2022; Dao, 2024; Shah et al., 2024; Zadouri et al., 2026) improves attention efficiency through IO-aware tiling, work partitioning, and hardware-specific pipelines. These techniques reduce intermediate memory traffic and improve GPU utilization for long sequences. FlashInfer (Ye et al., 2025) extends efficient attention execution to diverse KVcache layouts and dynamic inference workloads, with customizable attention variants supported through JIT compilation. Sparse attention methods such as MoBA (Lu et al., 2025), NSA (Yuan et al., 2025), and DeepSeek Sparse Attention (DeepSeek-AI et al., 2025) reduce attention computation through content-dependent token or block selection. Our focus is on efficiently executing user-specified visibility rules.

C ONCLUSION

We presented C OMPILE M ASK, which addresses the expressiveness-efficiency trade-off in attention masking through separate user-facing and kernel-facing representations connected by an offline compiler. The compiler translates a region-based MaskSpec into preprocess kernels that construct an Execution Descriptor from batch metadata at runtime. This design supports expressive mask specifications while enabling efficient attention execution through precomputed interval endpoints. Experiments demonstrate improvements in single-operator performance and end-toend training throughput over the evaluated baselines.

Attention Customization and Mask Representations. FlexAttention (Dong et al., 2025) exposes score and mask modifications that are compiled into fused attention kernels. AttentionEngine (Chen et al., 2025) combines programmable attention templates with cross-platform scheduling, while Flashlight (You et al., 2026) transforms standard PyTorch attention programs into fused kernels without requiring static templates. FlashMask (Wang et al., 2025) uses a compact column-wise interval representation for efficient masking, but its fixed four-vector format limits the patterns it can express. SPLAT (Gupta et al., 2025) exploits geometric regularity through an affine-compressed sparserow format and specialized code generation. SAS (Zhou et al., 2026) composes static and dynamic sparse patterns and generates inference kernels and KV-cache management functions. SplashAttention (The JAX Authors, 2024) also preprocesses masks into block-level information for sparse attention execution on TPUs. C OMPILE M ASK compiles user-specified visibility rules into preprocess kernels that instantiate an Execution Descriptor from batch metadata. The compiler derives interval counts from the rules rather than imposing a fixed count, and moves endpoint calculations out of attention execution. This separates expressive mask specification from efficient runtime visibility checks.

R EFERENCES Chen, F., Cheng, Y., Wang, L., Xia, Y., Miao, Z., Ma, L., Yang, F., Xue, J., Yang, Z., Yang, M., and Chen, H. AttentionEngine: A Versatile Framework for Efficient Attention Mechanisms on Diverse Hardware Platforms, February 2025. Dao, T. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. In Twelfth International Conference on Learning Representations (ICLR), pp. 1–14, 2024. Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 16344–16359. Curran Associates, Inc., 2022. DeepSeek-AI, Liu, A., Mei, A., Lin, B., Xue, B., Wang, B., Xu, B., Wu, B., Zhang, B., Lin, C., Dong, C., Lu, C., et al. DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models, December 2025.

Dong, J., Feng, B., Guessous, D., Liang, Y., and He, H. FlexAttention: A programming model for generating fused

C OMPILE M ASK: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking

attention variants. In Zaharia, M., Joshi, G., and Lin, Y. (eds.), Proceedings of Machine Learning and Systems, volume 7. MLSys, 2025.

jax/tree/main/jax/experimental/ pallas/ops/tpu/splash_attention. Tillet, P., Kung, H. T., and Cox, D. Triton: An intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pp. 10–19, Phoenix AZ USA, June 2019. ACM.

Gao, C., Li, S., Zhang, Y., Chen, J., Li, B., Lei, W., Jiang, P., and He, X. KuaiRand: An unbiased sequential recommendation dataset with randomly exposed videos. In Proceedings of the 31st ACM International Conference on Information & Knowledge Management, pp. 3953–3957, 2022.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is All you Need. In NeurIPS, 2017.

Gupta, A., Yuan, Y., Jain, D., Ge, Y., Aponte, D., Zhou, Y., and Mendis, C. SPLAT: A Framework for Optimised GPU Code-Generation for SParse reguLar ATtention. Proceedings of the ACM on Programming Languages, 9(OOPSLA1):1632–1660, April 2025. ISSN 2475-1421.

Wang, F. and Hegde, S. Accelerating direct preference optimization with prefix sharing. In NeurIPS 2024 Workshop on Fine-Tuning in Modern Machine Learning: Principles and Scalability, 2024.

Liang, W., Liu, T., Wright, L., Constable, W., Gu, A., Huang, C.-C., Zhang, I., Feng, W., Huang, H., Wang, J., Purandare, S., Nadathur, G., et al. TorchTitan: One-stop PyTorch native solution for production ready LLM pretraining, 2024.

Wang, G., Zeng, J., Xiao, X., Wu, S., Yang, J., Zheng, L., Chen, Z., Bian, J., Yu, D., and Wang, H. FlashMask: Efficient and rich mask extension of FlashAttention. In Yue, Y., Garg, A., Peng, N., Sha, F., and Yu, R. (eds.), International Conference on Learning Representations, volume 2025, pp. 27641–27661, 2025.

Liu, S., Tan, Z., and Wang, X. CLEAR: Conv-like linearization revs pre-trained diffusion transformers up. In Belgrave, D., Zhang, C., Lin, H., Pascanu, R., Koniusz, P., Ghassemi, M., and Chen, N. (eds.), Advances in Neural Information Processing Systems, volume 38, Main Conference, pp. 154623–154657. Curran Associates, Inc., 2025.

Wang, L., Cheng, Y., Shi, Y., Mo, Z., Tang, Z., Xie, W., Wu, T., Ma, L., Xia, Y., Xue, J., Yang, F., and Yang, Z. TileLang: Bridge programmability and performance in modern neural kernels. In The Fourteenth International Conference on Learning Representations, 2026.

Lu, E., Jiang, Z., Liu, J., Du, Y., Jiang, T., Hong, C., Liu, S., He, W., Yuan, E., Wang, Y., Huang, Z., Yuan, H., et al. MoBA: Mixture of Block Attention for Long-Context LLMs, February 2025.

Wu, M., Cheng, X., Liu, S., Shi, C., Ji, J., Ao, M. K., Velliengiri, P., Miao, X., Padon, O., and Jia, Z. Mirage: A Multi-Level superoptimizer for tensor programs. In 19th USENIX Symposium on Operating Systems Design and Implementation (OSDI 25), pp. 21–38, Boston, MA, July 2025. USENIX Association.

Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Ermon, S., and Finn, C. Direct preference optimization: Your language model is secretly a reward model. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 53728–53741. Curran Associates, Inc., 2023.

Yang, A., Li, A., Yang, B., Zhang, B., Hui, B., Zheng, B., Yu, B., Gao, C., Huang, C., Lv, C., Zheng, C., Liu, D., et al. Qwen3 Technical Report, 2025. Ye, Z., Lai, R., Shao, J., Chen, T., and Ceze, L. SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, pp. 660– 678, Vancouver BC Canada, March 2023. ACM.

Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., and Dao, T. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. In NeurIPS, 2024. Spector, B., Arora, S., Singhal, A., Parthasarathy, A., Fu, D., and Re, C. ThunderKittens: Simple, fast, and adorable kernels. In Yue, Y., Garg, A., Peng, N., Sha, F., and Yu, R. (eds.), International Conference on Learning Representations, volume 2025, pp. 2097–2130, 2025.

Ye, Z., Chen, L., Lai, R., Lin, W., Zhang, Y., Wang, S., Chen, T., Kasikci, B., Grover, V., Krishnamurthy, A., and Ceze, L. FlashInfer: Efficient and customizable attention engine for LLM inference serving. In Eighth Conference on Machine Learning and Systems, 2025.

The JAX Authors. SplashAttention: Sparse Flash Attention, 2024. URL https://github.com/jax-ml/ 11

C OMPILE M ASK: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking

You, B., Wang, I., Mustafaoglu, Z. S., Jangda, A., Moreira, A., Dathathri, R., Mahajan, D., and Pingali, K. Flashlight: PyTorch Compiler Extensions to Accelerate Attention Variants. In MLSys, 2026. Yuan, J., Gao, H., Dai, D., Luo, J., Zhao, L., Zhang, Z., Xie, Z., Wei, Y., Wang, L., Xiao, Z., Wang, Y., Ruan, C., et al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. In Che, W., Nabende, J., Shutova, E., and Pilehvar, M. T. (eds.), Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 23078– 23097, Vienna, Austria, July 2025. Association for Computational Linguistics. Zadouri, T., Hoehnerbach, M., Shah, J., Thakkar, V., and Dao, T. FlashAttention-4: Algorithm and kernel pipelining co-design for asymmetric hardware scaling. In Chowdhery, A. and Jia, Z. (eds.), Proceedings of Machine Learning and Systems, volume 8, pp. 912–926. MLSys, 2026. Zhai, J., Liao, L., Liu, X., Wang, Y., Li, R., Cao, X., Gao, L., Gong, Z., Gu, F., He, J., Lu, Y., and Shi, Y. Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations. In ICML, 2024. Zhou, C., Yu, L., Babu, A., Tirumala, K., Yasunaga, M., Shamis, L., Kahn, J., Ma, X., Zettlemoyer, L., and Levy, O. Transfusion: Predict the next token and diffuse images with one multi-modal model. In Yue, Y., Garg, A., Peng, N., Sha, F., and Yu, R. (eds.), International Conference on Learning Representations, volume 2025, pp. 6446–6469, 2025. Zhou, Y., Xiang, S., Yu, L., Song, Z., Mendis, C., and Wang, Y. SAS: Sparse attention synthesizer for efficient language model inference. In Proceedings of the 21st European Conference on Computer Systems, Eurosys ’26, pp. 1707–1721, New York, NY, USA, 2026. Association for Computing Machinery.

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

Analyze & explore related research in Conceptio

Related documents