@simulatte/doppler 0.1.0

package/src/gpu/kernels/fused_matmul_rmsnorm.js

@@ -0,0 +1,273 @@
+
+
+import { getDevice } from '../device.js';
+import { acquireBuffer, getBufferRequestedSize } from '../../memory/buffer-pool.js';
+import { createTensor } from '../tensor.js';
+import { getBuffer } from '../weight-buffer.js';
+import { dispatch, recordDispatch } from './dispatch.js';
+import { getPipelineFast, createUniformBufferWithView } from './utils.js';
+import { WORKGROUP_SIZES } from './constants.js';
+import { getKernelThresholds } from '../../config/schema/index.js';
+import { trace } from '../../debug/index.js';
+import { selectRuleValue } from './rule-registry.js';
+import { selectRuleValue as selectLoaderRule } from '../../rules/rule-registry.js';
+
+
+function resolveNormWeightDtype(byteSize, hiddenSize) {
+  if (!byteSize || hiddenSize == null) return null;
+  const f16Bytes = hiddenSize * 2;
+  const f32Bytes = hiddenSize * 4;
+  const sizeMatchesF32 = byteSize >= f32Bytes;
+  const sizeMatchesF16 = byteSize >= f16Bytes && !sizeMatchesF32;
+  if (!sizeMatchesF16 && !sizeMatchesF32) {
+    return null;
+  }
+  return selectLoaderRule('loader', 'weights', 'normWeightDtypeFromSize', {
+    sizeMatchesF16,
+    sizeMatchesF32,
+  });
+}
+
+
+export function selectMatmulRMSNormFusedVariant(N, dtype) {
+  if (!dtype) {
+    throw new Error('[MatmulRMSNormFused] dtype is required for variant selection.');
+  }
+  const isF16 = dtype === 'f16';
+  const isSmall = N <= WORKGROUP_SIZES.DEFAULT;
+  return selectRuleValue('fusedMatmulRmsnorm', 'variant', { isSmall, isF16 });
+}
+
+
+export async function runMatmulRMSNormFused(
+  input,
+  weight,
+  normWeight,
+  options
+) {
+  const device = getDevice();
+  const {
+    N,
+    K,
+    eps,
+    residual = null,
+    outputBuffer = null,
+    transposeB = true,  // Default: GGUF row-major weights
+    rmsNormWeightOffset = false,
+    label = null,
+  } = options;
+  if (eps == null) {
+    throw new Error('[MatmulRMSNormFused] eps is required.');
+  }
+
+  const { maxMediumN } = getKernelThresholds().fusedMatmul;
+  if (N > maxMediumN) {
+    throw new Error(`[MatmulRMSNormFused] N=${N} exceeds maxMediumN=${maxMediumN}; kernel only supports single-workgroup RMSNorm.`);
+  }
+
+  const weightBuffer = getBuffer(weight);
+  const normWeightBuffer = getBuffer(normWeight);
+  const normWeightSize = getBufferRequestedSize(normWeightBuffer);
+  const normWeightDtype = resolveNormWeightDtype(normWeightSize, N);
+  if (!normWeightDtype) {
+    throw new Error(
+      `[MatmulRMSNormFused] norm weight size (${normWeightSize} bytes) does not match ` +
+      `hiddenSize=${N} (expected ${N * 2} or ${N * 4} bytes).`
+    );
+  }
+
+  // Select variant based on output size and input dtype
+  if (!input.dtype) {
+    throw new Error('[MatmulRMSNormFused] input dtype is required.');
+  }
+  const dtype = input.dtype;
+  const variant = selectMatmulRMSNormFusedVariant(N, dtype);
+
+  trace.kernels(`MatmulRMSNormFused: N=${N}, K=${K}, variant=${variant}, dtype=${dtype}, hasResidual=${!!residual}, transposeB=${transposeB}, offset=${rmsNormWeightOffset}`);
+
+  const constants = { RMS_NORM_OFFSET: rmsNormWeightOffset, WEIGHT_IS_F16: normWeightDtype === 'f16' };
+  const pipeline = await getPipelineFast('fused_matmul_rmsnorm', variant, null, constants);
+
+  // Output buffer: [1, N] - size depends on dtype
+  const bytesPerElement = dtype === 'f16' ? 2 : 4;
+  const outputSize = N * bytesPerElement;
+  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'matmul_rmsnorm_fused_output');
+
+  // Create uniform buffer (8 u32/f32 = 32 bytes, padded for alignment)
+  const uniformBuffer = createUniformBufferWithView(
+    'matmul_rmsnorm_fused_uniforms',
+    32,
+    (view) => {
+      view.setUint32(0, N, true);
+      view.setUint32(4, K, true);
+      view.setFloat32(8, eps, true);
+      view.setUint32(12, residual ? 1 : 0, true);
+      view.setUint32(16, transposeB ? 1 : 0, true);
+      // Padding bytes 20-31 are zero-initialized
+    },
+    null,
+    device
+  );
+
+  // Create placeholder for residual if not provided
+  const residualBuffer = residual || device.createBuffer({
+    label: 'matmul_rmsnorm_residual_placeholder',
+    size: 4,
+    usage: GPUBufferUsage.STORAGE,
+  });
+
+  // Create bind group
+  const bindGroup = device.createBindGroup({
+    label: 'matmul_rmsnorm_fused_bind_group',
+    layout: pipeline.getBindGroupLayout(0),
+    entries: [
+      { binding: 0, resource: { buffer: uniformBuffer } },
+      { binding: 1, resource: { buffer: input.buffer } },
+      { binding: 2, resource: { buffer: weightBuffer } },
+      { binding: 3, resource: { buffer: normWeightBuffer } },
+      { binding: 4, resource: { buffer: output } },
+      { binding: 5, resource: { buffer: residualBuffer } },
+    ],
+  });
+
+  // Calculate workgroups
+  
+  const workgroups = 1;
+
+  const dispatchLabel = label ? `matmul_rmsnorm_fused:${label}` : 'matmul_rmsnorm_fused';
+  dispatch(device, pipeline, bindGroup, workgroups, dispatchLabel);
+
+  // Cleanup
+  uniformBuffer.destroy();
+  if (!residual) residualBuffer.destroy();
+
+  // Output dtype matches input dtype
+  return createTensor(output, input.dtype, [1, N], 'matmul_rmsnorm_fused_output');
+}
+
+
+export async function recordMatmulRMSNormFused(
+  recorder,
+  input,
+  weight,
+  normWeight,
+  options
+) {
+  const device = recorder.device;
+  const {
+    N,
+    K,
+    eps,
+    residual = null,
+    outputBuffer = null,
+    transposeB = true,  // Default: GGUF row-major weights
+    rmsNormWeightOffset = false,
+    label = null,
+  } = options;
+  if (eps == null) {
+    throw new Error('[MatmulRMSNormFused] eps is required.');
+  }
+
+  const { maxMediumN } = getKernelThresholds().fusedMatmul;
+  if (N > maxMediumN) {
+    throw new Error(`[MatmulRMSNormFused] N=${N} exceeds maxMediumN=${maxMediumN}; kernel only supports single-workgroup RMSNorm.`);
+  }
+
+  const weightBuffer = getBuffer(weight);
+  const normWeightBuffer = getBuffer(normWeight);
+  const normWeightSize = getBufferRequestedSize(normWeightBuffer);
+  const normWeightDtype = resolveNormWeightDtype(normWeightSize, N);
+  if (!normWeightDtype) {
+    throw new Error(
+      `[MatmulRMSNormFused] norm weight size (${normWeightSize} bytes) does not match ` +
+      `hiddenSize=${N} (expected ${N * 2} or ${N * 4} bytes).`
+    );
+  }
+
+  // Select variant based on dtype
+  if (!input.dtype) {
+    throw new Error('[MatmulRMSNormFused] input dtype is required.');
+  }
+  const dtype = input.dtype;
+  const variant = selectMatmulRMSNormFusedVariant(N, dtype);
+
+  trace.kernels(`recordMatmulRMSNormFused: N=${N}, K=${K}, variant=${variant}, dtype=${dtype}, hasResidual=${!!residual}, transposeB=${transposeB}, offset=${rmsNormWeightOffset}`);
+
+  const constants = { RMS_NORM_OFFSET: rmsNormWeightOffset, WEIGHT_IS_F16: normWeightDtype === 'f16' };
+  const pipeline = await getPipelineFast('fused_matmul_rmsnorm', variant, null, constants);
+
+  // Output buffer - size depends on dtype
+  const bytesPerElement = dtype === 'f16' ? 2 : 4;
+  const outputSize = N * bytesPerElement;
+  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'matmul_rmsnorm_fused_output');
+
+  // Uniform buffer via recorder (8 u32/f32 = 32 bytes, padded for alignment)
+  const uniformBuffer = createUniformBufferWithView(
+    'matmul_rmsnorm_fused_uniforms',
+    32,
+    (view) => {
+      view.setUint32(0, N, true);
+      view.setUint32(4, K, true);
+      view.setFloat32(8, eps, true);
+      view.setUint32(12, residual ? 1 : 0, true);
+      view.setUint32(16, transposeB ? 1 : 0, true);
+      // Padding bytes 20-31 are zero-initialized
+    },
+    recorder
+  );
+
+  // Placeholder for residual
+  const residualBuffer = residual || device.createBuffer({
+    label: 'matmul_rmsnorm_residual_placeholder',
+    size: 4,
+    usage: GPUBufferUsage.STORAGE,
+  });
+
+  // Bind group
+  const bindGroup = device.createBindGroup({
+    label: 'matmul_rmsnorm_fused_bind_group',
+    layout: pipeline.getBindGroupLayout(0),
+    entries: [
+      { binding: 0, resource: { buffer: uniformBuffer } },
+      { binding: 1, resource: { buffer: input.buffer } },
+      { binding: 2, resource: { buffer: weightBuffer } },
+      { binding: 3, resource: { buffer: normWeightBuffer } },
+      { binding: 4, resource: { buffer: output } },
+      { binding: 5, resource: { buffer: residualBuffer } },
+    ],
+  });
+
+  // Calculate workgroups
+  
+  const workgroups = 1;
+
+  const dispatchLabel = label ? `matmul_rmsnorm_fused:${label}` : 'matmul_rmsnorm_fused';
+  recordDispatch(recorder, pipeline, bindGroup, workgroups, dispatchLabel);
+
+  // Track placeholder for cleanup
+  if (!residual) {
+    recorder.trackTemporaryBuffer(residualBuffer);
+  }
+
+  // Output dtype matches input dtype
+  return createTensor(output, input.dtype, [1, N], 'matmul_rmsnorm_fused_output');
+}
+
+
+export function shouldUseFusedMatmulRMSNorm(M, N, K) {
+  // Only beneficial for decode (M=1)
+  if (M !== 1) {
+    return false;
+  }
+
+  const { maxMediumN, maxMediumK } = getKernelThresholds().fusedMatmul;
+  if (N > maxMediumN) {
+    return false;
+  }
+
+  if (typeof K === 'number' && K > maxMediumK) {
+    return false;
+  }
+
+  return true;
+}

package/src/gpu/kernels/fused_matmul_rmsnorm.wgsl

@@ -0,0 +1,324 @@
+// Fused GEMV + RMSNorm Kernel
+//
+// For decode (M=1), combines the down projection matmul with RMSNorm in a single kernel:
+// 1. Compute GEMV: C[1, N] = A[1, K] × B[K, N]  (down projection)
+// 2. Compute RMSNorm on C: output = C / sqrt(mean(C^2) + eps) * weight
+// 3. Optional residual: output = output + residual
+//
+// Benefits:
+// - Single GPU dispatch instead of 2
+// - No intermediate buffer for matmul output
+// - Better cache locality
+//
+// Expected speedup: 1.2-1.5x for post-FFN normalization path
+
+override WORKGROUP_SIZE: u32 = 256u;
+const MAX_WORKGROUP_SIZE: u32 = 256u;
+override COLS_PER_WG: u32 = 4u;
+override THREADS_PER_COL: u32 = 64u;  // WORKGROUP_SIZE / COLS_PER_WG
+override RMS_NORM_OFFSET: bool = false;
+override WEIGHT_IS_F16: bool = false;
+
+struct Uniforms {
+    N: u32,             // Output dimension (hidden_size)
+    K: u32,             // Input dimension (intermediate_size)
+    eps: f32,           // RMSNorm epsilon
+    has_residual: u32,  // 1 if residual provided, 0 otherwise
+    transpose_b: u32,   // 1 if weight is [N,K] (row-major), 0 if [K,N] (column-major)
+    _pad0: u32,         // Padding to 24 bytes for alignment
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> input: array<f32>;       // [1, K] - activation from FFN
+@group(0) @binding(2) var<storage, read> weight: array<f32>;      // [K, N] - down projection weight
+@group(0) @binding(3) var<storage, read> norm_weight: array<u32>; // [N] - RMSNorm weight (F16 or F32)
+@group(0) @binding(4) var<storage, read_write> output: array<f32>;// [1, N] - final output
+@group(0) @binding(5) var<storage, read> residual: array<f32>;    // [1, N] - optional residual
+
+// Shared memory for reduction
+var<workgroup> shared_partial: array<f32, MAX_WORKGROUP_SIZE>;  // Partial dot products
+var<workgroup> shared_output: array<f32, MAX_WORKGROUP_SIZE>;   // Output values for RMSNorm
+var<workgroup> shared_sum_sq: array<f32, MAX_WORKGROUP_SIZE>;   // Sum of squares for reduction
+
+fn apply_weight(w: f32) -> f32 {
+    if (RMS_NORM_OFFSET) {
+        return 1.0 + w;
+    }
+    return w;
+}
+
+fn load_norm_weight(idx: u32) -> f32 {
+    if (WEIGHT_IS_F16) {
+        let packed = norm_weight[idx >> 1u];
+        let pair = unpack2x16float(packed);
+        return select(pair.x, pair.y, (idx & 1u) == 1u);
+    }
+    return bitcast<f32>(norm_weight[idx]);
+}
+
+// Fused GEMV + RMSNorm: Main entry point
+// Each workgroup computes COLS_PER_WG output columns, then all workgroups
+// cooperate to compute the RMSNorm across the full output row.
+//
+// Phase 1: GEMV - compute dot products for each output column
+// Phase 2: Store outputs and compute partial sum of squares
+// Phase 3: Global reduction of sum of squares (via atomic or multi-pass)
+// Phase 4: Normalize outputs and optionally add residual
+//
+// Note: This kernel requires N <= WG_SIZE * COLS_PER_WG for single-pass RMSNorm.
+// For larger N, use the multi-workgroup variant.
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(
+    @builtin(global_invocation_id) global_id: vec3<u32>,
+    @builtin(local_invocation_id) local_id: vec3<u32>,
+    @builtin(workgroup_id) wg_id: vec3<u32>,
+    @builtin(num_workgroups) num_wgs: vec3<u32>
+) {
+    let tid = local_id.x;
+    let N = u.N;
+    let K = u.K;
+
+    // Each workgroup handles COLS_PER_WG columns of output
+    let col_in_wg = tid / THREADS_PER_COL;  // Which column within this workgroup (0-3)
+    let thread_in_col = tid % THREADS_PER_COL;  // Thread index within column (0-63)
+    let global_col = wg_id.x * COLS_PER_WG + col_in_wg;  // Global column index
+
+    // Phase 1: Compute dot product for this column
+    var dot_sum: f32 = 0.0;
+    if (global_col < N) {
+        // Each thread in a column processes a stripe of K
+        let k_per_thread = (K + THREADS_PER_COL - 1u) / THREADS_PER_COL;
+        let k_start = thread_in_col * k_per_thread;
+        let k_end = min(k_start + k_per_thread, K);
+
+        for (var k: u32 = k_start; k < k_end; k = k + 1u) {
+            // transpose_b=1: weight is [N,K], access weight[global_col * K + k]
+            // transpose_b=0: weight is [K,N], access weight[k * N + global_col]
+            let w_idx = select(k * N + global_col, global_col * K + k, u.transpose_b == 1u);
+            dot_sum = dot_sum + input[k] * weight[w_idx];
+        }
+    }
+
+    // Store partial dot product in shared memory
+    shared_partial[tid] = dot_sum;
+    workgroupBarrier();
+
+    // Reduce within each column (THREADS_PER_COL threads -> 1 value)
+    // Tree reduction: 64 -> 32 -> 16 -> 8 -> 4 -> 2 -> 1
+    for (var stride: u32 = THREADS_PER_COL / 2u; stride > 0u; stride = stride >> 1u) {
+        if (thread_in_col < stride) {
+            let idx = col_in_wg * THREADS_PER_COL + thread_in_col;
+            shared_partial[idx] = shared_partial[idx] + shared_partial[idx + stride];
+        }
+        workgroupBarrier();
+    }
+
+    // Thread 0 of each column has the final dot product
+    var matmul_output: f32 = 0.0;
+    if (thread_in_col == 0u && global_col < N) {
+        matmul_output = shared_partial[col_in_wg * THREADS_PER_COL];
+        shared_output[col_in_wg] = matmul_output;
+        shared_sum_sq[col_in_wg] = matmul_output * matmul_output;
+    }
+    workgroupBarrier();
+
+    // Phase 2: For single-workgroup case (N <= 1024), we can do RMSNorm directly
+    // For multi-workgroup, we'd need a second pass or atomic accumulation
+
+    // This version handles N <= WORKGROUP_SIZE * COLS_PER_WG
+    // For larger N, spawn multiple workgroups and do global reduction
+
+    // Reduce sum of squares across columns in this workgroup
+    // All threads must participate in barriers, so move barrier outside condition
+    for (var stride: u32 = COLS_PER_WG / 2u; stride > 0u; stride = stride >> 1u) {
+        if (tid < stride && tid + stride < COLS_PER_WG) {
+            shared_sum_sq[tid] = shared_sum_sq[tid] + shared_sum_sq[tid + stride];
+        }
+        workgroupBarrier();
+    }
+
+    // For single-workgroup: compute RMS and normalize
+    // Thread 0 has partial sum; need global sum across all workgroups
+    // For now, assume single workgroup (N <= COLS_PER_WG = 4)
+    // Single-workgroup kernel; selection code enforces N within limits.
+
+    let local_sum_sq = shared_sum_sq[0];
+
+    // Phase 3: Write normalized output
+    if (thread_in_col == 0u && global_col < N) {
+        let val = shared_output[col_in_wg];
+
+        // For multi-WG: would need global sum here
+        // For now: just use local (works for small N)
+        let mean_sq = local_sum_sq / f32(min(COLS_PER_WG, N));
+        let inv_rms = 1.0 / sqrt(mean_sq + u.eps);
+
+        var result = val * inv_rms * apply_weight(load_norm_weight(global_col));
+
+        if (u.has_residual == 1u) {
+            result = result + residual[global_col];
+        }
+
+        output[global_col] = result;
+    }
+}
+
+// Optimized single-workgroup variant for small N (hidden_size <= WORKGROUP_SIZE)
+// All output columns computed by one workgroup, RMSNorm in single pass
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn gemv_rmsnorm_small(
+    @builtin(local_invocation_id) local_id: vec3<u32>
+) {
+    let tid = local_id.x;
+    let N = u.N;
+    let K = u.K;
+
+    // Each thread computes one output column (for N <= WORKGROUP_SIZE)
+    var dot_sum: f32 = 0.0;
+    if (tid < N) {
+        for (var k: u32 = 0u; k < K; k = k + 1u) {
+            let w_idx = select(k * N + tid, tid * K + k, u.transpose_b == 1u);
+            dot_sum = dot_sum + input[k] * weight[w_idx];
+        }
+    }
+
+    // Store output value and compute sum of squares
+    shared_output[tid] = dot_sum;
+    shared_sum_sq[tid] = select(0.0, dot_sum * dot_sum, tid < N);
+    workgroupBarrier();
+
+    // Parallel reduction for sum of squares
+    for (var stride: u32 = WORKGROUP_SIZE / 2u; stride > 0u; stride = stride >> 1u) {
+        if (tid < stride) {
+            shared_sum_sq[tid] = shared_sum_sq[tid] + shared_sum_sq[tid + stride];
+        }
+        workgroupBarrier();
+    }
+
+    // Compute inverse RMS
+    let mean_sq = shared_sum_sq[0] / f32(N);
+    let inv_rms = 1.0 / sqrt(mean_sq + u.eps);
+
+    // Write normalized output
+    if (tid < N) {
+        var result = shared_output[tid] * inv_rms * apply_weight(load_norm_weight(tid));
+        if (u.has_residual == 1u) {
+            result = result + residual[tid];
+        }
+        output[tid] = result;
+    }
+}
+
+// Medium variant for N up to ~4096 (single workgroup, multiple elements per thread)
+// Handles Gemma 3's hidden_size=1152 and similar models
+// Each thread computes ceil(N/WORKGROUP_SIZE) output columns sequentially
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn gemv_rmsnorm_medium(
+    @builtin(local_invocation_id) local_id: vec3<u32>
+) {
+    let tid = local_id.x;
+    let N = u.N;
+    let K = u.K;
+
+    // Each thread handles multiple output columns
+    let cols_per_thread = (N + WORKGROUP_SIZE - 1u) / WORKGROUP_SIZE;
+
+    // Phase 1: Compute dot products for all columns assigned to this thread
+    // Store partial sum of squares as we go
+    var local_sum_sq: f32 = 0.0;
+
+    for (var c: u32 = 0u; c < cols_per_thread; c = c + 1u) {
+        let col = tid + c * WORKGROUP_SIZE;
+        if (col < N) {
+            // Compute dot product for this column
+            var dot_sum: f32 = 0.0;
+            for (var k: u32 = 0u; k < K; k = k + 1u) {
+                let w_idx = select(k * N + col, col * K + k, u.transpose_b == 1u);
+                dot_sum = dot_sum + input[k] * weight[w_idx];
+            }
+
+            // Store to output buffer (will normalize in-place later)
+            output[col] = dot_sum;
+            local_sum_sq = local_sum_sq + dot_sum * dot_sum;
+        }
+    }
+
+    // Phase 2: Parallel reduction for sum of squares
+    shared_sum_sq[tid] = local_sum_sq;
+    workgroupBarrier();
+
+    for (var stride: u32 = WORKGROUP_SIZE / 2u; stride > 0u; stride = stride >> 1u) {
+        if (tid < stride) {
+            shared_sum_sq[tid] = shared_sum_sq[tid] + shared_sum_sq[tid + stride];
+        }
+        workgroupBarrier();
+    }
+
+    // Compute inverse RMS (all threads read the same value)
+    let mean_sq = shared_sum_sq[0] / f32(N);
+    let inv_rms = 1.0 / sqrt(mean_sq + u.eps);
+
+    // Phase 3: Normalize outputs in-place
+    for (var c: u32 = 0u; c < cols_per_thread; c = c + 1u) {
+        let col = tid + c * WORKGROUP_SIZE;
+        if (col < N) {
+            var result = output[col] * inv_rms * apply_weight(load_norm_weight(col));
+            if (u.has_residual == 1u) {
+                result = result + residual[col];
+            }
+            output[col] = result;
+        }
+    }
+}
+
+// Multi-workgroup variant using two-phase approach
+// Phase 1: Each workgroup computes output columns and partial sum of squares
+// Phase 2: Separate kernel reduces partials and normalizes
+// This entry point is Phase 1 only
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn gemv_rmsnorm_phase1(
+    @builtin(global_invocation_id) global_id: vec3<u32>,
+    @builtin(local_invocation_id) local_id: vec3<u32>,
+    @builtin(workgroup_id) wg_id: vec3<u32>
+) {
+    let tid = local_id.x;
+    let N = u.N;
+    let K = u.K;
+
+    let col_in_wg = tid / THREADS_PER_COL;
+    let thread_in_col = tid % THREADS_PER_COL;
+    let global_col = wg_id.x * COLS_PER_WG + col_in_wg;
+
+    // Compute dot product
+    var dot_sum: f32 = 0.0;
+    if (global_col < N) {
+        let k_per_thread = (K + THREADS_PER_COL - 1u) / THREADS_PER_COL;
+        let k_start = thread_in_col * k_per_thread;
+        let k_end = min(k_start + k_per_thread, K);
+
+        for (var k: u32 = k_start; k < k_end; k = k + 1u) {
+            let w_idx = select(k * N + global_col, global_col * K + k, u.transpose_b == 1u);
+            dot_sum = dot_sum + input[k] * weight[w_idx];
+        }
+    }
+
+    shared_partial[tid] = dot_sum;
+    workgroupBarrier();
+
+    // Reduce within column
+    for (var stride: u32 = THREADS_PER_COL / 2u; stride > 0u; stride = stride >> 1u) {
+        if (thread_in_col < stride) {
+            let idx = col_in_wg * THREADS_PER_COL + thread_in_col;
+            shared_partial[idx] = shared_partial[idx] + shared_partial[idx + stride];
+        }
+        workgroupBarrier();
+    }
+
+    // Write matmul output (RMSNorm will be done in phase 2)
+    if (thread_in_col == 0u && global_col < N) {
+        output[global_col] = shared_partial[col_in_wg * THREADS_PER_COL];
+    }
+}