char-fork/llama.cpp

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

PascalGgml/cuda col2im 1d (#24417)3a3edc9ac

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.0 KiB81 linesraw
1#include "col2im-1d.cuh"
2#include "convert.cuh"
3
4// col2im_1d: scatter-add GEMM columns to 1D signal (gather approach)
5// columns: [K*OC, T_in]  ->  output: [T_out, OC]
6// Supports F32, F16, BF16 data with F32 accumulator.
7
8template <typename T>
9static __global__ void col2im_1d_kernel(
10        const T * __restrict__ col,
11        T       * __restrict__ dst,
12        const int T_in, const uint3 T_out_fd,
13        const int OC, const int K, const int K_OC,
14        const int s0, const int p0, const int total) {
15
16    const int idx = threadIdx.x + blockIdx.x * blockDim.x;
17    if (idx >= total) return;
18
19    // dst layout: [T_out, OC], ne[0]=T_out fastest
20    const uint2 qr  = fast_div_modulo((uint32_t)idx, T_out_fd);  // qr.x = idx / T_out, qr.y = idx % T_out
21    const int oc    = (int)qr.x;
22    const int t_out = (int)qr.y;
23    const int t_abs = t_out + p0;  // absolute position in uncropped signal
24
25    // Gather: find all (t_in, k) where t_in*s + k == t_abs, 0 <= k < K
26    int t_in_min = (t_abs - K + s0) / s0;  // ceil((t_abs - K + 1) / s)
27    if (t_in_min < 0) t_in_min = 0;
28    int t_in_max = t_abs / s0;
29    if (t_in_max >= T_in) t_in_max = T_in - 1;
30
31    float sum = 0.0f;
32    for (int t_in = t_in_min; t_in <= t_in_max; t_in++) {
33        const int k = t_abs - t_in * s0;
34        // col layout: [K*OC, T_in], column index = oc * K + k
35        sum += ggml_cuda_cast<float>(col[(oc * K + k) + t_in * K_OC]);
36    }
37
38    dst[idx] = ggml_cuda_cast<T>(sum);
39}
40
41void ggml_cuda_op_col2im_1d(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
42    const ggml_tensor * src0 = dst->src[0];
43    cudaStream_t stream = ctx.stream();
44
45    GGML_ASSERT(ggml_is_contiguous(src0));
46
47    const int32_t s0 = ((const int32_t *)(dst->op_params))[0];
48    const int32_t OC = ((const int32_t *)(dst->op_params))[1];
49    const int32_t p0 = ((const int32_t *)(dst->op_params))[2];
50
51    const int K_OC = (int) src0->ne[0];
52    const int T_in = (int) src0->ne[1];
53    const int K    = K_OC / OC;
54    const int T_out = (int) dst->ne[0];
55
56    const uint3 T_out_fd = init_fastdiv_values((uint32_t)T_out);
57
58    const int total = T_out * OC;
59    const int block_size = 256;
60    const int num_blocks = (total + block_size - 1) / block_size;
61
62    switch (src0->type) {
63        case GGML_TYPE_F32: {
64            col2im_1d_kernel<<<num_blocks, block_size, 0, stream>>>(
65                (const float *)src0->data, (float *)dst->data,
66                T_in, T_out_fd, OC, K, K_OC, s0, p0, total);
67        } break;
68        case GGML_TYPE_F16: {
69            col2im_1d_kernel<<<num_blocks, block_size, 0, stream>>>(
70                (const half *)src0->data, (half *)dst->data,
71                T_in, T_out_fd, OC, K, K_OC, s0, p0, total);
72        } break;
73        case GGML_TYPE_BF16: {
74            col2im_1d_kernel<<<num_blocks, block_size, 0, stream>>>(
75                (const nv_bfloat16 *)src0->data, (nv_bfloat16 *)dst->data,
76                T_in, T_out_fd, OC, K, K_OC, s0, p0, total);
77        } break;
78        default:
79            GGML_ABORT("col2im_1d: unsupported type");
80    }
81}