volvoxai 0.1.0

package/shaders/crossSDPA.wgsl

@@ -0,0 +1,74 @@
+@group(0) @binding(0) var<storage, read> q_in : array<f32>;
+            @group(0) @binding(1) var<storage, read> k_in : array<f32>;
+            @group(0) @binding(2) var<storage, read> v_in : array<f32>;
+            @group(0) @binding(3) 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 : f32,
+            }
+            @group(0) @binding(4) 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;
+                
+                // 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) {
+                    q_cache[d] = q_in[q_idx * d_model + (h_idx * head_dim) + d];
+                }
+                
+                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) {
+                        let k_val = k_in[k_idx * d_model + (h_idx * head_dim) + d];
+                        score = score + (q_cache[d] * k_val);
+                    }
+                    score = score * params.scale;
+                    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) {
+                        let k_val = k_in[k_idx * d_model + (h_idx * head_dim) + d];
+                        score = score + (q_cache[d] * k_val);
+                    }
+                    score = score * params.scale;
+                    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) {
+                            let k_val = k_in[k_idx * d_model + (h_idx * head_dim) + kd];
+                            score = score + (q_cache[kd] * k_val);
+                        }
+                        score = score * params.scale;
+                        
+                        let w = exp(score - max_logit) / sum_exp;
+                        let v_val = v_in[k_idx * d_model + (h_idx * head_dim) + d];
+                        out_val = out_val + (w * v_val);
+                    }
+                    output[q_idx * d_model + (h_idx * head_dim) + d] = out_val;
+                }
+            }

package/shaders/dequantizeLinear.wgsl

@@ -0,0 +1,14 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read> scale : array<f32>;
+        @group(0) @binding(2) var<storage, read> zero_point : array<f32>;
+        @group(0) @binding(3) var<storage, read_write> output : array<f32>;
+        struct Params { size : u32, has_zp : u32 }
+        @group(0) @binding(4) 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; }
+            var zp = 0.0;
+            if (params.has_zp == 1u) { zp = zero_point[0]; }
+            output[idx] = (input[idx] - zp) * scale[0];
+        }

package/shaders/div.wgsl

@@ -0,0 +1,34 @@
+@group(0) @binding(0) var<storage, read> a : array<f32>;
+        @group(0) @binding(1) var<storage, read> b : array<f32>;
+        @group(0) @binding(2) var<storage, read_write> output : array<f32>;
+        
+        struct Params { size : u32, is_b_scalar : u32, b_size : u32, a_size: u32, is_a_scalar: u32 }
+        @group(0) @binding(3) 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; }
+            
+            var a_val : f32 = 0.0;
+            if (params.is_a_scalar == 1u) {
+                a_val = a[0];
+            } else if (params.a_size < params.size && params.a_size > 0u) {
+                a_val = a[idx % params.a_size];
+            } else {
+                a_val = a[idx];
+            }
+            
+            var b_val : f32 = 0.0;
+            if (params.is_b_scalar == 1u) {
+                b_val = b[0];
+            } else if (params.b_size < params.size && params.b_size > 0u) {
+                b_val = b[idx % params.b_size];
+            } else {
+                b_val = b[idx];
+            }
+            
+            var out_val : f32 = 0.0;
+            out_val = a_val / b_val;
+            output[idx] = out_val;
+        }

package/shaders/elementwise.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;
+            undefined
+            output[idx] = out_val;
+        }

package/shaders/embedding.wgsl

@@ -0,0 +1,22 @@
+@group(0) @binding(0) var<storage, read> tokens : array<f32>;
+            @group(0) @binding(1) var<storage, read> weight : array<f32>;
+            @group(0) @binding(2) var<storage, read_write> output : array<f32>;
+            
+            struct Params { seq_len : u32, d_model : u32 }
+            @group(0) @binding(3) var<uniform> params : Params;
+
+            @compute @workgroup_size(64)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let token_idx = global_id.x;
+                if (token_idx >= params.seq_len) { return; }
+                
+                let d_model = params.d_model;
+                let token_id = u32(tokens[token_idx]);   // ids arrive as f32 (see CPU/WASM tiers)
+
+                let in_offset = token_id * d_model;
+                let out_offset = token_idx * d_model;
+                
+                for (var i = 0u; i < d_model; i = i + 1u) {
+                    output[out_offset + i] = weight[in_offset + i];
+                }
+            }

package/shaders/expand.wgsl

@@ -0,0 +1,18 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        struct Params { in_b: u32, in_h: u32, in_w: u32, in_c: u32, out_b: u32, out_h: u32, out_w: u32, out_c: 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;
+            let total = params.out_b * params.out_h * params.out_w * params.out_c;
+            if (idx >= total) { return; }
+            let oc = idx % params.out_c;
+            let ow = (idx / params.out_c) % params.out_w;
+            let oh = (idx / (params.out_c * params.out_w)) % params.out_h;
+            let ob = idx / (params.out_c * params.out_w * params.out_h);
+            let ib = ob % params.in_b;
+            let ih = oh % params.in_h; let iw = ow % params.in_w;
+            let ic = oc % params.in_c;
+            output[idx] = input[((ib * params.in_h + ih) * params.in_w + iw) * params.in_c + ic];
+        }

package/shaders/gELU.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];
+    let cdf = 0.5 * (1.0 + tanh(0.7978845608 * (x + 0.044715 * x * x * x)));
+    output[idx] = x * cdf;
+}

package/shaders/gather.wgsl

@@ -0,0 +1,17 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read> indices : array<f32>;
+@group(0) @binding(2) var<storage, read_write> output : array<f32>;
+// Gather along axis 0: output[i, ...] = input[indices[i], ...]. `row_size` is the
+// number of contiguous elements per gathered row (product of input dims after
+// axis 0); `num_idx` is the number of indices. total = num_idx * row_size.
+struct Params { row_size : u32, num_idx : u32, total : u32 }
+@group(0) @binding(3) var<uniform> p : Params;
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+  let idx = gid.x;
+  if (idx >= p.total) { return; }
+  let k = idx % p.row_size;
+  let i = idx / p.row_size;
+  let row = u32(indices[i]);
+  output[idx] = input[row * p.row_size + k];
+}

package/shaders/generalTranspose.wgsl

@@ -0,0 +1,19 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        @group(0) @binding(2) var<storage, read> md : array<u32>;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+            let idx = gid.x;
+            let total = md[0];
+            if (idx >= total) { return; }
+            let rank = md[1];
+            var rem = idx;
+            var in_idx = 0u;
+            for (var d = 0u; d < rank; d = d + 1u) {
+                let os = md[2u + d];
+                let coord = rem / os;
+                rem = rem - coord * os;
+                in_idx = in_idx + coord * md[2u + rank + d];
+            }
+            output[idx] = input[in_idx];
+        }

package/shaders/globalAveragePool.wgsl

@@ -0,0 +1,19 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read_write> output : array<f32>;
+
+struct Params { n : u32, h : u32, w : u32, c : u32 }
+@group(0) @binding(2) var<uniform> params : Params;
+
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let ch = gid.x;
+    let nb = gid.y;
+    if (nb >= params.n || ch >= params.c) { return; }
+    var sum = 0.0;
+    for (var y = 0u; y < params.h; y = y + 1u) {
+        for (var x = 0u; x < params.w; x = x + 1u) {
+            sum = sum + input[((nb * params.h + y) * params.w + x) * params.c + ch];
+        }
+    }
+    output[nb * params.c + ch] = sum / f32(params.h * params.w);
+}

package/shaders/hardSigmoid.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;
+            var v = x + 3.0; if (v < 0.0) { v = 0.0; } if (v > 6.0) { v = 6.0; } out_val = v / 6.0;
+            output[idx] = out_val;
+        }

package/shaders/hardSwish.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;
+            var v = x + 3.0; if (v < 0.0) { v = 0.0; } if (v > 6.0) { v = 6.0; } out_val = x * v / 6.0;
+            output[idx] = out_val;
+        }

package/shaders/interp1D.wgsl

@@ -0,0 +1,28 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+            @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+            
+            struct Params { in_c : u32, in_l : u32, out_l : u32 }
+            @group(0) @binding(2) var<uniform> params : Params;
+            
+            @compute @workgroup_size(64)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let i = global_id.x;
+                let c = global_id.y;
+                
+                if (i >= params.out_l || c >= params.in_c) { return; }
+                
+                let scale = f32(params.in_l) / f32(params.out_l);
+                var src = (f32(i) + 0.5) * scale - 0.5;
+                if (src < 0.0) { src = 0.0; }
+                if (src > f32(params.in_l - 1u)) { src = f32(params.in_l - 1u); }
+                
+                let lo = u32(src);
+                var hi = lo + 1u;
+                if (hi >= params.in_l) { hi = params.in_l - 1u; }
+                let t = src - f32(lo);
+                
+                let in_base = c * params.in_l;
+                let out_base = c * params.out_l;
+                
+                output[out_base + i] = input[in_base + lo] * (1.0 - t) + input[in_base + hi] * t;
+            }

package/shaders/layerNorm.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 { rows : u32, d_model : u32 }
+            @group(0) @binding(4) var<uniform> params : Params;
+
+            @compute @workgroup_size(64)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let row = global_id.x;
+                if (row >= params.rows) { return; }
+                let d_model = params.d_model;
+                let offset = row * d_model;
+                
+                var sum : f32 = 0.0;
+                var sq_sum : f32 = 0.0;
+                
+                for (var i = 0u; i < d_model; i = i + 1u) {
+                    let val = input[offset + i];
+                    sum = sum + val;
+                    sq_sum = sq_sum + (val * val);
+                }
+                
+                let mean = sum / f32(d_model);
+                let variance = (sq_sum / f32(d_model)) - (mean * mean);
+                let inv_std = inverseSqrt(variance + 1e-5);
+                
+                for (var i = 0u; i < d_model; i = i + 1u) {
+                    let norm_val = (input[offset + i] - mean) * inv_std;
+                    output[offset + i] = norm_val * weight[i] + bias[i];
+                }
+            }

package/shaders/leakyReLU.wgsl

@@ -0,0 +1,11 @@
+@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, alpha : f32 }
+        @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];
+            if (x > 0.0) { output[idx] = x; } else { output[idx] = x * params.alpha; }
+        }

package/shaders/linearF32.wgsl

@@ -0,0 +1,33 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+            @group(0) @binding(1) var<storage, read> weight_f32 : array<f32>;
+            @group(0) @binding(2) var<storage, read> dummyScale : array<f32>;
+            @group(0) @binding(3) var<storage, read> bias : array<f32>;
+            @group(0) @binding(4) var<storage, read_write> output : array<f32>;
+
+            struct Params {
+                seq_len : u32,
+                d_in : u32,
+                d_out : u32,
+            }
+            @group(0) @binding(5) var<uniform> params : Params;
+
+            @compute @workgroup_size(64, 1, 1)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let row = global_id.y;
+                let col = global_id.x;
+                
+                if (row >= params.seq_len || col >= params.d_out) { return; }
+                
+                var sum : f32 = 0.0;
+                
+                for (var k = 0u; k < params.d_in; k = k + 1u) {
+                    let in_val = input[row * params.d_in + k];
+                    let w_val = weight_f32[col * params.d_in + k];
+                    sum = sum + in_val * w_val;
+                }
+                
+                let b_val = bias[col];
+                // dummyScale is passed just to keep bindings consistent but not used here.
+                
+                output[row * params.d_out + col] = sum + b_val;
+            }

package/shaders/linearF32RowMajor.wgsl

@@ -0,0 +1,24 @@
+@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 {
+    seq_len : u32,
+    d_in : u32,
+    d_out : u32,
+}
+@group(0) @binding(4) var<uniform> params : Params;
+
+@compute @workgroup_size(16, 16, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let col = gid.x;
+    let row = gid.y;
+    if (row >= params.seq_len || col >= params.d_out) { return; }
+
+    var sum = bias[col];
+    for (var k = 0u; k < params.d_in; k = k + 1u) {
+        sum = sum + input[row * params.d_in + k] * weight[k * params.d_out + col];
+    }
+    output[row * params.d_out + col] = sum;
+}

package/shaders/linearInt8.wgsl

@@ -0,0 +1,42 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+            @group(0) @binding(1) var<storage, read> weight_int8_packed : array<u32>;
+            @group(0) @binding(2) var<storage, read> weight_scales : array<f32>;
+            @group(0) @binding(3) var<storage, read> bias : array<f32>;
+            @group(0) @binding(4) var<storage, read_write> output : array<f32>;
+
+            struct Params {
+                seq_len : u32,
+                d_in : u32,
+                d_out : u32,
+            }
+            @group(0) @binding(5) var<uniform> params : Params;
+
+            // Unpack one signed 8-bit integer from a 32-bit packed block
+            fn unpack_i8(packed: u32, byte_idx: u32) -> f32 {
+                let val_i32 = extractBits(i32(packed), byte_idx * 8u, 8u);
+                return f32(val_i32);
+            }
+
+            @compute @workgroup_size(64, 1, 1)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let row = global_id.y;
+                let col = global_id.x;
+                
+                if (row >= params.seq_len || col >= params.d_out) { return; }
+                
+                var sum : f32 = 0.0;
+                let d_in_4 = params.d_in / 4u;
+                
+                for (var i = 0u; i < d_in_4; i = i + 1u) {
+                    let w_packed = weight_int8_packed[col * d_in_4 + i];
+                    let in_base = row * params.d_in + i * 4u;
+                    
+                    sum = sum + input[in_base + 0u] * f32(extractBits(i32(w_packed), 0u, 8u));
+                    sum = sum + input[in_base + 1u] * f32(extractBits(i32(w_packed), 8u, 8u));
+                    sum = sum + input[in_base + 2u] * f32(extractBits(i32(w_packed), 16u, 8u));
+                    sum = sum + input[in_base + 3u] * f32(extractBits(i32(w_packed), 24u, 8u));
+                }
+                
+                let scale = weight_scales[col];
+                output[row * params.d_out + col] = (sum * scale) + bias[col];
+            }

package/shaders/logSoftmax.wgsl

@@ -0,0 +1,22 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        struct Params { b : u32, d : u32 }
+        @group(0) @binding(2) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let i = global_id.x;
+            if (i >= params.b) { return; }
+            let offset = i * params.d;
+            var max_val = -100000.0;
+            for (var j = 0u; j < params.d; j = j + 1u) {
+                if (input[offset + j] > max_val) { max_val = input[offset + j]; }
+            }
+            var sum = 0.0;
+            for (var j = 0u; j < params.d; j = j + 1u) {
+                sum = sum + exp(input[offset + j] - max_val);
+            }
+            let logSum = log(sum);
+            for (var j = 0u; j < params.d; j = j + 1u) {
+                output[offset + j] = (input[offset + j] - max_val) - logSum;
+            }
+        }

package/shaders/maxPool2D.wgsl

@@ -0,0 +1,37 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read_write> output : array<f32>;
+
+struct Params {
+    h : u32,
+    w : u32,
+    c : u32,
+    out_h : u32,
+    out_w : u32,
+    ky : u32,
+    kx : u32,
+    sy : u32,
+    sx : u32,
+    py : u32,
+    px : u32,
+}
+@group(0) @binding(2) var<uniform> params : Params;
+
+@compute @workgroup_size(8, 8, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let ox = gid.x;
+    let oy = gid.y;
+    let ch = gid.z;
+    if (ox >= params.out_w || oy >= params.out_h || ch >= params.c) { return; }
+
+    var best = -3.402823466e38;
+    for (var yy = 0u; yy < params.ky; yy = yy + 1u) {
+        let iy = i32(oy * params.sy + yy) - i32(params.py);
+        if (iy < 0 || iy >= i32(params.h)) { continue; }
+        for (var xx = 0u; xx < params.kx; xx = xx + 1u) {
+            let ix = i32(ox * params.sx + xx) - i32(params.px);
+            if (ix < 0 || ix >= i32(params.w)) { continue; }
+            best = max(best, input[(u32(iy) * params.w + u32(ix)) * params.c + ch]);
+        }
+    }
+    output[(oy * params.out_w + ox) * params.c + ch] = best;
+}

package/shaders/meanHeight.wgsl

@@ -0,0 +1,18 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+            @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+
+            struct Params { in_h : u32, in_w : u32, in_c : u32 }
+            @group(0) @binding(2) var<uniform> params : Params;
+
+            @compute @workgroup_size(64)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let x = global_id.x;
+                let c = global_id.y;
+                if (x >= params.in_w || c >= params.in_c) { return; }
+
+                var sum_val = 0.0;
+                for (var y = 0u; y < params.in_h; y = y + 1u) {
+                    sum_val = sum_val + input[(y * params.in_w + x) * params.in_c + c];
+                }
+                output[c * params.in_w + x] = sum_val / f32(params.in_h);
+            }

package/shaders/mul.wgsl

@@ -0,0 +1,32 @@
+@group(0) @binding(0) var<storage, read> a : array<f32>;
+@group(0) @binding(1) var<storage, read> b : array<f32>;
+@group(0) @binding(2) var<storage, read_write> output : array<f32>;
+
+struct Params { size : u32, is_b_scalar : u32, b_size : u32, a_size: u32, is_a_scalar: u32 }
+@group(0) @binding(3) 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; }
+
+    var a_val : f32 = 0.0;
+    if (params.is_a_scalar == 1u) {
+        a_val = a[0];
+    } else if (params.a_size < params.size && params.a_size > 0u) {
+        a_val = a[idx % params.a_size];
+    } else {
+        a_val = a[idx];
+    }
+
+    var b_val : f32 = 0.0;
+    if (params.is_b_scalar == 1u) {
+        b_val = b[0];
+    } else if (params.b_size < params.size && params.b_size > 0u) {
+        b_val = b[idx % params.b_size];
+    } else {
+        b_val = b[idx];
+    }
+
+    output[idx] = a_val * b_val;
+}