Metadata-Version: 2.4
Name: fastbsa
Version: 0.1.0
Summary: Block-sparse attention for consumer Blackwell, drop-in for LongCat-Video
Author-email: Vimanyu Taneja <taneja.vimanyu@gmail.com>
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/Occipital-Labs/fastbsa
Project-URL: Source, https://github.com/Occipital-Labs/fastbsa
Project-URL: Issues, https://github.com/Occipital-Labs/fastbsa/issues
Keywords: cuda,attention,block-sparse,blackwell,sm_120,video-generation
Classifier: Development Status :: 3 - Alpha
Classifier: Environment :: GPU :: NVIDIA CUDA :: 12
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: C++
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: torch>=2.7
Provides-Extra: test
Requires-Dist: pytest; extra == "test"
Requires-Dist: triton>=3.3; extra == "test"
Dynamic: license-file

# 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.

| phase | LongCat Triton | fastbsa | |
|---|---|---|---|
| layout | 8.20 ms | 0 ms | fused into the kernel |
| select | 5.48 ms | 5.55 ms | |
| attend | 132.85 ms | 108.12 ms | 1.22x |
| **total** | **146.38 ms** | **113.98 ms** | **1.28x** |

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](https://github.com/Occipital-Labs/fastbsa/releases); it is valid only
for the torch version it was built against.

## Use

```python
import fastbsa
fastbsa.patch_longcat()
```

Or call the kernel directly:

```python
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.
