float32 attention on inputs cast from bf16 crashes the builder or fails to build (casts folded into IAttention) #4936

import argparse, subprocess, sys
import numpy as np
import tensorrt as trt

p = argparse.ArgumentParser()
p.add_argument("--shape", default="8,8,50,805,64", help="B,H,L,S,D")
p.add_argument("--case", help="MASK,CAST_BACK: run one case per process, because some cases crash")
a = p.parse_args()
B, H, L, S, D = (int(x) for x in a.shape.split(","))
logger = trt.Logger(trt.Logger.WARNING)


def build(mask_kind, cast_back, input_type=trt.bfloat16):
    builder = trt.Builder(logger)
    net = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
    q = net.add_input("q", input_type, (B, H, L, D))
    k = net.add_input("k", input_type, (B, H, S, D))
    v = net.add_input("v", input_type, (B, H, S, D))
    if input_type != trt.float32:
        q, k, v = (net.add_cast(t, trt.float32).get_output(0) for t in (q, k, v))
    att = net.add_attention_v2(q, k, v, trt.AttentionNormalizationOp.SOFTMAX, trt.CausalMaskKind.NONE)
    att.decomposable = True
    if mask_kind == "float":
        att.mask = net.add_input("mask", trt.float32, (B, 1, 1, S))
    elif mask_kind == "bool":
        att.mask = net.add_input("mask", trt.bool, (B, 1, 1, S))
    out = att.get_output(0)
    if cast_back:
        out = net.add_cast(out, trt.bfloat16).get_output(0)
    net.mark_output(out)
    try:
        plan = builder.build_serialized_network(net, builder.create_builder_config())
        return "builds" if plan is not None else "BUILD FAILED"
    except Exception as error:
        return f"BUILD FAILED: {error}"


if a.case:
    mask_kind, cast_back, inputs = a.case.split(",")
    print(build(mask_kind, cast_back == "1", trt.float32 if inputs == "fp32" else trt.bfloat16), flush=True)
    sys.exit(0)
print("TensorRT", trt.__version__, "shape", (B, H, L, S, D), flush=True)
cases = [(m, c, "bf16") for m in ("none", "bool", "float") for c in (True, False)]
cases += [(m, False, "fp32") for m in ("none", "bool", "float")]  # control: float32 inputs, no casts to fold
for mask_kind, cast_back, inputs in cases:
    r = subprocess.run([sys.executable, __file__, "--shape", a.shape, "--case", f"{mask_kind},{int(cast_back)},{inputs}"],
                       capture_output=True, text=True)
    text = (r.stdout + r.stderr).splitlines()
    out = [l for l in r.stdout.splitlines() if l.startswith(("builds", "BUILD"))]
    errors = [l.strip() for l in text if "differs from" in l][:1]
    errors += [l.strip()[l.find("Error Code"):] for l in text if "[E]" in l and "Error Code" in l][-1:]
    result = out[-1] if r.returncode == 0 and out else f"CRASHED (exit code {r.returncode})"
    if result != "builds" and errors:
        result += " | " + " | ".join(e[:200] for e in errors)
    print(f"inputs={inputs} mask={mask_kind:5} output={'bf16' if cast_back else 'fp32'}: {result}", flush=True)