Hugging Face's logo

Efficient-Large-Model
/
Sol-Attn-Kernel-Source

Kernels
cuda
attention
triton
cute-dsl
Model card Files
xet
 Community
Sol-Attn-Kernel-Source / CARD.md
hp-l33's picture
hp-l33
Add Sol-Attn Kernel Builder source
8e9f35a verified
preview code
|
Raw
History Blame Contribute Delete
1.51 kB

Sol-Attn

Sol-Attn accelerates image and video generation with on-the-fly attention sparsification. The public API dispatches to CuTe DSL kernels on SM90, SM100, and SM120, and to Triton on SM80 and SM89 or when CuTe DSL is unavailable.

Usage 

from kernels import get_kernel

kernel = get_kernel("{{ repo_id }}", version={{ version }})

out = kernel.sol_attn(
    q,  # Contiguous BF16 CUDA tensor [batch, tokens, heads, 128].
    k,  # Same shape, dtype, layout, and device as q.
    v,  # Same shape, dtype, layout, and device as q.
    tau=1.0,
    thresh_type="exact",
)

The released implementation is noncausal and forward-only. Q/K/V must have the same BTHD shape. An optional exact KV sink is available through sink_start and sink_tokens.

Backends

Architecture Example GPU Backend
SM90 H100 CuTe DSL
SM100 GB200 CuTe DSL
SM120 RTX 5090 CuTe DSL
SM80 / SM89 A100 / RTX 4090 Triton

Paper

Accelerating Video Generation Inference via On-the-Fly Attention Sparsification

Source

The implementation is maintained in NVlabs/Sana. This release is pinned to commit 8a26fb0.