Nodes/KJNodes for ComfyUI/TorchCompileModelAdvanced
ComfyUI Node Runs on cloud

TorchCompileModelAdvanced

The fine-grained torch.compile knob for video models

By kijai·Created 3 years ago·Updated a day ago· 3,011
TorchCompileModelAdvanced
  • model
  • MODEL
backendinductor
fullgraphfalse
modedefault
dynamicfalse
compile_transformer_blocks_onlytrue
dynamo_cache_size_limit64
debug_compile_keysfalse
disable_dynamic_vramfalse

If you run Wan, LTX, or any of the big diffusion-transformer video models, you've probably seen torch.compile show up in the "how do I make this faster" threads next to SageAttention and TeaCache. This node is kijai's version of it with the lid off - the same JIT compilation as the vanilla TorchCompileModel node, but every setting exposed so you can tune it instead of taking the defaults on faith.

What it actually does

torch.compile traces your model's forward pass and compiles it into fused, optimized kernels the first time it runs. After that, every step is faster because you're no longer paying Python's per-op overhead. The community summary is blunt and accurate: "JIT compilation. Startup overhead but per-frame speed improvement." That's the whole trade. The first generation after you load the graph is slow - sometimes painfully so - because it's compiling. Every generation after that in the same session is quicker. On a 4090 running Wan with SageAttention plus TorchCompile plus TeaCache, people report 81 frames at 30 steps in around five and a half minutes, versus unusable without the stack. TorchCompile is one leg of that stool.

You wire it in like any model patch: MODEL in, patched MODEL out, and you drop it between your model loader and your sampler.

The inputs that matter

Most of the eight settings you can leave alone. The three worth understanding:

  • mode - default, max-autotune, max-autotune-no-cudagraphs, or reduce-overhead. default compiles fast and gives a solid speedup. max-autotune searches harder for the best kernels, so it's slower to compile but can run faster once it's done - worth it only if you're doing a long batch where the compile cost amortizes. For one-off gens, default is the sane pick.
  • compile_transformer_blocks_only (default true) - and leave it true. It compiles just the repeated transformer blocks instead of the whole graph, which is where nearly all the compute lives. The big win is that it recompiles far less often when things like resolution change, because the block shape stays stable even when the outer graph doesn't.
  • backend - inductor (the default, the one you want) or cudagraphs. Inductor is the standard PyTorch compiler backend. Don't switch unless you have a specific reason.

The rest - dynamic, dynamo_cache_size_limit, fullgraph, debug_compile_keys, and the optional disable_dynamic_vram - are escape hatches for when compilation fights you, covered below.

How to install it

It ships in kijai's KJNodes pack.

  • ComfyUI Manager - search KJNodes for ComfyUI, install, restart.
  • Manual - cd ComfyUI/custom_nodes && git clone https://github.com/kijai/ComfyUI-KJNodes, then pip install -r ComfyUI-KJNodes/requirements.txt, and restart. Portable-install users run the python_embeded\python.exe -m pip install -r ... variant.

There's no model to download. But torch.compile itself leans on Triton to build kernels, and on Windows that's the classic pain point - the same Triton dependency that makes SageAttention fiddly. On Linux and on cloud boxes it generally Just Works.

Common issues & troubleshooting

The first run is really slow, then it's fine. That's compilation, working as designed. Don't kill it. If you cancel mid-compile you just pay it again next time.

It recompiles every time you change resolution or frame count. torch.compile compiles per input shape, so a new shape triggers a fresh compile. Keeping compile_transformer_blocks_only on cuts most of this. If you're sweeping many shapes and blowing past the recompile budget, that's what dynamo_cache_size_limit (default 64) is for - raise it so the cache holds more compiled variants instead of evicting and recompiling. Setting dynamic to true can also let one compile cover a range of shapes.

fullgraph errors out. fullgraph=true demands the entire forward compile with zero graph breaks, and diffusion models frequently have a construct that won't. Leave it false unless you're deliberately hunting graph breaks with debug_compile_keys.

It's slower, not faster. If you're doing a single short gen, the compile overhead can eat the whole speedup - TorchCompile pays off across many steps or many gens, not on a one-shot. And if you stacked max-autotune for a single image, you paid the long compile for nothing. Use default, or don't compile at all for quick experiments.

CategoryKJNodes/torchcompile

Inputs (9)

NameTypeDefaultDescription
modelMODEL
backendCOMBOinductor2 options: inductor, cudagraphs
fullgraphBOOLEANfalseEnable full graph mode
modeCOMBOdefault4 options: default, max-autotune, max-autotune-no-cudagraphs, reduce-overhead
dynamicCOMBOfalseUse dynamic shape tracing.
compile_transformer_blocks_onlyBOOLEANtrueCompile only transformer blocks, faster compile and less error prone
dynamo_cache_size_limitINT640–1024torch._dynamo.config.cache_size_limit
debug_compile_keysBOOLEANfalsePrint the compile keys used for torch.compile
disable_dynamic_vramoptBOOLEANfalseDisable dynamic VRAM feature as it can cause issues with compile

Outputs (1)

NameTypeDescription
MODELMODEL