tensorRT: fp8 tier for the small DiTs (sm-music, sm-sfx)

optimized/tensorRT/README.md

@@ -126,7 +126,7 @@ Omit `--dit` / `--decoder` for an interactive arrow-key picker. Relative
 | `medium`         | `fp16mixed` | FMHA-fused (96 fused attention nodes) **and** fp32-accurate at every length |
 | `sm-music`/`sm-sfx` | `fp16mixed` | standard attention — already fuses in fp16-mixed           |
 
-`--precision` also takes `fp8` and `bf16` (both medium only) and `fp32` explicitly:
+`--precision` also takes `fp8` (all DiTs), `bf16` (medium only) and `fp32` explicitly:
 
 - **`fp16mixed`** — canonical: FP16 trunk, FP32 islands around RMSNorm and RoPE
   generation, and an FP16 attention core (QK^T → Softmax → P·V) so TRT's FMHA fuser
@@ -137,7 +137,8 @@ Omit `--dit` / `--decoder` for an interactive arrow-key picker. Relative
   before QK^T fixed that (**4.3× faster at L=4096**), which retired the reason to
   prefer `bf16`. Engines built before 2026-07 are the slow variant; rebuild with
   `build_from_onnx.py sa3-m`.
-- **`fp8`** — *medium only; the max-speed clean tier, calibrated.* fp8 E4M3 on the 176
+- **`fp8`** — *medium: the max-speed clean tier, calibrated; sm-music/sm-sfx: a clean
+  weight-halving tier (see the end of this bullet).* On **medium**: fp8 E4M3 on the 176
   linear GEMMs + bf16 fused FMHA (96 nodes) + a **baked fp32 RoPE constant table** (position
   cos/sin computed host-side at build and frozen as a graph constant — no in-graph trig,
   so the island is precision-policy-robust and cross-runtime-stable). **~1.3× faster than
@@ -159,6 +160,14 @@ Omit `--dit` / `--decoder` for an interactive arrow-key picker. Relative
   producer: `build/build_dit_bf16.py` RoPE-baker + `build/transplant_scales.py` calibrated-scale
   transplant; identity check: `scripts/verify_fp8_rope.py`; calibration by @ryanontheinside,
   [#47](https://github.com/Stability-AI/stable-audio-3/pull/47)).
+  <br>On **sm-music / sm-sfx** fp8 is a **different, simpler recipe** — fp8 E4M3 grafted onto
+  the linear GEMMs of the fp16mixed graph (attention stays fp16-fused, fp32 islands untouched;
+  no baked RoPE — these DiTs never had bf16's long-angle problem). It's a **clean weight-halving
+  tier** (engine 479 vs 936 MB, velocity-cos **~0.99** vs eager, clip% at/below fp16mixed) that
+  is only **marginally faster (~1.1×)**: a small DiT's ~5 ms forward at batch 1 is overhead-bound,
+  so fp8's GEMM savings barely show. Default stays fp16mixed. Rebuild: `build_from_onnx.py
+  sa3-sm-music-fp8` / `sa3-sm-sfx-fp8`; producer: `build/make_dit_fp8_smalldit.py`. Not
+  seed-reproducible vs fp16mixed.
 - **`bf16`** — *medium only.* Same `dit.onnx` as fp32, built with `BuilderFlag.BF16`;
   a uniform bf16 trunk also lets the FMHA fuser fire, and it is ~3% faster than
   `fp16mixed`. **But it drifts at long sequence**: weakly-typed BF16 lets TRT

optimized/tensorRT/build/README.md

@@ -324,6 +324,39 @@ repo) → `build_dit_fp8.py` (max-PTQ + per-channel weight scales; that builder
 merged here). Everyday consumers never recalibrate — they pull the published calibrated
 `dit_fp8.onnx`.
 
+## Small-DiT `fp8` — sm-music / sm-sfx (a different, simpler recipe)
+
+`sm-music` and `sm-sfx` also ship a selectable **`fp8`** engine (`--precision fp8`), but it is
+**not** the medium's baked-RoPE recipe — those DiTs use standard (non-differential) attention and
+never had the bf16 long-angle RoPE problem, so there is nothing to bake. Their fp8 is a straight
+**graft of fp8 E4M3 Q/DQ onto the linear GEMMs of the known-good `dit_fp16mixed.onnx`** — attention
+stays fp16-fused and the fp32 RMSNorm/RoPE islands are left exactly as the fp16mixed producer made
+them. Built `STRONGLY_TYPED` (`build_from_onnx.py sa3-sm-music-fp8` / `sa3-sm-sfx-fp8`); the QDQ
+carry the precision. Identity: 186 fp8 GEMMs + fp16 fused attention + the fp16mixed fp32 islands.
+
+Positioning is honest: this is a **clean weight-halving tier** (engine 479 vs 936 MB, velocity-cos
+~0.99 vs eager, clip% at or below fp16mixed), only **marginally faster** (~1.10–1.17×) — a small
+DiT's ~5 ms forward at batch 1 is overhead-bound, so fp8's GEMM-math savings barely show. Default
+stays `fp16mixed`; fp8 is for when the smaller engine / weight footprint helps. Not seed-reproducible
+vs fp16mixed.
+
+> ⚠ Do **not** produce these with `build_dit_fp8.py` (#47's ModelOpt path): on the small graphs its
+> island-flatten + reapply does not restore the fp32 islands correctly and the engine collapses to
+> velocity-cos ~0.69 with clipping (the GEMMs are fine — it's the islands). Grafting onto the
+> fp16mixed ONNX keeps the islands correct by construction.
+
+**Producer (refresh the ONNX).** `make_dit_fp8_smalldit.py` calibrates per-linear activation scales
+from the eager model (own-domain few-shot prompts + one full render) and grafts the fp8 Q/DQ. Two
+fp16-trunk specifics vs the medium inserter: Q/DQ scales are **FLOAT16** (fp16 trunk → DQ must output
+fp16) and floored at 1e-4 (fp16 underflows tiny scales to 0, which TRT rejects):
+
+```bash
+python make_dit_fp8_smalldit.py \
+    --model-config <ckpt>/model_config.json --checkpoint <ckpt>/model.safetensors \
+    --fp16mixed-onnx onnx/sa3-sm-music/dit_fp16mixed.onnx \
+    --domain Music --out onnx/sa3-sm-music/dit_fp8.onnx      # --domain SFX for sm-sfx
+```
+
 ## File map
 
 | File | Role | Flow |

optimized/tensorRT/build/build_from_onnx.py

@@ -142,6 +142,35 @@
         "profile":      _DIT_PROFILE,
         "plugin":       False,
     },
+    # SA3 small DiTs in fp8 — SELECTABLE (default stays fp16mixed). fp8 E4M3 on the
+    # 186 linear GEMMs, attention left fp16-fused and the fp32 RMSNorm/RoPE islands
+    # intact — an fp8-QDQ graft onto the fp16mixed graph (build/make_dit_fp8_smalldit.py),
+    # NOT the medium's baked-RoPE recipe (these DiTs never had the bf16 long-angle
+    # problem). Built STRONGLY_TYPED: the QDQ nodes carry fp8; TRT fires fp8 tensor-core
+    # GEMMs on the linears while the fp16 FMHA fuser still runs the attention. Same
+    # _DIT_PROFILE (batch=1, dynamic L∈[1,4096]) → identical CLI/feature surface.
+    # This is a CLEAN WEIGHT-HALVING tier (479 vs 936 MB, velocity-cos ~0.99 vs eager,
+    # clip% at/below fp16mixed), only marginally faster (~1.1×): the small DiTs' ~5 ms
+    # forward is overhead-bound at batch 1, so fp8's GEMM savings barely show. sm-* fp8
+    # is NOT seed-reproducible vs fp16mixed.
+    "sa3-sm-music-fp8": {
+        "onnx_hf":      ["sa3-sm-music/dit_fp8.onnx", "sa3-sm-music/dit_fp8.onnx.data"],
+        "trt_local":    "sa3-sm-music/dit_fp8.trt",
+        "flags":        set(),         # STRONGLY_TYPED + the fp8 QDQ carry the precision
+        "network":      "STRONGLY_TYPED",
+        "workspace_gb": 16,
+        "profile":      _DIT_PROFILE,
+        "plugin":       False,
+    },
+    "sa3-sm-sfx-fp8": {
+        "onnx_hf":      ["sa3-sm-sfx/dit_fp8.onnx", "sa3-sm-sfx/dit_fp8.onnx.data"],
+        "trt_local":    "sa3-sm-sfx/dit_fp8.trt",
+        "flags":        set(),
+        "network":      "STRONGLY_TYPED",
+        "workspace_gb": 16,
+        "profile":      _DIT_PROFILE,
+        "plugin":       False,
+    },
     "sa3-m": {
         # 2.9 GB external-data sidecar travels alongside.
         "onnx_hf":      ["sa3-m/dit_fp16mixed.onnx", "sa3-m/dit_fp16mixed.onnx.data"],

optimized/tensorRT/build/make_dit_fp8_smalldit.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+"""Produce the fp8 ONNX for the SMALL DiTs (sm-music / sm-sfx) by grafting fp8 E4M3 Q/DQ onto
+the LINEAR GEMMs of the model's fp16mixed ONNX — leaving attention fp16-fused and the fp32
+RMSNorm/RoPE islands intact. This is deliberately NOT the medium's fp8 recipe (baked RoPE +
+bf16 attention, build_dit_bf16.py + build_dit_fp8.py): the small DiTs never had the bf16
+long-angle RoPE problem, so their fp8 is a straight graft on the known-good fp16mixed graph.
+
+Why a dedicated script (vs ModelOpt/build_dit_fp8.py): running ModelOpt PTQ on these models
+flattens the fp32 islands and its island-reapply (tuned for the medium graph) does NOT restore
+them correctly here — the resulting engine drops to velocity-cos ~0.69 with clipping. Grafting
+onto the fp16mixed ONNX keeps the islands correct by construction (velocity-cos ~0.99 vs eager).
+
+Two fp16-trunk specifics vs dit_fp8_max/make_fp8_onnx.py (which targets the fp32-trunk medium):
+  * Q/DQ scales are FLOAT16 (the trunk is fp16 → DequantizeLinear must output fp16, else TRT
+    sees Half-vs-Float at the residual Adds).
+  * scales floored at 1e-4 (fp16 underflows anything <~6e-5 to 0, and TRT rejects a non-positive
+    scale; only bites zero-input layers e.g. to_local_embed when local_add_cond=0 — harmless).
+
+Activation scales are calibrated from the eager model on the model's own-domain few-shot prompts
+(Music for sm-music, SFX for sm-sfx) plus one full-length render, with a margin so nothing clips.
+Scale VALUES affect accuracy only, not whether fp8 fires or the latency.
+
+    python make_dit_fp8_smalldit.py \
+        --model-config <ckpt>/model_config.json --checkpoint <ckpt>/model.safetensors \
+        --fp16mixed-onnx onnx/sa3-sm-music/dit_fp16mixed.onnx \
+        --domain Music --out onnx/sa3-sm-music/dit_fp8.onnx
+
+Then compile with build_from_onnx.py sa3-sm-music-fp8 (STRONGLY_TYPED; the QDQ carry fp8).
+"""
+import argparse, os, re, time
+from collections import Counter
+from pathlib import Path
+import numpy as np
+import onnx
+from onnx import TensorProto, helper, numpy_helper
+
+E4M3_MAX = 448.0
+SCALE_DT = TensorProto.FLOAT16
+SCALE_FLOOR = 1e-4
+
+
+def calibrate_act_scales(model_config, checkpoint, fp16mixed_onnx, domain, margin, device):
+    """Per-linear activation max|x| from the eager model → {onnx_node_name: scale}. Maps ONNX
+    linear nodes to torch modules by weight name (ONNX 'dit.<path>.weight' vs torch 'model.<path>'
+    → match on the suffix after the first dotted component)."""
+    import torch
+    import torch.nn.functional as F
+    import make_calib as MK  # repo sibling; wraps load_diffusion_cond + StableAudioModel
+    from stable_audio_3.interface.reprompt import SYSTEM_PROMPTS, _extract_examples
+    torch.set_grad_enabled(False)
+
+    sa3 = MK._load_model(Path(model_config), Path(checkpoint), device)
+    dit = sa3.dit
+    suffix = lambda nm: nm.split(".", 1)[1] if "." in nm else nm
+    mods = {}
+
+    class Q:
+        def __init__(s, l): s.l = l; s.amax = 1e-9
+        def __call__(s, x): s.amax = max(s.amax, float(x.abs().amax())); return F.linear(x, s.l.weight, s.l.bias)
+    for name, mod in dit.named_modules():
+        if isinstance(mod, torch.nn.Linear) and min(mod.in_features, mod.out_features) >= 128:
+            mod._q = Q(mod); mod.forward = (lambda m: (lambda x: m._q(x)))(mod); mods[suffix(name)] = mod._q
+
+    prompts = _extract_examples(SYSTEM_PROMPTS[domain])[:14]
+    render = {"Music": "Genre: House, Subgenre: Deep House, BPM: 122 BPM, Tempo: Medium, "
+                       "VocalType: Instrumental, TrackType: Music, Grade: Neutral",
+              "SFX": "Heavy rain on a tin roof with distant thunder, steady continuous downpour"}.get(domain, prompts[0])
+    for i, p in enumerate(prompts):
+        sa3.generate(prompt=p, duration=MK.DEFAULT_DURATION_S, steps=8, cfg_scale=1.0,
+                     sampler_type="pingpong", seed=MK.DEFAULT_SEED + i, duration_padding_sec=0.0, return_latents=True)
+    sa3.generate(prompt=render, duration=1292 * 4096 / 44100.0, steps=8, cfg_scale=1.0,
+                 sampler_type="pingpong", seed=6000, duration_padding_sec=0.0, return_latents=True)
+    print(f"  calibrated {len(mods)} linears on {len(prompts)} {domain} prompts + one full render", flush=True)
+
+    m = onnx.load(fp16mixed_onnx, load_external_data=False); g = m.graph
+    inits = {i.name for i in g.initializer}; prod = {o: n for n in g.node for o in n.output}
+    def wsrc(n):
+        w = n.input[1]
+        if w in inits: return w
+        p = prod.get(w)
+        return p.input[0] if (p is not None and p.op_type == "Transpose" and p.input and p.input[0] in inits) else None
+    node_scale = {}
+    for n in g.node:
+        if n.op_type != "MatMul": continue
+        w = wsrc(n)
+        if w is None: continue
+        q = mods.get(suffix(w[:-len(".weight")] if w.endswith(".weight") else w))
+        if q is not None:
+            node_scale[n.name] = max(q.amax * margin / E4M3_MAX, SCALE_FLOOR)
+    gscale = max(max(q.amax for q in mods.values()) * margin / E4M3_MAX, SCALE_FLOOR)
+    return node_scale, gscale
+
+
+def topo_sort(g):
+    avail = {i.name for i in g.initializer} | {i.name for i in g.input} | {""}
+    remaining, result = list(g.node), []
+    while remaining:
+        nxt, prog = [], False
+        for n in remaining:
+            if all(i in avail for i in n.input):
+                result.append(n); [avail.add(o) for o in n.output]; prog = True
+            else:
+                nxt.append(n)
+        remaining = nxt
+        if not prog: raise RuntimeError(f"topo stuck: {len(remaining)}")
+    del g.node[:]; g.node.extend(result)
+
+
+def graft_fp8(fp16mixed_onnx, node_scale, gscale, out):
+    model = onnx.load(fp16mixed_onnx, load_external_data=True); g = model.graph
+    have = False
+    for op in model.opset_import:
+        if op.domain in ("", "ai.onnx"):
+            have = True
+            if op.version < 19: op.version = 19
+    if not have: model.opset_import.append(helper.make_opsetid("", 19))
+    if model.ir_version < 9: model.ir_version = 9
+    inits = {i.name: i for i in g.initializer}
+    prod = {o: n for n in g.node for o in n.output}
+    g.initializer.append(helper.make_tensor("fp8_zero", TensorProto.FLOAT8E4M3FN, [], [0.0]))
+    new_nodes, new_inits, made, skipped = [], [], 0, 0
+    for node in [n for n in g.node if n.op_type == "MatMul"]:
+        Wname = node.input[1]; via_t = tnode = w_src = Warr = None; via_t = False
+        if Wname in inits:
+            w_src = Wname; Warr = numpy_helper.to_array(inits[Wname])
+        else:
+            p = prod.get(Wname)
+            if p is not None and p.op_type == "Transpose" and p.input and p.input[0] in inits:
+                tnode = p; w_src = p.input[0]; via_t = True; Warr = numpy_helper.to_array(inits[w_src])
+        if Warr is None or Warr.ndim != 2:
+            skipped += 1; continue   # attention BMM (no weight initializer)
+        pfx = node.name.strip("/").replace("/", "_")
+        w_scale = float(max(np.abs(Warr.astype(np.float32)).max() / E4M3_MAX, SCALE_FLOOR))
+        a_scale = float(max(node_scale.get(node.name, gscale), SCALE_FLOOR))
+        new_inits += [helper.make_tensor(f"{pfx}_wscale", SCALE_DT, [], [w_scale]),
+                      helper.make_tensor(f"{pfx}_ascale", SCALE_DT, [], [a_scale])]
+        aq, adq = f"{pfx}_aq", f"{pfx}_adq"
+        new_nodes += [helper.make_node("QuantizeLinear", [node.input[0], f"{pfx}_ascale", "fp8_zero"], [aq], name=f"{pfx}_Qa"),
+                      helper.make_node("DequantizeLinear", [aq, f"{pfx}_ascale", "fp8_zero"], [adq], name=f"{pfx}_DQa")]
+        node.input[0] = adq
+        wq, wdq = f"{pfx}_wq", f"{pfx}_wdq"
+        new_nodes += [helper.make_node("QuantizeLinear", [w_src, f"{pfx}_wscale", "fp8_zero"], [wq], name=f"{pfx}_Qw"),
+                      helper.make_node("DequantizeLinear", [wq, f"{pfx}_wscale", "fp8_zero"], [wdq], name=f"{pfx}_DQw")]
+        if via_t:
+            for i, inp in enumerate(tnode.input):
+                if inp == w_src: tnode.input[i] = wdq
+        else:
+            node.input[1] = wdq
+        made += 1
+    g.initializer.extend(new_inits); g.node.extend(new_nodes); topo_sort(g)
+    if os.path.exists(out): os.remove(out)
+    if os.path.exists(out + ".data"): os.remove(out + ".data")
+    onnx.save(model, out, save_as_external_data=True, all_tensors_to_one_file=True,
+              location=os.path.basename(out) + ".data", size_threshold=1024)
+    c = Counter(n.op_type for n in g.node)
+    print(f"  fp8 Q/DQ on {made} linear MatMuls ({skipped} attention BMMs skipped); "
+          f"Q={c.get('QuantizeLinear',0)} DQ={c.get('DequantizeLinear',0)} Softmax={c.get('Softmax',0)}", flush=True)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--model-config", required=True)
+    ap.add_argument("--checkpoint", required=True)
+    ap.add_argument("--fp16mixed-onnx", required=True, help="the model's canonical dit_fp16mixed.onnx")
+    ap.add_argument("--out", required=True, help="output dit_fp8.onnx (a .data sidecar is written alongside)")
+    ap.add_argument("--domain", default="Music", choices=["Music", "SFX", "Instrument", "One-shot"],
+                    help="reprompt few-shot domain for activation calibration (Music for sm-music, SFX for sm-sfx)")
+    ap.add_argument("--margin", type=float, default=1.35, help="activation-scale headroom so nothing clips")
+    ap.add_argument("--device", default="cuda")
+    a = ap.parse_args()
+    t0 = time.time()
+    print(f"[make_dit_fp8_smalldit] calibrating ({a.domain}) ...", flush=True)
+    node_scale, gscale = calibrate_act_scales(a.model_config, a.checkpoint, a.fp16mixed_onnx, a.domain, a.margin, a.device)
+    print(f"  {len(node_scale)} node scales, global={gscale:.5f}; grafting fp8 ...", flush=True)
+    graft_fp8(a.fp16mixed_onnx, node_scale, gscale, a.out)
+    print(f"DONE -> {a.out} ({time.time()-t0:.0f}s)", flush=True)
+
+
+if __name__ == "__main__":
+    main()