Re-implement FlashAttention with new Xe atoms

[petercad]

This PR updates FlashAttention to the new copy/MMA atoms.

Changes:

• Prefill and decode unified into a single implementation, allowing simultaneous K and Q subgroup-level parallelization rather than an either-or.
• GEMMs and softmax grouped together and the full k loop consolidated into an FMHA mainloop class.
  • This will facilitate further manual pipelining/overlap of GEMM with softmax.
• Use new copy/MMA atoms and reorders to transparently support arbitrary data types.
• Automatic copy/MMA operator selection.

Current status: prefill/decode examples almost all working, similar/better performance to old examples.

Known issues:

• Head size 192 decode config doesn't compile yet -- to be fixed.
• Strange SYCL compiler behavior/bug with tSrS->tArP reorder. Apparently the compiler believes there is UB somewhere and will omit a large section of the kernel as a result. For the moment, there's a direct copy as a workaround while I pin down the issue. I'm not able to reproduce this behavior with the reorder in isolation.

Additional features (causal masking, variable sequence lengths, etc.) to be added later.

Reminder: the new atoms require a very recent driver due to necessary IGC fixes/enhancements. Recommended version: ci-comp_igc-30613.

[petercad]

I will break up this large commit into self-contained smaller commits after review is complete.

[rolandschulz]

why is this here? This isn't flash attention specific, is it?

[petercad]

No, it's not. These started as some simple helpers to make copying to/from SLM easier for the epilogue. We could move them, maybe to include/cute/algorithm/cute.hpp, though they should be made more sophisticated (use smaller/larger block sizes as appropriate, automatic fallback to scatter/gather, etc.).

[rolandschulz]

why don't you take i as compile time value to make this safer? The usage is on line 137 where the input comes from the unrolled loop index. If you replace the loop with for_each you have a compile time constant.

[petercad]

That is an option -- I did it this way since compile-time unrolling of the loop is IMO harder to use and harder to read.

I opened a compiler ticket for the lack of diagnostics, and they have a patch under review now to address it.

[rolandschulz]

I see. As long as we have diagnostic that's fine. Current solution won't compile for O0. Not sure whether it matters.

[petercad]

Good point on -O0. How important is it to support -O0 operation? Does the rest of CUTLASS work OK under -O0? (I know SYCL in general has had some functionality issues at -O0.)

[rolandschulz]

I believe the SYCL issues with -O0 have been resolved. I'm not aware of a good reason to compile with -O0. Even for debug O1/Og tends to be better. I think other parts of CUTLASS are correct with O0. But IGC used to crash for larger kernels compiled at O0. Haven't tried it with more recent IGC versions.