diff --git a/optimized/tensorRT/README.md b/optimized/tensorRT/README.md index d27c2da..8c74d3b 100644 --- a/optimized/tensorRT/README.md +++ b/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). diff --git a/optimized/tensorRT/build/build.py b/optimized/tensorRT/build/build.py index d1aa962..c5d8b11 100755 --- a/optimized/tensorRT/build/build.py +++ b/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)", diff --git a/optimized/tensorRT/build/build_from_onnx.py b/optimized/tensorRT/build/build_from_onnx.py index f75f70d..51e58f5 100755 --- a/optimized/tensorRT/build/build_from_onnx.py +++ b/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 diff --git a/optimized/tensorRT/scripts/sa3_trt.py b/optimized/tensorRT/scripts/sa3_trt.py index 199bc06..d710d19 100644 --- a/optimized/tensorRT/scripts/sa3_trt.py +++ b/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) diff --git a/optimized/tensorRT/scripts/sa3_trt_core.py b/optimized/tensorRT/scripts/sa3_trt_core.py index db92cd2..91f7011 100644 --- a/optimized/tensorRT/scripts/sa3_trt_core.py +++ b/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,7 +946,7 @@ 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}") @@ -911,6 +954,14 @@ def prompt_user_if_missing(args): 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"]