Grouped GEMM with split k

[masahi]

Is grouped GEMM with split k supported? I'm interested in both 2.x and 3.x API.

The motivation is the blog by Meta https://pytorch.org/blog/accelerating-moe-model/#30-work-decomposition---splitk, which found that for small M problems typical in LLM inference, split K works really well.

@hwu36

[masahi]

Sorry just found #692. But maybe good to hear an update on this issue, given the recent interest in small M problems from LLM domain. Mixtral and LoRA might benefit from this.

[masahi]

I just learned that TRT-LLM has implemented split-K grouped gemm in https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/splitk_gemm_grouped.h