@simulatte/doppler 0.1.7 → 0.1.9

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

@@ -14,13 +14,14 @@ import {
   recordCastF32ToF16,
 } from '../../../gpu/kernels/cast.js';
 import { createTensor } from '../../../gpu/tensor.js';
-import { releaseBuffer } from '../../../memory/buffer-pool.js';
+import { releaseBuffer, readBuffer, acquireBuffer, uploadData } 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';
+import { runGatedShortConvGPU } from '../../../gpu/kernels/gated-short-conv.js';
 
 
 export function isDecodeBuffer(decodeBuffers, buffer) {
@@ -174,17 +175,22 @@ export async function doConv(
     throw new Error('doConv requires hiddenSize > 0.');
   }
 
-  // Use the first 2x hidden projection channels as a gated conv-state projection.
+  // LFM2 gated short convolution (GPU-native):
+  // in_proj → 3×hidden → GPU kernel: split(B,C,x) + B*x + causal conv1d + C*conv_out → out_proj
   let inProj = null;
-  let activated = null;
-  let convInput = null;
+  let convOut = null;
   let outProj = null;
   try {
+    const convState = options.convState;
+    const hasConvState = Boolean(convState?.convWeightGPU && convState?.convStateGPU);
+    const projN = hasConvState ? hiddenSize * 3 : hiddenSize * 2;
+
+    // Project input
     inProj = await doMatmul(
       inputTensor,
       convInProj,
       numTokens,
-      hiddenSize * 2,
+      projN,
       hiddenSize,
       {
         transposeB: 'auto',
@@ -195,50 +201,32 @@ export async function doConv(
       },
       recorder
     );
-    activated = await doSiLURowSplit(inProj, {
-      numTokens,
-      dim: hiddenSize,
-      activation: 'silu',
-      swigluLimit: options.swigluLimit ?? null,
-      label: `${label}.activation`,
-      layerIdx,
-    }, recorder);
+
+    if (hasConvState) {
+      // GPU gated short conv kernel: B*x → conv1d → C*conv_out (all on GPU)
+      convOut = await runGatedShortConvGPU(inProj, convState, {
+        numTokens,
+        layerIdx,
+        recorder,
+      });
+    } else {
+      // SwiGLU gated activation fallback: silu(first_half) * second_half
+      convOut = 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;
-    }
-
+    // Output projection
     outProj = await doMatmul(
-      convInput,
+      convOut,
       convOutProj,
       numTokens,
       hiddenSize,
@@ -253,13 +241,8 @@ export async function doConv(
       recorder
     );
 
-    if (convInput && (!activated || convInput.buffer !== activated.buffer)) {
-      releaseOrTrack(recorder, convInput.buffer);
-      convInput = null;
-    } else if (activated) {
-      releaseOrTrack(recorder, activated.buffer);
-      activated = null;
-    }
+    releaseOrTrack(recorder, convOut.buffer);
+    convOut = null;
 
     if (kernelTrace.enabled && !recorder) {
       await traceStep('conv', label, layerIdx, outProj.buffer, [numTokens, hiddenSize]);
@@ -267,13 +250,100 @@ export async function doConv(
     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 (convOut) releaseOrTrack(recorder, convOut.buffer);
     if (inProj) releaseOrTrack(recorder, inProj.buffer);
     throw error;
   }
 }
 
+export async function initConvLayerState(convState, convKernel, convInProj, hiddenSize, label, layerIdx) {
+  const { isWeightBuffer } = await import('../../../gpu/weight-buffer.js');
+  const isWB = typeof isWeightBuffer === 'function' && isWeightBuffer(convKernel);
+  const kernelBuf = isWB ? convKernel.buffer : (convKernel instanceof GPUBuffer ? convKernel : convKernel.buffer ?? convKernel);
+  const kernelDtype = isWB ? String(convKernel.dtype ?? '').toLowerCase() : null;
+
+  // Determine kernel size from weight shape
+  let kernelSize = 3;
+  if (isWB && Array.isArray(convKernel.shape)) {
+    kernelSize = Number(convKernel.shape[convKernel.shape.length - 1]) || 3;
+  }
+
+  // Dequantize conv kernel weights to F32
+  const totalElements = hiddenSize * kernelSize;
+  const { QK_K, Q4K_BLOCK_BYTES } = await import('../../../config/schema/index.js');
+  const { dequantizeQ4KM } = await import('../../../converter/quantizer.js');
+  const { getDevice } = await import('../../../gpu/device.js');
+  const device = getDevice();
+
+  const isQ4K = kernelDtype === 'q4k' || kernelDtype === 'q4_k_m' || kernelDtype === 'q4_k';
+  let weightF32;
+
+  if (isQ4K) {
+    const numBlocks = Math.ceil(totalElements / QK_K);
+    const q4kBytes = numBlocks * Q4K_BLOCK_BYTES;
+    // GPU readBuffer returns zeros for some Q4K weight buffers, so prefer
+    // CPU-side rawBytes from the WeightBuffer when available.
+    const hasRawBytes = isWB && convKernel.rawBytes;
+    if (hasRawBytes) {
+      weightF32 = dequantizeQ4KM(new Uint8Array(convKernel.rawBytes), numBlocks, [totalElements]);
+    } else {
+      if (device) await device.queue.onSubmittedWorkDone();
+      const raw = await readBuffer(kernelBuf, q4kBytes);
+      weightF32 = dequantizeQ4KM(new Uint8Array(raw), numBlocks, [totalElements]);
+    }
+  } else if (kernelDtype === 'f16' || kernelDtype === 'bf16') {
+    if (device) await device.queue.onSubmittedWorkDone();
+    const raw = await readBuffer(kernelBuf, totalElements * 2);
+    const { decodeReadback } = await import('./debug-utils/index.js');
+    weightF32 = decodeReadback(raw, 'f16');
+  } else {
+    if (device) await device.queue.onSubmittedWorkDone();
+    const raw = await readBuffer(kernelBuf, totalElements * 4);
+    weightF32 = new Float32Array(raw);
+  }
+
+  // Validate dequantized weights are non-degenerate
+  let maxAbs = 0;
+  for (let i = 0; i < weightF32.length; i++) {
+    const abs = Math.abs(weightF32[i]);
+    if (abs > maxAbs) maxAbs = abs;
+  }
+  if (maxAbs === 0) {
+    const { log } = await import('../../../debug/index.js');
+    log.error('Pipeline', `${label} conv kernel weights are all zeros after dequantization (dtype=${kernelDtype}, elements=${totalElements}). Conv layers will produce degenerate output.`);
+  }
+
+  // Upload dequantized weights to GPU
+  const weightGPU = acquireBuffer(weightF32.byteLength, undefined, `${label}.conv_weight_f32`);
+  uploadData(weightGPU, weightF32);
+
+  // Create zeroed conv state buffer
+  const stateSize = hiddenSize * (kernelSize - 1) * Float32Array.BYTES_PER_ELEMENT;
+  const stateGPU = acquireBuffer(stateSize, undefined, `${label}.conv_state`);
+  uploadData(stateGPU, new Float32Array(hiddenSize * (kernelSize - 1)));
+
+  convState.convWeightGPU = weightGPU;
+  convState.convStateGPU = stateGPU;
+  convState.hiddenSize = hiddenSize;
+  convState.kernelSize = kernelSize;
+
+  // Pre-dequantize in_proj weight to F32 via CPU dequantization of the raw Q4K buffer.
+  // GPU readBuffer returns zeros for some Q4K weight buffers, so we dequantize from the
+  // WeightBuffer's raw bytes instead.
+  if (isWB && isWeightBuffer(convInProj)) {
+    const inProjDtype = String(convInProj.dtype ?? '').toLowerCase();
+    const isInProjQ4K = inProjDtype === 'q4k' || inProjDtype === 'q4_k_m' || inProjDtype === 'q4_k';
+    if (isInProjQ4K && convInProj.rawBytes) {
+      const inProjElements = hiddenSize * 3 * hiddenSize;
+      const inProjBlocks = Math.ceil(inProjElements / QK_K);
+      const inProjF32 = dequantizeQ4KM(new Uint8Array(convInProj.rawBytes), inProjBlocks, [inProjElements]);
+      const inProjGPU = acquireBuffer(inProjF32.byteLength, undefined, `${label}.in_proj_f32`);
+      uploadData(inProjGPU, inProjF32);
+      convState.inProjF32GPU = inProjGPU;
+    }
+  }
+}
+
 export async function doCast(input, toDtype, recorder) {
   if (toDtype !== 'f16' && toDtype !== 'f32') {
     throw new Error(`Unsupported cast target dtype "${toDtype}"`);

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

@@ -11,6 +11,7 @@ const STAGE_DEFAULT_CATEGORY = {
   embed_out: 'embed',
   // Attention stages (per-layer)
   attn_input: 'attn',
+  post_input_norm: 'attn',
   attn_normed: 'attn',
   linear_qkv_proj: 'attn',
   linear_z_proj: 'attn',

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

@@ -58,6 +58,30 @@ export function softmax(logits) {
   return exps;
 }
 
+function countFiniteCandidates(logits, padTokenId) {
+  let finiteCandidateCount = 0;
+  for (let i = 0; i < logits.length; i++) {
+    if (padTokenId != null && i === padTokenId) {
+      continue;
+    }
+    if (Number.isFinite(logits[i])) {
+      finiteCandidateCount += 1;
+    }
+  }
+  return finiteCandidateCount;
+}
+
+function assertFiniteSamplingCandidates(logits, padTokenId, label) {
+  const finiteCandidateCount = countFiniteCandidates(logits, padTokenId);
+  if (finiteCandidateCount > 0) {
+    return;
+  }
+  throw new Error(
+    `[Sampling] ${label} has no finite candidate logits after masking the pad token. ` +
+    'Upstream decode likely produced NaN/Inf or an all-masked distribution.'
+  );
+}
+
 
 export function sample(logits, opts) {
   const { temperature, topP, topK, decode, debug = false, padTokenId, seed } = opts;
@@ -66,16 +90,28 @@ export function sample(logits, opts) {
     logits[padTokenId] = -Infinity;
   }
 
+  assertFiniteSamplingCandidates(logits, padTokenId, 'Logits');
+
   // Greedy (argmax) when temperature = 0
   if (temperature === 0) {
-    let maxIdx = 0;
-    let maxVal = logits[0];
-    for (let i = 1; i < logits.length; i++) {
-      if (logits[i] > maxVal) {
-        maxVal = logits[i];
+    let maxIdx = -1;
+    let maxVal = -Infinity;
+    for (let i = 0; i < logits.length; i++) {
+      const value = logits[i];
+      if (!Number.isFinite(value)) {
+        continue;
+      }
+      if (value > maxVal) {
+        maxVal = value;
         maxIdx = i;
       }
     }
+    if (maxIdx < 0) {
+      throw new Error(
+        '[Sampling] Greedy sampling could not find a finite candidate logit. ' +
+        'Upstream decode likely produced NaN/Inf.'
+      );
+    }
     if (debug) {
       const text = decode?.([maxIdx]) ?? '?';
       trace.sample(`Greedy: id=${maxIdx} "${text}" logit=${maxVal.toFixed(4)}`);
@@ -96,7 +132,17 @@ export function sample(logits, opts) {
 
   let candidates = [];
   for (let i = 0; i < probs.length; i++) {
-    candidates.push({ token: i, prob: probs[i] });
+    const probability = probs[i];
+    if (!Number.isFinite(probability) || probability <= 0) {
+      continue;
+    }
+    candidates.push({ token: i, prob: probability });
+  }
+  if (candidates.length === 0) {
+    throw new Error(
+      '[Sampling] Softmax produced no finite candidate probabilities. ' +
+      'Upstream decode likely produced NaN/Inf logits.'
+    );
   }
   candidates.sort((a, b) => b.prob - a.prob);
 

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

@@ -15,6 +15,8 @@ export class PipelineState {
       layers: new Map(),
     };
 
+    this.convLayerStates = new Map();
+
     this.moeRouter = null;
 
     this.speculativeDecoder = null;

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

@@ -69,6 +69,11 @@ export declare class InferencePipeline extends PipelineState {
   // ==========================================================================
 
   generate(prompt: PromptInput, options?: GenerateOptions): AsyncGenerator<string, void, void>;
+  generateTokens(prompt: PromptInput, options?: GenerateOptions): AsyncGenerator<number, void, void>;
+  generateTokenIds(
+    prompt: PromptInput,
+    options?: GenerateOptions
+  ): Promise<{ tokenIds: number[]; stats: PipelineStats }>;
 
   decodeStepLogits(currentIds: number[], options?: GenerateOptions): Promise<LogitsStepResult>;
 

package/src/inference/pipelines/text.js

@@ -43,6 +43,7 @@ import {
 import { getDopplerLoader } from '../../loader/doppler-loader.js';
 import { registerPipeline, getPipelineFactory } from './registry.js';
 import { selectRuleValue } from '../../rules/rule-registry.js';
+import { initConvLayerState } from './text/ops.js';
 
 function destroyMoERouter(router) {
   if (router && typeof router.destroy === 'function') {
@@ -221,6 +222,9 @@ export class InferencePipeline extends PipelineState {
     // Initialize RoPE frequencies
     await this._initRoPE();
 
+    // Initialize conv layer states for gated short conv layers (LFM2)
+    await this._initConvLayerStates();
+
     this.isLoaded = true;
     log.info('Pipeline', 'Model loaded successfully');
   }
@@ -237,6 +241,7 @@ export class InferencePipeline extends PipelineState {
         resolvedKernelPath: this.resolvedKernelPath,
         kernelPathSource: this.kernelPathSource,
         keepF32Weights: this.runtimeConfig.inference.compute.keepF32Weights === true,
+        loaderDebug: this.runtimeConfig?.shared?.debug?.loader ?? null,
         onProgress: (info) => {
           if (info.stage !== 'layers' && info.stage !== 'shards') {
             log.verbose('Loader', `${info.stage}: ${Math.round(info.progress * 100)}%${info.message ? ` - ${info.message}` : ''}`);
@@ -310,7 +315,7 @@ export class InferencePipeline extends PipelineState {
       maxSeqLen,
       ropeTheta: config.ropeTheta,
       ropeLocalTheta: config.ropeLocalTheta,
-      mropeInterleaved: config.ropeInterleaved,
+      mropeInterleaved: config.mropeInterleaved,
       mropeSection: config.mropeSection,
       partialRotaryFactor: config.partialRotaryFactor,
       ropeScale: config.ropeScale,
@@ -327,6 +332,51 @@ export class InferencePipeline extends PipelineState {
   }
 
 
+  async _initConvLayerStates() {
+    const config = this.modelConfig;
+    if (!config?.layerTypes) return;
+    const { getDevice } = await import('../../gpu/device.js');
+    const device = getDevice();
+    if (!device) return;
+
+    const hiddenSize = config.hiddenSize;
+    const convStates = new Map();
+
+    for (let i = 0; i < config.layerTypes.length; i++) {
+      const lt = String(config.layerTypes[i] ?? '').toLowerCase();
+      if (lt !== 'conv' && lt !== 'convolution') continue;
+
+      const layerWeights = this.weights.get(`layer_${i}`);
+      if (!layerWeights) continue;
+      const convKernel = layerWeights?.convKernel;
+      if (!convKernel) continue;
+
+      const convState = {};
+      try {
+        await initConvLayerState(
+          convState,
+          convKernel,
+          layerWeights.convInProj ?? null,
+          hiddenSize,
+          `L${i}.conv`,
+          i
+        );
+        if (!convState.convWeightGPU || !convState.convStateGPU) {
+          continue;
+        }
+        convStates.set(i, convState);
+      } catch (e) {
+        log.warn('Pipeline', `Conv layer ${i} state init failed: ${e.message}`);
+      }
+    }
+
+    if (convStates.size > 0) {
+      this.convLayerStates = convStates;
+      log.info('Pipeline', `Initialized ${convStates.size} conv layer states (kernelSize=${convStates.values().next().value?.kernelSize})`);
+    }
+  }
+
+
   _resolveLayerPipeline() {
     if (!this.modelConfig) return;
     const runtimePlan = this.runtimeConfig.inference.pipeline ?? null;
@@ -349,6 +399,14 @@ export class InferencePipeline extends PipelineState {
     return this.generator.generate(prompt, options);
   }
 
+  generateTokens(prompt, options = {}) {
+    return this.generator.generateTokens(prompt, options);
+  }
+
+  generateTokenIds(prompt, options = {}) {
+    return this.generator.generateTokenIds(prompt, options);
+  }
+
   decodeStepLogits(currentIds, options = {}) {
     return this.generator.decodeStepLogits(currentIds, options);
   }