stabilityai

/

stable-audio-3-optimized

like 26

Follow

Stability AI 38.1k

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

TRT engines are portable across devices of the same compute capability, so the sm_120 engine runs on any sm_120 part (RTX 5090, RTX PRO 6000 Blackwell, RTX PRO 4500/4000, …). The latency figures below are specific to the card each engine was built and measured on — see the note under the timings.

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, batch 1, dedicated idle GPU.

H200 (sm_90 · 132 SM · 141 GB HBM3e · 700 W · TRT 10.15.1.29):

RTX PRO 4500 Blackwell Server Edition (sm_120 · 82 SM · 32 GB GDDR7 · 165 W · TRT 10.15.1.29):

Do not read the sm_120 absolute numbers as representative of sm_120 generally. The RTX PRO 4500 is a mid-range 165 W workstation part; other sm_120 cards (RTX 5090, RTX PRO 6000 Blackwell) carry roughly twice the SM count and considerably more memory bandwidth, so they will be substantially faster in absolute terms. The speedup ratio is the portable number here — both engines were measured on the same card, and the ratio held within 0.1× across two very different architectures, which is what you would expect from a change that is about kernel fusion rather than silicon.

For the same reason, the H200 and RTX PRO 4500 columns are not an architecture comparison: they are different product tiers (700 W datacenter vs 165 W workstation), and the ~4× absolute gap between them is silicon, not tuning.

Accuracy of each shipped engine against the dit_fp32.trt reference built and run on that same card (seed 6000, 8-step pingpong, TF32 off) — free-run latent std ratio and teacher-forced velocity cosine, mean over the 8 steps with the worst step in brackets:

Both are effectively fp32-exact. The H200 engine is a hair further from its reference at the 6-minute anchor (0.9997 vs 1.0000) — each card's TRT picks its own kernels and accumulation order, so per-arch references are the right comparison and small per-arch differences are expected.

380 s render through SAME-L (RTX PRO 4500, SAME-L windowed decoder, seed 6000): 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 compute-capability specific: these require TRT 10.15.x on sm_90 / sm_120 respectively. They are portable to any other device of the same compute capability (the sm_120 engine runs on RTX 5090 / RTX PRO 6000 Blackwell / etc.), but not across capabilities. Other architectures — including datacenter Blackwell (sm_100) — 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 — verified against the published ONNX: their QK^T operands are already FP16, where medium's were FP32 (which is why medium's pre-softmax Cast(FP32) was a no-op and its whole attention core sat unfused).
• ⚠ One forward-looking caveat for the small DiTs: although their attention already fuses, their Softmax is still inside an FP32 island, so the new bound_attention_core() pass matches and converts 40/40 on them too. Re-running the producer for sa3-sm-music / sa3-sm-sfx would therefore change their graphs (softmax FP32 → FP16). That variant is neither shipped nor validated — the engines in this PR are medium-only, and the published small-DiT ONNX/engines are untouched. Anyone regenerating the small DiTs should validate before shipping.

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 has not been run yet. Everything above is latent-space fidelity (teacher-forced velocity cosine, latent std ratio vs the per-card fp32 reference), amplitude metrics on full-length renders, and spot listening across 3 prompts x 3 lengths. That is a narrower bar than the dit_bf16 release cleared (FAD/floor + CLAP over the prompt bank at n=128). Do not merge on quality grounds until it lands — I will post it here.
• Pipeline-mode coverage is likewise partial: text-to-audio and varlen are exercised, but CFG, audio-to-audio, and inpaint have not been re-tested against these engines.

Validation complete — both blockers from "Not in this PR" are now closed

Full reports: FADCLAP_VALIDATION.md and MODE_COVERAGE.md.

1. FAD + CLAP, n=128 — clears the bar

FAD against each card's own dit_fp32.trt, expressed as a ratio to the fp32 re-seed floor (FAD between two fp32 runs at different seeds — anything below the floor is perceptually meaningless):

The dit_bf16 release cleared this same bar at 0.59× floor; this beats it by 5–74×. In 200 paired bootstrap resamples the new engine is inside the floor 100% of the time at the 2-minute anchor. CLAP shows no significant degradation at any arch or length (paired Δ vs fp32 = +0.00000 ± 0.00049 on H200 at 2 min; the only significant shifts are higher CLAP). Clipping and crest match fp32 (0.0180% vs 0.0182%, crest 6.95 vs 6.95).

Prompts are the 128 held-out entries (indices 512–639) of the same prompt bank the bf16 release was measured on, split by track id.

One honest qualification. "New is at-or-better than old" does not hold as a point estimate in every cell. The old engine keeps the whole O(L²) attention core in FP32, so it hugs the fp32 reference by construction; the new one runs that core in FP16. At 6.3 min on H200 the new engine is further from fp32 than the old by a margin that is statistically real (paired bootstrap P(new better) = 0.025) and perceptually irrelevant — +0.0115 FAD = 2.8% of the noise floor — mirroring the teacher-forced cosines already in this PR (0.9997 new vs 0.9998 old @L4092). At 6.3 min on sm_120 the ordering reverses (P = 0.84). So the defensible claim is "perceptually equivalent to fp32 and to the engine it replaces, at 4.3× the speed" — not "better than the old engine".

For scale: the two fp32 engines built for the two different cards are 100× further apart in FAD (0.4010) than the new engine is from its own card's fp32 (0.0039).

2. Pipeline modes — drop-in on both architectures

314 CLI invocations through the shipped sa3_trt.py: text-to-audio, CFG / negative-prompt / APG, audio-to-audio, inpaint, varlen across the whole advertised profile (L = 1 … 4096), both decoders, all three precisions, and both small DiTs. All run and produce finite, non-silent, correct-length audio with sane amplitude statistics; cross-arch stats agree to ~0.1 dB. The I/O contract is identical — same 6 inputs, names, dtypes, output, and single optimization profile.

Nothing in the mode matrix is broken by this change.

⚠ But the matrix surfaced five pre-existing runtime defects — none caused by this PR

Every one reproduces identically on the old published engine, and every one lives in the Python call helpers or CLI wrapper that this PR does not touch. Three are silent wrong-output bugs. These should gate the next code release, not this engine PR, but they need to be on the record:

Two hit the default configuration: P2/P3 mean --dit medium with its own default decoder (same-l) can return wrong audio with exit code 0, which is why the correctness matrix above was routed through --decoder same-s. P1 means every CFG measurement ever taken through this runtime was measuring the unconditional branch.

"CFG passes" above therefore means the dual-pass executes cleanly on both passes at every length with sane audio out — which is what the engine swap could plausibly have broken, and did not. It does not mean guidance is applied.

Also worth flagging: a measurement-infrastructure bug

The sad310 venv's numpy 1.23.5 (openblas64) computes threaded gemm incorrectly above ~64³ and returns non-orthogonal eigh eigenvectors, which was producing negative Fréchet distances (self-check FAD(X,X) = −17.57). VGGish's log-mel front-end is a numpy matmul, so embeddings were contaminated too. The numbers above were computed after pinning OPENBLAS_NUM_THREADS=1 with a hard self-test, moving all covariance/eigen work to numpy 2.2.6, and validating against closed-form values plus an independent torch-float64 implementation (agreement 0.00e+00 across all 16 cells). Historical FAD numbers computed in that venv are suspect — including, potentially, the 0.59×-floor bf16 reference quoted above.

Not covered

Human listening at n=128 (spot listening only, across 3 prompts × 3 lengths); CLAP text truncation to 77 tokens (identical across sets, so comparisons hold).

Validated: FAD 0.008x/0.113x of the fp32 re-seed floor at n=128, all pipeline modes drop-in, ear-checked. sm_120 completed to full 15-file parity with sm_90 so install.sh cannot 404 on it.