hanzo-ml 0.11.19

Minimalist ML framework.
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

#version 450
// Gated delta rule, SINGLE timestep (decode, seq_len==1). One invocation per (bh, v) state column,
// looping the k_dim in a local array. Matches gdn.rs::gated_delta_rule_recurrence for seq=1 and the
// CUDA gated_delta_rule_recurrence_kernel: state is [BH, K, V] with state[bh][j][v] at j*V + v.
//   decay = exp(g[bh]); s[j] *= decay; kv_mem = sum_j s[j]*k[j]
//   delta = (v[v] - kv_mem) * beta[bh]; s[j] += k[j]*delta; y[v] = sum_j s[j]*q[j]
// q is pre-scaled by the caller (1/sqrt(k_dim)), exactly as the CUDA recurrence wrapper does.
// MAX_K bounds the per-invocation k array; k_dim (head_k_dim) must be <= MAX_K (asserted host-side).
#define MAX_K 256u
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
layout(set = 0, binding = 0) readonly  buffer Q     { float q[]; };
layout(set = 0, binding = 1) readonly  buffer K     { float k[]; };
layout(set = 0, binding = 2) readonly  buffer V     { float v[]; };
layout(set = 0, binding = 3) readonly  buffer G     { float g[]; };
layout(set = 0, binding = 4) readonly  buffer Beta  { float beta[]; };
layout(set = 0, binding = 5)           buffer State { float state[]; };
layout(set = 0, binding = 6) writeonly buffer Out   { float outp[]; };
layout(push_constant) uniform Pc { uint bh_count; uint k_dim; uint v_dim; };
void main() {
    uint gid = gl_GlobalInvocationID.x;
    uint total = bh_count * v_dim;
    if (gid >= total) { return; }
    uint bh = gid / v_dim;
    uint v_idx = gid - bh * v_dim;

    uint q_base = bh * k_dim;          // q,k: [BH, K] (seq=1)
    uint state_base = bh * k_dim * v_dim;

    float decay = exp(g[bh]);
    float beta_t = beta[bh];
    float v_t = v[bh * v_dim + v_idx];

    float s[MAX_K];
    float kv_mem = 0.0;
    for (uint j = 0u; j < k_dim; j++) {
        float sj = state[state_base + j * v_dim + v_idx] * decay;
        s[j] = sj;
        kv_mem += sj * k[q_base + j];
    }
    float delta = (v_t - kv_mem) * beta_t;
    float y_t = 0.0;
    for (uint j = 0u; j < k_dim; j++) {
        float sj = s[j] + k[q_base + j] * delta;
        state[state_base + j * v_dim + v_idx] = sj;
        y_t += sj * q[q_base + j];
    }
    outp[bh * v_dim + v_idx] = y_t;
}