char-fork/llama.cpp

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

Gaurav GargRemove padding and multiple D2D copies for MTP (#24086)e95dae18d

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
6.7 KiB189 linesraw
1#version 450
2
3#extension GL_EXT_control_flow_attributes : require
4#extension GL_KHR_shader_subgroup_basic : enable
5#if USE_SUBGROUP_CLUSTERED
6#extension GL_KHR_shader_subgroup_clustered : enable
7#endif
8#if USE_SUBGROUP_ADD
9#extension GL_KHR_shader_subgroup_arithmetic : enable
10#endif
11
12// Caller guarantees valid spec constants: S_V % COLS_PER_WG == 0 and S_V % LANES_PER_COLUMN == 0,
13// so no bounds checking is needed.
14layout(constant_id = 0) const uint S_V = 128;
15layout(constant_id = 1) const uint KDA = 0;
16layout(constant_id = 2) const uint SUBGROUP_SIZE = 32;
17layout(constant_id = 3) const uint LANES_PER_COLUMN = 32;
18
19const uint COLS_PER_WG = SUBGROUP_SIZE / LANES_PER_COLUMN;
20const uint ROWS_PER_LANE = S_V / LANES_PER_COLUMN;
21
22layout(local_size_x_id = 2, local_size_y = 1, local_size_z = 1) in;
23
24layout(push_constant) uniform Parameters {
25    uint H;
26    uint n_tokens;
27    uint n_seqs;
28    uint s_off;
29    uint sq1, sq2, sq3;
30    uint sv1, sv2, sv3;
31    uint sb1, sb2, sb3;
32    uint neq1, rq3;
33    float scale;
34    uint K;
35};
36
37layout(binding = 0) readonly  buffer QBuf     { FLOAT_TYPE data_q[];     };
38layout(binding = 1) readonly  buffer KBuf     { FLOAT_TYPE data_k[];     };
39layout(binding = 2) readonly  buffer VBuf     { FLOAT_TYPE data_v[];     };
40layout(binding = 3) readonly  buffer GBuf     { FLOAT_TYPE data_g[];     };
41layout(binding = 4) readonly  buffer BetaBuf  { FLOAT_TYPE data_beta[];  };
42layout(binding = 5) readonly  buffer StateBuf { FLOAT_TYPE data_state[]; };
43layout(binding = 6)           buffer DstBuf   { FLOAT_TYPE data_dst[];   };
44
45#if !USE_SUBGROUP_ADD && !USE_SUBGROUP_CLUSTERED
46shared FLOAT_TYPE temp[SUBGROUP_SIZE];
47
48// This does a reduction across groups of LANES_PER_COLUMN
49FLOAT_TYPE reduce_add_shmem(FLOAT_TYPE partial) {
50    const uint lane = gl_SubgroupInvocationID;
51    temp[lane] = partial;
52    barrier();
53    [[unroll]] for (uint s = LANES_PER_COLUMN / 2u; s > 0; s >>= 1u) {
54        FLOAT_TYPE other = temp[lane ^ s];
55        barrier();
56        temp[lane] += other;
57        barrier();
58    }
59    const FLOAT_TYPE result = temp[lane];
60    barrier();
61    return result;
62}
63#endif
64
65// clusterSize for subgroupClusteredAdd must be a compile-time constant; branch on spec constant
66FLOAT_TYPE reduce_partial(FLOAT_TYPE partial) {
67    switch (LANES_PER_COLUMN) {
68        case 1u:
69            return partial;
70#if USE_SUBGROUP_CLUSTERED
71        // Workaround for GLSL requiring a literal constant for the cluster size.
72        // The branches should all fold away.
73        case 2u:
74            return subgroupClusteredAdd(partial, 2u);
75        case 4u:
76            return subgroupClusteredAdd(partial, 4u);
77        case 8u:
78            return subgroupClusteredAdd(partial, 8u);
79        case 16u:
80            return subgroupClusteredAdd(partial, 16u);
81        case 32u:
82            return subgroupClusteredAdd(partial, 32u);
83        case 64u:
84            return subgroupClusteredAdd(partial, 64u);
85#endif
86        default:
87#if USE_SUBGROUP_ADD
88            return subgroupAdd(partial);
89#else
90            return reduce_add_shmem(partial);
91#endif
92    }
93}
94
95void main() {
96    const uint head_id = gl_WorkGroupID.x;
97    const uint seq_id = gl_WorkGroupID.y;
98    const uint lane = gl_SubgroupInvocationID % LANES_PER_COLUMN;
99    const uint col = gl_WorkGroupID.z * COLS_PER_WG + (gl_SubgroupInvocationID / LANES_PER_COLUMN);
100
101    const uint iq1 = head_id % neq1;
102    const uint iq3 = seq_id / rq3;
103
104    const uint state_size = S_V * S_V;
105    // input state holds s0 only [S_v, S_v, H, n_seqs]: per-seq stride is H*D.
106    const uint state_in_base       = (seq_id * H + head_id) * state_size;
107    // output state layout per slot: same per-(seq,head) offset as the single-slot case.
108    const uint state_out_base      = (seq_id * H + head_id) * state_size;
109    const uint state_size_per_snap = state_size * H * n_seqs;
110
111    FLOAT_TYPE s_shard[ROWS_PER_LANE];
112    [[unroll]] for (uint r = 0; r < ROWS_PER_LANE; r++) {
113        s_shard[r] = FLOAT_TYPE(data_state[state_in_base + col * S_V + r * LANES_PER_COLUMN + lane]);
114    }
115
116    // snapshot slot mapping: slot 0 = most recent state, slot s = s tokens back.
117    // When n_tokens < K, only slots 0..n_tokens-1 are written; older slots are caller-owned.
118
119    uint attn_off = (seq_id * n_tokens * H + head_id) * S_V;
120
121    for (uint t = 0; t < n_tokens; t++) {
122        const uint q_off = iq3 * sq3 + t * sq2 + iq1 * sq1;
123        const uint k_off = q_off;
124        const uint v_off = seq_id * sv3 + t * sv2 + head_id * sv1;
125        const uint gb_off = seq_id * sb3 + t * sb2 + head_id * sb1;
126        const FLOAT_TYPE beta_val = FLOAT_TYPE(data_beta[gb_off]);
127
128        FLOAT_TYPE k_reg[ROWS_PER_LANE];
129        FLOAT_TYPE q_reg[ROWS_PER_LANE];
130        [[unroll]] for (uint r = 0; r < ROWS_PER_LANE; r++) {
131            const uint i = r * LANES_PER_COLUMN + lane;
132            k_reg[r] = FLOAT_TYPE(data_k[k_off + i]);
133            q_reg[r] = FLOAT_TYPE(data_q[q_off + i]);
134        }
135
136        FLOAT_TYPE g_exp[ROWS_PER_LANE];
137        if (KDA == 0) {
138            const FLOAT_TYPE g_val = exp(FLOAT_TYPE(data_g[gb_off]));
139            [[unroll]] for (uint r = 0; r < ROWS_PER_LANE; r++) {
140                g_exp[r] = g_val;
141            }
142        } else {
143            const uint g_base = gb_off * S_V;
144            [[unroll]] for (uint r = 0; r < ROWS_PER_LANE; r++) {
145                const uint i = r * LANES_PER_COLUMN + lane;
146                g_exp[r] = exp(FLOAT_TYPE(data_g[g_base + i]));
147            }
148        }
149
150        const FLOAT_TYPE v_val = FLOAT_TYPE(data_v[v_off + col]);
151
152        FLOAT_TYPE kv_shard = 0.0;
153        [[unroll]] for (uint r = 0; r < ROWS_PER_LANE; r++) {
154            kv_shard += g_exp[r] * s_shard[r] * k_reg[r];
155        }
156        FLOAT_TYPE kv_col = reduce_partial(kv_shard);
157
158        FLOAT_TYPE delta_col = (v_val - kv_col) * beta_val;
159
160        FLOAT_TYPE attn_partial = 0.0;
161        [[unroll]] for (uint r = 0; r < ROWS_PER_LANE; r++) {
162            s_shard[r] = g_exp[r] * s_shard[r] + k_reg[r] * delta_col;
163            attn_partial += s_shard[r] * q_reg[r];
164        }
165        FLOAT_TYPE attn_col = reduce_partial(attn_partial);
166
167        if (lane == 0) {
168            data_dst[attn_off + col] = attn_col * scale;
169        }
170
171        attn_off += S_V * H;
172
173        if (K > 1u) {
174            const int target_slot = int(n_tokens) - 1 - int(t);
175            if (target_slot >= 0 && target_slot < int(K)) {
176                const uint slot_base = s_off + uint(target_slot) * state_size_per_snap + state_out_base;
177                [[unroll]] for (uint r = 0; r < ROWS_PER_LANE; r++) {
178                    data_dst[slot_base + col * S_V + r * LANES_PER_COLUMN + lane] = s_shard[r];
179                }
180            }
181        }
182    }
183
184    if (K == 1u) {
185        [[unroll]] for (uint r = 0; r < ROWS_PER_LANE; r++) {
186            data_dst[s_off + state_out_base + col * S_V + r * LANES_PER_COLUMN + lane] = s_shard[r];
187        }
188    }
189}