@simulatte/doppler 0.1.0

package/src/gpu/kernels/matmul.js

@@ -0,0 +1,271 @@
+import { getDevice } from '../device.js';
+import { createTensor } from '../tensor.js';
+import { getBuffer, getLayout, getWeightDtype } from '../weight-buffer.js';
+import { log, trace, isTraceEnabled } from '../../debug/index.js';
+import { releaseBuffer } from '../../memory/buffer-pool.js';
+import { releaseUniformBuffer } from '../uniform-cache.js';
+import { castF16ToF32, recordCastF16ToF32 } from './cast.js';
+import {
+  resolveMatmulPhase,
+  resolveMatmulConstants,
+  getMatmulConfig,
+  isFusedQ4KDisabled,
+  toMatmulDtype,
+  resolveTransposeB,
+  validateMatmulDimensions,
+  validateMatmulOffsets,
+  getMatmulBindingSizes,
+  requiresF32Input,
+  selectMatmulVariantAndFlags,
+  resolveMatmulOutput,
+  selectMatmulKernel,
+} from './matmul-selection.js';
+import {
+  MatmulKernel,
+  calculateMatmulDispatch,
+  createMatmulUniformBuffer,
+  createMatmulBindGroupLayout,
+  getMatmulPipeline,
+} from './matmul-dispatch.js';
+
+export { isFusedQ4KDisabled, selectMatmulKernel };
+export { createMatmulBindGroupLayout };
+
+let _runMatmulDebugCount = 0;
+let _recordMatmulDebugCount = 0;
+
+function getDebugCounter(isRecord) {
+  return isRecord ? _recordMatmulDebugCount : _runMatmulDebugCount;
+}
+
+function incrementDebugCounter(isRecord) {
+  if (isRecord) {
+    _recordMatmulDebugCount += 1;
+    return;
+  }
+  _runMatmulDebugCount += 1;
+}
+
+function buildProfileLabel(options = {}) {
+  const layerLabel = Number.isFinite(options.layerIdx) ? `:L${options.layerIdx}` : '';
+  const roleLabel = options.role ? `:${options.role}` : '';
+  return `matmul${roleLabel}${layerLabel}`;
+}
+
+function createMatmulBindGroupEntries(variant, uniformBuffer, matmulInput, bBuffer, outputBuffer, offsets, bindingSizes) {
+  const isQ4KF16 = variant === 'q4_fused_multicol_f16'
+    || variant === 'q4_fused_f16a'
+    || variant === 'q4_fused_batched_f16'
+    || variant === 'q4_fused_multicol_f16a'
+    || variant === 'q4_fused_batched_f16a';
+
+  const entries = [
+    { binding: 0, resource: { buffer: uniformBuffer } },
+    { binding: 1, resource: { buffer: matmulInput.buffer, offset: offsets.aOffset, size: bindingSizes.aBindingSize } },
+    { binding: 2, resource: { buffer: bBuffer, offset: offsets.bOffset, size: bindingSizes.bBindingSize } },
+  ];
+
+  if (isQ4KF16) {
+    entries.push({
+      binding: 4,
+      resource: { buffer: outputBuffer, offset: offsets.cOffset, size: bindingSizes.cBindingSize },
+    });
+  } else {
+    entries.push({
+      binding: 3,
+      resource: { buffer: outputBuffer, offset: offsets.cOffset, size: bindingSizes.cBindingSize },
+    });
+  }
+
+  return entries;
+}
+
+async function executeMatmul(recorder, A, B, M, N, K, options = {}) {
+  const isRecord = Boolean(recorder);
+  const mode = isRecord ? 'record' : 'run';
+  const opLabel = isRecord ? 'recordMatmul' : 'runMatmul';
+  const device = recorder?.device || getDevice();
+
+  const {
+    alpha = 1.0,
+    outputBuffer = null,
+    transposeB: transposeBOption = true,
+    aOffset = 0,
+    bOffset = 0,
+    cOffset = 0,
+  } = options;
+
+  const bBuffer = getBuffer(B);
+  const weightDtype = getWeightDtype(B);
+  const weightLabel = (B && typeof B === 'object' ? B.label : null) ?? bBuffer?.label ?? null;
+  const weightLayout = getLayout(B);
+  const weightShape = B?.shape ? `[${B.shape.join(', ')}]` : null;
+
+  if (isTraceEnabled('kernels') && getDebugCounter(isRecord) < 20) {
+    incrementDebugCounter(isRecord);
+    const modeLabel = isRecord ? 'recordMatmul' : 'runMatmul';
+    trace.kernels(`${modeLabel}: M=${M}, N=${N}, K=${K}, transposeBOption=${transposeBOption}, weightLayout=${weightLayout}, weightDtype=${weightDtype}`);
+  }
+
+  const transposeB = resolveTransposeB(B, transposeBOption);
+  validateMatmulDimensions(opLabel, M, N, K);
+
+  const aDtype = toMatmulDtype(A.dtype);
+  const bDtype = toMatmulDtype(weightDtype ?? options.bDtype);
+  const requestedOutputDtype = options.outputDtype || A.dtype;
+
+  if (!isRecord && isTraceEnabled('kernels') && !weightDtype && !options.bDtype && M <= 2) {
+    log.warn('Matmul', `runMatmul: B buffer dtype unknown! size=${bBuffer.size}, M=${M}, N=${N}, K=${K}. Assuming f32.`);
+  }
+
+  validateMatmulOffsets(opLabel, aOffset, bOffset, cOffset);
+
+  const { variant, useQ4KFused, useGemv } = selectMatmulVariantAndFlags(
+    mode,
+    M,
+    N,
+    K,
+    aDtype,
+    bDtype,
+    transposeB,
+    requestedOutputDtype,
+    options
+  );
+
+  const phase = resolveMatmulPhase(M);
+  const constants = resolveMatmulConstants(options, phase);
+
+  let matmulInput = A;
+  let matmulADtype = aDtype;
+  let castedInput = null;
+  if (matmulADtype === 'f16' && requiresF32Input(variant)) {
+    if (isTraceEnabled('kernels')) {
+      trace.kernels(`Matmul: casting f16 activations to f32 for variant=${variant}`);
+    }
+    if (isRecord) {
+      castedInput = await recordCastF16ToF32(recorder, A);
+      recorder.trackTemporaryBuffer(castedInput.buffer);
+    } else {
+      castedInput = await castF16ToF32(A);
+    }
+    matmulInput = castedInput;
+    matmulADtype = 'f32';
+  }
+
+  let bindingSizes;
+  try {
+    bindingSizes = getMatmulBindingSizes(
+      opLabel,
+      matmulInput.buffer,
+      bBuffer,
+      M,
+      N,
+      K,
+      matmulADtype,
+      bDtype,
+      transposeB,
+      aOffset,
+      bOffset
+    );
+  } catch (err) {
+    if (!isRecord && err instanceof Error && err.message.includes('B buffer too small')) {
+      const detailParts = [];
+      if (weightLabel) detailParts.push(`label=${weightLabel}`);
+      if (weightDtype) detailParts.push(`weightDtype=${weightDtype}`);
+      if (weightLayout) detailParts.push(`layout=${weightLayout}`);
+      if (weightShape) detailParts.push(`shape=${weightShape}`);
+      if (Number.isFinite(bBuffer?.size)) detailParts.push(`bSize=${bBuffer.size}`);
+      const detail = detailParts.length ? ` (${detailParts.join(', ')})` : '';
+      throw new Error(`${err.message}${detail}`);
+    }
+    throw err;
+  }
+
+  if (!isRecord && isTraceEnabled('kernels') && bDtype === 'q4k') {
+    if (useQ4KFused) {
+      trace.kernels(`Q4K FUSED: M=${M}, N=${N}, K=${K}, variant=${variant} (WARNING: 2.3x slower than dequant)`);
+    } else {
+      trace.kernels(`Q4K DEQUANT: M=${M}, N=${N}, K=${K}, will dequant first then matmul with variant=${variant}`);
+    }
+  }
+
+  if (!isRecord && isTraceEnabled('kernels') && N > 100000) {
+    trace.kernels(`MATMUL_LARGE: N=${N}, variant=${variant}, aDtype=${aDtype}, bDtype=${bDtype}, transposeB=${transposeB}`);
+  }
+
+  const config = getMatmulConfig(variant, constants);
+  const kernel = new MatmulKernel(device);
+  const pipeline = await getMatmulPipeline(variant, constants);
+
+  const { output: C, outputSize, cBindingSize, actualOutputDtype } = resolveMatmulOutput(
+    variant,
+    M,
+    N,
+    outputBuffer
+  );
+
+  if (!Number.isFinite(outputSize) || outputSize <= 0) {
+    throw new Error(`[${opLabel}] Invalid output size: ${outputSize} (M=${M}, N=${N})`);
+  }
+
+  const cRequired = cOffset + cBindingSize;
+  if (C.size < cRequired) {
+    throw new Error(`[${opLabel}] Output buffer too small: ${C.size} < ${cRequired} (M=${M}, N=${N})`);
+  }
+
+  const dispatchPlan = calculateMatmulDispatch(variant, useQ4KFused, useGemv, M, N, config);
+  const uniformBuffer = createMatmulUniformBuffer(
+    'matmul_uniforms',
+    M,
+    N,
+    K,
+    alpha,
+    useQ4KFused,
+    transposeB,
+    dispatchPlan.uniformWorkgroupsX,
+    recorder || null,
+    device
+  );
+
+  const entries = createMatmulBindGroupEntries(
+    variant,
+    uniformBuffer,
+    matmulInput,
+    bBuffer,
+    C,
+    { aOffset, bOffset, cOffset },
+    {
+      aBindingSize: bindingSizes.aBindingSize,
+      bBindingSize: bindingSizes.bBindingSize,
+      cBindingSize,
+    }
+  );
+
+  const bindGroup = device.createBindGroup({
+    label: 'matmul_bind_group',
+    layout: pipeline.getBindGroupLayout(0),
+    entries,
+  });
+
+  if (isRecord) {
+    kernel.record(recorder, pipeline, bindGroup, dispatchPlan.workgroups, buildProfileLabel(options));
+  } else {
+    kernel.dispatch(pipeline, bindGroup, dispatchPlan.workgroups);
+    releaseUniformBuffer(uniformBuffer);
+    if (castedInput) {
+      releaseBuffer(castedInput.buffer);
+    }
+  }
+
+  return createTensor(C, actualOutputDtype, [M, N], 'matmul_output');
+}
+
+
+export async function runMatmul(A, B, M, N, K, options = {}) {
+  return executeMatmul(null, A, B, M, N, K, options);
+}
+
+
+export async function recordMatmul(recorder, A, B, M, N, K, options = {}) {
+  return executeMatmul(recorder, A, B, M, N, K, options);
+}

package/src/gpu/kernels/matmul_f16.wgsl

@@ -0,0 +1,170 @@
+// Matrix Multiplication Kernel - FP16 (Half Precision)
+//
+// Tiled matrix multiplication using FP16 for improved throughput.
+// Requires 'shader-f16' feature enabled.
+// C[M,N] = A[M,K] * B[K,N]  (or B^T when transpose_b=1)
+//
+// Uses FP16 for storage and computation, with optional FP32 accumulation
+// for better numerical stability.
+
+enable f16;
+
+// Tile dimensions - can use larger tiles with f16 due to smaller footprint
+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
+    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<f16>;
+@group(0) @binding(2) var<storage, read> B: array<f16>;
+@group(0) @binding(3) var<storage, read_write> C: array<f16>;
+
+// Shared memory tiles - f16 allows 2x data in same space
+var<workgroup> tileA: array<f16, TILE_AREA>;
+var<workgroup> tileB: array<f16, TILE_AREA>;
+var<workgroup> tileB_vec4: array<vec4<f16>, 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;
+
+    // Use f32 accumulator for numerical stability during summation
+    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;
+
+        // Load tile from A
+        if (row < u.M && a_col < u.K) {
+            tileA[tile_idx] = A[row * u.K + a_col];
+        } else {
+            tileA[tile_idx] = f16(0.0);
+        }
+
+        // Load tile from B (handle transpose)
+        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] = f16(0.0);
+        }
+
+        workgroupBarrier();
+
+        // Compute partial dot product with f32 accumulation
+        for (var k: u32 = 0u; k < TILE_SIZE; k = k + 1u) {
+            let a_val = f32(tileA[local_row * TILE_SIZE + k]);
+            let b_val = f32(tileB[k * TILE_SIZE + local_col]);
+            sum = sum + a_val * b_val;
+        }
+
+        workgroupBarrier();
+    }
+
+    // Write result back as f16
+    if (row < u.M && col < u.N) {
+        C[row * u.N + col] = f16(sum * u.alpha);
+    }
+}
+
+// Alternative entry point for vec4 column groups (best when N is a multiple of 4)
+@compute @workgroup_size(TILE_SIZE, TILE_SIZE, 1)
+fn main_vec4(
+    @builtin(global_invocation_id) global_id: vec3<u32>,
+    @builtin(local_invocation_id) local_id: vec3<u32>
+) {
+    if (TILE_SIZE > MAX_TILE_SIZE) {
+        return;
+    }
+
+    let row = global_id.x;
+    let col_base = global_id.y * 4u;  // Each thread handles 4 columns
+    let local_row = local_id.x;
+    let local_col = local_id.y;
+
+    var sum: vec4<f32> = vec4<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;
+
+        // Load A tile (single element per thread)
+        if (row < u.M && a_col < u.K) {
+            tileA[tile_idx] = A[row * u.K + a_col];
+        } else {
+            tileA[tile_idx] = f16(0.0);
+        }
+
+        // Load B tile (handle transpose)
+        var b_vec: vec4<f16> = vec4<f16>(f16(0.0));
+        if (b_row < u.K) {
+            if (u.transpose_b == 0u) {
+                if (col_base + 0u < u.N) { b_vec.x = B[b_row * u.N + col_base]; }
+                if (col_base + 1u < u.N) { b_vec.y = B[b_row * u.N + col_base + 1u]; }
+                if (col_base + 2u < u.N) { b_vec.z = B[b_row * u.N + col_base + 2u]; }
+                if (col_base + 3u < u.N) { b_vec.w = B[b_row * u.N + col_base + 3u]; }
+            } else {
+                if (col_base + 0u < u.N) { b_vec.x = B[(col_base + 0u) * u.K + b_row]; }
+                if (col_base + 1u < u.N) { b_vec.y = B[(col_base + 1u) * u.K + b_row]; }
+                if (col_base + 2u < u.N) { b_vec.z = B[(col_base + 2u) * u.K + b_row]; }
+                if (col_base + 3u < u.N) { b_vec.w = B[(col_base + 3u) * u.K + b_row]; }
+            }
+        }
+        tileB_vec4[tile_idx] = b_vec;
+
+        workgroupBarrier();
+
+        for (var k: u32 = 0u; k < TILE_SIZE; k = k + 1u) {
+            let a_val = f32(tileA[local_row * TILE_SIZE + k]);
+            let b_val = tileB_vec4[k * TILE_SIZE + local_col];
+            sum = sum + a_val * vec4<f32>(b_val);
+        }
+
+        workgroupBarrier();
+    }
+
+    // Write results
+    if (row < u.M && col_base < u.N) {
+        let base = row * u.N + col_base;
+        let scaled = sum * u.alpha;
+        if (col_base + 0u < u.N) { C[base] = f16(scaled.x); }
+        if (col_base + 1u < u.N) { C[base + 1u] = f16(scaled.y); }
+        if (col_base + 2u < u.N) { C[base + 2u] = f16(scaled.z); }
+        if (col_base + 3u < u.N) { C[base + 3u] = f16(scaled.w); }
+    }
+}

package/src/gpu/kernels/matmul_f16_tiled.wgsl

@@ -0,0 +1,165 @@
+// 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

@@ -0,0 +1,89 @@
+// 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;
+    }
+}