perf: memory efficient deepseek mla fused page-attention kernel

[yzh119]

Description:

This PR implements a memory efficient fused Deepseek MLA PageAttention kernel for decode, prefill, and chunked prefill operations.
MLA computation after matrix absorption can be described as a (multi-query attention) MQA kernel, with same K/V cache and special head dimensions: head_dim_qk=576, head_dim_vo=512.

Background:
For Deepseek v2/3, the large head_dim (512) makes it challenging to store the output tensor in registers when using tensor cores. A previous approach (#551) split head_dim across gridDim.y, which led to two main issues:

• Re-computation Overhead: Multiple blocks had to redundantly compute the Q*K operation.
• Memory Access Latency: Using multiple blocks prevented shared memory usage, forcing reliance on slower L2 for KV-Cache accesses.

New Design:
We use head-group fusion to increase the operational intensity of the kernel (appendix A in the paper ). To address the large head-dimension issue, we redesign the kernel (diagram below) to use two warp groups (WG1 and WG2, each with 4 warps) per CTA:

• QK Computation: Each warp group processes half of the CTA_TILE_KV dimension, eliminating redundant computations.
• Shared Memory Broadcast: Local QK results are written to shared memory (reusing the kpe buffer) to efficiently broadcast data for PV computation.
• PV Computation: The head dimension is split between the warp groups, with each computing half.

The maximum CTA_TILE_Q (bounded by register files) is 64, and in this case the maximum CTA_TILE_KV (bounded by shared memory limit) is 64 (for H100) and 32 (for A100) when number of pipeline stages is set to 2.

For large num_local_heads such as 128 (if no TP and no MTP), we create a cluster of size 2, and the upper half (64) and lower half (64) are dispatched to two SMs in a cluster, and we can use software-managed multicasting (for large page size), we leave it to later PR.

Benchmark results

Decoding Memory bandwidth on H100 SXM5 (Peak bandwidth = 3352 GB/s):

Config: batch_size=768, seq_len=1024, num_heads=16
Memory bandwidth: 2002.11 GB/s
Config: batch_size=768, seq_len=1024, num_heads=32
Memory bandwidth: 2035.59 GB/s
Config: batch_size=768, seq_len=1024, num_heads=64
Memory bandwidth: 2082.20 GB/s
Config: batch_size=768, seq_len=2048, num_heads=16
Memory bandwidth: 2064.97 GB/s
Config: batch_size=768, seq_len=2048, num_heads=32
Memory bandwidth: 2080.99 GB/s
Config: batch_size=768, seq_len=2048, num_heads=64
Memory bandwidth: 2082.78 GB/s

[seanxcwang]

test on H20（peak bandwidth 4TB/s）

Config: batch_size=768, seq_len=1024, num_heads=16
Memory bandwidth: 668.77 GB/s
Config: batch_size=768, seq_len=1024, num_heads=32
Memory bandwidth: 685.71 GB/s
Config: batch_size=768, seq_len=1024, num_heads=64
Memory bandwidth: 718.93 GB/s
Config: batch_size=768, seq_len=2048, num_heads=16
Memory bandwidth: 669.80 GB/s
Config: batch_size=768, seq_len=2048, num_heads=32
Memory bandwidth: 678.41 GB/s
Config: batch_size=768, seq_len=2048, num_heads=64
Memory bandwidth: 695.46 GB/s

Is there any performance improvement plan for H20, any suggestion will be helpful

[yzh119]

@seanxcwang I'm working on the FA3 version (this PR still uses FA2 template, not overlapping softmax and computation, and doesn't use wgmma) of this kernel, which should also benefit H20, please stay tuned :)

[yzh119]

Hi @seanxcwang , #887 should improve performance on H20.