Add torch.compile support to flash attention 3

[guilhermeleobas]

Enable torch.compile support for FlashAttention and improve testing

• Add support for torch.compile to recognize FlashAttention forward/backward functions.
• Update tests to use torch._subclasses.fake_tensor.fake_check to validate the FakeTensor implementation.
• Creates a file exposing the flags used to build FlashAttention at runtime (flash_attn_config.py)

[guilhermeleobas]

cc @zou3519 @anijain2305

[zou3519]

What's going on here? Adding additional clones is generally bad for performance

[zou3519]

do you have more context for this file? I don't see a mention of PHI-1 elsewhere

[guilhermeleobas]

I'll remove this file once the PR gets approved. I only keep it as it makes it easier to build and test flash attention.

[zou3519]

A bit jank but I'm fine with it

[guilhermeleobas]

@tridao, are you ok with running fake_check as part of the tests? I can invert the flag to be opt-in instead of opt-out.

[tridao]

would that slow down the tests much? Rn it takes 30-50mins to run all the tests

[guilhermeleobas]

I suspect it can double the time to run the tests, as fake_check would run the function twice to compare the fake version with the actual implementation. I can change this flag to be opt-in instead. So, to run fake_check, one should enable the flag.

[zou3519]

Do we have a sense of how comprehensive the existing tests in this file are? Are they good at exercising a variety of inputs?

[guilhermeleobas]

I guess @tridao would be a better person to answer this one.

[tridao]

we do test a lot of input shapes and different attn options (~100k tests iirc)

[zou3519]

Looks good to me, modulo the prints

[lantudou]

@guilhermeleobas Anything update? I tried your commits and found there are still some shape errors in the backward function, here is my test code.
`
import torch
from flash_attn_interface import flash_attn_func, _flash_attn_forward
from torch import nn
class EfficienctMultiHeadAttention(nn.Module):

def __init__(self, embed_size, num_heads, dropout=0.0, use_flash_attn=True):
    """
    初始化。

    参数:
        embed_size (int): 输入和输出的特征维度。
        num_heads (int): 注意力头的数量。
        dropout (float): 应用于注意力权重的 dropout 比率。
        use_flash_attn (bool): 是否尝试使用 FlashAttention (如果可用)。
    """
    super().__init__()
    assert embed_size % num_heads == 0, "embed_size 必须能被 num_heads 整除"

    self.embed_size = embed_size
    self.num_heads = num_heads
    self.head_dim = embed_size // num_heads
    self.use_flash_attn = use_flash_attn and (flash_attn_func is not None)

    # 使用一个线性层同时生成 Q, K, V，效率更高
    self.qkv_proj = nn.Linear(embed_size, 3 * embed_size)
    # 最终输出的线性层
    self.out_proj = nn.Linear(embed_size, embed_size)
    self.dropout = dropout

def forward(self, x, attention_mask=None):
    """
    前向传播。

    参数:
        x (torch.Tensor): 输入张量，形状为 (N, seq_length, embed_size)
        attention_mask (torch.Tensor, optional): 注意力掩码。
            对于 SDPA，布尔掩码中 True 的位置表示 *不* 被关注。
            对于 FlashAttention，通常使用 causal 参数，而非 mask。

    返回:
        torch.Tensor: 输出张量，形状为 (N, seq_length, embed_size)
    """
    N, seq_length, _ = x.shape

    # 1. 投影并切分 Q, K, V
    # (N, seq_length, embed_size) -> (N, seq_length, 3 * embed_size)
    qkv = self.qkv_proj(x)
    # (N, seq_length, 3 * embed_size) -> 3 x (N, seq_length, embed_size)
    q, k, v = qkv.chunk(3, dim=-1)

    # 2. 重塑 Q, K, V 以便多头计算
    # (N, seq_length, embed_size) -> (N, seq_length, num_heads, head_dim)
    q = q.view(N, seq_length, self.num_heads, self.head_dim)
    k = k.view(N, seq_length, self.num_heads, self.head_dim)
    v = v.view(N, seq_length, self.num_heads, self.head_dim)

    # --- 使用 FlashAttention 的路径 ---
    if self.use_flash_attn and attention_mask is None:
        # flash_attn_func 需要的形状是 (batch, seqlen, nheads, headdim)
        # 这与我们当前的形状完全匹配
        # flash_attn_func 的 dropout 在内部处理
        #print("正在使用 FlashAttention...")
        out = flash_attn_func(
            q, k, v
        )
    # 3. 合并多头的输出
    # (N, seq_length, num_heads, head_dim) -> (N, seq_length, embed_size)
    out = out.reshape(N, seq_length, self.embed_size)

    # 4. 通过最后的线性层
    out = self.out_proj(out)

    return out

batch_size = 16

sequence_length = 256
embedding_dim = 2048

test = EfficienctMultiHeadAttention(embedding_dim, num_heads=16).cuda().bfloat16()
input_tensor = torch.randn(batch_size, sequence_length, embedding_dim).cuda().bfloat16()
test = torch.compile(test, mode='max-autotune')

out = test(input_tensor)
loss = out.sum()
loss.backward()
`

[guilhermeleobas]

Hi @lantudou, thanks for the reproducer. Could you try it again?

[guilhermeleobas]

@tridao, could you take a look one more time once you have some cycles to spare?

[m3rcuriel]

I think the torch custom ops should be registered to a flash_attn_3 namespace to match the TORCH_LIBRARY registration in C. Would these conflict with FA2?

[guilhermeleobas]

I think the torch custom ops should be registered to a flash_attn_3 namespace to match the TORCH_LIBRARY registration in C. Would these conflict with FA2?

Is flash attention 3 an independent package of FA2? In the sense that in the future FA2 will be deprecated in favor of FA3?

[m3rcuriel]

...

Is flash attention 3 an independent package of FA2? In the sense that in the future FA2 will be deprecated in favor of FA3?

idk what the future plan is but at least right now transformers thinks it can have both at the same time

[guilhermeleobas]

idk what the future plan is but at least right now transformers thinks it can have both at the same time

Got it. I changed the namespace to flash_attn_3.

[Tomcli]

Thank you for the great work @guilhermeleobas. Just wondering have you tried to run training on the export FX graph? It seems like this implementation still missing register_autograd in order to run the training loop.

[guilhermeleobas]

Thanks for the feedback, @Tomcli. As for the second question. Probably it won't work. I based my implementation on what is implemented for FA2, which doesn't use register_autograd as well. I can do this in a follow-up PR if needed.

Edit: added in this one.

[guilhermeleobas]

Hi @Tomcli, could you try the last commit, please?

[supercharleszhu]

Suggested change

	is_varlen_q = bool(cu_seqlens_q)
	is_varlen_k = bool(cu_seqlens_k)
	is_varlen = is_varlen_q or is_varlen_k or bool(sequed_q) or bool(sequed_k)
	is_varlen_q = cu_seqlens_q is not None
	is_varlen_k = cu_seqlens_k is not None
	is_varlen = is_varlen_q or is_varlen_k or sequed_q is not None or sequed_k is not None