[SDPA] Allow user-specified priority order with context manager #140467

eqy · 2024-11-12T23:16:25Z

TODO: docs changes?
For better debuggability of issues like #139298

Better testing, current sketch:

import torch
from torch.nn.functional import scaled_dot_product_attention
from torch.nn.attention import SDPBackend, sdpa_kernel

q = torch.randn(64, 1024, 8, 64, dtype=torch.half, device='cuda')
print(torch._C._get_sdp_priority_order())

orders = [[SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION],
          [SDPBackend.MATH, SDPBackend.CUDNN_ATTENTION, SDPBackend.EFFICIENT_ATTENTION],
          [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH]]
import time
times = list()
for order in orders:
    print(order)
    with sdpa_kernel(order, set_priority=True):
        scaled_dot_product_attention(q, q, q)
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with sdpa_kernel(order, set_priority=True):
        scaled_dot_product_attention(q, q, q)
    torch.cuda.synchronize()
    t1 = time.perf_counter()
    times.append(t1 - t0)
print(times)
assert times[0] < times[1]
assert times[0] > times[2]
assert times[1] > times[2]
print(torch._C._get_sdp_priority_order())

cc @csarofeen @ptrblck @xwang233 @drisspg @mikaylagawarecki

pytorch-bot · 2024-11-12T23:16:29Z

🔗 Helpful Links

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

📄 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 da74e3d with merge base d0fd42e ():
💚 Looks good so far! There are no failures yet. 💚

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

drisspg · 2024-11-13T03:03:15Z

I agree w/ the premise here, will do a review tomorrow. I dont know if I like jamming the list into sdpa_kernel.. but let me sleep on it

Skylion007 · 2024-11-13T16:05:06Z

It would make sense for SDPA kernel to me. The old API is a little bit odd because the order of the backends has no effect on anything, wasting useful information. Really should take a set previously if the order did not matter.

Skylion007 · 2024-11-13T16:08:35Z

What if the user never wanted to fallback to MATH backend? For example:
order=[SDPBackend.EFFICIENT_ATTENTION, SDPBackend.CUDNN_ATTENTION,]]
seems like it should be valid / supported.

eqy · 2024-11-13T18:55:21Z

What if the user never wanted to fallback to MATH backend? For example: order=[SDPBackend.EFFICIENT_ATTENTION, SDPBackend.CUDNN_ATTENTION,]] seems like it should be valid / supported.

If a backend is not in the list it will not be selected as a fallback, this sets the priority order for all backends but preserves the current behavior of only enabling the ones in the list.

Skylion007 · 2024-11-14T18:16:21Z

What if the user never wanted to fallback to MATH backend? For example: order=[SDPBackend.EFFICIENT_ATTENTION, SDPBackend.CUDNN_ATTENTION,]] seems like it should be valid / supported.

If a backend is not in the list it will not be selected as a fallback, this sets the priority order for all backends but preserves the current behavior of only enabling the ones in the list.

Yeah, but is there case where you want priority, but want it to fallback to kernels outside of the list?

eqy · 2024-11-14T21:40:31Z

Not really, this is just a housekeeping thing as priority is currently a fixed-size array so there will be "something" in the priority list even if it's nonsensical

drisspg · 2024-11-18T22:57:14Z

My magic wand API is only accept the list for this context manager and have it specify the order.

drisspg · 2024-11-18T22:59:16Z

aten/src/ATen/native/transformers/cuda/sdp_utils.cpp

I doubt it would be perceptible, but I wonder if this adds any slowdown.

drisspg · 2024-11-18T22:59:43Z

aten/src/ATen/SDPBackend.h

This technically changes the namespace of the symbol, right?

drisspg · 2024-11-18T23:00:50Z

torch/nn/attention/__init__.py

This is a good middle ground though of the magic wand API.

eqy · 2024-11-25T20:54:22Z

@pytorchmergebot rebase

pytorchmergebot · 2024-11-25T20:55:58Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

pytorchmergebot · 2024-11-25T20:56:04Z

Successfully rebased prioritysdpa2 onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout prioritysdpa2 && git pull --rebase)

eqy · 2024-12-05T22:59:07Z

@pytorchmergebot merge

pytorchmergebot · 2024-12-05T23:00:58Z

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

pytorchmergebot · 2024-12-06T03:03:56Z

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / linux-focal-cuda12.4-py3.10-gcc9-sm86 / build

Details for Dev Infra team

Raised by workflow job

eqy · 2024-12-06T04:38:55Z

@pytorchmergebot merge

pytorchmergebot · 2024-12-06T04:41:34Z

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

…rch#140467) TODO: docs changes? For better debuggability of issues like pytorch#139298 Better testing, current sketch: ``` Python import torch from torch.nn.functional import scaled_dot_product_attention from torch.nn.attention import SDPBackend, sdpa_kernel q = torch.randn(64, 1024, 8, 64, dtype=torch.half, device='cuda') print(torch._C._get_sdp_priority_order()) orders = [[SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION], [SDPBackend.MATH, SDPBackend.CUDNN_ATTENTION, SDPBackend.EFFICIENT_ATTENTION], [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH]] import time times = list() for order in orders: print(order) with sdpa_kernel(order, set_priority=True): scaled_dot_product_attention(q, q, q) torch.cuda.synchronize() t0 = time.perf_counter() with sdpa_kernel(order, set_priority=True): scaled_dot_product_attention(q, q, q) torch.cuda.synchronize() t1 = time.perf_counter() times.append(t1 - t0) print(times) assert times[0] < times[1] assert times[0] > times[2] assert times[1] > times[2] print(torch._C._get_sdp_priority_order()) ``` Pull Request resolved: pytorch#140467 Approved by: https://github.com/drisspg

eqy added module: cudnn Related to torch.backends.cudnn, and CuDNN support open source topic: improvements topic category topic: not user facing topic category module: multi-headed-attention labels Nov 12, 2024

eqy requested a review from drisspg November 12, 2024 23:16

eqy requested review from albanD, jbschlosser and mikaylagawarecki as code owners November 12, 2024 23:16

soulitzer added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Nov 13, 2024

drisspg reviewed Nov 18, 2024

View reviewed changes

aten/src/ATen/SDPBackend.h Outdated

Copy link

Contributor

drisspg Nov 18, 2024

Choose a reason for hiding this comment

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

This technically changes the namespace of the symbol, right?

drisspg reviewed Nov 18, 2024

View reviewed changes

torch/nn/attention/__init__.py Outdated

Copy link

Contributor

drisspg Nov 18, 2024

Choose a reason for hiding this comment

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

This is a good middle ground though of the magic wand API.

eqy added 3 commits November 25, 2024 20:56

check in initial priority order impl

1825196

check in missing file

ac9f796

check in

71420c1

pytorchmergebot force-pushed the prioritysdpa2 branch from 4c59817 to 71420c1 Compare November 25, 2024 20:56

eqy added 2 commits November 25, 2024 21:52

update docstr

fd6e111

lint

64dbbb7

eqy added 5 commits November 25, 2024 23:43

lint

a47cef6

check in

ef06c64

more lint

8748540

lint

f8aea28

lint

da74e3d

drisspg added module: sdpa All things related to torch.nn.functional.scaled_dot_product_attentiion and removed module: multi-headed-attention labels Nov 27, 2024

drisspg approved these changes Dec 5, 2024

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Dec 5, 2024

pytorchmergebot added the merging label Dec 5, 2024

albanD removed their request for review December 5, 2024 23:06

pytorchmergebot removed the merging label Dec 6, 2024

pytorchmergebot added the merging label Dec 6, 2024

pytorchmergebot added the Merged label Dec 6, 2024

pytorchmergebot closed this in 8fc6d3a Dec 6, 2024

pytorchmergebot removed the merging label Dec 6, 2024

eqy mentioned this pull request Feb 24, 2025

[SDPA] Respect sdpa_kernel's priority_order setting in torch.compile #147768

Closed

[SDPA] Allow user-specified priority order with context manager #140467

[SDPA] Allow user-specified priority order with context manager #140467

Uh oh!

Conversation

eqy commented Nov 12, 2024 • edited by drisspg Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Nov 12, 2024 • 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/140467

✅ No Failures

Uh oh!

drisspg commented Nov 13, 2024

Uh oh!

Skylion007 commented Nov 13, 2024

Uh oh!

Skylion007 commented Nov 13, 2024

Uh oh!

eqy commented Nov 13, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Skylion007 commented Nov 14, 2024

Uh oh!

eqy commented Nov 14, 2024

Uh oh!

drisspg commented Nov 18, 2024

Uh oh!

drisspg Nov 18, 2024

Choose a reason for hiding this comment

Uh oh!

drisspg Nov 18, 2024

Choose a reason for hiding this comment

Uh oh!

drisspg Nov 18, 2024

Choose a reason for hiding this comment

Uh oh!

eqy commented Nov 25, 2024

Uh oh!

pytorchmergebot commented Nov 25, 2024

Uh oh!

pytorchmergebot commented Nov 25, 2024

Uh oh!

eqy commented Dec 5, 2024

Uh oh!

pytorchmergebot commented Dec 5, 2024

Merge started

Uh oh!

pytorchmergebot commented Dec 6, 2024

Merge failed

Uh oh!

eqy commented Dec 6, 2024

Uh oh!

pytorchmergebot commented Dec 6, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

eqy commented Nov 12, 2024 •

edited by drisspg

Loading

pytorch-bot bot commented Nov 12, 2024 •

edited

Loading

eqy commented Nov 13, 2024 •

edited

Loading