bf16 medium DiT: FMHA-fused TRT engine build recipe + CLI (bf16 = medium default, fp16-mixed kept) (#73)

Author: Cortexelus

optimized/tensorRT/README.md

@@ -117,6 +117,35 @@ package install + a ~5 GB engine download).
 Omit `--dit` / `--decoder` for an interactive arrow-key picker. Relative
 `--out` paths land in `output/`; absolute paths are honoured as-is.
 
+### DiT precision (`--precision`)
+
+The default resolves per model — no flag needed for the recommended setup:
+
+| model            | default     | why                                                        |
+|------------------|-------------|------------------------------------------------------------|
+| `medium`         | `bf16`      | FMHA-fused (0 → 96 fused attention nodes) → **1.8× @256 / 4.7× @4096** vs fp16-mixed, within the perceptual floor |
+| `sm-music`/`sm-sfx` | `fp16mixed` | standard attention — already fuses in fp16-mixed           |
+
+`--precision` also takes `fp16mixed` and `fp32` explicitly:
+
+- **`bf16`** — *medium only.* Same `dit.onnx` as fp32, built with `BuilderFlag.BF16`;
+  bf16 carries fp32's range so the FP32-softmax islands vanish and TRT's FMHA fuser
+  fires. **Not seed-reproducible vs fp16-mixed** — the medium DiT uses *differential*
+  attention (cancellation-sensitive), so bf16 is a different-but-equal take per seed
+  (quality within the re-seed floor; exact samples differ). Same varlen profile
+  (L 1..4096, opt 1292) and full mode/feature set as the other precisions.
+- **`fp16mixed`** — canonical FP16 trunk + FP32 islands. Bit-reproducible; use it when
+  you need exact per-seed reproducibility or max per-step fidelity.
+- **`fp32`** — pure FP32, bit-equivalent to PyTorch eager (~2× slower, ~2× VRAM).
+
+```bash
+# medium defaults to bf16 (fast); force the reproducible engine instead:
+./sa3 --prompt "..." --dit medium --decoder same-l --precision fp16mixed
+```
+
+The TRT DiT engines are static batch=1 (the ONNX bakes batch=1), so CFG runs as a
+sequential cond+uncond dual-pass at batch=1 for every precision.
+
 ## Speed & memory
 
 Measured on **H100 SXM 80 GB** at `--steps 8` (rf-denoiser sweet spot).

optimized/tensorRT/build/build.py

@@ -67,7 +67,13 @@ def _from_onnx(name):
      "outputs": ["same-l/dec_dynamic_triton_swa.trt"]},
     # DiT engines now build from pre-processed FP16-mixed ONNX on HF;
     # build_from_onnx.py does the simple STRONGLY_TYPED compile.
-    {"label": "DiT medium  (SA3-M, FP16-mixed)",
+    # Medium ships TWO engines: bf16 (FMHA-fused, the speed DEFAULT) built from
+    # the raw fp32 dit.onnx with BuilderFlag.BF16, and fp16-mixed (canonical,
+    # bit-reproducible, kept selectable). sm-music/sm-sfx are fp16-mixed only.
+    {"label": "DiT medium  (SA3-M, bf16 — FMHA-fused, medium DEFAULT)",
+     "command": _from_onnx("sa3-m-bf16"),
+     "outputs": ["sa3-m/dit_bf16.trt"]},
+    {"label": "DiT medium  (SA3-M, FP16-mixed — selectable, bit-reproducible)",
      "command": _from_onnx("sa3-m"),
      "outputs": ["sa3-m/dit_fp16mixed.trt"]},
     {"label": "DiT sm-music (FP16-mixed)",

optimized/tensorRT/build/build_from_onnx.py

@@ -141,6 +141,33 @@
         "profile":      _DIT_PROFILE,
         "plugin":       False,
     },
+    # SA3 medium DiT in bf16 — the medium SPEED DEFAULT (fp16-mixed above stays
+    # selectable). Reuses the SAME raw fp32 dit.onnx as the fp32 variant, but
+    # builds with BuilderFlag.BF16 + EXPLICIT_BATCH (NOT STRONGLY_TYPED): bf16
+    # carries fp32's dynamic range, so the FP32-softmax islands that fp16-mixed
+    # kept vanish and TRT 10.15's FMHA fuser fires (0 → 96 fused _gemm_mha_v2
+    # nodes) → measured ~1.76×@L=256 / 4.70×@L=4096 vs fp16-mixed, within the
+    # perceptual re-seed floor (FAD 0.59× floor, CLAP within spread, n=128).
+    # It is a build-time PRECISION RECIPE only — no new ONNX, no graph change.
+    # medium-only: sm-music / sm-sfx use standard attention and already fuse in
+    # their fp16-mixed engines. bf16 is NOT seed-reproducible vs fp16-mixed
+    # (differential attention is cancellation-sensitive) — a different-but-equal
+    # take per seed; use fp16-mixed for bit-reproducibility.
+    # Same _DIT_PROFILE (batch=1, dynamic L∈[1,4096], opt=1292) as every DiT
+    # engine → identical CLI/feature surface (varlen, CFG sequential dual-pass,
+    # neg-prompt/APG, a2a, inpaint). The DiT ONNX bakes batch=1, so there is no
+    # varbatch axis on ANY TRT DiT engine (CFG is a sequential dual-pass, as in
+    # fp16-mixed) — fusion is verified to survive the full L profile at batch=1.
+    "sa3-m-bf16": {
+        # 5.8 GB external-data sidecar (fp32 weights) travels alongside.
+        "onnx_hf":      ["sa3-m/dit.onnx", "sa3-m/dit.onnx.data"],
+        "trt_local":    "sa3-m/dit_bf16.trt",
+        "flags":        {"BF16"},
+        "network":      "EXPLICIT_BATCH",
+        "workspace_gb": 16,
+        "profile":      _DIT_PROFILE,
+        "plugin":       False,
+    },
     # ── FP32 variants ────────────────────────────────────────────────────
     # DiT FP32: read the unsurgered FP32 ONNX directly (dit.onnx), build
     # STRONGLY_TYPED. ~2× the engine size of FP16-mixed, ~2× slower, but

optimized/tensorRT/scripts/sa3_trt.py

@@ -393,7 +393,7 @@ class SA3Inference:
     DEFAULT_SIGMA_MAX = 1.0  # mega-graph path requires this
 
     def __init__(self, dit: str, decoder: str, *,
-                 precision: str = "fp16mixed",
+                 precision: str | None = None,
                  default_T_lat: int = 324, default_steps: int = 8,
                  default_seconds: float = 30.0,
                  models_dir: Path | None = None,
@@ -404,9 +404,13 @@ def __init__(self, dit: str, decoder: str, *,
         Args:
             dit:            one of DIT_CHOICES — "sm-music" / "sm-sfx" / "medium"
             decoder:        one of DECODER_PATHS — "same-s" / "same-l"
-            precision:      "fp16mixed" (default, fastest) or "fp32" (bit-equiv
-                            PyTorch eager, ~2× slower). Engines auto-download
-                            from HF if the requested precision file is missing.
+            precision:      None (default → per model: "bf16" for medium, else
+                            "fp16mixed"), or an explicit "bf16" (medium only;
+                            FMHA-fused, ~1.8-4.7× faster, within the perceptual
+                            floor, not seed-reproducible vs fp16mixed),
+                            "fp16mixed" (canonical, bit-reproducible), or "fp32"
+                            (bit-equiv PyTorch eager, ~2× slower). Engines auto-
+                            download from HF if the requested file is missing.
             default_T_lat:  latent length to build the initial graph at
             default_steps:  pingpong steps for the initial graph
             default_seconds: duration condition for the initial graph (used for
@@ -420,8 +424,15 @@ def __init__(self, dit: str, decoder: str, *,
             raise ValueError(f"unknown dit={dit!r}; valid: {list(DIT_CHOICES)}")
         if decoder not in DECODER_PATHS:
             raise ValueError(f"unknown decoder={decoder!r}; valid: {list(DECODER_PATHS)}")
+        # Resolve precision default per model: bf16 for medium (FMHA-fused speed
+        # default), fp16-mixed for sm-music/sm-sfx. bf16 is medium-only.
+        if precision is None:
+            precision = canon.default_precision(dit)
         if precision not in canon.PRECISIONS:
             raise ValueError(f"unknown precision={precision!r}; valid: {canon.PRECISIONS}")
+        if precision == "bf16" and dit != "medium":
+            raise ValueError("precision='bf16' is only available for dit='medium' "
+                             "(sm-music/sm-sfx already fuse in fp16mixed)")
 
         # Quiet: patch canon's stage/sub/_stage_vram to no-ops so loading
         # doesn't spam stdout (gradio in particular wants a clean log).
@@ -633,9 +644,12 @@ def main():
     ap.add_argument("--inpaint-range", default=None)
     ap.add_argument("--dit", choices=list(DIT_CHOICES.keys()), default=None)
     ap.add_argument("--decoder", choices=list(DECODER_PATHS.keys()), default=None)
-    ap.add_argument("--precision", choices=list(canon.PRECISIONS), default="fp16mixed",
-                    help="Engine precision: 'fp16mixed' (default, fast) or 'fp32' "
-                         "(bit-equiv PyTorch eager, slower). Auto-downloads from HF.")
+    ap.add_argument("--precision", choices=list(canon.PRECISIONS), default=None,
+                    help="DiT engine precision. Default resolves per model: 'bf16' for medium "
+                         "(FMHA-fused, ~1.8-4.7× faster, within perceptual floor; not "
+                         "seed-reproducible vs fp16mixed), 'fp16mixed' for sm-music/sm-sfx. "
+                         "'fp16mixed' = canonical/bit-reproducible. 'fp32' = bit-equiv PyTorch "
+                         "eager, slower. bf16 is medium-only. Auto-downloads from HF.")
     ap.add_argument("--models-dir", default=str(canon.MODELS_DIR))
     ap.add_argument("--seconds", type=float, default=30.0)
     ap.add_argument("--steps", type=int, default=8)

optimized/tensorRT/scripts/sa3_trt_core.py

@@ -103,7 +103,7 @@ def _detect_gpu_arch() -> str:
 DIT_ENGINE_FILES = {
     "sm-music": ["sa3-sm-music/dit_fp16mixed.trt"],
     "sm-sfx":   ["sa3-sm-sfx/dit_fp16mixed.trt"],
-    "medium":   ["sa3-m/dit_fp16mixed.trt"],
+    "medium":   ["sa3-m/dit_bf16.trt"],   # bf16 (FMHA-fused) is the medium default
 }
 DECODER_FILES = {
     "same-s": [
@@ -129,7 +129,7 @@ def _detect_gpu_arch() -> str:
                  "default_decoder": "same-s"},
     "sm-sfx":   {"engine": ARCH_DIR / "sa3-sm-sfx" / "dit_fp16mixed.trt",
                  "default_decoder": "same-s"},
-    "medium":   {"engine": ARCH_DIR / "sa3-m" / "dit_fp16mixed.trt",
+    "medium":   {"engine": ARCH_DIR / "sa3-m" / "dit_bf16.trt",   # bf16 = medium default
                  "default_decoder": "same-l"},
 }
 DECODER_PATHS = {
@@ -144,51 +144,94 @@ def _detect_gpu_arch() -> str:
 
 # ─── Precision-keyed engine maps ─────────────────────────────────────────
 #
-# The canonical engines are FP16-mixed (FP16 trunk + FP32 islands around
-# RMSNorm / Softmax / RoPE). Pure-FP32 variants are also published — same
-# numerical behavior as PyTorch eager FP32. Use `--precision fp32` on the
-# CLI to pick them; default is `fp16mixed`.
+# Three DiT precisions:
+#   bf16       — medium ONLY. Same dit.onnx as fp32, built with BuilderFlag.BF16
+#                (EXPLICIT_BATCH). bf16 carries fp32's range, so the FP32-softmax
+#                islands vanish and TRT 10.15's FMHA fuser fires (0 → 96 fused
+#                _gemm_mha_v2 nodes) → 1.76×@256 / 4.70×@4096 vs fp16-mixed,
+#                within the perceptual re-seed floor (FAD 0.59× floor, n=128).
+#                NOT seed-reproducible vs fp16-mixed (differential attention is
+#                cancellation-sensitive → a different-but-equal take per seed).
+#                THE MEDIUM DEFAULT. (sm-music/sm-sfx use standard attention and
+#                already fuse in fp16-mixed — no bf16 engine for them.)
+#   fp16mixed  — canonical (FP16 trunk + FP32 islands around RMSNorm/Softmax/
+#                RoPE). Kept selectable for bit-reproducibility / max per-step
+#                fidelity. The sm-music / sm-sfx default.
+#   fp32       — pure-FP32, bit-equivalent to PyTorch eager. ~2× size/latency.
 #
 # The lookup tables below resolve the engine filename per (dit/decoder,
-# precision). Encoders are FP16-mixed only.
+# precision). The bf16 DiT recipe is a build-time precision change only (no new
+# ONNX): reuse sa3-m/dit.onnx, build with BF16. Decoders/encoders are unchanged
+# by bf16 (it's a DiT-trunk fusion recipe), so decoder "bf16" reuses the
+# canonical decoder engine. Encoders are FP16-mixed only.
 DIT_ENGINE_FILENAME = {
+    "bf16":      "dit_bf16.trt",        # medium only (FMHA-fused; medium default)
     "fp16mixed": "dit_fp16mixed.trt",
     "fp32":      "dit_fp32.trt",
 }
+# DiT precisions actually built per model. bf16 is medium-only.
+_DIT_PRECISIONS = {
+    "sm-music": ("fp16mixed", "fp32"),
+    "sm-sfx":   ("fp16mixed", "fp32"),
+    "medium":   ("bf16", "fp16mixed", "fp32"),
+}
+# Per-DiT default precision (bf16 for medium, fp16mixed elsewhere).
+DIT_DEFAULT_PRECISION = {"sm-music": "fp16mixed", "sm-sfx": "fp16mixed", "medium": "bf16"}
 _DIT_SUBDIR = {"sm-music": "sa3-sm-music", "sm-sfx": "sa3-sm-sfx", "medium": "sa3-m"}
 DECODER_ENGINE_FILENAME = {
     "same-l": {
+        # bf16 is a DiT-only recipe → decoder reuses its canonical fp16-mixed engine.
+        "bf16":      "dec_dynamic_triton_swa.trt",
         "fp16mixed": "dec_dynamic_triton_swa.trt",
         "fp32":      "dec_dynamic_fp32.trt",
     },
     "same-s": {
+        "bf16":      "dec_dynamic_bf16.trt",
         "fp16mixed": "dec_dynamic_bf16.trt",
         "fp32":      "dec_dynamic_fp32.trt",
     },
 }
-PRECISIONS = ("fp16mixed", "fp32")
+PRECISIONS = ("bf16", "fp16mixed", "fp32")
+
+
+def default_precision(dit_name: str) -> str:
+    """Default DiT precision for a model: bf16 for medium (FMHA-fused speed
+    default), fp16-mixed otherwise."""
+    return DIT_DEFAULT_PRECISION.get(dit_name, "fp16mixed")
 
 
-def get_dit_engine_path(dit_name: str, precision: str = "fp16mixed") -> Path:
+def get_dit_engine_path(dit_name: str, precision: str = None) -> Path:
     if dit_name not in _DIT_SUBDIR:
         raise ValueError(f"unknown dit={dit_name!r}; valid: {list(_DIT_SUBDIR)}")
+    if precision is None:
+        precision = default_precision(dit_name)
     if precision not in DIT_ENGINE_FILENAME:
         raise ValueError(f"unknown precision={precision!r}; valid: {PRECISIONS}")
+    if precision == "bf16" and dit_name != "medium":
+        raise ValueError(
+            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)}")
     return ARCH_DIR / _DIT_SUBDIR[dit_name] / DIT_ENGINE_FILENAME[precision]
 
 
-def get_decoder_engine_path(decoder_name: str, precision: str = "fp16mixed") -> Path:
+def get_decoder_engine_path(decoder_name: str, precision: str = None) -> Path:
     if decoder_name not in DECODER_ENGINE_FILENAME:
         raise ValueError(f"unknown decoder={decoder_name!r}; valid: {list(DECODER_ENGINE_FILENAME)}")
+    if precision is None:
+        precision = "fp16mixed"   # decoders have no bf16-specific engine; canonical
     if precision not in DECODER_ENGINE_FILENAME[decoder_name]:
         raise ValueError(f"unknown precision={precision!r}; valid: {PRECISIONS}")
     return ARCH_DIR / decoder_name / DECODER_ENGINE_FILENAME[decoder_name][precision]
 
 
-def get_engine_files(dit_name: str, decoder_name: str, precision: str = "fp16mixed",
+def get_engine_files(dit_name: str, decoder_name: str, precision: str = None,
                        with_encoder: bool = False) -> list[str]:
     """Relative paths (under ARCH_DIR) needed for the chosen pipeline. Pass this
-    list to _ensure_files() to auto-download anything missing from HF."""
+    list to _ensure_files() to auto-download anything missing from HF.
+    precision=None resolves to the per-model default (bf16 for medium)."""
+    if precision is None:
+        precision = default_precision(dit_name)
     files = list(SHARED_FILES)
     files.append(f"{_DIT_SUBDIR[dit_name]}/{DIT_ENGINE_FILENAME[precision]}")
     files.append(f"{decoder_name}/{DECODER_ENGINE_FILENAME[decoder_name][precision]}")
@@ -903,14 +946,22 @@ def _arrow_pick(prompt: str, options: list[str], default: str | None = None) ->
 
 
 def prompt_user_if_missing(args):
-    """Fill in --dit / --decoder / --seed interactively if missing."""
+    """Fill in --dit / --decoder / --precision / --seed interactively if missing."""
     if args.dit is None:
         args.dit = _arrow_pick("Choose DiT model:", list(DIT_CHOICES.keys()), default="medium")
         print(f"  → {args.dit}")
     if args.decoder is None:
         suggested = DIT_CHOICES[args.dit]["default_decoder"]
         args.decoder = _arrow_pick("Choose audio decoder:", list(DECODER_PATHS.keys()), default=suggested)
         print(f"  → {args.decoder}")
+    # Resolve DiT precision default per model: bf16 for medium (FMHA-fused speed
+    # default), fp16-mixed for sm-music/sm-sfx. bf16 is medium-only (sm-music/
+    # sm-sfx use standard attention and already fuse in fp16mixed).
+    if getattr(args, "precision", None) is None:
+        args.precision = default_precision(args.dit)
+    if args.precision == "bf16" and args.dit != "medium":
+        sys.exit("error: --precision bf16 is only available for --dit medium "
+                 "(sm-music / sm-sfx already fuse in fp16mixed).")
     if args.seed is None:
         args.seed = random.randint(0, 2**31 - 1)
     return args
@@ -958,10 +1009,14 @@ def main():
     ap.add_argument("--decoder", choices=list(DECODER_PATHS.keys()), default=None,
                     help="Audio decoder. 'same-s' pairs with sm-* (110 MB engine). "
                          "'same-l' pairs with medium (1.2 GB engine). Interactive picker if omitted.")
-    ap.add_argument("--precision", choices=list(PRECISIONS), default="fp16mixed",
-                    help="Engine precision. 'fp16mixed' (default) = FP16 trunk + FP32 islands, "
-                         "fastest. 'fp32' = pure FP32, matches PyTorch eager bit-for-bit but ~2× "
-                         "slower and ~2× the VRAM. Engines auto-download from HF if missing.")
+    ap.add_argument("--precision", choices=list(PRECISIONS), default=None,
+                    help="DiT engine precision (default resolves per model: 'bf16' for medium, "
+                         "'fp16mixed' for sm-music/sm-sfx). 'bf16' (MEDIUM ONLY) = FMHA-fused, "
+                         "~1.8-4.7× faster than fp16mixed, within the perceptual floor, but not "
+                         "seed-reproducible vs fp16mixed (differential attention). 'fp16mixed' = "
+                         "FP16 trunk + FP32 islands (canonical, bit-reproducible). 'fp32' = pure "
+                         "FP32, matches PyTorch eager bit-for-bit but ~2× slower and ~2× the VRAM. "
+                         "Engines auto-download from HF if missing.")
     ap.add_argument("--models-dir", default=str(MODELS_DIR),
                     help=f"Directory containing the TRT engines. Default: {MODELS_DIR}")
     # ── Sampling ──
@@ -1165,8 +1220,12 @@ def _stage_vram(label): return 0
     _w_m  = torch.zeros(1, T5_MAX_LEN, device="cuda")
     _w_l  = torch.zeros(1, 257, T_lat, device="cuda")
     _w_lat   = torch.zeros(1, IO_CHANNELS, T_lat, device="cuda")
-    _w_audio = torch.zeros(1, 2, T_lat * SAMPLES_PER_LATENT, device="cuda") \
-                if "enc" in runners else None
+    # Encoder warmup is at the chunk_lat shape that encode_chunked actually
+    # uses; encoding at the full T_lat shape diverges past T_lat≈100-200 on
+    # both encoders (see encoder_encode docstring), so we never feed the
+    # engine that shape — encode_chunked stitches chunk_lat=50 windows.
+    _w_audio = torch.zeros(1, 2, DEFAULT_ENCODER_CHUNK_LAT * SAMPLES_PER_LATENT,
+                            device="cuda") if "enc" in runners else None
     # Optional: pre-allocate a pinned-memory destination buffer for the
     # Stage-5 narrow + DtoH path. With pinned dst + non_blocking=True the DMA
     # goes straight from GPU into RAM without the usual pageable→pinned
@@ -1275,7 +1334,7 @@ def _stage_vram(label): return 0
         audio_t = torch.from_numpy(audio_np).unsqueeze(0).cuda()   # (1, 2, T)
         sub(f"read+prep ({init_action})  {(time.time() - t0) * 1000:.0f} ms")
         t0 = time.time()
-        init_latents = encoder_encode(runners["enc"], audio_t)
+        init_latents = encode_chunked(runners["enc"], audio_t)
         sub(f"encode  {(time.time() - t0) * 1000:.0f} ms   latents {tuple(init_latents.shape)}")
         if args.free_models:
             runners["enc"].free(); del runners["enc"]