"""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()