Example 77 add blackwell fmha bwd for MLA shape

[uchihatmtkinu]

Description

Add MLA shape($d_{qk}=192$,$d_{vo}=128$) of attention bwd kernels in example 77. Both fp8 and fp16 are supported.

Design of MLA bwd kernel

fmha_bwd_mla has several changes to handle the larger $d_{qk} size:

• Q_TILE=64 and KV_TILE=128
• Calculate $S^{T}=QK$(instead of$S=KQ$),$dP^{T}=dOV$(instead of $dP=VdO$). $S^{T}$ and $dP^{T}$ are stored in TMEM with dim=64
• TMEM layout change: * $dK$ (128x192) | $dV$ (128x128) | $dQ$ (64x192)= $dP$ (64x128) & $S$ (64x128)
  • $dQ$ and $dP$ share the same address with 64dim and $S$ uses another 64dim
  • $P$ will be stored to SMEM instead of TMEM since the space is available.
• Shared Memory
  • Store $P$ in the smem
• Corresponding layout adaption of cute

Examples

The bwd MLA also uses 77_blackwell_fmha_bwd_fp(type), to trigger it, using --d=192 --d_vo=128.

Performance Data

setting	$d_{qk}$=$d_{vo}$=128	$d_{qk}$=192, $d_{vo}$=128
b=16,h=10,q=1024,k=1024	690T	576T
b=4,h=128,q=4096,k=4096	1015T	874T

[uchihatmtkinu]

@thakkarV Hi, I just add some new commits which fix the casual mask when IsQBegin==false, which means the Q and K are aligned at tail when seqlen_q != seqlen_kv. Please have a look.

[thakkarV]

@IonThruster @hwu36 please fwd to Richard?

[thakkarV]

@v0i0 too in case you have time to review