File size: 7,644 Bytes
ffc54ec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """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()
|