SDP Backend function fix #161169

ahkush · 2025-08-21T15:52:55Z

The issue cannot be reproduced using the original repro code provided in the issue description.

However, the underlying issue mentioned by the maintainer (missing functions in builder.py and trace_rules.py) was never addressed and can still be reproduced with this test case:

import torch
from torch.nn.attention import _cur_sdpa_kernel_backends

@torch.compile(fullgraph=True)
def test_function_that_triggers_error():
    return _cur_sdpa_kernel_backends()

print("Calling torch.compile function...")
try:
    result = test_function_that_triggers_error()
    print(f"Success: {result}")
except Exception as e:
    print(f"ERROR: {e}")
    print(f"Error type: {type(e)}")

The original repro likely no longer triggers the issue due to code path changes in the SDPA implementation, while the direct call to _cur_sdpa_kernel_backends() exposes the underlying problem where certain torch._C functions returning non-Tensor values aren't properly handled by dynamo tracing.

I have implemented the changes by adding the missing functions to both builder.py and trace_rules.py to properly handle these cases during compilation.

@guilhermeleobas

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

pytorch-bot · 2025-08-21T15:53:02Z

🔗 Helpful Links

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

📄 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

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

✅ No Failures

As of commit fa24351 with merge base 3d40642 ():
💚 Looks good so far! There are no failures yet. 💚

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

guilhermeleobas · 2025-08-21T16:13:39Z

@ahkush, can you add a test case?

ahkush · 2025-08-21T19:53:32Z

@guilhermeleobas, I've added a test case as requested. Please let me know if any other updates are needed. Thanks!

guilhermeleobas · 2025-08-22T15:46:36Z

test/dynamo/test_sdpa.py

+
+    def test_sdpa_c_functions_no_graph_break(self):
+
+        counter = CompileCounter()


Can you add a test using the repro from the issue?

import torch from torch.nn.attention import SDPBackend, sdpa_kernel SDPA_BACKEND_PRIORITY = [ SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.FLASH_ATTENTION, ] @sdpa_kernel(backends=SDPA_BACKEND_PRIORITY, set_priority=True) def scaled_dot_product_attention(q, k, v, *args, **kwargs): return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs) @torch.compile(fullgraph=True) def f(x): return scaled_dot_product_attention(x, x, x) x = torch.rand(128, 64, 64, 256, dtype=torch.float16, device='cuda') f(x)

Sure, will do that. Shall I replace the one I added with the repro in the issue or add it as another test?

Its up to you. If they test the same thing, then you can remove

I've added it as a separate test since they test different scenarios - one tests the direct function call, the other tests the original decorator usage pattern.

guilhermeleobas · 2025-08-27T13:25:48Z

I think this needs rebasing with main. Also, can you run lintrunner? First run lintrunner init and then lintrunner -a to format the changed files.

ahkush · 2025-08-28T15:51:56Z

@guilhermeleobas, Thanks for the feedback! I've rebased with main and ran lintrunner init followed by lintrunner -a to format the files.

guilhermeleobas · 2025-09-02T01:51:45Z

Hi @ahkush, thanks for your contribution. I just approved, but you need a second approval before it can be merged.

StrongerXi

Thanks, but please unlink this PR from #160691. The intent of that issue is to highlight bad error message during context manager tracing, and the example was mainly to illustrate that issue:).

ahkush · 2025-09-15T18:11:34Z

@pytorchbot merge

pytorch-bot · 2025-09-15T18:11:39Z

Pull workflow has not been scheduled for the PR yet. It could be because author doesn't have permissions to run those or skip-checks keywords were added to PR/commits, aborting merge. Please get/give approval for the workflows and/or remove skip ci decorators before next merge attempt. If you think this is a mistake, please contact PyTorch Dev Infra.

guilhermeleobas · 2025-09-15T19:01:42Z

@pytorchbot merge

pytorch-bot · 2025-09-15T19:01:47Z

Pull workflow has not been scheduled for the PR yet. It could be because author doesn't have permissions to run those or skip-checks keywords were added to PR/commits, aborting merge. Please get/give approval for the workflows and/or remove skip ci decorators before next merge attempt. If you think this is a mistake, please contact PyTorch Dev Infra.

guilhermeleobas · 2025-09-19T16:13:02Z

@StrongerXi, can you merge this one?

StrongerXi · 2025-09-19T17:18:17Z

@pytorchbot merge

pytorchmergebot · 2025-09-19T17:20:08Z

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

@guilhermeleobas

The issue cannot be reproduced using the original repro code provided in the issue description. However, the underlying issue mentioned by the maintainer (missing functions in `builder.py` and `trace_rules.py`) was never addressed and can still be reproduced with this test case: ```python import torch from torch.nn.attention import _cur_sdpa_kernel_backends @torch.compile(fullgraph=True) def test_function_that_triggers_error(): return _cur_sdpa_kernel_backends() print("Calling torch.compile function...") try: result = test_function_that_triggers_error() print(f"Success: {result}") except Exception as e: print(f"ERROR: {e}") print(f"Error type: {type(e)}") ``` The original repro likely no longer triggers the issue due to code path changes in the SDPA implementation, while the direct call to `_cur_sdpa_kernel_backends()` exposes the underlying problem where certain torch._C functions returning non-Tensor values aren't properly handled by dynamo tracing. I have implemented the changes by adding the missing functions to both `builder.py` and `trace_rules.py` to properly handle these cases during compilation. @guilhermeleobas Pull Request resolved: pytorch#161169 Approved by: https://github.com/guilhermeleobas, https://github.com/StrongerXi

@guilhermeleobas

The issue cannot be reproduced using the original repro code provided in the issue description. However, the underlying issue mentioned by the maintainer (missing functions in `builder.py` and `trace_rules.py`) was never addressed and can still be reproduced with this test case: ```python import torch from torch.nn.attention import _cur_sdpa_kernel_backends @torch.compile(fullgraph=True) def test_function_that_triggers_error(): return _cur_sdpa_kernel_backends() print("Calling torch.compile function...") try: result = test_function_that_triggers_error() print(f"Success: {result}") except Exception as e: print(f"ERROR: {e}") print(f"Error type: {type(e)}") ``` The original repro likely no longer triggers the issue due to code path changes in the SDPA implementation, while the direct call to `_cur_sdpa_kernel_backends()` exposes the underlying problem where certain torch._C functions returning non-Tensor values aren't properly handled by dynamo tracing. I have implemented the changes by adding the missing functions to both `builder.py` and `trace_rules.py` to properly handle these cases during compilation. @guilhermeleobas Pull Request resolved: pytorch#161169 Approved by: https://github.com/guilhermeleobas, https://github.com/StrongerXi

@guilhermeleobas

The issue cannot be reproduced using the original repro code provided in the issue description. However, the underlying issue mentioned by the maintainer (missing functions in `builder.py` and `trace_rules.py`) was never addressed and can still be reproduced with this test case: ```python import torch from torch.nn.attention import _cur_sdpa_kernel_backends @torch.compile(fullgraph=True) def test_function_that_triggers_error(): return _cur_sdpa_kernel_backends() print("Calling torch.compile function...") try: result = test_function_that_triggers_error() print(f"Success: {result}") except Exception as e: print(f"ERROR: {e}") print(f"Error type: {type(e)}") ``` The original repro likely no longer triggers the issue due to code path changes in the SDPA implementation, while the direct call to `_cur_sdpa_kernel_backends()` exposes the underlying problem where certain torch._C functions returning non-Tensor values aren't properly handled by dynamo tracing. I have implemented the changes by adding the missing functions to both `builder.py` and `trace_rules.py` to properly handle these cases during compilation. @guilhermeleobas Pull Request resolved: pytorch#161169 Approved by: https://github.com/guilhermeleobas, https://github.com/StrongerXi

pytorch-bot bot added the module: dynamo label Aug 21, 2025

pytorchbot added the open source label Aug 21, 2025

guilhermeleobas reviewed Aug 22, 2025

View reviewed changes

guilhermeleobas added the release notes: dynamo label Aug 22, 2025

ahkush added 6 commits August 27, 2025 14:29

160691: SDP Backend function fix

b36be5e

Added test case for 160691

3267bb9

changed test case location

e9ed3d9

removed extra space and comments

80bb869

added sdpa kernal decorator test

4107a5c

lintrunner formatting applied

fa24351

ahkush force-pushed the fix-dynamo-sdpa-backend-functions branch from 69d1c26 to fa24351 Compare August 27, 2025 18:48

guilhermeleobas approved these changes Sep 2, 2025

View reviewed changes

StrongerXi approved these changes Sep 5, 2025

View reviewed changes

pytorch-bot bot added the ciflow/inductor label Sep 5, 2025

ahkush changed the title ~~160691: SDP Backend function fix~~ SDP Backend function fix Sep 8, 2025

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

pytorchmergebot added the merging label Sep 19, 2025

pytorchmergebot added the Merged label Sep 19, 2025

pytorchmergebot closed this in ba3c2c8 Sep 19, 2025

pytorchmergebot removed the merging label Sep 19, 2025


		def test_sdpa_c_functions_no_graph_break(self):

		counter = CompileCounter()

SDP Backend function fix #161169

SDP Backend function fix #161169

Uh oh!

Conversation

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

Uh oh!

Uh oh!

pytorch-bot bot commented Aug 21, 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/161169

✅ No Failures

Uh oh!

guilhermeleobas commented Aug 21, 2025

Uh oh!

ahkush commented Aug 21, 2025

Uh oh!

guilhermeleobas Aug 22, 2025

Choose a reason for hiding this comment

Uh oh!

ahkush Aug 22, 2025

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ahkush Aug 22, 2025

Choose a reason for hiding this comment

Uh oh!

guilhermeleobas commented Aug 27, 2025

Uh oh!

ahkush commented Aug 28, 2025

Uh oh!

guilhermeleobas commented Sep 2, 2025

Uh oh!

StrongerXi left a comment

Choose a reason for hiding this comment

Uh oh!

ahkush commented Sep 15, 2025

Uh oh!

pytorch-bot bot commented Sep 15, 2025

Uh oh!

guilhermeleobas commented Sep 15, 2025

Uh oh!

pytorch-bot bot commented Sep 15, 2025

Uh oh!

guilhermeleobas commented Sep 19, 2025

Uh oh!

StrongerXi commented Sep 19, 2025

Uh oh!

pytorchmergebot commented Sep 19, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ahkush commented Aug 21, 2025 •

edited

Loading

pytorch-bot bot commented Aug 21, 2025 •

edited

Loading

guilhermeleobas Aug 22, 2025 •

edited

Loading