[QST] Split-k in hopper gather scatter gemm

What is your question?
How to use split-k in hopper gather scatter gemm?

I am trying to optimize a skinny gemm on hopper with the gather scatter kernel (adapting from example 52_hopper_gather_scatter_fusion). The default kernel doesn't perform the best for different index_sizes. I have the following two use cases.

use case 1:
M x N x K = 128 x 32768 x 8192; gather from N, index_size = 8192-16384

use case 2:
M x N x K = 128 x 8192 x 32768; gather from K, index_size = 8192-16384

The profiler suggests that using split-k of 1-4 is often sufficient to maximize the performance of skinny gemms (different k depending on the problem size). This worked well with Ampere when adapting example 36_gather_scatter_fusion.

I read https://github.com/NVIDIA/cutlass/issues/1079 to use split-k for hopper by
using cutlass::gemm::StreamKScheduler as the tile scheduler and splits argument to be larger than 1. This did not work with the gather scatter hopper kernel in example 52. I got the following assertion in the file gather_gemm.hpp.

  static_assert(cute::is_void_v<TileScheduler_> or cute::is_same_v<TileScheduler_, PersistentScheduler>,
    "Non-persistent warp-specialized kernel does not support specializing the tile scheduler.");

I removed the assertion, and the code did compile; but the performance did not improve as expected

Can we use split-k in hopper with the gather scatter gemm kernel? If so, how do I do this and what needs to be changed? Any help on this matter is greatly appreciated.

Start with example 52_gather_scatter_fusion and gather_gemm.hpp, then compare the split-k approach described in issue #1079. Investigate the tile-scheduler assertion and the reported performance behavior for the two skinny GEMM cases. Done means establishing whether split-k is supported for this Hopper kernel and documenting or implementing the required changes.