Skip to content

Add Blackwell MLA forward (shape: d=192, dv=128) implementation#2472

Merged
hwu36 merged 1 commit into
NVIDIA:mainfrom
dianzhangchen:feat/blackwell_mla_fwd
Jul 18, 2025
Merged

Add Blackwell MLA forward (shape: d=192, dv=128) implementation#2472
hwu36 merged 1 commit into
NVIDIA:mainfrom
dianzhangchen:feat/blackwell_mla_fwd

Conversation

@dianzhangchen

Copy link
Copy Markdown
Contributor

Description

Add the prefill phase MLA with shape $d_{qk}=192 , d_{vo}=128$ for the Blackwell architecture in example 77. It supports both FP16 and FP8.

Design of MLA fwd kernel

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

  • Add PipelineTmaAsyncMla to support producer to acquire specific bytes of data. since tileK is not equal to tileV when using MLA.
  • Optimize shared memory usage.
  • Optimize register usage to reduce register spilling.
  • Add TileScheduler to optimize the causal mask case.
  • Add 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.

Examples

Example 77 has new bins named 77_blackwell_mla_fwd_fp16 and 77_blackwell_mla_fwd_fp8.

@thakkarV

Copy link
Copy Markdown
Contributor

@v0i0 are you able to help review?

@hwu36
hwu36 merged commit 9baa06d into NVIDIA:main Jul 18, 2025
Albresky pushed a commit to Albresky/cutlass that referenced this pull request Oct 11, 2025
guocuimi pushed a commit to vectorch-ai/cutlass that referenced this pull request Nov 6, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants