Accelerate attention by SDPA

Zoe_Zhao: <p>Hi,</p>
<p>I tried to use SDPA to calculate Attention part fast. I used <code>time.time()</code> to record the time of training one batchsize (I used gradient accumlation). Here is how I record the time.</p>
<pre data-code-wrap="bash"><code class="lang-bash">    for epoch in range(config.epochs):
        if os.path.exists(config.model_save_path) and config.resume:
            losses = checkpoint['loss']
        else:
            losses = 0
       
        t = time.time()
        for idx, batch in enumerate(train_iter):
            if config.ddp:
                model.require_backward_grad_sync = (idx % config.gradient_accumulation_steps == config.gradient_accumulation_steps - 1)
            # batch_size seq
            b_token_ids = batch['input_ids'].t().to(config.device)  
            b_segs = batch['token_type_ids'].t().to(config.device)  
            b_mask = batch['attention_mask'].to(config.device)  
            b_mlm_label = batch['labels'].t().to(config.device)  


            with amp.autocast(device_type='cuda',dtype=torch.bfloat16):    # for forward function! bfloat16
                # with autocast(enabled=False):
                loss, mlm_logits = model(input_ids=b_token_ids,
                                         attention_mask=b_mask,
                                         token_type_ids=b_segs,
                                         masked_lm_labels=b_mlm_label,
                                         next_sentence_labels=None)

                loss = loss / config.gradient_accumulation_steps  # scale the 
            scaler.scale(loss).backward()

            if ((idx  % config.gradient_accumulation_steps) == (config.gradient_accumulation_steps - 1)):
                if config.grad_clip != 0.0:
                    scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad(set_to_none=True)
                scheduler.step()
                print(f'Time for one step(one batchsize):{time.time() - t}')
                t = time.time()

            losses += (loss.item() * config.gradient_accumulation_steps)

            mlm_acc, _, _ = accuracy(mlm_logits, b_mlm_label, config.pad_index)
            ....
       

</code></pre>
<p>And I change my attention to SDPA:</p>
<p>This is my version:</p>
<pre><code class="lang-auto">        .....
        if attn_mask is not None:
            attn_output_weights += attn_mask  # [batch_size * num_heads, tgt_len, src_len]

        .....

        attn_output_weights = F.softmax(attn_output_weights, dim=-1) 
        attn_output_weights = F.dropout(attn_output_weights, p=self.dropout, training=True) 
        attn_output = self.atten_qkv(v, attn_output_weights)
</code></pre>
<p>This is SDPA version:</p>
<pre><code class="lang-auto">    ....
       attn_output = F.scaled_dot_product_attention(q, k, v-self.atten_qkv.t_min_1, dropout_p=self.dropout)
....
</code></pre>
<p>Things are werid because these two version comsume almost the same time. I am not sure whether I use <code>F.scaled_dot_product_attention</code> in a right way. I tested on FP32, FP16, BF16, they all comsume almost the same time.</p>

ptrblck: <p>CUDA operations are executed asynchronously so you need to synchronize the code before starting and stopping host timers.</p>

Zoe_Zhao: <p>Thank you for your reply!</p>
<p>After knowing the dfierences between GPU time and CPU time, I tried to use <code>torch.cuda.synchronize()</code> before recording time.</p>
<p>My code is based on  gradient accumulation:</p>
<pre><code class="lang-auto">   for idx, batch in enumerate(train_iter):
        if config.ddp:
            model.require_backward_grad_sync = (idx % config.gradient_accumulation_steps == config.gradient_accumulation_steps - 1)
        b_token_ids = batch['input_ids'].t().to(config.device) 
        b_segs = batch['token_type_ids'].t().to(config.device)  
        b_mask = batch['attention_mask'].to(config.device)  
        b_mlm_label = batch['labels'].t().to(config.device)  

        if idx  % config.gradient_accumulation_steps == 0:#!!
            torch.cuda.synchronize()
            t0 = time.perf_counter()


        # with record_function("## forward ##"):
        # Use autocast context manager for automatic mixed precision
        with amp.autocast(device_type='cuda',dtype=torch.bfloat16):    # for forward function! bfloat16
            # with autocast(enabled=False):
            loss, mlm_logits = model(input_ids=b_token_ids,
                                     attention_mask=b_mask,
                                     token_type_ids=b_segs,
                                     masked_lm_labels=b_mlm_label,
                                     next_sentence_labels=None)

            loss = loss / config.gradient_accumulation_steps  # scale the loss to account for gradient accumulation

        scaler.scale(loss).backward()

        if ((idx  % config.gradient_accumulation_steps) == (config.gradient_accumulation_steps - 1)):
            if config.grad_clip != 0.0:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad(set_to_none=True)
            scheduler.step()
            torch.cuda.synchronize()
            p = time.perf_counter()
            print(f'rank:{ddp_local_rank}!!!!!!!Time for one step(one batchsize):{p - t0}b16')#!!
</code></pre>
<p><strong>Without sycronization (float32) CPU time:</strong></p>
<p>5.6577911376953125</p>
<p><strong>After sycronization (float32):</strong></p>
<p><code>torch.compile()</code> + <code>SDPA:</code></p>
<p>rank:2!!!.Time for one step(one batchsize):<strong>4.333474938757718!!!</strong><br>
rank:0!!!.Time for one step(one batchsize):4.331189756747335!!!<br>
rank:3!!!.Time for one step(one batchsize):4.331575924996287!!!<br>
rank:1!!!.Time for one step(one batchsize):4.332561417948455!!!</p>
<p><code>SDPA:</code></p>
<p>rank:1!!!.Time for one step(one batchsize):5.672491875011474!!!<br>
rank:0!!!.Time for one step(one batchsize):5.672714178916067!!!<br>
rank:2!!!.Time for one step(one batchsize):5.672652362845838!!!<br>
rank:3!!!.Time for one step(one batchsize):5.67308798385784!!!</p>
<p>without <code>SPDA</code> + without <code>torch.compile()</code>:</p>
<p>rank:1!!!Time for one step(one batchsize):5.78729198500514<br>
rank:2!!!Time for one step(one batchsize):5.791289611253887<br>
rank:0!!!Time for one step(one batchsize):5.7841532728634775<br>
rank:3!!!Time for one step(one batchsize):5.786291209049523</p>
<p><strong>Without sycronization (bfloat16) CPU time:</strong></p>
<p>2.8307549953460693</p>
<p><strong>After sycronization (bfloat16):</strong></p>
<p><code>SDPA:</code></p>
<p>rank:0!!!Time for one step(one batchsize):2.8250338532961905b16<br>
rank:3!!!Time for one step(one batchsize):2.8336645336821675b16<br>
rank:2!!!Time for one step(one batchsize):2.8370646508410573b16<br>
rank:1!!!Time for one step(one batchsize):2.834382927045226b16</p>
<p>without <code>SPDA</code> + without <code>torch.compile()</code>:</p>
<p>rank:1!!!Time for one step(one batchsize):2.828244622796774b16<br>
rank:0!!!Time for one step(one batchsize):2.8276006169617176b16<br>
rank:2!!!Time for one step(one batchsize):2.8268630262464285b16<br>
rank:3!!!Time for one step(one batchsize):2.8273697779513896b16</p>
<p>I’m confused, beacause it seems <strong>no acceleration and CPU time  looks like GPU time under most circumstances</strong>, and when floa32 wirh <code>torch.compile() </code>GPU time is smaller time, does that mean we need more CPU workers to process data faster?</p>

Zoe_Zhao: <p>And also, I compare the time between the atten part that defined by myself and SDPA:</p>
<pre><code class="lang-auto">import torch.utils.benchmark as benchmark
def benchmark_torch_function_in_microseconds(f, *args, **kwargs):
    t0 = benchmark.Timer(
        stmt="f(*args, **kwargs)", globals={"args": args, "kwargs": kwargs, "f": f}
    )
    return t0.blocked_autorange().mean * 1e6
</code></pre>
<p>SDPA:</p>
<pre><code class="lang-auto">print(
            f"The default implementation runs in {benchmark_torch_function_in_microseconds(F.scaled_dot_product_attention, q, k, v-self.atten_qkv.t_min_1):.3f} microseconds")

</code></pre>
<p>return:</p>
<p>The default implementation runs in 2108.397 microseconds<br>
The default implementation runs in 2111.590 microseconds<br>
The default implementation runs in 2108.828 microseconds<br>
The default implementation runs in 2112.032 microseconds<br>
The default implementation runs in 2109.155 microseconds<br>
The default implementation runs in 2110.308 microseconds</p>
<p>…</p>
<p>my version:</p>
<pre><code class="lang-auto">...
    def atten(self, q, k, v):
        attn_output_weights = torch.bmm(q, k.transpose(1, 2))
        attn_output_weights = F.softmax(attn_output_weights, dim=-1)  # [batch_size * num_heads, tgt_len, src_len]
        attn_output_weights = F.dropout(attn_output_weights, p=self.dropout, training=True)  # voltage
        attn_output = self.atten_qkv(v, attn_output_weights)
        return attn_output, attn_output_weights

    def forward(self, query, key, value, attn_mask=None, key_padding_mask=None):
...
         print(
            f"The default implementation runs in {benchmark_torch_function_in_microseconds(self.atten, q, k, v):.3f} microseconds")
....
</code></pre>
<p>return:</p>
<p>The default implementation runs in 2106.173 microseconds<br>
The default implementation runs in 2104.582 microseconds<br>
The default implementation runs in 2107.339 microseconds<br>
The default implementation runs in 2104.767 microseconds<br>
The default implementation runs in 2106.384 microseconds<br>
The default implementation runs in 2104.327 microseconds<br>
The default implementation runs in 2105.738 microseconds<br>
The default implementation runs in 2104.363 microseconds</p>
<p>…</p>
<p>I’m very confused, why my code cannot  be accelerated.</p>

Arunprakash-A: <aside class="quote no-group" data-username="Zoe_Zhao" data-post="1" data-topic="223211">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zoe_zhao/48/78238_2.png" class="avatar"> Zoe_Zhao:</div>
<blockquote>
<p><code>F.scaled_dot_product_attention</code></p>
</blockquote>
</aside>
<p>Did you use SDPA within the context manager? like</p>
<pre data-code-wrap="python"><code class="lang-python">from torch.nn.attention import SDPBackend, sdpa_kernel

# Only enable flash attention backend
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    scaled_dot_product_attention(...)

# Enable the Math or Efficient attention backends
with sdpa_kernel([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]):
    scaled_dot_product_attention(...)
</code></pre>
<p>Please check out the documentation <a href="https://docs.pytorch.org/docs/stable/generated/torch.nn.attention.sdpa_kernel.html#torch.nn.attention.sdpa_kernel" rel="noopener nofollow ugc">here</a>. Sometimes, not using SDPA correctly can actually slow down performance.</p>

Zoe_Zhao: <p>Thank you for your reply. But I don’t think conntext manager is necessary. Because <a href="https://docs.pytorch.org/tutorials/intermediate/scaled_dot_product_attention_tutorial.html" rel="noopener nofollow ugc">this</a> shows that <code>F.scaled_dot_product_attention</code> can be applied alone and it is fast.</p>

Arunprakash-A: <p>Yes, that’s correct. I was thinking that doing it explicitly might help. So I separated the concerns and ran the following code for a quick check and found that the implementation has no issues.</p>
<pre data-code-wrap="python"><code class="lang-python">import time
import math
import torch
import torch.nn.functional as F

query = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
key = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
value = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")

def 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
builtin_time = benchmark(F.scaled_dot_product_attention, query, key, value)
manual_time = benchmark(scaled_dot_product_attention, query, key, value)


print(f"Built-in time: {builtin_time*1000:.2f} ms")
print(f"manual time: {manual_time*1000:.2f} ms")
</code></pre>
<p>I got the following results,</p>
<pre><code class="lang-auto">Built-in time: 0.04 ms
manual time: 0.17 ms
</code></pre>
<p>There might be an issue (which is often subtle) in your code. Carefully check the shapes passed to the SDPA function and also the arguments used in the <code>MultiheadAttention</code> module (if you’re using that in your model. For reference, you might find the following threads helpful:</p>
<ul>
<li><a href="https://discuss.pytorch.org/t/causal-masking-hint/222699/3">causal-masking-hint</a></li>
<li><a href="https://discuss.pytorch.org/t/bug-scaled-dot-product-attention-slower-than-manual-multiplication/222377/3">bug-scaled-dot-product-attention-slower-than-manual-multiplication</a></li>
</ul>

Zoe_Zhao: <p>Thank you so much! I think I solved  the issue!<img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=14" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>
<p>My previous q,k,v’s shape is  [24*16, 512, 64] #[batch_size * num_heads,src_len, kdim].</p>
<p>And then I changed to <strong>[24, 16, 512, 64] # batch_size, num_heads, seq_length, head_dim. And it is really fast!</strong></p>
<p>return:</p>
<p>The default implementation runs in 184.328 microseconds<br>
The default implementation runs in 184.655 microseconds</p>
<p>compared to</p>
<aside class="quote no-group" data-username="Zoe_Zhao" data-post="4" data-topic="223211">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zoe_zhao/48/78238_2.png" class="avatar"> Zoe_Zhao:</div>
<blockquote>
<p>The default implementation runs in 2108.397 microseconds<br>
The default implementation runs in 2111.590 microseconds<br>
The default implementation runs in 2108.828 microseconds<br>
The default implementation runs in 2112.032 microseconds<br>
The default implementation runs in 2109.155 microseconds<br>
The default implementation runs in 2110.308 microseconds</p>
</blockquote>
</aside>
<p>A bit hard to believe, since the shapes of Q, K, and V really do matter! I’m curious why, because the fixed shape of fused kernel?</p>

Zoe_Zhao: <p>But things are weird, because although <code>F.scaled_dot_product_attention, query, key, value)</code> is fast, when I acculate one batch time, without  <code>F.scaled_dot_product_attention, query, key, value)</code> is a little fast.</p>
<p><strong>float32</strong> :</p>
<p>my attention:</p>
<pre><code class="lang-auto">    def atten(self, q, k, v):
        attn_output_weights = torch.bmm(q, k.transpose(1, 2))
        attn_output_weights = F.softmax(attn_output_weights, dim=-1)  
        attn_output_weights = F.dropout(attn_output_weights, p=self.dropout, training=True)  
        attn_output = self.atten_qkv(v, attn_output_weights)
        return attn_output, attn_output_weights
    def forward(..):
        ...
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        attn_output, attn_output_weights = self.atten(q, k, v-self.atten_qkv.t_min_1)    ## batch_size, num_heads, seq_length, head_dim
        torch.cuda.synchronize()
        t1 = time.perf_counter()
        elapsed_us = (t1 - t0) * 1e6  
        print(f'atten time:{elapsed_us}')
</code></pre>
<p>return:</p>
<p>atten time:1663.760282099247 ms<br>
atten time:1693.3688893914223 ms<br>
atten time:1687.0391555130482 ms</p>
<p><code>F.scaled_dot_product_attention, query, key, value)</code> version:</p>
<pre><code class="lang-auto">   def forward():
        ...
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        attn_output, attn_output_weights = self.atten(q, k, v-self.atten_qkv.t_min_1, dropout_p=self.dropout)    
        torch.cuda.synchronize()
        t1 = time.perf_counter()
        elapsed_us = (t1 - t0) * 1e6  
        print(f'atten time:{elapsed_us}')
</code></pre>
<p>return:</p>
<p>atten time:1054.6050034463406 ms<br>
atten time:1051.4757595956326 ms<br>
atten time:1051.2289591133595 ms</p>
<p>But when I caculate the batchtime:</p>
<pre><code class="lang-auto">
    for epoch in range(config.epochs):
        if os.path.exists(config.model_save_path) and config.resume:
            losses = checkpoint['loss']
        else:
            losses = 0


        for idx, batch in enumerate(train_iter):
            if config.ddp:
                model.require_backward_grad_sync = (idx % config.gradient_accumulation_steps == config.gradient_accumulation_steps - 1)
            b_token_ids = batch['input_ids'].t().to(config.device)  
            b_segs = batch['token_type_ids'].t().to(config.device)  
            b_mask = batch['attention_mask'].to(config.device)  
            b_mlm_label = batch['labels'].t().to(config.device)  

            if idx  % config.gradient_accumulation_steps == 0:
                torch.cuda.synchronize()
                t0 = time.perf_counter()

            with amp.autocast(device_type='cuda',dtype=torch.float32):    
                # with autocast(enabled=False):
                loss, mlm_logits = model(input_ids=b_token_ids,
                                         attention_mask=b_mask,
                                         token_type_ids=b_segs,
                                         masked_lm_labels=b_mlm_label,
                                         next_sentence_labels=None)

                loss = loss / config.gradient_accumulation_steps  # scale the loss to account for gradient accumulation
            scaler.scale(loss).backward()

            if ((idx  % config.gradient_accumulation_steps) == (config.gradient_accumulation_steps - 1)):

                if config.grad_clip != 0.0:
                    scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad(set_to_none=True)
                scheduler.step()
                torch.cuda.synchronize()
                p = time.perf_counter()
                print(f'rank:{ddp_local_rank}!!!!!!!Time for one step(one batchsize):{p - t0}b16')


            losses += (loss.item() * config.gradient_accumulation_steps)
            mlm_acc, _, _ = accuracy(mlm_logits, b_mlm_label, config.pad_index)
        ...
    
</code></pre>
<p>my attention return:<br>
rank:1!!!Time for one step(one batchsize):5.563375188037753s<br>
rank:0!!!Time for one step(one batchsize):5.564123823773116s<br>
rank:0!!!Time for one step(one batchsize):5.565241441130638s<br>
rank:1!!!Time for one step(one batchsize):5.5651940810494125s</p>
<p><code>F.scaled_dot_product_attention, query, key, value)</code> version return:</p>
<p>rank:0!!!Time for one step(one batchsize):5.580610099248588s<br>
rank:1!!!Time for one step(one batchsize):5.579201320186257s<br>
rank:0!!!Time for one step(one batchsize):5.576219431124628s<br>
rank:1!!!Time for one step(one batchsize):5.576044279150665s</p>

Arunprakash-A: <aside class="quote no-group" data-username="Zoe_Zhao" data-post="9" data-topic="223211">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zoe_zhao/48/78238_2.png" class="avatar"> Zoe_Zhao:</div>
<blockquote>
<p><code>attn_output, attn_output_weights = self.atten(q, k, v-self.atten_qkv.t_min_1, dropout_p=self.dropout)</code></p>
</blockquote>
</aside>
<p><code>F.scaled_dot_product_attention</code> returns only a tensor; however, your code unpacks two tensors (<code>attn_output</code> and <code>attn_output_weights</code>)? I think you are using the built-in <code>MultiHeadAttention</code> class. If that is the case, set <code>need_weights=False</code> in the <code>forward</code> method of the class.</p>

Zoe_Zhao: <aside class="quote no-group" data-username="Zoe_Zhao" data-post="9" data-topic="223211">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zoe_zhao/48/78238_2.png" class="avatar"> Zoe_Zhao:</div>
<blockquote>
<p><code>F.scaled_dot_product_attention, query, key, value)</code> version:</p>
<pre><code class="lang-auto">   def forward():
        ...
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        attn_output, attn_output_weights = self.atten(q, k, v-self.atten_qkv.t_min_1, dropout_p=self.dropout)    
        torch.cuda.synchronize()
        t1 = time.perf_counter()
        elapsed_us = (t1 - t0) * 1e6  
        print(f'atten time:{elapsed_us}')

</code></pre>
</blockquote>
</aside>
<p>Sorry！I copied the wrong version earlier. Here’s the correct one:</p>
<pre><code class="lang-auto">        torch.cuda.synchronize()
        t0 = time.perf_counter()    
        attn_output = F.scaled_dot_product_attention(q, k, v-self.atten_qkv.t_min_1, dropout_p=self.dropout)    ## batch_size, num_heads, seq_length, head_dim
        t1 = time.perf_counter()
        elapsed_us = (t1 - t0) * 1e6  
        print(f'atten time:{elapsed_us}')

</code></pre>
<p>My question is still why, when I measure the batch processing time, I don’t see any acceleration, but when I measure the time of a single attention layer, SDPA appears fast.<img src="https://discuss.pytorch.org/images/emoji/apple/worried.png?v=14" title=":worried:" class="emoji" alt=":worried:" loading="lazy" width="20" height="20"></p>

Arunprakash-A: <aside class="quote no-group" data-username="Zoe_Zhao" data-post="13" data-topic="223211">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zoe_zhao/48/78238_2.png" class="avatar"> Zoe_Zhao:</div>
<blockquote>
<p>My question is still why, when I measure the batch processing time, I don’t see any acceleration</p>
</blockquote>
</aside>
<p>I modified the batch size to 256 and the sequence length to 1024, then ran the code</p>
<aside class="quote no-group quote-modified" data-username="Arunprakash-A" data-post="7" data-topic="223211">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/arunprakash-a/48/78115_2.png" class="avatar"> Arunprakash-A:</div>
<blockquote>
<pre><code class="lang-auto">query = torch.rand(256, 8, 1024, 64, dtype=torch.float16, device="cuda")
key = torch.rand(256, 8, 1024, 64, dtype=torch.float16, device="cuda")
value = torch.rand(256, 8, 1024, 64, dtype=torch.float16, device="cuda")
</code></pre>
</blockquote>
</aside>
<p>I did not observe any decrease in performance and obtained the following results:</p>
<pre><code class="lang-auto">Built-in time: 17.29 ms
manual time: 152.88 ms
</code></pre>
<p>On the contrary, FlashAttention is 9x faster than the inefficient implementation.</p>

Zoe_Zhao: <p>Yes. I got it. I believe additional operations comsume the time in  my code not SDPA. Thank you!</p>