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

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.

License

MIT