[Bug][Performance Regression][Relax][DLight] FP16 Exp producer inlining across broadcast causes 30-60% slowdown on RTX A6000

[pinpinpoo]

Enabling FuseOps/FuseTIR causes a performance regression for an FP16 Exp chain followed by a broadcast Add on an RTX A6000. VM latency increases by 30% with 8 Exp layers and 60% with 16 layers compared with disabling these passes. Fusing the Exp chain while keeping the broadcast Add separate is faster than either configuration.

Workload

Inputs are x: [4096, 1] and b: [4096, 4096], both FP16:

z = x
for _ in range(depth):
    z = exp(x - float16(0.5) * z)
y = b + z

Inputs are sampled from [-0.25, 0.25] with fixed seeds. This recurrence keeps the Exp chain finite. The reference rounds to FP16 after each operator.

Results

The three configurations have the same inputs and a single [4096, 4096] output:

• OFF: skip FuseOps and FuseTIR in the CUDA pipeline.
• PARTIAL: place the Exp producer and broadcast Add in separate dataflow blocks, then run the full pipeline.
• ON: run the full pipeline on a single dataflow block.

Exp layers	ON/OFF	PARTIAL/OFF	ON/PARTIAL	Trials passing the stability and 10% slowdown thresholds
8	1.301768	0.793551	1.641148	5/5
16	1.595166	0.646946	2.444368	4/5

Ratios are latency ratios, summarized as the median of five per-trial ratios. For 16 layers, ON/OFF was 1.584963, 1.595166, 1.602671, 1.615977, 1.569858. One trial exceeded the 5% CV threshold; it is included in the table. GPU clocks were not locked.

Each trial runs in a separate process. Timing uses vm.time_evaluator("invoke_stateful", ...) with GPU-resident inputs: number=10, repeat=5, min_repeat_ms=50, and 40 measurement blocks, giving 200 samples per configuration. Warmup is 75 calls followed by 2 seconds of timed batches. The six configuration orders are cycled across blocks and offset across trials. The measured time includes VM allocations; compilation, transfers and correctness checks are outside the timer.

All three configurations passed reference and pairwise checks on three input draws per trial, with zero maximum absolute error. The saved IR confirms that ON fuses the Exp-to-Add edge and PARTIAL retains a fused producer followed by a separate Add. Cubin hashes match across the five trials for each configuration.

Environment

• TVM: 0.26.dev0+source.8f328e8
• Source commit recorded in the build manifest: 8f328e802cfe5e41fcc8f5c17e7582b1c28bfce4
• GPU: NVIDIA RTX A6000, sm_86; driver 550.120
• CUDA toolkit: 12.4.131; cuda-bindings: 12.9.7
• LLVM: 15.0.7; host target: llvm -mcpu=generic
• Linux x86_64, Python 3.11.16, NumPy 2.4.6

Reproduction

Save the code below as repro_broadcast_regression.py and run in a CUDA-enabled TVM environment with NumPy installed:

python repro_broadcast_regression.py --depth 8 --device 0 --out repro_depth8
python repro_broadcast_regression.py --depth 16 --device 0 --out repro_depth16

Use an idle GPU and a new output directory for each command. Each command runs five trials and saves per-pass IR, correctness results, timing samples and ratios. summary.json contains the median ratios and the number of stable trials with at least a 10% slowdown.

The script checks for 3 * depth + 1, 2, and 1 calls to relax.call_tir before VM lowering in OFF, PARTIAL, and ON. A failure here means the generated graph needs to be inspected before comparing timings.

I extracted this script from the benchmark used for the measurements above; I haven't rerun the standalone version on CUDA yet.

"""Standalone extraction for TVM 0.26.dev0+source.8f328e8 / CUDA.

Run twice: --depth 8 --out repro_d8, then --depth 16 --out repro_d16.
Each command runs five separate processes. No benchmark-package imports.
"""
import argparse
import itertools
import json
import math
from pathlib import Path
import platform
import statistics
import subprocess
import sys
import time

VARIANTS = ("off", "partial", "on")
ORDERS = list(itertools.permutations(VARIANTS))


def reference(x, b, depth):
    import numpy as np
    z = x
    for _ in range(depth):
        z = (x - (z * np.float16(0.5)).astype("float16")).astype("float16")
        z = np.exp(z.astype("float64")).astype("float16")
    return (z + b).astype("float16")


def graph(depth, partial):
    from tvm import relax as R
    from tvm.relax import op
    x = R.Var("x", R.TensorType([4096, 1], "float16"))
    b = R.Var("b", R.TensorType([4096, 4096], "float16"))
    bb = R.BlockBuilder()

    def producer():
        z = x
        for _ in range(depth):
            z = bb.emit(op.multiply(z, R.const(0.5, "float16")))
            z = bb.emit(op.subtract(x, z))
            z = bb.emit(op.exp(z))
        return z

    with bb.function("main", [x, b]):
        if partial:
            with bb.dataflow():
                z = bb.emit_output(producer())
            with bb.dataflow():
                y = bb.emit_output(bb.emit(op.add(z, b)))
        else:
            with bb.dataflow():
                y = bb.emit_output(bb.emit(op.add(producer(), b)))
        bb.emit_func_output(y)
    return bb.get()


def compile_vm(depth, mode, target, dev, out):
    import tvm
    from tvm import relax as R
    from tvm.relax.backend.cuda import pipeline
    out.mkdir()
    mod = graph(depth, mode == "partial")
    (out / "input.py").write_text(mod.script(), encoding="utf-8")

    @tvm.instrument.pass_instrument
    class Audit:
        def run_before_pass(self, mod, info):
            if mode == "off" and str(info.name).split(".")[-1] in ("FuseOps", "FuseTIR"):
                raise RuntimeError("Unexpected fusion pass in OFF")

    with target, tvm.transform.PassContext(opt_level=3, instruments=[Audit()]):
        steps = pipeline.library_dispatch_passes(target) + pipeline.legalize_passes(target)
        for i, p in enumerate(steps):
            name = str(p.info.name).split(".")[-1]
            if mode == "off" and name in ("FuseOps", "FuseTIR"):
                continue
            mod = p(mod)
            (out / f"{i:02d}_{name}.py").write_text(mod.script(), encoding="utf-8")
        calls = []

        def visit(expr):
            if isinstance(expr, R.Call) and isinstance(expr.op, tvm.ir.Op):
                if expr.op.name == "relax.call_tir":
                    calls.append(expr.args[0].name_hint)

        R.analysis.post_order_visit(mod["main"].body, visit)
        expected = {"off": 3 * depth + 1, "partial": 2, "on": 1}[mode]
        if len(calls) != expected:
            raise RuntimeError(f"{mode}: expected {expected} call_tir calls, got {calls}; inspect IR")
        for p in pipeline.dataflow_lower_passes(target) + pipeline.finalize_passes(target):
            mod = p(mod)
        ex = R.build(mod, target=target, relax_pipeline=None, tir_pipeline="default")
    return R.VirtualMachine(ex, dev)


def stats(samples):
    if len(samples) != 200 or any(not math.isfinite(x) or x <= 0 for x in samples):
        raise ValueError("Expected 200 finite positive samples")
    return dict(median_s=statistics.median(samples),
                cv=statistics.stdev(samples) / statistics.mean(samples), samples_s=samples)


def trial(args):
    import numpy as np
    import tvm
    dev = tvm.cuda(args.device)
    if not dev.exist:
        raise RuntimeError("CUDA device unavailable")
    target = tvm.target.Target(
        {"kind": "cuda", "arch": "sm_" + str(dev.compute_version).replace(".", "")},
        host={"kind": "llvm", "mcpu": "generic"})
    order = ORDERS[args.trial % 6]
    vms = {k: compile_vm(args.depth, k, target, dev, args.out / k) for k in order}
    errors = []
    for seed in range(3):
        rng = np.random.default_rng(seed)
        x = rng.uniform(-.25, .25, (4096, 1)).astype("float16")
        b = rng.uniform(-.25, .25, (4096, 4096)).astype("float16")
        arrays = [tvm.runtime.tensor(a, device=dev) for a in (x, b)]
        if seed == 0:
            timed_inputs = arrays
        expected, outputs = reference(x, b, args.depth), {}
        for k, vm in vms.items():
            vm.set_input("main", *arrays)
            vm.invoke_stateful("main")
            dev.sync()
            value = vm.get_outputs("main")
            while not hasattr(value, "numpy"):
                if len(value) != 1:
                    raise AssertionError("Expected exactly one output")
                value = value[0]
            outputs[k] = value.numpy()
        pairs = [(k, outputs[k], expected) for k in VARIANTS]
        pairs += [(f"{a}/{b}", outputs[a], outputs[b])
                  for a, b in itertools.combinations(VARIANTS, 2)]
        for label, actual, wanted in pairs:
            if actual.shape != wanted.shape or actual.dtype != wanted.dtype:
                raise AssertionError(f"{label}: output shape/dtype mismatch")
            if not np.isfinite(actual).all() or not np.isfinite(wanted).all():
                raise AssertionError(f"{label}: non-finite output")
            np.testing.assert_allclose(actual, wanted, rtol=.02, atol=.002)
            errors.append(dict(seed=seed, comparison=label,
                               max_abs_error=float(np.max(np.abs(
                                   actual.astype("float32") - wanted.astype("float32"))))))
    timers, samples = {}, {k: [] for k in VARIANTS}
    for k in order:
        vm = vms[k]
        vm.set_input("main", *timed_inputs)
        timers[k] = vm.time_evaluator("invoke_stateful", dev, number=10,
                                     repeat=5, min_repeat_ms=50)
        for _ in range(75):
            vm.invoke_stateful("main")
        dev.sync()
        deadline = time.monotonic() + 2.0
        while time.monotonic() < deadline:
            dev.sync()
            timers[k]("main")
    for block in range(40):
        for k in ORDERS[(args.trial + block) % 6]:
            dev.sync()
            samples[k].extend(float(v) for v in timers[k]("main").results)
    measured = {k: stats(v) for k, v in samples.items()}
    ratios = {f"{a}/{b}": measured[a]["median_s"] / measured[b]["median_s"]
              for a, b in (("on", "off"), ("partial", "off"), ("on", "partial"))}
    result = dict(depth=args.depth, trial=args.trial, tvm=tvm.__version__,
                  python=sys.version, platform=platform.platform(), target=str(target),
                  gpu=str(dev.device_name), correctness=errors, measured=measured, ratios=ratios,
                  on_off_stable=max(measured[k]["cv"] for k in ("on", "off")) <= .05)
    (args.out / "result.json").write_text(json.dumps(result, indent=2), encoding="utf-8")
    print(json.dumps({"trial": args.trial, "ratios": ratios,
                      "CV": {k: v["cv"] for k, v in measured.items()}}), flush=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--depth", type=int, choices=(8, 16), required=True)
    parser.add_argument("--device", type=int, default=0)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--trial", type=int, help=argparse.SUPPRESS)
    args = parser.parse_args()
    if args.out.exists():
        parser.error("--out must be a new directory")
    args.out.mkdir(parents=True)
    if args.trial is not None:
        trial(args)
        return
    results = []
    for i in range(5):
        dest = args.out / f"trial_{i}"
        print(f"Depth {args.depth}, trial {i + 1}/5", flush=True)
        subprocess.run([sys.executable, str(Path(__file__).resolve()), "--depth", str(args.depth),
                        "--device", str(args.device), "--out", str(dest), "--trial", str(i)], check=True)
        results.append(json.loads((dest / "result.json").read_text(encoding="utf-8")))
    summary = dict(median_ratios={k: statistics.median(r["ratios"][k] for r in results)
                                 for k in results[0]["ratios"]},
                   stable_on_off_slow_trials=sum(r["on_off_stable"] and r["ratios"]["on/off"] >= 1.10
                                                for r in results), total_trials=5)
    (args.out / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()

Where the extra computation appears

After FuseTIR, the fused PrimFunc still has [4096, 1] intermediate buffers and small producer loops, followed by the broadcast Add. After ApplyDefaultSchedule, the Exp chain is inlined into the [4096, 4096] output expression. The scheduling trace selects DLight Fallback.

For 16 layers:

• PARTIAL schedules the producer with 4 blocks of 1024 threads, then runs a separate Add kernel.
• ON schedules the output with 16384 blocks of 1024 threads and embeds the entire Exp chain in that kernel. Its SASS contains 16 MUFU.EX2 instruction sites.

This points to repeated producer evaluation along the broadcast dimension. The ON cubin uses 13 registers; PARTIAL uses 11 for the producer and 10 for Add. All report zero stack, local and shared memory, and the disassembly contains no LDL/STL instructions. These results do not suggest register spilling. Per-kernel timings and SFU utilization still need profiling.

Relevant code at the recorded commit:

• Fallback.apply calls try_inline before scheduling the remaining block.
• try_inline tries compute_inline and reverse_compute_inline.
• GraphPartitioner::RunFuse makes the earlier fusion decision.

Keeping the producer materialized reduces its output to an 8 KiB buffer and gives the fastest result in this test. Could the fusion or inlining decision account for producer cost and broadcast reuse? The PARTIAL configuration provides a working comparison for that change.

[KHARSHAVARDHAN-eng]

Hi @pinpinpoo, thanks for reporting this issue. I’d like to work on it.

I’ll investigate the FuseOps/FuseTIR interaction causing the FP16 performance regression, reproduce the issue, identify the underlying cause, and add appropriate regression/benchmark coverage.

Thanks!

[pinpinpoo]

Thanks @KHARSHAVARDHAN-eng — ping me if you need the harness.