llama-cpp-sys-4 0.10.0

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

enable f16;

struct MulMatIdGatherParams {
    offset_ids: u32,
    offset_dst: u32,

    n_expert: u32,
    n_expert_used: u32,
    n_tokens: u32,

    stride_ids_1: u32,

    m: u32,
};

@group(0) @binding(0) var<storage, read_write> ids: array<i32>;        // [n_expert_used, n_tokens]
@group(0) @binding(1) var<storage, read_write> ids_slot_pos: array<u32>; // [n_expert_used * n_tokens]
@group(0) @binding(2) var<storage, read_write> ids_tiles: array<u32>; // [n_tiles, (expert, first slot, n slots) * n_tiles]

@group(0) @binding(3) var<storage, read_write> dst: array<f32>; // [m, n_expert_used, n_tokens]

@group(0) @binding(4) var<uniform> params: MulMatIdGatherParams;

// per expert: the slot count, replaced by the first slot of the expert in ids_slot_pos by the prefix sum, then advanced as its write cursor
var<workgroup> ids_count: array<atomic<u32>, N_EXPERT>;

// launched as a single workgroup, where the threads process the elements of ids in parallel with atomics
@compute @workgroup_size(WG_SIZE)
fn main(@builtin(local_invocation_id) local_id: vec3<u32>) {

    let thread_id = local_id.x;
    let ids_slots_num = params.n_expert_used * params.n_tokens;

    for (var e = thread_id; e < N_EXPERT; e += WG_SIZE) {
        atomicStore(&ids_count[e], 0u);
    }

    workgroupBarrier();

    // count the slots of each expert
    for (var i = thread_id; i < ids_slots_num; i += WG_SIZE) {
        let row = i / params.n_expert_used;
        let col = i % params.n_expert_used;
        let expert = ids[params.offset_ids + row * params.stride_ids_1 + col];
        if (expert != -1) { // skipped slot, owned by no expert
            atomicAdd(&ids_count[u32(expert)], 1u);
        }
    }

    workgroupBarrier();

    // exclusive prefix sum over the counts, and emit one tile descriptor per `TILE_N_SLOTS` slots of each expert
    if (thread_id == 0u) {
        var base: u32 = 0;
        var tile: u32 = 0;

        for (var e: u32 = 0; e < N_EXPERT; e += 1) {
            let n = atomicLoad(&ids_count[e]);

            atomicStore(&ids_count[e], base);

            for (var k: u32 = 0; k < n; k += TILE_N_SLOTS) {
                ids_tiles[1 + 3 * tile]     = e;
                ids_tiles[1 + 3 * tile + 1] = base + k;
                ids_tiles[1 + 3 * tile + 2] = min(n - k, TILE_N_SLOTS);
                tile += 1;
            }

            base += n;
        }

        ids_tiles[0] = tile;
    }

    workgroupBarrier();

    // store the position in ids tensor to ids_slot_pos
    for (var i = thread_id; i < ids_slots_num; i += WG_SIZE) {
        let row = i / params.n_expert_used;
        let col = i % params.n_expert_used;
        let expert = ids[params.offset_ids + row * params.stride_ids_1 + col];
        if (expert != -1) {
            let pos = atomicAdd(&ids_count[u32(expert)], 1u);
            ids_slot_pos[pos] = i;
        }
    }

    // the matmul writes only the rows that an expert owns, so zero the rows of the skipped slots
    for (var i: u32 = 0; i < ids_slots_num; i += 1) {
        let row = i / params.n_expert_used;
        let col = i % params.n_expert_used;
        if (ids[params.offset_ids + row * params.stride_ids_1 + col] != -1) {
            continue;
        }
        let dst_row = params.offset_dst + i * params.m;
        for (var r = thread_id; r < params.m; r += WG_SIZE) {
            dst[dst_row + r] = 0.0f;
        }
    }
}