[diffusion] feat: support spargeattention (#37437)

This commit is contained in:
Mick
2026-09-02 09:42:33 +08:00
committed by GitHub
parent dc276264cb
commit dde0ecdb90
8 changed files with 463 additions and 1 deletions
@@ -64,6 +64,11 @@ For SGLang-native pipelines, the CLI accepts the lowercase names of `AttentionBa
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)", whiteSpace: "nowrap"}}>`SAGE_ATTN_3`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Requires SageAttention3 installed per upstream instructions.</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`sparge_attn`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)", whiteSpace: "nowrap"}}>`SPARGE_ATTN`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Training-free sparse SageAttention2. CUDA SM80/86/87/89/90, FP16/BF16, head dim 64/128, and square self-attention with sequence length at least 128. Other attention shapes use dense SDPA. Install <code>pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation</code>; configure retained block ratio with <code>topk</code>.</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`sol_attn`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)", whiteSpace: "nowrap"}}>`SOL_ATTN`</td>
@@ -213,6 +218,16 @@ Some backends require additional configuration. You can pass these parameters vi
</tbody>
</table>
**SpargeAttention (`sparge_attn`)**
| Parameter | Type | Description | Default |
|---|---|---|---|
| `topk` | `float` | Fraction of predicted attention blocks retained per head. Higher values preserve more attention work and generally improve quality. Must be in `(0, 1]`. | `0.5` |
SpargeAttention is approximate even when `topk=1`: the recommended upstream
kernel quantizes attention through SageAttention2. Validate output quality and
end-to-end latency on the target model and resolution before deployment.
**V-MoBA (`vmoba_attn`)**
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
@@ -494,6 +509,16 @@ Some backends require additional configuration. You can pass these parameters vi
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>CUDA-only (optional dependency).</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`sparge_attn`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Yes</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>No</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>No</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>No</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>CUDA SM80/86/87/89/90 only. Requires <code>SpargeAttn</code>; head dim 64/128 and square self-attention with sequence length at least 128.</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`sol_attn`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Yes</td>
@@ -678,6 +703,30 @@ dense switching; under ring parallelism the target must be ring-capable. Note
`sage_attn` / `sage_attn_3` are lossy (quantized attention) — validate quality
on your workload.
### Using SpargeAttention
Install the optional CUDA extension, then select the backend explicitly:
```bash
pip install ninja
pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation
sglang generate \
--model-path <MODEL_PATH_OR_ID> \
--prompt "..." \
--attention-backend sparge_attn \
--attention-backend-config topk=0.5
```
`sparge_attn` applies to compatible self-attention layers. Cross-attention uses
the existing dense fallback; short self-attention and asymmetric Q/KV also use
dense SDPA. LTX-2, LTX-2.3, and LTX-2.5 have compatible video (128) and audio
(64) head dimensions, so their sufficiently long, unmasked self-attention uses
SpargeAttention while prompt and audio-video cross-attention remain dense.
Ulysses sequence parallelism is supported, but ring attention is not because the
upstream kernel does not expose the softmax LSE needed for ring merging. K/V-gather
SP therefore follows the normal sparse-backend rule and uses Ulysses instead.
### Sage then Sol hybrid
`sol_attn` keeps the first `dense_steps` steps dense. Set