Flash Attention: Setup

[louisfd]

Add cubecl-attention crate for FlashAttention

• Introduces the cubecl-attention crate with the same hierarchical levels as matmul: batch, global, stage, tile.
• Tile level has the most differences compared to matmul. Currently, tile attention performs round trips to shared memory for operations within fragments. This will be optimized once we have finer control over MMA fragments.
• The algorithm currently only supports accelerated 8×8×8 matmul (used with Metal) and is hardcoded for this configuration.

TODOs:

• Generalize tile sizes.
• Stage: support partitioning with multiple planes.
• Global: iterate over multiple stage iterations.
• Batch: full-problem dispatch.

Despite these limitations, reusing the matmul architecture should make further development straightforward.

Validate your PR with burn.

It is important that you make sure that you don't introduce any bugs in burn.

Instructions

• Create a new branch or fork of the burn repo
• Update the main Cargo.toml with this PR hash.
• Fix any broken tests or compilation errors in burn.
•  Submit a PR in burn with your fixes and link it here.

[wingertge]

This will be optimized once we have finer control over MMA fragments.

Good news 😅