1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
|
# Owner(s): ["module: dynamo"]
import contextlib
import torch._dynamo.test_case
import torch._dynamo.testing
from torch._dynamo.testing import CompileCounter
from torch.backends.cuda import SDPAParams
@contextlib.contextmanager
def allow_in_graph_sdpa_params():
global SDPAParams
try:
old = SDPAParams
SDPAParams = torch._dynamo.allow_in_graph(SDPAParams)
yield
finally:
SDPAParams = old
class TestSDPA(torch._dynamo.test_case.TestCase):
def assert_ref_equals_params(self, actual, expected):
self.assertIs(actual.query, expected.query)
self.assertIs(actual.key, expected.key)
self.assertIs(actual.value, expected.value)
self.assertIs(actual.attn_mask, expected.attn_mask)
def test_returns_SDPAParams(self):
with allow_in_graph_sdpa_params():
counter = CompileCounter()
@torch.compile(fullgraph=True, backend=counter)
def fn(q, k, v, m):
return SDPAParams(q, k, v, m, 0.1, True, False)
q = torch.randn(10)
k = torch.randn(10)
v = torch.randn(10)
m = torch.randn(10)
o = fn(q, k, v, m)
self.assertTrue(isinstance(o, SDPAParams))
self.assert_ref_equals_params(o, SDPAParams(q, k, v, m, 0.1, True, False))
self.assertEqual(counter.frame_count, 1)
def test_graph_break_SDPAParams(self):
with allow_in_graph_sdpa_params():
counter = CompileCounter()
@torch.compile(backend=counter)
def fn(q, k, v, m):
z = SDPAParams(q, k, v, m, 0.1, True, False)
torch._dynamo.graph_break()
return z, q + 1
q = torch.randn(10)
k = torch.randn(10)
v = torch.randn(10)
m = torch.randn(10)
o, _ = fn(q, k, v, m)
self.assertTrue(isinstance(o, SDPAParams))
self.assert_ref_equals_params(o, SDPAParams(q, k, v, m, 0.1, True, False))
self.assertEqual(counter.frame_count, 2)
def test_input_SDPAParams(self):
with allow_in_graph_sdpa_params():
counter = CompileCounter()
@torch.compile(backend=counter)
def fn(sdpap, q):
torch._dynamo.graph_break()
return sdpap, sdpap.query + q
q = torch.randn(10)
k = torch.randn(10)
v = torch.randn(10)
m = torch.randn(10)
s = SDPAParams(q, k, v, m, 0.1, True, False)
o, _ = fn(s, q)
self.assertIs(o, s)
self.assertEqual(counter.frame_count, 1)
def test_intermediate_attr_access_SDPAParams(self):
with allow_in_graph_sdpa_params():
counter = CompileCounter()
@torch.compile(fullgraph=True, backend=counter)
def fn(q, k, v, m):
q += 1
z = SDPAParams(q, k, v, m, 0.1, True, False)
a = z.query
return a + 1, z, q
q = torch.randn(10)
k = torch.randn(10)
v = torch.randn(10)
m = torch.randn(10)
_, o, _ = fn(q, k, v, m)
expected = SDPAParams(q, k, v, m, 0.1, True, False)
self.assert_ref_equals_params(o, expected)
self.assertEqual(counter.frame_count, 1)
if __name__ == "__main__":
from torch._dynamo.test_case import run_tests
run_tests()
|
