@simulatte/doppler 0.1.6 → 0.1.8

package/src/inference/pipelines/text/attention/projections.d.ts

@@ -46,7 +46,16 @@ export function recordAttentionInputs(
   info: AttentionInputInfo | null | undefined
 ): void;
 
-export function resolveAttentionProjectionOutputDtype(attentionInputDtype: string): 'f16' | 'f32' | string;
+export function shouldForceF32AttentionProjectionForRoPE(options: {
+  attentionInputDtype: string;
+  headDim: number;
+  rotaryDim?: number;
+  interleaved?: boolean;
+}): boolean;
+export function resolveAttentionProjectionOutputDtype(
+  attentionInputDtype: string,
+  options?: { forceF32?: boolean }
+): 'f16' | 'f32' | string;
 export function resolveProjectionSliceOffsetBytes(
   weightBuffer: WeightBuffer | Tensor | GPUBuffer | null | undefined,
   outputRows: number,

package/src/inference/pipelines/text/attention/projections.js

@@ -1,10 +1,12 @@
-import { acquireBuffer } from '../../../../memory/buffer-pool.js';
+import { releaseBuffer } from '../../../../memory/buffer-pool.js';
 import { isWeightBuffer, getLayout, getWeightDtype } from '../../../../gpu/weight-buffer.js';
 import {
   runMatmul,
   recordMatmul,
   runSplitQKV,
   recordSplitQKV,
+  runSplitQG,
+  recordSplitQG,
   runRMSNorm,
   recordRMSNorm,
 } from '../../../../gpu/kernel-selector.js';
@@ -28,6 +30,13 @@ function getSplitRunner(recorder) {
   return (qkvTensor, options) => recordSplitQKV(recorder, qkvTensor, options);
 }
 
+function getSplitQGRunner(recorder) {
+  if (!recorder) {
+    return (qgTensor, options) => runSplitQG(qgTensor, options);
+  }
+  return (qgTensor, options) => recordSplitQG(recorder, qgTensor, options);
+}
+
 function getRmsNormRunner(recorder) {
   if (!recorder) {
     return (input, weight, eps, options) => runRMSNorm(input, weight, eps, options);
@@ -36,7 +45,7 @@ function getRmsNormRunner(recorder) {
 }
 
 function releaseOwnedWeightBuffer(layerWeight, resolvedWeightBuffer, releaseTemporary) {
-  if (layerWeight instanceof GPUBuffer || isWeightBuffer(layerWeight)) {
+  if ((typeof GPUBuffer !== 'undefined' && layerWeight instanceof GPUBuffer) || isWeightBuffer(layerWeight)) {
     return;
   }
   if (!resolvedWeightBuffer) {
@@ -66,10 +75,16 @@ async function projectSingleQkvTensor({
 }) {
   const runMatmulForMode = getMatmulRunner(recorder);
   const layerWeight = layerWeights?.[weightKey];
-  let projected;
+  if (!layerWeight) {
+    throw new Error(`Attention projection requires ${weightKey}.`);
+  }
+  if (!getWeightBuffer) {
+    throw new Error(`Attention projection requires getWeightBuffer for ${role}.`);
+  }
 
-  if (layerWeight && getWeightBuffer) {
-    const projBuffer = getWeightBuffer(layerWeight, role);
+  let projected;
+  const projBuffer = getWeightBuffer(layerWeight, role);
+  try {
     projected = await runMatmulForMode(normed, projBuffer, numTokens, outputSize, hiddenSize, {
       transposeB: 'auto',
       role,
@@ -77,26 +92,31 @@ async function projectSingleQkvTensor({
       kernelPath,
       outputDtype: matmulOutputDtype,
     });
+  } finally {
     releaseOwnedWeightBuffer(layerWeight, projBuffer, releaseTemporary);
-  } else {
-    const fallback = acquireBuffer(numTokens * outputSize * 4, undefined, outputLabel);
-    projected = createTensor(fallback, normed.dtype, [numTokens, outputSize], outputLabel);
   }
 
   const loraModule = getLoRAModule(lora, layerIdx, loraKey);
   if (loraModule && getWeightBuffer) {
-    const combined = await applyLoRA(
-      normed,
-      projected,
-      loraModule,
-      { M: numTokens, N: outputSize, K: hiddenSize },
-      getWeightBuffer,
-      recorder ?? undefined,
-      { kernelPath }
-    );
-    if (combined.buffer !== projected.buffer) {
-      releaseTemporary(projected.buffer);
-      projected = combined;
+    try {
+      const combined = await applyLoRA(
+        normed,
+        projected,
+        loraModule,
+        { M: numTokens, N: outputSize, K: hiddenSize },
+        getWeightBuffer,
+        recorder ?? undefined,
+        { kernelPath }
+      );
+      if (combined.buffer !== projected.buffer) {
+        releaseTemporary(projected.buffer);
+        projected = combined;
+      }
+    } catch (error) {
+      if (projected?.buffer) {
+        releaseTemporary(projected.buffer);
+      }
+      throw error;
     }
   }
 
@@ -190,13 +210,17 @@ async function projectQueryWithOptionalGate({
     return { qTensor, qGateTensor: null };
   }
 
+  // q_proj weights are stored with interleaved head layout: for head h,
+  // rows [h*headDim*2 : h*headDim*2+headDim] = Q, rows [h*headDim*2+headDim : (h+1)*headDim*2] = gate.
+  // Compute the full 2*qSize matmul, then de-interleave into separate Q and gate tensors.
   const runMatmulForMode = getMatmulRunner(recorder);
+  const runSplitQGForMode = getSplitQGRunner(recorder);
   const qWeightBuffer = getWeightBuffer(qWeight, 'q_proj');
-  const gateOffset = resolveProjectionSliceOffsetBytes(qWeightBuffer, qSize, hiddenSize);
+  let fullQGTensor = null;
   let qTensor = null;
   let qGateTensor = null;
   try {
-    qTensor = await runMatmulForMode(normed, qWeightBuffer, numTokens, qSize, hiddenSize, {
+    fullQGTensor = await runMatmulForMode(normed, qWeightBuffer, numTokens, qSize * 2, hiddenSize, {
       transposeB: 'auto',
       role: 'q_proj',
       layerIdx,
@@ -204,32 +228,54 @@ async function projectQueryWithOptionalGate({
       outputDtype: matmulOutputDtype,
     });
 
-    qGateTensor = await runMatmulForMode(normed, qWeightBuffer, numTokens, qSize, hiddenSize, {
-      transposeB: 'auto',
-      role: 'q_proj_gate',
-      layerIdx,
-      kernelPath,
-      bOffset: gateOffset,
-      outputDtype: matmulOutputDtype,
+    const split = await runSplitQGForMode(fullQGTensor, {
+      numTokens,
+      numHeads,
+      headDim,
     });
+    releaseTemporary(fullQGTensor.buffer);
+    fullQGTensor = null;
+    qTensor = split.Q;
+    qGateTensor = split.G;
+  } catch (error) {
+    if (fullQGTensor) {
+      releaseTemporary(fullQGTensor.buffer);
+    }
+    if (qTensor) {
+      releaseTemporary(qTensor.buffer);
+    }
+    if (qGateTensor) {
+      releaseTemporary(qGateTensor.buffer);
+    }
+    throw error;
   } finally {
     releaseOwnedWeightBuffer(qWeight, qWeightBuffer, releaseTemporary);
   }
 
   const loraModule = getLoRAModule(lora, layerIdx, 'q_proj');
   if (loraModule && getWeightBuffer) {
-    const combined = await applyLoRA(
-      normed,
-      qTensor,
-      loraModule,
-      { M: numTokens, N: qSize, K: hiddenSize },
-      getWeightBuffer,
-      recorder ?? undefined,
-      { kernelPath }
-    );
-    if (combined.buffer !== qTensor.buffer) {
-      releaseTemporary(qTensor.buffer);
-      qTensor = combined;
+    try {
+      const combined = await applyLoRA(
+        normed,
+        qTensor,
+        loraModule,
+        { M: numTokens, N: qSize, K: hiddenSize },
+        getWeightBuffer,
+        recorder ?? undefined,
+        { kernelPath }
+      );
+      if (combined.buffer !== qTensor.buffer) {
+        releaseTemporary(qTensor.buffer);
+        qTensor = combined;
+      }
+    } catch (error) {
+      if (qTensor?.buffer) {
+        releaseTemporary(qTensor.buffer);
+      }
+      if (qGateTensor?.buffer) {
+        releaseTemporary(qGateTensor.buffer);
+      }
+      throw error;
     }
   }
 
@@ -248,9 +294,22 @@ export function recordAttentionInputs(state, info) {
   state.stats.attentionInputs.push(info);
 }
 
-export function resolveAttentionProjectionOutputDtype(attentionInputDtype) {
+export function shouldForceF32AttentionProjectionForRoPE({
+  attentionInputDtype,
+  headDim,
+  rotaryDim = headDim,
+  interleaved = false,
+}) {
+  return attentionInputDtype === 'f16'
+    && Number.isFinite(headDim)
+    && Number.isFinite(rotaryDim)
+    && (rotaryDim !== headDim || interleaved === true);
+}
+
+export function resolveAttentionProjectionOutputDtype(attentionInputDtype, options = {}) {
   const useF16Activations = attentionInputDtype === 'f16';
-  return selectRuleValue('shared', 'dtype', 'f16OrFallbackByFlag', {
+  return selectRuleValue('inference', 'dtype', 'attentionProjectionOutputDtype', {
+    forceF32: options.forceF32 === true,
     useF16: useF16Activations,
     fallback: attentionInputDtype,
   });
@@ -289,82 +348,103 @@ export async function projectAttentionQKV({
   if (useFusedQKV && layerWeights.qkvProj && layerWeights.qkvSizes) {
     const [qSizeFused, kSizeFused, vSizeFused] = layerWeights.qkvSizes;
     const qkvSizeTotal = qSizeFused + kSizeFused + vSizeFused;
-    const qkvTensor = await runMatmulForMode(normed, layerWeights.qkvProj, numTokens, qkvSizeTotal, hiddenSize, {
-      transposeB: 'auto',
-      role: 'qkv_proj',
+    let qkvTensor = null;
+    try {
+      qkvTensor = await runMatmulForMode(normed, layerWeights.qkvProj, numTokens, qkvSizeTotal, hiddenSize, {
+        transposeB: 'auto',
+        role: 'qkv_proj',
+        layerIdx,
+        kernelPath,
+        outputDtype: matmulOutputDtype,
+      });
+      const split = await runSplitForMode(qkvTensor, {
+        numTokens,
+        qSize: qSizeFused,
+        kSize: kSizeFused,
+        vSize: vSizeFused,
+      });
+      releaseTemporary(qkvTensor.buffer);
+      if (onFusedQKV) {
+        onFusedQKV({ qSize: qSizeFused, kSize: kSizeFused, vSize: vSizeFused, totalSize: qkvSizeTotal });
+      }
+      return { qTensor: split.Q, qGateTensor: null, kTensor: split.K, vTensor: split.V, usedFusedQKV: true };
+    } catch (error) {
+      if (qkvTensor) {
+        releaseTemporary(qkvTensor.buffer);
+      }
+      throw error;
+    }
+  }
+
+  let qTensor = null;
+  let qGateTensor = null;
+  let kTensor = null;
+  let vTensor = null;
+  try {
+    ({ qTensor, qGateTensor } = await projectQueryWithOptionalGate({
+      recorder,
+      normed,
+      layerWeights,
+      numTokens,
+      numHeads,
+      headDim,
+      hiddenSize,
       layerIdx,
       kernelPath,
-      outputDtype: matmulOutputDtype,
+      matmulOutputDtype,
+      getWeightBuffer,
+      lora,
+      releaseTemporary,
+      attentionOutputGate,
+    }));
+
+    kTensor = await projectSingleQkvTensor({
+      recorder,
+      normed,
+      layerWeights,
+      weightKey: 'kProj',
+      role: 'k_proj',
+      outputSize: numKVHeads * headDim,
+      outputLabel: 'K',
+      loraKey: 'k_proj',
+      numTokens,
+      hiddenSize,
+      layerIdx,
+      kernelPath,
+      matmulOutputDtype,
+      getWeightBuffer,
+      lora,
+      releaseTemporary,
     });
-    const split = await runSplitForMode(qkvTensor, {
+
+    vTensor = await projectSingleQkvTensor({
+      recorder,
+      normed,
+      layerWeights,
+      weightKey: 'vProj',
+      role: 'v_proj',
+      outputSize: numKVHeads * headDim,
+      outputLabel: 'V',
+      loraKey: 'v_proj',
       numTokens,
-      qSize: qSizeFused,
-      kSize: kSizeFused,
-      vSize: vSizeFused,
+      hiddenSize,
+      layerIdx,
+      kernelPath,
+      matmulOutputDtype,
+      getWeightBuffer,
+      lora,
+      releaseTemporary,
     });
-    releaseTemporary(qkvTensor.buffer);
-    if (onFusedQKV) {
-      onFusedQKV({ qSize: qSizeFused, kSize: kSizeFused, vSize: vSizeFused, totalSize: qkvSizeTotal });
+
+    return { qTensor, qGateTensor, kTensor, vTensor, usedFusedQKV: false };
+  } catch (error) {
+    for (const tensor of [qTensor, qGateTensor, kTensor, vTensor]) {
+      if (tensor?.buffer) {
+        releaseTemporary(tensor.buffer);
+      }
     }
-    return { qTensor: split.Q, qGateTensor: null, kTensor: split.K, vTensor: split.V, usedFusedQKV: true };
+    throw error;
   }
-
-  const { qTensor, qGateTensor } = await projectQueryWithOptionalGate({
-    recorder,
-    normed,
-    layerWeights,
-    numTokens,
-    numHeads,
-    headDim,
-    hiddenSize,
-    layerIdx,
-    kernelPath,
-    matmulOutputDtype,
-    getWeightBuffer,
-    lora,
-    releaseTemporary,
-    attentionOutputGate,
-  });
-
-  const kTensor = await projectSingleQkvTensor({
-    recorder,
-    normed,
-    layerWeights,
-    weightKey: 'kProj',
-    role: 'k_proj',
-    outputSize: numKVHeads * headDim,
-    outputLabel: 'K',
-    loraKey: 'k_proj',
-    numTokens,
-    hiddenSize,
-    layerIdx,
-    kernelPath,
-    matmulOutputDtype,
-    getWeightBuffer,
-    lora,
-    releaseTemporary,
-  });
-
-  const vTensor = await projectSingleQkvTensor({
-    recorder,
-    normed,
-    layerWeights,
-    weightKey: 'vProj',
-    role: 'v_proj',
-    outputSize: numKVHeads * headDim,
-    outputLabel: 'V',
-    loraKey: 'v_proj',
-    numTokens,
-    hiddenSize,
-    layerIdx,
-    kernelPath,
-    matmulOutputDtype,
-    getWeightBuffer,
-    lora,
-    releaseTemporary,
-  });
-
-  return { qTensor, qGateTensor, kTensor, vTensor, usedFusedQKV: false };
 }
 
 export async function applyAttentionQKNorm({

package/src/inference/pipelines/text/attention/record.js

@@ -24,10 +24,12 @@ import { selectRuleValue } from '../../../../rules/rule-registry.js';
 import { SlidingWindowKVCache } from '../../../kv-cache.js';
 import {
   recordAttentionInputs,
+  shouldForceF32AttentionProjectionForRoPE,
   resolveAttentionProjectionOutputDtype,
   projectAttentionQKV,
   applyAttentionQKNorm,
 } from './projections.js';
+import { prepareAttentionProjectionInput } from './output-projection.js';
 
 import { releaseOrTrack, shouldDebugLayer } from './types.js';
 
@@ -90,9 +92,20 @@ export async function recordLayerAttentionGPU(
   const allowF16Attention = wantsF16Output && kvCacheDtype === 'f16';
   let attentionInput = input;
   let attentionInputTemp = false;
+  let normed = attentionInput;
+  let qTensor = null;
+  let qGateTensor = null;
+  let kTensor = null;
+  let vTensor = null;
+  let attnOutput = null;
+  let attnForProjection = null;
+  let output = null;
+  let finalOutput = null;
+  let oProjInputTemp = null;
   if (wantsF16Output && !allowF16Attention) {
     attentionInput = await recordCastF16ToF32(recorder, input);
     attentionInputTemp = true;
+    normed = attentionInput;
   }
 
   if (!layerWeights) {
@@ -108,7 +121,7 @@ export async function recordLayerAttentionGPU(
 
   // 1. Input norm
 
-  let normed = attentionInput;
+  try {
   if (!skipInputNorm && layerWeights.inputNorm && getNormWeightBuffer) {
     const normWeightBuf = getNormWeightBuffer(layerWeights.inputNorm, 'input_norm');
     normed = await recordRMSNorm(recorder, attentionInput, normWeightBuf, rmsNormEps, {
@@ -131,8 +144,16 @@ export async function recordLayerAttentionGPU(
   }
 
   // 2. Q/K/V projections
-  const matmulOutputDtype = resolveAttentionProjectionOutputDtype(desiredOutputDtype);
-  let { qTensor, qGateTensor, kTensor, vTensor, usedFusedQKV } = await projectAttentionQKV({
+  const matmulOutputDtype = resolveAttentionProjectionOutputDtype(desiredOutputDtype, {
+    forceF32: shouldForceF32AttentionProjectionForRoPE({
+      attentionInputDtype: desiredOutputDtype,
+      headDim,
+      rotaryDim: config.ropeRotaryDim,
+      interleaved: config.ropeInterleaved,
+    }),
+  });
+  let usedFusedQKV = false;
+  ({ qTensor, qGateTensor, kTensor, vTensor, usedFusedQKV } = await projectAttentionQKV({
     recorder,
     normed,
     layerWeights,
@@ -153,7 +174,7 @@ export async function recordLayerAttentionGPU(
         trace.attn(layerIdx, `Using fused QKV path: ${qSizeFused}+${kSizeFused}+${vSizeFused}=${totalSize}`);
       }
       : null,
-  });
+  }));
 
   // Optional per-head Q/K normalization.
   // Some models use RMSNorm with (1+weight) offset formula, controlled by rmsNormWeightOffset.
@@ -502,9 +523,9 @@ export async function recordLayerAttentionGPU(
     throw new Error(`Unsupported attention kernel variant "${attentionKernelVariant}" at layer ${layerIdx}`);
   }
 
-  const attnOutput = await runAttentionKernel();
+  attnOutput = await runAttentionKernel();
 
-  let attnForProjection = attnOutput;
+  attnForProjection = attnOutput;
   if (qGateTensor) {
     attnForProjection = await recordSiLU(recorder, attnOutput, {
       size: numTokens * numHeads * headDim,
@@ -518,19 +539,19 @@ export async function recordLayerAttentionGPU(
 
   // 6. Output projection (with optional fused residual for decode)
 
-  let output;
+  output = null;
   let residualFused = false;
   let oProjInput = attnForProjection;
-  let oProjInputTemp = null;
+  oProjInputTemp = null;
   if (layerWeights.oProj && getWeightBuffer) {
+    ({ oProjInput, oProjInputTemp } = await prepareAttentionProjectionInput(
+      attnForProjection,
+      matmulOutputDtype,
+      (tensor) => recordCastF32ToF16(recorder, tensor)
+    ));
     const oProjBuf = getWeightBuffer(layerWeights.oProj, 'o_proj');
     const loraO = getLoRAModule(lora, layerIdx, 'o_proj');
 
-    if (matmulOutputDtype === 'f16' && attnForProjection.dtype !== 'f16') {
-      oProjInput = await recordCastF32ToF16(recorder, attnForProjection);
-      oProjInputTemp = oProjInput;
-    }
-
     // Use fused o_proj + residual for decode when possible
     // Note: dtype from WeightBuffer metadata (buffer-dtypes WeakMap removed)
     const oProjDtype = getWeightDtype(oProjBuf);
@@ -589,7 +610,7 @@ export async function recordLayerAttentionGPU(
     }
   }
 
-  let finalOutput = output;
+  finalOutput = output;
 
   const buffersToTrack = [];
   if (output.buffer !== attnForProjection.buffer) {
@@ -619,4 +640,46 @@ export async function recordLayerAttentionGPU(
   }
 
   return { output: finalOutput, residualFused };
+  } catch (error) {
+    const tracked = new Set();
+    const trackOnce = (buffer) => {
+      if (!buffer || tracked.has(buffer)) return;
+      tracked.add(buffer);
+      recorder.trackTemporaryBuffer(buffer);
+    };
+    if (finalOutput?.buffer && finalOutput.buffer !== output?.buffer) {
+      trackOnce(finalOutput.buffer);
+    }
+    if (output?.buffer && output.buffer !== attnForProjection?.buffer) {
+      trackOnce(output.buffer);
+    }
+    if (oProjInputTemp?.buffer) {
+      trackOnce(oProjInputTemp.buffer);
+    }
+    if (attnForProjection?.buffer && attnForProjection.buffer !== attnOutput?.buffer) {
+      trackOnce(attnForProjection.buffer);
+    }
+    if (attnOutput?.buffer) {
+      trackOnce(attnOutput.buffer);
+    }
+    if (qGateTensor?.buffer) {
+      trackOnce(qGateTensor.buffer);
+    }
+    if (qTensor?.buffer) {
+      trackOnce(qTensor.buffer);
+    }
+    if (kTensor?.buffer) {
+      trackOnce(kTensor.buffer);
+    }
+    if (vTensor?.buffer) {
+      trackOnce(vTensor.buffer);
+    }
+    if (normed?.buffer && normed.buffer !== attentionInput?.buffer) {
+      trackOnce(normed.buffer);
+    }
+    if (attentionInputTemp && attentionInput?.buffer) {
+      trackOnce(attentionInput.buffer);
+    }
+    throw error;
+  }
 }