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
40
41
42
43
44

#version 450
// Causal conv1d single-step update (decode path), mirroring cuda/gdn.cu
// causal_conv1d_update_kernel. One invocation per (channel, batch): shift the
// per-channel conv_state window left by one, append the new input sample, dot
// with the channel's weight, apply SiLU. conv_state is updated in place.
//
// x: [B, conv_dim, 1]  weight: [conv_dim, kernel_size]
// conv_state: [B, conv_dim, kernel_size] (in/out)  output: [B, conv_dim, 1]
//
// Vulkan computes in f32 (f16/bf16 tensors are stored as f32 on this backend),
// so the kernel is f32 throughout; SiLU matches the CUDA float math exactly.
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;

const uint MAX_KS = 8u;

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)          buffer Cs    { float cs[]; };
layout(set = 0, binding = 3) writeonly buffer Out  { float outp[]; };
layout(push_constant) uniform Pc { uint batch_size; uint conv_dim; uint kernel_size; };

void main() {
    uint ch = gl_GlobalInvocationID.x;
    uint b = gl_GlobalInvocationID.y;
    if (ch >= conv_dim || b >= batch_size) { return; }

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

    float win[MAX_KS];
    for (uint i = 0u; i < kernel_size - 1u; i++) {
        win[i] = cs[cs_base + i + 1u];
    }
    win[kernel_size - 1u] = x[b * conv_dim + ch];

    float acc = 0.0;
    for (uint i = 0u; i < kernel_size; i++) {
        cs[cs_base + i] = win[i];
        acc = fma(win[i], w[w_base + i], acc);
    }

    float sig = 1.0 / (1.0 + exp(-acc));
    outp[b * conv_dim + ch] = acc * sig;
}