Skip to content

Masks#

Mask functions define which query-key pairs can attend to each other. Each function returns a mask_mod — a callable with signature (b, h, q_idx, kv_idx) -> bool — that can be passed to create_block_mask to produce a BlockMask. See Concepts for details on how masks and block sparsity work together.

Causal#

Standard lower-triangular causal mask. Each position attends only to itself and earlier positions.

from attn_gym.masks import causal_mask

block_mask = create_block_mask(causal_mask, B, H, S, S, device=device)

attn_gym.masks.causal.causal_mask(b, h, q_idx, kv_idx) #

attn_gym.masks.causal.create_causal_block_mask_fast(batch_size, num_heads, q_seq_len, kv_seq_len, device, block_size=128, separate_full_blocks=True) #

Create a causal block mask efficiently without materializing the full mask.

This function generates the block mask data structure directly for causal attention, avoiding the need to create and process a full dense mask. This is much more efficient for long sequences.

Parameters:

Name Type Description Default
q_seq_len int

Query sequence length

 required
kv_seq_len int

Key/value sequence length

 required
device device

Device to create tensors on

 required
batch_size int | None

Batch size (defaults to 1 if None)

 required
num_heads int | None

Number of attention heads (defaults to 1 if None)

 required
block_size int

Block size for the block mask (both Q and KV use same size)

 128
separate_full_blocks bool

Whether to separate full blocks from partial blocks

 True

Returns:

Name Type Description
BlockMask BlockMask

Block mask object for causal attention

Sliding Window#

Each position attends to a fixed-size window of preceding tokens (combined with causal masking).

from attn_gym.masks import generate_sliding_window

mask_mod = generate_sliding_window(window_size=1024)
block_mask = create_block_mask(mask_mod, B, H, S, S, device=device)

attn_gym.masks.sliding_window.generate_sliding_window(window_size) #

Generates a sliding window attention mask with a given window size. Args: window_size: The size of the sliding window.

Note

We assume that the window size represents the lookback size and we mask out all future tokens similar to causal masking.

Dilated Sliding Window#

Sliding window with dilation — attends to every dilation-th token within the window.

attn_gym.masks.dilated_sliding_window.generate_dilated_sliding_window(window_size, dilation) #

Generates a dilated sliding window attention mask. Args: window_size: The size of the sliding window. dilation: The dilation factor for the sliding window.

Note

Query at position i can only attend to keys within a window of size window_size centered around i, where the keys are at positions j such that: * abs(i - j) <= window_size * abs(i - j) % dilation == 0

Global + Sliding Window#

Longformer-style attention (paper): a bidirectional sliding window plus designated global tokens (e.g. CLS) that attend to and are attended by every position.

import torch
from attn_gym.masks import generate_global_sliding_window

is_global = torch.zeros(S, dtype=torch.bool, device=device)
is_global[0] = True  # CLS token
mask_mod = generate_global_sliding_window(window_size=512, is_global=is_global)
block_mask = create_block_mask(mask_mod, B, H, S, S, device=device)

attn_gym.masks.global_sliding_window.generate_global_sliding_window(window_size, is_global) #

Generates a Longformer-style global + sliding window attention mask.

Parameters:

Name Type Description Default
window_size int

The symmetric sliding window radius; position i attends to positions j with abs(i - j) <= window_size.

 required
is_global Tensor

Boolean tensor of shape [SEQ_LEN] marking global tokens (e.g. CLS or task tokens). Global tokens attend to all positions and all positions attend to them.

 required
Note

Following the Longformer paper, attention is bidirectional: local attention is a window centered on each position, and global attention is symmetric. Compose with a causal mask via and_masks for decoder use.

Prefix LM#

Bidirectional attention over a prefix, causal attention over the rest.

from attn_gym.masks import generate_prefix_lm_mask

mask_mod = generate_prefix_lm_mask(prefix_length=512)

attn_gym.masks.prefix_lm.generate_prefix_lm_mask(prefix_length) #

Generates a prefix LM causal attention mask.

Parameters:

Name Type Description Default
prefix_length int

The length of the prefix.

 required
Note

This mask allows full attention within the prefix (first PREFIX_LENGTH tokens) and causal attention for the rest of the sequence.

Block Diffusion#

Hybrid attention mask for Block Diffusion language models (BD3-LM, paper). Training attends over a concatenated [x_t; x_0] sequence of length 2 * S: noised tokens attend bidirectionally within their block and to clean tokens in previous blocks, while clean tokens attend block-causally to other clean tokens. The mask is highly block-sparse, so FlexAttention can achieve large speedups over dense SDPA.

from attn_gym.masks import generate_block_diffusion_mask

mask_mod = generate_block_diffusion_mask(seq_len=S, block_size=16)
block_mask = create_block_mask(mask_mod, B, H, 2 * S, 2 * S, device=device)

attn_gym.masks.block_diffusion.generate_block_diffusion_mask(seq_len, block_size) #

Generates the Block Diffusion training mask from BD3-LM section 3.1.

Training runs attention over a length 2 * seq_len sequence: the noised tokens x_t occupy positions [0, seq_len) and the clean tokens x_0 occupy positions [seq_len, 2 * seq_len). The mask is the union of three pieces:

  • Block Diagonal: noised tokens attend bidirectionally within their own block.
  • Offset Block Causal: noised tokens attend to clean tokens in strictly previous blocks.
  • Block Causal: clean tokens attend to clean tokens in their own and previous blocks.

Clean tokens never attend to noised tokens.

Parameters:

Name Type Description Default
seq_len int

Length of the clean sequence x_0. The returned mask_mod is defined over indices in [0, 2 * seq_len).

 required
block_size int

Diffusion block size L'; the sequence is partitioned into contiguous blocks of this many tokens.

 required

JetSpec Tree Attention#

Tree-verification attention for JetSpec: each tree query attends to cached prefix keys and to flattened tree keys only when they are ancestors of that query node, including self. The example below uses the Figure 3 candidate-tree order: return, a, -, +, B, b, sum.

from attn_gym.masks import build_tree_ancestor_matrix, generate_jetspec_tree_causal_mask_mod

parent_indices = [-1, 0, 1, 1, 3, 3, 0]
ancestor = build_tree_ancestor_matrix(parent_indices, device="cuda")
mask_mod = generate_jetspec_tree_causal_mask_mod(prefix_length=1024, ancestor_matrix=ancestor)
block_mask = create_block_mask(
    mask_mod,
    B,
    H,
    ancestor.shape[0],
    1024 + ancestor.shape[0],
    device="cuda",
)

attn_gym.masks.jetspec.build_tree_ancestor_matrix(parent_indices, device='cuda') #

Build the ancestor matrix for the paper's tree-causal mask.

Parameters:

Name Type Description Default
parent_indices Sequence[int]

Parent index for each flattened tree node in parent-before-child order. The root entry is ignored and is conventionally -1.

 required
device str | device

Output device.

 'cuda'

Returns:

Type Description
Tensor

A boolean tensor where ancestor[i, j] is true when node j is on node
Tensor

i's root-to-self path.

attn_gym.masks.jetspec.generate_jetspec_tree_causal_mask_mod(prefix_length, ancestor_matrix) #

Generate the paper's tree-causal mask for parallel tree drafting/verification.

Parameters:

Name Type Description Default
prefix_length int

Number of prefix keys before the flattened tree keys.

 required
ancestor_matrix Tensor

Boolean (num_tree_nodes, num_tree_nodes) tensor. Entry [i, j] is true when tree node j is an ancestor of query node i, including self.

 required

Returns:

Type Description
_mask_mod_signature

A mask_mod for Q_LEN=num_tree_nodes and
_mask_mod_signature

KV_LEN=prefix_length + num_tree_nodes.

JetSpec's draft-head training path uses a different multi-block causal mask: sampled-block queries attend to all verified prefix keys and causally within their own sampled block.

from attn_gym.masks import generate_jetspec_training_mask_mod

prefix_length = 4096
block_size = 16
num_blocks = 3
mask_mod = generate_jetspec_training_mask_mod(prefix_length, block_size)
block_mask = create_block_mask(
    mask_mod,
    B,
    H,
    num_blocks * block_size,
    prefix_length + num_blocks * block_size,
    device="cuda",
)

attn_gym.masks.jetspec.generate_jetspec_training_mask_mod(prefix_length, block_size) #

Generate the paper's multi-block causal draft-head training mask.

Query rows are sampled-block tokens. KV columns are prefix followed by one or more sampled blocks, each laid out as [anchor, future_1, ..., future_N]. Each query can attend to the full verified prefix and to the anchor plus earlier positions in its own block, but not to future positions or other sampled blocks.

Parameters:

Name Type Description Default
prefix_length int

Number of verified prefix keys before sampled training blocks.

 required
block_size int

Number of tokens in each sampled block, including the anchor.

 required

Returns:

Type Description
_mask_mod_signature

A mask_mod for Q_LEN=num_blocks * block_size and
_mask_mod_signature

KV_LEN=prefix_length + num_blocks * block_size.

Document Mask#

For packed sequences: restrict attention to within document boundaries by wrapping a base mask with document offsets.

from attn_gym.masks import causal_mask, generate_doc_mask_mod
from attn_gym.masks.document_mask import length_to_offsets

lengths = [3, 2, 5]
offsets = length_to_offsets(lengths, device="cuda")
mask_mod = generate_doc_mask_mod(causal_mask, offsets)

attn_gym.masks.document_mask.generate_doc_mask_mod(mask_mod, offsets) #

Generates mask mods that apply to inputs to flex attention in the sequence stacked format.

Parameters:

Name Type Description Default
mask_mod _mask_mod_signature

The mask mod to apply to the documents

 required
offsets Tensor

This tensor should be of shape(num_documents + 1) this should contain the cumulative counts of document tokens. e.g. if you have 3 documents of length 2, 4, 3 then offsets = [0, 2, 6, 9]

 required
Note

What is the sequence stacked format? When assembling batches of inputs, we take multiple sequences and stack them together to form 1 large sequence. We then use masking to ensure that the attention scores are only applied to tokens within the same document.

For causal packed sequences where every batch element shares the same document layout, use the causal-only offset form. It is equivalent to wrapping causal_mask, but exposes each query row as one contiguous KV interval.

from attn_gym.masks import generate_packed_causal_doc_mask_mod
from attn_gym.masks.document_mask import length_to_offsets

lengths = [3, 2, 5]
offsets = length_to_offsets(lengths, device="cuda")
mask_mod = generate_packed_causal_doc_mask_mod(offsets)

attn_gym.masks.document_mask.generate_packed_causal_doc_mask_mod(offsets) #

Generates a causal document mask for packed sequences sharing one offset layout.

Parameters:

Name Type Description Default
offsets Tensor

Cumulative document token counts with shape (num_documents + 1,).

 required
Note

This is equivalent to generate_doc_mask_mod(causal_mask, offsets) for packed causal attention, but expresses each query row as one contiguous KV interval.

Neighborhood Attention (NATTEN)#

Multi-dimensional neighborhood attention patterns.

attn_gym.masks.natten.generate_natten(canvas_w, canvas_h, kernel_w, kernel_h) #

Generates a NATTEN attention mask with a given kernel size. Args: canvas_w: The width of the canvas. canvas_h: The height of the canvas. kernel_w: The width of the kernel. kernel_h: The height of the kernel.

attn_gym.masks.natten.generate_tiled_natten(W, H, K_W, K_H, T_W, T_H) #

Generates a NATTEN attention mask with a given kernel size and static tiling. Args: W: The width of the canvas. H: The height of the canvas. K_W: The width of the kernel. K_H: The height of the kernel. T_W: The width of the tile. T_H: The height of the tile.

attn_gym.masks.natten.generate_morton_natten(canvas_w, canvas_h, kernel_w, kernel_h) #

Generates a NATTEN attention mask with a given kernel size under morton curve layout. Args: canvas_w: The width of the canvas. canvas_h: The height of the canvas. kernel_w: The width of the kernel. kernel_h: The height of the kernel.

STA (Sparse Temporal Attention)#

attn_gym.masks.sta.generate_sta_mask_mod_2d(canvas_hw, kernel_hw, tile_hw, text_seq_len=0) #

Generates a 2D STA mask with a given kernel size.

Parameters:

Name Type Description Default
canvas_hw Tuple[int, int]

The shape of the canvas (height, width).

 required
kernel_hw Tuple[int, int]

The shape of the kernel (height, width).

 required
tile_hw Tuple[int, int]

The shape of the tile (height, width).

 required
text_seq_len int

The length of the text sequence for masking.

 0

attn_gym.masks.sta.generate_sta_mask_mod_3d(canvas_twh, kernel_twh, tile_twh, text_seq_len=0) #

Generates a 3D STA mask with a given kernel size.

Parameters:

Name Type Description Default
canvas_twh Tuple[int, int, int]

The shape of the canvas (time, height, width).

 required
kernel_twh Tuple[int, int, int]

The shape of the kernel (time, height, width).

 required
tile_twh Tuple[int, int, int]

The shape of the tile (time, height, width).

 required
text_seq_len int

The length of the text sequence for masking.

 0

VSA (Video Sparse Attention)#

VSA's fine sparse pass can be represented by precomputing each query tile's top-k KV tiles and turning those tile ids into a BlockMask. Use create_vsa_block_mask for the direct Triton construction path; deriving the same mask with generic create_block_mask is intended only for visualization or small correctness checks. For current Flex FLASH all-full tile paths, use create_vsa_flash_block_mask. The coarse top-k selection and FastVideo-style additive coarse/fine output combine run outside the FlexAttention kernel.

from attn_gym.masks import compute_vsa_coarse_attention, create_vsa_block_mask

coarse = compute_vsa_coarse_attention(q, k, v, tile_numel=64, top_k=78)
block_mask = create_vsa_block_mask(
    coarse.topk_indices,
    tile_numel=64,
    num_kv_tiles=k.shape[-2] // 64,
)
out = flex_attention(q, k, v, block_mask=block_mask)

attn_gym.masks.vsa.compute_vsa_coarse_attention(query, key, value, tile_numel, top_k, scale=None, sort_indices=True, include_self=False, q_variable_block_sizes=None, kv_variable_block_sizes=None) #

Run VSA's coarse attention pass and extract fine-pass top-k tiles.

Parameters:

Name Type Description Default
query Tensor

Query tensor with shape (B, H, S_Q, D) in tile-major order.

 required
key Tensor

Key tensor with shape (B, H, S_KV, D) in tile-major order.

 required
value Tensor

Value tensor with shape (B, H, S_KV, D) in tile-major order.

 required
tile_numel int

Number of fine tokens in one VSA tile.

 required
top_k int

Number of KV tiles selected for each query tile.

 required
scale float | None

Optional scale applied to the coarse dot products. Defaults to 1 / sqrt(D).

 None
sort_indices bool

Whether to sort selected KV tile ids ascending before returning.

 True
include_self bool

Whether to force each query tile to include the same-id KV tile.

 False
q_variable_block_sizes Tensor | None

Optional real-token counts for padded query tiles.

 None
kv_variable_block_sizes Tensor | None

Optional real-token counts for padded KV tiles.

 None

Returns:

Type Description
VSACoarseResult

VSACoarseResult containing top-k indices, tile-level coarse output, and
VSACoarseResult

coarse log-sum-exp.

Raises:

Type Description
ValueError

If value shape is incompatible or top_k is invalid.

attn_gym.masks.vsa.create_vsa_block_mask(topk_indices, tile_numel, num_kv_tiles, variable_block_sizes=None) #

Build the runtime FlexAttention BlockMask for VSA's fine sparse pass.

Parameters:

Name Type Description Default
topk_indices Tensor

Selected KV tile ids with shape (B, H, Q_TILES, TOP_K).

 required
tile_numel int

Number of fine tokens in one VSA tile.

 required
num_kv_tiles int

Total number of KV tiles before top-k pruning.

 required
variable_block_sizes Tensor | None

Optional real-token counts for each padded KV tile.

 None

Returns:

Type Description
BlockMask

A directly constructed block-sparse mask. Full selected tiles use
BlockMask

full_kv_* metadata; padded edge tiles use partial metadata plus a padding
BlockMask

predicate. The top-k membership itself is encoded by block metadata.

Raises:

Type Description
ValueError

If rank or static shape arguments are invalid.

attn_gym.masks.vsa.create_vsa_flash_block_mask(topk_indices, tile_numel, num_kv_tiles, kv_block_size=128, variable_block_sizes=None) #

Build the VSA block representation required by Flex FLASH.

Parameters:

Name Type Description Default
topk_indices Tensor

Selected KV tile ids with shape (B, H, Q_TILES, TOP_K).

 required
tile_numel int

Number of fine tokens in one VSA tile.

 required
num_kv_tiles int

Total number of KV tiles before top-k pruning.

 required
kv_block_size int

KV block size of the returned mask. FA4 requires this to equal its KV tile size (tile_n, 128 on SM100 for head_dim 128), so each selected VSA tile is expanded into tile_numel // kv_block_size consecutive KV blocks. The Q block size stays tile_numel because FA4 only requires it to be a multiple of its effective Q tile.

 128
variable_block_sizes Tensor | None

Optional real-token counts for each padded KV tile. Fully valid FA4 KV subblocks are emitted as full_kv_* blocks, while partially valid subblocks are emitted as normal kv_* blocks so the padding mask_mod is applied inside Flex FLASH.

 None

Returns:

Type Description
BlockMask

A BlockMask suitable for kernel_options={"BACKEND": "FLASH"}.

Raises:

Type Description
ValueError

If rank or static shape arguments are invalid.

attn_gym.masks.vsa.generate_vsa_mask_mod(topk_indices, tile_numel, variable_block_sizes=None) #

Create a VSA mask_mod from precomputed top-k tile indices.

Parameters:

Name Type Description Default
topk_indices Tensor

Selected KV tile ids with shape (B, H, Q_TILES, TOP_K).

 required
tile_numel int

Number of fine tokens in one VSA tile.

 required
variable_block_sizes Tensor | None

Optional real-token counts for each padded KV tile.

 None

Returns:

Type Description
_mask_mod_signature

A mask_mod that allows a query token to attend to KV tokens whose tile
_mask_mod_signature

appears in that query tile's top-k list and, if provided, excludes padded KV
_mask_mod_signature

positions in edge tiles.

Raises:

Type Description
ValueError

If topk_indices has the wrong rank or tile_numel is invalid.

Batchify#

Groups tokens into batches where attention is only allowed within the same group.

attn_gym.masks.batchify.batchify_mask_mod(mask_mod, batchify_size) #

Given arbirary mask_mod, batchify it to only allow attention within the same batch.

Parameters:

Name Type Description Default
mask_mod _mask_mod_signature

The mask mod to apply within each batch.

 required
batchify_size int

The number of tokens in each batch.

 required

Flamingo Cross-Attention#

Cross-attention mask for Flamingo-style vision-language models.

attn_gym.masks.flamingo.generate_vision_cross_attention_mask_mod(intervals, image_token_length) #

Generates a mask mod for VisionCrossAttention.

Parameters:

Name Type Description Default
intervals Tensor

Tensor of shape (num_images, 2) containing the start and end indices for each image.

 required
image_token_length int

Number of tokens per image.

 required

Sparse VideoGen#

Spatial and temporal attention masks following the Sparse VideoGen paper.

attn_gym.masks.svg.generate_spatial_head_mask_mod(prompt_length=226, num_frames=13, token_per_frame=1350, width=2, attn_sink=False, round_width=128) #

Generates a spatial head mask as specified in SVG. The same mask can also be used for temporal attention mask after applying the layout transformation by setting attn_sink to False.

Parameters:

Name Type Description Default
prompt_length int

The length of the prompt.

 226
num_frames int

The number of frames in the video.

 13
token_per_frame int

The number of tokens per frame.

 1350
width int

The width of the spatial head mask, determine the number of frames that can be attended.

 2
attn_sink bool

Whether to use the attention sink for the first column.

 False
round_width int

The number to round to for better hardware utilization, usually set to 128.

 128

attn_gym.masks.svg.generate_temporal_head_mask_mod(prompt_length=226, num_frames=13, token_per_frame=1350, width=2) #

Generates a temporal head mask as specified in SVG.

Parameters:

Name Type Description Default
prompt_length int

The length of the prompt.

 226
num_frames int

The number of frames in the video.

 13
token_per_frame int

The number of tokens per frame.

 1350
width int

The width of the temporal head mask, determine the number of frames that can be attended.

 2