Gather indices for batched matrix

[masahi]

Hi, I'm looking at gather_scatter_fusion.cu example https://github.com/NVIDIA/cutlass/blob/main/examples/36_gather_scatter_fusion/gather_scatter_fusion.cu and wondering if it is possible to gather along the batch dimension.

For example, I have a batched matrix of shape (B1, K, N) and in a GEMM kernel I want to operate on a "virtual" matrix of shape (B2, K, N), which is gathered from the first matrix along the batch dimension. I have a 1D index vector of shape (B2,) which tells which matrix slice to pick among the batch B1.

Is this possible? cc @hwu36

[hwu36]

yes.

https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/gemm/kernel/gemm_universal.h#L506

threadblock_tile_offset.k() is essentially batch id.

[masahi]

Do you mean I need to modify this code to support gathering from arbitrary index? i.e., I cannot use ptr_gather_B_indices https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/gemm/kernel/gemm_universal.h#L144 for this use case?

[hwu36]

you can reuse the existing index array, but you need to load batch id as some_array[threadblock_tile_offset.k()]

[masahi]

Yes, my ask was whether or not I can do that without modifying any CUTLASS source code itself. It seems source code modification is required.

[hwu36]

correct. have fun. :)