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 size_t * src0_nb, 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) {
// TODO: handle this in a consistent and simpler way after AMD MFMA support has been added
// The ideal one shall be: #if defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
#if (!defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)) || defined(AMD_WMMA_AVAILABLE)
#if defined(AMD_WMMA_AVAILABLE)
    // Special case for tf32, just dummy mma layout as wmma doesn't support it.
    constexpr int tile_B_I = std::is_same_v<T, float> ? 8 : 16;
    constexpr int tile_C_J = std::is_same_v<T, float> ? 8 : 16;
    typedef tile<16,       8, T>     tile_A;
    typedef tile<tile_B_I, 8, T>     tile_B;
    typedef tile<16,       tile_C_J, float> tile_C;

    constexpr bool a_supported = tile_A::supported();
    constexpr bool b_supported = tile_B::supported();
    constexpr bool c_supported = tile_C::supported();
    constexpr bool supported = a_supported && b_supported && c_supported;
#else
    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();
    constexpr bool supported = I_16_supported || I_32_supported;

    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;
#endif // defined(AMD_WMMA_AVAILABLE)
    if constexpr (!supported) {

Review comment (suggestion):
```suggestion
```

I'm not yet sure how exactly I want to do it.

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 size_t * src0_nb, 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) {
// TODO: handle this in a consistent and simpler way after AMD MFMA support has been added
#if (!defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)) || defined(AMD_WMMA_AVAILABLE)
#if defined(AMD_WMMA_AVAILABLE)
    // Special case for tf32, just dummy mma layout as wmma doesn't support it.
    constexpr int tile_B_I = std::is_same_v<T, float> ? 8 : 16;
    constexpr int tile_C_J = std::is_same_v<T, float> ? 8 : 16;
    typedef tile<16,       8, T>     tile_A;
    typedef tile<tile_B_I, 8, T>     tile_B;
    typedef tile<16,       tile_C_J, float> tile_C;

    constexpr bool a_supported = tile_A::supported();
    constexpr bool b_supported = tile_B::supported();
    constexpr bool c_supported = tile_C::supported();
    constexpr bool supported = a_supported && b_supported && c_supported;
#else
    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();
    constexpr bool supported = I_16_supported || I_32_supported;

    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;
#endif // defined(AMD_WMMA_AVAILABLE)
    if constexpr (!supported) {
        NO_DEVICE_CODE;