@simulatte/doppler 0.1.9 → 0.2.0

package/src/gpu/kernels/matmul_f16_tiled.wgsl

@@ -1,165 +0,0 @@
-// Register-Tiled Matrix Multiplication Kernel - FP16
-//
-// High-throughput matmul for prefill using 4x4 register tiling.
-// Each thread computes a 4x4 tile of C (16 outputs), vs 1 output in matmul_f16.
-// Workgroup tile: 64x64 with 16x16 = 256 threads.
-// This yields 16x more outputs per workgroup than the base f16 kernel.
-//
-// C[M,N] = A[M,K] * B^T[N,K]  (transpose_b=1, SafeTensors/RDRR row-major layout)
-//
-// Design:
-// - f16 storage, f32 accumulators for numerical stability
-// - coalesced scalar loads into shared tiles
-// - Shared memory tiles: 64 x TILE_K for A and B
-// - B is always assumed row-major (transpose_b=1) for RDRR models
-
-enable f16;
-
-// Tile dimensions
-const TILE_M: u32 = 64u;   // Rows of C per workgroup
-const TILE_N: u32 = 64u;   // Cols of C per workgroup
-const TILE_K: u32 = 16u;   // K-dimension tile step
-
-// Thread tile: each thread computes a 4x4 block of C
-const THREAD_M: u32 = 4u;
-const THREAD_N: u32 = 4u;
-
-// Workgroup dimensions: 16x16 = 256 threads
-// 16 threads along M covers 64 rows (4 per thread)
-// 16 threads along N covers 64 cols (4 per thread)
-const WG_M: u32 = 16u;
-const WG_N: u32 = 16u;
-
-// Shared memory: TILE_M * TILE_K + TILE_N * TILE_K
-// = 64*16 + 64*16 = 2048 f16 elements = 4096 bytes total
-var<workgroup> tileA: array<f16, 1024>;  // TILE_M * TILE_K = 64 * 16
-var<workgroup> tileB: array<f16, 1024>;  // TILE_N * TILE_K = 64 * 16
-
-struct Uniforms {
-    M: u32,
-    N: u32,
-    K: u32,
-    alpha: f32,
-    transpose_b: u32, // reserved for uniform compatibility; tiled path assumes transpose_b=1
-    _pad0: u32,
-    _pad1: u32,
-    _pad2: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> A: array<f16>;
-@group(0) @binding(2) var<storage, read> B: array<f16>;
-@group(0) @binding(3) var<storage, read_write> C: array<f16>;
-
-@compute @workgroup_size(WG_M, WG_N, 1)
-fn main(
-    @builtin(local_invocation_id) local_id: vec3<u32>,
-    @builtin(workgroup_id) wg_id: vec3<u32>
-) {
-    let tx = local_id.x;  // 0..15, maps to M dimension
-    let ty = local_id.y;  // 0..15, maps to N dimension
-    let tid = tx * WG_N + ty;  // Linear thread index 0..255
-
-    // Global row/col base for this thread's 4x4 output tile
-    let row_base = wg_id.x * TILE_M + tx * THREAD_M;
-    let col_base = wg_id.y * TILE_N + ty * THREAD_N;
-
-    // f32 accumulators: 4x4 = 16 per thread
-    var acc00: f32 = 0.0; var acc01: f32 = 0.0; var acc02: f32 = 0.0; var acc03: f32 = 0.0;
-    var acc10: f32 = 0.0; var acc11: f32 = 0.0; var acc12: f32 = 0.0; var acc13: f32 = 0.0;
-    var acc20: f32 = 0.0; var acc21: f32 = 0.0; var acc22: f32 = 0.0; var acc23: f32 = 0.0;
-    var acc30: f32 = 0.0; var acc31: f32 = 0.0; var acc32: f32 = 0.0; var acc33: f32 = 0.0;
-
-    let num_tiles = (u.K + TILE_K - 1u) / TILE_K;
-
-    for (var t: u32 = 0u; t < num_tiles; t = t + 1u) {
-        let k_offset = t * TILE_K;
-
-        // Load A tile: 64 rows x 16 cols = 1024 elements, 256 threads -> 4 elements each
-        // Thread tid loads elements at indices: tid*4, tid*4+1, tid*4+2, tid*4+3
-        // Mapping: element i -> row = i / TILE_K, col = i % TILE_K
-        {
-            let load_base = tid * 4u;
-            for (var i: u32 = 0u; i < 4u; i = i + 1u) {
-                let elem_idx = load_base + i;
-                let load_row = elem_idx / TILE_K;
-                let load_col = elem_idx % TILE_K;
-                let global_row = wg_id.x * TILE_M + load_row;
-                let global_col = k_offset + load_col;
-                if (global_row < u.M && global_col < u.K) {
-                    tileA[elem_idx] = A[global_row * u.K + global_col];
-                } else {
-                    tileA[elem_idx] = f16(0.0);
-                }
-            }
-        }
-
-        // Load B tile: 64 rows x 16 cols = 1024 elements
-        // B is [N, K] (row-major/transposed): B[col, k] = B[col * K + k]
-        // tileB layout: tileB[n_local * TILE_K + k_local] = B[global_n, global_k]
-        {
-            let load_base = tid * 4u;
-            for (var i: u32 = 0u; i < 4u; i = i + 1u) {
-                let elem_idx = load_base + i;
-                let load_row = elem_idx / TILE_K;  // N-dimension local offset
-                let load_col = elem_idx % TILE_K;  // K-dimension local offset
-                let global_n = wg_id.y * TILE_N + load_row;
-                let global_k = k_offset + load_col;
-                if (global_n < u.N && global_k < u.K) {
-                    tileB[elem_idx] = B[global_n * u.K + global_k];
-                } else {
-                    tileB[elem_idx] = f16(0.0);
-                }
-            }
-        }
-
-        workgroupBarrier();
-
-        // Compute: each thread accumulates its 4x4 output tile
-        // A tile: tileA[m_local * TILE_K + k] for rows tx*4..tx*4+3
-        // B tile: tileB[n_local * TILE_K + k] for cols ty*4..ty*4+3
-        for (var k: u32 = 0u; k < TILE_K; k = k + 1u) {
-            // Load 4 A values for this thread's rows
-            let a0 = f32(tileA[(tx * THREAD_M + 0u) * TILE_K + k]);
-            let a1 = f32(tileA[(tx * THREAD_M + 1u) * TILE_K + k]);
-            let a2 = f32(tileA[(tx * THREAD_M + 2u) * TILE_K + k]);
-            let a3 = f32(tileA[(tx * THREAD_M + 3u) * TILE_K + k]);
-
-            // Load 4 B values for this thread's columns
-            let b0 = f32(tileB[(ty * THREAD_N + 0u) * TILE_K + k]);
-            let b1 = f32(tileB[(ty * THREAD_N + 1u) * TILE_K + k]);
-            let b2 = f32(tileB[(ty * THREAD_N + 2u) * TILE_K + k]);
-            let b3 = f32(tileB[(ty * THREAD_N + 3u) * TILE_K + k]);
-
-            // Outer product: 4x4 accumulation
-            acc00 += a0 * b0; acc01 += a0 * b1; acc02 += a0 * b2; acc03 += a0 * b3;
-            acc10 += a1 * b0; acc11 += a1 * b1; acc12 += a1 * b2; acc13 += a1 * b3;
-            acc20 += a2 * b0; acc21 += a2 * b1; acc22 += a2 * b2; acc23 += a2 * b3;
-            acc30 += a3 * b0; acc31 += a3 * b1; acc32 += a3 * b2; acc33 += a3 * b3;
-        }
-
-        workgroupBarrier();
-    }
-
-    // Write 4x4 output tile
-    let alpha = u.alpha;
-    if (row_base + 0u < u.M && col_base + 0u < u.N) { C[(row_base + 0u) * u.N + col_base + 0u] = f16(acc00 * alpha); }
-    if (row_base + 0u < u.M && col_base + 1u < u.N) { C[(row_base + 0u) * u.N + col_base + 1u] = f16(acc01 * alpha); }
-    if (row_base + 0u < u.M && col_base + 2u < u.N) { C[(row_base + 0u) * u.N + col_base + 2u] = f16(acc02 * alpha); }
-    if (row_base + 0u < u.M && col_base + 3u < u.N) { C[(row_base + 0u) * u.N + col_base + 3u] = f16(acc03 * alpha); }
-
-    if (row_base + 1u < u.M && col_base + 0u < u.N) { C[(row_base + 1u) * u.N + col_base + 0u] = f16(acc10 * alpha); }
-    if (row_base + 1u < u.M && col_base + 1u < u.N) { C[(row_base + 1u) * u.N + col_base + 1u] = f16(acc11 * alpha); }
-    if (row_base + 1u < u.M && col_base + 2u < u.N) { C[(row_base + 1u) * u.N + col_base + 2u] = f16(acc12 * alpha); }
-    if (row_base + 1u < u.M && col_base + 3u < u.N) { C[(row_base + 1u) * u.N + col_base + 3u] = f16(acc13 * alpha); }
-
-    if (row_base + 2u < u.M && col_base + 0u < u.N) { C[(row_base + 2u) * u.N + col_base + 0u] = f16(acc20 * alpha); }
-    if (row_base + 2u < u.M && col_base + 1u < u.N) { C[(row_base + 2u) * u.N + col_base + 1u] = f16(acc21 * alpha); }
-    if (row_base + 2u < u.M && col_base + 2u < u.N) { C[(row_base + 2u) * u.N + col_base + 2u] = f16(acc22 * alpha); }
-    if (row_base + 2u < u.M && col_base + 3u < u.N) { C[(row_base + 2u) * u.N + col_base + 3u] = f16(acc23 * alpha); }
-
-    if (row_base + 3u < u.M && col_base + 0u < u.N) { C[(row_base + 3u) * u.N + col_base + 0u] = f16(acc30 * alpha); }
-    if (row_base + 3u < u.M && col_base + 1u < u.N) { C[(row_base + 3u) * u.N + col_base + 1u] = f16(acc31 * alpha); }
-    if (row_base + 3u < u.M && col_base + 2u < u.N) { C[(row_base + 3u) * u.N + col_base + 2u] = f16(acc32 * alpha); }
-    if (row_base + 3u < u.M && col_base + 3u < u.N) { C[(row_base + 3u) * u.N + col_base + 3u] = f16(acc33 * alpha); }
-}

package/src/gpu/kernels/matmul_f16w_f32a.wgsl

@@ -1,89 +0,0 @@
-// Matrix Multiplication Kernel - f16 weights, f32 activations
-//
-// A is f32 (activations), B is f16 (weights), C is f32.
-// C[M,N] = A[M,K] * B[K,N]  (or B^T when transpose_b=1)
-// This reduces weight bandwidth and leverages f16 load/throughput
-// while keeping f32 outputs for compatibility with f32-only ops.
-
-enable f16;
-
-const MAX_TILE_SIZE: u32 = 16u;
-const TILE_AREA: u32 = MAX_TILE_SIZE * MAX_TILE_SIZE;
-
-override TILE_SIZE: u32 = 16u;  // Must be <= MAX_TILE_SIZE
-
-struct Uniforms {
-    M: u32,
-    N: u32,
-    K: u32,
-    alpha: f32,
-    transpose_b: u32,  // 0 = normal, 1 = B is stored transposed [N,K]
-    _pad0: u32,
-    _pad1: u32,
-    _pad2: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> A: array<f32>;
-@group(0) @binding(2) var<storage, read> B: array<f16>;
-@group(0) @binding(3) var<storage, read_write> C: array<f32>;
-
-var<workgroup> tileA: array<f32, TILE_AREA>;
-// Keep workgroup storage in f32 for broader Vulkan/Dawn compatibility on mobile.
-var<workgroup> tileB: array<f32, TILE_AREA>;
-
-@compute @workgroup_size(TILE_SIZE, TILE_SIZE, 1)
-fn main(
-    @builtin(global_invocation_id) gid: vec3<u32>,
-    @builtin(local_invocation_id) lid: vec3<u32>
-) {
-    if (TILE_SIZE > MAX_TILE_SIZE) {
-        return;
-    }
-
-    let row = gid.x;
-    let col = gid.y;
-    let local_row = lid.x;
-    let local_col = lid.y;
-
-    var sum: f32 = 0.0;
-
-    let num_tiles = (u.K + TILE_SIZE - 1u) / TILE_SIZE;
-    let tile_idx = local_row * TILE_SIZE + local_col;
-
-    for (var t: u32 = 0u; t < num_tiles; t = t + 1u) {
-        let a_col = t * TILE_SIZE + local_col;
-        let b_row = t * TILE_SIZE + local_row;
-
-        if (row < u.M && a_col < u.K) {
-            tileA[tile_idx] = A[row * u.K + a_col];
-        } else {
-            tileA[tile_idx] = 0.0;
-        }
-
-        if (b_row < u.K && col < u.N) {
-            if (u.transpose_b == 0u) {
-                tileB[tile_idx] = f32(B[b_row * u.N + col]);
-            } else {
-                // B is [N, K], access element [col, b_row]
-                tileB[tile_idx] = f32(B[col * u.K + b_row]);
-            }
-        } else {
-            tileB[tile_idx] = 0.0;
-        }
-
-        workgroupBarrier();
-
-        for (var k: u32 = 0u; k < TILE_SIZE; k = k + 1u) {
-            let a_val = tileA[local_row * TILE_SIZE + k];
-            let b_val = tileB[k * TILE_SIZE + local_col];
-            sum = sum + a_val * b_val;
-        }
-
-        workgroupBarrier();
-    }
-
-    if (row < u.M && col < u.N) {
-        C[row * u.N + col] = sum * u.alpha;
-    }
-}

package/src/gpu/kernels/matmul_f16w_f32a_tiled.wgsl

@@ -1,154 +0,0 @@
-// Register-Tiled Matrix Multiplication Kernel - f16 weights, f32 activations
-//
-// High-throughput prefill matmul: f32 activations (A), f16 weights (B), f32 output (C).
-// Each thread computes a 4x4 tile of C (16 outputs), vs 1 output in matmul_f16w_f32a.
-// Workgroup tile: 64x64 with 16x16 = 256 threads.
-// This yields 16x more outputs per workgroup than the base f16w_f32a kernel.
-//
-// C[M,N] = A[M,K] * B^T[N,K]  (transpose_b=1, RDRR row-major layout)
-//
-// Design:
-// - f32 A storage (activations), f16 B storage (weights)
-// - f32 shared tiles for broader Vulkan/Dawn compatibility
-// - f32 accumulators for numerical stability
-// - B loaded from f16 global into f32 shared tile on staging
-
-enable f16;
-
-// Tile dimensions
-const TILE_M: u32 = 64u;   // Rows of C per workgroup
-const TILE_N: u32 = 64u;   // Cols of C per workgroup
-const TILE_K: u32 = 16u;   // K-dimension tile step
-
-// Thread tile: each thread computes a 4x4 block of C
-const THREAD_M: u32 = 4u;
-const THREAD_N: u32 = 4u;
-
-// Workgroup dimensions: 16x16 = 256 threads
-const WG_M: u32 = 16u;
-const WG_N: u32 = 16u;
-
-// Shared memory: TILE_M * TILE_K + TILE_N * TILE_K
-// tileA: 64*16 f32 = 4096 bytes, tileB: 64*16 f32 = 4096 bytes → 8192 bytes total
-var<workgroup> tileA: array<f32, 1024>;  // TILE_M * TILE_K = 64 * 16
-var<workgroup> tileB: array<f32, 1024>;  // TILE_N * TILE_K = 64 * 16
-
-struct Uniforms {
-    M: u32,
-    N: u32,
-    K: u32,
-    alpha: f32,
-    transpose_b: u32,
-    _pad0: u32,
-    _pad1: u32,
-    _pad2: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> A: array<f32>;
-@group(0) @binding(2) var<storage, read> B: array<f16>;
-@group(0) @binding(3) var<storage, read_write> C: array<f32>;
-
-@compute @workgroup_size(WG_M, WG_N, 1)
-fn main(
-    @builtin(local_invocation_id) local_id: vec3<u32>,
-    @builtin(workgroup_id) wg_id: vec3<u32>
-) {
-    let tx = local_id.x;  // 0..15, maps to M dimension
-    let ty = local_id.y;  // 0..15, maps to N dimension
-    let tid = tx * WG_N + ty;  // Linear thread index 0..255
-
-    // Global row/col base for this thread's 4x4 output tile
-    let row_base = wg_id.x * TILE_M + tx * THREAD_M;
-    let col_base = wg_id.y * TILE_N + ty * THREAD_N;
-
-    // f32 accumulators: 4x4 = 16 per thread
-    var acc00: f32 = 0.0; var acc01: f32 = 0.0; var acc02: f32 = 0.0; var acc03: f32 = 0.0;
-    var acc10: f32 = 0.0; var acc11: f32 = 0.0; var acc12: f32 = 0.0; var acc13: f32 = 0.0;
-    var acc20: f32 = 0.0; var acc21: f32 = 0.0; var acc22: f32 = 0.0; var acc23: f32 = 0.0;
-    var acc30: f32 = 0.0; var acc31: f32 = 0.0; var acc32: f32 = 0.0; var acc33: f32 = 0.0;
-
-    let num_tiles = (u.K + TILE_K - 1u) / TILE_K;
-
-    for (var t: u32 = 0u; t < num_tiles; t = t + 1u) {
-        let k_offset = t * TILE_K;
-
-        // Load A tile (f32): 64 rows x 16 cols = 1024 elements, 256 threads → 4 each
-        {
-            let load_base = tid * 4u;
-            for (var i: u32 = 0u; i < 4u; i = i + 1u) {
-                let elem_idx = load_base + i;
-                let load_row = elem_idx / TILE_K;
-                let load_col = elem_idx % TILE_K;
-                let global_row = wg_id.x * TILE_M + load_row;
-                let global_col = k_offset + load_col;
-                if (global_row < u.M && global_col < u.K) {
-                    tileA[elem_idx] = A[global_row * u.K + global_col];
-                } else {
-                    tileA[elem_idx] = 0.0;
-                }
-            }
-        }
-
-        // Load B tile (f16 → f32): B is [N, K] (row-major/transposed)
-        {
-            let load_base = tid * 4u;
-            for (var i: u32 = 0u; i < 4u; i = i + 1u) {
-                let elem_idx = load_base + i;
-                let load_row = elem_idx / TILE_K;  // N-dimension local offset
-                let load_col = elem_idx % TILE_K;  // K-dimension local offset
-                let global_n = wg_id.y * TILE_N + load_row;
-                let global_k = k_offset + load_col;
-                if (global_n < u.N && global_k < u.K) {
-                    tileB[elem_idx] = f32(B[global_n * u.K + global_k]);
-                } else {
-                    tileB[elem_idx] = 0.0;
-                }
-            }
-        }
-
-        workgroupBarrier();
-
-        // Compute: each thread accumulates its 4x4 output tile
-        for (var k: u32 = 0u; k < TILE_K; k = k + 1u) {
-            let a0 = tileA[(tx * THREAD_M + 0u) * TILE_K + k];
-            let a1 = tileA[(tx * THREAD_M + 1u) * TILE_K + k];
-            let a2 = tileA[(tx * THREAD_M + 2u) * TILE_K + k];
-            let a3 = tileA[(tx * THREAD_M + 3u) * TILE_K + k];
-
-            let b0 = tileB[(ty * THREAD_N + 0u) * TILE_K + k];
-            let b1 = tileB[(ty * THREAD_N + 1u) * TILE_K + k];
-            let b2 = tileB[(ty * THREAD_N + 2u) * TILE_K + k];
-            let b3 = tileB[(ty * THREAD_N + 3u) * TILE_K + k];
-
-            acc00 += a0 * b0; acc01 += a0 * b1; acc02 += a0 * b2; acc03 += a0 * b3;
-            acc10 += a1 * b0; acc11 += a1 * b1; acc12 += a1 * b2; acc13 += a1 * b3;
-            acc20 += a2 * b0; acc21 += a2 * b1; acc22 += a2 * b2; acc23 += a2 * b3;
-            acc30 += a3 * b0; acc31 += a3 * b1; acc32 += a3 * b2; acc33 += a3 * b3;
-        }
-
-        workgroupBarrier();
-    }
-
-    // Write 4x4 output tile (f32)
-    let alpha = u.alpha;
-    if (row_base + 0u < u.M && col_base + 0u < u.N) { C[(row_base + 0u) * u.N + col_base + 0u] = acc00 * alpha; }
-    if (row_base + 0u < u.M && col_base + 1u < u.N) { C[(row_base + 0u) * u.N + col_base + 1u] = acc01 * alpha; }
-    if (row_base + 0u < u.M && col_base + 2u < u.N) { C[(row_base + 0u) * u.N + col_base + 2u] = acc02 * alpha; }
-    if (row_base + 0u < u.M && col_base + 3u < u.N) { C[(row_base + 0u) * u.N + col_base + 3u] = acc03 * alpha; }
-
-    if (row_base + 1u < u.M && col_base + 0u < u.N) { C[(row_base + 1u) * u.N + col_base + 0u] = acc10 * alpha; }
-    if (row_base + 1u < u.M && col_base + 1u < u.N) { C[(row_base + 1u) * u.N + col_base + 1u] = acc11 * alpha; }
-    if (row_base + 1u < u.M && col_base + 2u < u.N) { C[(row_base + 1u) * u.N + col_base + 2u] = acc12 * alpha; }
-    if (row_base + 1u < u.M && col_base + 3u < u.N) { C[(row_base + 1u) * u.N + col_base + 3u] = acc13 * alpha; }
-
-    if (row_base + 2u < u.M && col_base + 0u < u.N) { C[(row_base + 2u) * u.N + col_base + 0u] = acc20 * alpha; }
-    if (row_base + 2u < u.M && col_base + 1u < u.N) { C[(row_base + 2u) * u.N + col_base + 1u] = acc21 * alpha; }
-    if (row_base + 2u < u.M && col_base + 2u < u.N) { C[(row_base + 2u) * u.N + col_base + 2u] = acc22 * alpha; }
-    if (row_base + 2u < u.M && col_base + 3u < u.N) { C[(row_base + 2u) * u.N + col_base + 3u] = acc23 * alpha; }
-
-    if (row_base + 3u < u.M && col_base + 0u < u.N) { C[(row_base + 3u) * u.N + col_base + 0u] = acc30 * alpha; }
-    if (row_base + 3u < u.M && col_base + 1u < u.N) { C[(row_base + 3u) * u.N + col_base + 1u] = acc31 * alpha; }
-    if (row_base + 3u < u.M && col_base + 2u < u.N) { C[(row_base + 3u) * u.N + col_base + 2u] = acc32 * alpha; }
-    if (row_base + 3u < u.M && col_base + 3u < u.N) { C[(row_base + 3u) * u.N + col_base + 3u] = acc33 * alpha; }
-}

package/src/gpu/kernels/matmul_f32.wgsl

@@ -1,100 +0,0 @@
-// Matrix Multiplication Kernel - FP32 Fallback
-//
-// Tiled matrix multiplication using shared memory (workgroup storage)
-// C[M,N] = A[M,K] * B[K,N]  (or B^T when transpose_b=1)
-//
-// This is the fallback kernel when shader-f16 is unavailable.
-// Uses 16x16 tiles for good occupancy across devices.
-
-// Tile dimensions - optimized for 256 threads per workgroup
-const MAX_TILE_SIZE: u32 = 16u;
-const TILE_AREA: u32 = MAX_TILE_SIZE * MAX_TILE_SIZE;
-
-override TILE_SIZE: u32 = 16u;  // Must be <= MAX_TILE_SIZE
-
-// Uniforms for matrix dimensions
-struct Uniforms {
-    M: u32,           // Rows of A and C
-    N: u32,           // Cols of B and C (or rows of B when transposed)
-    K: u32,           // Cols of A, Rows of B (or cols of B when transposed)
-    alpha: f32,       // Scaling factor (typically 1.0)
-    transpose_b: u32, // 0 = normal, 1 = B is stored transposed [N,K] -> treat as [K,N]
-    _pad0: u32,
-    _pad1: u32,
-    _pad2: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> A: array<f32>;
-@group(0) @binding(2) var<storage, read> B: array<f32>;
-@group(0) @binding(3) var<storage, read_write> C: array<f32>;
-
-// Shared memory tiles for A and B
-var<workgroup> tileA: array<f32, TILE_AREA>;
-var<workgroup> tileB: array<f32, TILE_AREA>;
-
-@compute @workgroup_size(TILE_SIZE, TILE_SIZE, 1)
-fn main(
-    @builtin(global_invocation_id) global_id: vec3<u32>,
-    @builtin(local_invocation_id) local_id: vec3<u32>,
-    @builtin(workgroup_id) workgroup_id: vec3<u32>
-) {
-    if (TILE_SIZE > MAX_TILE_SIZE) {
-        return;
-    }
-
-    let row = global_id.x;
-    let col = global_id.y;
-    let local_row = local_id.x;
-    let local_col = local_id.y;
-
-    // Accumulator for dot product
-    var sum: f32 = 0.0;
-
-    // Number of tiles needed to cover K dimension
-    let num_tiles = (u.K + TILE_SIZE - 1u) / TILE_SIZE;
-
-    // Iterate over tiles
-    for (var t: u32 = 0u; t < num_tiles; t = t + 1u) {
-        // Global indices for loading tiles
-        let a_col = t * TILE_SIZE + local_col;
-        let b_row = t * TILE_SIZE + local_row;
-
-        // Load tile from A into shared memory (with bounds check)
-        let tile_idx = local_row * TILE_SIZE + local_col;
-        if (row < u.M && a_col < u.K) {
-            tileA[tile_idx] = A[row * u.K + a_col];
-        } else {
-            tileA[tile_idx] = 0.0;
-        }
-
-        // Load tile from B into shared memory (with bounds check)
-        // When transpose_b=1, B is stored as [N,K] so we access B[col,b_row] instead of B[b_row,col]
-        if (b_row < u.K && col < u.N) {
-            if (u.transpose_b == 0u) {
-                tileB[tile_idx] = B[b_row * u.N + col];
-            } else {
-                // B is [N, K], access element [col, b_row]
-                tileB[tile_idx] = B[col * u.K + b_row];
-            }
-        } else {
-            tileB[tile_idx] = 0.0;
-        }
-
-        // Synchronize to ensure tile is fully loaded
-        workgroupBarrier();
-
-        // Compute partial dot product for this tile
-        for (var k: u32 = 0u; k < TILE_SIZE; k = k + 1u) {
-            sum = sum + tileA[local_row * TILE_SIZE + k] * tileB[k * TILE_SIZE + local_col];
-        }
-
-        // Synchronize before loading next tile
-        workgroupBarrier();
-    }
-
-    // Write result (with bounds check)
-    if (row < u.M && col < u.N) {
-        C[row * u.N + col] = sum * u.alpha;
-    }
-}

package/src/gpu/kernels/matmul_gemv.wgsl

@@ -1,80 +0,0 @@
-// Optimized GEMV Kernel - Parallel Reduction over K
-// For M=1 decode: C[N] = A[K] * B^T[K,N]
-//
-// Key optimizations:
-// 1. Each workgroup computes ONE output element C[col]
-// 2. Threads cooperatively reduce over K dimension (parallel reduction)
-// 3. Coalesced memory access - adjacent threads read adjacent K elements
-// 4. Final reduction in shared memory
-//
-// A is f32 (activations), B is f16 (weights transposed [N,K]), C is f32.
-
-enable f16;
-
-override WORKGROUP_SIZE: u32 = 256u;
-const MAX_WORKGROUP_SIZE: u32 = 256u;
-
-struct Uniforms {
-    M: u32,           // Always 1 for GEMV
-    N: u32,           // Output dimension (# of output columns)
-    K: u32,           // Inner dimension (dot product length)
-    alpha: f32,       // Scaling factor
-    transpose_b: u32, // Expected to be 1 (B stored as [N,K])
-    _pad0: u32,
-    _pad1: u32,
-    _pad2: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> A: array<f32>;
-@group(0) @binding(2) var<storage, read> B: array<f16>;
-@group(0) @binding(3) var<storage, read_write> C: array<f32>;
-
-// Shared memory for parallel reduction
-var<workgroup> shared_sum: array<f32, MAX_WORKGROUP_SIZE>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(
-    @builtin(global_invocation_id) gid: vec3<u32>,
-    @builtin(local_invocation_id) lid: vec3<u32>,
-    @builtin(workgroup_id) wg_id: vec3<u32>
-) {
-    let col = wg_id.x;  // Output column this workgroup computes
-    let local_id = lid.x;
-
-    if (col >= u.N) {
-        return;
-    }
-
-    // Each thread computes partial sum for k = local_id, local_id+WORKGROUP_SIZE, ...
-    var partial_sum: f32 = 0.0;
-
-    // B is stored transposed [N, K], so B[col, k] = B[col * K + k]
-    let b_row_offset = col * u.K;
-
-    // Stride through K with workgroup-sized steps
-    var k: u32 = local_id;
-    for (; k < u.K; k = k + WORKGROUP_SIZE) {
-        let a_val = A[k];
-        let b_val = f32(B[b_row_offset + k]);
-        partial_sum = partial_sum + a_val * b_val;
-    }
-
-    // Store partial sum to shared memory
-    shared_sum[local_id] = partial_sum;
-    workgroupBarrier();
-
-    // Parallel reduction in shared memory
-    // Tree reduction: 256 -> 128 -> 64 -> 32 -> 16 -> 8 -> 4 -> 2 -> 1
-    for (var stride: u32 = WORKGROUP_SIZE / 2u; stride > 0u; stride = stride / 2u) {
-        if (local_id < stride) {
-            shared_sum[local_id] = shared_sum[local_id] + shared_sum[local_id + stride];
-        }
-        workgroupBarrier();
-    }
-
-    // Thread 0 writes the final result
-    if (local_id == 0u) {
-        C[col] = shared_sum[0] * u.alpha;
-    }
-}

package/src/gpu/kernels/matmul_gemv_f16a.wgsl

@@ -1,81 +0,0 @@
-// Optimized GEMV Kernel - Parallel Reduction over K
-// For M=1 decode: C[N] = A[K] * B^T[K,N]
-//
-// Key optimizations:
-// 1. Each workgroup computes ONE output element C[col]
-// 2. Threads cooperatively reduce over K dimension (parallel reduction)
-// 3. Coalesced memory access - adjacent threads read adjacent K elements
-// 4. Final reduction in shared memory
-//
-// A is f16 (activations), B is f16 (weights transposed [N,K]), C is f16.
-// Accumulation is performed in f32 for stability.
-
-enable f16;
-
-override WORKGROUP_SIZE: u32 = 256u;
-const MAX_WORKGROUP_SIZE: u32 = 256u;
-
-struct Uniforms {
-    M: u32,           // Always 1 for GEMV
-    N: u32,           // Output dimension (# of output columns)
-    K: u32,           // Inner dimension (dot product length)
-    alpha: f32,       // Scaling factor
-    transpose_b: u32, // Expected to be 1 (B stored as [N,K])
-    _pad0: u32,
-    _pad1: u32,
-    _pad2: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> A: array<f16>;
-@group(0) @binding(2) var<storage, read> B: array<f16>;
-@group(0) @binding(3) var<storage, read_write> C: array<f16>;
-
-// Shared memory for parallel reduction
-var<workgroup> shared_sum: array<f32, MAX_WORKGROUP_SIZE>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(
-    @builtin(global_invocation_id) gid: vec3<u32>,
-    @builtin(local_invocation_id) lid: vec3<u32>,
-    @builtin(workgroup_id) wg_id: vec3<u32>
-) {
-    let col = wg_id.x;  // Output column this workgroup computes
-    let local_id = lid.x;
-
-    if (col >= u.N) {
-        return;
-    }
-
-    // Each thread computes partial sum for k = local_id, local_id+WORKGROUP_SIZE, ...
-    var partial_sum: f32 = 0.0;
-
-    // B is stored transposed [N, K], so B[col, k] = B[col * K + k]
-    let b_row_offset = col * u.K;
-
-    // Stride through K with workgroup-sized steps
-    var k: u32 = local_id;
-    for (; k < u.K; k = k + WORKGROUP_SIZE) {
-        let a_val = f32(A[k]);
-        let b_val = f32(B[b_row_offset + k]);
-        partial_sum = partial_sum + a_val * b_val;
-    }
-
-    // Store partial sum to shared memory
-    shared_sum[local_id] = partial_sum;
-    workgroupBarrier();
-
-    // Parallel reduction in shared memory
-    // Tree reduction: 256 -> 128 -> 64 -> 32 -> 16 -> 8 -> 4 -> 2 -> 1
-    for (var stride: u32 = WORKGROUP_SIZE / 2u; stride > 0u; stride = stride / 2u) {
-        if (local_id < stride) {
-            shared_sum[local_id] = shared_sum[local_id] + shared_sum[local_id + stride];
-        }
-        workgroupBarrier();
-    }
-
-    // Thread 0 writes the final result
-    if (local_id == 0u) {
-        C[col] = f16(shared_sum[0] * u.alpha);
-    }
-}