Bug? scaled_dot_product_attention slower than manual multiplication?

JosAr: <p>Hi,</p>
<p>I have observed that in my script the usage of <code>scaled_dot_product_attention</code> slows my code down a bit.</p>
<p>Then I have tested if the implementation of scaled_dot_product_attention is actually more efficient in my case, turn out: no. The manual implementation is faster. I this a bug, or should it be expected? I have written something similar yesterday ( <a href="https://discuss.pytorch.org/t/bug-rmsnorm-slower-than-layernorm/222330/3" class="inline-onebox">Bug? RMSNorm slower than LayerNorm - #3 by JosAr</a> ) for RMSNorm and the newest torch version fixed it; but it does not fix the behaviour I will describe now:</p>
<p>I tested a very easy manual dot product attention, vs the unoptimised code given in <a href="https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html" class="inline-onebox" rel="noopener nofollow ugc">torch.nn.functional.scaled_dot_product_attention — PyTorch 2.8 documentation</a> vs the optimised one with the following script:</p>
<pre><code class="lang-auto">import torch
import torch.nn.functional as F
import time
import math

device = 'cuda'
batch_size, seq_len, d_model = 32, 512, 768
shape = [12, 6, 80, 60, 32]
shape = [6, 12, 40, 60, 32]
num_heads = 8
d_k = 32 // num_heads

# Generate random data
q = torch.randn(shape, device=device)
k = torch.randn(shape, device=device)
v = torch.randn(shape, device=device)

def manual_attention(q, k, v):
    scores = torch.matmul(q, k.transpose(-2, -1)) / (d_k ** 0.5)
    attn_weights = F.softmax(scores, dim=-1)
    return torch.matmul(attn_weights, v)

# Inefficient implementation equivalent to SDPA:
def manual_scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0,
        is_causal=False, scale=None, enable_gqa=False) -&gt; torch.Tensor:
    L, S = query.size(-2), key.size(-2)
    scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
    attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
    if is_causal:
        assert attn_mask is None
        temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
        attn_bias.to(query.dtype)

    if attn_mask is not None:
        if attn_mask.dtype == torch.bool:
            attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
        else:
            attn_bias = attn_mask + attn_bias

    if enable_gqa:
        key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
        value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)

    attn_weight = query @ key.transpose(-2, -1) * scale_factor
    attn_weight += attn_bias
    attn_weight = torch.softmax(attn_weight, dim=-1)
    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
    return attn_weight @ value

def benchmark(func, *args, num_runs=500, warmup=10):
    # Warmup
    for _ in range(warmup):
        _ = func(*args)

    torch.cuda.synchronize()

    start = time.time()
    for _ in range(num_runs):
        _ = func(*args)

    torch.cuda.synchronize()

    return (time.time() - start) / num_runs

# Benchmark both approaches
manual_time = benchmark(manual_attention, q, k, v)
builtin_time = benchmark(F.scaled_dot_product_attention, q, k, v)
inefficient_time = benchmark(manual_scaled_dot_product_attention, q, k, v)


print(torch.__version__)
print(f"Manual implementation: {manual_time*1000:.2f} ms")
print(f"Built-in SDPA: {builtin_time*1000:.2f} ms")
print(f"inefficient time: {inefficient_time*1000:.2f} ms")
</code></pre>
<p>And have received the following results on an A100, using <strong>torch 2.9.0.dev20250813+cu129</strong></p>
<div class="md-table">
<table>
<thead>
<tr>
<th>shape</th>
<th>Manual</th>
<th>Inefficient</th>
<th>SDPA</th>
</tr>
</thead>
<tbody>
<tr>
<td>[6, 12, 40, 60, 32]</td>
<td>0.51 ms</td>
<td>0.53 ms</td>
<td>0.65 ms</td>
</tr>
<tr>
<td>[12, 6, 80, 60, 32]</td>
<td>0.93 ms</td>
<td>1.03 ms</td>
<td>1.26 ms</td>
</tr>
</tbody>
</table>
</div><p>Shouldn’t the efficient SDPA be faster; at least compared to the “inefficient” implementation?</p>
<p>Furthermore Id like to make a feature request (if that’s possible?): Is it possible to make the already implemented <code>attn_bias</code> an optional argument of SDPA? Papers like the SWIN transformer (<a href="https://arxiv.org/pdf/2103.14030" rel="noopener nofollow ugc">https://arxiv.org/pdf/2103.14030</a>)  give a positional Bias (equation (4)). It would be great to have an efficient implementation of that as well!</p>
<p>Thanks a lot already! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=14" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

JosAr: <p>Edit: If its due to the shapes of key, query, value tensors; there are also efficient implementations for that (I’ve seen <a href="https://github.com/zzd1992/FlashWindowAttention" class="inline-onebox" rel="noopener nofollow ugc">GitHub - zzd1992/FlashWindowAttention: Speedup the attention computation of Swin Transformer</a> ) - will those be integrated in the future? Still I thought SDPA should be on-par with the “inefficient implementation”.</p>

Arunprakash-A: <p>No, it is <strong>not</strong> a bug.</p>
<p>SDPA is much faster (efficient) with causal attention</p>
<p>For correct benchmarking, the <code>shape</code> variable in your script must be modified to the appropriate format.</p>
<p><code>shape = [60, 8, 512, 32] # batch_size, num_heads, seq_length, head_dim</code></p>
<p>Now execute it. The table below shows the comparison (in ms) as the sequence length increases from 512 to 1024 to 4096</p>
<div class="md-table">
<table>
<thead>
<tr>
<th></th>
<th>512</th>
<th>1024</th>
<th>4096</th>
</tr>
</thead>
<tbody>
<tr>
<td>Manual Implementation (<strong>no</strong> causal masking)</td>
<td>13.67</td>
<td>53.86</td>
<td>140.16</td>
</tr>
<tr>
<td>Inefficient Implementation (causal masking)</td>
<td>18.03</td>
<td>70.76</td>
<td>214.05</td>
</tr>
<tr>
<td>SDPA (causal masking)</td>
<td>3.95</td>
<td>15.4</td>
<td>40.38</td>
</tr>
<tr>
<td></td>
<td></td>
<td></td>
<td></td>
</tr>
</tbody>
</table>
</div><p>Your implementation <strong>did not</strong> include a causal mask. However, adding a causal mask to the <code>score</code> adds additional FLOPS (latency) that increase quadratically with increasing sequence length. Therefore, compared to an inefficient implementation, the efficient SDPA is five times faster!</p>
<p>Here are the results from running your script on my device (L40 GPU). The only change I made was modifying the <code>shape</code> variable as mentioned above</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/9/7943e048fa9e411eea9b045b6a7421e464ffd09b.png" data-download-href="https://discuss.pytorch.org/uploads/default/7943e048fa9e411eea9b045b6a7421e464ffd09b" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/9/7943e048fa9e411eea9b045b6a7421e464ffd09b.png" alt="image" data-base62-sha1="hiLcgrHP3mV2vuHIDKHYMyhE5B1" width="305" height="255"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">305×255 65.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Hope it helps!</p>
<p>Regards</p>
<p>~Arun</p>

JosAr: <p>Hi!</p>
<p>Thanks for the answer! I have progressed in finding the issue, but it still exists.</p>
<p>Indeed the shape variables I have used earlier were the ones I want to use; in the examplary code there was unfortunately one line of legacy code that might have confused you (sorry for that), but I am interested in the runtimes of <code>shape = [12, 6, 80, 60, 32]</code>and <code>shape = [6, 12, 40, 60, 32]</code>. The Last dimension is the feature dimension E=32, the second last is the sequence length (L=60).</p>
<p>All other dimension (6 heads, 12 batches, 80 “sliding windows” - i am working on some kind of image project) should be treated in my understanding of attention as a kind of batch; if heads or batches  it should computation-wise not matter if I am mistaken.</p>
<p>I further broke it down to be closer to your example, and the documentation:</p>
<p>If you run <code>shape = [24, 6, 2048, 32]</code>  (N, H, L, E) or <code>shape = [24, 1, 6, 2048, 32]</code> (N, …, H,L,E) should not make any difference, if I read from <a href="https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html" class="inline-onebox" rel="noopener nofollow ugc">torch.nn.functional.scaled_dot_product_attention — PyTorch 2.8 documentation</a>  correctly, right? There is only one additional dimension of size 1.</p>
<p>For me it does though, but only for the build in SPDA, not for the others.</p>
<p>What am I missing?</p>