[AG-GEMM] Add warp specialized ag-gemm that supports KDA and MLA up projections

[ShawnWeiChew]

This PR implements AG-GEMM for KDA and MLA up projection specific shapes.

09/09 Findings
We initially had results where TK did much worse than it does now. I think it was because I was running the kernels in succession with TK last. After reverting back to the old commit and adding a timeout between different kernels, the TK timings remained within ~5% of previous measurements

Notes:

• Compared to tuned configurations of Cutlass and TK's implementations of ag-gemm, we are within ±10% of their performance.
• TK uses an SM-specialized TMA load and multicast store to perform all gather, while our implementation and cutlass' is a warp-specialized version performs an all-gather with ancudaAsyncMemcpy operation
  • Confident that asyncMemcpy always lands before the corresponding TMA store, because the time spent waiting on the memcpy completion flag is the same as the time spent on loading an already-set flag (0.2us)
• Based on the above design choices, I think the general trends are that:
  • TK's SM specialized design performs better when the amount of communication work > computation work
  • Ours and cutlass' warp specialized design performs better when computation > communication. In our case, I also think it's because the bulk memcpy's are done in a per device manner, so it imposes some sort of restriction on the way we can traverse the A matrix.
• I decided to make a separate file as opposed to using the original ag-gemm / ag-gemm-blackwell path because:
  • This kernel use a different all-gather mechanism
  • This kernel accommodates a 1-CTA MMA (and potentially other KDA / MLA specific fusions down the road)
• My strategy for dealing with odd shapes (M=3072, 3584) was to pad them up to the tuned tile shape

KDA projection shapes (updated 09/09)

MLA projection shapes (updated 09/09)

logs

[MaoZiming]

@ShawnWeiChew Great work! " In our case, I also think it's because the bulk memcpy's are done in a per device manner, so it imposes some sort of restriction on the way we can traverse the A matrix."

Could you elaborate on this point?

[ShawnWeiChew]

@MaoZiming Because my implementation handles the copies serially, we have to work on device n's SP chunk before device n + 1's SP chunk. (I tried adjusting the number of copy streams that were running at once but that, at best, did not have a negative impact on performance)

For example, in a TP=8, global_m = 2048 case, we handle tile A in (256, K) chunk, rather than a (2048, K) chunk

Whereas in a regular GEMM, we could potentially make use of a traversal pattern like this:

I mentioned this because relative to a regular GEMM, the traversal pattern imposed by our implementation reduces the opportunity for reuse of the same chunk of B-tile in L2 cache, since we have to traverse the full B tile for each device's SP chunk.

[MaoZiming]

Others LGTM