Device-side grouped GEMM

[osayamenja]

Hello!

Is there a way to do a device-side grouped (or batched) gemm? That is, my goal is to perform a set of GEMMs with non uniform sizes, flexible constraint as we can pad zeros, inside a kernel. I could serially invoke N cute::cooperative_gemm calls, but that may not be as optimal as highlighted in this GTC'24 talk.

Any suggestions for actualizing the grouped gemm, preferably for capability >= 70? Thanks.

[jackkosaian]

By "device-side," do you mean that you'd like to perform the grouped-GEMM from within an existing kernel, rather than launching two separate kernels?

[jackkosaian]

There is not currently a way to do this.

cuBLASDx provides some methods for doing this for basic GEMMs. However, I'm not sure that it has been extended to support grouped GEMM.

[thakkarV]

Note that CUTLASS collective already support this. There's no such thing as a separate "device side grouped GEMM". That's exactly what the collectives are

[osayamenja]

@thakkarV thanks for the suggestion, I had only been looking at cute gemm and had no idea about the Collective API in cutlass. I'll be doing some reading about that here :)

So, achieving this would mean constructing a collective in device code and configuring for grouped like this example, is this similar to what you had in mind?

[thakkarV]

yes, that example ultimately uses this collective: https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized.hpp

which you can use directly in your code on the device side.

[osayamenja]

thanks! @thakkarV @jackkosaian