Add more GPU architctures support

[uchihatmtkinu]

Hi FlashMLA Team,

This PR, submitted by NVIDIA, introduces support for additional GPU architectures in FlashMLA.

Key Points

More GPU Support

• Refactored the file structure to improve extensibility.
• Added forward and backward FMHA kernels and tests for other GPU Architectures(e.g. SM100), supporting both non-MLA ($d_{qk}=d_{vo}=128$) and MLA shapes ($d_{qk}=192$, $d_{vo}=128$). These kernels are based on the latest CUTLASS release Example 77.

Design for MLA Shape($d_{qk}=192$, $d_{vo}=128$)

Forward Kernel[PR]

Mainly based on the implementation of FMHA. The main changes are as follows:

• Added PipelineTmaAsyncMla to support producer to acquire specific bytes of data. since tileK is not equal to tileV when using MLA.
• Optimized shared memory usage.
• Optimized register usage to reduce register spilling.
• Introduced TileScheduler to optimize the causal mask case.
• Added a second kind of causal mask implementation (with Q at the end of the matrix) to handle N_Q ≠ N_K, which is more commonly used.

Backward Kernel[PR]

The backward FMHA kernel has several changes to handle the larger $d_{qk}$ size:

• Q_TILE=64 and KV_TILE=128
• Computing $S^{T}=QK$(instead of $S=KQ$), $dP^{T}=dOV$(instead of $dP=VdO$). $S^{T}$ and $dP^{T}$ are stored in TMEM with dimension 64.
• TMEM layout updates:
  • $dK$(128x192) | $dV$(128x128) | $dQ$(64x192)=$dP$(64x128) & $S$(64x128)
  • $dQ$ and $dP$ share the same 64-dimension address region, while $S$ uses a separate 64-dimension block.
  • $P$ is now stored in shared memory instead of TMEM to utilize available space.
• Corresponding layout updates of cute usage.

On-going and coming soon

• GQA support for backward FMHA kernels.
• Fix the leading padding buffer in forward kernels. In the current version of the CUTLASS forward implementation, a leading buffer is required for the Q, K, V, and O matrices. We plan to address the leading padding buffer in a future update.
• MLA decoding kernel support fore more GPU architectures(e.g. Ada)
• Incontiguous tensor support

Long-term optimization

• Performance improvements in both forward and backward MLA FMHA kernels.
• More GPU architectures support(e.g. Ada).

We are actively working on these optimizations. It aims to continuously provide support for NVIDIA's latest GPU hardware in FlashMLA, ensuring the community can perform efficient and scalable training and inference with DeepSeek models for next-generation AI workloads.

Thanks!

[beginlner]

Test in README needs update.

[beginlner]

LGTM!

[kiwi3shark]

Hi guys @uchihatmtkinu, I noticed there is a SM90 FMHA example offered by cutlass, so I wonder the feasibility to build corresponding SM90 dense prefill ops.

Besides, How hard would it be to support GQA?

[uchihatmtkinu]

Hi @kiwi3shark , it should be feasible on SM90 with CUTLASS example for MLA and GQA. However, we have no plan to implement them at current.