@simulatte/doppler 0.1.6 → 0.1.8

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

@@ -122,6 +122,20 @@ function resolveTokenText(tokenizer, tokenIds, fallbackText = '?', renderTokenTe
   return fallbackText;
 }
 
+export function shouldRetryWithFinitenessFallback(error) {
+  if (error?.name === 'FinitenessError') {
+    return true;
+  }
+  const message = typeof error?.message === 'string'
+    ? error.message
+    : (typeof error === 'string' ? error : '');
+  if (!message.startsWith('[Sampling]')) {
+    return false;
+  }
+  return message.includes('no finite candidate logits after masking the pad token')
+    || message.includes('Softmax produced no finite candidate probabilities');
+}
+
 export class PipelineGenerator {
 
   #state;
@@ -351,7 +365,7 @@ export class PipelineGenerator {
       try {
         prefillLogits = await this._prefill(inputIds, opts);
       } catch (error) {
-        if (error.name === 'FinitenessError') {
+        if (shouldRetryWithFinitenessFallback(error)) {
           log.warn('Pipeline', `FinitenessGuard caught NaN/Inf during prefill. Retrying with F32 precision.`);
           prefillLogits = await this._retryWithFinitenessFallback(
             opts,
@@ -395,13 +409,34 @@ export class PipelineGenerator {
         log.debug('Pipeline', `After rep penalty top-5: ${topAfterPenalty.map(t => `"${t.text}"(${(t.prob * 100).toFixed(1)}%)`).join(', ')}`);
       }
 
-      const firstToken = sample(prefillLogits, {
-        temperature: opts.temperature,
-        topP: opts.topP,
-        topK: opts.topK,
-        padTokenId,
-        seed: opts.seed,
-      });
+      let firstToken;
+      try {
+        firstToken = sample(prefillLogits, {
+          temperature: opts.temperature,
+          topP: opts.topP,
+          topK: opts.topK,
+          padTokenId,
+          seed: opts.seed,
+        });
+      } catch (error) {
+        if (!shouldRetryWithFinitenessFallback(error)) {
+          throw error;
+        }
+        log.warn('Pipeline', 'FinitenessGuard caught non-finite prefill logits at sampling. Retrying with F32 precision.');
+        prefillLogits = await this._retryWithFinitenessFallback(
+          opts,
+          'prefill-sample',
+          () => this._prefill(inputIds, opts)
+        );
+        applyRepetitionPenalty(prefillLogits, generatedIds, opts.repetitionPenalty);
+        firstToken = sample(prefillLogits, {
+          temperature: opts.temperature,
+          topP: opts.topP,
+          topK: opts.topK,
+          padTokenId,
+          seed: opts.seed,
+        });
+      }
 
       if (opts.debug) {
         const firstTokenText = resolveTokenText(this.#state.tokenizer, [firstToken], `[${firstToken}]`, (tokens) => this.#state.tokenizer?.decode?.(tokens, true, false));
@@ -479,7 +514,7 @@ export class PipelineGenerator {
     try {
       prefillResult = await this._prefillToHidden(inputIds, opts);
     } catch (error) {
-      if (error.name === 'FinitenessError') {
+      if (shouldRetryWithFinitenessFallback(error)) {
         log.warn('Pipeline', `FinitenessGuard caught NaN/Inf during prefillKVOnly. Retrying with F32 precision.`);
         prefillResult = await this._retryWithFinitenessFallback(
           opts,
@@ -544,7 +579,7 @@ export class PipelineGenerator {
     try {
       prefillResult = await this._prefillToHidden(inputIds, opts);
     } catch (error) {
-      if (error.name === 'FinitenessError') {
+      if (shouldRetryWithFinitenessFallback(error)) {
         log.warn('Pipeline', `FinitenessGuard caught NaN/Inf during prefillWithEmbedding. Retrying with F32 precision.`);
         prefillResult = await this._retryWithFinitenessFallback(
           opts,
@@ -833,7 +868,7 @@ export class PipelineGenerator {
           try {
             nextToken = await this._decodeStep(generatedIds, opts);
           } catch (singleTokenError) {
-            if (singleTokenError.name === 'FinitenessError') {
+            if (shouldRetryWithFinitenessFallback(singleTokenError)) {
               log.warn('Pipeline', `FinitenessGuard caught NaN/Inf at batch step ${tokensGenerated}. Truncating KV cache and retrying token with F32 precision.`);
               nextToken = await this._retryDecodeStepWithFinitenessWindow(
                 generatedIds,
@@ -858,7 +893,7 @@ export class PipelineGenerator {
         try {
           nextToken = await this._decodeStep(generatedIds, opts);
         } catch (error) {
-          if (error.name === 'FinitenessError') {
+          if (shouldRetryWithFinitenessFallback(error)) {
             log.warn('Pipeline', `FinitenessGuard caught NaN/Inf at step ${tokensGenerated}. Truncating KV cache and retrying token with F32 precision.`);
             nextToken = await this._retryDecodeStepWithFinitenessWindow(
               generatedIds,
@@ -918,11 +953,9 @@ export class PipelineGenerator {
       throw new Error('Embed buffer not found or not a supported buffer type');
     }
     const embedBuffer = isWeightBuffer(embedBufferRaw) ? embedBufferRaw.buffer : embedBufferRaw;
-    const embedDtype = isWeightBuffer(embedBufferRaw)
-      ? getWeightDtype(embedBufferRaw)
-      : isCpuWeightBuffer(embedBufferRaw)
-        ? embedBufferRaw.dtype
-        : null;
+    const embedDtype = isCpuWeightBuffer(embedBufferRaw)
+      ? embedBufferRaw.dtype
+      : getWeightDtype(embedBufferRaw);
     if (opts.debug) {
       const embedSize = embedBuffer instanceof GPUBuffer ? embedBuffer.size : 'N/A';
       log.debug('Pipeline', `Embed buffer: type=${embedBuffer?.constructor?.name}, size=${embedSize}, dtype=${embedDtype}`);
@@ -1043,18 +1076,9 @@ export class PipelineGenerator {
           if (allowReadback(`pipeline.prefill.layer-${l}`)) {
             try {
               const sampleSize = config.hiddenSize * activationBytes;
-              const staging = device.createBuffer({
-                size: sampleSize,
-                usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
-              });
-              const enc = device.createCommandEncoder();
               const lastTokenOffset = (numTokens - 1) * config.hiddenSize * activationBytes;
-              enc.copyBufferToBuffer(currentHiddenBuffer, lastTokenOffset, staging, 0, sampleSize);
-              device.queue.submit([enc.finish()]);
-              await staging.mapAsync(GPUMapMode.READ);
-              const data = decodeReadback(staging.getMappedRange().slice(0), activationDtype);
-              staging.unmap();
-              staging.destroy();
+              const readback = await readBufferSlice(currentHiddenBuffer, lastTokenOffset, sampleSize);
+              const data = decodeReadback(readback, activationDtype);
               let min = Infinity;
               let max = -Infinity;
               let maxAbs = 0;
@@ -1112,20 +1136,12 @@ export class PipelineGenerator {
     if (opts.debug) {
       log.debug('Pipeline', `LAYER_LOOP_DONE, currentHiddenBuffer type=${currentHiddenBuffer?.constructor?.name}`);
       if (currentHiddenBuffer && allowReadback('pipeline.prefill.final-hidden')) {
-        const device = getDevice();
         const lastTokenOffset = (numTokens - 1) * config.hiddenSize * activationBytes;
         const sampleSize = config.hiddenSize * activationBytes;
-        const staging = device.createBuffer({
-          size: sampleSize,
-          usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
-        });
-        const enc = device.createCommandEncoder();
-        enc.copyBufferToBuffer(currentHiddenBuffer, lastTokenOffset, staging, 0, sampleSize);
-        device.queue.submit([enc.finish()]);
-        await staging.mapAsync(GPUMapMode.READ);
-        const data = decodeReadback(staging.getMappedRange().slice(0), activationDtype);
-        staging.unmap();
-        staging.destroy();
+        const data = decodeReadback(
+          await readBufferSlice(currentHiddenBuffer, lastTokenOffset, sampleSize),
+          activationDtype
+        );
         const nanCount = Array.from(data).filter(x => !Number.isFinite(x)).length;
         const nonZero = Array.from(data).filter(x => Number.isFinite(x) && x !== 0).slice(0, 5);
         log.debug('Pipeline', `FINAL_HIDDEN[pos=${numTokens - 1}]: nan=${nanCount}/${data.length}, sample=[${nonZero.map(x => x.toFixed(4)).join(', ')}]`);

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

@@ -190,6 +190,12 @@ export interface WeightLoadResult {
   layerRouterWeights: Map<number, RouterWeights>;
 }
 
+export interface ResolvedQ4KConfig {
+  useFusedQ4K: boolean;
+  q4kLayout: 'row' | 'col' | null;
+  keepF32Weights: boolean;
+}
+
 /** Options for loadWeights */
 export interface LoadWeightsOptions {
   storageContext?: PipelineStorageContext;
@@ -211,6 +217,13 @@ export function loadWeights(
   options?: LoadWeightsOptions
 ): Promise<WeightLoadResult>;
 
+export function resolveQ4KConfig(
+  manifest: Manifest,
+  kernelPath?: KernelPathSchema | null,
+  kernelPathSource?: KernelPathSource,
+  keepF32Weights?: boolean
+): ResolvedQ4KConfig;
+
 /**
  * Apply Gemma chat template to a prompt.
  */

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

@@ -2,7 +2,7 @@
 
 import { parseModelConfig } from './config.js';
 import { getDevice, getDeviceLimits, getKernelCapabilities } from '../../../gpu/device.js';
-import { acquireBuffer } from '../../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../../memory/buffer-pool.js';
 import { KVCache, SlidingWindowKVCache, TieredKVCache, BasisDecomposedPagedCache } from '../../kv-cache.js';
 import { Tokenizer } from '../../tokenizer.js';
 import { MoERouter } from '../../moe-router.js';
@@ -11,9 +11,13 @@ import { getDopplerLoader } from '../../../loader/doppler-loader.js';
 import { log, setGPUDevice, trace as debugTrace } from '../../../debug/index.js';
 import { getRuntimeConfig } from '../../../config/runtime.js';
 import { PAGED_LAYOUT_SEQ_LEN_THRESHOLD } from '../../../config/schema/index.js';
-import { isKernelPathFusedQ4K } from '../../../config/kernel-path-loader.js';
+import { isKernelPathFusedQ4K, kernelPathRequiresF32MatmulWeights } from '../../../config/kernel-path-loader.js';
 import { createWeightBuffer, getWeightDtype, isWeightBuffer } from '../../../gpu/weight-buffer.js';
 import { selectRuleValue } from '../../../rules/rule-registry.js';
+import {
+  createSourceStorageContext,
+  getSourceRuntimeMetadata,
+} from '../../../tooling/source-runtime-bundle.js';
 
 function resolveErrorMessage(error) {
   if (error && typeof error === 'object' && typeof error.message === 'string') {
@@ -56,12 +60,61 @@ function normalizeBaseUrl(baseUrl) {
   return baseUrl.replace(/\/$/, '');
 }
 
+async function fetchBytes(url, offset = null, length = null) {
+  const headers = {};
+  if (Number.isFinite(offset) && Number.isFinite(length) && length > 0) {
+    const start = Math.max(0, Math.floor(offset));
+    const end = start + Math.max(0, Math.floor(length)) - 1;
+    headers.Range = `bytes=${start}-${end}`;
+  }
+  const response = await fetch(url, { headers });
+  if (!response.ok) {
+    throw new Error(`Failed to fetch ${url}: ${response.status}`);
+  }
+  return new Uint8Array(await response.arrayBuffer());
+}
+
 function createRemoteStorageContext(baseUrl, manifest) {
   const root = normalizeBaseUrl(baseUrl);
   if (!root || !isRDRRManifest(manifest)) {
     return null;
   }
 
+  const sourceRuntime = getSourceRuntimeMetadata(manifest);
+  if (sourceRuntime) {
+    const readRange = async (relativePath, offset, length) => {
+      const filename = String(relativePath || '').replace(/^\/+/, '');
+      if (!filename) {
+        throw new Error('Direct-source artifact path is required.');
+      }
+      const url = `${root}/${filename}`;
+      return fetchBytes(url, offset, length);
+    };
+    const readText = async (relativePath) => {
+      const filename = String(relativePath || '').replace(/^\/+/, '');
+      if (!filename) return null;
+      const response = await fetch(`${root}/${filename}`);
+      if (!response.ok) {
+        throw new Error(`Failed to fetch ${filename} from ${root}: ${response.status}`);
+      }
+      return response.text();
+    };
+    const readBinary = async (relativePath) => {
+      const filename = String(relativePath || '').replace(/^\/+/, '');
+      if (!filename) {
+        throw new Error('Direct-source binary asset path is required.');
+      }
+      return fetchBytes(`${root}/${filename}`);
+    };
+    return createSourceStorageContext({
+      manifest,
+      readRange,
+      readText,
+      readBinary,
+      verifyHashes: true,
+    });
+  }
+
   return {
     async loadShard(index) {
       const shard = manifest.shards[index];
@@ -69,17 +122,13 @@ function createRemoteStorageContext(baseUrl, manifest) {
       if (!filename) {
         throw new Error(`Manifest shard ${index} is missing filename.`);
       }
-      const response = await fetch(`${root}/${filename.replace(/^\/+/, '')}`);
-      if (!response.ok) {
-        throw new Error(`Failed to fetch shard ${index} from ${root}: ${response.status}`);
-      }
-      return new Uint8Array(await response.arrayBuffer());
+      return fetchBytes(`${root}/${filename.replace(/^\/+/, '')}`);
     },
   };
 }
 
 
-function resolveQ4KConfig(
+export function resolveQ4KConfig(
   manifest,
   kernelPath,
   kernelPathSource = 'none',
@@ -101,18 +150,23 @@ function resolveQ4KConfig(
     );
   }
   let useFused = kernelPath ? isKernelPathFusedQ4K(kernelPath) : hasSubgroups;
+  const kernelPathKeepsF32Weights = kernelPathRequiresF32MatmulWeights(kernelPath);
   if (q4kLayout === 'col') {
     useFused = false;
   }
+  const resolvedKeepF32Weights = keepF32Weights || kernelPathKeepsF32Weights;
 
   const pathLabel = kernelPath?.id ?? 'auto';
   const layoutLabel = q4kLayout ?? 'none';
-  debugTrace.loader(`Q4K config: fused=${useFused}, kernelPath=${pathLabel}, source=${kernelPathSource}, layout=${layoutLabel}, subgroups=${hasSubgroups}`);
+  debugTrace.loader(
+    `Q4K config: fused=${useFused}, kernelPath=${pathLabel}, source=${kernelPathSource}, ` +
+    `layout=${layoutLabel}, keepF32Weights=${resolvedKeepF32Weights}, subgroups=${hasSubgroups}`
+  );
 
   return {
     useFusedQ4K: useFused,
     q4kLayout,
-    keepF32Weights,
+    keepF32Weights: resolvedKeepF32Weights,
   };
 }
 
@@ -326,20 +380,29 @@ export async function initRoPEFrequencies(config, useGPU) {
   // Upload to GPU if available
   const device = getDevice();
   if (device && useGPU) {
-    const cosBuffer = acquireBuffer(globalFreqs.cos.byteLength, undefined, 'rope_cos');
-    const sinBuffer = acquireBuffer(globalFreqs.sin.byteLength, undefined, 'rope_sin');
-    device.queue.writeBuffer(cosBuffer, 0, globalFreqs.cos.buffer, globalFreqs.cos.byteOffset, globalFreqs.cos.byteLength);
-    device.queue.writeBuffer(sinBuffer, 0, globalFreqs.sin.buffer, globalFreqs.sin.byteOffset, globalFreqs.sin.byteLength);
-
-    
-    let localCosBuffer;
-    
-    let localSinBuffer;
-    if (localFreqs) {
-      localCosBuffer = acquireBuffer(localFreqs.cos.byteLength, undefined, 'rope_local_cos');
-      localSinBuffer = acquireBuffer(localFreqs.sin.byteLength, undefined, 'rope_local_sin');
-      device.queue.writeBuffer(localCosBuffer, 0, localFreqs.cos.buffer, localFreqs.cos.byteOffset, localFreqs.cos.byteLength);
-      device.queue.writeBuffer(localSinBuffer, 0, localFreqs.sin.buffer, localFreqs.sin.byteOffset, localFreqs.sin.byteLength);
+    let cosBuffer = null;
+    let sinBuffer = null;
+    let localCosBuffer = null;
+    let localSinBuffer = null;
+    try {
+      cosBuffer = acquireBuffer(globalFreqs.cos.byteLength, undefined, 'rope_cos');
+      sinBuffer = acquireBuffer(globalFreqs.sin.byteLength, undefined, 'rope_sin');
+      device.queue.writeBuffer(cosBuffer, 0, globalFreqs.cos.buffer, globalFreqs.cos.byteOffset, globalFreqs.cos.byteLength);
+      device.queue.writeBuffer(sinBuffer, 0, globalFreqs.sin.buffer, globalFreqs.sin.byteOffset, globalFreqs.sin.byteLength);
+
+      if (localFreqs) {
+        localCosBuffer = acquireBuffer(localFreqs.cos.byteLength, undefined, 'rope_local_cos');
+        localSinBuffer = acquireBuffer(localFreqs.sin.byteLength, undefined, 'rope_local_sin');
+        device.queue.writeBuffer(localCosBuffer, 0, localFreqs.cos.buffer, localFreqs.cos.byteOffset, localFreqs.cos.byteLength);
+        device.queue.writeBuffer(localSinBuffer, 0, localFreqs.sin.buffer, localFreqs.sin.byteOffset, localFreqs.sin.byteLength);
+      }
+    } catch (error) {
+      for (const buffer of [cosBuffer, sinBuffer, localCosBuffer, localSinBuffer]) {
+        if (buffer) {
+          releaseBuffer(buffer);
+        }
+      }
+      throw error;
     }
 
     log.debug(
@@ -444,6 +507,12 @@ export function createKVCache(modelConfig, useGPU, debug = false, runtimeConfig)
     cacheLayout = 'paged';
     layoutSource = 'threshold';
   }
+  if (forceContiguousKVCache && cacheLayout === 'paged') {
+    throw new Error(
+      'Paged KV cache layout is not supported for models with full-attention layers. ' +
+      'Set runtime.inference.kvcache.layout to "contiguous" instead.'
+    );
+  }
   if (debug && cacheLayout !== runtimeKV.layout) {
     log.debug('Pipeline', `KV cache layout override: ${runtimeKV.layout} -> ${cacheLayout} (${layoutSource})`);
   }
@@ -541,7 +610,7 @@ export function createKVCache(modelConfig, useGPU, debug = false, runtimeConfig)
 
   if (debug) {
     if (forceContiguousKVCache && modelConfig.layerTypes) {
-      log.debug('Pipeline', 'Layer pattern includes full-attention layers; forcing contiguous KV cache.');
+      log.debug('Pipeline', 'Layer pattern includes full-attention layers; paged layout blocked, contiguous enforced.');
     }
     const isSliding = kvCache instanceof SlidingWindowKVCache;
     log.debug('Pipeline', `KV cache: type=${kvCache?.constructor?.name || 'unknown'}, kvDtype=${kvCache.kvDtype}, layout=${kvCache.layout}, maxSeqLen=${kvCache.maxSeqLen}, windowSize=${isSliding ? kvCache.windowSize : null}`);

package/src/inference/pipelines/text/kernel-path-auto-select.js

@@ -78,6 +78,7 @@ export function resolveCapabilityKernelPathRef(configuredKernelPathRef, kernelPa
 
   const normalizedPolicy = resolveKernelPathPolicy(kernelPathPolicy);
   const hasSubgroups = capabilities?.hasSubgroups === true;
+  const hasF16 = capabilities?.hasF16 === true;
   const normalizedSource = normalizeKernelPathSource(kernelPathSource);
   const allowCapabilityAutoSelection = normalizedPolicy.mode === 'capability-aware'
     && normalizedPolicy.sourceScope.includes(normalizedSource);
@@ -85,6 +86,7 @@ export function resolveCapabilityKernelPathRef(configuredKernelPathRef, kernelPa
   return selectRuleValue('inference', 'kernelPath', 'autoSelect', {
     kernelPathRef: configuredKernelPathRef,
     hasSubgroups,
+    hasF16,
     allowCapabilityAutoSelection,
   });
 }

package/src/inference/pipelines/text/kernel-trace.d.ts

@@ -12,6 +12,8 @@
  * Snapshot of a tensor's statistics (no full data, just stats).
  */
 export interface TensorSnapshot {
+  ok: boolean;
+  error: string | null;
   shape: number[];
   dtype: string;
   stats: {

package/src/inference/pipelines/text/kernel-trace.js

@@ -283,6 +283,9 @@ export async function traceStep(name, label, layer, outputBuffer, outputShape, o
   if (layer >= 0 && !kernelTrace.shouldTraceLayer(layer)) return;
 
   const output = await snapshotTensor(outputBuffer, outputShape);
+  if (!output.ok) {
+    throw new Error(`[TRACE] Failed to snapshot output for ${label}: ${output.error}`);
+  }
 
   // Snapshot inputs if provided (expensive - only do if tracing)
   
@@ -290,6 +293,9 @@ export async function traceStep(name, label, layer, outputBuffer, outputShape, o
   if (options?.inputs && options?.inputShapes) {
     for (let i = 0; i < options.inputs.length; i++) {
       const snap = await snapshotTensor(options.inputs[i], options.inputShapes[i]);
+      if (!snap.ok) {
+        throw new Error(`[TRACE] Failed to snapshot input ${i} for ${label}: ${snap.error}`);
+      }
       inputs.push(snap);
     }
   }

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

@@ -2,7 +2,7 @@
 
 import { log, trace } from '../../../debug/index.js';
 import { getDevice } from '../../../gpu/device.js';
-import { releaseBuffer } from '../../../memory/buffer-pool.js';
+import { releaseBuffer, readBuffer } from '../../../memory/buffer-pool.js';
 import { allowReadback } from '../../../gpu/perf-guards.js';
 import { createTensor } from '../../../gpu/tensor.js';
 import {
@@ -228,6 +228,7 @@ export async function processLayerGPU(layerIdx, inputBuffer, numTokens, isPrefil
       linearRuntime: context.linearAttentionRuntime ?? null,
       getWeightBuffer: (weight, label) => getWeightBuffer(weight, label),
       getNormWeightBuffer: (weight, label) => getNormWeightBuffer(weight, label, weightConfig, debugFlags),
+      debugProbes: context.debugProbes,
       recorder: recorder ?? null,
     });
   } else {
@@ -275,6 +276,7 @@ export async function processLayerGPU(layerIdx, inputBuffer, numTokens, isPrefil
         : (ropeFreqsSin),
       kvCache: ((kvCache)),
       stats: context.stats,
+      debugProbes: context.debugProbes,
       linearRuntime: context.linearAttentionRuntime ?? null,
     };
 
@@ -314,14 +316,7 @@ export async function processLayerGPU(layerIdx, inputBuffer, numTokens, isPrefil
       if (allowReadback(`layer.attn-out.${layerIdx}`)) {
         try {
           const sampleSize = Math.min(128, attnOutput.buffer.size);
-          const staging = device.createBuffer({ size: sampleSize, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ });
-          const enc = device.createCommandEncoder();
-          enc.copyBufferToBuffer(attnOutput.buffer, 0, staging, 0, sampleSize);
-          device.queue.submit([enc.finish()]);
-          await staging.mapAsync(GPUMapMode.READ);
-          const data = new Float32Array(staging.getMappedRange().slice(0));
-          staging.unmap();
-          staging.destroy();
+          const data = new Float32Array(await readBuffer(attnOutput.buffer, sampleSize));
           let maxAbs = 0;
           for (let i = 0; i < data.length; i++) {
             const abs = Math.abs(data[i]);

package/src/inference/pipelines/text/linear-attention.d.ts

@@ -3,6 +3,7 @@ import type { Tensor } from '../../../gpu/tensor.js';
 import type { WeightBuffer } from '../../../gpu/weight-buffer.js';
 import type { CommandRecorder } from '../../../gpu/command-recorder.js';
 import type { LinearNormMode } from '../../../config/schema/index.js';
+import type { ProbeConfigSchema } from '../../../config/schema/index.js';
 
 export interface LinearLayerRuntimeState {
   layerIdx: number;
@@ -67,6 +68,7 @@ export interface RunLinearAttentionLayerOptions {
     weight: GPUBuffer | Float32Array | ArrayBuffer,
     label: string
   ) => GPUBuffer;
+  debugProbes?: ProbeConfigSchema[] | null;
   recorder?: CommandRecorder | null;
 }
 
@@ -74,6 +76,19 @@ export declare function hasLinearAttentionLayers(layerTypes: unknown): boolean;
 
 export declare function createLinearAttentionRuntime(): LinearAttentionRuntime;
 
+export declare function inferLinearNormMode(
+  weight: { size?: number; dtype?: string } | GPUBuffer | WeightBuffer | ArrayBufferView | ArrayBuffer | null | undefined,
+  projectionLayout: {
+    headVDim: number;
+    valueDim: number;
+  }
+): LinearNormMode | null;
+
+export declare function applyLinearNormWeightOffset(
+  values: Float32Array,
+  rmsNormWeightOffset: boolean
+): Float32Array;
+
 export declare function resetLinearAttentionRuntime(
   runtime: LinearAttentionRuntime | null | undefined
 ): LinearAttentionRuntime;