diff --git a/optimized/tensorRT/README.md b/optimized/tensorRT/README.md
index 0a84dcf..8d55226 100644
--- a/optimized/tensorRT/README.md
+++ b/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)).
+
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
diff --git a/optimized/tensorRT/build/README.md b/optimized/tensorRT/build/README.md
index 4c23938..de5b1af 100644
--- a/optimized/tensorRT/build/README.md
+++ b/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 /model_config.json --checkpoint /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 |
diff --git a/optimized/tensorRT/build/build_from_onnx.py b/optimized/tensorRT/build/build_from_onnx.py
index e59c6fb..2c5d337 100755
--- a/optimized/tensorRT/build/build_from_onnx.py
+++ b/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"],
diff --git a/optimized/tensorRT/build/make_dit_fp8_smalldit.py b/optimized/tensorRT/build/make_dit_fp8_smalldit.py
new file mode 100644
index 0000000..a56219a
--- /dev/null
+++ b/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 /model_config.json --checkpoint /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..weight' vs torch 'model.'
+ → 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()
diff --git a/optimized/tensorRT/scripts/sa3_trt_core.py b/optimized/tensorRT/scripts/sa3_trt_core.py
index 20a0f54..96df006 100644
--- a/optimized/tensorRT/scripts/sa3_trt_core.py
+++ b/optimized/tensorRT/scripts/sa3_trt_core.py
@@ -163,7 +163,16 @@ def _detect_gpu_arch() -> str:
# decoder clips 2–3% of samples on a 6-min render. Fine at short
# lengths (clean at L=256); prefer fp16mixed for anything long.
# Also not seed-reproducible vs fp16mixed.
-# fp8 — medium ONLY, MAX-SPEED clean tier. fp8 E4M3 on the 176 linear
+# fp8 — all DiTs, fp8 E4M3 on the linear GEMMs (attention + RoPE kept
+# higher precision). On MEDIUM it's a max-speed clean tier (~1.3×
+# over fp16mixed) via the recipe described below. On sm-music /
+# sm-sfx it is a CLEAN WEIGHT-HALVING tier (engine 479 vs 936 MB),
+# only marginally faster (~1.10–1.17×): those DiTs' ~5 ms forward is
+# overhead-bound at batch 1, so fp8's GEMM-math savings barely show.
+# Their fp8 is an fp8-QDQ graft onto the fp16mixed graph (fp8 linears
+# + fp16 fused attention + the fp16mixed fp32 islands, STRONGLY_TYPED),
+# velocity-cos ~0.99 vs eager, clip% at/below fp16mixed. medium recipe:
+# MAX-SPEED clean tier. fp8 E4M3 on the 176 linear
# GEMMs + bf16 fused FMHA (96 nodes) + a BAKED fp32 RoPE constant
# table: position cos/sin are computed host-side at build time and
# frozen as a graph Constant (no in-graph trig), so the island is
@@ -190,14 +199,16 @@ def _detect_gpu_arch() -> str:
# canonical decoder engine. Encoders are FP16-mixed only.
DIT_ENGINE_FILENAME = {
"bf16": "dit_bf16.trt", # medium only; drifts at long sequence
- "fp8": "dit_fp8.trt", # medium only; max-speed clean tier (fp8 lin + baked RoPE)
+ "fp8": "dit_fp8.trt", # all DiTs; fp8 linears (medium: +baked RoPE; small: graft on fp16mixed)
"fp16mixed": "dit_fp16mixed.trt",
"fp32": "dit_fp32.trt",
}
-# DiT precisions actually built per model. bf16 and fp8 are medium-only.
+# DiT precisions actually built per model. bf16 is medium-only; fp8 is available
+# for all three (medium via baked-RoPE/bf16-attn; sm-music/sm-sfx via an fp8-QDQ
+# graft onto their fp16mixed graph — fp8 linears + fp16 fused attn + fp32 islands).
_DIT_PRECISIONS = {
- "sm-music": ("fp16mixed", "fp32"),
- "sm-sfx": ("fp16mixed", "fp32"),
+ "sm-music": ("fp16mixed", "fp8", "fp32"),
+ "sm-sfx": ("fp16mixed", "fp8", "fp32"),
"medium": ("bf16", "fp8", "fp16mixed", "fp32"),
}
# Per-DiT default precision — fp16mixed everywhere. Medium moved off bf16 once
@@ -256,11 +267,8 @@ def get_dit_engine_path(dit_name: str, precision: str = None) -> Path:
f"precision='bf16' is only available for --dit medium (FMHA-fused); "
f"{dit_name} uses standard attention and already fuses in fp16mixed. "
f"Valid for {dit_name}: {_DIT_PRECISIONS.get(dit_name)}")
- if precision == "fp8" and dit_name != "medium":
- raise ValueError(
- f"precision='fp8' is only available for --dit medium (fp8 linears + "
- f"bf16 fused FMHA + baked fp32 RoPE); {dit_name} ships fp16mixed only. "
- f"Valid for {dit_name}: {_DIT_PRECISIONS.get(dit_name)}")
+ # fp8 is available for all three DiTs (medium: baked-RoPE + bf16 attn; sm-music/
+ # sm-sfx: fp8-QDQ graft on their fp16mixed graph). Only bf16 stays medium-only.
return ARCH_DIR / _DIT_SUBDIR[dit_name] / DIT_ENGINE_FILENAME[precision]
@@ -1133,10 +1141,11 @@ def main():
help="DiT engine precision (default is 'fp16mixed' for every model). "
"'fp16mixed' = FP16 trunk + FP32 RMSNorm/RoPE islands with an FMHA-fused "
"FP16 attention core (canonical; fp32-accurate at every length). "
- "'fp8' (MEDIUM ONLY) = the max-speed clean tier: fp8 linears + bf16 fused "
- "FMHA + a baked fp32 RoPE constant table, ~1.3× faster than fp16mixed at "
- "every length and clean at long sequence (std 0.86, 0.000%% clip @6-min); "
- "capped at L<=4096 (the baked table = the SAME-L decoder's own cap). "
+ "'fp8' = fp8 E4M3 linears (attention + RoPE kept higher precision). On MEDIUM "
+ "it's a max-speed clean tier (bf16 fused FMHA + baked fp32 RoPE, ~1.3× faster, "
+ "clean at long sequence, capped at L<=4096). On sm-music/sm-sfx it's a clean "
+ "weight-halving tier (479 vs 936 MB engine, velocity-cos ~0.99 vs eager) that is "
+ "only marginally faster (~1.1×; their small forward is overhead-bound at batch 1). "
"'bf16' (MEDIUM ONLY) = ~3%% faster still, but it evaluates RoPE's angle "
"in bf16 and drifts at long sequence (clips 2-3%% of samples on a 6-min "
"render); fine for short clips, not seed-reproducible vs fp16mixed. "