llama-cpp-sys-2 0.1.158

Low Level Bindings to llama.cpp
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

#version 450

#extension GL_EXT_control_flow_attributes : enable
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
#ifdef USE_SUBGROUPS
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_ballot : require
#endif

layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
layout(constant_id = 0) const uint BLOCK_SIZE = 128;
layout(constant_id = 1) const uint NUM_SUBGROUPS = 1;

layout (binding = 0) readonly buffer M {float16_t data_m[];};
layout (binding = 1) writeonly buffer I {int32_t data_i[];};

layout (push_constant) uniform parameter {
    uint KV;
    uint nem1;
    uint nem2;
    uint nbm1;
    uint nbm2;
    uint nbm3;
    uint n_kv_max;
} p;

#ifdef USE_SUBGROUPS
shared uvec4 ballots_sh[NUM_SUBGROUPS];
#else
shared uint scan[BLOCK_SIZE];
#endif

// One workgroup per mask row: compact the finite-mask KV positions into a
// per-row index list of length n_kv_max, -1 padded. Emitted in ascending KV
// order so the downstream attention accumulation is deterministic.
void main() {
    const uint i1 = gl_WorkGroupID.x;
    const uint i2 = gl_WorkGroupID.y;
    const uint i3 = gl_WorkGroupID.z;
    const uint tid = gl_LocalInvocationIndex;

    const uint m_base   = i3 * p.nbm3 + i2 * p.nbm2 + i1 * p.nbm1;
    const uint out_base = ((i3 * p.nem2 + i2) * p.nem1 + i1) * p.n_kv_max;

    uint base = 0;
    for (uint chunk = 0; chunk < p.KV; chunk += BLOCK_SIZE) {
        const uint k = chunk + tid;
        bool selected = false;
        if (k < p.KV) {
            const float v = float(data_m[m_base + k]);
            selected = !isinf(v) && !isnan(v);
        }

#ifdef USE_SUBGROUPS
        const uvec4 ballot = subgroupBallot(selected);
        if (subgroupElect()) {
            ballots_sh[gl_SubgroupID] = ballot;
        }
        barrier();

        uint subgroup_base = 0;
        uint total = 0;
        [[unroll]] for (uint s = 0; s < gl_NumSubgroups; ++s) {
            if (s == gl_SubgroupID) {
                subgroup_base = total;
            }
            total += subgroupBallotBitCount(ballots_sh[s]);
        }
        barrier();

        const uint slot = base + subgroup_base + subgroupBallotExclusiveBitCount(ballot);
#else
        // Hillis-Steele inclusive prefix sum over the workgroup.
        scan[tid] = selected ? 1u : 0u;
        barrier();
        for (uint off = 1; off < BLOCK_SIZE; off <<= 1) {
            uint add = 0;
            if (tid >= off) {
                add = scan[tid - off];
            }
            barrier();
            scan[tid] += add;
            barrier();
        }

        const uint inclusive = scan[tid];
        const uint total = scan[BLOCK_SIZE - 1];
        const uint slot = base + inclusive - 1u;
#endif

        if (selected && slot < p.n_kv_max) {
            data_i[out_base + slot] = int32_t(k);
        }
        base += total;
        barrier();
    }

    for (uint s = min(base, p.n_kv_max) + tid; s < p.n_kv_max; s += BLOCK_SIZE) {
        data_i[out_base + s] = int32_t(-1);
    }
}