[dynamo] Fix graph break on calling functions decorated with special context manager #160703

StrongerXi · 2025-08-15T01:55:05Z

Stack from ghstack (oldest at bottom):

-> [dynamo] Fix graph break on calling functions decorated with special context manager #160703
[dynamo] Fix crash and silent incorrectness issues in attention.sdpa_kernel calls with kwargs #160684

As title. This is a follow-up of the previous patch, with the goal of
supporting a new pattern that showed up in ComfyUI:
https://github.com/comfyanonymous/ComfyUI/blob/644b23ac0b92442b475e44397c62aa8de929d546/comfy/ops.py#L44

Effectively, the semantics of calling a function decorated with a
context manager is:

@ctx_manager(args)
def f(x):
    ...

f(x)
# ----->
with ctx_manager(args):
    f.__wrapped__(x)

Yes, a fresh context manager instance per invokation, see CPython source code:
https://github.com/python/cpython/blob/3.12/Lib/contextlib.py#L119-L122

So Dynamo already

knows how to handle the with ctx_manager(args) syntax, and has
special handling for a few torch native context managers, like
sdpa_kernel in this patch.
can trace through a good chunk (at least the ones that matter in this
case) of contextlib.

This patch just let Dynamo trace a bit more into contextlib, and then
keep the torch-native special cases by moving their handling a bit down
the stack, so that no additional logic is introduced -- it's only
refactored.

This also allows us to get rid of some _sdpa_kernel_variadic special
handling, since now we will trace through its code, and it boils down to
sdpa_kernel anyways.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames @Lucaskabela @mlazos

[ghstack-poisoned]

pytorch-bot · 2025-08-15T01:55:08Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/160703

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 98f308a with merge base ee1b041 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

…context manager As title. This is a follow-up of the previous patch, with the goal of supporting a new pattern that showed up in ComfyUI: https://github.com/comfyanonymous/ComfyUI/blob/644b23ac0b92442b475e44397c62aa8de929d546/comfy/ops.py#L44 Effectively, the semantics of calling a function decorated with a context manager is: ```python ctx_manager(args) def f(x): ... f(x) with ctx_manager(args): f(x) ``` Yes, a fresh context manager instance per invokation, see CPython source code: https://github.com/python/cpython/blob/3.12/Lib/contextlib.py#L119-L122 So Dynamo already 1. knows how to handle the `with ctx_manager(args)` syntax, and has special handling for a few torch native context managers, like `sdpa_kernel` in this patch. 2. can trace through a good chunk (at least the ones that matter in this case) of contextlib. This patch just let Dynamo trace a bit more into contextlib, and then keep the torch-native special cases by moving their handling a bit down the stack, so that no additional logic is introduced -- it's only refactored. This also allows us to get rid of some `_sdpa_kernel_variadic` special handling, since now we will trace through its code, and it boils down to `sdpa_kernel` anyways. ghstack-source-id: 130f92b Pull-Request-resolved: #160703

[ghstack-poisoned]

…context manager As title. This is a follow-up of the previous patch, with the goal of supporting a new pattern that showed up in ComfyUI: https://github.com/comfyanonymous/ComfyUI/blob/644b23ac0b92442b475e44397c62aa8de929d546/comfy/ops.py#L44 Effectively, the semantics of calling a function decorated with a context manager is: ```python ctx_manager(args) def f(x): ... f(x) with ctx_manager(args): f(x) ``` Yes, a fresh context manager instance per invokation, see CPython source code: https://github.com/python/cpython/blob/3.12/Lib/contextlib.py#L119-L122 So Dynamo already 1. knows how to handle the `with ctx_manager(args)` syntax, and has special handling for a few torch native context managers, like `sdpa_kernel` in this patch. 2. can trace through a good chunk (at least the ones that matter in this case) of contextlib. This patch just let Dynamo trace a bit more into contextlib, and then keep the torch-native special cases by moving their handling a bit down the stack, so that no additional logic is introduced -- it's only refactored. This also allows us to get rid of some `_sdpa_kernel_variadic` special handling, since now we will trace through its code, and it boils down to `sdpa_kernel` anyways. ghstack-source-id: 5be04ef Pull-Request-resolved: #160703

guilhermeleobas · 2025-08-15T18:30:07Z

torch/_dynamo/variables/user_defined.py

+                if isinstance(args[0], UserFunctionVariable):
+                    fn = args[0].fn
+                    fn_source = args[0].source
+                    # At this point `fn` is the "unwrapped" version of the
+                    # context-managerified function.
+                    if fn is torch.nn.attention.sdpa_kernel.__wrapped__:
+                        args, kwargs = args[1], args[2]
+                        name_to_arg_map = bind_args_cached(
+                            fn, tx, fn_source, args.items, kwargs.items
+                        )
+                        backends = name_to_arg_map["backends"].as_python_constant()
+                        set_priority = name_to_arg_map[
+                            "set_priority"
+                        ].as_python_constant()
+                        return SDPAKernelVariable.create(tx, backends, set_priority)


There's a TorchCtxManagerClassVariable that handles special cases of context managers. Could you try to implement this context manager using it?

pytorch/torch/_dynamo/variables/torch.py

Lines 13 to 14 in 30d2f98

TorchCtxManagerClassVariable: Handles torch context managers like torch.no_grad(), autocast, etc.

Provides implementations for entering/exiting these contexts during tracing.

I was basically moving the sdpa_kernels special casing out of it, since it's no longer a "class value" anymore after this patch (we'll trace into the actual wrapped sdpa_kernels as a function value), and I was intercepting the actual context manager construction right here.

Regardless, keeping the logic in TorchCtxManagerClassVariable might also make things more organized anyways, despite the misnomer. Ideally I'm hoping we just trace through the python code as much as possible eventually, without special casing, to avoid issues like #160684.

[ghstack-poisoned]

…context manager As title. This is a follow-up of the previous patch, with the goal of supporting a new pattern that showed up in ComfyUI: https://github.com/comfyanonymous/ComfyUI/blob/644b23ac0b92442b475e44397c62aa8de929d546/comfy/ops.py#L44 Effectively, the semantics of calling a function decorated with a context manager is: ```python ctx_manager(args) def f(x): ... f(x) #-------> with ctx_manager(args): f(x) ``` Yes, a fresh context manager instance per invokation, see CPython source code: https://github.com/python/cpython/blob/3.12/Lib/contextlib.py#L119-L122 So Dynamo already 1. knows how to handle the `with ctx_manager(args)` syntax, and has special handling for a few torch native context managers, like `sdpa_kernel` in this patch. 2. can trace through a good chunk (at least the ones that matter in this case) of contextlib. This patch just let Dynamo trace a bit more into contextlib, and then keep the torch-native special cases by moving their handling a bit down the stack, so that no additional logic is introduced -- it's only refactored. This also allows us to get rid of some `_sdpa_kernel_variadic` special handling, since now we will trace through its code, and it boils down to `sdpa_kernel` anyways. ghstack-source-id: 364a63b Pull-Request-resolved: #160703

StrongerXi · 2025-08-18T17:42:18Z

@pytorchbot merge

pytorchmergebot · 2025-08-18T17:44:03Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

…context manager (pytorch#160703) As title. This is a follow-up of the previous patch, with the goal of supporting a new pattern that showed up in ComfyUI: https://github.com/comfyanonymous/ComfyUI/blob/644b23ac0b92442b475e44397c62aa8de929d546/comfy/ops.py#L44 Effectively, the semantics of calling a function decorated with a context manager is: ```python @ctx_manager(args) def f(x): ... f(x) # -----> with ctx_manager(args): f.__wrapped__(x) ``` Yes, a fresh context manager instance per invokation, see CPython source code: https://github.com/python/cpython/blob/3.12/Lib/contextlib.py#L119-L122 So Dynamo already 1. knows how to handle the `with ctx_manager(args)` syntax, and has special handling for a few torch native context managers, like `sdpa_kernel` in this patch. 2. can trace through a good chunk (at least the ones that matter in this case) of contextlib. This patch just let Dynamo trace a bit more into contextlib, and then keep the torch-native special cases by moving their handling a bit down the stack, so that no additional logic is introduced -- it's only refactored. This also allows us to get rid of some `_sdpa_kernel_variadic` special handling, since now we will trace through its code, and it boils down to `sdpa_kernel` anyways. Pull Request resolved: pytorch#160703 Approved by: https://github.com/guilhermeleobas, https://github.com/mlazos ghstack dependencies: pytorch#160684

Update

555dba1

[ghstack-poisoned]

pytorch-bot bot added ciflow/inductor module: dynamo labels Aug 15, 2025

This was referenced Aug 15, 2025

[dynamo] Fix crash and silent incorrectness issues in attention.sdpa_kernel calls with kwargs #160684

Closed

Avoid torch compile graphbreak for older pytorch versions comfyanonymous/ComfyUI#9344

Merged

StrongerXi added the release notes: dynamo label Aug 15, 2025

StrongerXi requested review from guilhermeleobas and zou3519 August 15, 2025 16:43

Update

ae6ed3b

[ghstack-poisoned]

guilhermeleobas reviewed Aug 15, 2025

View reviewed changes

Update

98f308a

[ghstack-poisoned]

guilhermeleobas approved these changes Aug 15, 2025

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Aug 18, 2025

pytorchmergebot added the merging label Aug 18, 2025

mlazos approved these changes Aug 18, 2025

View reviewed changes

pytorchmergebot added the Merged label Aug 18, 2025

pytorchmergebot closed this in a1a555e Aug 18, 2025

pytorchmergebot removed the merging label Aug 18, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[dynamo] Fix graph break on calling functions decorated with special context manager #160703

[dynamo] Fix graph break on calling functions decorated with special context manager #160703

Uh oh!

StrongerXi commented Aug 15, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Aug 15, 2025 •

edited

Loading

Uh oh!

guilhermeleobas Aug 15, 2025

Uh oh!

StrongerXi Aug 15, 2025

Uh oh!

StrongerXi commented Aug 18, 2025

Uh oh!

pytorchmergebot commented Aug 18, 2025

Uh oh!

Uh oh!

	TorchCtxManagerClassVariable: Handles torch context managers like torch.no_grad(), autocast, etc.
	Provides implementations for entering/exiting these contexts during tracing.

[dynamo] Fix graph break on calling functions decorated with special context manager #160703

[dynamo] Fix graph break on calling functions decorated with special context manager #160703

Uh oh!

Conversation

StrongerXi commented Aug 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Aug 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/160703

✅ No Failures

Uh oh!

guilhermeleobas Aug 15, 2025

Choose a reason for hiding this comment

Uh oh!

StrongerXi Aug 15, 2025

Choose a reason for hiding this comment

Uh oh!

StrongerXi commented Aug 18, 2025

Uh oh!

pytorchmergebot commented Aug 18, 2025

Merge started

Uh oh!

Uh oh!

StrongerXi commented Aug 15, 2025 •

edited

Loading

pytorch-bot bot commented Aug 15, 2025 •

edited

Loading