Attention

attention

Kind

Bases: Enum

Discrete head-pattern label.

causal_mask

causal_mask(seq_len: int, offset: int = 0, dtype: Dtype = float32) -> array

Create a causal (lower-triangular) attention mask.

Parameters:

Name	Type	Description	Default
seq_len	int	Sequence length.	required
offset	int	Offset for KV cache (total KV length = offset + seq_len).	0
dtype	Dtype	Output dtype. Masked positions are -inf.	float32

Returns:

Type	Description
array	Mask of shape (1, 1, seq_len, offset + seq_len).

sliding_window_mask

sliding_window_mask(seq_len: int, window_size: int, offset: int = 0, dtype: Dtype = float32) -> array

Create a sliding window causal attention mask.

Each position can attend to at most window_size previous positions (including itself).

Parameters:

Name	Type	Description	Default
seq_len	int	Sequence length.	required
window_size	int	Size of the attention window.	required
offset	int	Offset for KV cache.	0
dtype	Dtype	Output dtype. Masked positions are -inf.	float32

Returns:

Type	Description
array	Mask of shape (1, 1, seq_len, offset + seq_len).

block_contiguous_permutation

block_contiguous_permutation(scores: array, *, block_size: int, descending: bool = True) -> tuple[array, array]

Sort tokens by score so high-importance ones cluster into early blocks.

Parameters:

Name	Type	Description	Default
scores	array	(S,) per-token importance score.	required
block_size	int	Block size of the downstream sparse kernel. Informational only — this function does not pad S to a multiple of block_size. The caller is responsible if they need exact alignment.	required
descending	bool	If True (default), highest scores land at position 0. If False, lowest first.	True

Returns:

Type	Description
array	(perm, inv_perm) of shape (S,), dtype int32.
array	perm[i] is the original index of the token now at position i.
tuple[array, array]	inv_perm[i] is the new position of the token originally at i.
tuple[array, array]	Tie-breaking among equal scores is stable (preserves original order),
tuple[array, array]	which is the MLX argsort contract.

invert_permutation

invert_permutation(perm: array) -> array

Compute the inverse of a 1D permutation.

Equivalent to mx.argsort(perm). Caller is responsible for ensuring perm is a valid permutation of [0, S); misuse silently produces wrong results.

Parameters:

Name	Type	Description	Default
perm	array	(S,) int permutation array.	required

Returns:

Type	Description
array	(S,) int32 inverse permutation.

classify

classify(scores: array, *, spatial_threshold: float = 0.5, temporal_threshold: float = 0.5) -> list[Kind]

Convert raw per-head mass scores to discrete labels.

Parameters:

Name	Type	Description	Default
scores	array	(num_heads, 2) array. Column 0 = mass on same-frame keys, column 1 = mass on same-spatial-position keys.	required
spatial_threshold	float	Min column-0 mass to label a head SPATIAL.	0.5
temporal_threshold	float	Min column-1 mass to label a head TEMPORAL.	0.5

Returns:

Type	Description
list[Kind]	List of Kind of length num_heads. Tie-break: if both columns
list[Kind]	exceed their threshold, SPATIAL wins.

classify_heads_from_probs

classify_heads_from_probs(probs: array, T: int, H: int, W: int) -> array

Per-head attention-mass fractions on same-frame and same-position keys.

Uses all queries (no sampling) — assumes the caller has already paid the cost of materializing (B, num_heads, S, S) softmaxed probabilities.

Parameters:

Name	Type	Description	Default
probs	array	(B, num_heads, S, S) attention probabilities. The caller is responsible for ensuring these are valid (non-negative, sum-1 along the key axis); not validated here because it would cost O(S²).	required
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required

Returns:

Type	Description
array	(num_heads, 2) float array. Column 0 = mass on same-frame keys,
array	column 1 = mass on same-spatial-position keys. Averaged over batch
array	and queries.

classify_heads_from_qk

classify_heads_from_qk(q: array, k: array, T: int, H: int, W: int, *, n_samples: int = 64, key: array | None = None) -> array

Per-head attention-mass fractions, sampled from Q,K.

Avoids materializing the full (B, num_heads, S, S) attention by sampling n_samples queries uniformly per call. Reproducible: with a fixed key, returns identical results.

Parameters:

Name	Type	Description	Default
q	array	(B, num_heads, S, D) query tensor.	required
k	array	(B, num_heads, S, D) key tensor, same shape as q.	required
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
n_samples	int	How many queries to sample uniformly per call. Must satisfy 0 < n_samples <= S.	64
key	array | None	mx.random key for the sampler. If None, uses mx.random.key(0) so default behavior is deterministic across calls. Callers who want variance must pass their own key.	None

Returns:

Type	Description
array	(num_heads, 2) float array. Column 0 = mass on same-frame keys,
array	column 1 = mass on same-spatial-position keys.

frame_stride_diagonal_mask

frame_stride_diagonal_mask(T: int, H: int, W: int, *, num_diagonals: int, dtype: Dtype = float32) -> array

Multi-diagonal mask at frame-stride offsets (Sparse-vDiT M3).

Token at flat index i attends to token at flat index j iff (j - i) is a multiple of the per-frame stride H*W in {-(k-1)*HW, ..., -HW, 0, HW, ..., (k-1)*HW} where k = num_diagonals. Captures the "multi-diagonal" head pattern from Sparse-vDiT (Chen et al. 2025): same (h, w) across nearby frames.

Setting num_diagonals=1 reduces to the main diagonal (self-attention only). Setting num_diagonals=T is equivalent to temporal_only_mask.

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
num_diagonals	int	Strictly positive number of diagonal bands (counting the main diagonal once).	required
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

radial_box_mask

radial_box_mask(T: int, H: int, W: int, *, radius_t: int, radius_s: float, dtype: Dtype = float32) -> array

Hard-cutoff radial spatiotemporal mask.

Query (t, h, w) attends to (t', h', w') iff |t-t'| <= radius_t AND sqrt((h-h')**2 + (w-w')**2) <= radius_s.

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
radius_t	int	Non-negative temporal radius (frames, inclusive).	required
radius_s	float	Non-negative Euclidean spatial radius (latent units).	required
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

radial_gaussian_mask

radial_gaussian_mask(T: int, H: int, W: int, *, sigma_t: float, sigma_s: float, cutoff: float = -6.0, dtype: Dtype = float32) -> array

Exponential-decay radial mask (dense log-weights).

Value at (i, j) is -(dt**2 / (2 sigma_t**2) + ds**2 / (2 sigma_s**2)) where ds**2 = (h-h')**2 + (w-w')**2. Values below cutoff are clamped to -inf so the mask is usable in fp16 without underflow.

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
sigma_t	float	Temporal scale, strictly positive.	required
sigma_s	float	Spatial scale, strictly positive.	required
cutoff	float	Strictly negative log-weight floor; values below are replaced by -inf. Default -6.0 ≈ exp(-6) ≈ 0.0025.	-6.0
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

sliding_tile_block_mask

sliding_tile_block_mask(T: int, H: int, W: int, *, tile: tuple[int, int, int], window: tuple[int, int, int] = (1, 1, 1), dtype: Dtype = float32) -> array

Tile-block sliding attention (STA, ICML 2025).

Tokens are grouped into non-overlapping tiles of shape tile = (tt, th, tw). Every query in a tile attends to all keys in the ±window neighboring tiles (window in tile units, inclusive).

Parameters:

Name	Type	Description	Default
T	int	Number of frames. Must be divisible by tile[0].	required
H	int	Latent height. Must be divisible by tile[1].	required
W	int	Latent width. Must be divisible by tile[2].	required
tile	tuple[int, int, int]	(tt, th, tw) tile dims, all positive.	required
window	tuple[int, int, int]	(wt, wh, ww) non-negative tile-unit radii.	(1, 1, 1)
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

sliding_tile_centered_mask

sliding_tile_centered_mask(T: int, H: int, W: int, *, window: tuple[int, int, int], dtype: Dtype = float32) -> array

Per-query centered spatiotemporal window mask.

Token (t, h, w) attends to (t', h', w') iff |t-t'| <= window[0] AND |h-h'| <= window[1] AND |w-w'| <= window[2].

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
window	tuple[int, int, int]	(dt, dh, dw) non-negative inclusive radii.	required
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

spatial_only_mask

spatial_only_mask(T: int, H: int, W: int, *, dtype: Dtype = float32) -> array

Mask that restricts attention to tokens in the same frame.

Each token at frame t attends only to other tokens whose frame index equals t. Captures the "spatial-locality" head pattern from Sparse VideoGen.

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

temporal_only_mask

temporal_only_mask(T: int, H: int, W: int, *, dtype: Dtype = float32) -> array

Mask that restricts attention to tokens at the same spatial position.

Each token at (h, w) attends only to tokens whose (h, w) matches, across all frames. Captures the "temporal-locality" head pattern from Sparse VideoGen.

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

vertical_stripe_mask

vertical_stripe_mask(T: int, H: int, W: int, *, key_indices: array, dtype: Dtype = float32) -> array

Anchor-column mask (Sparse-vDiT M4).

Every query attends only to a fixed set of "sink" key tokens identified by key_indices (flat indices into the T-major sequence). Captures the "vertical-stripe" head pattern from Sparse-vDiT, where a small set of anchor positions act as global memory.

The set must be non-empty and contain unique in-range indices. The main diagonal is not added automatically — include it in key_indices if self-attention is desired.

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
key_indices	array	1-D mx.array of int indices in [0, T*H*W).	required
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

masks

Attention mask utilities.

causal_mask

causal_mask(seq_len: int, offset: int = 0, dtype: Dtype = float32) -> array

Create a causal (lower-triangular) attention mask.

Parameters:

Name	Type	Description	Default
seq_len	int	Sequence length.	required
offset	int	Offset for KV cache (total KV length = offset + seq_len).	0
dtype	Dtype	Output dtype. Masked positions are -inf.	float32

Returns:

Type	Description
array	Mask of shape (1, 1, seq_len, offset + seq_len).

sliding_window_mask

sliding_window_mask(seq_len: int, window_size: int, offset: int = 0, dtype: Dtype = float32) -> array

Create a sliding window causal attention mask.

Each position can attend to at most window_size previous positions (including itself).

Parameters:

Name	Type	Description	Default
seq_len	int	Sequence length.	required
window_size	int	Size of the attention window.	required
offset	int	Offset for KV cache.	0
dtype	Dtype	Output dtype. Masked positions are -inf.	float32

Returns:

Type	Description
array	Mask of shape (1, 1, seq_len, offset + seq_len).

permute

Block-contiguous token permutation (SVG2 semantic permutation).

Reorders a sequence of tokens so that high-importance ones fall into the first contiguous blocks, which is what block-sparse attention kernels actually need to realize their savings. Pair with mx.take(x, perm, axis=...) to permute Q/K/V tensors and mx.take(y, inv_perm, axis=...) to undo.

block_contiguous_permutation

block_contiguous_permutation(scores: array, *, block_size: int, descending: bool = True) -> tuple[array, array]

Sort tokens by score so high-importance ones cluster into early blocks.

Parameters:

Name	Type	Description	Default
scores	array	(S,) per-token importance score.	required
block_size	int	Block size of the downstream sparse kernel. Informational only — this function does not pad S to a multiple of block_size. The caller is responsible if they need exact alignment.	required
descending	bool	If True (default), highest scores land at position 0. If False, lowest first.	True

Returns:

Type	Description
array	(perm, inv_perm) of shape (S,), dtype int32.
array	perm[i] is the original index of the token now at position i.
tuple[array, array]	inv_perm[i] is the new position of the token originally at i.
tuple[array, array]	Tie-breaking among equal scores is stable (preserves original order),
tuple[array, array]	which is the MLX argsort contract.

invert_permutation

invert_permutation(perm: array) -> array

Compute the inverse of a 1D permutation.

Equivalent to mx.argsort(perm). Caller is responsible for ensuring perm is a valid permutation of [0, S); misuse silently produces wrong results.

Parameters:

Name	Type	Description	Default
perm	array	(S,) int permutation array.	required

Returns:

Type	Description
array	(S,) int32 inverse permutation.

profile

Head-pattern profiler for video DiTs.

Classify each attention head as SPATIAL (mass concentrated on same-frame keys), TEMPORAL (same-position cross-frame), or OTHER (neither).

All functions assume T-major token flattening — same convention as mlx_arsenal.attention.video_masks: tokens flatten as [t0(h0w0..hHwW), t1(...), ..., tT(...)] to a sequence of length S = T*H*W.

Kind

Bases: Enum

Discrete head-pattern label.

classify

classify(scores: array, *, spatial_threshold: float = 0.5, temporal_threshold: float = 0.5) -> list[Kind]

Convert raw per-head mass scores to discrete labels.

Parameters:

Name	Type	Description	Default
scores	array	(num_heads, 2) array. Column 0 = mass on same-frame keys, column 1 = mass on same-spatial-position keys.	required
spatial_threshold	float	Min column-0 mass to label a head SPATIAL.	0.5
temporal_threshold	float	Min column-1 mass to label a head TEMPORAL.	0.5

Returns:

Type	Description
list[Kind]	List of Kind of length num_heads. Tie-break: if both columns
list[Kind]	exceed their threshold, SPATIAL wins.

classify_heads_from_probs

classify_heads_from_probs(probs: array, T: int, H: int, W: int) -> array

Per-head attention-mass fractions on same-frame and same-position keys.

Uses all queries (no sampling) — assumes the caller has already paid the cost of materializing (B, num_heads, S, S) softmaxed probabilities.

Parameters:

Name	Type	Description	Default
probs	array	(B, num_heads, S, S) attention probabilities. The caller is responsible for ensuring these are valid (non-negative, sum-1 along the key axis); not validated here because it would cost O(S²).	required
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required

Returns:

Type	Description
array	(num_heads, 2) float array. Column 0 = mass on same-frame keys,
array	column 1 = mass on same-spatial-position keys. Averaged over batch
array	and queries.

classify_heads_from_qk

classify_heads_from_qk(q: array, k: array, T: int, H: int, W: int, *, n_samples: int = 64, key: array | None = None) -> array

Per-head attention-mass fractions, sampled from Q,K.

Avoids materializing the full (B, num_heads, S, S) attention by sampling n_samples queries uniformly per call. Reproducible: with a fixed key, returns identical results.

Parameters:

Name	Type	Description	Default
q	array	(B, num_heads, S, D) query tensor.	required
k	array	(B, num_heads, S, D) key tensor, same shape as q.	required
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
n_samples	int	How many queries to sample uniformly per call. Must satisfy 0 < n_samples <= S.	64
key	array | None	mx.random key for the sampler. If None, uses mx.random.key(0) so default behavior is deterministic across calls. Callers who want variance must pass their own key.	None

Returns:

Type	Description
array	(num_heads, 2) float array. Column 0 = mass on same-frame keys,
array	column 1 = mass on same-spatial-position keys.

video_masks

Spatiotemporal attention masks for video diffusion transformers.

All functions in this module assume T-major token flattening: a video tensor of shape (T, H, W) is flattened to S = T*H*W tokens in the order [t0(h0w0..hHwW), t1(...), ..., tT(...)]. This matches LTX-Video, CogVideoX, and the convention used by mlx_arsenal.spatial.patchify.

Each function returns a mask of shape (1, 1, S, S) with float values: 0.0 means the query is allowed to attend to the key, -inf means it is blocked. The shape broadcasts over batch and head axes expected by mx.fast.scaled_dot_product_attention.

For typical LTX latents (T=8, H=32, W=32 → S=8192) the mask is S² ≈ 67M entries. Use dtype=mx.float16 to halve memory.

spatial_only_mask

spatial_only_mask(T: int, H: int, W: int, *, dtype: Dtype = float32) -> array

Mask that restricts attention to tokens in the same frame.

Each token at frame t attends only to other tokens whose frame index equals t. Captures the "spatial-locality" head pattern from Sparse VideoGen.

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

temporal_only_mask

temporal_only_mask(T: int, H: int, W: int, *, dtype: Dtype = float32) -> array

Mask that restricts attention to tokens at the same spatial position.

Each token at (h, w) attends only to tokens whose (h, w) matches, across all frames. Captures the "temporal-locality" head pattern from Sparse VideoGen.

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

sliding_tile_centered_mask

sliding_tile_centered_mask(T: int, H: int, W: int, *, window: tuple[int, int, int], dtype: Dtype = float32) -> array

Per-query centered spatiotemporal window mask.

Token (t, h, w) attends to (t', h', w') iff |t-t'| <= window[0] AND |h-h'| <= window[1] AND |w-w'| <= window[2].

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
window	tuple[int, int, int]	(dt, dh, dw) non-negative inclusive radii.	required
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

sliding_tile_block_mask

sliding_tile_block_mask(T: int, H: int, W: int, *, tile: tuple[int, int, int], window: tuple[int, int, int] = (1, 1, 1), dtype: Dtype = float32) -> array

Tile-block sliding attention (STA, ICML 2025).

Tokens are grouped into non-overlapping tiles of shape tile = (tt, th, tw). Every query in a tile attends to all keys in the ±window neighboring tiles (window in tile units, inclusive).

Parameters:

Name	Type	Description	Default
T	int	Number of frames. Must be divisible by tile[0].	required
H	int	Latent height. Must be divisible by tile[1].	required
W	int	Latent width. Must be divisible by tile[2].	required
tile	tuple[int, int, int]	(tt, th, tw) tile dims, all positive.	required
window	tuple[int, int, int]	(wt, wh, ww) non-negative tile-unit radii.	(1, 1, 1)
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

radial_box_mask

radial_box_mask(T: int, H: int, W: int, *, radius_t: int, radius_s: float, dtype: Dtype = float32) -> array

Hard-cutoff radial spatiotemporal mask.

Query (t, h, w) attends to (t', h', w') iff |t-t'| <= radius_t AND sqrt((h-h')**2 + (w-w')**2) <= radius_s.

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
radius_t	int	Non-negative temporal radius (frames, inclusive).	required
radius_s	float	Non-negative Euclidean spatial radius (latent units).	required
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

frame_stride_diagonal_mask

frame_stride_diagonal_mask(T: int, H: int, W: int, *, num_diagonals: int, dtype: Dtype = float32) -> array

Multi-diagonal mask at frame-stride offsets (Sparse-vDiT M3).

Token at flat index i attends to token at flat index j iff (j - i) is a multiple of the per-frame stride H*W in {-(k-1)*HW, ..., -HW, 0, HW, ..., (k-1)*HW} where k = num_diagonals. Captures the "multi-diagonal" head pattern from Sparse-vDiT (Chen et al. 2025): same (h, w) across nearby frames.

Setting num_diagonals=1 reduces to the main diagonal (self-attention only). Setting num_diagonals=T is equivalent to temporal_only_mask.

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
num_diagonals	int	Strictly positive number of diagonal bands (counting the main diagonal once).	required
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

vertical_stripe_mask

vertical_stripe_mask(T: int, H: int, W: int, *, key_indices: array, dtype: Dtype = float32) -> array

Anchor-column mask (Sparse-vDiT M4).

Every query attends only to a fixed set of "sink" key tokens identified by key_indices (flat indices into the T-major sequence). Captures the "vertical-stripe" head pattern from Sparse-vDiT, where a small set of anchor positions act as global memory.

The set must be non-empty and contain unique in-range indices. The main diagonal is not added automatically — include it in key_indices if self-attention is desired.

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
key_indices	array	1-D mx.array of int indices in [0, T*H*W).	required
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).

radial_gaussian_mask

radial_gaussian_mask(T: int, H: int, W: int, *, sigma_t: float, sigma_s: float, cutoff: float = -6.0, dtype: Dtype = float32) -> array

Exponential-decay radial mask (dense log-weights).

Value at (i, j) is -(dt**2 / (2 sigma_t**2) + ds**2 / (2 sigma_s**2)) where ds**2 = (h-h')**2 + (w-w')**2. Values below cutoff are clamped to -inf so the mask is usable in fp16 without underflow.

Parameters:

Name	Type	Description	Default
T	int	Number of frames.	required
H	int	Latent height.	required
W	int	Latent width.	required
sigma_t	float	Temporal scale, strictly positive.	required
sigma_s	float	Spatial scale, strictly positive.	required
cutoff	float	Strictly negative log-weight floor; values below are replaced by -inf. Default -6.0 ≈ exp(-6) ≈ 0.0025.	-6.0
dtype	Dtype	Output dtype.	float32

Returns:

Type	Description
array	Additive mask of shape (1, 1, T*H*W, T*H*W).