char-fork/llama.cpp

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

Titaniumtownsycl: coalesce the ssm_conv window loads (#26612)f8e30266d

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
4.4 KiB136 linesraw
1#include "ssm_conv.hpp"
2#include "common.hpp"
3
4#include <cstdio>
5
6using namespace sycl;
7
8static void kernel_ssm_conv(
9    queue &q,
10    const float *src_data,
11    const float *weights,
12    float *dst_data,
13    int d_conv,
14    int d_inner,
15    int n_t,
16    int n_s,
17    int ncs __attribute__((unused)),
18    int src_stride_inner,
19    int src_stride_seq,
20    int dst_stride_token,
21    int dst_stride_seq
22) {
23    const size_t total_work = static_cast<size_t>(d_inner) * static_cast<size_t>(n_t) * static_cast<size_t>(n_s);
24    const size_t work_group_size = 256;
25    const size_t num_work_groups = (total_work + work_group_size - 1) / work_group_size;
26
27    const range<1> global_range(num_work_groups * work_group_size);
28    const range<1> local_range(work_group_size);
29
30    q.submit([&](handler &h) {
31        h.parallel_for(
32            nd_range<1>(global_range, local_range),
33            [=](nd_item<1> item) {
34                const size_t idx = item.get_global_id(0);
35                if (idx >= total_work) {
36                    return;
37                }
38
39                // src has the tokens of one channel contiguous, dst has the channels of one
40                // token contiguous, so either the loads or the store must be strided. Indexing
41                // token-fastest coalesces the d_conv loads, which measured faster except for
42                // short, cache-resident rows.
43                const int token   = static_cast<int>(idx % n_t);
44                const int channel = static_cast<int>((idx / n_t) % d_inner);
45                const int seq     = static_cast<int>(idx / (static_cast<size_t>(n_t) * static_cast<size_t>(d_inner)));
46
47                const float *s = src_data
48                    + static_cast<size_t>(seq) * static_cast<size_t>(src_stride_seq)
49                    + static_cast<size_t>(channel) * static_cast<size_t>(src_stride_inner)
50                    + static_cast<size_t>(token);
51
52                const float *c = weights + static_cast<size_t>(channel) * static_cast<size_t>(d_conv);
53
54                float sumf = 0.0f;
55                for (int i0 = 0; i0 < d_conv; ++i0) {
56                    sumf += s[i0] * c[i0];
57                }
58
59                const size_t dst_idx =
60                    static_cast<size_t>(seq) * static_cast<size_t>(dst_stride_seq) +
61                    static_cast<size_t>(token) * static_cast<size_t>(dst_stride_token) +
62                    static_cast<size_t>(channel);
63
64                dst_data[dst_idx] = sumf;
65            }
66        );
67    });
68}
69
70inline void ggml_sycl_op_ssm_conv(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
71    ggml_tensor * src0 = dst->src[0];
72    ggml_tensor * src1 = dst->src[1];
73
74    GGML_ASSERT(src0->type == GGML_TYPE_F32);
75    GGML_ASSERT(src1->type == GGML_TYPE_F32);
76    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
77
78    const int d_conv   = src1->ne[0];
79    const int ncs      = src0->ne[0];
80    const int d_inner  = src0->ne[1];
81    const int n_t      = dst->ne[1];
82    const int n_s      = dst->ne[2];
83
84    GGML_ASSERT(src0->ne[0] == d_conv - 1 + n_t);
85    GGML_ASSERT(src0->ne[1] == d_inner);
86    GGML_ASSERT(src1->ne[1] == d_inner);
87
88    GGML_ASSERT(dst->ne[0] == d_inner);
89    GGML_ASSERT(dst->ne[1] == n_t);
90    GGML_ASSERT(dst->ne[2] == n_s);
91
92    GGML_ASSERT(src0->nb[0] == sizeof(float));
93    GGML_ASSERT(src1->nb[0] == sizeof(float));
94
95    GGML_ASSERT(src0->nb[1] == src0->ne[0] * sizeof(float));
96
97    const int src_stride_inner = ncs;
98    const int src_stride_seq   = ncs * d_inner;
99    const int dst_stride_token = d_inner;
100    const int dst_stride_seq   = d_inner * n_t;
101
102    try {
103        queue *q = ctx.stream();
104
105        const float *src_data = static_cast<const float *>(src0->data);
106        const float *weights  = static_cast<const float *>(src1->data);
107        float *dst_data       = static_cast<float *>(dst->data);
108
109        GGML_ASSERT(src_data && weights && dst_data);
110
111        kernel_ssm_conv(
112            *q,
113            src_data,
114            weights,
115            dst_data,
116            d_conv,
117            d_inner,
118            n_t,
119            n_s,
120            ncs,
121            src_stride_inner,
122            src_stride_seq,
123            dst_stride_token,
124            dst_stride_seq
125        );
126
127    } catch (const std::exception &e) {
128        std::fprintf(stderr, "[SYCL-SSM_CONV] ERROR: %s\n", e.what());
129        throw;
130    }
131}
132
133void ggml_sycl_ssm_conv(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
134    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
135    ggml_sycl_op_ssm_conv(ctx, dst);
136}