Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/topk-moe.cu (CUDA)

Code before:
            if (val > max_val || (val == max_val && expert < max_expert)) {
                max_val    = val;
                max_expert = expert;
            }
        }

        if ((max_expert & (WARP_SIZE - 1)) == threadIdx.x) {
            wt[max_expert / WARP_SIZE] = -INFINITY;

            wt_shared_ptr[k] = max_val;
            ids[k]           = max_expert;
            if constexpr (with_norm) {
                wt_sum += max_val;
            }
        }
    }

    if constexpr (with_norm) {
        wt_sum              = warp_reduce_sum(wt_sum);
        const float inv_sum = 1.0f / wt_sum;

        if (threadIdx.x == 0) {
            for (int i = 0; i < n_expert_used; i++) {
                wt_shared_ptr[i] = wt_shared_ptr[i] * inv_sum;
            }
        }
    }

    if (threadIdx.x == 0) {
        for (int i = 0; i < n_expert_used; i++) {
            weights[i] = wt_shared_ptr[i];
        }
    }
}

template <bool with_norm>
static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx,
                                 const float *               logits,
                                 float *                     weights,
                                 int32_t *                   ids,
                                 const int                   n_rows,
                                 const int                   n_expert,
                                 const int                   n_expert_used) {
    const int    rows_per_block = 4;
    dim3         grid_dims((n_rows + rows_per_block - 1) / rows_per_block, 1, 1);
    dim3         block_dims(WARP_SIZE, rows_per_block, 1);
    cudaStream_t stream = ctx.stream();

    const int nbytes_shared = n_expert_used * rows_per_block * sizeof(float);

    switch (n_expert) {

Review comment (suggestion):
```suggestion
        for (int i = threadIdx.x; i < n_expert_used; i += WARP_SIZE) {
            wt_shared_ptr[i] = wt_shared_ptr[i] * inv_sum;
        }
```

Use this pattern to reduce memory accesses by a factor of up to 32, you can also apply it to the writeback to VRAM.

Code after:
            if (val > max_val || (val == max_val && expert < max_expert)) {
                max_val    = val;
                max_expert = expert;
            }
        }

        if ((max_expert & (WARP_SIZE - 1)) == threadIdx.x) {
            wt[max_expert / WARP_SIZE] = -INFINITY;

            wt_shared_ptr[k] = max_val;
            ids[k]           = max_expert;
            if constexpr (with_norm) {
                wt_sum += max_val;
            }
        }
    }

    if constexpr (with_norm) {
        wt_sum              = warp_reduce_sum(wt_sum);
        const float inv_sum = 1.0f / wt_sum;

        for (int i = threadIdx.x; i < n_expert_used; i += WARP_SIZE) {
            wt_shared_ptr[i] = wt_shared_ptr[i] * inv_sum;
        }
    }

    for (int i = threadIdx.x; i < n_expert_used; i += WARP_SIZE) {
        weights[i] = wt_shared_ptr[i];
    }
}

template <bool with_norm>
static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx,
                                 const float *               logits,
                                 float *                     weights,
                                 int32_t *                   ids,
                                 const int                   n_rows,
                                 const int                   n_expert,
                                 const int                   n_expert_used) {
    const int    rows_per_block = 4;
    dim3         grid_dims((n_rows + rows_per_block - 1) / rows_per_block, 1, 1);
    dim3         block_dims(WARP_SIZE, rows_per_block, 1);
    cudaStream_t stream = ctx.stream();

    const int nbytes_shared = n_expert_used * rows_per_block * sizeof(float);

    switch (n_expert) {
        case 1:
            topk_moe_cuda<1, with_norm>
                <<<grid_dims, block_dims, nbytes_shared, stream>>>(logits, weights, ids, n_rows, n_expert_used);
            break;