@simulatte/doppler 0.1.8 → 0.1.9

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

@@ -28,6 +28,7 @@ import type {
   SpeculativeConfigSchema,
   KernelPathSchema,
 } from '../../../config/schema/index.js';
+import type { LoaderDebugConfigSchema } from '../../../config/schema/debug.schema.js';
 import type { KernelPathSource } from '../../../config/kernel-path-loader.js';
 
 export interface PipelineStorageContext {
@@ -206,6 +207,7 @@ export interface LoadWeightsOptions {
   resolvedKernelPath?: KernelPathSchema | null;
   kernelPathSource?: KernelPathSource;
   keepF32Weights?: boolean;
+  loaderDebug?: LoaderDebugConfigSchema | null;
 }
 
 /**

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

@@ -309,13 +309,21 @@ export async function initRoPEFrequencies(config, useGPU) {
   if (!Number.isFinite(ropeScale) || ropeScale <= 0) {
     throw new Error(`RoPE scale must be a positive number; got "${ropeScale}".`);
   }
-  const resolvedLocalScale = ropeLocalScale ?? ropeScale;
-  if (!Number.isFinite(resolvedLocalScale) || resolvedLocalScale <= 0) {
+  const resolvedLocalScale = ropeLocalScale;
+  if (resolvedLocalScale != null && (!Number.isFinite(resolvedLocalScale) || resolvedLocalScale <= 0)) {
     throw new Error(`Local RoPE scale must be a positive number; got "${resolvedLocalScale}".`);
   }
   const resolvedLocalTheta = ropeLocalTheta ?? ropeTheta;
-  const resolvedLocalScalingType = ropeLocalScalingType ?? ropeScalingType;
-  const resolvedLocalScaling = ropeLocalScaling ?? ropeScaling;
+  const resolvedLocalScalingType = (
+    ropeLocalScalingType === undefined
+      ? ropeScalingType
+      : ropeLocalScalingType
+  );
+  const resolvedLocalScaling = (
+    ropeLocalScalingType === undefined
+      ? ropeScaling
+      : ropeLocalScaling
+  );
   const resolvedRotaryDim = resolveRotaryDim(headDim, rotaryDim, partialRotaryFactor);
   const halfDim = resolvedRotaryDim / 2;
   if (mropeInterleaved === true && Array.isArray(mropeSection)) {
@@ -646,7 +654,12 @@ export async function initTokenizer(manifest, options = {}) {
 
 
 export async function loadWeights(manifest, modelConfig, options = {}) {
-  const { onProgress, loadingConfig, baseUrl } = options;
+  const {
+    onProgress,
+    loadingConfig,
+    baseUrl,
+    loaderDebug,
+  } = options;
   const runtimeStorageContext = options.storageContext
     ?? createRemoteStorageContext(baseUrl, manifest);
   const verifyHashes = (
@@ -668,6 +681,7 @@ export async function loadWeights(manifest, modelConfig, options = {}) {
       keepF32Weights
     )
   );
+  dopplerLoader.setLoaderDebugConfig(loaderDebug ?? null);
 
   const tensorsFile = isRDRRManifest(manifest) ? manifest.tensorsFile : null;
   if (baseUrl && tensorsFile) {

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

@@ -43,19 +43,16 @@ export function detectSandwichNorm(config) {
 }
 
 
-export function isMoELayer(layerIdx, config, layerWeights) {
+export function isMoELayer(layerIdx, config) {
   if (!config.useMoE) return false;
 
-  // Check if layer has router weights
-  if (layerWeights?.routerWeight) return true;
-
-  // Fall back to layer_types array if available
+  // Manifest-first: check layerTypes from config (derived from manifest.inference.layerPattern)
   const layerTypes = config.layerTypes;
   if (Array.isArray(layerTypes) && layerIdx < layerTypes.length) {
     return layerTypes[layerIdx] === 'moe';
   }
 
-  // Default: assume all layers are MoE if model uses MoE
+  // No layerTypes available: assume all layers are MoE
   return true;
 }
 
@@ -87,6 +84,11 @@ function assertSupportedLayerRuntime(layerIdx, config) {
   }
 }
 
+function getConvLayerState(convLayerStates, layerIdx) {
+  if (!convLayerStates) return {};
+  return convLayerStates.get(layerIdx) ?? {};
+}
+
 function isSlidingLayerType(layerType) {
   const normalized = normalizeLayerType(layerType);
   return normalized === 'sliding_attention'
@@ -103,6 +105,14 @@ function isConvLayerType(layerType) {
     || normalized === 'liv_convolution';
 }
 
+export function hasConvLayers(layerTypes) {
+  if (!Array.isArray(layerTypes)) return false;
+  for (let i = 0; i < layerTypes.length; i++) {
+    if (isConvLayerType(layerTypes[i])) return true;
+  }
+  return false;
+}
+
 function isLinearLayerType(layerType) {
   const normalized = normalizeLayerType(layerType);
   return normalized === 'linear_attention'
@@ -201,8 +211,22 @@ export async function processLayerGPU(layerIdx, inputBuffer, numTokens, isPrefil
       );
     }
     const convKernel = layerWeights?.convKernel ?? null;
+    // Apply input norm (operator_norm) before conv mixer — matches HF Lfm2 forward pass
+    let normedTensor = inputTensor;
+    const inputNormWeight = layerWeights?.inputNorm ?? null;
+    if (inputNormWeight) {
+      const normWeightBuf = getNormWeightBuffer(inputNormWeight, `L${layerIdx}.conv_input_norm`);
+      normedTensor = await doRMSNorm(inputTensor, normWeightBuf, rmsNormEps, {
+        batchSize: numTokens,
+        hiddenSize,
+        rmsNormWeightOffset: config.rmsNormWeightOffset,
+        label: `L${layerIdx}.conv_input_norm`,
+        layerIdx,
+      }, recorder);
+      if (!(inputNormWeight instanceof GPUBuffer)) releaseOrTrack(recorder, normWeightBuf);
+    }
     attnOutput = await doConv(
-      inputTensor,
+      normedTensor,
       getWeightBuffer(convInProj, `L${layerIdx}.conv_in_proj`),
       convKernel ? getWeightBuffer(convKernel, `L${layerIdx}.conv_kernel`) : null,
       getWeightBuffer(convOutProj, `L${layerIdx}.conv_out_proj`),
@@ -213,9 +237,13 @@ export async function processLayerGPU(layerIdx, inputBuffer, numTokens, isPrefil
         label: `L${layerIdx}.conv`,
         swigluLimit: config.swigluLimit,
         kernelPath: context.kernelPath ?? null,
+        convState: getConvLayerState(context.convLayerStates, layerIdx),
       },
       recorder
     );
+    if (normedTensor !== inputTensor) {
+      releaseOrTrack(recorder, normedTensor.buffer);
+    }
   } else if (isLinearLayer) {
     attnOutput = await runLinearAttentionLayer(inputTensor, layerWeights ?? null, {
       layerIdx,
@@ -721,6 +749,7 @@ async function processLayerPlanGPU(layerIdx, inputBuffer, numTokens, isPrefill,
               label: `L${layerIdx}.plan_conv`,
               swigluLimit: config.swigluLimit,
               kernelPath: context.kernelPath ?? null,
+              convState: getConvLayerState(context.convLayerStates, layerIdx),
             },
             recorder
           );
@@ -782,7 +811,7 @@ async function processLayerPlanGPU(layerIdx, inputBuffer, numTokens, isPrefill,
           let outputTensor;
           const { runMoEFFNGPU, runDenseFFNGPU } = await import('./ffn/index.js');
 
-          const canAutoMoe = config.useMoE && isMoELayer(layerIdx, config, layerWeights);
+          const canAutoMoe = config.useMoE && isMoELayer(layerIdx, config);
           const useMoe = selectRuleValue(
             'inference',
             'layer',

package/src/inference/pipelines/text/moe-gpu.js

@@ -23,6 +23,7 @@ import {
   validateMoeShape,
   resolveMoeVendorProfile,
   resolveGptOssKernelPathProfile,
+  resolveMixtralKernelPathProfile,
 } from './moe-shape-validator.js';
 
 export async function moeFeedForwardGPU(
@@ -52,7 +53,10 @@ export async function moeFeedForwardGPU(
   if (topK == null) {
     throw new Error('MoE topK is required in config.');
   }
-  const modelType = config.modelType ?? (expertFormat === 'gpt-oss' ? 'gpt-oss' : 'mixtral');
+  if (config.modelType == null) {
+    throw new Error('MoE config.modelType is required; got null/undefined.');
+  }
+  const modelType = config.modelType;
   validateMoeShape(
     { hiddenSize, intermediateSize, moeTopK: topK, numExperts, expertFormat },
     { modelType }
@@ -130,7 +134,13 @@ export async function moeFeedForwardGPU(
     trace.buffers(`MoE L${layerIdx} router_logits`, { min, max, nanCount, dtype: logitsDtype });
   }
 
+  // Profile resolution: routerTopK/dequantExpert are resolved for tracing and
+  // forward validation. Actual kernel dispatch uses the generic softmax.rules.json
+  // topkVariant rules (keyed by modelType) and format-specific dequant paths.
+  // GPT-OSS: dequantTileShape actively steers MXFP4 dequant; routerTopK is trace-only.
+  // Mixtral: expert weights are pre-loaded (no runtime dequant); both fields are trace-only.
   let gptOssKernelPathProfile = null;
+  let mixtralKernelPathProfile = null;
   if (modelType === 'gpt-oss') {
     gptOssKernelPathProfile = await resolveGptOssKernelPathProfile({
       hasF16: caps.hasF16,
@@ -141,6 +151,14 @@ export async function moeFeedForwardGPU(
       groupSize: 32,
       tileShape: vendorProfile.dequantTileShape,
     });
+  } else if (modelType === 'mixtral') {
+    mixtralKernelPathProfile = await resolveMixtralKernelPathProfile({
+      hasF16: caps.hasF16,
+      hasSubgroups: caps.hasSubgroups,
+      routerDtype: logitsDtype,
+      weightsDtype: activationDtype,
+      outputDtype: activationDtype,
+    });
   }
 
   stepStart = perfMark();
@@ -159,7 +177,7 @@ export async function moeFeedForwardGPU(
   perfLog(`MoE L${layerIdx} topk`, stepStart, {
     topK,
     modelType,
-    routerTopKKernel: gptOssKernelPathProfile?.routerTopK ?? null,
+    routerTopKKernel: gptOssKernelPathProfile?.routerTopK ?? mixtralKernelPathProfile?.routerTopK ?? null,
   });
 
   if (isTraceEnabled('buffers')) {
@@ -211,7 +229,7 @@ export async function moeFeedForwardGPU(
   const bytesPerElement = selectRuleValue('shared', 'dtype', 'bytesFromDtype', { dtype: activationDtype });
   const bytesPerToken = hiddenSize * bytesPerElement;
   let maxTokensPerExpert = resolveMaxTokensPerExpert(numTokens, numExperts, topK, hiddenSize, activationDtype);
-  if (modelType === 'gpt-oss') {
+  if (vendorProfile.maxTokensPerExpertScale !== 1.0) {
     maxTokensPerExpert = Math.max(
       1,
       Math.round(maxTokensPerExpert * vendorProfile.maxTokensPerExpertScale)

package/src/inference/pipelines/text/moe-shape-validator.d.ts

@@ -29,3 +29,12 @@ export interface GptOssKernelPathProfile {
 export declare function resolveGptOssKernelPathProfile(
   context: Record<string, unknown>
 ): Promise<GptOssKernelPathProfile>;
+
+export interface MixtralKernelPathProfile {
+  routerTopK: string;
+  dequantExpert: string;
+}
+
+export declare function resolveMixtralKernelPathProfile(
+  context: Record<string, unknown>
+): Promise<MixtralKernelPathProfile>;

package/src/inference/pipelines/text/moe-shape-validator.js

@@ -7,17 +7,15 @@ function asVendorString(caps) {
 }
 
 export function resolveMoeVendorProfile(modelType) {
-  if (modelType !== 'gpt-oss') {
-    return {
-      preferVec4Dequant: false,
-      dequantTileShape: 'scalar',
-      routerWorkgroupSize: 128,
-      maxTokensPerExpertScale: 1.0,
-    };
-  }
   const caps = getKernelCapabilities();
   const vendor = asVendorString(caps);
-  return selectRuleValue('kernels', 'moeGptoss', 'vendorQuirkProfile', { vendor });
+  if (modelType === 'gpt-oss') {
+    return selectRuleValue('kernels', 'moeGptoss', 'vendorQuirkProfile', { vendor });
+  }
+  if (modelType === 'mixtral') {
+    return selectRuleValue('kernels', 'moeMixtral', 'vendorQuirkProfile', { vendor });
+  }
+  throw new Error(`[MoE] Unknown modelType "${modelType}" for vendor profile resolution.`);
 }
 
 function resolveGptOssRuleContext(context) {
@@ -41,6 +39,25 @@ export async function resolveGptOssKernelPathProfile(context) {
   };
 }
 
+function resolveMixtralRuleContext(context) {
+  return {
+    modelType: 'mixtral',
+    hasF16: context?.hasF16,
+    hasSubgroups: context?.hasSubgroups,
+    routerDtype: context?.routerDtype ?? 'f32',
+    weightsDtype: context?.weightsDtype,
+    outputDtype: context?.outputDtype ?? context?.weightsDtype,
+  };
+}
+
+export async function resolveMixtralKernelPathProfile(context) {
+  const ruleContext = resolveMixtralRuleContext(context);
+  return {
+    routerTopK: selectRuleValue('kernels', 'moeMixtral', 'routerTopKVariant', ruleContext),
+    dequantExpert: selectRuleValue('kernels', 'moeMixtral', 'dequantVariant', ruleContext),
+  };
+}
+
 export function validateMoeShape(config, options = {}) {
   const {
     hiddenSize,
@@ -66,8 +83,11 @@ export function validateMoeShape(config, options = {}) {
 
   if (modelType === 'gpt-oss') {
     const policy = selectRuleValue('kernels', 'moeGptoss', 'shapePolicy', { modelType });
-    const hiddenDivisor = policy.hiddenSizeDivisor ?? 32;
-    const intermediateDivisor = policy.intermediateSizeDivisor ?? 32;
+    if (policy.hiddenSizeDivisor == null || policy.intermediateSizeDivisor == null) {
+      throw new Error('[MoE] GPT-OSS shapePolicy is missing hiddenSizeDivisor or intermediateSizeDivisor.');
+    }
+    const hiddenDivisor = policy.hiddenSizeDivisor;
+    const intermediateDivisor = policy.intermediateSizeDivisor;
     if (hiddenSize % hiddenDivisor !== 0 || intermediateSize % intermediateDivisor !== 0) {
       throw new Error(
         `[MoE] GPT-OSS shape policy violation: hiddenSize (${hiddenSize}) % ${hiddenDivisor} = ${hiddenSize % hiddenDivisor}, ` +

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/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);
   }