-
-
Notifications
You must be signed in to change notification settings - Fork 7.1k
Expand file tree
/
Copy pathsdpa_mask_backend_probe.py
More file actions
98 lines (81 loc) · 3.63 KB
/
Copy pathsdpa_mask_backend_probe.py
File metadata and controls
98 lines (81 loc) · 3.63 KB
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
#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Which SDPA backends tolerate a dense bool attn_mask, and at what cost, at Hunyuan's real
joint shape (B=1, H=16, N=50345, D=128, bf16)? Decides whether nulling the all-True mask is
the real win on the PRODUCTION cuDNN path (not just the native math fallback)."""
import time
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
B, H, N, D = 1, 16, 50345, 128
dev, dt = "cuda:0", torch.bfloat16
def mk():
return torch.randn(B, H, N, D, device = dev, dtype = dt)
def timed(fn, iters = 20):
try:
torch.cuda.synchronize()
for _ in range(3):
fn()
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(iters):
fn()
torch.cuda.synchronize()
return (time.perf_counter() - t0) / iters * 1e3
except torch.OutOfMemoryError:
# OOM on the dense NxN mask is a memory limit, not a backend rejecting it; don't mislabel UNSUPPORTED.
torch.cuda.empty_cache()
return "OOM"
except Exception as e: # noqa: BLE001
return f"UNSUPPORTED ({type(e).__name__})"
def _identity(run, reference):
"""Whether ``run``'s dense output is bitwise-identical to the default dispatch's.
"yes" on exactly one backend names the kernel the dispatcher selected. A backend that cannot
run the dense mask at all reports why instead, so the column never silently reads as a
mismatch when nothing ran."""
if reference is None:
return "n/a"
try:
out = run()
except torch.OutOfMemoryError:
torch.cuda.empty_cache()
return "OOM"
except Exception: # noqa: BLE001
return "unsupported"
return "yes" if torch.equal(reference, out) else f"no ({(reference - out).abs().max():.1e})"
q, k, v = mk(), mk(), mk()
dense = torch.ones(B, 1, N, N, dtype = torch.bool, device = dev)
backends = {
"default(dispatch)": None,
"MATH": [SDPBackend.MATH],
"FLASH": [SDPBackend.FLASH_ATTENTION],
"EFFICIENT": [SDPBackend.EFFICIENT_ATTENTION],
"CUDNN": [SDPBackend.CUDNN_ATTENTION],
}
# The default dispatch's own dense output, so each forced backend can be checked against it.
# Timings alone cannot say WHICH backend the dispatcher picked, because forcing one adds sdpa_kernel overhead and two
# different kernels can land at similar times. Bitwise identity can: the forced backend that reproduces this tensor
# exactly is the one the dispatcher chose.
try:
reference = F.scaled_dot_product_attention(q, k, v, attn_mask = dense)
except Exception: # noqa: BLE001 -- no reference: the identity column just reports n/a
reference = None
print(f"shape B={B} H={H} N={N} D={D} {dt}\n")
print(f"{'backend':<20}{'mask=dense(ms)':>18}{'mask=None(ms)':>18}{'==default(dense)':>19}")
for name, bk in backends.items():
def run_dense():
if bk is None:
return F.scaled_dot_product_attention(q, k, v, attn_mask = dense)
with sdpa_kernel(bk):
return F.scaled_dot_product_attention(q, k, v, attn_mask = dense)
def run_none():
if bk is None:
return F.scaled_dot_product_attention(q, k, v, attn_mask = None)
with sdpa_kernel(bk):
return F.scaled_dot_product_attention(q, k, v, attn_mask = None)
dms = timed(run_dense)
nms = timed(run_none)
d_s = f"{dms:.2f}" if isinstance(dms, float) else dms
n_s = f"{nms:.2f}" if isinstance(nms, float) else nms
print(f"{name:<20}{d_s:>18}{n_s:>18}{_identity(run_dense, reference):>19}")