Add files via upload
This commit is contained in:
eddy
@@ -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
|
||||
@@ -0,0 +1,4 @@
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
@@ -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"]
|
||||
|
||||
Reference in New Issue
Block a user