fastbsa

A CUDA block-sparse attention kernel for consumer Blackwell (sm_120), a drop-in replacement for the Triton BSA that ships with LongCat-Video.

The number

LongCat-Video 720p refinement at its shipped BSA config, on one RTX PRO 6000 Blackwell Server. Per DiT layer per denoising step, median of 50, against an autotuned Triton baseline.

End to end through the real model, the same 96 BSA calls cost 17.10 s on Triton and 11.43 s on fastbsa, a 1.50x speedup.

Install

pip install fastbsa

This compiles the extension against your own torch, so it needs torch >= 2.7, the CUDA 12.8 toolkit and an sm_120 card. A prebuilt sm_120 wheel is attached to each release; it is valid only for the torch version it was built against.

Use

import fastbsa
fastbsa.patch_longcat()

Or call the kernel directly:

from fastbsa import flash_attn_bsa_3d
out = flash_attn_bsa_3d(q, k, v, latent_shape_q, latent_shape_k, sparsity=0.9375,
                        chunk_3d_shape_q=(4, 4, 4), chunk_3d_shape_k=(4, 4, 4))

How it works

LongCat permutes Q, K and V into block-contiguous order before every call; fastbsa evaluates that index map inside the kernel's tile loop, so nothing moves through HBM. Block selection stays in shared memory instead of round-tripping a 446 MB score matrix through torch.topk. The attention kernel is hand-written for sm_120 with mma.sync and cp.async.

Selection reproduces LongCat's exactly, so output is validated against fp32 dense-masked attention. Tests run with pytest tests/.

License

Apache-2.0, see LICENSE. The test suite compares against LongCat-Video's own Triton kernel, vendored under tests/reference/ at the revision in tests/reference/UPSTREAM_REV and licensed under MIT by Meituan. It is test-only and is not part of the installed package.