hanzo-ml 0.11.63

Fast multi-backend tensor & ML framework for Rust (CPU/CUDA/Metal/Vulkan/ROCm) with quantization — the compute core of the Hanzo stack.
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

#version 450
// Causal conv1d over a full sequence (prefill path), mirroring cuda/gdn.cu
// causal_conv1d_full_kernel. One invocation per output element (channel, pos,
// batch): a left-padded causal window of kernel_size ending at pos, dotted with
// the channel weight, then SiLU. Zero-pads positions before the start.
//
// x: [B, conv_dim, S]  weight: [conv_dim, kernel_size]  output: [B, conv_dim, S]
//
// Grid is flattened to 1D over total = B*conv_dim*S so it uses the backend's
// standard ceil(total/64) dispatch; index math recovers (b, ch, pos).
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;

layout(set = 0, binding = 0) readonly  buffer X    { float x[]; };
layout(set = 0, binding = 1) readonly  buffer W    { float w[]; };
layout(set = 0, binding = 2) writeonly buffer Out  { float outp[]; };
layout(push_constant) uniform Pc { uint batch_size; uint conv_dim; uint seq_len; uint kernel_size; };

void main() {
    uint gid = gl_GlobalInvocationID.x;
    uint total = batch_size * conv_dim * seq_len;
    if (gid >= total) { return; }

    uint pos = gid % seq_len;
    uint ch = (gid / seq_len) % conv_dim;
    uint b = gid / (seq_len * conv_dim);

    uint x_base = (b * conv_dim + ch) * seq_len;
    uint w_base = ch * kernel_size;

    float acc = 0.0;
    for (uint i = 0u; i < kernel_size; i++) {
        int src_pos = int(pos) - int(kernel_size - 1u) + int(i);
        float x_val = (src_pos >= 0) ? x[x_base + uint(src_pos)] : 0.0;
        acc = fma(x_val, w[w_base + i], acc);
    }

    float sig = 1.0 / (1.0 + exp(-acc));
    outp[x_base + pos] = acc * sig;
}