@simulatte/doppler 0.1.8 → 0.2.0

package/src/inference/pipelines/text/logits/cpu.d.ts

@@ -1,81 +0,0 @@
-/**
- * CPU reference implementations for logits computation.
- *
- * Provides CPU fallback implementations for RMSNorm, matmul, and softcapping.
- * These are used when GPU is unavailable or for validation.
- *
- * @module inference/pipelines/text/logits/cpu
- */
-
-/**
- * CPU RMSNorm implementation.
- *
- * Computes: output[i] = (x[i] / rms) * weight[i] (or (1 + weight[i]) when enabled)
- * where rms = sqrt(mean(x^2) + eps)
- *
- * @param x - Input tensor
- * @param weight - Norm weights
- * @param eps - Epsilon for numerical stability
- * @param rmsNormWeightOffset - Use (1 + weight) scaling
- * @returns Normalized tensor
- */
-export function rmsNormCPU(
-  x: Float32Array,
-  weight: Float32Array,
-  eps?: number,
-  rmsNormWeightOffset?: boolean
-): Float32Array;
-
-/**
- * Convert a single float16 value to float32.
- *
- * @param h - Float16 value as uint16
- * @returns Float32 value
- */
-export function f16ToF32(h: number): number;
-
-/**
- * Convert a buffer of float16 values to float32.
- *
- * @param data - ArrayBuffer containing float16 data
- * @returns Float32Array with converted values
- */
-export function f16BufferToF32(data: ArrayBuffer): Float32Array;
-
-/**
- * CPU matmul implementation (fallback for non-GPU).
- *
- * Computes: output = input @ weight^T
- * Input: [M, K], Weight: [N, K] (row) or [K, N] (column), Output: [M, N]
- *
- * @param input - Input tensor [M, K]
- * @param weight - Weight tensor [N, K] (row) or [K, N] (column)
- * @param M - Batch size (num tokens)
- * @param N - Output size (vocab size)
- * @param K - Hidden size
- * @param layout - Weight layout ('row' or 'column')
- * @param weightStride - Optional stride override for weight indexing
- * @returns Output tensor [M, N]
- */
-export function matmulCPU(
-  input: Float32Array,
-  weight: Float32Array,
-  M: number,
-  N: number,
-  K: number,
-  layout?: 'row' | 'column',
-  weightStride?: number | null
-): Float32Array;
-
-/**
- * Apply softcapping to logits (Gemma 2 style).
- *
- * Computes: logits = tanh(logits / cap) * cap
- *
- * This bounds logits to [-cap, cap] with smooth transitions,
- * preventing extreme values from dominating softmax.
- *
- * @param logits - Logits tensor to modify in-place
- * @param cap - Softcap value (Gemma 2 default: 30.0)
- */
-export function applySoftcapping(logits: Float32Array, cap: number): void;

package/src/inference/pipelines/text/logits/cpu.js

@@ -1,91 +0,0 @@
-
-
-
-// Caller should pass eps from model config; inline default for safety
-export function rmsNormCPU(
-  x,
-  weight,
-  eps = 1e-5,
-  rmsNormWeightOffset = false
-) {
-  const n = x.length;
-  let sumSq = 0;
-  for (let i = 0; i < n; i++) {
-    sumSq += x[i] * x[i];
-  }
-  const rms = Math.sqrt(sumSq / n + eps);
-
-  const result = new Float32Array(n);
-  for (let i = 0; i < n; i++) {
-    const w = weight[i % weight.length];
-    const scale = rmsNormWeightOffset ? 1 + w : w;
-    result[i] = (x[i] / rms) * scale;
-  }
-  return result;
-}
-
-
-export function f16ToF32(h) {
-  const sign = (h >> 15) & 0x1;
-  const exp = (h >> 10) & 0x1f;
-  const mant = h & 0x3ff;
-
-  if (exp === 0) {
-    if (mant === 0) return sign ? -0 : 0;
-    const f = mant / 1024 * Math.pow(2, -14);
-    return sign ? -f : f;
-  }
-  if (exp === 31) {
-    return mant ? NaN : (sign ? -Infinity : Infinity);
-  }
-
-  const f = (1 + mant / 1024) * Math.pow(2, exp - 15);
-  return sign ? -f : f;
-}
-
-
-export function f16BufferToF32(data) {
-  const u16 = new Uint16Array(data);
-  const out = new Float32Array(u16.length);
-  for (let i = 0; i < u16.length; i++) {
-    out[i] = f16ToF32(u16[i]);
-  }
-  return out;
-}
-
-
-export function matmulCPU(
-  input,
-  weight,
-  M,
-  N,
-  K,
-  layout = 'row',
-  weightStride
-) {
-  const result = new Float32Array(M * N);
-  const stride = weightStride ?? (layout === 'row' ? K : N);
-
-  for (let m = 0; m < M; m++) {
-    for (let n = 0; n < N; n++) {
-      let sum = 0;
-      for (let k = 0; k < K; k++) {
-        // Row layout: weight is [N, K] (vocab x hidden).
-        // Column layout: weight is [K, N] (hidden x vocab).
-        const weightIndex = layout === 'row'
-          ? n * stride + k
-          : k * stride + n;
-        sum += input[m * K + k] * weight[weightIndex];
-      }
-      result[m * N + n] = sum;
-    }
-  }
-  return result;
-}
-
-
-export function applySoftcapping(logits, cap) {
-  for (let i = 0; i < logits.length; i++) {
-    logits[i] = Math.tanh(logits[i] / cap) * cap;
-  }
-}

package/src/inference/pipelines/text/logits/gpu.d.ts

@@ -1,113 +0,0 @@
-/**
- * GPU implementations for logits computation.
- *
- * Provides GPU-accelerated implementations for computing logits,
- * including both immediate execution and recorded (batched) variants.
- *
- * @module inference/pipelines/text/logits/gpu
- */
-
-import type { Tensor } from '../../../../gpu/tensor.js';
-import type { CpuWeightBuffer } from '../../../../gpu/weight-buffer.js';
-import type { CommandRecorder } from '../../../../gpu/command-recorder.js';
-import type { LargeWeightConfigSchema, ProbeConfigSchema, KernelPathSchema } from '../../../../config/schema/index.js';
-import type { LogitsConfig, LogitsWeights, LogitsDebugFlags } from './types.js';
-
-/**
- * Resolve CPU weight buffer dimensions for LM head.
- */
-export function resolveCpuWeightDims(lmHead: CpuWeightBuffer): { vocabSize: number; hiddenSize: number };
-
-/**
- * Calculate the maximum rows per chunk for LM head matmul.
- */
-export function resolveLmHeadChunkRows(
-  device: GPUDevice,
-  numTokens: number,
-  hiddenSize: number,
-  config?: LargeWeightConfigSchema
-): number;
-
-/**
- * Extract a chunk of the LM head weight matrix.
- */
-export function extractLmHeadChunk(
-  data: Float32Array,
-  layout: 'row' | 'column',
-  hiddenSize: number,
-  vocabSize: number,
-  rowOffset: number,
-  rowCount: number
-): Float32Array;
-
-/**
- * Write chunk logits to the full logits buffer.
- */
-export function writeChunkLogits(
-  target: Float32Array,
-  chunk: Float32Array,
-  numTokens: number,
-  vocabSize: number,
-  rowOffset: number,
-  rowCount: number
-): void;
-
-/**
- * Compute logits using chunked GPU matmul for large LM heads.
- *
- * Used when LM head weights are CPU-resident and too large
- * to fit in a single GPU buffer binding.
- */
-export function computeChunkedLogitsGPU(
-  normedTensor: Tensor,
-  lmHead: CpuWeightBuffer,
-  numTokens: number,
-  hiddenSize: number,
-  vocabSize: number,
-  weightVocabSize: number,
-  debugProbes?: ProbeConfigSchema[] | null,
-  largeWeightConfig: LargeWeightConfigSchema,
-  kernelPath?: KernelPathSchema | null
-): Promise<Float32Array>;
-
-/**
- * Compute logits and return GPU buffer directly (deferred readback).
- *
- * This variant avoids the ~1MB readback per token, enabling GPU-side sampling.
- * Use with runGPUSample or runArgmax to sample directly on GPU.
- *
- * @param hiddenStates - Hidden states from transformer [numTokens, hiddenSize]
- * @param numTokens - Number of tokens
- * @param weights - Final norm and LM head weights
- * @param config - Model configuration for logits
- * @param debugFlags - Debug flags to prevent repeated logging (optional)
- * @returns GPU buffer containing logits [numTokens, vocabSize]
- */
-export function computeLogitsGPU(
-  hiddenStates: GPUBuffer | Float32Array,
-  numTokens: number,
-  weights: LogitsWeights,
-  config: LogitsConfig,
-  debugFlags?: LogitsDebugFlags,
-): Promise<{ logitsBuffer: GPUBuffer; vocabSize: number; logitsDtype: 'f16' | 'f32' } | null>;
-
-/**
- * Record logits computation (batched, no submit).
- *
- * This variant uses the CommandRecorder to batch logits computation with
- * preceding layer operations, avoiding a GPU sync point.
- *
- * @param recorder - CommandRecorder for batched operations
- * @param hiddenStates - Hidden states from transformer [numTokens, hiddenSize]
- * @param numTokens - Number of tokens
- * @param weights - Final norm and LM head weights
- * @param config - Model configuration for logits
- * @returns GPU buffer containing logits [numTokens, vocabSize] and vocab size
- */
-export function recordLogitsGPU(
-  recorder: CommandRecorder,
-  hiddenStates: GPUBuffer,
-  numTokens: number,
-  weights: LogitsWeights,
-  config: LogitsConfig,
-): Promise<{ logitsBuffer: GPUBuffer; vocabSize: number; logitsDtype: 'f16' | 'f32' }>;

package/src/inference/pipelines/text/logits/gpu.js

@@ -1,411 +0,0 @@
-
-
-import { getDevice, getKernelCapabilities } from '../../../../gpu/device.js';
-import { acquireBuffer, releaseBuffer } from '../../../../memory/buffer-pool.js';
-import { runMatmul, runRMSNorm } from '../../../../gpu/kernel-selector.js';
-import { recordMatmul } from '../../../../gpu/kernels/matmul.js';
-import { recordRMSNorm } from '../../../../gpu/kernels/rmsnorm.js';
-import { createTensor } from '../../../../gpu/tensor.js';
-import { castF16ToF32, castF32ToF16, recordCastF16ToF32 } from '../../../../gpu/kernels/cast.js';
-import { createWeightBuffer, isWeightBuffer, isCpuWeightBuffer } from '../../../../gpu/weight-buffer.js';
-import { log, trace, isTraceEnabled } from '../../../../debug/index.js';
-import { getRuntimeConfig } from '../../../../config/runtime.js';
-import { selectRuleValue } from '../../../../rules/rule-registry.js';
-import { runProbes } from '../probes.js';
-import { f16BufferToF32 } from './cpu.js';
-import { readBufferWithCleanup } from './utils.js';
-
-function shouldForceStableF32Logits(config, inputDtype) {
-  // Small Gemma-family checkpoints can overflow in pure F16 logits path after RMSNorm offset.
-  return inputDtype === 'f16'
-    && config.rmsNormWeightOffset === true
-    && Number.isFinite(config.hiddenSize)
-    && config.hiddenSize <= 768;
-}
-
-
-export function resolveCpuWeightDims(lmHead) {
-  if (lmHead.shape.length !== 2) {
-    throw new Error(`[Logits] CPU LM head shape must be 2D, got [${lmHead.shape.join(', ')}]`);
-  }
-  if (lmHead.layout === 'column') {
-    return { hiddenSize: lmHead.shape[0], vocabSize: lmHead.shape[1] };
-  }
-  return { vocabSize: lmHead.shape[0], hiddenSize: lmHead.shape[1] };
-}
-
-
-export function resolveLmHeadChunkRows(
-  device,
-  numTokens,
-  hiddenSize,
-  config
-) {
-  const resolved = config ?? getRuntimeConfig().inference.largeWeights;
-  if (resolved.safetyRatio == null) {
-    throw new Error('runtime.inference.largeWeights.safetyRatio is required.');
-  }
-  const safety = Math.min(Math.max(resolved.safetyRatio, 0.1), 1);
-  const maxBinding = Math.min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize);
-  const maxBytes = Math.floor(maxBinding * safety);
-
-  const maxRowsByWeight = Math.floor(maxBytes / (hiddenSize * 4));
-  const maxRowsByOutput = Math.floor(maxBytes / (numTokens * 4));
-  const maxRows = Math.min(maxRowsByWeight, maxRowsByOutput);
-
-  if (!Number.isFinite(maxRows) || maxRows <= 0) {
-    throw new Error(
-      `[Logits] LM head chunk size underflow (maxBytes=${maxBytes}, hiddenSize=${hiddenSize}, numTokens=${numTokens}).`
-    );
-  }
-
-  const override = resolved.lmHeadChunkRows ?? null;
-  if (override && override > 0) {
-    return Math.min(override, maxRows);
-  }
-  return maxRows;
-}
-
-
-export function extractLmHeadChunk(
-  data,
-  layout,
-  hiddenSize,
-  vocabSize,
-  rowOffset,
-  rowCount
-) {
-  if (layout === 'row') {
-    const start = rowOffset * hiddenSize;
-    return data.subarray(start, start + rowCount * hiddenSize);
-  }
-
-  const chunk = new Float32Array(hiddenSize * rowCount);
-  for (let k = 0; k < hiddenSize; k++) {
-    const srcOffset = k * vocabSize + rowOffset;
-    const dstOffset = k * rowCount;
-    chunk.set(data.subarray(srcOffset, srcOffset + rowCount), dstOffset);
-  }
-  return chunk;
-}
-
-
-export function writeChunkLogits(
-  target,
-  chunk,
-  numTokens,
-  vocabSize,
-  rowOffset,
-  rowCount
-) {
-  for (let t = 0; t < numTokens; t++) {
-    const srcOffset = t * rowCount;
-    const dstOffset = t * vocabSize + rowOffset;
-    target.set(chunk.subarray(srcOffset, srcOffset + rowCount), dstOffset);
-  }
-}
-
-
-export async function computeChunkedLogitsGPU(
-  normedTensor,
-  lmHead,
-  numTokens,
-  hiddenSize,
-  vocabSize,
-  weightVocabSize,
-  debugProbes,
-  largeWeightConfig,
-  kernelPath = null
-) {
-  const device = getDevice();
-  if (!device) {
-    throw new Error('[Logits] GPU device not available for chunked LM head.');
-  }
-  if (!largeWeightConfig) {
-    throw new Error('[Logits] largeWeights config is required for chunked LM head.');
-  }
-
-  const chunkRows = resolveLmHeadChunkRows(device, numTokens, hiddenSize, largeWeightConfig);
-  const caps = getKernelCapabilities();
-  const weightDtype = selectRuleValue('inference', 'dtype', 'lmHeadChunkWeightDtype', {
-    preferF16: largeWeightConfig.preferF16,
-    lmHeadDtype: lmHead.dtype,
-    hasF16: caps.hasF16,
-  });
-  const preferF16 = weightDtype === 'f16';
-  const logits = new Float32Array(numTokens * vocabSize);
-
-  if (isTraceEnabled('logits')) {
-    trace.logits(`LM_HEAD_CHUNKED: vocab=${vocabSize}, chunkRows=${chunkRows}, layout=${lmHead.layout}, f16=${preferF16}`);
-  }
-
-  for (let rowOffset = 0; rowOffset < vocabSize; rowOffset += chunkRows) {
-    const rowCount = Math.min(chunkRows, vocabSize - rowOffset);
-    const chunkData = extractLmHeadChunk(
-      lmHead.data,
-      lmHead.layout,
-      hiddenSize,
-      weightVocabSize,
-      rowOffset,
-      rowCount
-    );
-
-    const f32Buffer = acquireBuffer(chunkData.byteLength, undefined, 'lm_head_chunk_f32');
-    device.queue.writeBuffer(
-      f32Buffer,
-      0,
-      chunkData.buffer,
-      chunkData.byteOffset,
-      chunkData.byteLength
-    );
-
-    const chunkShape = lmHead.layout === 'column'
-      ? [hiddenSize, rowCount]
-      : [rowCount, hiddenSize];
-
-    let weightBuffer = createWeightBuffer(f32Buffer, 'f32', lmHead.layout, chunkShape, 'lm_head_chunk_f32');
-
-    if (preferF16) {
-      const f32Tensor = createTensor(f32Buffer, 'f32', chunkShape, 'lm_head_chunk_f32');
-      const f16Tensor = await castF32ToF16(f32Tensor);
-      releaseBuffer(f32Buffer);
-      weightBuffer = createWeightBuffer(f16Tensor.buffer, 'f16', lmHead.layout, chunkShape, 'lm_head_chunk_f16');
-    }
-
-    const logitsTensor = await runMatmul(normedTensor, weightBuffer, numTokens, rowCount, hiddenSize, {
-      transposeB: 'auto',
-      role: 'lm_head',
-      kernelPath,
-    });
-
-    if (debugProbes?.length) {
-      await runProbes('logits', logitsTensor.buffer, {
-        numTokens,
-        hiddenSize: rowCount,
-        probes: debugProbes,
-        dtype: logitsTensor.dtype,
-      });
-    }
-
-    const logitsBytes = selectRuleValue('shared', 'dtype', 'bytesFromDtype', { dtype: logitsTensor.dtype });
-    const chunkLogitsData = await readBufferWithCleanup(
-      logitsTensor.buffer,
-      numTokens * rowCount * logitsBytes,
-      () => {
-        releaseBuffer(logitsTensor.buffer);
-        releaseBuffer(weightBuffer.buffer);
-      }
-    );
-    const chunkLogits = logitsTensor.dtype === 'f16'
-      ? f16BufferToF32(chunkLogitsData)
-      : new Float32Array(chunkLogitsData);
-    writeChunkLogits(logits, chunkLogits, numTokens, vocabSize, rowOffset, rowCount);
-  }
-
-  return logits;
-}
-
-
-export async function computeLogitsGPU(
-  hiddenStates,
-  numTokens,
-  weights,
-  config,
-  debugFlags,
-) {
-  const {
-    hiddenSize,
-    vocabSize,
-    rmsNormEps,
-    useTiedEmbeddings,
-    embeddingVocabSize,
-    activationDtype,
-  } = config;
-  const { finalNorm, lmHead } = weights;
-  const device = getDevice();
-
-  if (!device) {
-    return null;
-  }
-  if (!activationDtype) {
-    throw new Error('[Logits] activationDtype is required.');
-  }
-
-  if (!finalNorm || !lmHead) {
-    log.warn('Pipeline', 'Final norm or LM head not loaded');
-    return null;
-  }
-  if (isCpuWeightBuffer(lmHead)) {
-    return null;
-  }
-
-  // Get or create input buffer
-  
-  let inputBuffer;
-  let inputBufferOwned = false;
-  if (hiddenStates instanceof GPUBuffer) {
-    inputBuffer = hiddenStates;
-  } else {
-    inputBuffer = acquireBuffer( (hiddenStates).byteLength, undefined, 'logits_input');
-    device.queue.writeBuffer(inputBuffer, 0,  (hiddenStates));
-    inputBufferOwned = true;
-  }
-
-  // Apply final RMSNorm
-  
-  let normWeightBuffer;
-  let normWeightBufferOwned = false;
-  if (finalNorm instanceof GPUBuffer) {
-    normWeightBuffer = finalNorm;
-  } else {
-    normWeightBuffer = acquireBuffer( (finalNorm).byteLength, undefined, 'final_norm_w');
-    device.queue.writeBuffer(normWeightBuffer, 0,  (finalNorm));
-    normWeightBufferOwned = true;
-  }
-
-  
-  const inputDtype = hiddenStates instanceof GPUBuffer ? activationDtype : 'f32';
-  // Wrap input buffer as Tensor for RMSNorm
-  const inputTensor = createTensor(inputBuffer, inputDtype, [numTokens, hiddenSize], 'logits_input');
-  const forceStableF32Logits = shouldForceStableF32Logits(config, inputDtype);
-  let normInputTensor = inputTensor;
-  let normInputOwned = false;
-  if (forceStableF32Logits) {
-    normInputTensor = await castF16ToF32(inputTensor);
-    normInputOwned = true;
-  }
-  const normedTensor = await runRMSNorm(normInputTensor, normWeightBuffer, rmsNormEps, {
-    batchSize: numTokens,
-    hiddenSize,
-    rmsNormWeightOffset: config.rmsNormWeightOffset,
-  });
-  if (normInputOwned) {
-    releaseBuffer(normInputTensor.buffer);
-  }
-
-  // Project to vocab via LM head
-  
-  let lmHeadBuffer;
-  let lmHeadBufferOwned = false;
-  if (lmHead instanceof GPUBuffer) {
-    lmHeadBuffer = lmHead;
-  } else if (isWeightBuffer(lmHead)) {
-    lmHeadBuffer = lmHead;
-  } else {
-    const rawBuffer = acquireBuffer( (lmHead).byteLength, undefined, 'lm_head_w');
-    device.queue.writeBuffer(rawBuffer, 0,  (lmHead));
-    lmHeadBuffer = rawBuffer;
-    lmHeadBufferOwned = true;
-  }
-
-  const matmulVocabSize = useTiedEmbeddings && embeddingVocabSize
-    ? embeddingVocabSize
-    : vocabSize;
-
-  const logitsTensor = await runMatmul(normedTensor, lmHeadBuffer, numTokens, matmulVocabSize, hiddenSize, {
-    transposeB: 'auto',
-    role: 'lm_head',
-    kernelPath: config.kernelPath ?? null,
-  });
-
-  // Cleanup intermediate buffers (but keep logitsBuffer)
-  if (inputBufferOwned) releaseBuffer(inputBuffer);
-  releaseBuffer(normedTensor.buffer);
-  if (normWeightBufferOwned) releaseBuffer(normWeightBuffer);
-  if (lmHeadBufferOwned) releaseBuffer(isWeightBuffer(lmHeadBuffer) ? lmHeadBuffer.buffer : lmHeadBuffer);
-
-  return { logitsBuffer: logitsTensor.buffer, vocabSize: matmulVocabSize, logitsDtype: logitsTensor.dtype };
-}
-
-
-export async function recordLogitsGPU(
-  recorder,
-  hiddenStates,
-  numTokens,
-  weights,
-  config,
-) {
-  const {
-    hiddenSize,
-    vocabSize,
-    rmsNormEps,
-    useTiedEmbeddings,
-    embeddingVocabSize,
-    activationDtype = 'f32',
-  } = config;
-  const { finalNorm, lmHead } = weights;
-  const matmulVocabSize = useTiedEmbeddings && embeddingVocabSize ? embeddingVocabSize : vocabSize;
-
-  if (!finalNorm || !lmHead) {
-    throw new Error('[recordLogitsGPU] Final norm or LM head not loaded');
-  }
-  if (isCpuWeightBuffer(lmHead)) {
-    throw new Error('[recordLogitsGPU] CPU-resident LM head not supported in recorded path');
-  }
-
-  // Get norm weight buffer
-  
-  let normWeightBuffer;
-  let normWeightOwned = false;
-  if (finalNorm instanceof GPUBuffer) {
-    normWeightBuffer = finalNorm;
-  } else {
-    normWeightBuffer = acquireBuffer( (finalNorm).byteLength, undefined, 'final_norm_w');
-    recorder.device.queue.writeBuffer(normWeightBuffer, 0,  (finalNorm));
-    normWeightOwned = true;
-  }
-
-  
-  const inputDtype = activationDtype;
-  // Wrap input buffer as Tensor for RMSNorm
-  const inputTensor = createTensor(hiddenStates, inputDtype, [numTokens, hiddenSize], 'logits_input');
-  const forceStableF32Logits = shouldForceStableF32Logits(config, inputDtype);
-  let normInputTensor = inputTensor;
-  let normInputOwned = false;
-  if (forceStableF32Logits) {
-    normInputTensor = await recordCastF16ToF32(recorder, inputTensor);
-    normInputOwned = true;
-  }
-  // Record RMSNorm (no submit)
-  const normedTensor = await recordRMSNorm(recorder, normInputTensor, normWeightBuffer, rmsNormEps, {
-    batchSize: numTokens,
-    hiddenSize,
-    rmsNormWeightOffset: config.rmsNormWeightOffset,
-  });
-
-  // Get LM head buffer
-  
-  let lmHeadBuffer;
-  let lmHeadBufferOwned = false;
-  if (lmHead instanceof GPUBuffer) {
-    lmHeadBuffer = lmHead;
-  } else if (isWeightBuffer(lmHead)) {
-    lmHeadBuffer = lmHead;
-  } else {
-    const rawBuffer = acquireBuffer( (lmHead).byteLength, undefined, 'lm_head_w');
-    recorder.device.queue.writeBuffer(rawBuffer, 0,  (lmHead));
-    lmHeadBuffer = rawBuffer;
-    lmHeadBufferOwned = true;
-  }
-
-  // Record matmul (no submit)
-  const logitsTensor = await recordMatmul(recorder, normedTensor, lmHeadBuffer, numTokens, matmulVocabSize, hiddenSize, {
-    transposeB: 'auto',
-    role: 'lm_head',
-    kernelPath: config.kernelPath ?? null,
-  });
-
-  // Track intermediate buffer for cleanup after submit
-  recorder.trackTemporaryBuffer(normedTensor.buffer);
-  if (normWeightOwned) {
-    recorder.trackTemporaryBuffer(normWeightBuffer);
-  }
-  if (normInputOwned) {
-    recorder.trackTemporaryBuffer(normInputTensor.buffer);
-  }
-  if (lmHeadBufferOwned) {
-    recorder.trackTemporaryBuffer(isWeightBuffer(lmHeadBuffer) ? lmHeadBuffer.buffer : lmHeadBuffer);
-  }
-
-  return { logitsBuffer: logitsTensor.buffer, vocabSize: matmulVocabSize, logitsDtype: logitsTensor.dtype };
-}