Fix Dynamo `lru_cache` warnings during `torch.compile` by jiqing-feng · Pull Request #13384 · huggingface/diffusers · GitHub
Skip to content

Fix Dynamo lru_cache warnings during torch.compile - #13384

Merged
sayakpaul merged 4 commits into
huggingface:mainfrom
jiqing-feng:compile
Apr 3, 2026
Merged

Fix Dynamo lru_cache warnings during torch.compile#13384
sayakpaul merged 4 commits into
huggingface:mainfrom
jiqing-feng:compile

Conversation

@jiqing-feng

@jiqing-feng jiqing-feng commented Apr 2, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes Dynamo lru_cache warnings when using torch.compile on diffusion pipelines. Two changes:

  1. attention_dispatch.py: dispatch_attention_fn calls is_torch_version(">=", "2.5.0") at runtime, which is @lru_cache-wrapped. Replace with the existing module-level constant _CAN_USE_FLEX_ATTN so Dynamo never traces into it.

  2. torch_utils.py: lru_cache_unless_export only bypasses lru_cache during torch.export (is_exporting). Add is_compiling check so torch.compile also bypasses the cache wrapper.

Reproduce

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-4B", torch_dtype=torch.bfloat16).to("cpu")
pipe.transformer = torch.compile(pipe.transformer, backend="inductor")
pipe(prompt="a cat", height=256, width=256, num_inference_steps=1, generator=torch.Generator().manual_seed(0))
# Before: UserWarning about lru_cache in attention_dispatch.py / torch_utils.py
# After: no warning

out before fix:

/opt/venv/lib/python3.12/site-packages/torch/_dynamo/variables/functions.py:2435: UserWarning: Dynamo detected a call to a `functools.lru_cache`-wrapped function at 'attention_dispatch.py:426'. Dynamo ignores the cache wrapper and directly traces the wrapped function. Silent incorrectness is only a *potential* risk, not something we have observed. Enable TORCH_LOGS=+dynamo for a DEBUG stack trace.

This call originates from:
  File "/home/jiqing/diffusers/src/diffusers/models/attention_dispatch.py", line 426, in dispatch_attention_fn
    if is_torch_version(">=", "2.5.0"):

  torch._dynamo.utils.warn_once(msg)

Flux2PipelineOutput(images=[<PIL.Image.Image image mode=RGB size=256x256 at 0x7A540282FDA0>])

out after fix:

Flux2PipelineOutput(images=[<PIL.Image.Image image mode=RGB size=256x256 at 0x7A540282FDA0>])

@jiqing-feng
fix compile issue
13bcfe5 
Signed-off-by: jiqing-feng <jiqing.feng@intel.com>
@jiqing-feng
compile friendly
6a94fc6 
Signed-off-by: jiqing-feng <jiqing.feng@intel.com>
@jiqing-feng

jiqing-feng commented Apr 2, 2026

Copy link
Copy Markdown
Contributor Author

sayakpaul
sayakpaul reviewed Apr 2, 2026
View reviewed changes
Comment thread src/diffusers/models/attention_dispatch.py
"_parallel_config": parallel_config,
}
if is_torch_version(">=", "2.5.0"):
if _CAN_USE_FLEX_ATTN:

@sayakpaul sayakpaul Apr 2, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this a safe replacement? If so, could you elaborate further?

@jiqing-feng jiqing-feng Apr 2, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, exactly the same. https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_dispatch.py#L62-L72

@jiqing-feng jiqing-feng Apr 2, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added comments for it.

sayakpaul
sayakpaul reviewed Apr 2, 2026
View reviewed changes

@sayakpaul sayakpaul left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR. Left one comment.

@jiqing-feng
add comments
a8366a0 
Signed-off-by: jiqing-feng <jiqing.feng@intel.com>
@HuggingFaceDocBuilderDev

HuggingFaceDocBuilderDev commented Apr 2, 2026

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@jiqing-feng

jiqing-feng commented Apr 3, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

sayakpaul
sayakpaul approved these changes Apr 3, 2026
View reviewed changes

@sayakpaul sayakpaul left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR! Failing test is unrelated.

@sayakpaul
sayakpaul merged commit a05c8e9 into huggingface:main Apr 3, 2026
10 of 11 checks passed
terarachang pushed a commit to terarachang/diffusers that referenced this pull request Apr 30, 2026
@jiqing-feng @sayakpaul
Fix Dynamo lru_cache warnings during torch.compile (huggingface#1…
be1fda4 
…3384)

* fix compile issue

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

* compile friendly

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

* add comments

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

---------

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
@jiqing-feng
jiqing-feng deleted the compile branch May 27, 2026 06:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@sayakpaul sayakpaul sayakpaul approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@jiqing-feng @HuggingFaceDocBuilderDev @sayakpaul