From 51198538de342bfc0959a64a652f8e739ee54ff2 Mon Sep 17 00:00:00 2001 From: eddy Date: Sat, 11 Oct 2025 06:48:21 -0700 Subject: [PATCH] Add files via upload --- README.md | 93 +++++++++++++++++++++ __init__.py | 4 + nodes.py | 214 +++++++++++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 15 ++++ 4 files changed, 326 insertions(+) create mode 100644 README.md create mode 100644 __init__.py create mode 100644 nodes.py create mode 100644 pyproject.toml diff --git a/README.md b/README.md new file mode 100644 index 0000000..b3a60b0 --- /dev/null +++ b/README.md @@ -0,0 +1,93 @@ +# ComfyUI-TorchCompileSpeed + +A compact, non-intrusive ComfyUI node set that boosts torch.compile performance and cache hit rate. Designed to plug into WanVideo Cython Model Loader without touching its source. + +Author: eddy + +## Highlights + +- Speed mode (recommended) + - inductor + max-autotune-no-cudagraphs + - dynamic=True for better shape tolerance and cache reuse + - CUDA Graphs disabled to cut capture overhead + - Triton autotune fully enabled +- Smarter reuse (inside our nodes only) + - Optional reuse_if_similar to skip recompiles for the same model + config +- Experimental PTX assist (for no-CUDA-Graphs paths) + - When experimental_ptx is on, a light warmup triggers PTX/kernel cache early + - Fast-math toggle and optional TRITON_CACHE_DIR for cross-session reuse +- Drop‑in with WanVideo Cython Model Loader + - Output type is WANCOMPILEARGS, so you can wire it straight into compile_args + +## Installation + +- Place this folder under ComfyUI/custom_nodes and restart ComfyUI + +## Nodes and parameters + +### Torch Compile Speed Settings +Output: WANCOMPILEARGS (torch_compile_args) + +Required: +- backend: inductor/cudagraphs (default inductor) +- fullgraph: enable fullgraph (default False) +- mode: default/max-autotune/max-autotune-no-cudagraphs/reduce-overhead/speed (default speed) +- dynamic: dynamic compilation (default False; set True for speed mode) +- dynamo_cache_size_limit: torch._dynamo.config.cache_size_limit (default 64) +- compile_transformer_blocks_only: compile transformer blocks only (default True) +- reuse_if_similar: reuse compiled results for same model + config (default True) +- experimental_ptx: enable experimental PTX assist (default False) +- ptx_fast_math: enable fast-math when available (default True) +- warmup_runs: warmup iterations to trigger PTX/kernel cache (default 1, range 0–5) + +Optional: +- dynamo_recompile_limit: torch._dynamo.config.recompile_limit (default 128) +- ptx_cache_dir: set TRITON_CACHE_DIR to help cross‑session reuse + +Recommended (speed‑first): +- mode=speed, dynamic=True, fullgraph=False +- compile_transformer_blocks_only=True +- reuse_if_similar=True +- experimental_ptx=True, warmup_runs=1–2, ptx_fast_math=True + +### Apply Torch Compile (optional) +Input: MODEL + WANCOMPILEARGS, Output: MODEL +- Wraps model forward with torch.compile +- With reuse_if_similar=True, repeated calls for the same model + config reuse the compiled forward +- With experimental_ptx=True: + - Tries a lightweight triton.ops matmul warmup; if unavailable, falls back to a small torch.compile matmul warmup + - You can set TRITON_CACHE_DIR; combined with warmup_runs it tends to make PTX/kernel cache available sooner + +## WanVideo integration (non‑intrusive) +- Wire directly: Torch Compile Speed Settings → WanVideo Cython Model Loader.compile_args +- Or use with other models: Settings → Apply Torch Compile → MODEL + +## Logs and verification +- With experimental_ptx enabled, console prints either: + - [TorchCompileSpeed] PTX warmup via triton.ops.matmul, or + - [TorchCompileSpeed] PTX warmup via torch.compile(matmul) +- When speed mode applies, you’ll see inductor config messages. If a knob isn’t available in your PyTorch/Triton build, a warning is printed and safely ignored. + +## Troubleshooting +- Slow first run: includes compilation + autotune. Second and later runs should be much faster. +- triton.ops missing: falls back to torch.compile warmup and still produces PTX/kernel cache. +- Can’t connect to WanVideo Loader: ensure the Settings output type is WANCOMPILEARGS (this node already uses it). +- OOM or pressure: lower dynamo_cache_size_limit, or keep compile_transformer_blocks_only=True. + +## Performance reference (RTX 5090, CUDA 12.0) +- default: first ~14.38s, second ~0ms, about 48,301× +- max-autotune-no-cudagraphs: first ~7.73s, second ~0ms, about 26,603× +- speed: first ~10.52s, second ~0ms, about 84,554× + +Note: Numbers vary by system. experimental_ptx often reduces first‑run cost on no‑CUDA‑Graphs paths and improves second‑run hits. + +## Disclaimer +- experimental_ptx is experimental. Behavior depends on your PyTorch/Triton build; unsupported knobs are ignored. + +## Changelog +- v1.1.0 + - Added experimental_ptx, ptx_fast_math, warmup_runs, ptx_cache_dir + - Added reuse_if_similar and compile_transformer_blocks_only controls + - Kept non‑intrusive design and WanVideo Cython Model Loader compatibility +- v1.0.0 + - Initial release with speed mode and core torch.compile integration diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..af0819b --- /dev/null +++ b/__init__.py @@ -0,0 +1,4 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..60404d5 --- /dev/null +++ b/nodes.py @@ -0,0 +1,214 @@ +import torch +import weakref +COMPILED_FORWARD_CACHE = weakref.WeakKeyDictionary() + + +class TorchCompileSpeedSettings: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "backend": (["inductor","cudagraphs"], {"default": "inductor"}), + "fullgraph": ("BOOLEAN", {"default": False, "tooltip": "Enable full graph mode"}), + "mode": (["default", "max-autotune", "max-autotune-no-cudagraphs", "reduce-overhead", "speed"], {"default": "speed"}), + "dynamic": ("BOOLEAN", {"default": False, "tooltip": "Enable dynamic mode"}), + "dynamo_cache_size_limit": ("INT", {"default": 64, "min": 0, "max": 1024, "step": 1, "tooltip": "torch._dynamo.config.cache_size_limit"}), + "compile_transformer_blocks_only": ("BOOLEAN", {"default": True, "tooltip": "Compile only transformer blocks"}), + "reuse_if_similar": ("BOOLEAN", {"default": True, "tooltip": "Reuse compiled artifacts when similar"}), + "experimental_ptx": ("BOOLEAN", {"default": False, "tooltip": "Enable experimental PTX acceleration"}), + "ptx_fast_math": ("BOOLEAN", {"default": True, "tooltip": "Enable fast math if available"}), + "warmup_runs": ("INT", {"default": 1, "min": 0, "max": 5, "step": 1}), + }, + "optional": { + "ptx_cache_dir": ("STRING", {"default": ""}), + "dynamo_recompile_limit": ("INT", {"default": 128, "min": 0, "max": 1024, "step": 1, "tooltip": "torch._dynamo.config.recompile_limit"}), + }, + } + RETURN_TYPES = ("WANCOMPILEARGS",) + RETURN_NAMES = ("torch_compile_args",) + FUNCTION = "set_args" + CATEGORY = "optimization" + DESCRIPTION = """torch.compile settings for maximum speed optimization. + +Speed Mode Features: +- Uses inductor backend with max-autotune-no-cudagraphs +- Enables dynamic compilation for better cache reuse +- Disables CUDA graphs for flexibility +- Enables all Triton autotune optimizations +- First run: comprehensive autotune (slower) +- Second run: cached execution (extremely fast) + +Author: eddy +""" + + def set_args(self, backend, fullgraph, mode, dynamic, dynamo_cache_size_limit, compile_transformer_blocks_only, reuse_if_similar, experimental_ptx, ptx_fast_math, warmup_runs, ptx_cache_dir="", dynamo_recompile_limit=128): + + if mode == "speed": + backend = "inductor" + fullgraph = False + dynamic = True + effective_mode = "max-autotune-no-cudagraphs" + speed_preset = True + else: + effective_mode = mode + speed_preset = False + + compile_args = { + "backend": backend, + "fullgraph": fullgraph, + "mode": effective_mode, + "dynamic": dynamic, + "dynamo_cache_size_limit": dynamo_cache_size_limit, + "dynamo_recompile_limit": dynamo_recompile_limit, + "compile_transformer_blocks_only": compile_transformer_blocks_only, + "reuse_if_similar": reuse_if_similar, + "experimental_ptx": experimental_ptx, + "ptx_fast_math": ptx_fast_math, + "warmup_runs": warmup_runs, + "ptx_cache_dir": ptx_cache_dir, + "speed_preset": speed_preset, + } + + return (compile_args, ) + +class ApplyTorchCompile: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "compile_args": ("WANCOMPILEARGS",), + }, + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "apply_compile" + CATEGORY = "optimization" + DESCRIPTION = """Apply torch.compile to model with specified settings. + +This node wraps the model's forward pass with torch.compile for acceleration. +Use with TorchCompileSpeedSettings node for optimal configuration. + +Author: eddy +""" + + def apply_compile(self, model, compile_args): + backend = compile_args["backend"] + mode = compile_args["mode"] + dynamic = compile_args["dynamic"] + fullgraph = compile_args["fullgraph"] + + if compile_args.get("speed_preset", False) or compile_args.get("experimental_ptx", False): + try: + from torch._inductor import config as inductor_config + inductor_config.triton.cudagraphs = False + inductor_config.max_autotune = True + inductor_config.max_autotune_pointwise = True + inductor_config.max_autotune_gemm = True + if hasattr(inductor_config, 'max_autotune_conv'): + inductor_config.max_autotune_conv = True + if compile_args.get("experimental_ptx", False): + if hasattr(inductor_config, 'coordinate_descent_tuning'): + inductor_config.coordinate_descent_tuning = True + if hasattr(inductor_config, 'triton') and hasattr(inductor_config.triton, 'use_fast_math'): + inductor_config.triton.use_fast_math = compile_args.get("ptx_fast_math", True) + print("[TorchCompileSpeed] Applied inductor config") + except Exception as e: + print(f"[TorchCompileSpeed] Warning: Could not apply inductor config: {e}") + + try: + torch._dynamo.config.cache_size_limit = compile_args["dynamo_cache_size_limit"] + torch._dynamo.config.recompile_limit = compile_args.get("dynamo_recompile_limit", 128) + except Exception as e: + print(f"[TorchCompileSpeed] Warning: Could not set dynamo config: {e}") + + if compile_args.get("experimental_ptx", False): + try: + import os + if compile_args.get("ptx_cache_dir"): + os.environ["TRITON_CACHE_DIR"] = str(compile_args.get("ptx_cache_dir")) + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + device = next(iter(model.model.parameters())).device if any(True for _ in model.model.parameters()) else torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.bfloat16 if hasattr(torch, 'bfloat16') else torch.float16 + warmed = False + try: + import triton + import triton.ops as tops + if torch.cuda.is_available(): + a = torch.randn(256, 256, device=device, dtype=dtype) + b = torch.randn(256, 256, device=device, dtype=dtype) + for _ in range(int(compile_args.get("warmup_runs", 1))): + _ = tops.matmul(a, b) + torch.cuda.synchronize() + warmed = True + print("[TorchCompileSpeed] PTX warmup via triton.ops.matmul") + except Exception: + pass + if not warmed and torch.cuda.is_available(): + class _Matmul(torch.nn.Module): + def forward(self, x, y): + return torch.matmul(x, y) + mod = _Matmul().to(device=device, dtype=dtype).eval() + cmod = torch.compile(mod, backend=backend, fullgraph=fullgraph, mode=mode, dynamic=dynamic) + a = torch.randn(512, 512, device=device, dtype=dtype) + b = torch.randn(512, 512, device=device, dtype=dtype) + for _ in range(int(compile_args.get("warmup_runs", 1))): + with torch.no_grad(): + _ = cmod(a, b) + torch.cuda.synchronize() + print("[TorchCompileSpeed] PTX warmup via torch.compile(matmul)") + except Exception as e: + print(f"[TorchCompileSpeed] Warning: PTX warmup failed: {e}") + + model_clone = model.clone() + + try: + cache_enabled = compile_args.get("reuse_if_similar", True) + sig = (id(model_clone.model), backend, mode, dynamic, fullgraph) + if cache_enabled: + cached_map = COMPILED_FORWARD_CACHE.get(model_clone.model) + if cached_map is None: + cached_map = {} + COMPILED_FORWARD_CACHE[model_clone.model] = cached_map + if sig in cached_map: + compiled_forward = cached_map[sig] + print("[TorchCompileSpeed] Reused compiled forward from cache") + else: + original_forward = model_clone.model.forward + compiled_forward = torch.compile( + original_forward, + backend=backend, + fullgraph=fullgraph, + mode=mode, + dynamic=dynamic + ) + cached_map[sig] = compiled_forward + print(f"[TorchCompileSpeed] Compiled and cached forward backend={backend}, mode={mode}, dynamic={dynamic}") + else: + original_forward = model_clone.model.forward + compiled_forward = torch.compile( + original_forward, + backend=backend, + fullgraph=fullgraph, + mode=mode, + dynamic=dynamic + ) + model_clone.model.forward = compiled_forward + except Exception as e: + print(f"[TorchCompileSpeed] ERROR: Compilation failed: {e}") + return (model,) + + return (model_clone,) + +NODE_CLASS_MAPPINGS = { + "TorchCompileSpeedSettings": TorchCompileSpeedSettings, + "ApplyTorchCompile": ApplyTorchCompile, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "TorchCompileSpeedSettings": "Torch Compile Speed Settings", + "ApplyTorchCompile": "Apply Torch Compile", +} + diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..b60f5b4 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,15 @@ +[project] +name = "comfyui-torchcompilespeed" +version = "1.0.0" +description = "Standalone torch.compile optimization node for ComfyUI with speed-focused presets" +authors = [{name = "eddy"}] +readme = "README.md" +requires-python = ">=3.8" +license = {text = "MIT"} + +[project.urls] +Repository = "https://github.com/yourusername/ComfyUI-TorchCompileSpeed" + +[tool.setuptools] +packages = ["comfyui_torchcompilespeed"] +