Skip to content

bsa_attn

Bidirectional Sparse Attention (BSA) backend for FastVideo.

Pure-PyTorch reference implementation from: "Bidirectional Sparse Attention for Faster Video Diffusion Training" (arXiv:2509.01085)

BSA sparsifies both queries (pruning redundant tokens per block) and key-value pairs (keeping only relevant KV blocks per query block).

This is a training-free inference backend: it works with any model trained with full attention by applying BSA sparsity at inference time.

Classes

fastvideo.attention.backends.bsa_attn.BSAAttentionImpl 

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

Bases: AttentionImpl

Source code in fastvideo/attention/backends/bsa_attn.py 
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
def __init__(
    self,
    num_heads: int,
    head_size: int,
    causal: bool,
    softmax_scale: float,
    num_kv_heads: int | None = None,
    prefix: str = "",
    **extra_impl_args,
) -> None:
    self.prefix = prefix
    self.num_heads = num_heads
    self.head_size = head_size
    if num_kv_heads is not None and num_kv_heads != num_heads:
        raise ValueError("BSA backend does not support grouped-query attention")
    if causal:
        raise ValueError("BSA backend is bidirectional; causal=True is unsupported")
    if softmax_scale is not None:
        expected_scale = 1.0 / math.sqrt(self.head_size)
        if not math.isclose(softmax_scale, expected_scale, rel_tol=1e-4, abs_tol=1e-5):
            raise ValueError("softmax_scale must be default (1/sqrt(d)) for BSA")
    try:
        sp_group = get_sp_group()
        self.sp_size = sp_group.world_size
    except (AssertionError, RuntimeError):
        self.sp_size = 1

Methods:

fastvideo.attention.backends.bsa_attn.BSAAttentionImpl.forward 
forward(query: Tensor, key: Tensor, value: Tensor, attn_metadata: BSAAttentionMetadata) -> Tensor

BSA attention forward pass.

Input tensors are already in tile-contiguous order from preprocess_qkv.

Parameters:

Name Type Description Default
query Tensor

[B, L, num_heads, D] (tile-ordered)

 required
key Tensor

[B, L, num_heads, D] (tile-ordered)

 required
value Tensor

[B, L, num_heads, D] (tile-ordered)

 required
attn_metadata BSAAttentionMetadata

BSA metadata

 required

Returns:

Name Type Description
output Tensor

[B, L, num_heads, D] (tile-ordered)
Source code in fastvideo/attention/backends/bsa_attn.py 
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
def forward(
    self,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attn_metadata: BSAAttentionMetadata,
) -> torch.Tensor:
    """
    BSA attention forward pass.

    Input tensors are already in tile-contiguous order from preprocess_qkv.

    Args:
        query: [B, L, num_heads, D] (tile-ordered)
        key:   [B, L, num_heads, D] (tile-ordered)
        value: [B, L, num_heads, D] (tile-ordered)
        attn_metadata: BSA metadata

    Returns:
        output: [B, L, num_heads, D] (tile-ordered)
    """
    B, L, H, D = query.shape
    block_size = attn_metadata.block_size
    num_blocks = attn_metadata.num_blocks
    assert num_blocks * block_size == L, "Sequence length must match tiling"

    # Reshape to [B, H, L, D] for attention computation
    q = query.transpose(1, 2).contiguous()  # [B, H, L, D]
    k = key.transpose(1, 2).contiguous()
    v = value.transpose(1, 2).contiguous()

    # Reshape into blocks: [B, H, num_blocks, block_size, D]
    q_blocks = q.view(B, H, num_blocks, block_size, D)
    k_blocks = k.view(B, H, num_blocks, block_size, D)
    v_blocks = v.view(B, H, num_blocks, block_size, D)

    # --- Query sparsification ---
    sparse_q, keep_indices, keep_size = _prune_queries(q_blocks, attn_metadata.query_keep_ratio)

    # --- KV block selection ---
    kv_mask = _select_kv_blocks(
        sparse_q,
        k_blocks,
        attn_metadata.kv_cumulative_threshold,
        attn_metadata.min_kv_blocks,
    )

    # --- Sparse attention ---
    sparse_output = _compute_sparse_attention(sparse_q, k_blocks, v_blocks, kv_mask)

    # --- Reconstruct pruned positions ---
    full_output = _reconstruct_pruned(sparse_output, keep_indices, block_size)

    # Reshape back: [B, H, num_blocks, block_size, D] -> [B, H, L, D] -> [B, L, H, D]
    hidden_states = full_output.view(B, H, L, D).transpose(1, 2)

    return hidden_states
fastvideo.attention.backends.bsa_attn.BSAAttentionImpl.postprocess_output 
postprocess_output(output: Tensor, attn_metadata: BSAAttentionMetadata) -> Tensor

Reorder tokens from tile-contiguous order back to raster order.

Source code in fastvideo/attention/backends/bsa_attn.py 
676
677
678
679
680
681
682
def postprocess_output(
    self,
    output: torch.Tensor,
    attn_metadata: BSAAttentionMetadata,
) -> torch.Tensor:
    """Reorder tokens from tile-contiguous order back to raster order."""
    return output[:, attn_metadata.reverse_tile_partition_indices]
fastvideo.attention.backends.bsa_attn.BSAAttentionImpl.preprocess_qkv 
preprocess_qkv(qkv: Tensor, attn_metadata: BSAAttentionMetadata) -> Tensor

Reorder tokens from raster order to tile-contiguous order.

Source code in fastvideo/attention/backends/bsa_attn.py 
667
668
669
670
671
672
673
674
def preprocess_qkv(
    self,
    qkv: torch.Tensor,
    attn_metadata: BSAAttentionMetadata,
) -> torch.Tensor:
    """Reorder tokens from raster order to tile-contiguous order."""
    # qkv: [B, L, num_heads, D]
    return qkv[:, attn_metadata.tile_partition_indices]

Functions:

fastvideo.attention.backends.bsa_attn.get_reverse_tile_partition_indices cached 

get_reverse_tile_partition_indices(dit_seq_shape: tuple[int, int, int], tile_size: tuple[int, int, int], device: device) -> LongTensor

Inverse mapping: tile-contiguous order back to raster order.

Source code in fastvideo/attention/backends/bsa_attn.py 
73
74
75
76
77
78
79
80
@functools.lru_cache(maxsize=10)
def get_reverse_tile_partition_indices(
    dit_seq_shape: tuple[int, int, int],
    tile_size: tuple[int, int, int],
    device: torch.device,
) -> torch.LongTensor:
    """Inverse mapping: tile-contiguous order back to raster order."""
    return torch.argsort(get_tile_partition_indices(dit_seq_shape, tile_size, device))

fastvideo.attention.backends.bsa_attn.get_tile_partition_indices cached 

get_tile_partition_indices(dit_seq_shape: tuple[int, int, int], tile_size: tuple[int, int, int], device: device) -> LongTensor

Map raster-order tokens to tile-contiguous order.

Source code in fastvideo/attention/backends/bsa_attn.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
@functools.lru_cache(maxsize=10)
def get_tile_partition_indices(
    dit_seq_shape: tuple[int, int, int],
    tile_size: tuple[int, int, int],
    device: torch.device,
) -> torch.LongTensor:
    """Map raster-order tokens to tile-contiguous order."""
    T, H, W = dit_seq_shape
    ts, hs, ws = tile_size
    indices = torch.arange(T * H * W, device=device, dtype=torch.long).reshape(T, H, W)
    ls = []
    for t in range(math.ceil(T / ts)):
        for h in range(math.ceil(H / hs)):
            for w in range(math.ceil(W / ws)):
                ls.append(indices[
                    t * ts:min(t * ts + ts, T),
                    h * hs:min(h * hs + hs, H),
                    w * ws:min(w * ws + ws, W),
                ].flatten())
    return torch.cat(ls, dim=0)