cortexelus commited on
Commit
f3263ec
·
verified ·
1 Parent(s): fb17708

Rebuild medium DiT fp16-mixed engines (sm_90 + new sm_120): attention now fuses, 4.3x faster

Browse files

Rebuilds the medium DiT's FP16-mixed TRT engines for **sm_90** and **sm_120** from the corrected ONNX merged in #3, and adds sm_120 as a published architecture.

Companion to #3 (graph fix) and to the code PR on `Stability-AI/stable-audio-3` (producer pass, runtime default, docs).

## What changed

| file | arch | size | fused MHA |
|---|---|---|---|
| `tensorRT/sm_90/sa3-m/dit_fp16mixed.trt` | H200 (sm_90) | 2.92 GB | 96 |
| `tensorRT/sm_120/sa3-m/dit_fp16mixed.trt` | Blackwell (sm_120) | 2.92 GB | 96 |

Both built with the public consumer recipe — `build_from_onnx.py sa3-m`, i.e. parse + `STRONGLY_TYPED`, no builder flags, profile `L ∈ [1, 4096]` opt 1292, TF32 cleared. No pinning, no plugins, nothing a consumer cannot reproduce. TRT 10.15.1.29.

sm_90 replaces the published engine in place (same path, same filename). sm_120 is new — Blackwell workstation users previously had to compile locally.

## Why

The published engine kept the entire O(L²) attention core in FP32 and therefore could not use TRT's FMHA fuser: **0 fused MHA nodes**. Bounding the RoPE island (#3) lets QK^T → Softmax → P·V run FP16, so the fuser fires on all 96 attentions.

## Measured

Per DiT forward, CUDA-event median of 7 after 3 warmup, TF32 off:

| L | sm_90 published | sm_90 this PR | sm_120 published | sm_120 this PR |
|---|---|---|---|---|
| 256 | 11.1 ms | **7.4 ms** | 23.4 ms | **17.6 ms** |
| 1292 | 27.9 ms | **12.2 ms** | 99.2 ms | **44.4 ms** |
| 4096 | 180.9 ms | **41.4 ms** (4.4×) | 714.0 ms | **166.1 ms** (4.3×) |

Accuracy vs the same-arch `dit_fp32.trt` reference (seed 6000, 8-step pingpong, TF32 off) — free-run latent std ratio and teacher-forced velocity cosine:

| L | std ratio | teacher-forced cos |
|---|---|---|
| 256 | 1.00× | **1.0000** |
| 1292 | 1.00× | **1.0000** |
| 4092 | 1.00× | **1.0000** (min 0.9999) |

380 s render through SAME-L: 0.001% samples clipped, crest 6.51, against the FP32 engine's 0.000% / 6.25.

So it is both **~4.3× faster and marginally more accurate** than the engine it replaces (the old one scored 0.9998) — closer to eager semantics, since PyTorch also runs the attention core at the autocast dtype inside fused SDPA.

## Compatibility

- Filename, precision name (`fp16mixed`), profile, and I/O contract are all unchanged. `sa3_trt.py` picks these up with no config change.
- **Output changes at a fixed seed.** This is a bug fix, not a re-tune: renders differ from the previous engine. Deterministic and reproducible going forward, but not bit-identical to the old engine.
- Engines are TRT-version and arch specific: these require TRT 10.15.x on sm_90 / sm_120 respectively. Other architectures compile from the corrected ONNX with `build_from_onnx.py sa3-m`.
- Only the medium DiT is affected. `sa3-sm-music` / `sa3-sm-sfx` use standard (non-differential) attention and already fused 40/40 MHA in their published engines.

## Not in this PR

- **`sa3-m/dit_bf16.trt` is unchanged and still drifts at long sequence.** It evaluates RoPE's rotation angle in bf16, where the angle reaches ~4155 rad at L=4092 and bf16's spacing is 32 rad (> 2π) — position information for the fast-rotating dimensions is destroyed, the latent inflates ~2.5× over 8 steps, and the decoder clips 2–3% of samples on a 6-min render. It is clean at short lengths. The companion code PR documents this and moves the medium default to `fp16mixed`; a corrected bf16 engine is in progress separately.
- FAD/CLAP perceptual validation at n=128 is still running; the numbers above are latent-space and amplitude metrics plus spot renders. **Do not merge on quality grounds until that lands** — I will post it here.

tensorRT/sm_120/sa3-m/dit_fp16mixed.trt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ff6fdfa5f02fdfd8dbe6f5b1c2940afb44b508ae4b773b89d8bc5339e38dbed2
3
+ size 2917847172
tensorRT/sm_90/sa3-m/dit_fp16mixed.trt CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:399f3fa18e21f86528322a4543fed17999f6bd95589886f2d4f8f3e2c77fc425
3
- size 2914585244
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d924c1fb8e0b5c5e19c82da96609eed2c999818af029d606f2f9c644a12da0da
3
+ size 2921827284