stabilityai

/

stable-audio-3-optimized

like 27

Follow

Stability AI 38.2k

Replaces the defective sa3-m/dit_bf16.trt on both architectures, and adds the ONNX it is built from. Companion to code PR Stability-AI/stable-audio-3#81.

What changed

The defect being fixed

RoPE's rotation angle is t · inv_freq with inv_freq[0] == 1.0 exactly, so it reaches ~4155 rad at L=4092. bf16 has 8 significand bits, so its spacing up there is 32 rad — more than a full 2π rotation. The published engine was built from the raw dit.onnx with BuilderFlag.BF16, which leaves TRT free to evaluate that angle in bf16, destroying 9/16 frequency pairs. Over 8 sampling steps the latent inflates ~2.5× and the decoder renders it that much too hot:

It is clean at short lengths, which is why the original validation passed: that was run at short and standard lengths, where the defect does not bite.

The fix

The cos/sin tables depend on nothing but position — not activations, not the timestep, not the prompt — so they are compile-time constants. The new producer precomputes them in fp64 on the host, embeds fp32 initializers, rewires all 96 trig sites to two shared slices, and lets dead-code elimination delete the angle chain entirely.

No fp32 island is then needed in the RoPE region at all, because there is no in-graph trig left to protect. And the baked table values are bf16-safe — they live in [−1,1], where rounding them measures latent std 0.9161 against fp32's 0.9162.

The seconds_total (62,000 rad) and timestep (55,000 rad) Fourier chains are functions of runtime inputs, so they cannot be baked and still get a small fp32 island — 25 layers on scalar and [1,128] tensors, ~0.1%. The consumer builder selects them structurally (after baking, every remaining Cos/Sin in the graph is one of those two chains) rather than by node name, so the recipe survives a re-export.

Measured

Both engines built through the public consumer path — build_from_onnx.py sa3-m-bf16, nothing bespoke.

Timings are per DiT forward, CUDA-event median of 7 after 3 warmup, TF32 off, batch 1, idle GPU. The RTX PRO 4500 is a mid-range 165 W part — other sm_120 cards (RTX 5090, RTX PRO 6000 Blackwell) have roughly twice its SM count and will be faster in absolute terms; the ratio is the portable number.

How to read this against fp16mixed

dit_fp16mixed.trt remains the medium default and this does not change that. fp16-mixed is fp32-exact (teacher-forced cos 1.0000 / 0.9997) for ~4–7% more time. The corrected bf16 is a genuine ~200× improvement on the engine it replaces, but it is not fp32-equivalent, and it is weaker on sm_90 than on sm_120 (0.9923 vs 0.9983; worst single step 0.9484 vs 0.9927) — same ONNX, same recipe, different kernel selection per card.

So: use it if you want the fastest medium tier and can accept a small, audibly-checked deviation from fp32. Use fp16mixed (the default) if you want exactness or bit-reproducibility. bf16 is still not seed-reproducible against fp16-mixed.

Validation

Latent-space fidelity (teacher-forced velocity cosine and free-run latent std against each card's own dit_fp32.trt), amplitude metrics on full-length 380 s renders, and a 6.3-minute listening check on H200 against both PyTorch eager fp32 and dit_fp16mixed — the same standard that caught the original defect, since aggregate metrics at short length did not.

Not run for this engine: the n=128 FAD/CLAP battery that fp16mixed cleared. Given bf16 is a selectable non-default tier that is measurably behind fp16-mixed by design, that gap is stated rather than closed.