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

Code before:
    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;
        }
    }


    for (int col = threadIdx.y*warp_size + threadIdx.x; col < ncols; col += nwarps*warp_size) {
        tile_A A[ntA][warp_size / tile_A::J];
#pragma unroll
        for (int itA = 0; itA < ntA; ++itA) {
#pragma unroll

Review comment (refactor):
Take a look at how `iex_used` is determined in `mmq_ids_helper`. You can exploit the fact that the expert you're looking for always appears either 0 or 1 times in `ids`. So if you condition the write to `slot_map` on the expert being found you have a guarantee that only a single thread will write to that location.

Code after:
    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) {
        tile_A A[ntA][warp_size / tile_A::J];
#pragma unroll
        for (int itA = 0; itA < ntA; ++itA) {
#pragma unroll
            for (int i = 0; i < tile_A::I; ++i) {
                tile_xy[i*tile_k_padded + threadIdx.x] = x[(itA*tile_A::I + i)*stride_row  + col];
            }
#pragma unroll