@simulatte/doppler 0.1.8 → 0.2.0

package/src/inference/pipelines/text/ops.js

@@ -1,367 +0,0 @@
-
-
-import {
-  runRMSNorm, runResidualAdd, runMatmul, runSiLU, runGeLU,
-  recordRMSNorm, recordResidualAdd, recordMatmul, recordSiLU, recordGeLU,
-  runSiLURowSplit, recordSiLURowSplit,
-  runMatmulRMSNormFused, recordMatmulRMSNormFused,
-  runConv2D, recordConv2D,
-} from '../../../gpu/kernel-selector.js';
-import {
-  castF16ToF32,
-  castF32ToF16,
-  recordCastF16ToF32,
-  recordCastF32ToF16,
-} from '../../../gpu/kernels/cast.js';
-import { createTensor } from '../../../gpu/tensor.js';
-import { releaseBuffer } from '../../../memory/buffer-pool.js';
-import { kernelTrace, traceStep } from './kernel-trace.js';
-import {
-  runLayerAttentionGPU,
-  recordLayerAttentionGPU,
-} from './attention/index.js';
-import { runLinearAttentionLayer } from './linear-attention.js';
-
-
-export function isDecodeBuffer(decodeBuffers, buffer) {
-  return !!decodeBuffers?.ownsBuffer(buffer);
-}
-
-
-export function releaseOrTrack(recorder, buffer, decodeBuffers) {
-  if (isDecodeBuffer(decodeBuffers, buffer)) {
-    return;
-  }
-  if (recorder) {
-    recorder.trackTemporaryBuffer(buffer);
-  } else {
-    releaseBuffer(buffer);
-  }
-}
-
-
-export async function doRMSNorm(input, weight, eps, options, recorder) {
-  const result = recorder
-    ? await recordRMSNorm(recorder, input, weight, eps, options)
-    : await runRMSNorm(input, weight, eps, options);
-
-  // Trace the kernel output
-  if (kernelTrace.enabled && !recorder) {
-    const layer = options.layerIdx ?? -1;
-    const label = options.label ?? 'rmsnorm';
-    await traceStep('rmsnorm', label, layer, result.buffer, [options.batchSize, options.hiddenSize]);
-  }
-
-  return result;
-}
-
-
-export async function doResidualAdd(a, b, size, recorder, traceOptions) {
-  const options = traceOptions?.outputBuffer ? { outputBuffer: traceOptions.outputBuffer } : {};
-  const result = recorder
-    ? await recordResidualAdd(recorder, a, b, size, options)
-    : await runResidualAdd(a, b, size, options);
-
-  // Trace the kernel output
-  if (kernelTrace.enabled && !recorder && traceOptions) {
-    await traceStep('residual_add', traceOptions.label ?? 'residual', traceOptions.layerIdx ?? -1, result.buffer, [size]);
-  }
-
-  return result;
-}
-
-
-export async function doMatmul(A, B, M, N, K, options = {}, recorder) {
-  const result = recorder
-    ? await recordMatmul(recorder, A, B, M, N, K, options)
-    : await runMatmul(A, B, M, N, K, options);
-
-  // Trace the kernel output
-  if (kernelTrace.enabled && !recorder) {
-    const layer = options.layerIdx ?? -1;
-    const label = options.label ?? 'matmul';
-    await traceStep('matmul', label, layer, result.buffer, [M, N]);
-  }
-
-  return result;
-}
-
-
-export async function doSiLU(input, options = {}, recorder) {
-  const result = recorder
-    ? await recordSiLU(recorder, input, options)
-    : await runSiLU(input, options);
-
-  // Trace the kernel output
-  if (kernelTrace.enabled && !recorder && options.size) {
-    await traceStep('silu', options.label ?? 'silu', options.layerIdx ?? -1, result.buffer, [options.size]);
-  }
-
-  return result;
-}
-
-
-export async function doGeLU(input, options = {}, recorder) {
-  const result = recorder
-    ? await recordGeLU(recorder, input, options)
-    : await runGeLU(input, options);
-
-  // Trace the kernel output
-  if (kernelTrace.enabled && !recorder && options.size) {
-    await traceStep('gelu', options.label ?? 'gelu', options.layerIdx ?? -1, result.buffer, [options.size]);
-  }
-
-  return result;
-}
-
-
-export async function doSiLURowSplit(input, options, recorder) {
-  const result = recorder
-    ? await recordSiLURowSplit(recorder, input, options)
-    : await runSiLURowSplit(input, options);
-
-  // Trace the kernel output
-  if (kernelTrace.enabled && !recorder) {
-    await traceStep('silu_row_split', options.label ?? 'ffn_activation', options.layerIdx ?? -1, result.buffer, [options.numTokens, options.dim]);
-  }
-
-  return result;
-}
-
-
-export async function doMatmulRMSNormFused(input, weight, normWeight, options, recorder) {
-  // The fused kernel takes Tensor input but residual is still GPUBuffer
-  const fusedOptions = {
-    N: options.N,
-    K: options.K,
-    eps: options.eps,
-    residual: options.residual?.buffer ?? null,
-    outputBuffer: options.outputBuffer,
-    transposeB: options.transposeB,
-    rmsNormWeightOffset: options.rmsNormWeightOffset,
-    label: options.label ?? null,
-  };
-  const resultTensor = recorder
-    ? await recordMatmulRMSNormFused(recorder, input, weight, normWeight, fusedOptions)
-    : await runMatmulRMSNormFused(input, weight, normWeight, fusedOptions);
-
-  // Trace the kernel output
-  if (kernelTrace.enabled && !recorder) {
-    await traceStep('fused_matmul_rmsnorm', options.label ?? 'fused_matmul_rmsnorm', options.layerIdx ?? -1, resultTensor.buffer, [1, options.N]);
-  }
-
-  return resultTensor;
-}
-
-export async function doConv(
-  inputTensor,
-  convInProj,
-  convKernel,
-  convOutProj,
-  options = {},
-  recorder
-) {
-  const numTokens = Number(options.numTokens);
-  const hiddenSize = Number(options.hiddenSize);
-  const layerIdx = Number.isFinite(options.layerIdx) ? options.layerIdx : -1;
-  const label = options.label ?? 'conv';
-  const kernelPath = options.kernelPath ?? null;
-
-  if (!Number.isFinite(numTokens) || numTokens <= 0) {
-    throw new Error('doConv requires numTokens > 0.');
-  }
-  if (!Number.isFinite(hiddenSize) || hiddenSize <= 0) {
-    throw new Error('doConv requires hiddenSize > 0.');
-  }
-
-  // Use the first 2x hidden projection channels as a gated conv-state projection.
-  let inProj = null;
-  let activated = null;
-  let convInput = null;
-  let outProj = null;
-  try {
-    inProj = await doMatmul(
-      inputTensor,
-      convInProj,
-      numTokens,
-      hiddenSize * 2,
-      hiddenSize,
-      {
-        transposeB: 'auto',
-        label: `${label}.in_proj`,
-        layerIdx,
-        kernelPath,
-        role: 'conv_in_proj',
-      },
-      recorder
-    );
-    activated = await doSiLURowSplit(inProj, {
-      numTokens,
-      dim: hiddenSize,
-      activation: 'silu',
-      swigluLimit: options.swigluLimit ?? null,
-      label: `${label}.activation`,
-      layerIdx,
-    }, recorder);
-
-    releaseOrTrack(recorder, inProj.buffer);
-    inProj = null;
-
-    convInput = activated;
-    if (convKernel && options.conv2d && options.conv2d.enabled === true) {
-      const convTensorInput = createTensor(activated.buffer, activated.dtype, [
-        options.conv2d.inChannels,
-        options.conv2d.height,
-        options.conv2d.width,
-      ], `${label}.conv_input`);
-      const convOptions = {
-        inChannels: options.conv2d.inChannels,
-        outChannels: options.conv2d.outChannels,
-        height: options.conv2d.height,
-        width: options.conv2d.width,
-        kernelH: options.conv2d.kernelH,
-        kernelW: options.conv2d.kernelW,
-        stride: options.conv2d.stride ?? 1,
-        pad: options.conv2d.pad ?? 0,
-      };
-      const convResult = recorder
-        ? await recordConv2D(recorder, convTensorInput, convKernel, null, convOptions)
-        : await runConv2D(convTensorInput, convKernel, null, convOptions);
-      convInput = createTensor(
-        convResult.buffer,
-        convResult.dtype,
-        [numTokens, hiddenSize],
-        `${label}.conv_output`
-      );
-      releaseOrTrack(recorder, activated.buffer);
-      activated = null;
-    }
-
-    outProj = await doMatmul(
-      convInput,
-      convOutProj,
-      numTokens,
-      hiddenSize,
-      hiddenSize,
-      {
-        transposeB: 'auto',
-        label: `${label}.out_proj`,
-        layerIdx,
-        kernelPath,
-        role: 'conv_out_proj',
-      },
-      recorder
-    );
-
-    if (convInput && (!activated || convInput.buffer !== activated.buffer)) {
-      releaseOrTrack(recorder, convInput.buffer);
-      convInput = null;
-    } else if (activated) {
-      releaseOrTrack(recorder, activated.buffer);
-      activated = null;
-    }
-
-    if (kernelTrace.enabled && !recorder) {
-      await traceStep('conv', label, layerIdx, outProj.buffer, [numTokens, hiddenSize]);
-    }
-    return outProj;
-  } catch (error) {
-    if (outProj) releaseOrTrack(recorder, outProj.buffer);
-    if (convInput && (!activated || convInput.buffer !== activated.buffer)) releaseOrTrack(recorder, convInput.buffer);
-    if (activated) releaseOrTrack(recorder, activated.buffer);
-    if (inProj) releaseOrTrack(recorder, inProj.buffer);
-    throw error;
-  }
-}
-
-export async function doCast(input, toDtype, recorder) {
-  if (toDtype !== 'f16' && toDtype !== 'f32') {
-    throw new Error(`Unsupported cast target dtype "${toDtype}"`);
-  }
-  if (input.dtype === toDtype) {
-    return input;
-  }
-  if (input.dtype === 'f16' && toDtype === 'f32') {
-    return recorder
-      ? recordCastF16ToF32(recorder, input)
-      : castF16ToF32(input);
-  }
-  if (input.dtype === 'f32' && toDtype === 'f16') {
-    return recorder
-      ? recordCastF32ToF16(recorder, input)
-      : castF32ToF16(input);
-  }
-  throw new Error(`Unsupported cast path ${input.dtype} -> ${toDtype}`);
-}
-
-
-export async function doAttention(
-  inputTensor,
-  layerWeights,
-  config,
-  state,
-  debug,
-  debugFlags,
-  getWeightBufferFn,
-  getNormWeightBufferFn,
-  debugCheckBuffer,
-  recorder,
-  lora
-) {
-  const normalizedLayerType = String(config?.layerType ?? '').trim().toLowerCase();
-  const isLinearLayer = normalizedLayerType === 'linear_attention'
-    || normalizedLayerType === 'linear'
-    || normalizedLayerType === 'gated_delta'
-    || normalizedLayerType === 'gated_delta_net';
-  if (isLinearLayer) {
-    return {
-      output: await runLinearAttentionLayer(inputTensor, layerWeights, {
-        layerIdx: config.layerIdx,
-        numTokens: config.numTokens,
-        hiddenSize: config.hiddenSize,
-        config,
-        currentSeqLen: config.currentSeqLen,
-        activationDtype: config.activationDtype,
-        kernelPath: config.kernelPath ?? null,
-        linearRuntime: state?.linearRuntime ?? null,
-        getWeightBuffer: getWeightBufferFn,
-        getNormWeightBuffer: getNormWeightBufferFn,
-        recorder: recorder ?? null,
-      }),
-      residualFused: false,
-    };
-  }
-
-  const isBDPA = state?.kvCache?.layout === 'bdpa_paged';
-  if (recorder && isBDPA) {
-    throw new Error('BDPA attention does not support command recorder mode. Disable command batching for BDPA.');
-  }
-
-  if (recorder) {
-    return recordLayerAttentionGPU(
-      recorder,
-      inputTensor,
-      layerWeights,
-      config,
-      state,
-      debug,
-      debugFlags,
-      getWeightBufferFn,
-      getNormWeightBufferFn,
-      debugCheckBuffer,
-      lora
-    );
-  }
-  return runLayerAttentionGPU(
-    inputTensor,
-    layerWeights,
-    config,
-    state,
-    debug,
-    debugFlags,
-    getWeightBufferFn,
-    getNormWeightBufferFn,
-    debugCheckBuffer,
-    lora
-  );
-}

package/src/inference/pipelines/text/probes.d.ts

@@ -1,31 +0,0 @@
-/**
- * Config-driven probe helpers for targeted GPU buffer inspection.
- *
- * Probes let us read specific token/dimension values without adding
- * ad-hoc logs in the code. Probes are configured via runtime.shared.debug.probes.
- */
-
-import type { ProbeConfigSchema, ProbeStage } from '../../../config/schema/index.js';
-import type { CommandRecorder } from '../../../gpu/command-recorder.js';
-
-/**
- * Run configured probes for a specific stage.
- */
-export function runProbes(
-  stage: ProbeStage,
-  buffer: GPUBuffer | Float32Array,
-  options: {
-    layerIdx?: number;
-    numTokens: number;
-    hiddenSize: number;
-    probes?: ProbeConfigSchema[] | null;
-    recorder?: CommandRecorder | null;
-    dtype?: 'f16' | 'f32';
-  }
-): Promise<void>;
-
-export function hasProbeStage(
-  probes: ProbeConfigSchema[] | null | undefined,
-  stage: ProbeStage,
-  layerIdx?: number
-): boolean;

package/src/inference/pipelines/text/probes.js

@@ -1,170 +0,0 @@
-
-
-import { trace } from '../../../debug/index.js';
-import { getDevice } from '../../../gpu/device.js';
-import { allowReadback } from '../../../gpu/perf-guards.js';
-import { f16ToF32 } from '../../../loader/dtype-utils.js';
-import { readBufferSlice } from '../../../memory/buffer-pool.js';
-
-
-const STAGE_DEFAULT_CATEGORY = {
-  embed_out: 'embed',
-  // Attention stages (per-layer)
-  attn_input: 'attn',
-  attn_normed: 'attn',
-  linear_qkv_proj: 'attn',
-  linear_z_proj: 'attn',
-  linear_a_proj: 'attn',
-  linear_b_proj: 'attn',
-  linear_core_out: 'attn',
-  q_proj: 'attn',
-  k_proj: 'attn',
-  v_proj: 'attn',
-  q_rope: 'attn',
-  k_rope: 'attn',
-  attn_scores: 'attn',
-  attn_out: 'attn',
-  o_proj: 'attn',
-  post_attn: 'attn',
-  // FFN stages (per-layer)
-  ffn_normed: 'ffn',
-  ffn_in: 'ffn',
-  ffn_gate: 'ffn',
-  ffn_up: 'ffn',
-  ffn_act: 'ffn',
-  ffn_out: 'ffn',
-  layer_out: 'ffn',
-  // Final stages
-  pre_final_norm: 'logits',
-  final_norm: 'logits',
-  logits: 'logits',
-  logits_final: 'logits',
-};
-
-
-function matchesLayer(layers, layerIdx) {
-  if (!layers || layers.length === 0) return true;
-  if (layerIdx === undefined || layerIdx === null) return false;
-  return layers.includes(layerIdx);
-}
-
-
-function resolveTokens(tokens, numTokens) {
-  const raw = tokens && tokens.length > 0 ? tokens : [0];
-  
-  const resolved = [];
-  for (const t of raw) {
-    const idx = t < 0 ? numTokens + t : t;
-    if (idx >= 0 && idx < numTokens) {
-      resolved.push(idx);
-    }
-  }
-  return resolved;
-}
-
-
-function getTraceLogger(category, layerIdx) {
-  switch (category) {
-    case 'attn':
-      return ( message) =>  (trace.attn)(layerIdx ?? -1, message);
-    case 'ffn':
-      return ( message) =>  (trace.ffn)(layerIdx ?? -1, message);
-    case 'kv':
-      return ( message) =>  (trace.kv)(layerIdx ?? -1, message);
-    case 'loader':
-      return ( message) => trace.loader(message);
-    case 'kernels':
-      return ( message) => trace.kernels(message);
-    case 'logits':
-      return ( message) => trace.logits(message);
-    case 'embed':
-      return ( message) => trace.embed(message);
-    case 'sample':
-      return ( message) => trace.sample(message);
-    case 'buffers':
-      return ( message) => trace.buffers(message);
-    case 'perf':
-      return ( message) => trace.perf(message);
-    case 'all':
-    default: {
-      return ( message) => trace.embed(message);
-    }
-  }
-}
-
-
-export async function runProbes(stage, buffer, options) {
-  const { layerIdx, numTokens, hiddenSize, probes, recorder, dtype = 'f32' } = options;
-  if (!probes || probes.length === 0) return;
-  if (!buffer) return;
-  if (recorder) return;
-
-  const isCpuBuffer = buffer instanceof Float32Array;
-  const device = isCpuBuffer ? null : getDevice();
-  if (!isCpuBuffer && !device) return;
-
-  const stageProbes = probes.filter((probe) => probe.stage === stage);
-  if (stageProbes.length === 0) return;
-  if (!isCpuBuffer && !allowReadback(`probe.${stage}`)) return;
-
-  // Determine bytes per element based on dtype
-  const bytesPerElement = dtype === 'f16' ? 2 : 4;
-
-  for (const probe of stageProbes) {
-    if (!matchesLayer(probe.layers, layerIdx)) continue;
-
-    const dims = probe.dims ?? [];
-    if (dims.length === 0) continue;
-
-    const tokens = resolveTokens(probe.tokens, numTokens);
-    if (tokens.length === 0) continue;
-
-    const category = probe.category && probe.category !== 'all'
-      ? probe.category
-      : STAGE_DEFAULT_CATEGORY[stage];
-    const logger = getTraceLogger(category, layerIdx);
-    const probeId = probe.id ? ` ${probe.id}` : '';
-
-    for (const tokenIdx of tokens) {
-
-      const values = [];
-      for (const dimIdx of dims) {
-        if (dimIdx < 0 || dimIdx >= hiddenSize) {
-          values.push(`${dimIdx}=out_of_range`);
-          continue;
-        }
-        if (isCpuBuffer) {
-          const idx = tokenIdx * hiddenSize + dimIdx;
-          const value =  (buffer)[idx];
-          values.push(`${dimIdx}=${value.toFixed(4)}`);
-          continue;
-        }
-        const elementOffset = tokenIdx * hiddenSize + dimIdx;
-        const byteOffset = elementOffset * bytesPerElement;
-        // WebGPU requires offset and size to be multiples of 4
-        const alignedOffset = Math.floor(byteOffset / 4) * 4;
-        const offsetWithinRead = byteOffset - alignedOffset;
-        const readSize = 4; // Always read 4 bytes (aligned)
-        const readback = await readBufferSlice(buffer, alignedOffset, readSize);
-        let value;
-        if (dtype === 'f16') {
-          // offsetWithinRead is 0 or 2 for F16 - extract correct u16
-          const u16Array = new Uint16Array(readback);
-          const u16Index = offsetWithinRead / 2;
-          value = f16ToF32(u16Array[u16Index]);
-        } else {
-          value = new Float32Array(readback)[0];
-        }
-        values.push(`${dimIdx}=${value.toFixed(4)}`);
-      }
-
-      logger(`PROBE${probeId} stage=${stage} token=${tokenIdx} values=[${values.join(', ')}]`);
-    }
-  }
-}
-
-
-export function hasProbeStage(probes, stage, layerIdx) {
-  if (!probes || probes.length === 0) return false;
-  return probes.some((probe) => probe.stage === stage && matchesLayer(probe.layers, layerIdx));
-}

package/src/inference/pipelines/text/sampling.d.ts

@@ -1,54 +0,0 @@
-/**
- * Token sampling from logits with temperature, top-k, top-p, and repetition penalty.
- */
-
-export interface SamplingOptions {
-  temperature: number;
-  topP: number;
-  topK: number;
-  decode?: (tokens: number[]) => string;
-  debug?: boolean;
-  padTokenId?: number;
-}
-
-export interface TokenCandidate {
-  token: number;
-  prob: number;
-}
-
-export interface TopKResult {
-  token: number;
-  logit: number;
-  prob: number;
-  text: string;
-}
-
-export interface LogitStats {
-  min: number;
-  max: number;
-  nanCount: number;
-  infCount: number;
-  top5: TopKResult[];
-}
-
-export function applyRepetitionPenalty(
-  logits: Float32Array,
-  previousTokens: number[],
-  penalty: number
-): void;
-
-export function softmax(logits: Float32Array): Float32Array;
-
-export function sample(logits: Float32Array, opts: SamplingOptions): number;
-
-export function getTopK(
-  logits: Float32Array,
-  k?: number,
-  decode?: (tokens: number[]) => string
-): TopKResult[];
-
-export function logitsSanity(
-  logits: Float32Array,
-  label: string,
-  decode?: (tokens: number[]) => string
-): LogitStats;