char-fork/llama.cpp

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

Kevin Hoppervulkan: add DeepSeek-V4 hyper-connection fused ops (DSV4_HC_COMB/PRE/POST) (#26578)7a333e724

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.7 KiB90 linesraw
1#version 450
2
3#extension GL_EXT_control_flow_attributes : require
4#extension GL_KHR_shader_subgroup_basic : require
5#extension GL_KHR_shader_subgroup_shuffle : require
6
7// 16 lanes per token, indexed idst + hc*isrc: idst in bits 0..1, isrc in bits 2..3,
8// so subgroupShuffleXor by 1|2 reduces a row and by 4|8 a column.
9
10layout(constant_id = 0) const uint SUBGROUP_SIZE = 32;
11
12layout(local_size_x_id = 0, local_size_y = 4, local_size_z = 1) in;
13
14layout(push_constant) uniform parameter
15{
16    uint n_tokens;
17
18    uint nbm0; uint nbm1;   // mixes
19    uint nbs0;              // scale
20    uint nbb0;              // base
21    uint nbd0; uint nbd1; uint nbd2;   // dst
22
23    uint m_offset;
24    uint s_offset;
25    uint b_offset;
26    uint d_offset;
27
28    float eps;
29    uint n_iter;
30};
31
32layout(binding = 0, std430) readonly buffer M { float data_m[]; };
33layout(binding = 1, std430) readonly buffer S { float data_s[]; };
34layout(binding = 2, std430) readonly buffer B { float data_b[]; };
35layout(binding = 3, std430) writeonly buffer D { float data_d[]; };
36
37const uint hc          = 4;
38const uint comb_offset = 2 * hc;
39
40const uint TOKENS_PER_SUBGROUP = SUBGROUP_SIZE / 16;
41
42void main() {
43    const uint lane = gl_SubgroupInvocationID;
44    const uint blk  = lane >> 4;    // which 16-lane block, i.e. which token
45    const uint idx  = lane & 15;    // idst + hc*isrc
46
47    const uint sg = gl_WorkGroupID.x * gl_WorkGroupSize.y + gl_SubgroupID;
48    const uint it = sg * TOKENS_PER_SUBGROUP + blk;
49
50    // no early return, the shuffles need every lane; out-of-range blocks compute a discarded value
51    const bool in_range = it < n_tokens;
52
53    const float scale_comb = data_s[s_offset + 2 * nbs0];
54
55    float v = 0.0f;
56    if (in_range) {
57        v = data_m[m_offset + (comb_offset + idx) * nbm0 + it * nbm1] * scale_comb
58          + data_b[b_offset + (comb_offset + idx) * nbb0];
59    }
60
61    // Softmax across destinations: the four lanes sharing an isrc.
62    float vmax = max(v, subgroupShuffleXor(v, 1));
63    vmax = max(vmax, subgroupShuffleXor(vmax, 2));
64    v = exp(v - vmax);
65
66    float sum = v + subgroupShuffleXor(v, 1);
67    sum += subgroupShuffleXor(sum, 2);
68    v = v / sum + eps;
69
70    // Normalize columns: equal destination indices are four lanes apart.
71    sum = v + subgroupShuffleXor(v, 4);
72    sum += subgroupShuffleXor(sum, 8);
73    v /= sum + eps;
74
75    for (uint i = 1; i < n_iter; ++i) {
76        sum = v + subgroupShuffleXor(v, 1);
77        sum += subgroupShuffleXor(sum, 2);
78        v /= sum + eps;
79
80        sum = v + subgroupShuffleXor(v, 4);
81        sum += subgroupShuffleXor(sum, 8);
82        v /= sum + eps;
83    }
84
85    if (in_range) {
86        const uint idst = idx & 3;
87        const uint isrc = idx >> 2;
88        data_d[d_offset + idst * nbd0 + isrc * nbd1 + it * nbd2] = v;
89    }
90}