stabilityai

/

stable-audio-3-optimized

like 27

Follow

Stability AI 38.2k

Fixes 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

NVIDIA RTX PRO 4500 Blackwell Server Edition (sm_120 · 82 SM · 32 GB GDDR7 · 165 W) · TRT 10.15.1.29 · TF32 off · batch 1 · CUDA-event median of 7 after 3 warmup, dedicated idle GPU.

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 will be faster in absolute terms — read the ratio, not the absolute ms, as the portable result.

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.