source/benchmark.py

"""Warm end-to-end timing + stage breakdown + optional kernel profile for one variant.

CUDA-synchronized wall time per generation includes prompt encoding, all 8 denoising steps and
VAE decode; excludes model load, first-use warmup/compilation and PNG writing. Batch 1, CFG 1,
prefix KV cache, saved Turbo schedule.
"""
import argparse, collections, json, statistics, time
from pathlib import Path
import torch
from common import ROOT, BASE, load_bf16_pipeline
from prompts import EVALUATION


def load(a):
    if a.precision == 'bf16':
        return load_bf16_pipeline()
    if a.precision == 'fp8':
        from fp8_runtime import load_pipeline
    else:
        from nvfp4_runtime import load_pipeline
    return load_pipeline(str(BASE), a.quant or str(ROOT / a.precision / 'release' / 'transformer'))


@torch.inference_mode()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('precision', choices=['bf16', 'fp8', 'nvfp4'])
    ap.add_argument('--quant'); ap.add_argument('--accelerate', action='store_true')
    ap.add_argument('--attention', default='bf16', choices=['bf16', 'sage']); ap.add_argument('--fp8-first-steps', type=int, default=0)
    ap.add_argument('--sizes', default='1024,2048'); ap.add_argument('--n', type=int, default=5)
    ap.add_argument('--profile', action='store_true'); ap.add_argument('--out')
    ap.add_argument('--prompt-index', type=int, default=15)
    a = ap.parse_args()
    t = time.perf_counter(); pipe = load(a); torch.cuda.synchronize(); load_s = time.perf_counter() - t
    if a.accelerate:
        from acceleration import accelerate_pipeline
        accelerate_pipeline(pipe, attention=a.attention)
    if a.fp8_first_steps:
        from nvfp4_runtime import enable_fp8_first_steps
        enable_fp8_first_steps(pipe, str(ROOT / 'fp8' / 'release' / 'transformer'), a.fp8_first_steps,
                               accelerate=a.accelerate, attention=a.attention)
    stages = collections.defaultdict(list)
    def timed(name, fn):
        def wrap(*args, **kw):
            torch.cuda.synchronize(); s = time.perf_counter(); r = fn(*args, **kw); torch.cuda.synchronize()
            stages[name].append(time.perf_counter() - s); return r
        return wrap
    prompt = EVALUATION[a.prompt_index]
    result = {'precision': a.precision, 'accelerated': a.accelerate, 'attention': a.attention,
              'fp8_first_steps': a.fp8_first_steps, 'load_seconds': load_s, 'torch': torch.__version__,
              'gpu': torch.cuda.get_device_name(), 'timing': {}}
    for size in [int(s) for s in a.sizes.split(',')]:
        t = time.perf_counter()
        pipe(prompt=prompt, width=size, height=size, generator=torch.Generator('cuda').manual_seed(30000))
        torch.cuda.synchronize(); warm = time.perf_counter() - t
        vals = []
        for j in range(a.n):
            torch.cuda.reset_peak_memory_stats(); torch.cuda.synchronize(); t = time.perf_counter()
            pipe(prompt=prompt, width=size, height=size, generator=torch.Generator('cuda').manual_seed(30001 + j))
            torch.cuda.synchronize(); vals.append(time.perf_counter() - t)
        # Separate instrumented run (adds synchronizations) for the stage breakdown.
        stages.clear()
        targets = [(pipe, 'encode_prompt', 'encode'), (pipe.vae, 'decode', 'vae_decode'), (pipe.transformer, 'forward', 'transformer')]
        saved = [(obj, name, obj.__dict__.get(name, None), name in obj.__dict__) for obj, name, _ in targets]
        for obj, name, label in targets:
            setattr(obj, name, timed(label, getattr(obj, name)))
        pipe(prompt=prompt, width=size, height=size, generator=torch.Generator('cuda').manual_seed(30001))
        for obj, name, value, had in saved:  # restore exactly (keeps runtime patches such as split prefill/hybrid)
            if had:
                obj.__dict__[name] = value
            else:
                obj.__dict__.pop(name, None)
        result['timing'][str(size)] = {'seconds': vals, 'mean': statistics.mean(vals), 'median': statistics.median(vals),
                                       'min': min(vals), 'first_use_seconds': warm, 'peak_gb': torch.cuda.max_memory_allocated() / 1e9,
                                       'stages': {'encode': stages['encode'], 'vae_decode': stages['vae_decode'],
                                                  'transformer_calls': stages['transformer']}}
        print(json.dumps({'size': size, 'mean': result['timing'][str(size)]['mean'], 'min': min(vals), 'first_use': warm,
                          'encode': stages['encode'], 'vae': stages['vae_decode'],
                          'transformer': [round(x, 4) for x in stages['transformer']]}), flush=True)
    if a.profile:
        prof = torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA])
        count = {'n': 0}
        def before(m, args, kw):
            if count['n'] == 3: prof.start()
        def after(m, args, kw, r):
            if count['n'] == 3: torch.cuda.synchronize(); prof.stop()
            count['n'] += 1
        h1 = pipe.transformer.register_forward_pre_hook(before, with_kwargs=True)
        h2 = pipe.transformer.register_forward_hook(after, with_kwargs=True)
        pipe(prompt=prompt, width=1024, height=1024, generator=torch.Generator('cuda').manual_seed(40000))
        h1.remove(); h2.remove()
        totals = collections.defaultdict(lambda: [0., 0])
        for e in prof.events():
            if e.device_type == torch.autograd.DeviceType.CUDA:
                totals[e.name][0] += e.device_time_total / 1e3; totals[e.name][1] += 1
        top = sorted(totals.items(), key=lambda kv: -kv[1][0])[:40]
        result['cached_step_kernels_ms'] = [{'kernel': k[:160], 'ms': v[0], 'calls': v[1]} for k, v in top]
        result['cached_step_gpu_ms'] = sum(v[0] for v in totals.values())
        print(json.dumps({'cached_step_gpu_ms': result['cached_step_gpu_ms']}))
        for row in result['cached_step_kernels_ms'][:25]:
            print(f"{row['ms']:8.3f} ms {row['calls']:5d}  {row['kernel'][:120]}")
        # Whole 1024 generation: count native quantized GEMM and attention kernel launches per call type.
        prof = torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA])
        calls = {'n': 0}
        h = pipe.transformer.register_forward_pre_hook(lambda m, a_, k_: calls.__setitem__('n', calls['n'] + 1), with_kwargs=True)
        with prof:
            pipe(prompt=prompt, width=1024, height=1024, generator=torch.Generator('cuda').manual_seed(40001))
            torch.cuda.synchronize()
        h.remove()
        names = collections.Counter(e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA)
        result['generation_kernel_evidence'] = {
            'nvfp4_blockscaled_gemm_sm120': sum(v for k, v in names.items() if 'blockscaled_gemm_sm120' in k and 'f4E2M1FN' in k),
            'fp8_e4m3_cutlass_sm120_gemm': sum(v for k, v in names.items() if 'Sm120' in k and 'float_e4m3' in k),
            'sage_attention_kernels': sum(v for k, v in names.items() if 'qk_int8' in k.lower() or 'sageattn' in k.lower() or 'qk_int_sv_f8' in k.lower()),
            'bf16_flash_attention': sum(v for k, v in names.items() if 'flash_fwd' in k),
            'transformer_forward_calls_seen_by_hook': calls['n'],
            'all_kernels': dict(names)}
        print(json.dumps({k: v for k, v in result['generation_kernel_evidence'].items() if k != 'all_kernels'}))
    if a.out:
        Path(a.out).parent.mkdir(parents=True, exist_ok=True); Path(a.out).write_text(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()