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

Code before:
    const int channel_x   = has_ids ? expert_idx : (channel_dst / channel_ratio);
    const int channel_y   = channel_dst;
    const int sample_dst  = blockIdx.z;
    const int sample_x    = sample_dst / sample_ratio;
    const int sample_y    = sample_dst;

    x   += int64_t(sample_x)  *stride_sample_x   + channel_x  *stride_channel_x  + row0*stride_row ;
    y   += int64_t(sample_y)  *stride_sample_y   + (has_ids ? 0 : channel_y  *stride_channel_y);
    dst += int64_t(sample_dst)*stride_sample_dst + (has_ids ? 0 : channel_dst*stride_channel_dst);

    const float2 * y2 = (const float2 *) y;

    extern __shared__ char data_mmv[];

    char * shmem_base = data_mmv;
    int  * slot_map   = (int *) shmem_base;
    char * compute_base = has_ids ? (shmem_base + GGML_PAD(cols_per_block, 16) * sizeof(int)) : shmem_base;

    tile_C C[ntA][ntB];

    T * tile_xy = (T *) compute_base + threadIdx.y*(tile_A::I * tile_k_padded);

    if constexpr (has_ids) {
        int found = 0;

        for (int j = threadIdx.y; j < cols_per_block; j += nwarps) {
            const int32_t * __restrict__ id_row = ids + j*stride_row_id;

            if (threadIdx.x == 0) {
                slot_map[j] = -1;
            }

            for (int k_base = 0; k_base < nchannels_dst; k_base += warp_size) {
                int k = k_base + threadIdx.x;
                int match = (k < nchannels_dst) && (id_row[k*stride_col_id] == expert_idx);

                unsigned mask = __ballot_sync(0xffffffff, match);
                if (mask) {
                    int leader = __ffs(mask) - 1;
                    if (threadIdx.x == leader) {
                        slot_map[j] = k_base + leader;
                    }
                    found = 1;
                    break;
                }
            }
        }

        if (!__syncthreads_or(found)) {
            return;
        }

Review comment (suggestion):
This is a bad pattern because the compiler cannot unroll the loop. You should do it like this instead:

```
#pragma unroll
for (int j0 = 0; j0 < cols_per_block; j0 += nwarps) {
    const int j = j0 + threadIdx.y;
}
```

Code after:
    const int channel_x   = has_ids ? expert_idx : (channel_dst / channel_ratio);
    const int channel_y   = channel_dst;
    const int sample_dst  = blockIdx.z;
    const int sample_x    = sample_dst / sample_ratio;
    const int sample_y    = sample_dst;

    x   += int64_t(sample_x)  *stride_sample_x   + channel_x  *stride_channel_x  + row0*stride_row ;
    y   += int64_t(sample_y)  *stride_sample_y   + (has_ids ? 0 : channel_y  *stride_channel_y);
    dst += int64_t(sample_dst)*stride_sample_dst + (has_ids ? 0 : channel_dst*stride_channel_dst);

    const float2 * y2 = (const float2 *) y;

    extern __shared__ char data_mmv[];

    char * shmem_base = data_mmv;
    int  * slot_map   = (int *) shmem_base;
    char * compute_base = has_ids ? (shmem_base + GGML_PAD(cols_per_block, 16) * sizeof(int)) : shmem_base;

    tile_C C[ntA][ntB];

    T * tile_xy = (T *) compute_base + threadIdx.y*(tile_A::I * tile_k_padded);

    if constexpr (has_ids) {
        int found = 0;

        for (int j0 = 0; j0 < cols_per_block; j0 += nwarps) {
            const int j = j0 + threadIdx.y;
            const int32_t * __restrict__ id_row = ids + j*stride_row_id;

            if (threadIdx.x == 0) {
                slot_map[j] = -1;
            }

            for (int k = threadIdx.x; k < nchannels_dst; k += warp_size) {
                int match = id_row[k*stride_col_id] == expert_idx;

                if (match) {
                    slot_map[j] = k;
                    found = 1;
                    break;
                }
            }
        }

        if (!__syncthreads_or(found)) {
            return;
        }
    }


    for (int col = threadIdx.y*warp_size + threadIdx.x; col < ncols; col += nwarps*warp_size) {