char-fork/llama.cpp

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

Alexander Heislerggml-cuda: provide static workspace for cuBLAS handles (#26574)d9b6be07d

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
4.8 KiB125 linesraw
1#include "out-prod.cuh"
2
3#include <cstdint>
4
5static __global__ void k_compute_out_prod_ptrs(
6        const float * src0_d, const float * src1_d, float * dst_d,
7        const float ** ptrs_a, const float ** ptrs_b, float ** ptrs_c,
8        const int64_t ne2, const int64_t ne3,
9        const int64_t dps2, const int64_t dps3,
10        const size_t s02, const size_t s03,
11        const size_t s12, const size_t s13,
12        const size_t s2,  const size_t s3) {
13    const int64_t i2 = blockIdx.x*blockDim.x + threadIdx.x;
14    const int64_t i3 = blockIdx.y*blockDim.y + threadIdx.y;
15
16    if (i2 >= ne2 || i3 >= ne3) {
17        return;
18    }
19
20    const int64_t idx = i3*ne2 + i2;
21
22    ptrs_a[idx] = src0_d + (i3/dps3)*s03 + (i2/dps2)*s02;
23    ptrs_b[idx] = src1_d +  i3      *s13 +  i2      *s12;
24    ptrs_c[idx] = dst_d  +  i3      *s3  +  i2      *s2;
25}
26
27void ggml_cuda_out_prod(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
28    const ggml_tensor * src0 = dst->src[0];
29    const ggml_tensor * src1 = dst->src[1];
30
31    GGML_TENSOR_BINARY_OP_LOCALS
32
33    GGML_ASSERT(src0->type == GGML_TYPE_F32);
34    GGML_ASSERT(src1->type == GGML_TYPE_F32);
35    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
36
37    GGML_ASSERT(ne01 == ne11);
38    GGML_ASSERT(ne0 == ne00);
39    GGML_ASSERT(ne1 == ne10);
40
41    GGML_ASSERT(ne2 % src0->ne[2] == 0);
42    GGML_ASSERT(ne3 % src0->ne[3] == 0);
43
44    GGML_ASSERT(ne2 == src1->ne[2]);
45    GGML_ASSERT(ne3 == src1->ne[3]);
46
47    const float * src0_d = (const float *) src0->data;
48    const float * src1_d = (const float *) src1->data;
49    float       *  dst_d = (float       *)  dst->data;
50
51    cudaStream_t   stream = ctx.stream();
52    cublasHandle_t handle = ctx.cublas_handle();
53
54    const float alpha = 1.0f;
55    const float beta = 0.0f;
56
57    const int64_t lda = nb01 / sizeof(float);
58    const int64_t ldc = nb1  / sizeof(float);
59
60    const bool src1_T = ggml_is_transposed(src1);
61    const cublasOperation_t src1_cublas_op =  src1_T ? CUBLAS_OP_N : CUBLAS_OP_T;
62    const int64_t           ldb            = (src1_T ?        nb10 :        nb11) /  sizeof(float);
63    GGML_ASSERT(                             (src1_T ?        nb11 :        nb10) == sizeof(float));
64
65    // data strides in dimensions 2/3
66    const size_t s02 = nb02 / sizeof(float);
67    const size_t s03 = nb03 / sizeof(float);
68    const size_t s12 = nb12 / sizeof(float);
69    const size_t s13 = nb13 / sizeof(float);
70    const size_t s2  = nb2  / sizeof(float);
71    const size_t s3  = nb3  / sizeof(float);
72
73    // dps == dst per src0, used for group query attention
74    const int64_t dps2 = ne2 / ne02;
75    const int64_t dps3 = ne3 / ne03;
76
77    if (dps2 == 1 && ne2 > 1) {
78        // src0 has uniform stride s02 along dim 2; batch the inner loop with a strided GEMM
79        GGML_ASSERT(ne2 <= std::numeric_limits<int>::max());
80        const int batch_count = (int) ne2;
81        for (int64_t i3 = 0; i3 < ne3; ++i3) {
82            CUBLAS_CHECK(
83                cublasSgemmStridedBatched(handle, CUBLAS_OP_N, src1_cublas_op,
84                        ne0, ne1, ne01,
85                        &alpha, src0_d + (i3/dps3)*s03, lda, s02,
86                                src1_d +  i3     *s13, ldb, s12,
87                        &beta,  dst_d  +  i3     *s3,  ldc, s2,
88                        batch_count));
89        }
90    } else if (ne2 > 1 || ne3 > 1) {
91        // dps2 > 1 (src0 broadcast along dim 2 with non-uniform stride) or multiple GEMMs
92        // along dim 3: compute per-GEMM pointers on the device and use a single batched GEMM.
93        GGML_ASSERT(ne3 > 0);
94        GGML_ASSERT(ne2 <= (int64_t) std::numeric_limits<int>::max() / ne3);
95        const int batch_count = (int) (ne2 * ne3);
96
97        ggml_cuda_pool_alloc<const float *> ptrs_a(ctx.pool(), batch_count);
98        ggml_cuda_pool_alloc<const float *> ptrs_b(ctx.pool(), batch_count);
99        ggml_cuda_pool_alloc<      float *> ptrs_c(ctx.pool(), batch_count);
100
101        const dim3 block_dims(16, 16);
102        const dim3 grid_dims((ne2 + block_dims.x - 1)/block_dims.x, (ne3 + block_dims.y - 1)/block_dims.y);
103        k_compute_out_prod_ptrs<<<grid_dims, block_dims, 0, stream>>>(
104            src0_d, src1_d, dst_d,
105            ptrs_a.get(), ptrs_b.get(), ptrs_c.get(),
106            ne2, ne3, dps2, dps3, s02, s03, s12, s13, s2, s3);
107        CUDA_CHECK(cudaGetLastError());
108
109        CUBLAS_CHECK(
110            cublasSgemmBatched(handle, CUBLAS_OP_N, src1_cublas_op,
111                    ne0, ne1, ne01,
112                    &alpha, ptrs_a.get(), lda,
113                            ptrs_b.get(), ldb,
114                    &beta,  ptrs_c.get(), ldc,
115                    batch_count));
116    } else {
117        // ne2 == 1 && ne3 == 1: single GEMM
118        CUBLAS_CHECK(
119            cublasSgemm(handle, CUBLAS_OP_N, src1_cublas_op,
120                    ne0, ne1, ne01,
121                    &alpha, src0_d, lda,
122                            src1_d, ldb,
123                    &beta,  dst_d,  ldc));
124    }
125}