h3-blackwell-runtime/FLASH4.md

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:

{
  "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:

/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.