perf: dynamic split-k for MLA

[yzh119]

#804 didn't implement split-k, which might result in performance degradation if concurrency is not large enough. This PR fixes issue.

We implemented the v2 scheduler and write-through optimization mentioned in our paper (section 3.3 and appendix in D.2) for load-balancing.

In an early PR (#72), we turned off cudaLaunchCooperativeKernels and grid.sync() because we are not sure whether it's compatible with CUDAGraph. This PR adds them back again for grid synchronization, to save some kernel launch overhead.

Benchmark

On H100 SXM5 80GB (3352 GB/s), this PR:

Config: batch_size=1, seq_len=1024, num_heads=16
Memory bandwidth: 22.33 GB/s
Config: batch_size=16, seq_len=1024, num_heads=16
Memory bandwidth: 330.72 GB/s
Config: batch_size=32, seq_len=1024, num_heads=16
Memory bandwidth: 638.73 GB/s
Config: batch_size=64, seq_len=1024, num_heads=16
Memory bandwidth: 1188.90 GB/s
Config: batch_size=1, seq_len=2048, num_heads=16
Memory bandwidth: 40.74 GB/s
Config: batch_size=16, seq_len=2048, num_heads=16
Memory bandwidth: 592.77 GB/s
Config: batch_size=32, seq_len=2048, num_heads=16
Memory bandwidth: 1112.83 GB/s
Config: batch_size=64, seq_len=2048, num_heads=16
Memory bandwidth: 1506.01 GB/s
Config: batch_size=1, seq_len=4096, num_heads=16
Memory bandwidth: 72.53 GB/s
Config: batch_size=16, seq_len=4096, num_heads=16
Memory bandwidth: 1007.80 GB/s
Config: batch_size=32, seq_len=4096, num_heads=16
Memory bandwidth: 1438.99 GB/s
Config: batch_size=64, seq_len=4096, num_heads=16
Memory bandwidth: 1730.62 GB/s
Config: batch_size=1, seq_len=8192, num_heads=16
Memory bandwidth: 120.74 GB/s
Config: batch_size=16, seq_len=8192, num_heads=16
Memory bandwidth: 1340.86 GB/s
Config: batch_size=32, seq_len=8192, num_heads=16
Memory bandwidth: 1689.36 GB/s
Config: batch_size=64, seq_len=8192, num_heads=16
Memory bandwidth: 1901.26 GB/s
Config: batch_size=1, seq_len=16384, num_heads=16
Memory bandwidth: 177.94 GB/s
Config: batch_size=16, seq_len=16384, num_heads=16
Memory bandwidth: 1619.51 GB/s
Config: batch_size=32, seq_len=16384, num_heads=16
Memory bandwidth: 1876.50 GB/s
Config: batch_size=64, seq_len=16384, num_heads=16
Memory bandwidth: 2010.58 GB/s
Config: batch_size=1, seq_len=32768, num_heads=16
Memory bandwidth: 231.70 GB/s
Config: batch_size=16, seq_len=32768, num_heads=16
Memory bandwidth: 1835.16 GB/s
Config: batch_size=32, seq_len=32768, num_heads=16
Memory bandwidth: 1997.24 GB/s
Config: batch_size=64, seq_len=32768, num_heads=16
Memory bandwidth: 2067.99 GB/s

Before this PR:

Config: batch_size=1, seq_len=1024, num_heads=16
Memory bandwidth: 15.46 GB/s
Config: batch_size=16, seq_len=1024, num_heads=16
Memory bandwidth: 238.49 GB/s
Config: batch_size=32, seq_len=1024, num_heads=16
Memory bandwidth: 472.44 GB/s
Config: batch_size=64, seq_len=1024, num_heads=16
Memory bandwidth: 929.12 GB/s
Config: batch_size=1, seq_len=2048, num_heads=16
Memory bandwidth: 15.47 GB/s
Config: batch_size=16, seq_len=2048, num_heads=16
Memory bandwidth: 250.71 GB/s
Config: batch_size=32, seq_len=2048, num_heads=16
Memory bandwidth: 500.21 GB/s
Config: batch_size=64, seq_len=2048, num_heads=16
Memory bandwidth: 996.37 GB/s
Config: batch_size=1, seq_len=4096, num_heads=16
Memory bandwidth: 16.36 GB/s
Config: batch_size=16, seq_len=4096, num_heads=16
Memory bandwidth: 257.59 GB/s
Config: batch_size=32, seq_len=4096, num_heads=16
Memory bandwidth: 515.88 GB/s
Config: batch_size=64, seq_len=4096, num_heads=16
Memory bandwidth: 1035.55 GB/s
Config: batch_size=1, seq_len=8192, num_heads=16
Memory bandwidth: 16.37 GB/s
Config: batch_size=16, seq_len=8192, num_heads=16
Memory bandwidth: 261.47 GB/s
Config: batch_size=32, seq_len=8192, num_heads=16
Memory bandwidth: 524.76 GB/s
Config: batch_size=64, seq_len=8192, num_heads=16
Memory bandwidth: 1054.54 GB/s
Config: batch_size=1, seq_len=16384, num_heads=16
Memory bandwidth: 16.50 GB/s
Config: batch_size=16, seq_len=16384, num_heads=16
Memory bandwidth: 263.69 GB/s
Config: batch_size=32, seq_len=16384, num_heads=16
Memory bandwidth: 528.89 GB/s
Config: batch_size=64, seq_len=16384, num_heads=16
Memory bandwidth: 1064.87 GB/s
Config: batch_size=1, seq_len=32768, num_heads=16
Memory bandwidth: 16.45 GB/s
Config: batch_size=16, seq_len=32768, num_heads=16
Memory bandwidth: 264.66 GB/s
Config: batch_size=32, seq_len=32768, num_heads=16
Memory bandwidth: 530.87 GB/s
Config: batch_size=64, seq_len=32768, num_heads=16
Memory bandwidth: 1070.93 GB/s

[yzh119]

Note that this PR uses f32 as the intermediate data type for output matrix in split-k.

Using 16bit data type (e.g. f16/bf16) might be faster, but will result in some accuracy loss, especially for bf16+long context.

[zhyncs]

This PR is amazing!!