@simulatte/doppler 0.1.5 → 0.1.7

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

@@ -1043,18 +1043,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 +1103,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

@@ -71,9 +71,13 @@ export interface PipelineContexts {
  */
 export interface RoPEConfig {
   headDim: number;
+  rotaryDim?: number;
   maxSeqLen: number;
   ropeTheta: number;
   ropeLocalTheta?: number | null;
+  mropeInterleaved?: boolean;
+  mropeSection?: number[] | null;
+  partialRotaryFactor?: number | null;
   ropeScale: number;
   ropeLocalScale?: number;
   ropeScalingType?: string | null;

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';
@@ -14,6 +14,10 @@ import { PAGED_LAYOUT_SEQ_LEN_THRESHOLD } from '../../../config/schema/index.js'
 import { isKernelPathFusedQ4K } 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,11 +122,7 @@ 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(/^\/+/, '')}`);
     },
   };
 }
@@ -206,13 +255,45 @@ function isSameRoPEScalingConfig(
       === (rightScaling?.original_max_position_embeddings ?? null);
 }
 
+function resolveRotaryDim(headDim, rotaryDim, partialRotaryFactor) {
+  if (rotaryDim != null) {
+    if (!Number.isFinite(rotaryDim) || rotaryDim <= 0 || (rotaryDim % 2) !== 0) {
+      throw new Error(`RoPE rotary dim must be a positive even integer; got "${rotaryDim}".`);
+    }
+    if (rotaryDim > headDim) {
+      throw new Error(`RoPE rotary dim ${rotaryDim} cannot exceed headDim ${headDim}.`);
+    }
+    return rotaryDim;
+  }
+  if (partialRotaryFactor == null) {
+    return headDim;
+  }
+  if (!Number.isFinite(partialRotaryFactor) || partialRotaryFactor <= 0 || partialRotaryFactor > 1) {
+    throw new Error(
+      `RoPE partialRotaryFactor must be a number in (0, 1]; got "${partialRotaryFactor}".`
+    );
+  }
+  const resolved = Math.trunc(headDim * partialRotaryFactor);
+  if (resolved <= 0 || (resolved % 2) !== 0) {
+    throw new Error(
+      `RoPE partialRotaryFactor=${partialRotaryFactor} with headDim=${headDim} resolves ` +
+      `to rotaryDim=${resolved}, but rotaryDim must be a positive even integer.`
+    );
+  }
+  return resolved;
+}
+
 
 export async function initRoPEFrequencies(config, useGPU) {
   const {
     headDim,
+    rotaryDim,
     maxSeqLen,
     ropeTheta,
     ropeLocalTheta,
+    mropeInterleaved,
+    mropeSection,
+    partialRotaryFactor,
     ropeScale,
     ropeLocalScale,
     ropeScalingType,
@@ -230,14 +311,23 @@ export async function initRoPEFrequencies(config, useGPU) {
   const resolvedLocalTheta = ropeLocalTheta ?? ropeTheta;
   const resolvedLocalScalingType = ropeLocalScalingType ?? ropeScalingType;
   const resolvedLocalScaling = ropeLocalScaling ?? ropeScaling;
+  const resolvedRotaryDim = resolveRotaryDim(headDim, rotaryDim, partialRotaryFactor);
+  const halfDim = resolvedRotaryDim / 2;
+  if (mropeInterleaved === true && Array.isArray(mropeSection)) {
+    const expandedDim = mropeSection.reduce((sum, entry) => sum + entry, 0) * 2;
+    if (expandedDim !== resolvedRotaryDim) {
+      throw new Error(
+        `RoPE mropeSection expands to ${expandedDim} dims, but rotaryDim is ${resolvedRotaryDim}.`
+      );
+    }
+  }
 
-  const halfDim = headDim / 2;
   const isYarn = ropeScalingType === 'yarn';
   const isLocalYarn = resolvedLocalScalingType === 'yarn';
 
   // Compute global (full_attention) frequencies
   const globalFreqs = computeRoPEFreqsForTheta(
-    ropeTheta, headDim, maxSeqLen, ropeScale, ropeScalingType, ropeScaling
+    ropeTheta, resolvedRotaryDim, maxSeqLen, ropeScale, ropeScalingType, ropeScaling
   );
 
   // Compute local (sliding_attention) frequencies if different from global.
@@ -256,7 +346,7 @@ export async function initRoPEFrequencies(config, useGPU) {
   if (hasDistinctLocalTheta || hasDistinctLocalScaling) {
     localFreqs = computeRoPEFreqsForTheta(
       resolvedLocalTheta,
-      headDim,
+      resolvedRotaryDim,
       maxSeqLen,
       resolvedLocalScale,
       resolvedLocalScalingType,
@@ -285,27 +375,37 @@ 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(
       'Pipeline',
-      `RoPE frequencies initialized (GPU): ${maxSeqLen} positions, dim=${halfDim}, headDim=${headDim}, ` +
+      `RoPE frequencies initialized (GPU): ${maxSeqLen} positions, dim=${halfDim}, headDim=${headDim}, rotaryDim=${resolvedRotaryDim}, ` +
       `theta=${ropeTheta}${hasDistinctLocalTheta ? `, localTheta=${resolvedLocalTheta}` : ''}, ` +
-      `scaling=${ropeScalingType ?? 'none'}:${ropeScale}${hasDistinctLocalScaling ? `, localScaling=${resolvedLocalScalingType ?? 'none'}:${resolvedLocalScale}` : ''}`
+      `scaling=${ropeScalingType ?? 'none'}:${ropeScale}${hasDistinctLocalScaling ? `, localScaling=${resolvedLocalScalingType ?? 'none'}:${resolvedLocalScale}` : ''}, ` +
+      `interleaved=${mropeInterleaved === true}`
     );
 
     return {
@@ -318,9 +418,10 @@ export async function initRoPEFrequencies(config, useGPU) {
 
   log.debug(
     'Pipeline',
-    `RoPE frequencies initialized (CPU): ${maxSeqLen} positions, dim=${halfDim}, headDim=${headDim}, ` +
+    `RoPE frequencies initialized (CPU): ${maxSeqLen} positions, dim=${halfDim}, headDim=${headDim}, rotaryDim=${resolvedRotaryDim}, ` +
     `theta=${ropeTheta}${hasDistinctLocalTheta ? `, localTheta=${resolvedLocalTheta}` : ''}, ` +
-    `scaling=${ropeScalingType ?? 'none'}:${ropeScale}${hasDistinctLocalScaling ? `, localScaling=${resolvedLocalScalingType ?? 'none'}:${resolvedLocalScale}` : ''}`
+    `scaling=${ropeScalingType ?? 'none'}:${ropeScale}${hasDistinctLocalScaling ? `, localScaling=${resolvedLocalScalingType ?? 'none'}:${resolvedLocalScale}` : ''}, ` +
+    `interleaved=${mropeInterleaved === true}`
   );
 
   return {
@@ -688,6 +789,10 @@ function applyChatMLTemplate(prompt) {
   return `<|im_start|>user\n${prompt}<|im_end|>\n<|im_start|>assistant\n`;
 }
 
+function applyQwenTemplate(prompt) {
+  return `<|im_start|>user\n${prompt}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n`;
+}
+
 function applyTranslateGemmaTemplate() {
   throw new Error(
     'TranslateGemma template requires structured messages. ' +
@@ -702,7 +807,7 @@ const PROMPT_TEMPLATES = {
   'llama3': applyHeaderBasedTemplate,
   'gpt-oss': applyChannelBasedTemplate,
   'chatml': applyChatMLTemplate,
-  'qwen': applyChatMLTemplate,  // Qwen uses ChatML format
+  'qwen': applyQwenTemplate,
   'translategemma': applyTranslateGemmaTemplate,
 };
 
@@ -721,7 +826,7 @@ export function applyChatTemplate(prompt, templateType) {
 export const applyGemmaChatTemplate = applyTurnBasedTemplate;
 export const applyLlama3ChatTemplate = applyHeaderBasedTemplate;
 export const applyGptOssChatTemplate = applyChannelBasedTemplate;
-export const applyQwenChatTemplate = applyChatMLTemplate;
+export const applyQwenChatTemplate = applyQwenTemplate;
 
 
 export function isStopToken(token, stopTokenIds, eosTokenId) {

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 {
@@ -259,6 +260,8 @@ export async function processLayerGPU(layerIdx, inputBuffer, numTokens, isPrefil
       attentionOutputGate: config.attentionOutputGate,
       causalAttention: config.causalAttention,
       rmsNormWeightOffset: config.rmsNormWeightOffset,
+      ropeRotaryDim: config.ropeRotaryDim,
+      ropeInterleaved: config.ropeInterleaved,
       tokenIds: context.currentTokenIds ?? null,
       kernelPath: context.kernelPath ?? null,
       disableRoPE,
@@ -312,14 +315,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]);
@@ -661,6 +657,8 @@ async function processLayerPlanGPU(layerIdx, inputBuffer, numTokens, isPrefill,
             attentionOutputGate: config.attentionOutputGate,
             causalAttention: config.causalAttention,
             rmsNormWeightOffset: config.rmsNormWeightOffset,
+            ropeRotaryDim: config.ropeRotaryDim,
+            ropeInterleaved: config.ropeInterleaved,
             tokenIds: context.currentTokenIds ?? null,
             skipInputNorm: step.skipInputNorm === true,
             activationDtype,
@@ -690,6 +688,7 @@ async function processLayerPlanGPU(layerIdx, inputBuffer, numTokens, isPrefill,
               hiddenSize,
               probes: context.debugProbes,
               recorder,
+              dtype: outputDtype,
             });
           }
           break;
@@ -733,6 +732,7 @@ async function processLayerPlanGPU(layerIdx, inputBuffer, numTokens, isPrefill,
               hiddenSize,
               probes: context.debugProbes,
               recorder,
+              dtype: outputDtype,
             });
           }
           break;
@@ -767,6 +767,7 @@ async function processLayerPlanGPU(layerIdx, inputBuffer, numTokens, isPrefill,
               hiddenSize,
               probes: context.debugProbes,
               recorder,
+              dtype: outputDtype,
             });
           }
           break;
@@ -801,6 +802,7 @@ async function processLayerPlanGPU(layerIdx, inputBuffer, numTokens, isPrefill,
               hiddenSize,
               probes: context.debugProbes,
               recorder,
+              dtype: outputDtype,
             });
           }
           break;
@@ -825,6 +827,7 @@ async function processLayerPlanGPU(layerIdx, inputBuffer, numTokens, isPrefill,
               hiddenSize,
               probes: context.debugProbes,
               recorder,
+              dtype: outputDtype,
             });
           }
           break;
@@ -851,6 +854,7 @@ async function processLayerPlanGPU(layerIdx, inputBuffer, numTokens, isPrefill,
               hiddenSize,
               probes: context.debugProbes,
               recorder,
+              dtype: toDtype,
             });
           }
           break;
@@ -880,6 +884,7 @@ async function processLayerPlanGPU(layerIdx, inputBuffer, numTokens, isPrefill,
     hiddenSize,
     probes: context.debugProbes,
     recorder,
+    dtype: getSlotDtype('state') ?? activationDtype,
   });
 
   const computeConfig = context.runtimeComputeConfig ?? null;

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,14 @@ 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 resetLinearAttentionRuntime(
   runtime: LinearAttentionRuntime | null | undefined
 ): LinearAttentionRuntime;

package/src/inference/pipelines/text/linear-attention.js

@@ -4,6 +4,7 @@ import { readBuffer, releaseBuffer, uploadData, acquireBuffer } from '../../../m
 import { log } from '../../../debug/index.js';
 import { decodeReadback } from './debug-utils/index.js';
 import { runLinearAttentionCoreGPU } from '../../../gpu/kernels/linear-attention-core.js';
+import { runProbes } from './probes.js';
 
 const LINEAR_RUNTIME_SCHEMA_VERSION = 1;
 const QK_L2NORM_EPS = 1e-6;
@@ -173,9 +174,22 @@ function inferLinearNormModeFromWeight(weight, projectionLayout) {
   if (weight instanceof ArrayBuffer) {
     return classify(Math.trunc(weight.byteLength / Float32Array.BYTES_PER_ELEMENT));
   }
+  const explicitDtype = typeof weight?.dtype === 'string' ? weight.dtype.toLowerCase() : null;
+  const trackedDtype = isGpuBuffer(weight) ? String(getBufferDtype(weight) ?? '').toLowerCase() : '';
+  const bytesPerElement = bytesFromDtype(explicitDtype || trackedDtype || null);
+  const sizedElements = Number.isFinite(weight?.size)
+    ? Math.trunc(Number(weight.size) / bytesPerElement)
+    : null;
+  if (sizedElements && Number(weight.size) % bytesPerElement === 0) {
+    return classify(sizedElements);
+  }
   return null;
 }
 
+export function inferLinearNormMode(weight, projectionLayout) {
+  return inferLinearNormModeFromWeight(weight, projectionLayout);
+}
+
 function resolveLinearNormMode(configNormMode, normWeight, projectionLayout, layerIdx) {
   const configuredMode = normalizeLinearNormMode(configNormMode);
   const inferredMode = inferLinearNormModeFromWeight(normWeight, projectionLayout);
@@ -185,7 +199,15 @@ function resolveLinearNormMode(configNormMode, normWeight, projectionLayout, lay
       `but norm.weight shape implies "${inferredMode}".`
     );
   }
-  return configuredMode ?? inferredMode ?? 'shared';
+  if (configuredMode) {
+    return configuredMode;
+  }
+  if (inferredMode) {
+    return inferredMode;
+  }
+  throw new Error(
+    `linear_attention layer ${layerIdx} requires explicit linearNormMode or a norm.weight shape that resolves it.`
+  );
 }
 
 async function readWeightAsF32(weight, expectedElements, label) {
@@ -395,10 +417,17 @@ async function createLayerRuntimeState(
 
   let convKernelSize = toPositiveInt(config.linearConvKernelDim) ?? null;
   if (isWeightBuffer(convKernel) && Array.isArray(convKernel.shape) && convKernel.shape.length >= 3) {
-    convKernelSize = toPositiveInt(convKernel.shape[2]) ?? convKernelSize;
+    const shapeKernelSize = toPositiveInt(convKernel.shape[2]) ?? null;
+    if (convKernelSize != null && shapeKernelSize != null && convKernelSize !== shapeKernelSize) {
+      throw new Error(
+        `linear_attention layer ${layerIdx} declares linearConvKernelDim=${convKernelSize}, ` +
+        `but conv1d weight shape implies ${shapeKernelSize}.`
+      );
+    }
+    convKernelSize = shapeKernelSize ?? convKernelSize;
   }
   if (!convKernelSize) {
-    convKernelSize = 4;
+    throw new Error(`linear_attention layer ${layerIdx} requires linearConvKernelDim.`);
   }
 
   const convWeight = await readWeightAsF32(
@@ -435,6 +464,11 @@ async function createLayerRuntimeState(
   const recurrentState = new Float32Array(
     projectionLayout.numVHeads * projectionLayout.headKDim * projectionLayout.headVDim
   );
+  const rmsNormEps = Number(config.rmsNormEps);
+  if (!Number.isFinite(rmsNormEps) || rmsNormEps <= 0) {
+    throw new Error(`linear_attention layer ${layerIdx} requires a positive rmsNormEps.`);
+  }
+
   const layerState = {
     layerIdx,
     seqLen: currentSeqLen,
@@ -452,7 +486,7 @@ async function createLayerRuntimeState(
     vSize: projectionLayout.vSize,
     qRep: projectionLayout.qRep,
     normMode,
-    rmsNormEps: Number(config.rmsNormEps) || 1e-6,
+    rmsNormEps,
     convWeight,
     dtBias,
     aNegExp,
@@ -681,13 +715,13 @@ export async function runLinearAttentionLayer(inputTensor, layerWeights, options
     const normWeightBuffer = getNormWeightBuffer(layerWeights.inputNorm, `L${layerIdx}.linear_input_norm`);
     try {
       if (recorder) {
-        normedTensor = await recordRMSNorm(recorder, inputTensor, normWeightBuffer, Number(config.rmsNormEps) || 1e-6, {
+        normedTensor = await recordRMSNorm(recorder, inputTensor, normWeightBuffer, layerState.rmsNormEps, {
           batchSize: numTokens,
           hiddenSize,
           rmsNormWeightOffset: config.rmsNormWeightOffset,
         });
       } else {
-        normedTensor = await runRMSNorm(inputTensor, normWeightBuffer, Number(config.rmsNormEps) || 1e-6, {
+        normedTensor = await runRMSNorm(inputTensor, normWeightBuffer, layerState.rmsNormEps, {
           batchSize: numTokens,
           hiddenSize,
           rmsNormWeightOffset: config.rmsNormWeightOffset,
@@ -755,6 +789,38 @@ export async function runLinearAttentionLayer(inputTensor, layerWeights, options
   });
 
   try {
+    await runProbes('linear_qkv_proj', qkvTensor.buffer, {
+      layerIdx,
+      numTokens,
+      hiddenSize: projectionLayout.convDim,
+      probes: options.debugProbes,
+      recorder,
+      dtype: qkvTensor.dtype,
+    });
+    await runProbes('linear_z_proj', zTensor.buffer, {
+      layerIdx,
+      numTokens,
+      hiddenSize: projectionLayout.valueDim,
+      probes: options.debugProbes,
+      recorder,
+      dtype: zTensor.dtype,
+    });
+    await runProbes('linear_a_proj', aTensor.buffer, {
+      layerIdx,
+      numTokens,
+      hiddenSize: projectionLayout.numVHeads,
+      probes: options.debugProbes,
+      recorder,
+      dtype: aTensor.dtype,
+    });
+    await runProbes('linear_b_proj', bTensor.buffer, {
+      layerIdx,
+      numTokens,
+      hiddenSize: projectionLayout.numVHeads,
+      probes: options.debugProbes,
+      recorder,
+      dtype: bTensor.dtype,
+    });
     const coreTensor = await runLinearAttentionCoreGPU(
       qkvTensor,
       zTensor,
@@ -768,6 +834,14 @@ export async function runLinearAttentionLayer(inputTensor, layerWeights, options
         recorder,
       }
     );
+    await runProbes('linear_core_out', coreTensor.buffer, {
+      layerIdx,
+      numTokens,
+      hiddenSize: projectionLayout.valueDim,
+      probes: options.debugProbes,
+      recorder,
+      dtype: coreTensor.dtype,
+    });
     layerState.seqLen = currentSeqLen + numTokens;
     const outProjWeight = getWeightBuffer(layerWeights.oProj, `L${layerIdx}.linear_out_proj`);
     try {