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 
240
241
242
243
244
245
246
247
248
249
250
251
252
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 
290
291
292
293
294
295
296
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 
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
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:
        return scatter_into_tile_buf(x, target_shape, attn_metadata.non_pad_index, None,
                                     attn_metadata.tile_partition_indices)

    # Buffer scoped to the per-step metadata (lazily allocated on the
    # first VSA layer's call within a denoising step), which keeps reuse
    # safe across concurrent requests.
    buf = scatter_into_tile_buf(x, target_shape, attn_metadata.non_pad_index, attn_metadata.tile_buf,
                                attn_metadata.tile_partition_indices)
    attn_metadata.tile_buf = buf
    return buf

Functions:

fastvideo.attention.backends.video_sparse_attn.compute_topk 

compute_topk(sparsity: float, num_blocks: int) -> int

Blocks to keep for a sparsity level, clamped to [1, num_blocks].

Source code in fastvideo/attention/backends/video_sparse_attn.py 
161
162
163
def compute_topk(sparsity: float, num_blocks: int) -> int:
    """Blocks to keep for a sparsity level, clamped to [1, num_blocks]."""
    return max(1, min(math.ceil((1 - sparsity) * num_blocks), num_blocks))

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, tile_size: tuple[int, int, int] = VSA_TILE_SIZE) -> 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
99
@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,
    tile_size: tuple[int, int, int] = VSA_TILE_SIZE,
) -> 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 = 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

fastvideo.attention.backends.video_sparse_attn.scatter_into_tile_buf 

scatter_into_tile_buf(x: Tensor, target_shape: tuple[int, ...], dst_index: Tensor, buf: Tensor | None, src_index: Tensor | None = None) -> Tensor

Zero-padded tile scatter shared by the VSA backends.

Allocates (zeros) when buf is missing or mismatched; otherwise reuses it — pad slots are never written and every non-pad slot is fully overwritten per call, so a reused buffer stays valid. Callers own the buffer's lifetime (per-metadata for Wan VSA, per-builder for VSA-H3) and its aliasing contract: the result is only valid until the next call with the same buffer.

Source code in fastvideo/attention/backends/video_sparse_attn.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def scatter_into_tile_buf(
    x: torch.Tensor,
    target_shape: tuple[int, ...],
    dst_index: torch.Tensor,
    buf: torch.Tensor | None,
    src_index: torch.Tensor | None = None,
) -> torch.Tensor:
    """Zero-padded tile scatter shared by the VSA backends.

    Allocates (zeros) when ``buf`` is missing or mismatched; otherwise reuses
    it — pad slots are never written and every non-pad slot is fully
    overwritten per call, so a reused buffer stays valid. Callers own the
    buffer's lifetime (per-metadata for Wan VSA, per-builder for VSA-H3) and
    its aliasing contract: the result is only valid until the next call with
    the same buffer.
    """
    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)
    buf[:, dst_index] = x if src_index is None else x[:, src_index]
    return buf