cortexelus commited on
Commit
b4300a9
·
verified ·
1 Parent(s): 08c64b9

Fix medium DiT FP16-mixed ONNX: bound the RoPE island so attention fuses (4.3x faster, cos 1.0000)

Browse files

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 (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`.

Files changed (1) hide show
  1. onnx/sa3-m/dit_fp16mixed.onnx +2 -2
onnx/sa3-m/dit_fp16mixed.onnx CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:ccc51b4e067ad1132b84be71f11797da4eae2b3544977bcd52466d23e7a824d7
3
- size 4248874
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:018448748ac272a8a7e701fc52e9a3f3b04909e04e478a1c27e3c0f685702687
3
+ size 4247760