sm_120 sa3-m: add the fp8 medium-DiT tier (176 fp8 GEMM + 96 bf16 fused MHA, 1.43-1.95x vs fp16mixed)
Adds the fp8 medium-DiT tier for sm_120 (Blackwell). Built from the calibratedonnx/sa3-m/dit_fp8.onnx (commit 8383b7d0) with the same recipe as sm_90 — no code
changes needed, the sa3-m-fp8 build target is arch-agnostic.
Selectable with --dit medium --precision fp8; the default stays fp16mixed.
Identity matches sm_90 exactly
Verified with scripts/verify_fp8_rope.py on DETAILED-verbosity twins of both engines:
| sm_120 | sm_90 (published) | |
|---|---|---|
| fp8 E4M3 GEMMs | 176 / 181 | 176 / 181 |
fused MHA (_gemm_mha_v2) |
96, all BF16 | 96, all BF16 |
| runtime cos/sin layers | 2 | 2 |
| engine size | 1421.19 MB | 1423.75 MB |
Speed — the gain is larger on sm_120 than on sm_90
ms per forward, vs the published fp16mixed on each arch, median of 9, CUDA events on the
work's own stream:
| L | sm_90 fp8 | vs fp16mixed | sm_120 fp8 | vs fp16mixed |
|---|---|---|---|---|
| 256 | 4.99 | 1.48x | 8.32 | 1.95x |
| 1292 | 9.34 | 1.33x | 28.42 | 1.58x |
| 4096 | 31.31 | 1.34x | 118.08 | 1.43x |
sm_90 reproduces its documented 1.40/1.32/1.34. sm_120 does better, which is the expected
direction: it is the bandwidth-starved part (734 vs 4084 GB/s measured), and fp8 halves
weight traffic.
End-to-end 120 s render on an RTX PRO 4500: 493 ms vs 620 ms for fp16mixed (243x
realtime), no tail-energy collapse, 0.05% clipped samples.
Fidelity
Velocity cosine vs the FP32 engine at L=1292, swept across the timestep, because
calibration effects concentrate near t=1.0 and are invisible mid-schedule:
t=1.0 t=0.98 t=0.95 t=0.9 t=0.7 t=0.5 t=0.2
sm_90 fp8 0.9744 0.9714 0.9580 0.9522 0.9733 0.9855 0.9911
sm_120 fp8 0.9845 0.9778 0.9767 0.9661 0.9766 0.9876 0.9917
sm_120 is at or above sm_90 at every point. The t=1.0 reading also confirms this engine
came from the calibrated ONNX rather than the earlier uncalibrated one, whose worst step
sits far lower.
As on sm_90, this is a SPEED tier, not a fidelity upgrade — fp16mixed remains ~1.0 against
fp32 and stays the default.
Tooling note: verify_fp8_rope.py's CHECK 3 previously reported False for this engine on
both architectures; fixed in https://github.com/Stability-AI/stable-audio-3/pull/88.