Skip to content

video_sparse_attn

Classes

fastvideo.attention.backends.video_sparse_attn.VideoSparseAttentionImpl 

VideoSparseAttentionImpl(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/video_sparse_attn.py 
214
215
216
217
218
219
220
221
222
223
224
225
226
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
    sp_group = get_sp_group()
    self.sp_size = sp_group.world_size

Methods:

fastvideo.attention.backends.video_sparse_attn.VideoSparseAttentionImpl.preprocess_qkv 
preprocess_qkv(qkv: Tensor, attn_metadata: VideoSparseAttentionMetadata) -> Tensor

Tile QKV; aliasing contract: see tile().

Source code in fastvideo/attention/backends/video_sparse_attn.py 
272
273
274
275
276
277
278
def preprocess_qkv(
    self,
    qkv: torch.Tensor,
    attn_metadata: VideoSparseAttentionMetadata,
) -> torch.Tensor:
    """Tile QKV; aliasing contract: see ``tile()``."""
    return self.tile(qkv, attn_metadata)
fastvideo.attention.backends.video_sparse_attn.VideoSparseAttentionImpl.tile 
tile(x: Tensor, attn_metadata: VideoSparseAttentionMetadata) -> Tensor

Tile x into attn_metadata.tile_buf and return it.

The returned tensor aliases the per-metadata buffer and is only valid until the next tile() / preprocess_qkv call on the same attn_metadata. Callers must consume (or copy) the result before invoking another VSA layer with the same metadata. Today both call sites materialize copies via .transpose(...).contiguous() inside forward(), so the contract holds; future callers must preserve it.

Source code in fastvideo/attention/backends/video_sparse_attn.py 
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
def tile(self, x: torch.Tensor, attn_metadata: VideoSparseAttentionMetadata) -> torch.Tensor:
    """Tile ``x`` into ``attn_metadata.tile_buf`` and return it.

    The returned tensor aliases the per-metadata buffer and is only
    valid until the next ``tile()`` / ``preprocess_qkv`` call on the
    same ``attn_metadata``.  Callers must consume (or copy) the
    result before invoking another VSA layer with the same metadata.
    Today both call sites materialize copies via
    ``.transpose(...).contiguous()`` inside ``forward()``, so the
    contract holds; future callers must preserve it.
    """
    num_tiles = attn_metadata.num_tiles
    t_padded_size = num_tiles[0] * VSA_TILE_SIZE[0]
    h_padded_size = num_tiles[1] * VSA_TILE_SIZE[1]
    w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2]
    target_shape = (x.shape[0], t_padded_size * h_padded_size * w_padded_size, x.shape[-2], x.shape[-1])

    if not attn_metadata.cache_tile_buf:
        buf = torch.zeros(target_shape, device=x.device, dtype=x.dtype)
        buf[:, attn_metadata.non_pad_index] = x[:, attn_metadata.tile_partition_indices]
        return buf

    # Reuse the per-step buffer stashed on metadata (lazily allocated
    # on the first VSA layer's call within a denoising step).  Pad
    # positions are zero from the initial torch.zeros and never
    # written to.  Scoping to metadata makes reuse safe across
    # concurrent requests and keeps the "pad positions are zero"
    # invariant trivially true: ``non_pad_index`` is fixed within
    # a single metadata instance.
    buf = attn_metadata.tile_buf
    if (buf is None or buf.shape != target_shape or buf.dtype != x.dtype or buf.device != x.device):
        buf = torch.zeros(target_shape, device=x.device, dtype=x.dtype)
        attn_metadata.tile_buf = buf

    buf[:, attn_metadata.non_pad_index] = x[:, attn_metadata.tile_partition_indices]
    return buf

Functions:

fastvideo.attention.backends.video_sparse_attn.construct_variable_block_sizes cached 

construct_variable_block_sizes(dit_seq_shape: tuple[int, int, int], num_tiles: tuple[int, int, int], device: device) -> LongTensor

Compute the number of valid (non‑padded) tokens inside every (ts_t × ts_h × ts_w) tile after padding ‑‑ flattened in the order (t‑tile, h‑tile, w‑tile) that rearrange uses.

Returns

torch.LongTensor # shape: [∏ full_window_size]

Source code in fastvideo/attention/backends/video_sparse_attn.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
@functools.lru_cache(maxsize=10)
def construct_variable_block_sizes(
    dit_seq_shape: tuple[int, int, int],
    num_tiles: tuple[int, int, int],
    device: torch.device,
) -> torch.LongTensor:
    """
    Compute the number of valid (non‑padded) tokens inside every
    (ts_t × ts_h × ts_w) tile after padding ‑‑ flattened in the order
    (t‑tile, h‑tile, w‑tile) that `rearrange` uses.

    Returns
    -------
    torch.LongTensor  # shape: [∏ full_window_size]
    """
    # unpack
    t, h, w = dit_seq_shape
    ts_t, ts_h, ts_w = VSA_TILE_SIZE
    n_t, n_h, n_w = num_tiles

    def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.LongTensor:
        """Vector with the size of each tile along one dimension."""
        sizes = torch.full((n_tiles, ), tile, dtype=torch.int, device=device)
        # size of last (possibly partial) tile
        remainder = dim_len - (n_tiles - 1) * tile
        sizes[-1] = remainder if remainder > 0 else tile
        return sizes

    t_sizes = _sizes(t, ts_t, n_t)  # [n_t]
    h_sizes = _sizes(h, ts_h, n_h)  # [n_h]
    w_sizes = _sizes(w, ts_w, n_w)  # [n_w]

    # broadcast‑multiply to get voxels per tile, then flatten
    block_sizes = (
        t_sizes[:, None, None]  # [n_t, 1,   1]
        * h_sizes[None, :, None]  # [1,   n_h, 1]
        * w_sizes[None, None, :]  # [1,   1,   n_w]
    ).reshape(-1)  # [n_t * n_h * n_w]

    return block_sizes