char-fork/llama.cpp

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

Pranav UttarkarImplemented vulkan cross_entropy_loss and cross_entropy_loss_back (#27216)bf9421646

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.1 KiB78 linesraw
1#version 450
2
3#include "generic_head.glsl"
4#include "types.glsl"
5
6#extension GL_EXT_control_flow_attributes : enable
7
8layout(constant_id = 0) const uint BLOCK_SIZE = 32;
9layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
10
11layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
12layout (binding = 1) readonly buffer B {B_TYPE data_b[];};
13layout (binding = 2) writeonly buffer D {D_TYPE data_d[];};
14
15shared FLOAT_TYPE tmp[BLOCK_SIZE];
16
17FLOAT_TYPE wg_reduce_max(FLOAT_TYPE v) {
18    const uint tid = gl_LocalInvocationID.x;
19    tmp[tid] = v;
20    barrier();
21    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
22        if (tid < s) {
23            tmp[tid] = max(tmp[tid], tmp[tid + s]);
24        }
25        barrier();
26    }
27    v = tmp[0];
28    barrier();
29    return v;
30}
31
32FLOAT_TYPE wg_reduce_sum(FLOAT_TYPE v) {
33    const uint tid = gl_LocalInvocationID.x;
34    tmp[tid] = v;
35    barrier();
36    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
37        if (tid < s) {
38            tmp[tid] += tmp[tid + s];
39        }
40        barrier();
41    }
42    v = tmp[0];
43    barrier();
44    return v;
45}
46
47void main() {
48    const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
49    const uint tid = gl_LocalInvocationID.x;
50
51    if (row >= p.KY) {
52        return;
53    }
54
55    const uint off = row * p.KX;
56
57    FLOAT_TYPE max_logit = FLOAT_TYPE(uintBitsToFloat(0xFF800000));
58    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
59        max_logit = max(max_logit, FLOAT_TYPE(data_a[off + i]));
60    }
61    max_logit = wg_reduce_max(max_logit);
62
63    FLOAT_TYPE sum_exp = FLOAT_TYPE(0.0f);
64    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
65        sum_exp += exp(FLOAT_TYPE(data_a[off + i]) - max_logit);
66    }
67    const FLOAT_TYPE log_sum = log(wg_reduce_sum(sum_exp));
68
69    FLOAT_TYPE loss = FLOAT_TYPE(0.0f);
70    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
71        loss += (FLOAT_TYPE(data_a[off + i]) - max_logit - log_sum) * FLOAT_TYPE(data_b[off + i]);
72    }
73    loss = -wg_reduce_sum(loss) / FLOAT_TYPE(p.KY);
74
75    if (tid == 0) {
76        data_d[row] = D_TYPE(loss);
77    }
78}