Instructions to use stabilityai/stable-audio-3-optimized with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Stable Audio 3
How to use stabilityai/stable-audio-3-optimized with Stable Audio 3:
# No code snippets available yet for this library. # To use this model, check the repository files and the library's documentation. # Want to help? PRs adding snippets are welcome at: # https://github.com/huggingface/huggingface.js
- Notebooks
- Google Colab
- Kaggle
Fix medium DiT FP16-mixed ONNX: bound the RoPE island so attention fuses (4.3x faster, cos 1.0000)
Browse filesFixes the medium DiT's FP16-mixed ONNX so the attention core runs FP16 and TensorRT's FMHA fuser can fire. **Graph-only change — the 2.9 GB weight sidecar is byte-identical (sha256 `a038ae6d…`) and is not re-uploaded.**
## The bug
`build_dit_fp16mixed.py`'s island passes kept the entire O(L²) attention core in FP32:
- `find_fp32_islands()` blocks **every** `Softmax` (section C), and
- the RoPE island emits q/k in FP32 with nothing casting them back, so `fix_dtype_mismatches()` then kept the QK^T `MatMul` FP32 to match its operands.
Result: **0 fused MHA nodes** and **4.3× slower at L=4096** than necessary — for no accuracy gain.
Neither FP32 region was ever required:
- PyTorch does not keep the attention core in FP32. `apply_rotary_pos_emb` ends with `t.to(out_dtype)` — it computes in FP32 and casts *back* to the autocast dtype — and attention is then `F.scaled_dot_product_attention(q, k, v)`, a fused kernel taking FP16 inputs. The FP32 softmax island has no counterpart in the source.
- TRT's FMHA kernels accumulate softmax in FP32 internally, so an explicit FP32 softmax buys nothing the kernel does not already do.
## The fix
A new `bound_attention_core()` pass (in the code repo) terminates the RoPE island before QK^T, mirroring the source's cast-back:
```
before: Mul(fp32) → MatMul(fp32,fp32) → Cast(fp32)[no-op] → Softmax(fp32) → Cast(fp16) → MatMul(P·V)
after: Mul(fp32) → Cast(fp16)×2 → MatMul(fp16) → Softmax(fp16) → MatMul(P·V)
```
The RMSNorm islands are untouched — those guard variance *overflow* (FP16's 5-bit exponent), a different failure mode from precision, and are still required. 96/96 attentions convert; initializer dtypes are unchanged at 428 FP16 / 98 FP32.
## Measured (RTX PRO 4500 Blackwell sm_120 · TRT 10.15.1.29 · TF32 off)
| medium DiT @L=4096 | published | this PR |
|---|---|---|
| fused MHA nodes | 0 | **96** |
| ms / forward | 714 | **167** (4.3×) |
| teacher-forced velocity cos vs FP32 | 0.9998 | **1.0000** |
| free-run latent std vs FP32 | 1.00× | 1.00× |
| 380 s render: clipped / crest | 0.000% / 6.62 | 0.001% / 6.51 |
Faster **and** marginally more accurate — it is closer to eager semantics than the graph it replaces. Cross-checked on H200 (sm_90): 41.4 ms vs 180.9 ms, same 96 fused MHA, so this is not arch-specific.
## Compatibility
- Engine filename, ONNX filename, and the `fp16mixed` precision name are all unchanged; no consumer code changes are required to *use* it.
- **Output changes.** This is a bug fix, so renders differ from the previous engine at the same seed. The engine remains deterministic and reproducible going forward, but it is not bit-identical to the old one.
- Rebuild affected engines with `build_from_onnx.py sa3-m` (unchanged invocation). Must be built `STRONGLY_TYPED` — weakly-typed + `BuilderFlag.FP16` lets TRT re-cast the FP32 RMSNorm islands and reintroduces the overflow this recipe exists to prevent (measured: teacher-forced cos collapses to 0.88).
- Only the medium DiT is affected. `sa3-sm-music` / `sa3-sm-sfx` use standard (non-differential) attention and already fuse 40/40 MHA in their published FP16-mixed engines, so their RoPE island is already correctly bounded.
Code-side changes (the `bound_attention_core()` pass, `--no-bound-attn` escape hatch for reproducing the old engines, and docs) are in the companion PR to `Stability-AI/stable-audio-3`.
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:018448748ac272a8a7e701fc52e9a3f3b04909e04e478a1c27e3c0f685702687
|
| 3 |
+
size 4247760
|