--- title: nabla ---

NABLA block-sparse flex-attention backend (Kandinsky5 "nabla" checkpoints).

The block mask is data-dependent: nablaT_v2 mean-pools 64-token blocks of the fractal-ordered sequence, thresholds the softmaxed block map, and ORs it with a precomputed spatio-temporal-window (STA) mask carried on the attention metadata. The mask spans the full sequence, so this backend does not support sequence parallelism — use it via LocalAttention only.

fastvideo.attention.backends.nabla.AttentionBackend

AttentionBackend()

Abstract class for attention backends.

fastvideo.attention.backends.nabla.AttentionBackend.get_name

get_name() -> str

fastvideo.attention.backends.nabla.AttentionBackend.get_impl_cls

get_impl_cls() -> type[AttentionImpl]

fastvideo.attention.backends.nabla.AttentionBackend.get_metadata_cls

get_metadata_cls() -> type[AttentionMetadata]

fastvideo.attention.backends.nabla.AttentionBackend.get_builder_cls

get_builder_cls() -> type[AttentionMetadataBuilder]

fastvideo.attention.backends.nabla.AttentionImpl

AttentionImpl(num_heads: int, head_size: int, softmax_scale: float, causal: bool = False, num_kv_heads: int | None = None, prefix: str = '', **extra_impl_args)

fastvideo.attention.backends.nabla.AttentionImpl.preprocess_qkv

preprocess_qkv(qkv: torch.Tensor, attn_metadata: T) -> torch.Tensor

Preprocess QKV tensor before performing attention operation.

Default implementation returns the tensor unchanged. Subclasses can override this to implement custom preprocessing like reshaping, tiling, scaling, or other transformations.

Called AFTER all_to_all for distributed attention

Parameters

qkv: torch.Tensor

The query-key-value tensor

attn_metadata: T

Metadata for the attention operation

Returns

torch.Tensor

Processed QKV tensor

fastvideo.attention.backends.nabla.AttentionImpl.postprocess_output

postprocess_output(output: torch.Tensor, attn_metadata: T) -> torch.Tensor

Postprocess the output tensor after the attention operation.

Default implementation returns the tensor unchanged. Subclasses can override this to implement custom postprocessing like untiling, scaling, or other transformations.

Called BEFORE all_to_all for distributed attention

Parameters

output: torch.Tensor

The output tensor from the attention operation

attn_metadata: T

Metadata for the attention operation

Returns

torch.Tensor

Postprocessed output tensor

fastvideo.attention.backends.nabla.AttentionImpl.forward

forward(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: T) -> torch.Tensor

fastvideo.attention.backends.nabla.AttentionMetadata

AttentionMetadata(current_timestep: int, *, VSA_sparsity: float = 0.0)

Attention metadata for prefill and decode batched together.

fastvideo.attention.backends.nabla.AttentionMetadata.asdict_zerocopy

asdict_zerocopy(skip_fields: set[str] | None = None) -> dict[str, Any]

Similar to dataclasses.asdict, but avoids deepcopying.

fastvideo.attention.backends.nabla.AttentionMetadataBuilder

AttentionMetadataBuilder()

Abstract class for attention metadata builders.

Create the builder, remember some configuration and parameters.

fastvideo.attention.backends.nabla.AttentionMetadataBuilder.prepare

prepare() -> None

Prepare for one batch.

fastvideo.attention.backends.nabla.AttentionMetadataBuilder.build

build(**kwargs: Any) -> AttentionMetadata

Build attention metadata with on-device tensors.

fastvideo.attention.backends.nabla.nablaT_v2

nablaT_v2(q: torch.Tensor, k: torch.Tensor, sta: torch.Tensor, thr: float = 0.9) -> BlockMask
Source
def nablaT_v2(
    q: torch.Tensor,
    k: torch.Tensor,
    sta: torch.Tensor,
    thr: float = 0.9,
) -> "BlockMask":
    q = q.transpose(1, 2).contiguous()
    k = k.transpose(1, 2).contiguous()

    # Map estimation
    B, h, S, D = q.shape
    s1 = S // 64
    qa = q.reshape(B, h, s1, 64, D).mean(-2)
    ka = k.reshape(B, h, s1, 64, D).mean(-2).transpose(-2, -1)
    map = qa @ ka

    map = torch.softmax(map / math.sqrt(D), dim=-1)
    # Map binarization
    vals, inds = map.sort(-1)
    cvals = vals.cumsum_(-1)
    mask = (cvals >= 1 - thr).int()
    mask = mask.gather(-1, inds.argsort(-1))

    mask = torch.logical_or(mask, sta)

    # BlockMask creation
    kv_nb = mask.sum(-1).to(torch.int32)
    kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32)
    return BlockMask.from_kv_blocks(torch.zeros_like(kv_nb), kv_inds, kv_nb, kv_inds, BLOCK_SIZE=64, mask_mod=None)

fastvideo.attention.backends.nabla.NablaAttentionBackend

NablaAttentionBackend()
Source
class NablaAttentionBackend(AttentionBackend):

    @staticmethod
    def get_name() -> str:
        return "NABLA_ATTN"

    @staticmethod
    def get_impl_cls() -> type["NablaAttentionImpl"]:
        return NablaAttentionImpl

    @staticmethod
    def get_metadata_cls() -> type["NablaAttentionMetadata"]:
        return NablaAttentionMetadata

    @staticmethod
    def get_builder_cls() -> type["NablaAttentionMetadataBuilder"]:
        return NablaAttentionMetadataBuilder

fastvideo.attention.backends.nabla.NablaAttentionBackend.get_name

get_name() -> str
Source
@staticmethod
def get_name() -> str:
    return "NABLA_ATTN"

fastvideo.attention.backends.nabla.NablaAttentionBackend.get_impl_cls

get_impl_cls() -> type[NablaAttentionImpl]
Source
@staticmethod
def get_impl_cls() -> type["NablaAttentionImpl"]:
    return NablaAttentionImpl

fastvideo.attention.backends.nabla.NablaAttentionBackend.get_metadata_cls

get_metadata_cls() -> type[NablaAttentionMetadata]
Source
@staticmethod
def get_metadata_cls() -> type["NablaAttentionMetadata"]:
    return NablaAttentionMetadata

fastvideo.attention.backends.nabla.NablaAttentionBackend.get_builder_cls

get_builder_cls() -> type[NablaAttentionMetadataBuilder]
Source
@staticmethod
def get_builder_cls() -> type["NablaAttentionMetadataBuilder"]:
    return NablaAttentionMetadataBuilder

fastvideo.attention.backends.nabla.NablaAttentionMetadata

NablaAttentionMetadata(current_timestep: int, sta_mask: torch.Tensor = None, P: float = 0.9, visual_shape: tuple[int, int, int] = (0, 0, 0), *, VSA_sparsity: float = 0.0)
Source
@dataclass
class NablaAttentionMetadata(AttentionMetadata):
    # Block-level STA window mask [1, 1, S/64, S/64], precomputed once per run.
    sta_mask: torch.Tensor = None  # type: ignore[assignment]
    # Cumulative-probability threshold for block-map binarization.
    P: float = 0.9
    visual_shape: tuple[int, int, int] = (0, 0, 0)

fastvideo.attention.backends.nabla.NablaAttentionMetadataBuilder

NablaAttentionMetadataBuilder()
Source
class NablaAttentionMetadataBuilder(AttentionMetadataBuilder):

    def __init__(self) -> None:
        pass

    def prepare(self) -> None:
        pass

    def build(
        self,
        current_timestep: int,
        sta_mask: torch.Tensor,
        P: float,
        visual_shape: tuple[int, int, int],
        **kwargs: Any,
    ) -> NablaAttentionMetadata:
        return NablaAttentionMetadata(
            current_timestep=current_timestep,
            sta_mask=sta_mask,
            P=P,
            visual_shape=visual_shape,
        )

fastvideo.attention.backends.nabla.NablaAttentionMetadataBuilder.prepare

prepare() -> None
Source
def prepare(self) -> None:
    pass

fastvideo.attention.backends.nabla.NablaAttentionMetadataBuilder.build

build(current_timestep: int, sta_mask: torch.Tensor, P: float, visual_shape: tuple[int, int, int], **kwargs: Any) -> NablaAttentionMetadata
Source
def build(
    self,
    current_timestep: int,
    sta_mask: torch.Tensor,
    P: float,
    visual_shape: tuple[int, int, int],
    **kwargs: Any,
) -> NablaAttentionMetadata:
    return NablaAttentionMetadata(
        current_timestep=current_timestep,
        sta_mask=sta_mask,
        P=P,
        visual_shape=visual_shape,
    )

fastvideo.attention.backends.nabla.NablaAttentionImpl

NablaAttentionImpl(num_heads: int, head_size: int, softmax_scale: float, causal: bool = False, num_kv_heads: int | None = None, prefix: str = '', **extra_impl_args)
Source
class NablaAttentionImpl(AttentionImpl):

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        softmax_scale: float,
        causal: bool = False,
        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
        if not CAN_USE_FLEX_ATTN:
            raise RuntimeError("NABLA attention requires torch.nn.attention.flex_attention, "
                               "which is unavailable in this PyTorch build.")
        if causal:
            raise ValueError("NABLA attention does not support causal masking.")

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: NablaAttentionMetadata,
    ) -> torch.Tensor:
        # q/k/v: [B, S, heads, head_dim], fractal-ordered by the model; S % 64 == 0.
        block_mask = nablaT_v2(query, key, attn_metadata.sta_mask, thr=attn_metadata.P)
        return flex_attention(
            query=query.transpose(1, 2),
            key=key.transpose(1, 2),
            value=value.transpose(1, 2),
            block_mask=block_mask,
        ).transpose(1, 2)

fastvideo.attention.backends.nabla.NablaAttentionImpl.forward

forward(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: NablaAttentionMetadata) -> torch.Tensor
Source
def forward(
    self,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attn_metadata: NablaAttentionMetadata,
) -> torch.Tensor:
    # q/k/v: [B, S, heads, head_dim], fractal-ordered by the model; S % 64 == 0.
    block_mask = nablaT_v2(query, key, attn_metadata.sta_mask, thr=attn_metadata.P)
    return flex_attention(
        query=query.transpose(1, 2),
        key=key.transpose(1, 2),
        value=value.transpose(1, 2),
        block_mask=block_mask,
    ).transpose(1, 2)