File size: 1,509 Bytes
8e9f35a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
---
library_name: kernels
{% if license %}license: {{ license }}
{% endif %}tags:
- kernels
- cuda
- attention
- triton
- cute-dsl
---

# 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

```python
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](https://arxiv.org/abs/2607.24027)

## Source

The implementation is maintained in
[`NVlabs/Sana`](https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn).
This release is pinned to commit
[`8a26fb0`](https://github.com/NVlabs/Sana/commit/8a26fb0ec9e353125ead798cb2e312d5ce48cded).