memory_efficient_attention faster than flash attention 2 backend?

[asahni04]

❓ Questions and Help

expected other way around. what is the fastest kernel i can use here?

            q = q.to(dtype)  # XFormers needs manual casting of the operators
            k = k.to(dtype)
            v = v.to(dtype)
            x = memory_efficient_attention(
                q,
                k,
                v,
                p=self.attn_drop.p if self.training else 0.0,
                op=self.efficient_attention_ops,
            )
            vs

            q,k,v = map(lambda t: rearrange(t, "b n h d -> b h n d", d=self.head_dim), (q, k, v))
            with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
                    x = scaled_dot_product_attention(
                        q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0
                    )  # Scale is automatically computed by the torch implementation