cortexelus commited on
Commit
af5680d
·
verified ·
1 Parent(s): 8383b7d

sm_120 sa3-m: add the fp8 medium-DiT tier (176 fp8 GEMM + 96 bf16 fused MHA, 1.43-1.95x vs fp16mixed)

Browse files

Adds the fp8 medium-DiT tier for sm_120 (Blackwell). Built from the calibrated
`onnx/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.

Files changed (1) hide show
  1. tensorRT/sm_120/sa3-m/dit_fp8.trt +3 -0
tensorRT/sm_120/sa3-m/dit_fp8.trt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c2613da5f8b5434b9ca33f543668ca260c5bf1252ba74a1d1633e936e8fb40d1
3
+ size 1490230196