Metadata-Version: 2.4
Name: h3-sparse-attn
Version: 0.1.0
Summary: Block-sparse attention for the MiniMax H3 video DiT
Author-email: Vimanyu Taneja <taneja.vimanyu@gmail.com>
License-Expression: MIT
Project-URL: Homepage, https://github.com/Occipital-Labs/h3-sparse-attn
Project-URL: Repository, https://github.com/Occipital-Labs/h3-sparse-attn
Project-URL: Issues, https://github.com/Occipital-Labs/h3-sparse-attn/issues
Keywords: attention,sparse-attention,cuda,triton,diffusion,video-generation,minimax-h3
Classifier: Development Status :: 4 - Beta
Classifier: Environment :: GPU :: NVIDIA CUDA
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.6
Requires-Dist: triton>=3.2
Requires-Dist: diffusers>=0.40
Requires-Dist: ninja
Provides-Extra: dev
Requires-Dist: pytest>=8; extra == "dev"
Requires-Dist: ruff>=0.6; extra == "dev"
Dynamic: license-file

# h3-sparse-attn

Block-sparse self-attention for the MiniMax H3 video DiT. Each query block attends to its top-k key blocks of 128
tokens; the skipped blocks are folded back in as one summary row each. Training-free, so the released H3 checkpoint
runs unchanged.

Requires sm_120 (RTX 5090, RTX PRO 6000 Blackwell), bf16, head_dim 128, forward only. Anything else falls back to
the attention H3 already ships.

This is an approximation: same quality, but not the same sample as full attention at a fixed seed. `top_k` sets how
far it drifts.

## Install

```
pip install h3-sparse-attn
```

## Use

```python
import h3_sparse_attn

h3_sparse_attn.patch(model, top_k=64)
```

`model` is a `MiniMaxH3Transformer3DModel` or a pipeline holding one. `top_k` counts 128-token key blocks.

## Measured

RTX PRO 6000 Blackwell Server Edition, bf16, 56 heads, head_dim 128, `top_k=64`. Sequence lengths are what H3 packs
at 24 fps. "Step" is one denoising step of the 50-layer DiT.

| Video | Sequence | Attention | Speedup | Step | Speedup |
| --- | --- | --- | --- | --- | --- |
| 768p 5s | 37,740 | 112.4 → 30.7 ms | 3.7x | 10.67 → 6.47 s | 1.6x |
| 768p 10s | 73,416 | 426.6 → 61.8 ms | 6.9x | 31.26 → 12.70 s | 2.5x |
| 768p 15s | 103,996 | 862.7 → 91.0 ms | 9.5x | 57.12 → 18.11 s | 3.2x |
| 2K 5s | 133,644 | 1435.6 → 121.8 ms | 11.8x | 89.22 → 23.48 s | 3.8x |

## License

MIT
