Training with flex attention is extremely slow due to torch.compile settings

yijiong_yu: <p>When I was running an LLM supervised fine-tuning code, where each training sample has a different sequence length,<br>
I found that when using <code>flex_attention</code>, the training speed was nearly <strong>twice slower</strong> than using <code>sdpa</code> or <code>eager</code>.</p>
<p>So I conducted a series of tests to analyze the speed of <code>flex_attention</code> under different compilation settings, when the sequence length and the logic of <code>mask_mod</code> change.</p>
<p>Environment:</p>
<pre><code class="lang-auto">torch==2.8.0
cuda==12.8
transformers==4.54.1
python 3.10
GPU: L40 48GB 
ubuntu~20.04.2
</code></pre>
<h2><a name="p-474587-setting-1-1" class="anchor" href="#p-474587-setting-1-1"></a>setting 1</h2>
<pre data-code-wrap="python"><code class="lang-python">torch.compile(flex_attention, dynamic=True)
torch.compile(create_block_mask, dynamic=True)
</code></pre>
<div class="md-table">
<table>
<thead>
<tr>
<th>changing parameter</th>
<th>forward time</th>
<th>backward time</th>
</tr>
</thead>
<tbody>
<tr>
<td>none</td>
<td>0.0013</td>
<td>0.0006</td>
</tr>
<tr>
<td>seq len</td>
<td>0.0013</td>
<td>0.0011</td>
</tr>
<tr>
<td>mask mod logic</td>
<td>10.2329</td>
<td>0.0117</td>
</tr>
</tbody>
</table>
</div><h2><a name="p-474587-setting-2-2" class="anchor" href="#p-474587-setting-2-2"></a>setting 2</h2>
<pre data-code-wrap="python"><code class="lang-python">torch.compile(flex_attention, dynamic=False)
torch.compile(create_block_mask, dynamic=False)
</code></pre>
<div class="md-table">
<table>
<thead>
<tr>
<th>changing parameter</th>
<th>forward time</th>
<th>backward time</th>
</tr>
</thead>
<tbody>
<tr>
<td>none</td>
<td>0.0008</td>
<td>0.0006</td>
</tr>
<tr>
<td>seq len</td>
<td>1.3453</td>
<td>5.4088</td>
</tr>
<tr>
<td>mask mod logic</td>
<td>0.7264</td>
<td>0.2412</td>
</tr>
</tbody>
</table>
</div><h2><a name="p-474587-setting-3-3" class="anchor" href="#p-474587-setting-3-3"></a>setting 3</h2>
<pre data-code-wrap="python"><code class="lang-python">torch.compile(flex_attention, dynamic=None)
torch.compile(create_block_mask, dynamic=None)
</code></pre>
<div class="md-table">
<table>
<thead>
<tr>
<th>changing parameter</th>
<th>forward time</th>
<th>backward time</th>
</tr>
</thead>
<tbody>
<tr>
<td>none</td>
<td>0.0013</td>
<td>0.0008</td>
</tr>
<tr>
<td>seq len</td>
<td>0.0012</td>
<td>0.0013</td>
</tr>
<tr>
<td>mask mod logic</td>
<td>1.8366</td>
<td>5.4306</td>
</tr>
</tbody>
</table>
</div><h2><a name="p-474587-setting-4-4" class="anchor" href="#p-474587-setting-4-4"></a>setting 4</h2>
<pre data-code-wrap="python"><code class="lang-python">torch.compile(flex_attention, dynamic=True)
torch.compile(create_block_mask, dynamic=False)
</code></pre>
<div class="md-table">
<table>
<thead>
<tr>
<th>changing parameter</th>
<th>forward time</th>
<th>backward time</th>
</tr>
</thead>
<tbody>
<tr>
<td>none</td>
<td>0.0009</td>
<td>0.0004</td>
</tr>
<tr>
<td>seq len</td>
<td>1.2185</td>
<td>0.0010</td>
</tr>
<tr>
<td>mask mod logic</td>
<td>7.3509</td>
<td>0.0121</td>
</tr>
</tbody>
</table>
</div><h2><a name="p-474587-setting-5-5" class="anchor" href="#p-474587-setting-5-5"></a>setting 5</h2>
<pre data-code-wrap="python"><code class="lang-python">torch.compile(flex_attention, dynamic=False)
torch.compile(create_block_mask, dynamic=True)
</code></pre>
<div class="md-table">
<table>
<thead>
<tr>
<th>changing parameter</th>
<th>forward time</th>
<th>backward time</th>
</tr>
</thead>
<tbody>
<tr>
<td>none</td>
<td>0.0014</td>
<td>0.0005</td>
</tr>
<tr>
<td>seq len</td>
<td>6.4919</td>
<td>0.0016</td>
</tr>
<tr>
<td>mask mod logic</td>
<td>7.5734</td>
<td>0.0012</td>
</tr>
</tbody>
</table>
</div><h1><a name="p-474587-findings-6" class="anchor" href="#p-474587-findings-6"></a>Findings</h1>
<ol>
<li>Changing the logic of <code>mask_mod</code> or the sequence length will both lead to recompilation, which will significantly slow down the speed of flex attention under some compilation settings.</li>
<li>The logic or the parameters of <code>mask_mod</code> <strong>should not change</strong> when running different samples, otherwise the speed will be very slow, no matter whether <code>torch.compile</code> is set to dynamic or not. That means data-dependent mask_mod logic is not recommended.</li>
<li>If you need to change the sequence length, it is recommended to set <code>torch.compile</code> to <code>dynamic=True</code> or <code>dynamic=None</code> for <code>flex_attention</code> and <code>create_block_mask</code>. Otherwise, the speed will be very slow.</li>
</ol>
<h1><a name="p-474587-conclusion-7" class="anchor" href="#p-474587-conclusion-7"></a>Conclusion</h1>
<p>By default, the <code>torch.nn.attention.flex_attention.create_block_mask(...,_compile=True)</code> compiles with<br>
<code>dynamic=False</code>. By default, the <code>transformers.integrations.flex_attention.compile_friendly_flex_attention</code> compiles with<br>
<code>dynamic=False</code>.</p>
<p>Therefore, if using <code>transformers</code> with <code>flex_attention</code> to train samples with various lengths, the speed is inevitably very slow.</p>
<h1><a name="p-474587-my-test-code-8" class="anchor" href="#p-474587-my-test-code-8"></a>my test code</h1>
<pre data-code-wrap="python"><code class="lang-python">import os
os.environ["CUDA_VISIBLE_DEVICES"] = "5"

import torch
import time

from torch.nn.attention.flex_attention import flex_attention, create_block_mask, and_masks
from transformers.modeling_utils import flex_attention_forward
from transformers.integrations.flex_attention import compile_friendly_flex_attention

def causal_mask(b, h, q_idx, kv_idx):
    return q_idx &gt;= kv_idx

if __name__ == '__main__':

    head_dim=128
    softmax_scale= head_dim ** -0.5

    flex_attention = torch.compile(flex_attention, dynamic=False)
    create_block_mask= torch.compile(create_block_mask, dynamic=True)


    for query_length in [4000,500,6000,4000]:
        query_states = torch.randn(1, query_length, 32, head_dim, dtype=torch.bfloat16, device="cuda",
                                   requires_grad=True).transpose(1, 2)
        key_states = torch.randn(1, query_length, 4, head_dim, dtype=torch.bfloat16, device="cuda",
                                 requires_grad=True).transpose(1, 2)
        value_states = torch.randn(1, query_length, 4, head_dim, dtype=torch.bfloat16, device="cuda",
                                   requires_grad=True).transpose(1, 2)
        # the label is a tensor of zeros with the same shape as the output of flex_attention
        label = torch.randn(1, query_length, 32, 128, dtype=torch.bfloat16, device="cuda") * 0


        for sliding_window in [1000,2000,3000,1000]:

                time1= time.time()
                # forward
                def sliding_window_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -&gt; bool:
                    return (q_idx - kv_idx) &lt; sliding_window

                mask_mod=and_masks(sliding_window_mask,causal_mask)
                if sliding_window ==-1:
                    mask_mod=causal_mask

                block_mask = create_block_mask(mask_mod, 1, 1, query_length, query_length, device="cuda")
                out_flex = flex_attention(query_states,
                                          key_states,
                                          value_states,
                                          block_mask=block_mask,
                                          enable_gqa=True,
                                          scale=softmax_scale,
                                          kernel_options={"BLOCK_M": 32,"BLOCK_N": 32,}).transpose(1,2)
                time2 = time.time()
                # backward
                loss = torch.nn.functional.mse_loss(out_flex, label)
                loss.backward()

                time3 = time.time()

                print(f"Sliding window: {sliding_window}, Query length: {query_length}")
                print(f"Flex attention forward time: {time2 - time1:.4f} seconds, backward time: {time3 - time2:.4f} seconds")
</code></pre>