windowed_attention ¶

Chunked non-causal sliding-window self-attention for MLX scaling studies.

This module is intentionally standalone (mlx.core + stdlib only) so it can be micro-benchmarked without pulling in the DiT / FastVideo stack.

Window policy¶

Symmetric sliding window (non-causal). For query index i the allowed key indices are:

sinks:   ``j in [0, sink)``  (always visible to every query, if ``sink > 0``)
local:   ``j in [max(0, i - half), min(S, i + half + 1))``
         where ``half = window // 2``

so each query sees roughly window + 1 local keys (plus any sinks outside that range). This is appropriate for a dense, bidirectional DiT denoise pass.

Implementation note (FLOPs)¶

A full-size additive attention mask still materialises an O(S^2) score matrix inside SDPA and does not reduce work. Instead we tile the sequence into query blocks and run mx.fast.scaled_dot_product_attention only against the union of keys that block needs (local slice ± sinks). That makes per-block work O(chunk * (window + sink) * D) and total work O(S * (window + sink) * D).

Functions:¶

fastvideo.mlx_runtime.windowed_attention.full_attention ¶

full_attention(q: array, k: array, v: array, scale: float | None = None) -> array

Compute dense scaled dot-product attention over the full sequence.

Parameters:

Name	Type	Description	Default
`scale`	`float`	Attention scaling factor. If omitted, uses the inverse square root of the head dimension.	`None`

Returns:

Type	Description
`array`	mx.array: Attention output with shape `(B, H, S, D)`.

Source code in fastvideo/mlx_runtime/windowed_attention.py

def full_attention(
    q: mx.array,
    k: mx.array,
    v: mx.array,
    scale: float | None = None,
) -> mx.array:
    """
    Compute dense scaled dot-product attention over the full sequence.

    Parameters:
        scale (float, optional): Attention scaling factor. If omitted, uses the
            inverse square root of the head dimension.

    Returns:
        mx.array: Attention output with shape ``(B, H, S, D)``.
    """
    _, _, _, d = _validate_qkv(q, k, v)
    sc = _default_scale(d, scale)
    return mx.fast.scaled_dot_product_attention(q, k, v, scale=sc)

fastvideo.mlx_runtime.windowed_attention.windowed_attention ¶

windowed_attention(q: array, k: array, v: array, window: int, sink: int = 0, scale: float | None = None, *, chunk_size: int | None = None) -> array

Apply symmetric sliding-window self-attention with optional global sink positions.

Parameters:

Name	Type	Description	Default
`q`	`array`	Query tensor shaped `(B, H, S, D)`.	required
`k`	`array`	Key tensor shaped `(B, H, S, D)`.	required
`v`	`array`	Value tensor shaped `(B, H, S, D)`.	required
`window`	`int`	Symmetric attention window width in tokens; must be at least 1.	required
`sink`	`int`	Number of leading key positions available to every query; must be between 0 and the sequence length.	`0`
`scale`	`Optional[float]`	Softmax scale. Defaults to `1 / sqrt(D)`.	`None`
`chunk_size`	`Optional[int]`	Query block length used for chunked processing. Defaults to the smaller of `window` and 512.	`None`

Returns:

Type	Description
`array`	mx.array: Attention output with the same shape as `q`.

Raises:

Type	Description
`ValueError`	If the inputs or attention parameters are invalid.
`RuntimeError`	If a query block has no available keys.

Source code in fastvideo/mlx_runtime/windowed_attention.py

def windowed_attention(
    q: mx.array,
    k: mx.array,
    v: mx.array,
    window: int,
    sink: int = 0,
    scale: float | None = None,
    *,
    chunk_size: int | None = None,
) -> mx.array:
    """
    Apply symmetric sliding-window self-attention with optional global sink positions.

    Parameters:
        q (mx.array): Query tensor shaped `(B, H, S, D)`.
        k (mx.array): Key tensor shaped `(B, H, S, D)`.
        v (mx.array): Value tensor shaped `(B, H, S, D)`.
        window (int): Symmetric attention window width in tokens; must be at least 1.
        sink (int): Number of leading key positions available to every query; must
            be between 0 and the sequence length.
        scale (Optional[float]): Softmax scale. Defaults to `1 / sqrt(D)`.
        chunk_size (Optional[int]): Query block length used for chunked processing.
            Defaults to the smaller of `window` and 512.

    Returns:
        mx.array: Attention output with the same shape as `q`.

    Raises:
        ValueError: If the inputs or attention parameters are invalid.
        RuntimeError: If a query block has no available keys.
    """
    _, _, seq_len, d = _validate_qkv(q, k, v)

    if window < 1:
        raise ValueError(f"window must be >= 1, got {window}")
    if sink < 0:
        raise ValueError(f"sink must be >= 0, got {sink}")
    if sink > seq_len:
        raise ValueError(f"sink ({sink}) cannot exceed sequence length ({seq_len})")

    sc = _default_scale(d, scale)
    half = window // 2

    # When the requested window is at least the sequence length, every query can
    # see every key under a symmetric policy — fall back to one dense SDPA.
    # (Sinks are redundant once the full key set is used.)
    if window >= seq_len:
        return mx.fast.scaled_dot_product_attention(q, k, v, scale=sc)

    chunk = min(window, 512) if chunk_size is None else int(chunk_size)
    if chunk < 1:
        raise ValueError(f"chunk_size must be >= 1, got {chunk}")
    chunk = min(chunk, seq_len)

    outputs: list[mx.array] = []
    for qs in range(0, seq_len, chunk):
        qe = min(seq_len, qs + chunk)
        q_block = q[:, :, qs:qe, :]
        ranges = _key_ranges_for_block(qs, qe, seq_len, half, sink)
        k_block, v_block = _concat_kv_slices(k, v, ranges)
        if k_block.shape[2] == 0:
            raise RuntimeError(f"empty key set for query block [{qs}, {qe}) with window={window}, sink={sink}")
        query_positions = mx.arange(qs, qe)[:, None]
        key_positions = mx.array([position for start, end in ranges for position in range(start, end)])[None, :]
        mask = mx.abs(query_positions - key_positions) <= half
        if sink > 0:
            mask = mask | (key_positions < sink)
        out_block = mx.fast.scaled_dot_product_attention(q_block, k_block, v_block, scale=sc, mask=mask)
        outputs.append(out_block)

    return mx.concatenate(outputs, axis=2)