Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/mmf.cuh (CUDA)

Code before:

using namespace ggml_cuda_mma;

#define MMF_ROWS_PER_BLOCK 32

struct mmf_ids_data {
    const int32_t * ids_src_compact = nullptr;
    const int32_t * ids_dst_compact = nullptr;
    const int32_t * expert_bounds_dev = nullptr;
    int n_experts = 0;
    int sis1 = 0;
};

void ggml_cuda_mul_mat_f(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst);

bool ggml_cuda_should_use_mmf(enum ggml_type type, int cc, int warp_size, const int64_t * scr0_ne, const int src1_ncols, bool mul_mat_id);

template <typename T, int rows_per_block, int cols_per_block, int nwarps, bool has_ids>
static __device__ void mul_mat_f_impl(
        const T * __restrict__ x, const float * __restrict__ y, const int32_t * __restrict__ ids, float * __restrict__ dst,
        const int ncols, const int ncols_dst_total, const int nchannels_dst, const int stride_row, const int stride_col_y, const int stride_col_dst,
        const int stride_col_id, const int stride_row_id,
        const int channel_ratio, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
        const int sample_ratio, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst) {
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
#if __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA
    typedef tile<32, 8, T>     tile_A;
    typedef tile< 8, 8, T>     tile_B;
    typedef tile<32, 8, float> tile_C;
#else
    // In principle also possible to use tiles with I == 32, the performance difference is ~1%.
    typedef tile<16, 8, T>     tile_A;
    typedef tile< 8, 8, T>     tile_B;
    typedef tile<16, 8, float> tile_C;
#endif // __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA

    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
    constexpr int tile_k_padded = warp_size + 4;
    constexpr int ntA = rows_per_block / tile_A::I;
    constexpr int ntB = (cols_per_block + tile_B::I - 1) / tile_B::I;

    const int row0        = blockIdx.x * rows_per_block;

    int expert_idx = 0;
    int col_base = 0;

    const int channel_dst = has_ids ? 0 : blockIdx.y;

    if constexpr (has_ids) {
        // experts + tiles of ncols_dst are packed in the y dimension
        int col_tiles = (ncols_dst_total + cols_per_block - 1) / cols_per_block;

Review comment (suggestion):
Can we make this more generic like `using tile_A = mma_traits::get_mmf_tile()::tileA`. In the future when we add more architectures like AMD or blackwell, it might be easier to experiment with different tile configurations

Code after:

using namespace ggml_cuda_mma;

#define MMF_ROWS_PER_BLOCK 32

struct mmf_ids_data {
    const int32_t * ids_src_compact = nullptr;
    const int32_t * ids_dst_compact = nullptr;
    const int32_t * expert_bounds_dev = nullptr;
    int n_experts = 0;
    int sis1 = 0;
};

void ggml_cuda_mul_mat_f(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst);

bool ggml_cuda_should_use_mmf(enum ggml_type type, int cc, int warp_size, const int64_t * scr0_ne, const int src1_ncols, bool mul_mat_id);

template <typename T, int rows_per_block, int cols_per_block, int nwarps, bool has_ids>
__launch_bounds__(ggml_cuda_get_physical_warp_size()*nwarps, 1)
static __global__ void mul_mat_f(
        const T * __restrict__ x, const float * __restrict__ y, const int32_t * __restrict__ ids, float * __restrict__ dst,
        const int ncols, const int ncols_dst_total, const int nchannels_dst, const int stride_row, const int stride_col_y, const int stride_col_dst,
        const int stride_col_id, const int stride_row_id,
        const int channel_ratio, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
        const int sample_ratio, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst) {
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
    constexpr bool I_16_supported = tile<16, 8, T>::supported() && tile<16, 8, float>::supported();
    constexpr bool I_32_supported = tile<32, 8, T>::supported() && tile<32, 8, float>::supported();

    if (!I_16_supported && !I_32_supported) {
        NO_DEVICE_CODE;
        return;
    }

    constexpr int I_preferred = I_16_supported ? 16 : 32; // For Turing MMA both work but 16 is ~1% faster.

    typedef tile<I_preferred, 8, T>     tile_A;
    typedef tile<8,           8, T>     tile_B;
    typedef tile<I_preferred, 8, float> tile_C;

    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
    constexpr int tile_k_padded = warp_size + 4;
    constexpr int ntA = rows_per_block / tile_A::I;
    constexpr int ntB = (cols_per_block + tile_B::I - 1) / tile_B::I;

    const int row0        = blockIdx.x * rows_per_block;

    int expert_idx = 0;
    int col_base = 0;

    const int channel_dst = has_ids ? 0 : blockIdx.y;