SoundGame/tools/stable-audio-3/optimized/tensorRT/build/build_dit_fp16.py

#!/usr/bin/env python3
"""Build a SA3 DiT TRT engine in FP16-mixed precision: FP16 trunk with FP32
islands around RMSNorm and (by default) the rotary position embedding
*generation* that PyTorch marks @autocast(enabled=False). The attention core
(QK^T -> Softmax -> P.V) runs FP16 so TRT's FMHA fuser can fire — step 6 below.

Background:
- The BF16 engine drifts at long sequence, and the cause is narrower than
  "bf16 accumulates error over the 8 sampling steps": a per-op ablation showed
  bf16 arithmetic everywhere else in the DiT is clean (GEMMs, RMSNorm, the
  differential-attention subtraction, even the cos/sin tables). The single
  culprit is the RoPE rotation ANGLE, t * inv_freq, which reaches ~4155 rad at
  L=4092 — where bf16's spacing is 32 rad, five whole rotations — so cos/sin of
  it is uncorrelated with the true value and position information for the
  fast-rotating dims is destroyed. The axis that separates a good engine from a
  bad one is therefore "fp32 angle island present", not fp16-vs-bf16.
- Naive FP16 (just BuilderFlag.FP16) catastrophically diverges: TRT promotes
  RMSNorm's Pow/ReduceMean/Sqrt to FP16, which overflows (variance > 65504).
  The Softmax is a different story — see bound_attention_core(): FP16 softmax
  is fine, and an FP32 one costs 4.3x at L=4096 for nothing.
- MLX gets away with FP16 because it casts those subgraphs to FP32 internally.

This script reproduces that recipe on the ONNX side:

  1. Identify FP32 "islands":
       - Every RMSNorm chain (Pow -> ReduceMean -> Add -> Sqrt -> Div -> Mul -> Mul)
       - Every Softmax (transient — step 6 hands the attention core back to FP16)
       - (mode=rope, default) Every region reachable from a Cast(to=FP32) that
         leads to Cos/Sin or Einsum — i.e. the rotary position embedding
         (@autocast(enabled=False) in PyTorch source)
  2. Strip the 569 no-op Cast(to=FP32) nodes that the PyTorch ONNX export
     left as semantic markers — we already have the islands explicitly.
  3. Wrap each island input/output with Cast(to=FP32)/Cast(to=FP16). These
     are no-ops in the still-FP32 graph; they become real boundaries after
     step 4.
  4. Manually convert non-island initializers/Constant/value_info to FP16.
     RMSNorm gamma weights and other initializers feeding islands stay FP32.
  5. Post-process: walk the graph inserting autocast nodes wherever two
     operands of a Mul/MatMul/Add/etc disagree on dtype (this handles the
     boundary between FP32 t5_hidden input and FP16 trunk).
  6. bound_attention_core(): terminate the RoPE island before QK^T so the
     attention core (QK^T -> Softmax -> P.V) runs FP16 and TRT's FMHA fuser can
     fire. Mirrors `apply_rotary_pos_emb`'s closing `t.to(out_dtype)`, which the
     island passes above do not reproduce. WITHOUT this step the whole O(L^2)
     attention core stays FP32: 0 fused MHA nodes and 4.3x slower at L=4096, for
     no accuracy gain (teacher-forced cos vs FP32 is 1.0000 WITH the pass and
     0.9998 without). See the function docstring. Disable with --no-bound-attn
     only to reproduce the pre-2026-07 engines.
  7. Build a TRT engine with NetworkDefinitionCreationFlag.STRONGLY_TYPED
     so TRT respects the per-tensor dtypes exactly (no auto-promotion).
     STRONGLY_TYPED is mandatory, not stylistic: weakly-typed + BuilderFlag.FP16
     lets TRT re-cast the FP32 RMSNorm islands and reintroduces the overflow.

Inputs/outputs stay FP32 so the runtime can swap engines transparently.

Validated results (sa3-sm-music, L=1292):
  - Single-step cos vs PT FP32: 0.99997
  - 8-step final-latent cos vs PT FP32: 0.998
  - Decoded audio RMS-curve correlation vs PT FP32: 0.999
  - Inference time per 8-step: ~44 ms (BF16: 42 ms, FP32: 116 ms)
  - Engine size: 926 MB (BF16: 935 MB, FP32: 1842 MB)

Usage:
    python build_dit_fp16.py
        [--mode {minimal,rope,full}]      # default: rope
        [--input  /tmp/dit_sm-music_fixed_v2.onnx]
        [--onnx   /tmp/dit_sm-music_fp16.onnx]
        [--engine .../models/sm_90/sa3-sm-music/dit_fp16.trt]
"""
import argparse
import os
import sys
import time
from pathlib import Path

SCRIPTS_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPTS_DIR))


T5_TOKENS = 256
T5_HIDDEN_DIM = 768

# Optimization profile — same as sa3-sm-music canonical build.
# Min=1 lets sub-trained short forms run; opt=1292 is the canonical tactic point.
_DIT_PROFILE = {
    "x":              [(1, 256, 1),     (1, 256, 1292),   (1, 256, 4096)],
    "t":              [(1,),            (1,),             (1,)],
    "t5_hidden":      [(1, T5_TOKENS, T5_HIDDEN_DIM)] * 3,
    "t5_mask":        [(1, T5_TOKENS)] * 3,
    "seconds_total":  [(1,)] * 3,
    "local_add_cond": [(1, 257, 1),     (1, 257, 1292),   (1, 257, 4096)],
}


def find_fp32_islands(model, mode="minimal"):
    """Identify nodes that must stay in FP32.

    mode="minimal": only RMSNorm chains + Softmax. Smallest island set —
        produces fastest engine but may not preserve all the precision-
        sensitive paths PyTorch's autocast(enabled=False) was protecting.
    mode="rope": minimal + the rotary embedding generation + apply_rotary
        islands (matches PyTorch's @autocast(enabled=False) on RoPE).
    mode="full": every node reachable from a Cast(to=FP32) seed in the
        original FP32 ONNX. Largest island set — preserves every PyTorch
        .float() intent but most of the network ends up FP32.

    All modes additionally block:
      - RMSNorm chains (Pow -> ReduceMean -> Add -> Sqrt -> Div -> Mul -> Mul)
      - Softmax
    """
    out2node = {}
    in2nodes = {}
    for n in model.graph.node:
        for o in n.output:
            out2node[o] = n
        for inp in n.input:
            in2nodes.setdefault(inp, []).append(n)

    blocked = set()

    # Walk forward from a node, stopping at any Cast op. Add every node
    # we visit to `blocked`. Used for tracing FP32 islands from Cast(to=FP32)
    # seeds.
    def bfs_forward_until_cast(seed):
        queue = list(seed.output)
        seen = set([seed.name])
        while queue:
            out_name = queue.pop(0)
            for consumer in in2nodes.get(out_name, []):
                if consumer.name in seen:
                    continue
                seen.add(consumer.name)
                if consumer.op_type == "Cast":
                    continue  # boundary
                blocked.add(consumer.name)
                queue.extend(consumer.output)

    # (A) Original Cast(to=FP32) seeds — varies by mode.
    fp32_cast_seeds = []
    for c in model.graph.node:
        if c.op_type != "Cast":
            continue
        to_val = None
        for a in c.attribute:
            if a.name == "to":
                to_val = a.i
        if to_val != 1:  # FP32
            continue
        fp32_cast_seeds.append(c)

    # Heuristic: classify Cast(to=FP32) seeds by what they're protecting.
    # RoPE-island casts are those near Range, Einsum, Cos, Sin.
    # RMSNorm-pre casts are those right before a Pow (and the chain head).
    # apply_rotary casts are those right before Mul/Add inside attention.
    def is_rope_island_seed(seed):
        # The 2 rope seeds (one for q, one for k) lead toward Cos/Sin via
        # Range -> Cast -> Div -> Einsum -> Cos/Sin.
        # Also each apply_rotary call casts q, k to FP32 before the rotation
        # multiply/add. Those are 2 per layer.
        # Cheap heuristic: walk downstream up to 3 steps and check if we
        # encounter Cos/Sin or hit the Q/K projection chain.
        queue = [(seed, 0)]
        seen = set([seed.name])
        while queue:
            n, d = queue.pop(0)
            if d > 6:
                continue
            for o in n.output:
                for c in in2nodes.get(o, []):
                    if c.name in seen:
                        continue
                    seen.add(c.name)
                    if c.op_type in ("Cos", "Sin", "Einsum"):
                        return True
                    queue.append((c, d + 1))
        return False

    rope_seeds = []
    other_seeds = []
    if mode in ("rope", "full"):
        for s in fp32_cast_seeds:
            if mode == "full" or is_rope_island_seed(s):
                blocked.add(s.name)
                bfs_forward_until_cast(s)
                rope_seeds.append(s)
            else:
                other_seeds.append(s)

    # (B) Explicit RMSNorm chains
    RMS_STEPS = ["ReduceMean", "Add", "Sqrt", "Div", "Mul", "Mul"]
    rmsnorm_chains = 0
    for n in model.graph.node:
        if n.op_type != "Pow":
            continue
        chain = [n]
        cur = n
        ok = True
        for step in RMS_STEPS:
            ops = in2nodes.get(cur.output[0], [])
            match = [c for c in ops if c.op_type == step]
            if not match:
                ok = False
                break
            cur = match[0]
            chain.append(cur)
        if not ok:
            continue
        rmsnorm_chains += 1
        for c in chain:
            blocked.add(c.name)

    # (C) Explicit Softmax
    softmax_count = 0
    for n in model.graph.node:
        if n.op_type == "Softmax":
            blocked.add(n.name)
            softmax_count += 1

    print(f"  FP32 islands (mode={mode}): "
          f"{len(rope_seeds)} RoPE/island Cast seeds (of {len(fp32_cast_seeds)}), "
          f"{rmsnorm_chains} RMSNorm chains, "
          f"{softmax_count} Softmax")
    print(f"  total blocked nodes: {len(blocked)}")

    # Op-type breakdown for sanity
    from collections import Counter
    name2node = {n.name: n for n in model.graph.node}
    op_counts = Counter(name2node[name].op_type for name in blocked
                        if name in name2node)
    print(f"  blocked op-type top10:")
    for k, v in op_counts.most_common(10):
        print(f"    {k}: {v}")

    return blocked


def strip_noop_fp32_casts(model):
    """Remove Cast(to=FP32) nodes whose input is already FP32. The PyTorch
    ONNX export inserts these to mirror `x.float()` calls in FP32 modules
    (where they're no-ops). Removing them before the FP16 conversion lets
    the converter correctly insert FP32->FP16 transitions; without removal,
    these stale casts produce dtype contradictions (TRT: "A is Float, B is
    Half"). The leading Cast(to=FP32) for RMSNorm pre_norm chains is one
    such example."""
    import onnx as _onnx
    from onnx import shape_inference, TensorProto
    # In-memory shape_inference fails on models >2GB (protobuf serialization
    # limit) — happens for sa3-m. The fallback (file-based inference + tempdir
    # external data) is fragile because inferred initializers retain refs to
    # the tempdir even after explicit clearing. Simpler escape: skip the
    # no-op-Cast-stripping pass entirely. It's an optimization (slightly
    # smaller graph, easier TRT inspection), not required for correctness —
    # the downstream conversion handles redundant Casts fine.
    try:
        inferred = shape_inference.infer_shapes(model)
    except Exception as e:
        if "Failed to serialize" in str(e) or "EncodeError" in type(e).__name__:
            print(f"  shape_inference failed (model >2GB, protobuf limit) — "
                  f"skipping no-op-Cast stripping. Engine will still build correctly.")
            return model
        raise
    vi_dtype = {}
    for vi in list(inferred.graph.value_info) + list(inferred.graph.input):
        vi_dtype[vi.name] = vi.type.tensor_type.elem_type
    for init in inferred.graph.initializer:
        vi_dtype[init.name] = init.data_type

    # Collect no-op Casts to remove
    to_remove = []
    rename_map = {}  # output -> input (rewire downstream)
    for n in model.graph.node:
        if n.op_type != "Cast":
            continue
        to_attr = None
        for a in n.attribute:
            if a.name == "to":
                to_attr = a.i
        if to_attr != TensorProto.FLOAT:  # only FP32 casts
            continue
        in_dtype = vi_dtype.get(n.input[0])
        if in_dtype == TensorProto.FLOAT:
            to_remove.append(n)
            rename_map[n.output[0]] = n.input[0]
    print(f"  stripping {len(to_remove)} no-op Cast(to=FP32) nodes")

    # Transitively close rename_map so chains like Cast_6 -> Cast_8 -> consumer
    # rewire all the way back to the pre-Cast_6 source. Otherwise we end up
    # with consumers referring to the intermediate Cast_6_output_0 which has
    # also been removed.
    def resolve(name):
        while name in rename_map:
            name = rename_map[name]
        return name
    rename_map = {k: resolve(v) for k, v in rename_map.items()}

    # Rewire downstream consumers
    for n in model.graph.node:
        for i, inp in enumerate(n.input):
            if inp in rename_map:
                n.input[i] = rename_map[inp]
    # Also fix graph outputs (rare)
    for o in model.graph.output:
        if o.name in rename_map:
            o.name = rename_map[o.name]

    # Remove the dead Cast nodes
    for n in to_remove:
        model.graph.node.remove(n)

    return model


def wrap_islands_with_casts(model, blocked_names):
    """Surround every FP32-island (Pow/ReduceMean/.../Mul) and Softmax node
    with explicit Cast(FP16->FP32) on inputs and Cast(FP32->FP16) on outputs
    BEFORE the FP16 conversion runs. This ensures the converter sees clear
    dtype boundaries — every input to a blocked node is preceded by a Cast,
    every output is followed by a Cast.

    Why pre-insert: the converter's auto-Cast-insertion has fragile
    interactions with shape inference; pre-inserting our own Casts (whose
    `to` attribute we'll later rewrite via convert_float_to_float16) gives
    deterministic boundaries.

    Strategy:
      - For each blocked-island INPUT edge (input of a blocked node whose
        source is NOT another blocked node), insert
        Cast(in -> in__fp32) and rewire the blocked node to consume in__fp32.
        Initially in__fp32 is Cast(to=FP32) — a no-op in the FP32 graph;
        after convert_float_to_float16, surrounding becomes FP16 and this
        Cast performs the actual upcast.
      - For each blocked-island OUTPUT edge (a non-blocked consumer of a
        blocked node), insert Cast(out -> out__fp16) and rewire that consumer.
    """
    from onnx import TensorProto, helper

    in2nodes = {}
    for n in model.graph.node:
        for inp in n.input:
            in2nodes.setdefault(inp, []).append(n)
    out2node = {}
    for n in model.graph.node:
        for o in n.output:
            out2node[o] = n

    new_casts = []
    n_input_wraps = 0
    n_output_wraps = 0

    # We need to know each tensor's dtype to skip non-float inputs (INT64
    # axes/indices/etc). Build a dtype map from initializers + Constant
    # outputs + graph inputs, then do forward inference for the rest.
    tensor_dtype = {}
    for init in model.graph.initializer:
        tensor_dtype[init.name] = init.data_type
    for gi in model.graph.input:
        tensor_dtype[gi.name] = gi.type.tensor_type.elem_type
    for node in model.graph.node:
        if node.op_type == "Constant":
            for a in node.attribute:
                if a.name == "value":
                    for o in node.output:
                        tensor_dtype[o] = a.t.data_type
        elif node.op_type == "Cast":
            for a in node.attribute:
                if a.name == "to":
                    for o in node.output:
                        tensor_dtype[o] = a.i
        elif node.op_type == "Shape":
            for o in node.output:
                tensor_dtype[o] = TensorProto.INT64
        elif node.op_type == "ConstantOfShape":
            dt = TensorProto.FLOAT
            for a in node.attribute:
                if a.name == "value":
                    dt = a.t.data_type
            for o in node.output:
                tensor_dtype[o] = dt
        elif node.op_type == "Where":
            # Where(condition, x, y) — output dtype matches x/y (inputs[1])
            if len(node.input) >= 2:
                dt = tensor_dtype.get(node.input[1])
                if dt is not None:
                    for o in node.output:
                        tensor_dtype.setdefault(o, dt)
        elif node.op_type in ("Equal", "Greater", "Less", "And", "Or", "Not",
                              "GreaterOrEqual", "LessOrEqual"):
            # Boolean output
            for o in node.output:
                tensor_dtype.setdefault(o, TensorProto.BOOL)
        else:
            # Default: scan inputs for the first FP-typed input dtype, fall
            # back to first input dtype. This avoids leaking BOOL from
            # boolean conditions etc.
            picked = None
            for inp in node.input:
                if not inp:
                    continue
                dt = tensor_dtype.get(inp)
                if dt in (TensorProto.FLOAT, TensorProto.FLOAT16,
                          TensorProto.BFLOAT16, TensorProto.DOUBLE):
                    picked = dt
                    break
            if picked is None and node.input:
                picked = tensor_dtype.get(node.input[0])
            if picked is not None:
                for o in node.output:
                    tensor_dtype.setdefault(o, picked)

    # === Inputs ===
    # For each blocked node, every FLOAT input not coming from another
    # blocked node needs a Cast(to=FP32) preface. We skip non-float inputs
    # (INT64 axes/indices/etc) and initializer inputs (those stay FP32
    # directly when feeding an island).
    init_names = {init.name for init in model.graph.initializer}
    input_cast_map = {}
    for node in model.graph.node:
        if node.name not in blocked_names:
            continue
        for i, inp in enumerate(node.input):
            if not inp:
                continue
            src = out2node.get(inp)
            if src is not None and src.name in blocked_names:
                continue  # interior edge
            if inp in init_names:
                continue  # initializer — let it stay FP32, no Cast needed
            # Skip non-float inputs (INT64 axes, etc.)
            dt = tensor_dtype.get(inp)
            if dt is not None and dt not in (TensorProto.FLOAT,
                                              TensorProto.FLOAT16,
                                              TensorProto.BFLOAT16):
                continue
            if inp in input_cast_map:
                node.input[i] = input_cast_map[inp]
                continue
            cast_out = f"{inp}__island_in_fp32"
            cast_name = f"{inp}__island_in_cast"
            new_casts.append(helper.make_node(
                "Cast", inputs=[inp], outputs=[cast_out],
                name=cast_name, to=TensorProto.FLOAT))
            input_cast_map[inp] = cast_out
            node.input[i] = cast_out
            n_input_wraps += 1

    # === Outputs ===
    # For each blocked node, every consumer not in blocked needs a
    # Cast(to=FP16) wrapper, with the blocked node's output rewired
    # through the new cast for that consumer set.
    # Dedupe: one cast per output that has FP16-consumers.
    output_cast_map = {}  # original output name -> new fp16 cast output name
    for node in model.graph.node:
        if node.name not in blocked_names:
            continue
        for out_name in node.output:
            consumers = in2nodes.get(out_name, [])
            non_blocked = [c for c in consumers if c.name not in blocked_names]
            if not non_blocked:
                continue
            cast_out = f"{out_name}__island_out_fp16"
            cast_name = f"{out_name}__island_out_cast"
            new_casts.append(helper.make_node(
                "Cast", inputs=[out_name], outputs=[cast_out],
                name=cast_name, to=TensorProto.FLOAT16))
            output_cast_map[out_name] = cast_out
            for c in non_blocked:
                for i, inp in enumerate(c.input):
                    if inp == out_name:
                        c.input[i] = cast_out
            n_output_wraps += 1

    # Also rewire graph outputs that reference blocked node outputs.
    for o in model.graph.output:
        if o.name in output_cast_map:
            # We can't really rename graph outputs, but the cast was added,
            # so we leave it. The downstream consumer chain has been rewired.
            pass

    model.graph.node.extend(new_casts)
    print(f"  wrapped {len(blocked_names)} blocked nodes: "
          f"{n_input_wraps} input casts, {n_output_wraps} output casts")
    return model


def bound_attention_core(model):
    """Terminate the RoPE island before QK^T so the attention core runs FP16.

    Why this pass exists
    --------------------
    find_fp32_islands() blocks every Softmax (section C) *and* the RoPE region.
    Because RoPE emits q/k in FP32 and nothing casts them back,
    fix_dtype_mismatches then keeps the QK^T MatMul in FP32 to match its operands.
    Net effect: the entire O(L^2) attention core — QK^T, the scale Mul, and Softmax
    — runs FP32, and TensorRT's FMHA fuser cannot fire on it. Measured on the
    medium DiT: **0 fused MHA nodes and 4.3x slower at L=4096** than it needs to be.

    Neither FP32 region is actually required:

      * PyTorch does not keep the attention core in FP32. `apply_rotary_pos_emb`
        ends with `t, t_unrotated = t.to(out_dtype), t_unrotated.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 the explicit
        FP32 softmax buys nothing the kernel does not already do.

    So for each attention this rewrites

        Mul(fp32) -> MatMul(fp32,fp32) -> Cast(fp32)[no-op] -> Softmax(fp32) -> Cast(fp16) -> MatMul(P.V)
    to
        Mul(fp32) -> Cast(fp16) x2 -> MatMul(fp16) ->            Softmax(fp16) ->            MatMul(P.V)

    mirroring the source's cast-back. The RMSNorm islands are untouched: those
    guard variance *overflow* (FP16's 5-bit exponent), a different failure mode
    from precision, and they are still required.

    Verified on sa3-m: 96/96 attentions converted, 96 fused MHA nodes at FP16,
    teacher-forced velocity cos 1.0000 vs the FP32 engine at L=256/1292/4092
    (the unbounded graph matches it at 256/1292 but drops to 0.9998 at 4092), and
    4.3x faster at L=4096.

    Scope: only the medium DiT has the unbounded island. In sa3-sm-music /
    sa3-sm-sfx the RoPE island already terminates before QK^T (verified: their
    QK^T operands are FP16, so their published engines fuse 40/40 MHA), and what
    the pass converts there is the Softmax island alone — a real change to those
    graphs, but not the fix this exists for. Their published ONNX/engines predate
    the pass; --no-bound-attn skips it and reproduces the graph they were built
    from.

    NOTE the resulting graph must be built STRONGLY_TYPED. Building it weakly-typed
    with BuilderFlag.FP16 lets TRT re-cast the FP32 RMSNorm islands too, which is
    the naive-FP16 overflow this recipe exists to prevent (measured: teacher-forced
    cos collapses to 0.88).
    """
    from collections import defaultdict
    from onnx import TensorProto, helper

    g = model.graph
    producer = {}
    consumers = defaultdict(list)
    for n in g.node:
        for o in n.output:
            producer[o] = n
        for i in n.input:
            if i:
                consumers[i].append(n)
    graph_outputs = {o.name for o in g.output}

    def cast_to(node):
        for a in node.attribute:
            if a.name == "to":
                return a.i
        return None

    drop, rename, new_casts = set(), {}, []
    converted, converted_softmax = 0, []
    skipped = defaultdict(int)

    for sm in [n for n in g.node if n.op_type == "Softmax"]:
        pre = producer.get(sm.input[0])
        if pre is None or pre.op_type != "Cast" or cast_to(pre) != TensorProto.FLOAT:
            skipped["no_fp32_cast_before_softmax"] += 1
            continue
        if len(consumers[sm.input[0]]) != 1:
            skipped["pre_cast_output_shared"] += 1
            continue
        qk = producer.get(pre.input[0])
        if qk is None or qk.op_type != "MatMul":
            skipped["upstream_not_matmul"] += 1
            continue
        posts = consumers[sm.output[0]]
        if not posts or any(p.op_type != "Cast" or cast_to(p) != TensorProto.FLOAT16
                            for p in posts):
            skipped["no_fp16_cast_after_softmax"] += 1
            continue
        if sm.output[0] in graph_outputs or any(p.output[0] in graph_outputs for p in posts):
            skipped["touches_graph_output"] += 1
            continue
        if len(consumers[qk.output[0]]) != 1:
            skipped["qk_output_shared"] += 1
            continue

        for idx, src in enumerate(list(qk.input)):
            out_name = f"{src}__attnfp16_{converted}_{idx}"
            new_casts.append(helper.make_node(
                "Cast", [src], [out_name], to=TensorProto.FLOAT16,
                name=f"{src}__attnfp16_cast_{converted}_{idx}"))
            qk.input[idx] = out_name
        sm.input[0] = qk.output[0]
        drop.add(id(pre))
        for p in posts:
            rename[p.output[0]] = sm.output[0]
            drop.add(id(p))
        converted += 1
        converted_softmax.append(sm)

    print(f"  bound_attention_core: {converted} attention(s) -> FP16 core")
    if skipped:
        print(f"    skipped: {dict(skipped)}")
    if converted == 0:
        raise RuntimeError("bound_attention_core converted nothing — the island "
                           "passes upstream must have changed shape")

    def resolve(name):
        seen = set()
        while name in rename and name not in seen:
            seen.add(name)
            name = rename[name]
        return name

    kept = [n for n in g.node if id(n) not in drop]
    for n in kept:
        for i, t in enumerate(n.input):
            r = resolve(t)
            if r != t:
                n.input[i] = r

    # Each new Cast must precede its consumer to keep the graph topologically sorted.
    cast_by_output = {c.output[0]: c for c in new_casts}
    ordered = []
    for n in kept:
        for t in n.input:
            c = cast_by_output.pop(t, None)
            if c is not None:
                ordered.append(c)
        ordered.append(n)
    if cast_by_output:
        raise RuntimeError(f"{len(cast_by_output)} inserted casts could not be placed")
    del g.node[:]
    g.node.extend(ordered)

    # Retype/drop stale value_info: QK^T and Softmax outputs are FP16 now, and the
    # spliced Cast outputs no longer exist. A stale FP32 declaration on an FP16
    # tensor makes STRONGLY_TYPED reject the graph. Only the attentions we actually
    # converted move — a skipped Softmax keeps its FP32 declaration.
    retype = set()
    for sm in converted_softmax:
        retype.add(sm.output[0])
        retype.add(sm.input[0])     # rewired above to the QK^T MatMul output
    dead = set(rename)
    kept_vi = []
    for vi in g.value_info:
        if vi.name in dead:
            continue
        if vi.name in retype and vi.type.tensor_type.elem_type == TensorProto.FLOAT:
            vi.type.tensor_type.elem_type = TensorProto.FLOAT16
        kept_vi.append(vi)
    del g.value_info[:]
    g.value_info.extend(kept_vi)
    return model


def fix_dtype_mismatches(model, blocked_names, low_dt=None):
    """Insert Cast nodes wherever two operands of an op disagree on float dtype.

    Forward dtype propagation from graph inputs and initializers, computing each
    tensor's dtype as it would appear at runtime. For binary/multi-input ops
    (MatMul, Add, Mul, Div, Sub, Pow, Where, ...) all float operands are aligned
    onto one dtype: FP32 when the consuming node is in `blocked_names` (an FP32
    island — never downcast those), otherwise the trunk dtype, and the mismatched
    operands get a Cast.

    `low_dt` is that trunk dtype: FLOAT16 for the fp16-mixed recipe (the default
    when None) or BFLOAT16 for a bf16 trunk. Both recipes need the same graph-wide
    reconciliation — a bf16 trunk with fp32 islands hits exactly the same class of
    mismatch (`ElementWiseOperation DIV must have same input types`,
    `IMatrixMultiplyLayer must have same input types`) — so this pass is shared
    rather than duplicated.
    """
    from onnx import TensorProto, helper

    if low_dt is None:
        low_dt = TensorProto.FLOAT16    # default trunk dtype: the fp16-mixed recipe

    init_dtype = {init.name: init.data_type for init in model.graph.initializer}
    graph_input_dtype = {gi.name: gi.type.tensor_type.elem_type for gi in model.graph.input}

    def cast_to(node):
        for a in node.attribute:
            if a.name == "to":
                return a.i
        return None

    # Forward dtype propagation. Returns dtype for each tensor.
    # Topological order: process nodes in graph order (ONNX guarantees topo).
    tensor_dtype = {}
    tensor_dtype.update(init_dtype)
    tensor_dtype.update(graph_input_dtype)
    # Constant nodes — their outputs are determined by their .t value
    for n in model.graph.node:
        if n.op_type == "Constant":
            for a in n.attribute:
                if a.name == "value":
                    for o in n.output:
                        tensor_dtype[o] = a.t.data_type

    def get_dt(name):
        return tensor_dtype.get(name, None)

    def set_dt(name, dt):
        tensor_dtype[name] = dt

    # Walk nodes in order, infer their output dtype.
    new_casts = []
    n_inserted = 0

    # Ops that require all float inputs to share dtype
    SHARED_FLOAT_DT_OPS = {
        "MatMul", "Gemm", "Add", "Mul", "Div", "Sub", "Pow", "Min", "Max",
        "Where", "Equal", "Greater", "Less", "And", "Or",
        "Conv", "ConvTranspose",
    }
    # Ops where the output dtype = first float input's dtype
    PRESERVE_OPS = {
        "Reshape", "Transpose", "Unsqueeze", "Squeeze", "Slice", "Concat",
        "Gather", "Identity", "Tile", "Expand", "Pad", "Split",
        "ReduceMean", "ReduceSum", "ReduceMax", "ReduceMin", "ReduceProd",
        "Softmax", "Sigmoid", "Tanh", "Relu", "Gelu", "Sqrt", "Sin", "Cos",
        "Neg", "Abs", "Exp", "Log", "Conv", "ConvTranspose",
    }

    nodes_list = list(model.graph.node)
    for node in nodes_list:
        if node.op_type == "Cast":
            for o in node.output:
                set_dt(o, cast_to(node))
            continue
        if node.op_type == "Constant":
            continue
        if node.op_type == "Shape":
            # output is int64
            for o in node.output:
                set_dt(o, TensorProto.INT64)
            continue
        if node.op_type == "ConstantOfShape":
            # output dtype = attr value dtype (default FP32)
            dt = TensorProto.FLOAT
            for a in node.attribute:
                if a.name == "value":
                    dt = a.t.data_type
            for o in node.output:
                set_dt(o, dt)
            continue
        if node.op_type == "Range":
            # output is whatever the start dtype is
            dt = get_dt(node.input[0]) if node.input else TensorProto.INT64
            for o in node.output:
                set_dt(o, dt)
            continue
        if node.op_type in ("Equal", "Greater", "Less", "And", "Or", "Not",
                            "GreaterOrEqual", "LessOrEqual"):
            for o in node.output:
                set_dt(o, TensorProto.BOOL)
            continue

        # For shared-dtype ops, find the dominant float dtype and cast
        # any mismatched inputs.
        float_input_dts = []
        for inp in node.input:
            if not inp:
                continue
            dt = get_dt(inp)
            if dt in (TensorProto.FLOAT, low_dt):
                float_input_dts.append((inp, dt))
        if node.op_type in SHARED_FLOAT_DT_OPS and len(float_input_dts) > 1:
            unique_dts = set(dt for _, dt in float_input_dts)
            if len(unique_dts) > 1:
                # Mismatch — need to align. Rule: if blocked node, force FP32.
                # Otherwise prefer low_dt (the trunk dtype, FP16 by default).
                if node.name in blocked_names:
                    target = TensorProto.FLOAT
                else:
                    target = low_dt
                for inp, dt in float_input_dts:
                    if dt == target:
                        continue
                    # Insert a Cast(to=target) before this input
                    cast_out = f"{inp}__autocast_{n_inserted}"
                    cast_node = helper.make_node(
                        "Cast", inputs=[inp], outputs=[cast_out],
                        name=f"{inp}__autocast_node_{n_inserted}",
                        to=target)
                    new_casts.append(cast_node)
                    set_dt(cast_out, target)
                    # Rewire this node's input
                    for i, ninp in enumerate(node.input):
                        if ninp == inp:
                            node.input[i] = cast_out
                    n_inserted += 1

        # Compute output dtype
        if node.op_type in SHARED_FLOAT_DT_OPS or node.op_type in PRESERVE_OPS:
            # Output dtype: first float input dtype (after possible casts)
            for inp in node.input:
                if not inp:
                    continue
                dt = get_dt(inp)
                if dt in (TensorProto.FLOAT, low_dt):
                    for o in node.output:
                        set_dt(o, dt)
                    break
        else:
            # Default: preserve first input dtype
            if node.input:
                dt = get_dt(node.input[0])
                if dt is not None:
                    for o in node.output:
                        set_dt(o, dt)

    model.graph.node.extend(new_casts)
    print(f"  fix_dtype_mismatches: inserted {n_inserted} autocast nodes")
    return model


def manual_convert_to_fp16(model, blocked_names):
    """Convert FP32 ONNX to FP16 trunk + FP32 islands.

    Algorithm:
    1. Compute the set of "FP32 tensors" — outputs of blocked nodes, plus
       outputs of our Cast(to=FP32) wrappers.
    2. For every initializer NOT consumed exclusively by blocked nodes,
       convert from FP32 to FP16. (If exclusively in islands, keep FP32.)
    3. For every value_info NOT in the FP32-tensor set, set dtype FP16.
    4. Convert Constant nodes (op_type=Constant) outside blocked: change
       value tensor to FP16.

    Inputs of the graph stay FP32 (so the runtime doesn't need to change).
    Output of the graph: we'll Cast to FP32 at the very end so the runtime
    sees FP32.
    """
    from onnx import TensorProto, numpy_helper, helper
    import numpy as np

    in2nodes = {}
    out2node = {}
    for n in model.graph.node:
        for inp in n.input:
            in2nodes.setdefault(inp, []).append(n)
        for o in n.output:
            out2node[o] = n

    # Tensors that must stay FP32: every output of a blocked node + every
    # output of our wrap_islands Cast(to=FP32) nodes.
    fp32_tensors = set()
    for node in model.graph.node:
        if node.name in blocked_names:
            for o in node.output:
                fp32_tensors.add(o)
        # Cast(to=FP32) outputs — these are the island input wrappers
        if node.op_type == "Cast":
            for a in node.attribute:
                if a.name == "to" and a.i == TensorProto.FLOAT:
                    for o in node.output:
                        fp32_tensors.add(o)

    print(f"  FP32 tensors: {len(fp32_tensors)}")

    # === Convert initializers ===
    # An initializer "feeds an FP32 island" if all of its consumers are
    # either blocked themselves, or are Cast(to=FP32) wrappers whose
    # eventual consumers are all blocked. We walk through Cast(to=FP32)
    # nodes transparently.
    def initializer_feeds_island_only(name):
        consumers = in2nodes.get(name, [])
        if not consumers:
            return False
        for c in consumers:
            if c.name in blocked_names:
                continue
            # Allow walking through Cast(to=FP32) wrapper to see what
            # ultimately consumes it.
            if c.op_type == "Cast":
                to_val = None
                for a in c.attribute:
                    if a.name == "to":
                        to_val = a.i
                if to_val == TensorProto.FLOAT:
                    # All downstream of this Cast must also be blocked
                    cast_outputs = c.output
                    deep_consumers = []
                    for o in cast_outputs:
                        deep_consumers.extend(in2nodes.get(o, []))
                    if all(dc.name in blocked_names for dc in deep_consumers) and deep_consumers:
                        continue
            return False
        return True

    init_consumed_by_blocked = {init.name: initializer_feeds_island_only(init.name)
                                 for init in model.graph.initializer}

    n_init_fp16 = 0
    n_init_fp32_keep = 0
    for init in model.graph.initializer:
        if init.data_type != TensorProto.FLOAT:
            continue
        if init_consumed_by_blocked.get(init.name, False):
            n_init_fp32_keep += 1
            continue
        # Convert FP32 -> FP16
        arr = numpy_helper.to_array(init).astype(np.float16)
        new_init = numpy_helper.from_array(arr, name=init.name)
        init.CopyFrom(new_init)
        n_init_fp16 += 1
    print(f"  initializers: {n_init_fp16} converted to FP16, "
          f"{n_init_fp32_keep} kept FP32")

    # === Convert Constant nodes (these are inline tensor literals) ===
    def constant_feeds_island_only(consumers):
        """Like initializer_feeds_island_only but for Constant outputs."""
        if not consumers:
            return False
        for c in consumers:
            if c.name in blocked_names:
                continue
            if c.op_type == "Cast":
                to_val = None
                for a in c.attribute:
                    if a.name == "to":
                        to_val = a.i
                if to_val == TensorProto.FLOAT:
                    deep_consumers = []
                    for o in c.output:
                        deep_consumers.extend(in2nodes.get(o, []))
                    if all(dc.name in blocked_names for dc in deep_consumers) and deep_consumers:
                        continue
            return False
        return True

    n_const_fp16 = 0
    n_const_fp32_keep = 0
    for node in model.graph.node:
        if node.op_type != "Constant":
            continue
        # Find the value attribute
        for a in node.attribute:
            if a.name == "value" and a.t.data_type == TensorProto.FLOAT:
                # Where is this Constant consumed?
                consumers = []
                for o in node.output:
                    consumers.extend(in2nodes.get(o, []))
                if constant_feeds_island_only(consumers):
                    n_const_fp32_keep += 1
                    continue
                arr = numpy_helper.to_array(a.t).astype(np.float16)
                new_t = numpy_helper.from_array(arr)
                a.t.CopyFrom(new_t)
                n_const_fp16 += 1
                break
    print(f"  Constant nodes: {n_const_fp16} converted to FP16, "
          f"{n_const_fp32_keep} kept FP32")

    # === Update value_info dtypes ===
    n_vi_fp16 = 0
    for vi in model.graph.value_info:
        if vi.type.tensor_type.elem_type != TensorProto.FLOAT:
            continue
        if vi.name in fp32_tensors:
            continue
        vi.type.tensor_type.elem_type = TensorProto.FLOAT16
        n_vi_fp16 += 1
    print(f"  value_info: {n_vi_fp16} set to FP16")

    # === Update graph output dtype to match interior dtype ===
    # The original graph has output "velocity" as FP32. After our changes,
    # the graph's last layer (likely a MatMul or Add) is FP16 (it's not
    # blocked). So the actual output is FP16. We cast back to FP32 to
    # preserve the runtime contract.
    for go in model.graph.output:
        if go.type.tensor_type.elem_type != TensorProto.FLOAT:
            continue
        # Find producer
        prod = out2node.get(go.name)
        if prod is None:
            continue
        if prod.name in blocked_names:
            continue  # output is FP32 from blocked, no cast needed
        if go.name in fp32_tensors:
            continue  # already considered FP32
        # The output is now FP16 (produced by an FP16 trunk node). Insert
        # Cast(to=FP32) at the very end: rename the producer's output and
        # add a Cast node feeding the original output name.
        renamed_out = f"{go.name}__pre_final_fp32cast"
        cast_name = f"{go.name}__final_fp32cast"
        # Rewire the producer
        for i, oname in enumerate(prod.output):
            if oname == go.name:
                prod.output[i] = renamed_out
        # Add a value_info for the renamed (still FP16) output
        new_vi = helper.make_tensor_value_info(
            renamed_out, TensorProto.FLOAT16,
            [d.dim_param if d.HasField("dim_param") else d.dim_value
             for d in go.type.tensor_type.shape.dim])
        model.graph.value_info.append(new_vi)
        # Add the Cast node
        cast_node = helper.make_node(
            "Cast", inputs=[renamed_out], outputs=[go.name],
            name=cast_name, to=TensorProto.FLOAT)
        model.graph.node.append(cast_node)
        print(f"  added final Cast(FP16->FP32) for output {go.name}")

    return model


def convert_to_fp16(input_onnx, output_onnx, mode="minimal", bound_attn=True):
    """Load FP32 ONNX, identify FP32 islands, convert everything else to
    FP16, and save."""
    import onnx

    print(f"  loading {input_onnx} ({os.path.getsize(input_onnx)/1e6:.0f} MB)")
    model = onnx.load(input_onnx)
    # Materialise any external-data tensors INLINE so subsequent
    # numpy_helper.to_array calls don't try to re-read the .data sidecar
    # (which may be deleted, or located elsewhere from a tempdir-based
    # shape_inference fallback below). Walks initializers + Constant nodes.
    from onnx import numpy_helper as _nh
    _onnx_TP = onnx.TensorProto
    def _inline_tensor(t):
        if t.data_location != _onnx_TP.EXTERNAL:
            return
        # Read bytes via the existing external_data ref, then promote to raw_data.
        arr = _nh.to_array(t, base_dir=os.path.dirname(input_onnx))
        t.raw_data = arr.tobytes()
        t.data_location = _onnx_TP.DEFAULT
        t.ClearField("external_data")
    for init in model.graph.initializer:
        _inline_tensor(init)
    for node in model.graph.node:
        if node.op_type == "Constant":
            for attr in node.attribute:
                if attr.name == "value":
                    _inline_tensor(attr.t)
    print(f"  inlined external-data tensors")

    # Find FP32 islands FIRST (so RoPE mode can use the original Cast(to=FP32)
    # seeds that mark PyTorch's @autocast(enabled=False) regions).
    blocked_names = find_fp32_islands(model, mode=mode)

    # The original FP32 ONNX has 652 Cast(to=FP32) nodes — all no-ops
    # because the entire graph is already FP32. These were inserted by
    # PyTorch ONNX export to mark `.float()` and @autocast(enabled=False)
    # call sites. We strip them now (after finding islands by them).
    model = strip_noop_fp32_casts(model)

    # Pre-insert Cast(FP32) at every island input edge and Cast(FP16) at
    # every island output edge. These Casts are still no-ops in the
    # currently-FP32 graph; they become real boundaries after manual_convert.
    model = wrap_islands_with_casts(model, blocked_names)

    print(f"  manual conversion to FP16 (keeping {len(blocked_names)} nodes FP32)...")
    t0 = time.time()
    fp16_model = manual_convert_to_fp16(model, blocked_names)
    print(f"  converted in {time.time()-t0:.0f}s")

    # Walk the graph and insert Cast nodes wherever two operands of a
    # node disagree on dtype (FP16 vs FP32). This handles all the random
    # places where t5_hidden (FP32 input) feeds into an op alongside an
    # FP16 tensor, or where Cast(to=FP32) outputs reach FP16 consumers, etc.
    print(f"  fixing dtype mismatches with autocast insertion...")
    fp16_model = fix_dtype_mismatches(fp16_model, blocked_names)

    # The island passes above leave the whole O(L^2) attention core in FP32 —
    # see bound_attention_core() for why that is wrong and what it costs.
    if bound_attn:
        fp16_model = bound_attention_core(fp16_model)

    print(f"  saving to {output_onnx}")
    # Try a plain save first; if the protobuf exceeds 2GB, fall back to
    # external-data save (needed for sa3-m: its FP16-mixed proto is ~3GB
    # because most FP16 weights are inline but FP32 island weights doubled
    # them in size).
    try:
        onnx.save(fp16_model, output_onnx)
    except Exception as e:
        if "Failed to serialize" not in str(e) and "EncodeError" not in type(e).__name__:
            raise
        print(f"  (proto >2GB; using external-data save)")
        onnx.save_model(fp16_model, output_onnx, save_as_external_data=True,
                         all_tensors_to_one_file=True,
                         location=os.path.basename(output_onnx) + ".data",
                         size_threshold=1024, convert_attribute=False)
    print(f"  saved {os.path.getsize(output_onnx)/1e6:.0f} MB")

    # Quick dtype audit
    from collections import Counter
    init_dtypes = Counter()
    for init in fp16_model.graph.initializer:
        init_dtypes[init.data_type] += 1
    print(f"  initializer dtype counts (1=FP32, 10=FP16, 7=INT64): {dict(init_dtypes)}")

    cast_to_dtypes = Counter()
    for n in fp16_model.graph.node:
        if n.op_type == "Cast":
            for a in n.attribute:
                if a.name == "to":
                    cast_to_dtypes[a.i] += 1
    print(f"  Cast.to counts (1=FP32, 10=FP16): {dict(cast_to_dtypes)}")

    return output_onnx


def build_trt_engine(onnx_path, engine_path, workspace_gb=16):
    """Build TRT engine with STRONGLY_TYPED + sa3-sm-music profile."""
    import tensorrt as trt
    print(f"\n  building TRT engine -> {engine_path}")

    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    net_flags = 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
    network = builder.create_network(net_flags)
    parser = trt.OnnxParser(network, logger)

    if not parser.parse_from_file(str(onnx_path)):
        for i in range(parser.num_errors):
            print(f"  parse error: {parser.get_error(i)}")
        sys.exit(2)

    cfg = builder.create_builder_config()
    cfg.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_gb << 30)

    profile = builder.create_optimization_profile()
    for input_name, (lo, opt, hi) in _DIT_PROFILE.items():
        profile.set_shape(input_name, lo, opt, hi)
    cfg.add_optimization_profile(profile)
    print(f"  optimization profile: {len(_DIT_PROFILE)} input(s)")

    print(f"  building TRT (workspace {workspace_gb} GB, STRONGLY_TYPED, FP16-mixed)...")
    t0 = time.time()
    serialized = builder.build_serialized_network(network, cfg)
    if serialized is None:
        print(f"  BUILD FAILED")
        sys.exit(3)
    print(f"  built in {time.time()-t0:.0f}s ({serialized.nbytes/1e6:.0f} MB)")

    Path(engine_path).parent.mkdir(parents=True, exist_ok=True)
    with open(engine_path, "wb") as f:
        f.write(serialized)
    print(f"  wrote {engine_path}")
    return engine_path


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--input",  default="/tmp/dit_sm-music_fixed_v2.onnx",
                    help="Input FP32 ONNX")
    ap.add_argument("--onnx",   default="/tmp/dit_sm-music_fp16.onnx",
                    help="Output FP16-mixed ONNX (intermediate)")
    ap.add_argument("--engine", default="/weka2/cj/clod/sa3s/stable-audio-3/"
                                       "optimized/tensorRT/models/sm_90/"
                                       "sa3-sm-music/dit_fp16.trt",
                    help="Output TRT engine path")
    ap.add_argument("--workspace-gb", type=int, default=16)
    ap.add_argument("--mode", choices=("minimal", "rope", "full"),
                    default="rope",
                    help="FP32 island coverage: minimal=RMSNorm+Softmax only, "
                         "rope=+RoPE freq/apply_rotary, full=every Cast(FP32) "
                         "downstream region from the original ONNX.")
    ap.add_argument("--no-bound-attn", action="store_true",
                    help="Skip bound_attention_core(). NOT recommended: leaving the "
                         "RoPE island unbounded keeps QK^T + Softmax in FP32, which "
                         "blocks TRT's FMHA fuser (0 fused MHA nodes, 4.3x slower at "
                         "L=4096) for no accuracy gain. Only for reproducing the "
                         "pre-2026-07 engines.")
    ap.add_argument("--skip-convert", action="store_true",
                    help="Skip ONNX conversion (reuse existing --onnx file)")
    ap.add_argument("--skip-build", action="store_true",
                    help="Skip TRT build (just produce the ONNX)")
    args = ap.parse_args()

    if not args.skip_convert:
        print(f"━━━ Convert FP32 ONNX -> FP16-mixed ONNX (mode={args.mode}) ━━━")
        convert_to_fp16(args.input, args.onnx, mode=args.mode,
                             bound_attn=not args.no_bound_attn)

    if not args.skip_build:
        print("\n━━━ Build TRT engine ━━━")
        build_trt_engine(args.onnx, args.engine, args.workspace_gb)


if __name__ == "__main__":
    main()