volvoxai 0.1.0

package/shaders/conv2DPointwise16Tile.wgsl

@@ -0,0 +1,86 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read> weight : array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> bias : array<vec4<f32>>;
+@group(0) @binding(3) var<storage, read_write> output : array<vec4<f32>>;
+
+struct Params {
+    n : u32,
+    in_h : u32,
+    in_w : u32,
+    in_c : u32,
+    out_c : u32,
+    out_h : u32,
+    out_w : u32,
+    kh : u32,
+    kw : u32,
+    sy : u32,
+    sx : u32,
+    pt : u32,
+    pl : u32,
+    groups : u32,
+    relu : u32,
+    dy : u32,
+    dx : u32,
+}
+@group(0) @binding(4) var<uniform> params : Params;
+
+const TILE_C : u32 = 64u;
+var<workgroup> tile_weight : array<vec4<f32>, 256>;
+
+fn apply_relu4(v_in : vec4<f32>) -> vec4<f32> {
+    var v = v_in;
+    if (params.relu == 1u) {
+        v = max(v, vec4<f32>(0.0));
+    } else if (params.relu >= 2u) {
+        v = min(max(v, vec4<f32>(0.0)), vec4<f32>(6.0));
+    }
+    return v;
+}
+
+@compute @workgroup_size(8, 8, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>,
+        @builtin(local_invocation_id) lid3 : vec3<u32>) {
+    let ox = gid.x;
+    let oy = gid.y;
+    let oc16_count = params.out_c / 16u;
+    let nb = gid.z / oc16_count;
+    let oc0 = (gid.z - nb * oc16_count) * 16u;
+    let lid = lid3.y * 8u + lid3.x;
+    let in_bounds = nb < params.n && ox < params.out_w && oy < params.out_h;
+
+    var sum0 = vec4<f32>(0.0);
+    var sum1 = vec4<f32>(0.0);
+    var sum2 = vec4<f32>(0.0);
+    var sum3 = vec4<f32>(0.0);
+    let input_base = ((nb * params.in_h + oy) * params.in_w + ox) * params.in_c;
+
+    for (var tile_start = 0u; tile_start < params.in_c; tile_start = tile_start + TILE_C) {
+        let tile_count = min(TILE_C, params.in_c - tile_start);
+        for (var wi = lid; wi < tile_count * 4u; wi = wi + 64u) {
+            let ic = tile_start + wi / 4u;
+            let oc_vec = wi - (wi / 4u) * 4u;
+            tile_weight[wi] = weight[(ic * params.out_c + oc0) / 4u + oc_vec];
+        }
+        workgroupBarrier();
+
+        if (in_bounds) {
+            for (var ti = 0u; ti < tile_count; ti = ti + 1u) {
+                let x = input[input_base + tile_start + ti];
+                let wbase = ti * 4u;
+                sum0 = sum0 + x * tile_weight[wbase];
+                sum1 = sum1 + x * tile_weight[wbase + 1u];
+                sum2 = sum2 + x * tile_weight[wbase + 2u];
+                sum3 = sum3 + x * tile_weight[wbase + 3u];
+            }
+        }
+        workgroupBarrier();
+    }
+
+    if (!in_bounds) { return; }
+    let output_base = (((nb * params.out_h + oy) * params.out_w + ox) * params.out_c + oc0) / 4u;
+    let bias_base = oc0 / 4u;
+    output[output_base] = apply_relu4(sum0 + bias[bias_base]);
+    output[output_base + 1u] = apply_relu4(sum1 + bias[bias_base + 1u]);
+    output[output_base + 2u] = apply_relu4(sum2 + bias[bias_base + 2u]);
+    output[output_base + 3u] = apply_relu4(sum3 + bias[bias_base + 3u]);
+}

package/shaders/conv2DPointwise8.wgsl

@@ -0,0 +1,85 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read> weight : array<f32>;
+@group(0) @binding(2) var<storage, read> bias : array<f32>;
+@group(0) @binding(3) var<storage, read_write> output : array<f32>;
+
+struct Params {
+    n : u32,
+    in_h : u32,
+    in_w : u32,
+    in_c : u32,
+    out_c : u32,
+    out_h : u32,
+    out_w : u32,
+    kh : u32,
+    kw : u32,
+    sy : u32,
+    sx : u32,
+    pt : u32,
+    pl : u32,
+    groups : u32,
+    relu : u32,
+    dy : u32,
+    dx : u32,
+}
+@group(0) @binding(4) var<uniform> params : Params;
+
+fn apply_relu(v_in : f32) -> f32 {
+    var v = v_in;
+    if (params.relu == 1u) {
+        v = max(v, 0.0);
+    } else if (params.relu >= 2u) {
+        v = min(max(v, 0.0), 6.0);
+    }
+    return v;
+}
+
+@compute @workgroup_size(8, 8, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let ox = gid.x;
+    let oy = gid.y;
+    let oc8_count = (params.out_c + 7u) / 8u;
+    let nb = gid.z / oc8_count;
+    let oc0 = (gid.z - nb * oc8_count) * 8u;
+    if (nb >= params.n || ox >= params.out_w || oy >= params.out_h || oc0 >= params.out_c) { return; }
+
+    let has1 = oc0 + 1u < params.out_c;
+    let has2 = oc0 + 2u < params.out_c;
+    let has3 = oc0 + 3u < params.out_c;
+    let has4 = oc0 + 4u < params.out_c;
+    let has5 = oc0 + 5u < params.out_c;
+    let has6 = oc0 + 6u < params.out_c;
+    let has7 = oc0 + 7u < params.out_c;
+    var sum0 = 0.0;
+    var sum1 = 0.0;
+    var sum2 = 0.0;
+    var sum3 = 0.0;
+    var sum4 = 0.0;
+    var sum5 = 0.0;
+    var sum6 = 0.0;
+    var sum7 = 0.0;
+
+    let input_base = ((nb * params.in_h + oy) * params.in_w + ox) * params.in_c;
+    for (var ic = 0u; ic < params.in_c; ic = ic + 1u) {
+        let x = input[input_base + ic];
+        let wbase = ic * params.out_c + oc0;
+        sum0 = sum0 + x * weight[wbase];
+        if (has1) { sum1 = sum1 + x * weight[wbase + 1u]; }
+        if (has2) { sum2 = sum2 + x * weight[wbase + 2u]; }
+        if (has3) { sum3 = sum3 + x * weight[wbase + 3u]; }
+        if (has4) { sum4 = sum4 + x * weight[wbase + 4u]; }
+        if (has5) { sum5 = sum5 + x * weight[wbase + 5u]; }
+        if (has6) { sum6 = sum6 + x * weight[wbase + 6u]; }
+        if (has7) { sum7 = sum7 + x * weight[wbase + 7u]; }
+    }
+
+    let output_base = ((nb * params.out_h + oy) * params.out_w + ox) * params.out_c + oc0;
+    output[output_base] = apply_relu(sum0 + bias[oc0]);
+    if (has1) { output[output_base + 1u] = apply_relu(sum1 + bias[oc0 + 1u]); }
+    if (has2) { output[output_base + 2u] = apply_relu(sum2 + bias[oc0 + 2u]); }
+    if (has3) { output[output_base + 3u] = apply_relu(sum3 + bias[oc0 + 3u]); }
+    if (has4) { output[output_base + 4u] = apply_relu(sum4 + bias[oc0 + 4u]); }
+    if (has5) { output[output_base + 5u] = apply_relu(sum5 + bias[oc0 + 5u]); }
+    if (has6) { output[output_base + 6u] = apply_relu(sum6 + bias[oc0 + 6u]); }
+    if (has7) { output[output_base + 7u] = apply_relu(sum7 + bias[oc0 + 7u]); }
+}

package/shaders/conv2DPointwise8Vec2.wgsl

@@ -0,0 +1,70 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read> weight : array<vec2<f32>>;
+@group(0) @binding(2) var<storage, read> bias : array<vec2<f32>>;
+@group(0) @binding(3) var<storage, read_write> output : array<vec2<f32>>;
+
+struct Params {
+    n : u32,
+    in_h : u32,
+    in_w : u32,
+    in_c : u32,
+    out_c : u32,
+    out_h : u32,
+    out_w : u32,
+    kh : u32,
+    kw : u32,
+    sy : u32,
+    sx : u32,
+    pt : u32,
+    pl : u32,
+    groups : u32,
+    relu : u32,
+    dy : u32,
+    dx : u32,
+}
+@group(0) @binding(4) var<uniform> params : Params;
+
+fn apply_relu2(v_in : vec2<f32>) -> vec2<f32> {
+    var v = v_in;
+    if (params.relu == 1u) {
+        v = max(v, vec2<f32>(0.0));
+    } else if (params.relu >= 2u) {
+        v = min(max(v, vec2<f32>(0.0)), vec2<f32>(6.0));
+    }
+    return v;
+}
+
+@compute @workgroup_size(8, 8, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let ox = gid.x;
+    let oy = gid.y;
+    let oc8_count = (params.out_c + 7u) / 8u;
+    let nb = gid.z / oc8_count;
+    let oc0 = (gid.z - nb * oc8_count) * 8u;
+    if (nb >= params.n || ox >= params.out_w || oy >= params.out_h || oc0 >= params.out_c) { return; }
+
+    let has1 = oc0 + 2u < params.out_c;
+    let has2 = oc0 + 4u < params.out_c;
+    let has3 = oc0 + 6u < params.out_c;
+    var sum0 = vec2<f32>(0.0);
+    var sum1 = vec2<f32>(0.0);
+    var sum2 = vec2<f32>(0.0);
+    var sum3 = vec2<f32>(0.0);
+
+    let input_base = ((nb * params.in_h + oy) * params.in_w + ox) * params.in_c;
+    for (var ic = 0u; ic < params.in_c; ic = ic + 1u) {
+        let x = input[input_base + ic];
+        let wbase = (ic * params.out_c + oc0) / 2u;
+        sum0 = sum0 + x * weight[wbase];
+        if (has1) { sum1 = sum1 + x * weight[wbase + 1u]; }
+        if (has2) { sum2 = sum2 + x * weight[wbase + 2u]; }
+        if (has3) { sum3 = sum3 + x * weight[wbase + 3u]; }
+    }
+
+    let output_base = (((nb * params.out_h + oy) * params.out_w + ox) * params.out_c + oc0) / 2u;
+    let bias_base = oc0 / 2u;
+    output[output_base] = apply_relu2(sum0 + bias[bias_base]);
+    if (has1) { output[output_base + 1u] = apply_relu2(sum1 + bias[bias_base + 1u]); }
+    if (has2) { output[output_base + 2u] = apply_relu2(sum2 + bias[bias_base + 2u]); }
+    if (has3) { output[output_base + 3u] = apply_relu2(sum3 + bias[bias_base + 3u]); }
+}

package/shaders/conv2DPointwise8Vec4.wgsl

@@ -0,0 +1,65 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read> weight : array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> bias : array<vec4<f32>>;
+@group(0) @binding(3) var<storage, read_write> output : array<vec4<f32>>;
+
+struct Params {
+    n : u32,
+    in_h : u32,
+    in_w : u32,
+    in_c : u32,
+    out_c : u32,
+    out_h : u32,
+    out_w : u32,
+    kh : u32,
+    kw : u32,
+    sy : u32,
+    sx : u32,
+    pt : u32,
+    pl : u32,
+    groups : u32,
+    relu : u32,
+    dy : u32,
+    dx : u32,
+}
+@group(0) @binding(4) var<uniform> params : Params;
+
+fn apply_relu4(v_in : vec4<f32>) -> vec4<f32> {
+    var v = v_in;
+    if (params.relu == 1u) {
+        v = max(v, vec4<f32>(0.0));
+    } else if (params.relu >= 2u) {
+        v = min(max(v, vec4<f32>(0.0)), vec4<f32>(6.0));
+    }
+    return v;
+}
+
+@compute @workgroup_size(8, 8, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let ox = gid.x;
+    let oy = gid.y;
+    let oc8_count = (params.out_c + 7u) / 8u;
+    let nb = gid.z / oc8_count;
+    let oc0 = (gid.z - nb * oc8_count) * 8u;
+    if (nb >= params.n || ox >= params.out_w || oy >= params.out_h || oc0 >= params.out_c) { return; }
+
+    var sum0 = vec4<f32>(0.0);
+    var sum1 = vec4<f32>(0.0);
+    let input_base = ((nb * params.in_h + oy) * params.in_w + ox) * params.in_c;
+    let has_second = oc0 + 4u < params.out_c;
+    for (var ic = 0u; ic < params.in_c; ic = ic + 1u) {
+        let x = input[input_base + ic];
+        let wbase = (ic * params.out_c + oc0) / 4u;
+        sum0 = sum0 + x * weight[wbase];
+        if (has_second) {
+            sum1 = sum1 + x * weight[wbase + 1u];
+        }
+    }
+
+    let output_base = (((nb * params.out_h + oy) * params.out_w + ox) * params.out_c + oc0) / 4u;
+    let bias_base = oc0 / 4u;
+    output[output_base] = apply_relu4(sum0 + bias[bias_base]);
+    if (has_second) {
+        output[output_base + 1u] = apply_relu4(sum1 + bias[bias_base + 1u]);
+    }
+}

package/shaders/conv2DRegularC3Out16.wgsl

@@ -0,0 +1,75 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read> weight : array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> bias : array<vec4<f32>>;
+@group(0) @binding(3) var<storage, read_write> output : array<vec4<f32>>;
+
+struct Params {
+    n : u32,
+    in_h : u32,
+    in_w : u32,
+    in_c : u32,
+    out_c : u32,
+    out_h : u32,
+    out_w : u32,
+    kh : u32,
+    kw : u32,
+    sy : u32,
+    sx : u32,
+    pt : u32,
+    pl : u32,
+    groups : u32,
+    relu : u32,
+    dy : u32,
+    dx : u32,
+}
+@group(0) @binding(4) var<uniform> params : Params;
+
+fn apply_relu4(v_in : vec4<f32>) -> vec4<f32> {
+    var v = v_in;
+    if (params.relu == 1u) {
+        v = max(v, vec4<f32>(0.0));
+    } else if (params.relu >= 2u) {
+        v = min(max(v, vec4<f32>(0.0)), vec4<f32>(6.0));
+    }
+    return v;
+}
+
+@compute @workgroup_size(8, 8, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let ox = gid.x;
+    let oy = gid.y;
+    let oc16_count = params.out_c / 16u;
+    let nb = gid.z / oc16_count;
+    let oc0 = (gid.z - nb * oc16_count) * 16u;
+    if (nb >= params.n || ox >= params.out_w || oy >= params.out_h) { return; }
+
+    var sum0 = vec4<f32>(0.0);
+    var sum1 = vec4<f32>(0.0);
+    var sum2 = vec4<f32>(0.0);
+    var sum3 = vec4<f32>(0.0);
+    for (var yy = 0u; yy < params.kh; yy = yy + 1u) {
+        let iy = i32(oy * params.sy + yy * params.dy) - i32(params.pt);
+        if (iy < 0 || iy >= i32(params.in_h)) { continue; }
+        for (var xx = 0u; xx < params.kw; xx = xx + 1u) {
+            let ix = i32(ox * params.sx + xx * params.dx) - i32(params.pl);
+            if (ix < 0 || ix >= i32(params.in_w)) { continue; }
+            let input_base = ((nb * params.in_h + u32(iy)) * params.in_w + u32(ix)) * 3u;
+            let weight_base = (((yy * params.kw + xx) * 3u) * params.out_c + oc0) / 4u;
+            for (var ic = 0u; ic < 3u; ic = ic + 1u) {
+                let x = input[input_base + ic];
+                let wbase = weight_base + ic * (params.out_c / 4u);
+                sum0 = sum0 + x * weight[wbase];
+                sum1 = sum1 + x * weight[wbase + 1u];
+                sum2 = sum2 + x * weight[wbase + 2u];
+                sum3 = sum3 + x * weight[wbase + 3u];
+            }
+        }
+    }
+
+    let output_base = (((nb * params.out_h + oy) * params.out_w + ox) * params.out_c + oc0) / 4u;
+    let bias_base = oc0 / 4u;
+    output[output_base] = apply_relu4(sum0 + bias[bias_base]);
+    output[output_base + 1u] = apply_relu4(sum1 + bias[bias_base + 1u]);
+    output[output_base + 2u] = apply_relu4(sum2 + bias[bias_base + 2u]);
+    output[output_base + 3u] = apply_relu4(sum3 + bias[bias_base + 3u]);
+}

package/shaders/convTranspose2D.wgsl

@@ -0,0 +1,33 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read> weight : array<f32>;
+        @group(0) @binding(2) var<storage, read> bias : array<f32>;
+        @group(0) @binding(3) var<storage, read_write> output : array<f32>;
+        struct Params { b: u32, in_h: u32, in_w: u32, in_c: u32, out_h: u32, out_w: u32, out_c: u32, kh: u32, kw: u32, sh: u32, sw: u32, ph: u32, pw: u32, has_bias: u32 }
+        @group(0) @binding(4) var<uniform> params : Params;
+        @compute @workgroup_size(8, 8, 1)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let x = global_id.x; let y = global_id.y; let batch_c = global_id.z;
+            if (x >= params.out_w || y >= params.out_h || batch_c >= (params.b * params.out_c)) { return; }
+            let ob = batch_c / params.out_c;
+            let oc = batch_c % params.out_c;
+            var sum = 0.0;
+            if (params.has_bias == 1u) { sum = bias[oc]; }
+            for (var ic = 0u; ic < params.in_c; ic = ic + 1u) {
+                for (var ky = 0u; ky < params.kh; ky = ky + 1u) {
+                    for (var kx = 0u; kx < params.kw; kx = kx + 1u) {
+                        let oy_shifted = i32(y) + i32(params.ph) - i32(ky);
+                        let ox_shifted = i32(x) + i32(params.pw) - i32(kx);
+                        if (oy_shifted % i32(params.sh) == 0 && ox_shifted % i32(params.sw) == 0) {
+                            let iy = oy_shifted / i32(params.sh);
+                            let ix = ox_shifted / i32(params.sw);
+                            if (iy >= 0 && iy < i32(params.in_h) && ix >= 0 && ix < i32(params.in_w)) {
+                                let in_val = input[((ob * params.in_h + u32(iy)) * params.in_w + u32(ix)) * params.in_c + ic];
+                                let w_val = weight[ic * (params.out_c * params.kh * params.kw) + oc * (params.kh * params.kw) + ky * params.kw + kx];
+                                sum = sum + in_val * w_val;
+                            }
+                        }
+                    }
+                }
+            }
+            output[((ob * params.out_h + y) * params.out_w + x) * params.out_c + oc] = sum;
+        }

package/shaders/copy.wgsl

@@ -0,0 +1,13 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        struct Params { size : u32 }
+        @group(0) @binding(2) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let idx = global_id.x;
+            if (idx >= params.size) { return; }
+            let x = input[idx];
+            var out_val = x;
+            out_val = x;
+            output[idx] = out_val;
+        }

package/shaders/crossAttention.wgsl

@@ -0,0 +1,140 @@
+@group(0) @binding(0) var<storage, read> q_in : array<f32>;
+            @group(0) @binding(1) var<storage, read> kv_in : array<f32>;
+            @group(0) @binding(2) var<storage, read> weight_int8_packed : array<u32>;
+            @group(0) @binding(3) var<storage, read> w_scale : array<f32>;
+            @group(0) @binding(4) var<storage, read> bias : array<f32>;
+            @group(0) @binding(5) var<storage, read_write> output : array<f32>;
+            
+            struct Params {
+                seq_len_q : u32,
+                seq_len_kv : u32,
+                d_model : u32,
+                num_heads : u32,
+                head_dim : u32,
+                scale_factor : f32,
+                has_scale : u32,
+                has_bias : u32,
+            }
+            @group(0) @binding(6) var<uniform> params : Params;
+            
+            @compute @workgroup_size(64, 1, 1)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let q_idx = global_id.x; 
+                let h_idx = global_id.y; 
+                
+                if (q_idx >= params.seq_len_q || h_idx >= params.num_heads) { return; }
+                
+                let head_dim = params.head_dim;
+                let d_model = params.d_model;
+                let d_model_4 = d_model / 4u;
+                
+                // Cache Q for this head and this q_idx
+                var q_cache : array<f32, 64>; // max head_dim 64
+                for (var d = 0u; d < head_dim; d = d + 1u) {
+                    var sum = 0.0;
+                    let out_col = h_idx * head_dim + d;
+                    for (var i = 0u; i < d_model_4; i = i + 1u) {
+                        let w_packed = weight_int8_packed[out_col * d_model_4 + i];
+                        let in_base = q_idx * d_model + i * 4u;
+                        sum = sum + q_in[in_base + 0u] * f32(extractBits(i32(w_packed), 0u, 8u));
+                        sum = sum + q_in[in_base + 1u] * f32(extractBits(i32(w_packed), 8u, 8u));
+                        sum = sum + q_in[in_base + 2u] * f32(extractBits(i32(w_packed), 16u, 8u));
+                        sum = sum + q_in[in_base + 3u] * f32(extractBits(i32(w_packed), 24u, 8u));
+                    }
+                    if (params.has_scale > 0u) { sum = sum * w_scale[out_col]; }
+                    if (params.has_bias > 0u) { sum = sum + bias[out_col]; }
+                    q_cache[d] = sum;
+                }
+                
+                var max_logit : f32 = -1e38;
+                
+                // Pass 1: find max
+                for (var k_idx = 0u; k_idx < params.seq_len_kv; k_idx = k_idx + 1u) {
+                    var score : f32 = 0.0;
+                    for (var d = 0u; d < head_dim; d = d + 1u) {
+                        var k_val = 0.0;
+                        let out_col = d_model + h_idx * head_dim + d;
+                        for (var i = 0u; i < d_model_4; i = i + 1u) {
+                            let w_packed = weight_int8_packed[out_col * d_model_4 + i];
+                            let in_base = k_idx * d_model + i * 4u;
+                            k_val = k_val + kv_in[in_base + 0u] * f32(extractBits(i32(w_packed), 0u, 8u));
+                            k_val = k_val + kv_in[in_base + 1u] * f32(extractBits(i32(w_packed), 8u, 8u));
+                            k_val = k_val + kv_in[in_base + 2u] * f32(extractBits(i32(w_packed), 16u, 8u));
+                            k_val = k_val + kv_in[in_base + 3u] * f32(extractBits(i32(w_packed), 24u, 8u));
+                        }
+                        if (params.has_scale > 0u) { k_val = k_val * w_scale[out_col]; }
+                        if (params.has_bias > 0u) { k_val = k_val + bias[out_col]; }
+                        
+                        score = score + (q_cache[d] * k_val);
+                    }
+                    score = score * params.scale_factor;
+                    if (score > max_logit) { max_logit = score; }
+                }
+                
+                // Pass 2: sum exp
+                var sum_exp : f32 = 0.0;
+                for (var k_idx = 0u; k_idx < params.seq_len_kv; k_idx = k_idx + 1u) {
+                    var score : f32 = 0.0;
+                    for (var d = 0u; d < head_dim; d = d + 1u) {
+                        var k_val = 0.0;
+                        let out_col = d_model + h_idx * head_dim + d;
+                        for (var i = 0u; i < d_model_4; i = i + 1u) {
+                            let w_packed = weight_int8_packed[out_col * d_model_4 + i];
+                            let in_base = k_idx * d_model + i * 4u;
+                            k_val = k_val + kv_in[in_base + 0u] * f32(extractBits(i32(w_packed), 0u, 8u));
+                            k_val = k_val + kv_in[in_base + 1u] * f32(extractBits(i32(w_packed), 8u, 8u));
+                            k_val = k_val + kv_in[in_base + 2u] * f32(extractBits(i32(w_packed), 16u, 8u));
+                            k_val = k_val + kv_in[in_base + 3u] * f32(extractBits(i32(w_packed), 24u, 8u));
+                        }
+                        if (params.has_scale > 0u) { k_val = k_val * w_scale[out_col]; }
+                        if (params.has_bias > 0u) { k_val = k_val + bias[out_col]; }
+                        
+                        score = score + (q_cache[d] * k_val);
+                    }
+                    score = score * params.scale_factor;
+                    sum_exp = sum_exp + exp(score - max_logit);
+                }
+                
+                // Pass 3: output
+                for (var d = 0u; d < head_dim; d = d + 1u) {
+                    var out_val : f32 = 0.0;
+                    for (var k_idx = 0u; k_idx < params.seq_len_kv; k_idx = k_idx + 1u) {
+                        var score : f32 = 0.0;
+                        for (var kd = 0u; kd < head_dim; kd = kd + 1u) {
+                            var k_val = 0.0;
+                            let out_col = d_model + h_idx * head_dim + kd;
+                            for (var i = 0u; i < d_model_4; i = i + 1u) {
+                                let w_packed = weight_int8_packed[out_col * d_model_4 + i];
+                                let in_base = k_idx * d_model + i * 4u;
+                                k_val = k_val + kv_in[in_base + 0u] * f32(extractBits(i32(w_packed), 0u, 8u));
+                                k_val = k_val + kv_in[in_base + 1u] * f32(extractBits(i32(w_packed), 8u, 8u));
+                                k_val = k_val + kv_in[in_base + 2u] * f32(extractBits(i32(w_packed), 16u, 8u));
+                                k_val = k_val + kv_in[in_base + 3u] * f32(extractBits(i32(w_packed), 24u, 8u));
+                            }
+                            if (params.has_scale > 0u) { k_val = k_val * w_scale[out_col]; }
+                            if (params.has_bias > 0u) { k_val = k_val + bias[out_col]; }
+                            
+                            score = score + (q_cache[kd] * k_val);
+                        }
+                        score = score * params.scale_factor;
+                        
+                        let w = exp(score - max_logit) / sum_exp;
+                        
+                        var v_val = 0.0;
+                        let v_col = d_model * 2u + h_idx * head_dim + d;
+                        for (var i = 0u; i < d_model_4; i = i + 1u) {
+                            let w_packed = weight_int8_packed[v_col * d_model_4 + i];
+                            let in_base = k_idx * d_model + i * 4u;
+                            v_val = v_val + kv_in[in_base + 0u] * f32(extractBits(i32(w_packed), 0u, 8u));
+                            v_val = v_val + kv_in[in_base + 1u] * f32(extractBits(i32(w_packed), 8u, 8u));
+                            v_val = v_val + kv_in[in_base + 2u] * f32(extractBits(i32(w_packed), 16u, 8u));
+                            v_val = v_val + kv_in[in_base + 3u] * f32(extractBits(i32(w_packed), 24u, 8u));
+                        }
+                        if (params.has_scale > 0u) { v_val = v_val * w_scale[v_col]; }
+                        if (params.has_bias > 0u) { v_val = v_val + bias[v_col]; }
+                        
+                        out_val = out_val + (w * v_val);
+                    }
+                    output[q_idx * d_model + (h_idx * head_dim) + d] = out_val;
+                }
+            }

package/shaders/crossAttentionF32.wgsl

@@ -0,0 +1,98 @@
+@group(0) @binding(0) var<storage, read> q_in : array<f32>;
+@group(0) @binding(1) var<storage, read> kv_in : array<f32>;
+@group(0) @binding(2) var<storage, read> weight : array<f32>;
+@group(0) @binding(3) var<storage, read> w_scale : array<f32>;
+@group(0) @binding(4) var<storage, read> bias : array<f32>;
+@group(0) @binding(5) var<storage, read_write> output : array<f32>;
+
+struct Params {
+    seq_len_q : u32,
+    seq_len_kv : u32,
+    d_model : u32,
+    num_heads : u32,
+    head_dim : u32,
+    scale_factor : f32,
+    has_scale : u32,
+    has_bias : u32,
+}
+@group(0) @binding(6) var<uniform> params : Params;
+
+fn project(src : ptr<function, array<f32, 64>>, row : u32, out_col : u32) -> f32 {
+    var sum = 0.0;
+    let w_base = out_col * params.d_model;
+    for (var i = 0u; i < params.d_model; i = i + 1u) {
+        sum = sum + (*src)[i] * weight[w_base + i];
+    }
+    if (params.has_scale > 0u) { sum = sum * w_scale[out_col]; }
+    if (params.has_bias > 0u) { sum = sum + bias[out_col]; }
+    return sum;
+}
+
+@compute @workgroup_size(64, 1, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let q_idx = gid.x;
+    let h_idx = gid.y;
+    if (q_idx >= params.seq_len_q || h_idx >= params.num_heads) { return; }
+    if (params.d_model > 64u || params.head_dim > 64u) { return; }
+
+    var q_src : array<f32, 64>;
+    for (var i = 0u; i < params.d_model; i = i + 1u) {
+        q_src[i] = q_in[q_idx * params.d_model + i];
+    }
+
+    var q_proj : array<f32, 64>;
+    for (var d = 0u; d < params.head_dim; d = d + 1u) {
+        let out_col = h_idx * params.head_dim + d;
+        q_proj[d] = project(&q_src, q_idx, out_col);
+    }
+
+    var max_logit = -1e38;
+    for (var k_idx = 0u; k_idx < params.seq_len_kv; k_idx = k_idx + 1u) {
+        var kv_src : array<f32, 64>;
+        for (var i = 0u; i < params.d_model; i = i + 1u) {
+            kv_src[i] = kv_in[k_idx * params.d_model + i];
+        }
+        var score = 0.0;
+        for (var d = 0u; d < params.head_dim; d = d + 1u) {
+            let out_col = params.d_model + h_idx * params.head_dim + d;
+            score = score + q_proj[d] * project(&kv_src, k_idx, out_col);
+        }
+        score = score * params.scale_factor;
+        if (score > max_logit) { max_logit = score; }
+    }
+
+    var sum_exp = 0.0;
+    for (var k_idx = 0u; k_idx < params.seq_len_kv; k_idx = k_idx + 1u) {
+        var kv_src : array<f32, 64>;
+        for (var i = 0u; i < params.d_model; i = i + 1u) {
+            kv_src[i] = kv_in[k_idx * params.d_model + i];
+        }
+        var score = 0.0;
+        for (var d = 0u; d < params.head_dim; d = d + 1u) {
+            let out_col = params.d_model + h_idx * params.head_dim + d;
+            score = score + q_proj[d] * project(&kv_src, k_idx, out_col);
+        }
+        score = score * params.scale_factor;
+        sum_exp = sum_exp + exp(score - max_logit);
+    }
+
+    for (var d = 0u; d < params.head_dim; d = d + 1u) {
+        var out_val = 0.0;
+        for (var k_idx = 0u; k_idx < params.seq_len_kv; k_idx = k_idx + 1u) {
+            var kv_src : array<f32, 64>;
+            for (var i = 0u; i < params.d_model; i = i + 1u) {
+                kv_src[i] = kv_in[k_idx * params.d_model + i];
+            }
+            var score = 0.0;
+            for (var kd = 0u; kd < params.head_dim; kd = kd + 1u) {
+                let k_col = params.d_model + h_idx * params.head_dim + kd;
+                score = score + q_proj[kd] * project(&kv_src, k_idx, k_col);
+            }
+            score = score * params.scale_factor;
+            let attn = exp(score - max_logit) / sum_exp;
+            let v_col = params.d_model * 2u + h_idx * params.head_dim + d;
+            out_val = out_val + attn * project(&kv_src, k_idx, v_col);
+        }
+        output[q_idx * params.d_model + h_idx * params.head_dim + d] = out_val;
+    }
+}