Meta's xFormers has Attention kernel written in CUTLASS

[hwu36]

So glad to see people using CUTLASS to implement their own kernels, especially something complex like Attention.

@danthe3rd @fmassa

[danthe3rd]

Hi @hwu36
Our attention kernel heavily relies on cutlass yes. As our goals were to (1) be quite efficient (2) support many GPUs and (3) support many data types (fp16, bf16, f32), we found that cutlass helped us meet all of these goals :) It's amazing how optimised cutlass code is "out of the box" on many situations, compared to anything we would have written ourselves. Reading through the code was also a way to learn tricks on how to write optimised CUDA code.
Our use-case was a bit different from what CUTLASS targets (eg prologue+mma+epilogue), as we want to chain multiple matmuls on the same kernel. The fused gemm example especially was really useful on that regard.

We found ourselves extending cutlass's classes on several occasions:
(1) We subclassed the threadblock::Mma classes - mainly to support doing the prologue early (so when we start our second matmul, data is already loaded in memory and we don't have to wait for gmem)
(2) We took the classes from the fused gemm example, and adapted them to work for V100/SIMT
(3) We also had to change how iterators handle residuals, to handle arbitrary sequence lengths for Q and K -
(4) We added a prefetch method to the epilogue's tile iterator
(5) Finally, we also modified the threadblock::Epilogue class to pipeline memory loads from source (was important for performance in our case as we did many matmul on matrices with small K), support source having a different dtype than the target (we want to accumulate in f32, but the output can be f16), and passing the row id to our custom epilogue
(6) For proper performance in the future for K <= 32, we would need to implement some more code, as @tridao pointed out

Several times we had to modify existing classes, which properties were private (rather than protected) - I assumed it's because you might not want users to rely on implementation details - so we found ourselves copy-pasting cutlass classes just to modify a few things. Maybe there is a better way?
I'm happy to chat if it makes sense to you to upstream some of these changes. Supporting chained matmul in CUTLASS's core would be really useful for multiple kernels for instance.
Finally, we would love to collaborate more closely on this :)

[hwu36]

@danthe3rd , thank you very much for your summary. It is very cool. From what you said, it sounds like a fused attention kernel. You already did most of things I would do if I wrote attention.

What is the relation of this one and your flash attention?

The collaboration is also very welcome. You can tell from our 2.10 that cutlass needs a good fused attention.

[danthe3rd]

This is indeed a fused attention kernel.

What is the relation of this one and your flash attention?

In xFormers, we expose a single memory_efficient_attention function for end users, which dispatches to either Flash attention, or our own kernel based on cutlass. Historically, we started exploring fused kernels in April to save on memory. Flash demonstrated impressive performance gains as well, and motivated us to push it further - CUTLASS made our life easier to do that :)

In terms of implementation, there are a few key differences with Flash:
(1) We store intermediate results/operands in shared memory - Flash stores some results/operands in RF. This allows us to support arbitrary embedding sizes K
(2) FW only: We parallelise across the number of queries. For input sizes [B, M, H, K], Flash will launch O(B*H) blocks, when we launch O(B*M*H)

In terms of performance, we outperform Flash on the FW for K>64 by up to 1.8x. We're still lagging behind for the BW pass [1] - further optimisation will be needed.

The collaboration is also very welcome

I'll send you an email to discuss that further

[1] Performance comparison for BW: facebookresearch/xformers#469
Interesting columns: flash[flshatt] (Flash), 48_chunk3_31735f9 (our implementation) and vanilla (eager pytorch)

[hwu36]

I'll send you an email to discuss that further

You know my email address, right?

Everything you said above are correct choices. Two more things that I would like to add are 1) cutlass gemm epilogue permutation to keep reducing the memory traffic. 2) cutlass group gemm interface so that both fixed sequence length and variable sequence length can be supported.

[danthe3rd]

I just sent it to you and @mnicely . I found your @nvidia.com address on a commit if it's correct

[mnicely]

I'll send you an email to discuss that further

@danthe3rd would you mind adding me to that email?

[tridao]

@danthe3rd, @fmassa and I have been collaborating and swapping notes about FlashAttention / Mem-efficient attention.
They've been doing amazing work using Cutlass in the attention implementation.
I've also started to rewrite FlashAttention to use Cutlass (just the forward pass for now).
I could imagine further collab where we merge the implementations (e.g. currently Cutlass has a fused attention implementation as well).

[hwu36]

We store intermediate results/operands in shared memory -

Our b2b gemm also has a version that stores the results in the registers. It can be used in the short sequence length case.

[hwu36]

currently Cutlass has a fused attention implementation as well

I am not going to call our current ex41 a fused implementation. This one uses group gemm to support variable sequence length well, but it is not fast in fixed sequence length because the kernel is not fully fused. In the long run, we want to use group gemm together with fmha to support both fixed sequence and variable sequence length efficiently in one kernel.

[hwu36]

The front pass is merged in #662 .