cortexelus commited on
Commit
59bd967
Β·
verified Β·
1 Parent(s): 06debc3

Fix the bf16 medium DiT: bake RoPE's tables into the graph (clipping 3.112% -> 0.014%)

Browse files

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](https://github.com/Stability-AI/stable-audio-3/pull/81).

## What changed

| file | size | note |
|---|---|---|
| `onnx/sa3-m/dit_bf16.onnx` | **4.8 MB** | new β€” RoPE tables baked as fp32 constants. Shares `dit.onnx.data`, so no new weight copy |
| `tensorRT/sm_90/sa3-m/dit_bf16.trt` | 2.92 GB | replaced |
| `tensorRT/sm_120/sa3-m/dit_bf16.trt` | 2.92 GB | replaced |

## 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:

| shipped engine @L=4092 | value |
|---|---|
| teacher-forced velocity cos vs fp32 | **0.5477** |
| 380 s render clipping | **3.112%** |
| crest (peak Γ· RMS) | **2.50** β€” fp32 is 6.25 |

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.

| | sm_120 Β· RTX PRO 4500 Blackwell (82 SM, 165 W) | sm_90 Β· H200 (132 SM, 700 W) |
|---|---|---|
| teacher-forced cos vs fp32 @L=256 / 1292 / 4092 | 0.9993 / 0.9990 / **0.9983** | 0.9987 / 0.9990 / **0.9923** |
| latent std ratio vs fp32 @L=4092 | 1.13Γ— | 1.03Γ— |
| 380 s clipping / crest | **0.014%** / 4.54 | 0.001–0.117% / 4.23–8.64 |
| ms per forward @L=4096 | **160.8** (fp16mixed 168.2) | **39.3** (fp16mixed 41.9) |
| fused MHA nodes | 96 | 96 |

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.

onnx/sa3-m/dit_bf16.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c70d3e0db208a570bb163518e4a63ee59eb61508038bc11e874e081c5c449ebc
3
+ size 4765460
tensorRT/sm_120/sa3-m/dit_bf16.trt CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:dc2fcfd7f1e94729164bbde2a5e129ce0bfa6bb0499b37f24b3d9ae4a05518c6
3
- size 2924359388
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:43a2c0600c1949255c546fa6107a1f5a4f828fb36a7efcb59ff19625645204c7
3
+ size 2924806340
tensorRT/sm_90/sa3-m/dit_bf16.trt CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:2c8655a1b67ff11c6a5b51433e7a2323a87eb8752f5469f78fe9499e83aefa52
3
- size 2919426020
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d5ac9e56dde6cd7e3693f79c771ba1a3770000c072dc77412e1f936a70c3c9c4
3
+ size 2919609428