From 7214ee9d42952e7aee51ce47fe9a5b88c5fd5a71 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 15 Jul 2026 15:06:40 +0200 Subject: [PATCH] triton/dynamo/inductor cache location and timer stats Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 2 ++ modules/devices.py | 19 ++++++++++++++-- modules/loader.py | 3 +++ modules/processing.py | 25 +++++++++++++++++++-- modules/sd_hijack_te.py | 2 ++ modules/sd_models_compile.py | 35 ++++++++++++++++++++++++++++++ modules/sdnq/common.py | 9 +++++++- modules/sdnq/timers.py | 10 +++++++++ modules/timer.py | 10 ++++++--- modules/video_models/video_load.py | 2 -- ui/gallery.ts | 1 - 11 files changed, 107 insertions(+), 11 deletions(-) create mode 100644 modules/sdnq/timers.py diff --git a/CHANGELOG.md b/CHANGELOG.md index f2673490a..9f3b21e84 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ - **Compute** - torch: update to `2.13.0` for CUDA, ROCm, IPEX + - torch: explicitly set inductor and triton cache locations + - torch: log triton/dynamo/inductor timer stats - cuda: update to `13.2` - sdnq quantization optimizations - sdnq attention optimizations diff --git a/modules/devices.py b/modules/devices.py index d7776594f..17dddc1ee 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -432,7 +432,6 @@ def test_triton(early: bool = False): if triton_version is None: try: import torch._inductor.triton as torch_triton - triton_version = torch_triton.__version__ except Exception: pass @@ -626,7 +625,23 @@ def set_cuda_params(): except Exception: tunable = [False, False] log.info(f'Torch parameters: backend={backend} device={device_name} config={opts.cuda_dtype} dtype={dtype} fp16={"pass" if fp16_ok else "fail"} bf16={"pass" if bf16_ok else "fail"} triton={"pass" if triton_ok else "fail"} optimization="{opts.cross_attention_optimization}"') - log.info(f'Torch compute: context={inference_context.__name__} nohalf={opts.no_half} nohalfvae={opts.no_half_vae} upcast={opts.upcast_sampling} deterministic={opts.cudnn_deterministic} tunable={tunable}') + try: + num_threads = torch._inductor.config.compile_threads # pylint: disable=protected-access + except Exception: + num_threads = None + log.info(f'Torch compute: context={inference_context.__name__} nohalf={opts.no_half} nohalfvae={opts.no_half_vae} upcast={opts.upcast_sampling} deterministic={opts.cudnn_deterministic} tunable={tunable} threads={num_threads}') + + try: + from torch._inductor.runtime.runtime_utils import cache_dir + inductor_cache = cache_dir() + except Exception: + inductor_cache = os.getenv("TORCHINDUCTOR_CACHE_DIR", None) + try: + from triton import knobs + triton_cache = knobs.cache.dir + except Exception: + triton_cache = os.getenv("TRITON_CACHE_DIR", None) + log.info(f'Torch cache: inductor="{inductor_cache}" triton="{triton_cache}"') def randn(seed, shape=None): diff --git a/modules/loader.py b/modules/loader.py index 290781495..ed47ab628 100644 --- a/modules/loader.py +++ b/modules/loader.py @@ -56,6 +56,9 @@ except Exception as e: report(f'scipy=={scipy.__version__ if scipy is not None else None}', e) timer.startup.record("scipy") +inductor_cache = os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "inductor")) +triton_cache = os.environ.setdefault("TRITON_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "triton")) + try: import atexit import torch._inductor.async_compile as ac diff --git a/modules/processing.py b/modules/processing.py index f25bdf512..40a1394a5 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -416,6 +416,27 @@ def process_samples(p: StableDiffusionProcessing, samples): return out_images, out_infotexts +def print_stats(): + log.debug(f'Processed: timers={timer.process.dct()}') + log.debug(f'Processed: memory={memstats.memory_stats()}') + + if timer.compiler.get_total() > 0.0001: + log.debug(f'Processed: compile={timer.compiler.dct(min_time=0)}') + timer.compiler.reset() + + if shared.opts.sdnq_dequantize_compile: + from modules.sdnq.timers import update_sdnq_attention_timers + update_sdnq_attention_timers() + if timer.autotune.get_total() > 0.0001: + log.debug(f'Processed: autotune={timer.autotune.dct(min_time=0)}') + + if devices.triton_ok: + from modules.sd_models_compile import update_compile_times + update_compile_times() + if timer.dynamo.get_total() > 0.0001: + log.debug(f'Processed: dynamo={timer.dynamo.dct(min_time=1.0)}') + + def process_images_inner(p: StableDiffusionProcessing) -> Processed: if type(p.prompt) == list: assert len(p.prompt) > 0 @@ -570,10 +591,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: p.scripts.postprocess(p, results) timer.process.record('post') p.ops = list(set(p.ops)) + if not p.disable_extra_networks: log.info(f'Processed: images={len(output_images)} its={(p.steps * len(output_images)) / (t1 - t0):.2f} ops={p.ops}') - log.debug(f'Processed: timers={timer.process.dct()}') - log.debug(f'Processed: memory={memstats.memory_stats()}') + print_stats() if shared.cmd_opts.lowvram or shared.cmd_opts.medvram: devices.torch_gc(force=True, reason='final') diff --git a/modules/sd_hijack_te.py b/modules/sd_hijack_te.py index ff52c399e..0e9111c09 100644 --- a/modules/sd_hijack_te.py +++ b/modules/sd_hijack_te.py @@ -80,6 +80,8 @@ def hijack_encode_prompt(*args, **kwargs): errors.display(e, 'Encode prompt') t1 = time.time() timer.process.add('te', t1-t0) + if t1 - t0 > 10: + log.warning(f'Encode: time={t1-t0:.3f} long encode prompt') shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) shared.state.end(jobid) # from modules import memstats diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 55978c210..ea2ed669e 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -370,3 +370,38 @@ def openvino_post_compile(op="base"): # delete unet after OpenVINO compile if not shared.opts.openvino_disable_memory_cleanup and hasattr(shared.sd_refiner, "unet"): shared.sd_refiner.unet.apply(sd_models_utils.convert_to_faketensors) devices.torch_gc(force=True) + + +def update_compile_times(): + try: + from torch._dynamo.utils import compile_times, reset_frame_count + raw_str = compile_times() + reset_frame_count() + except Exception: + return {} + from modules.timer import dynamo + dynamo.reset() + lines = raw_str.strip().split('\n') + # parsed = [] + for line in lines: + if not line or 'TorchDynamo compilation metrics' in line or 'Function, Runtimes' in line: + continue + parts = line.split(',') + fn = parts[0].strip() + try: + times = [float(t.strip()) for t in parts[1:] if t.strip()] + if times: + # parsed.append((fn, sum(times), len(times), max(times))) + dynamo.add(fn, round(sum(times), 2)) + except ValueError: + continue + """ + parsed.sort(key=lambda x: x[1], reverse=True) + results = {} + min_time = 0.1 + for fn, total, count, max_val in parsed: + if total > min_time: + dynamo.ts(fn, total) + results[fn] = { "total": round(total, 2), "count": count, "avg": round(total / count, 2), "max": round(max_val, 2) } + return results + """ diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py index 6d4031d37..c2021ea7d 100644 --- a/modules/sdnq/common.py +++ b/modules/sdnq/common.py @@ -2,6 +2,7 @@ import os import json +import time import torch from modules import shared @@ -338,10 +339,13 @@ weights_dtype_order = [ use_torch_compile = shared.opts.sdnq_dequantize_compile # this setting requires a full restart of the webui to apply + def check_torch_compile() -> bool: # dynamo can be disabled after startup return use_torch_compile and not torch._dynamo.config.disable # pylint: disable=protected-access + if use_torch_compile: + from modules.timer import compiler as compile_timer if hasattr(torch._dynamo.config, "recompile_limit"): torch._dynamo.config.recompile_limit = max(8192, getattr(torch._dynamo.config, "recompile_limit", 0)) if hasattr(torch._dynamo.config, "cache_size_limit"): @@ -361,7 +365,10 @@ if use_torch_compile: if os.environ.get("SDNQ_COMPILE_KWARGS", None) is not None: for key, value in json.loads(os.environ.get("SDNQ_COMPILE_KWARGS")).items(): kwargs[key] = value - return torch.compile(fn, **kwargs) + t0 = time.time() + res = torch.compile(fn, **kwargs) + compile_timer.ts(fn.__name__, t0) + return res else: def compile_func(fn, **kwargs): # pylint: disable=unused-argument return fn diff --git a/modules/sdnq/timers.py b/modules/sdnq/timers.py new file mode 100644 index 000000000..4d3216c52 --- /dev/null +++ b/modules/sdnq/timers.py @@ -0,0 +1,10 @@ +def update_sdnq_attention_timers(): + from modules.timer import autotune + autotune.reset() + from modules.sdnq.kernels import triton_atten, triton_mm, triton_scaled_mm + if getattr(triton_atten.sdnq_attn_kernel, 'bench_time', None) is not None: + autotune.add('sdnq_attn_kernel', getattr(triton_atten.sdnq_attn_kernel, 'bench_time', 0)) + if getattr(triton_mm.sdnq_triton_mm, 'bench_time', None) is not None: + autotune.add('sdnq_triton_mm', getattr(triton_mm.sdnq_triton_mm, 'bench_time', 0)) + if getattr(triton_scaled_mm.sdnq_scaled_mm, 'bench_time', None) is not None: + autotune.add('sdnq_scaled_mm', getattr(triton_scaled_mm.sdnq_scaled_mm, 'bench_time', 0)) diff --git a/modules/timer.py b/modules/timer.py index 0d650e283..d5d94df01 100644 --- a/modules/timer.py +++ b/modules/timer.py @@ -10,11 +10,11 @@ except Exception: class Timer: - def __init__(self): + def __init__(self, profile=False): self.start = time.time() self.records = {} self.total = 0 - self.profile = False + self.profile = profile def elapsed(self, reset=True): end = time.time() @@ -71,10 +71,14 @@ class Timer: return res def reset(self): - self.__init__() + self.records.clear() + self.__init__(self.profile) startup = Timer() process = Timer() launch = Timer() init = Timer() load = Timer() +dynamo = Timer() +compiler = Timer(profile=True) +autotune = Timer(profile=True) diff --git a/modules/video_models/video_load.py b/modules/video_models/video_load.py index a0d003ada..5744e68e6 100644 --- a/modules/video_models/video_load.py +++ b/modules/video_models/video_load.py @@ -155,9 +155,7 @@ def load_model(selected: models_def.Model): shared.sd_model = load_custom(selected.repo) else: log.debug(f'Load video: module=pipe repo="{selected.repo}" cls={selected.repo_cls.__name__}') - print('HERE1') sd_models.hf_prefetch_configs(selected.repo, {}, 'video') - print('HERE2') shared.sd_model = selected.repo_cls.from_pretrained( pretrained_model_name_or_path=selected.repo, revision=selected.repo_revision, diff --git a/ui/gallery.ts b/ui/gallery.ts index f82839277..e09e52cec 100644 --- a/ui/gallery.ts +++ b/ui/gallery.ts @@ -1013,7 +1013,6 @@ export async function gallerySort(key) { const folderNames = Array.from(folderGroups.keys()); const sortedFolderNames = currentSort.endsWith('A') ? folderNames.sort((a, b) => a.localeCompare(b)) : folderNames.sort((a, b) => b.localeCompare(a)); - console.log('HERE', sortedFolderNames); for (const folderName of sortedFolderNames) { const files = folderGroups.get(folderName); files.sort(sortMode.func);