char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

anujjcuda: fuse MoE weighted expert reduction (#25952)3466812d1

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.4 KiB65 linesraw
1#include "moe-weighted-reduction.cuh"
2
3static __global__ void moe_weighted_reduction_f32(const float * __restrict__ experts,
4                                                  const float * __restrict__ expert_scale,
5                                                  const float * __restrict__ weights,
6                                                  float * __restrict__ dst,
7                                                  const int64_t n_embd,
8                                                  const int     n_expert_used) {
9    const int64_t token = blockIdx.x;
10    const int64_t col   = (int64_t) blockIdx.y * blockDim.x + threadIdx.x;
11    if (col >= n_embd) {
12        return;
13    }
14
15    const uint64_t first_row   = (uint64_t) token * n_expert_used;
16    const float    first_scale = expert_scale != nullptr ? expert_scale[first_row] : 1.0f;
17    float          sum         = (experts[first_row * n_embd + col] * first_scale) * weights[first_row];
18
19    for (int expert = 1; expert < n_expert_used; ++expert) {
20        const uint64_t row   = first_row + expert;
21        const float   scale = expert_scale != nullptr ? expert_scale[row] : 1.0f;
22        sum += (experts[row * n_embd + col] * scale) * weights[row];
23    }
24    dst[token * n_embd + col] = sum;
25}
26
27static void launch_moe_weighted_reduction(const float * experts,
28                                          const float * expert_scale,
29                                          const float * weights,
30                                          float *       dst,
31                                          int64_t       n_embd,
32                                          int64_t       n_tokens,
33                                          int           n_expert_used,
34                                          cudaStream_t  stream) {
35    constexpr int threads = 256;
36    const dim3 blocks(n_tokens, (n_embd + threads - 1) / threads, 1);
37    moe_weighted_reduction_f32
38        <<<blocks, threads, 0, stream>>>(experts, expert_scale, weights, dst, n_embd, n_expert_used);
39}
40
41void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx,
42                                         const ggml_tensor *         experts,
43                                         const ggml_tensor *         expert_scale,
44                                         const ggml_tensor *         weights,
45                                         ggml_tensor *               dst) {
46    GGML_ASSERT(experts->type == GGML_TYPE_F32);
47    GGML_ASSERT(weights->type == GGML_TYPE_F32);
48    GGML_ASSERT(expert_scale == nullptr || expert_scale->type == GGML_TYPE_F32);
49    GGML_ASSERT(dst->type == GGML_TYPE_F32);
50    GGML_ASSERT(ggml_is_contiguous(experts));
51    GGML_ASSERT(ggml_is_contiguous(weights));
52    GGML_ASSERT(expert_scale == nullptr || ggml_is_contiguous(expert_scale));
53    GGML_ASSERT(ggml_is_contiguous(dst));
54
55    const int64_t n_embd        = experts->ne[0];
56    const int64_t n_expert_used = experts->ne[1];
57    const int64_t n_tokens      = experts->ne[2] * experts->ne[3];
58    cudaStream_t  stream        = ctx.stream();
59
60    launch_moe_weighted_reduction((const float *) experts->data,
61                                  expert_scale ? (const float *) expert_scale->data : nullptr,
62                                  (const float *) weights->data,
63                                  (float *) dst->data, n_embd, n_tokens, (int) n_expert_used, stream);
64    CUDA_CHECK(cudaGetLastError());
65}