72 lines
2.2 KiB
Markdown
72 lines
2.2 KiB
Markdown
|
# FlashAttention-4
|
||
|
|
|
||
|
|
The hot runtime exposes the official FlashAttention-4 CuTeDSL Blackwell kernel
|
||
|
|
as the strict `flash4` attention backend. It does not silently fall back to
|
||
|
|
SageAttention or SDPA.
|
||
|
|
|
||
|
|
## Versions
|
||
|
|
|
||
|
|
- `flash-attn-4==4.0.0b27`
|
||
|
|
- `nvidia-cutlass-dsl==4.6.2` with the CUDA 13 extra
|
||
|
|
- `quack-kernels==0.6.4`
|
||
|
|
- Upstream: <https://github.com/Dao-AILab/flash-attention>
|
||
|
|
- License: BSD-3-Clause
|
||
|
|
|
||
|
|
The package is a beta release. The Spark image pins these versions rather than
|
||
|
|
tracking the latest prerelease.
|
||
|
|
|
||
|
|
## API
|
||
|
|
|
||
|
|
Select Flash4 per request:
|
||
|
|
|
||
|
|
```json
|
||
|
|
{
|
||
|
|
"prompt": "A simple cinematic scene.",
|
||
|
|
"output": "/output/flash4-example.mp4",
|
||
|
|
"width": 960,
|
||
|
|
"height": 544,
|
||
|
|
"frames": 124,
|
||
|
|
"steps": 12,
|
||
|
|
"attention": "flash4"
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
H3's main attention blocks call Flash4 in native
|
||
|
|
`[batch, sequence, heads, head_dim]` layout. The adapter requires CUDA, matching
|
||
|
|
FP16/BF16 Q/K/V tensors, matching shapes, and H3's head dimension of 128. Any
|
||
|
|
violation raises an error instead of changing kernels.
|
||
|
|
|
||
|
|
## GB10 Validation
|
||
|
|
|
||
|
|
Flash4 dispatches its SM120 forward implementation on the GB10, which PyTorch
|
||
|
|
reports as compute capability 12.1.
|
||
|
|
|
||
|
|
| Validation | Result |
|
||
|
|
| --- | --- |
|
||
|
|
| BF16 `[1,257,8,128]` vs SDPA | max abs `0.001953125`, mean abs `5.04e-05` |
|
||
|
|
| BF16 `[1,20480,56,128]` vs SDPA | max abs `0.000244141`, mean abs `1.96e-08` |
|
||
|
|
| Tiny joint AV H3 smoke | passed, no fallback |
|
||
|
|
| 28 runtime contracts | passed |
|
||
|
|
|
||
|
|
Matched base 12-step talking benchmark at `960x544x124`, seed `440410`:
|
||
|
|
|
||
|
|
| Backend | Sampling | Relative to Flash4 |
|
||
|
|
| --- | ---: | ---: |
|
||
|
|
| Sage2 | `114.67s` | `-19.4%` |
|
||
|
|
| Flash4 | `142.23s` | baseline |
|
||
|
|
| PyTorch SDPA | `144.40s` | `+1.5%` |
|
||
|
|
|
||
|
|
Flash4 is therefore close to SDPA and materially slower than Sage2 for this H3
|
||
|
|
workload on GB10. Its value is as another exact-style attention and audio-quality
|
||
|
|
comparison, not as the current speed default.
|
||
|
|
|
||
|
|
The matched Flash4 proof contains H.264 video and stereo 32 kHz AAC audio:
|
||
|
|
|
||
|
|
```text
|
||
|
|
/home/daniel/StoryStudioAssets/H3-output/h3-native-upscale-t2v/talking-woman-base12-flash4-960x544.mp4
|
||
|
|
```
|
||
|
|
|
||
|
|
The current CuTeDSL package emits a one-time `AuxData` JIT argument warning.
|
||
|
|
Compilation and inference still complete successfully, but the warning should
|
||
|
|
be rechecked when upgrading FlashAttention-4 or CUTLASS DSL.
|