Flex attention - gaps in profiler

username: tugot17
cooked: <p>I have recently been playing with Flex attention, trying to replace some of my custom triton kernels. Unfortunately Flex attention ones were substantially (around 5x) slower than the custom one I tried. I use something similar to the <a href="https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py" rel="noopener nofollow ugc">Tri Dao flash attention kernel</a>, but with bias. I started looking at the torch profile for a single forward pass on the complied version after the model warmup and it looks very strange.</p>
<pre><code class="lang-auto">import torch
from typing import Optional
from torch.nn.attention.flex_attention import (
    create_block_mask,
    flex_attention,
)
from torch.nn.attention.flex_attention import _mask_mod_signature
import random
from functools import partial
from triton.testing import do_bench
from torch.profiler import profile, ProfilerActivity, record_function

torch._dynamo.config.cache_size_limit = 1000
flex_attention = torch.compile(flex_attention, dynamic=False)

def flexattention_varlen_bias(
    query: torch.Tensor,  # shape (total_tokens_q, n_heads, head_dim)
    key: torch.Tensor,    # shape (total_tokens_k, n_heads, head_dim)
    value: torch.Tensor,  # shape (total_tokens_v, n_heads, head_dim)
    *,
    bias: Optional[torch.Tensor] = None,  # shape (total_tokens_q,)
    document_id: torch.Tensor,  # shape (total_tokens,)
    use_causal_mask: bool = False,
) -&gt; torch.Tensor:
    device = query.device 
    _, _, total_tokens, _ =  query.size()

    def additive_attention_fn(score, b, h, q_idx, kv_idx):
        return score + bias[kv_idx]

    def document_causal_mask(b, h, q_idx, kv_idx):
        causal_mask = q_idx &gt;= kv_idx if use_causal_mask else True
        document_mask = document_id[q_idx] == document_id[kv_idx]
        return causal_mask &amp; document_mask 

    block_mask = create_block_mask(
        document_causal_mask,
        1,
        1,
        total_tokens,
        total_tokens,
        device=device,
    )

    return flex_attention(
        query,
        key,
        value,
        block_mask=block_mask,
        score_mod=additive_attention_fn if bias is not None else None
    )


if __name__ == "__main__":
    random.seed(0)
    device = "cuda"
    dtype = torch.float16
    causal = True
    max_seq_len, doc_count = 2048, 8
    nheads = 32
    HEAD_DIM = 128

    def generate_random_lengths(total_length, num_documents):
        lengths = [1] * num_documents
        remaining_length = total_length - num_documents
        for _ in range(remaining_length):
            index = random.randint(0, num_documents - 1)
            lengths[index] += 1
        return lengths

    lengths = generate_random_lengths(max_seq_len, doc_count)

    document_id = torch.zeros(max_seq_len, dtype=torch.int32, device=device)
    current_idx = 0
    for doc_id, length in enumerate(lengths):
        document_id[current_idx:current_idx+length] = doc_id
        current_idx += length

    query = torch.randn(1, max_seq_len, nheads, HEAD_DIM, device=device, dtype=dtype)
    key = torch.randn_like(query)
    value = torch.randn_like(query)
    bias = torch.rand(max_seq_len, device=device, dtype=dtype)

    flex_attention_fn = partial(flexattention_varlen_bias,
        query,
        key,
        value,
        bias=bias,
        document_id=document_id,
        use_causal_mask=causal
    )

    benchmark_time = do_bench(flex_attention_fn, warmup=25, rep=100)
    print(f"Benchmark time: {benchmark_time:.4f} ms")

    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
        with record_function("flexattention_varlen_bias"):
            out = flex_attention_fn()
    prof.export_chrome_trace("flexattention_trace.json")

    print("Output shape:", out.size())
    # print("Document ID tensor:", document_id)
    print("Lengths of each document:", lengths)
</code></pre>
<p>The profile:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/e/4e0b874bd562b6b5fea356ab4b764777bdb86b44.png" data-download-href="https://discuss.pytorch.org/uploads/default/4e0b874bd562b6b5fea356ab4b764777bdb86b44" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/e/4e0b874bd562b6b5fea356ab4b764777bdb86b44_2_690x200.png" alt="image" data-base62-sha1="b8pWqyVZ14sXfSJf9JNoJRmxvs8" width="690" height="200" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/e/4e0b874bd562b6b5fea356ab4b764777bdb86b44_2_690x200.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/e/4e0b874bd562b6b5fea356ab4b764777bdb86b44_2_1035x300.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/4/e/4e0b874bd562b6b5fea356ab4b764777bdb86b44_2_1380x400.png 2x" data-dominant-color="8EADB1"><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">1756×510 10.6 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>There are a lot of gaps, it looks as if the flex attention was not properly compiled. Do you have any ideas how to reduce these? I tried a bunch of different compilation modes but unfortunately to no avail.</p>