@simulatte/doppler 0.1.6 → 0.1.8

package/src/gpu/kernels/fused_matmul_rmsnorm.js

@@ -1,7 +1,7 @@
 
 
 import { getDevice } from '../device.js';
-import { acquireBuffer, getBufferRequestedSize } from '../../memory/buffer-pool.js';
+import { acquireBuffer, getBufferRequestedSize, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor } from '../tensor.js';
 import { getBuffer } from '../weight-buffer.js';
 import { dispatch, recordDispatch } from './dispatch.js';
@@ -91,7 +91,8 @@ export async function runMatmulRMSNormFused(
   // 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');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'matmul_rmsnorm_fused_output');
+  const output = outputBuffer || ownedOutput;
 
   // Create uniform buffer (8 u32/f32 = 32 bytes, padded for alignment)
   const uniformBuffer = createUniformBufferWithView(
@@ -110,36 +111,44 @@ export async function runMatmulRMSNormFused(
   );
 
   // Create placeholder for residual if not provided
+  const ownsResidualBuffer = !residual;
   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);
+  try {
+    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 } },
+      ],
+    });
+
+    const workgroups = 1;
+    const dispatchLabel = label ? `matmul_rmsnorm_fused:${label}` : 'matmul_rmsnorm_fused';
+    dispatch(device, pipeline, bindGroup, workgroups, dispatchLabel);
+  } catch (error) {
+    uniformBuffer.destroy();
+    if (ownsResidualBuffer) {
+      residualBuffer.destroy();
+    }
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 
   // Cleanup
   uniformBuffer.destroy();
-  if (!residual) residualBuffer.destroy();
+  if (ownsResidualBuffer) residualBuffer.destroy();
 
   // Output dtype matches input dtype
   return createTensor(output, input.dtype, [1, N], 'matmul_rmsnorm_fused_output');
@@ -199,7 +208,8 @@ export async function recordMatmulRMSNormFused(
   // 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');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'matmul_rmsnorm_fused_output');
+  const output = outputBuffer || ownedOutput;
 
   // Uniform buffer via recorder (8 u32/f32 = 32 bytes, padded for alignment)
   const uniformBuffer = createUniformBufferWithView(
@@ -217,35 +227,42 @@ export async function recordMatmulRMSNormFused(
   );
 
   // Placeholder for residual
+  const ownsResidualBuffer = !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);
+  try {
+    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 } },
+      ],
+    });
+
+    const workgroups = 1;
+    const dispatchLabel = label ? `matmul_rmsnorm_fused:${label}` : 'matmul_rmsnorm_fused';
+    recordDispatch(recorder, pipeline, bindGroup, workgroups, dispatchLabel);
+  } catch (error) {
+    if (ownsResidualBuffer) {
+      residualBuffer.destroy();
+    }
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 
   // Track placeholder for cleanup
-  if (!residual) {
+  if (ownsResidualBuffer) {
     recorder.trackTemporaryBuffer(residualBuffer);
   }
 

package/src/gpu/kernels/gather.js

@@ -1,5 +1,5 @@
 import { getKernelCapabilities } from '../device.js';
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { WORKGROUP_SIZES, VEC4_ELEMENTS_PER_WG } from './constants.js';
 import { unifiedKernelWrapper } from './utils.js';
 import { trace } from '../../debug/index.js';
@@ -26,7 +26,6 @@ async function _gather(
   options = {}
 ) {
   const {
-    useVec4 = true,
     outputBuffer = null,
     embeddingDtype,
     outputDtype,
@@ -43,9 +42,22 @@ async function _gather(
   if (outputDtype == null) {
     throw new Error('[Gather] outputDtype is required.');
   }
+  if (embeddingDtype === 'f16' && !caps.hasF16) {
+    throw new Error('[Gather] embeddingDtype=f16 requires shader-f16 support.');
+  }
+  if (outputDtype === 'f16' && !caps.hasF16) {
+    throw new Error('[Gather] outputDtype=f16 requires shader-f16 support.');
+  }
 
-  const useF16Input = embeddingDtype === 'f16' && caps.hasF16;
-  const useF16Output = outputDtype === 'f16' && caps.hasF16;
+  const requestedVec4 = options.useVec4;
+  const wantsVec4 = requestedVec4 ?? true;
+  if (requestedVec4 === true && hiddenSize % 4 !== 0) {
+    throw new Error('[Gather] useVec4=true requires hiddenSize to be divisible by 4.');
+  }
+
+  const useF16Input = embeddingDtype === 'f16';
+  const useF16Output = outputDtype === 'f16';
+  const useVec4 = wantsVec4 && hiddenSize % 4 === 0;
 
   trace.embed(
     `Gather: numTokens=${numTokens}, hiddenSize=${hiddenSize}, vocabSize=${vocabSize}, ` +
@@ -64,6 +76,7 @@ async function _gather(
   const paddedHiddenSize = padToQ4KBlock(hiddenSize);
   const outputSize = numTokens * paddedHiddenSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'gather_output');
+  const ownedOutput = outputBuffer ? null : output;
 
   const uniforms = {
     num_tokens: numTokens,
@@ -82,16 +95,22 @@ async function _gather(
       ? Math.ceil((numTokens * hiddenSize) / VEC4_ELEMENTS_PER_WG)
       : Math.ceil((numTokens * hiddenSize) / WORKGROUP_SIZES.DEFAULT));
 
-  await unifiedKernelWrapper(
-    'gather',
-    target,
-    variant,
-    [indices, embeddings, output],
-    uniforms,
-    workgroups
-  );
-
-  return createTensor(output, actualDtype, [numTokens, hiddenSize], 'gather_output');
+  try {
+    await unifiedKernelWrapper(
+      'gather',
+      target,
+      variant,
+      [indices, embeddings, output],
+      uniforms,
+      workgroups
+    );
+    return createTensor(output, actualDtype, [numTokens, hiddenSize], 'gather_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function runGather(
@@ -116,4 +135,3 @@ export async function recordGather(
 ) {
   return _gather(recorder, indices, embeddings, numTokens, hiddenSize, vocabSize, options);
 }
-

package/src/gpu/kernels/gelu.js

@@ -1,5 +1,5 @@
 
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor, dtypeBytes } from '../tensor.js';
 import { WORKGROUP_SIZES } from './constants.js';
 import { unifiedKernelWrapper } from './utils.js';
@@ -26,16 +26,24 @@ async function _gelu(target, input, options = {}) {
   const outputSize = inferredSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'gelu_output');
   const gateBuffer = gate ?? input;
-
-  await unifiedKernelWrapper(
-    'gelu', target, variant,
-    [input, output, gateBuffer],
-    { size: inferredSize, rowsplit_dim: 0 },
-    Math.ceil(inferredSize / WORKGROUP_SIZES.DEFAULT),
-    overrides
-  );
-
-  return createTensor(output, input.dtype, [inferredSize], 'gelu_output');
+  const ownedOutput = outputBuffer ? null : output;
+
+  try {
+    await unifiedKernelWrapper(
+      'gelu', target, variant,
+      [input, output, gateBuffer],
+      { size: inferredSize, rowsplit_dim: 0 },
+      Math.ceil(inferredSize / WORKGROUP_SIZES.DEFAULT),
+      overrides
+    );
+
+    return createTensor(output, input.dtype, [inferredSize], 'gelu_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function runGeLU(input, options = {}) {

package/src/gpu/kernels/grouped_pointwise_conv2d.js

@@ -55,33 +55,43 @@ async function _groupedPointwiseConv2D(target, input, weight, bias, options = {}
     device.queue.writeBuffer(biasBuffer, 0, new Uint8Array(paddedSize));
   }
 
-  await unifiedKernelWrapper(
-    'grouped_pointwise_conv2d',
-    target,
-    variant,
-    [input, weightBuffer, biasBuffer, output],
-    {
-      in_channels: inChannels,
-      out_channels: outChannels,
-      height,
-      width,
-      groups,
-      _pad0: 0,
-      _pad1: 0,
-      _pad2: 0,
-    },
-    [Math.ceil(spatial / WORKGROUP_SIZES.DEFAULT), outChannels, 1]
-  );
+  try {
+    await unifiedKernelWrapper(
+      'grouped_pointwise_conv2d',
+      target,
+      variant,
+      [input, weightBuffer, biasBuffer, output],
+      {
+        in_channels: inChannels,
+        out_channels: outChannels,
+        height,
+        width,
+        groups,
+        _pad0: 0,
+        _pad1: 0,
+        _pad2: 0,
+      },
+      [Math.ceil(spatial / WORKGROUP_SIZES.DEFAULT), outChannels, 1]
+    );
+
+    if (tempBias) {
+      if (recorder) {
+        recorder.trackTemporaryBuffer(tempBias);
+      } else {
+        releaseBuffer(tempBias);
+      }
+    }
 
-  if (tempBias) {
-    if (recorder) {
-      recorder.trackTemporaryBuffer(tempBias);
-    } else {
+    return createTensor(output, input.dtype, [outChannels, height, width], 'grouped_pointwise_conv2d_output');
+  } catch (error) {
+    if (tempBias) {
       releaseBuffer(tempBias);
     }
+    if (!outputBuffer) {
+      releaseBuffer(output);
+    }
+    throw error;
   }
-
-  return createTensor(output, input.dtype, [outChannels, height, width], 'grouped_pointwise_conv2d_output');
 }
 
 export async function runGroupedPointwiseConv2D(input, weight, bias, options = {}) {

package/src/gpu/kernels/groupnorm.js

@@ -17,6 +17,9 @@ function validateOptions(options) {
   if (!Number.isFinite(numGroups) || numGroups <= 0) {
     throw new Error('GroupNorm requires numGroups > 0.');
   }
+  if (channels % numGroups !== 0) {
+    throw new Error('GroupNorm requires channels to be divisible by numGroups.');
+  }
   if (!Number.isFinite(eps)) {
     throw new Error('GroupNorm requires eps.');
   }
@@ -44,34 +47,42 @@ async function _groupNorm(target, input, weight, bias, options = {}) {
 
   const statsSize = numGroups * 2 * 4;
   const statsBuffer = acquireBuffer(statsSize, undefined, 'groupnorm_stats');
-
-  await unifiedKernelWrapper(
-    'groupnorm_stats',
-    target,
-    statsVariant,
-    [input, statsBuffer],
-    uniforms,
-    numGroups
-  );
-
   const bytesPerElement = dtypeBytes(input.dtype);
   const outputSize = channels * height * width * bytesPerElement;
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'groupnorm_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'groupnorm_output');
+  const output = outputBuffer || ownedOutput;
 
-  const weightBuffer = getBuffer(weight);
-  const biasBuffer = getBuffer(bias);
+  try {
+    await unifiedKernelWrapper(
+      'groupnorm_stats',
+      target,
+      statsVariant,
+      [input, statsBuffer],
+      uniforms,
+      numGroups
+    );
 
-  const total = channels * height * width;
-  const workgroups = Math.ceil(total / WORKGROUP_SIZES.DEFAULT);
+    const weightBuffer = getBuffer(weight);
+    const biasBuffer = getBuffer(bias);
 
-  await unifiedKernelWrapper(
-    'groupnorm_apply',
-    target,
-    applyVariant,
-    [input, statsBuffer, weightBuffer, biasBuffer, output],
-    uniforms,
-    workgroups
-  );
+    const total = channels * height * width;
+    const workgroups = Math.ceil(total / WORKGROUP_SIZES.DEFAULT);
+
+    await unifiedKernelWrapper(
+      'groupnorm_apply',
+      target,
+      applyVariant,
+      [input, statsBuffer, weightBuffer, biasBuffer, output],
+      uniforms,
+      workgroups
+    );
+  } catch (error) {
+    releaseBuffer(statsBuffer);
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 
   if (recorder) {
     recorder.trackTemporaryBuffer(statsBuffer);

package/src/gpu/kernels/index.d.ts

@@ -326,6 +326,14 @@ export {
   type SplitQKVResult,
 } from './split_qkv.js';
 
+// Split Q and Gate (de-interleave attentionOutputGate q_proj output)
+export {
+  runSplitQG,
+  recordSplitQG,
+  type SplitQGOptions,
+  type SplitQGResult,
+} from './split_qg.js';
+
 // Transpose
 export {
   runTranspose,

package/src/gpu/kernels/index.js

@@ -268,6 +268,12 @@ export {
   recordSplitQKV,
 } from './split_qkv.js';
 
+// Split Q and Gate (de-interleave attentionOutputGate q_proj output)
+export {
+  runSplitQG,
+  recordSplitQG,
+} from './split_qg.js';
+
 // Transpose
 export {
   runTranspose,

package/src/gpu/kernels/kv-quantize.js

@@ -78,8 +78,11 @@ export async function runKVQuantize(
   });
 
   const workgroups = [numKVHeads, numTokens, 1];
-  dispatch(device, pipeline, bindGroup, workgroups, 'kv_quantize');
-  uniformBuffer.destroy();
+  try {
+    dispatch(device, pipeline, bindGroup, workgroups, 'kv_quantize');
+  } finally {
+    uniformBuffer.destroy();
+  }
 }
 
 

package/src/gpu/kernels/layernorm.js

@@ -1,6 +1,6 @@
 
 import { getKernelCapabilities } from '../device.js';
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor } from '../tensor.js';
 import { padToQ4KBlock } from '../../config/schema/index.js';
 import { selectRuleValue } from './rule-registry.js';
@@ -36,17 +36,25 @@ export async function runLayerNorm(
   const paddedHiddenSize = padToQ4KBlock(inferredHiddenSize);
   const outputSize = batchSize * paddedHiddenSize * bytesPerElement;
   const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'layernorm_output');
+  const ownedOutput = outputBuffer ? null : outputBuf;
 
-  await unifiedKernelWrapper(
-    'layernorm',
-    null,
-    variant,
-    [input, weight, bias, outputBuf],
-    { hidden_size: inferredHiddenSize, num_tokens: batchSize, eps },
-    batchSize
-  );
+  try {
+    await unifiedKernelWrapper(
+      'layernorm',
+      null,
+      variant,
+      [input, weight, bias, outputBuf],
+      { hidden_size: inferredHiddenSize, num_tokens: batchSize, eps },
+      batchSize
+    );
 
-  return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'layernorm_output');
+    return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'layernorm_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function recordLayerNorm(
@@ -66,15 +74,23 @@ export async function recordLayerNorm(
   const paddedHiddenSize = padToQ4KBlock(inferredHiddenSize);
   const outputSize = batchSize * paddedHiddenSize * bytesPerElement;
   const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'layernorm_output');
+  const ownedOutput = outputBuffer ? null : outputBuf;
 
-  await unifiedKernelWrapper(
-    'layernorm',
-    recorder,
-    variant,
-    [input, weight, bias, outputBuf],
-    { hidden_size: inferredHiddenSize, num_tokens: batchSize, eps },
-    batchSize
-  );
+  try {
+    await unifiedKernelWrapper(
+      'layernorm',
+      recorder,
+      variant,
+      [input, weight, bias, outputBuf],
+      { hidden_size: inferredHiddenSize, num_tokens: batchSize, eps },
+      batchSize
+    );
 
-  return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'layernorm_output');
+    return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'layernorm_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }

package/src/gpu/kernels/logit-merge.js

@@ -266,9 +266,11 @@ export class LogitMergeKernel {
     pass.end();
 
     this.#device.queue.submit([encoder.finish()]);
-
-    // Cleanup temporary buffer
-    paramsBuffer.destroy();
+    this.#device.queue.onSubmittedWorkDone()
+      .catch(() => {})
+      .finally(() => {
+        paramsBuffer.destroy();
+      });
 
     return mergedBuffer;
   }

package/src/gpu/kernels/matmul-selection.js

@@ -29,7 +29,13 @@ function selectQ4KFusedVariant(isM1, wantF16Output, aDtype) {
 }
 
 
-export function resolveMatmulPhase(M) {
+export function resolveMatmulPhase(M, phaseOverride = null) {
+  if (phaseOverride != null) {
+    if (phaseOverride !== 'decode' && phaseOverride !== 'prefill') {
+      throw new Error(`[Matmul] Invalid phase override "${phaseOverride}". Expected "decode" or "prefill".`);
+    }
+    return phaseOverride;
+  }
   return selectKernelRuleValue('matmul', 'phase', { isDecode: M === 1 });
 }
 
@@ -125,7 +131,9 @@ export function selectMatmulKernel(options = {}) {
   const { tiledPrefillMinRows } = getKernelThresholds().matmul;
 
   const inputsAreF16 = aDtype === 'f16' && bDtype === 'f16';
-  const weightsAreF16 = bDtype === 'f16' && aDtype !== 'f16';
+  // F16 weights needing F32a path: weights are F16 and either activation is already F32,
+  // or both inputs are F16 but output is F32 (activation will be cast to F32 by executeMatmul)
+  const weightsAreF16 = bDtype === 'f16' && (aDtype !== 'f16' || outputDtype !== 'f16');
   const useF16Matmul = outputDtype === 'f16' && preferF16 && inputsAreF16 && capabilities.hasF16;
   const useF16wF32a = preferF16 && weightsAreF16 && capabilities.hasF16;
   const useTiled = isPrefill
@@ -244,6 +252,30 @@ export function requiresF32Input(variant) {
   return !supportsF16Input(variant);
 }
 
+function resolveRequiredWeightDtype(config) {
+  const shaderFile = String(config?.shaderFile ?? config?.wgsl ?? '');
+  if (!shaderFile) {
+    return null;
+  }
+  if (shaderFile.startsWith('fused_matmul_q4')) {
+    return 'q4k';
+  }
+  if (
+    shaderFile === 'matmul_f16.wgsl'
+    || shaderFile === 'matmul_f16_tiled.wgsl'
+    || shaderFile === 'matmul_f16w_f32a.wgsl'
+    || shaderFile === 'matmul_f16w_f32a_tiled.wgsl'
+    || shaderFile === 'matmul_gemv_subgroup.wgsl'
+    || shaderFile === 'matmul_gemv_subgroup_f16a.wgsl'
+  ) {
+    return 'f16';
+  }
+  if (shaderFile === 'matmul_f32.wgsl') {
+    return 'f32';
+  }
+  return null;
+}
+
 
 function resolveMatmulOverride(
   variantOverride,
@@ -287,6 +319,16 @@ function resolveMatmulOverride(
     );
   }
 
+  const requiredWeightDtype = resolveRequiredWeightDtype(config);
+  const weightDtypeOk = !requiredWeightDtype
+    || bDtype === requiredWeightDtype
+    || (requiredWeightDtype === 'f16' && bDtype === 'q4k');
+  if (!weightDtypeOk) {
+    return failOrWarn(
+      `Matmul kernel "${variantOverride}" requires ${requiredWeightDtype} weights but B dtype is ${bDtype}.`
+    );
+  }
+
   if (supportsF16Input(override) && aDtype !== 'f16') {
     return failOrWarn(`Matmul kernel "${variantOverride}" requires f16 activations but A dtype is ${aDtype}.`);
   }
@@ -341,7 +383,7 @@ function selectGemvVariant(useF16Gemv, useF32Gemv, hasSubgroups, useVec4, N, mul
 export function selectMatmulVariantAndFlags(mode, M, N, K, aDtype, bDtype, transposeB, requestedOutputDtype, options) {
   const capabilities = getKernelCapabilities();
   const strict = getKernelPathStrict();
-  const phase = resolveMatmulPhase(M);
+  const phase = resolveMatmulPhase(M, options.phaseOverride ?? null);
   let pathVariant = getKernelPathMatmulVariant(options.role, phase, options.layerIdx, options.kernelPath);
   const hadPathVariant = Boolean(pathVariant);
 
@@ -426,7 +468,8 @@ export function selectMatmulVariantAndFlags(mode, M, N, K, aDtype, bDtype, trans
 
   const canGemv = M === 1 && effectiveBDtype === 'f16' && capabilities.hasF16;
   const useF16Gemv = canGemv && aDtype === 'f16' && wantF16Output;
-  const useF32Gemv = canGemv && aDtype === 'f32';
+  // F32 GEMV: activation is F32, or activation is F16 with F32 output (will be cast to F32)
+  const useF32Gemv = canGemv && (aDtype === 'f32' || (aDtype === 'f16' && !wantF16Output));
   const useGemv = useF16Gemv || useF32Gemv;
   const useVec4 = (K % 4 === 0);
   const { multicolThreshold } = getKernelThresholds().matmul;

package/src/gpu/kernels/matmul.d.ts

@@ -23,6 +23,8 @@ export interface MatmulOptions extends OutputBufferOptions, OutputDtypeOptions,
   layerIdx?: number;
   /** Explicit kernel path context for variant selection (avoids global path state). */
   kernelPath?: KernelPathSchema | null;
+  /** Optional explicit phase for kernel-path lookup when the runtime rewrites rows (for example prefill last-position logits). */
+  phaseOverride?: 'decode' | 'prefill' | null;
   /**
    * Whether B matrix is stored transposed.
    * - true: B is [N,K] (SafeTensors/row-major), needs transpose