Skip to content
FastVideo
Esc
↑↓navigate↵open⌘Jpreview
On this page

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)</code></pre></details>


fastvideo.attention.backends.nabla.NablaAttentionBackend


NablaAttentionBackend()


Source
class NablaAttentionBackend(AttentionBackend):

@staticmethod
def get_name() -&gt; str:
    return &quot;NABLA_ATTN&quot;

@staticmethod
def get_impl_cls() -&gt; type&#91;&quot;NablaAttentionImpl&quot;&#93;:
    return NablaAttentionImpl

@staticmethod
def get_metadata_cls() -&gt; type&#91;&quot;NablaAttentionMetadata&quot;&#93;:
    return NablaAttentionMetadata

@staticmethod
def get_builder_cls() -&gt; type&#91;&quot;NablaAttentionMetadataBuilder&quot;&#93;:
    return NablaAttentionMetadataBuilder</code></pre></details>


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) -&gt; None:
    pass

def prepare(self) -&gt; None:
    pass

def build(
    self,
    current_timestep: int,
    sta_mask: torch.Tensor,
    P: float,
    visual_shape: tuple&#91;int, int, int&#93;,
    **kwargs: Any,
) -&gt; NablaAttentionMetadata:
    return NablaAttentionMetadata(
        current_timestep=current_timestep,
        sta_mask=sta_mask,
        P=P,
        visual_shape=visual_shape,
    )</code></pre></details>


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 = &quot;&quot;,
    **extra_impl_args,
) -&gt; None:
    if not CAN_USE_FLEX_ATTN:
        raise RuntimeError(&quot;NABLA attention requires torch.nn.attention.flex_attention, &quot;
                           &quot;which is unavailable in this PyTorch build.&quot;)
    if causal:
        raise ValueError(&quot;NABLA attention does not support causal masking.&quot;)

def forward(
    self,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attn_metadata: NablaAttentionMetadata,
) -&gt; torch.Tensor:
    # q/k/v: &#91;B, S, heads, head_dim&#93;, 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)</code></pre></details>


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)

Was this page helpful?