vsa_kernel/__init__.py

"""Vendored FastVideo block-sparse attention Triton kernels.

Copied verbatim from `fastvideo-kernel` (https://github.com/hao-ai-lab/FastVideo,
`fastvideo-kernel/python/fastvideo_kernel/triton_kernels/`), Apache License 2.0.

Only the two pure-Triton modules are vendored: the package's default route is a
CUDA extension that has to be compiled per architecture (and whose fastest entry,
`block_sparse_attn_sm100a`, is GB200-only), while FastVideo's own
`--vsa-kernel triton` route -- these two files -- runs anywhere Triton does and
computes exactly the same masked attention.
"""

from .block_sparse_attn_triton import triton_block_sparse_attn_forward
from .index import map_to_index

__all__ = ["triton_block_sparse_attn_forward", "map_to_index"]