@simulatte/doppler 0.1.5 → 0.1.7

package/src/inference/network-evolution.js

@@ -1,33 +1,29 @@
-let fallbackRandomState = (Date.now() >>> 0) || 0x6d2b79f5;
-
-function unseededRandom() {
-  if (typeof crypto !== 'undefined' && typeof crypto.getRandomValues === 'function') {
-    const values = new Uint32Array(1);
-    crypto.getRandomValues(values);
-    return values[0] / 4294967296;
+function requireRandomSource(random) {
+  if (typeof random !== 'function') {
+    throw new Error('network evolution requires an explicit random() source.');
   }
-  fallbackRandomState = (fallbackRandomState + 0x6d2b79f5) >>> 0;
-  return fallbackRandomState / 4294967296;
+  return random;
 }
 
-export const mutateGenome = (genome, mutationRate = 0.1) => {
+export const mutateGenome = (genome, mutationRate = 0.1, random = null) => {
+  const sample = requireRandomSource(random);
   
   const mutated = JSON.parse(JSON.stringify(genome));
-  if (unseededRandom() < mutationRate) {
+  if (sample() < mutationRate) {
     
     const types = ['chain', 'tree', 'mesh', 'dag'];
-    mutated.topology.type = types[Math.floor(unseededRandom() * types.length)];
+    mutated.topology.type = types[Math.floor(sample() * types.length)];
   }
 
   for (const node of mutated.nodes) {
-    if (unseededRandom() < mutationRate && typeof node.temperature === 'number') {
-      node.temperature = Math.min(1, Math.max(0, node.temperature + (unseededRandom() - 0.5) * 0.2));
+    if (sample() < mutationRate && typeof node.temperature === 'number') {
+      node.temperature = Math.min(1, Math.max(0, node.temperature + (sample() - 0.5) * 0.2));
     }
   }
 
   for (const edge of mutated.edges) {
-    if (unseededRandom() < mutationRate) {
-      edge.weight = Math.min(1, Math.max(0, edge.weight + (unseededRandom() - 0.5) * 0.4));
+    if (sample() < mutationRate) {
+      edge.weight = Math.min(1, Math.max(0, edge.weight + (sample() - 0.5) * 0.4));
     }
   }
 
@@ -35,8 +31,9 @@ export const mutateGenome = (genome, mutationRate = 0.1) => {
 };
 
 
-export const crossoverGenome = (a, b) => {
-  return unseededRandom() < 0.5 ? JSON.parse(JSON.stringify(a)) : JSON.parse(JSON.stringify(b));
+export const crossoverGenome = (a, b, random = null) => {
+  const sample = requireRandomSource(random);
+  return sample() < 0.5 ? JSON.parse(JSON.stringify(a)) : JSON.parse(JSON.stringify(b));
 };
 
 
@@ -48,7 +45,9 @@ export async function evolveNetwork(config) {
     mutationRate = 0.1,
     evaluate,
     randomGenome,
+    random,
   } = config;
+  const sample = requireRandomSource(random);
 
   let population = Array.from({ length: populationSize }, () => randomGenome());
 
@@ -63,9 +62,9 @@ export async function evolveNetwork(config) {
     const offspring = [];
 
     while (offspring.length < populationSize - eliteCount) {
-      const parentA = scored[Math.floor(unseededRandom() * scored.length)].genome;
-      const parentB = scored[Math.floor(unseededRandom() * scored.length)].genome;
-      const child = mutateGenome(crossoverGenome(parentA, parentB), mutationRate);
+      const parentA = scored[Math.floor(sample() * scored.length)].genome;
+      const parentB = scored[Math.floor(sample() * scored.length)].genome;
+      const child = mutateGenome(crossoverGenome(parentA, parentB, sample), mutationRate, sample);
       offspring.push(child);
     }
 

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

@@ -8,6 +8,8 @@ export type PipelineContextOptions = {
   assignProgress?: boolean;
 };
 
+export declare function restorePipelineContexts(target: Record<string, unknown>): boolean;
+
 export declare function applyPipelineContexts(
   target: Record<string, unknown>,
   contexts?: Record<string, unknown>,
@@ -15,4 +17,5 @@ export declare function applyPipelineContexts(
 ): {
   runtimeConfig: Record<string, unknown>;
   sharedDebug: Record<string, unknown> | null | undefined;
+  restore: () => void;
 };

package/src/inference/pipelines/context.js

@@ -1,8 +1,115 @@
-import { getDevice, setDevice } from '../../gpu/device.js';
+import {
+  getDevice,
+  getKernelCapabilities,
+  getPlatformConfig,
+  setDevice,
+} from '../../gpu/device.js';
 import { applyDebugConfig, setGPUDevice } from '../../debug/index.js';
 import { getRuntimeConfig, setRuntimeConfig } from '../../config/runtime.js';
+import {
+  getLogLevel,
+  getTrace,
+  isSilentMode,
+  setLogLevel,
+  setSilentMode,
+  setTrace,
+} from '../../debug/config.js';
+import {
+  gpuDevice as debugGpuDevice,
+  traceBreakOnAnomaly,
+  traceLayerFilter,
+  traceMaxDecodeSteps,
+} from '../../debug/config.js';
+
+const RESTORE_PIPELINE_CONTEXTS = Symbol('restorePipelineContexts');
+
+function captureTargetField(target, key) {
+  return {
+    present: Object.prototype.hasOwnProperty.call(target, key),
+    value: target[key],
+  };
+}
+
+function restoreTargetField(target, key, snapshot) {
+  if (snapshot.present) {
+    target[key] = snapshot.value;
+    return;
+  }
+  delete target[key];
+}
+
+function captureDebugState() {
+  return {
+    logLevel: getLogLevel(),
+    traceCategories: getTrace(),
+    traceLayers: [...traceLayerFilter],
+    traceMaxDecodeSteps,
+    traceBreakOnAnomaly,
+    silentMode: isSilentMode(),
+    gpuDevice: debugGpuDevice,
+  };
+}
+
+function restoreDebugState(snapshot) {
+  if (snapshot.silentMode !== isSilentMode()) {
+    setSilentMode(snapshot.silentMode);
+  }
+  if (getLogLevel() !== snapshot.logLevel) {
+    setLogLevel(snapshot.logLevel);
+  }
+
+  const traceCategories = getTrace();
+  const traceChanged = traceCategories.length !== snapshot.traceCategories.length
+    || traceCategories.some((category, idx) => category !== snapshot.traceCategories[idx])
+    || traceLayerFilter.length !== snapshot.traceLayers.length
+    || traceLayerFilter.some((layer, idx) => layer !== snapshot.traceLayers[idx])
+    || traceMaxDecodeSteps !== snapshot.traceMaxDecodeSteps
+    || traceBreakOnAnomaly !== snapshot.traceBreakOnAnomaly;
+
+  if (traceChanged) {
+    if (snapshot.traceCategories.length > 0) {
+      setTrace(snapshot.traceCategories.join(','), {
+        layers: snapshot.traceLayers.length > 0 ? snapshot.traceLayers : undefined,
+        maxDecodeSteps: snapshot.traceMaxDecodeSteps > 0 ? snapshot.traceMaxDecodeSteps : undefined,
+        breakOnAnomaly: snapshot.traceBreakOnAnomaly,
+      });
+    } else {
+      setTrace(false);
+    }
+  }
+
+  setGPUDevice(snapshot.gpuDevice ?? null);
+}
+
+export function restorePipelineContexts(target) {
+  const restore = target?.[RESTORE_PIPELINE_CONTEXTS];
+  if (typeof restore !== 'function') {
+    return false;
+  }
+  delete target[RESTORE_PIPELINE_CONTEXTS];
+  restore();
+  return true;
+}
 
 export function applyPipelineContexts(target, contexts = {}, options = {}) {
+  restorePipelineContexts(target);
+
+  const previousRuntimeConfig = getRuntimeConfig();
+  const previousDevice = getDevice();
+  const previousPlatformConfig = getPlatformConfig();
+  const previousAdapterInfo = previousDevice
+    ? (getKernelCapabilities().adapterInfo ?? null)
+    : null;
+  const previousDebugState = captureDebugState();
+  const targetSnapshot = {
+    gpuContext: captureTargetField(target, 'gpuContext'),
+    useGPU: captureTargetField(target, 'useGPU'),
+    memoryContext: captureTargetField(target, 'memoryContext'),
+    storageContext: captureTargetField(target, 'storageContext'),
+    baseUrl: captureTargetField(target, 'baseUrl'),
+    _onProgress: captureTargetField(target, '_onProgress'),
+  };
+
   const runtimeConfig = contexts.runtimeConfig
     ? setRuntimeConfig(contexts.runtimeConfig)
     : getRuntimeConfig();
@@ -40,5 +147,38 @@ export function applyPipelineContexts(target, contexts = {}, options = {}) {
     target._onProgress = contexts.onProgress;
   }
 
-  return { runtimeConfig, sharedDebug };
+  let restored = false;
+  const restore = () => {
+    if (restored) {
+      return;
+    }
+    restored = true;
+    delete target[RESTORE_PIPELINE_CONTEXTS];
+
+    setRuntimeConfig(previousRuntimeConfig);
+    if (previousDevice) {
+      setDevice(previousDevice, {
+        platformConfig: previousPlatformConfig,
+        adapterInfo: previousAdapterInfo,
+      });
+    } else {
+      setDevice(null);
+    }
+    restoreDebugState(previousDebugState);
+    restoreTargetField(target, 'gpuContext', targetSnapshot.gpuContext);
+    restoreTargetField(target, 'useGPU', targetSnapshot.useGPU);
+    restoreTargetField(target, 'memoryContext', targetSnapshot.memoryContext);
+    restoreTargetField(target, 'storageContext', targetSnapshot.storageContext);
+    restoreTargetField(target, 'baseUrl', targetSnapshot.baseUrl);
+    restoreTargetField(target, '_onProgress', targetSnapshot._onProgress);
+  };
+
+  Object.defineProperty(target, RESTORE_PIPELINE_CONTEXTS, {
+    value: restore,
+    configurable: true,
+    enumerable: false,
+    writable: false,
+  });
+
+  return { runtimeConfig, sharedDebug, restore };
 }

package/src/inference/pipelines/diffusion/helpers.js

@@ -54,8 +54,13 @@ export function createDiffusionIndexBuffer(device, indices, label) {
     size: indices.byteLength,
     usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
   });
-  device.queue.writeBuffer(buffer, 0, indices);
-  return buffer;
+  try {
+    device.queue.writeBuffer(buffer, 0, indices);
+    return buffer;
+  } catch (error) {
+    buffer.destroy();
+    throw error;
+  }
 }
 
 export function expectDiffusionWeight(weight, label) {

package/src/inference/pipelines/diffusion/pipeline.js

@@ -1,7 +1,7 @@
 import { getDevice, getKernelCapabilities } from '../../../gpu/device.js';
 import { log, trace } from '../../../debug/index.js';
 import { registerPipeline } from '../registry.js';
-import { applyPipelineContexts } from '../context.js';
+import { applyPipelineContexts, restorePipelineContexts } from '../context.js';
 import { createInitializedPipeline } from '../factory.js';
 import { createRng, sampleNormal } from '../rng.js';
 import { initializeDiffusion } from './init.js';
@@ -52,6 +52,18 @@ function generateLatents(width, height, channels, latentScale, seed) {
   return { latents, latentWidth, latentHeight };
 }
 
+function generateNoiseVector(size, seed) {
+  if (!Number.isFinite(size) || size <= 0) {
+    throw new Error(`generateNoiseVector requires a positive size, got ${size}.`);
+  }
+  const out = new Float32Array(size);
+  const rand = createRng(seed ?? createRandomSeed());
+  for (let i = 0; i < size; i++) {
+    out[i] = sampleNormal(rand);
+  }
+  return out;
+}
+
 function extractTokenSet(tokensByEncoder, key) {
   const output = {};
   for (const [name, entry] of Object.entries(tokensByEncoder || {})) {
@@ -195,13 +207,10 @@ async function applySchedulerStep(latentsTensor, scheduler, stepIndex, timestep,
     const isFinalStep = stepIndex + 1 >= scheduler.timesteps.length - 1;
     const noise = isFinalStep
       ? null
-      : generateLatents(
-          runtime.latent.width,
-          runtime.latent.height,
-          runtime.latent.channels,
-          runtime.latent.scale,
+      : generateNoiseVector(
+          sample.length,
           (options.seedBase ?? createRandomSeed()) + stepIndex + 1
-        ).latents;
+        );
     const step = stepScmScheduler(scheduler, modelOutput, timestep, sample, stepIndex, noise);
     return createLatentTensor(step.prevSample, [...latentsTensor.shape], runtime);
   }
@@ -310,6 +319,7 @@ export class DiffusionPipeline {
     this.vaeWeights = null;
     this.textEncoderWeights = null;
     this.transformerWeights = null;
+    restorePipelineContexts(this);
   }
 
   async ensureVaeWeights() {

package/src/inference/pipelines/diffusion/sd3-transformer.js

@@ -299,26 +299,26 @@ function resolveModulationSegments(weight, hiddenSize, fallbackSegments, resolve
     if (Number.isInteger(segments) && segments > 0) {
       return segments;
     }
-    log.warn(
-      'Diffusion',
-      `Modulation segments mismatch for ${name || 'unknown'}: rows=${rows}, hidden=${hiddenSize}, fallback=${fallbackSegments}`
+    throw new Error(
+      `Modulation segments mismatch for ${name || 'unknown'}: rows=${rows}, hidden=${hiddenSize}, ` +
+      `expected an integer multiple instead of falling back to ${fallbackSegments}.`
     );
   }
-  return fallbackSegments;
+  throw new Error(
+    `Modulation tensor "${name || 'unknown'}" is missing shape metadata. ` +
+    `Runtime cannot fall back to ${fallbackSegments} segments.`
+  );
 }
 
 function resolveModulationOffsets(segments, hiddenSize) {
-  if (segments >= 9) {
+  if (segments === 9) {
     return {
       attn: { scale: 0, shift: hiddenSize, gate: hiddenSize * 2 },
       attn2: { scale: hiddenSize * 3, shift: hiddenSize * 4, gate: hiddenSize * 5 },
       ff: { scale: hiddenSize * 6, shift: hiddenSize * 7, gate: hiddenSize * 8 },
     };
   }
-  if (segments >= 6) {
-    if (segments !== 6) {
-      log.warn('Diffusion', `Unexpected modulation segment count=${segments}; using 6-segment layout.`);
-    }
+  if (segments === 6) {
     const attn = { scale: 0, shift: hiddenSize, gate: hiddenSize * 2 };
     return {
       attn,
@@ -326,7 +326,7 @@ function resolveModulationOffsets(segments, hiddenSize) {
       ff: { scale: hiddenSize * 3, shift: hiddenSize * 4, gate: hiddenSize * 5 },
     };
   }
-  throw new Error(`Unsupported modulation segments=${segments} (expected >= 6).`);
+  throw new Error(`Unsupported modulation segments=${segments} (expected 6 or 9).`);
 }
 
 async function buildModulation(timeText, weight, bias, hiddenSize, segments, runtime, matmul, weightName, ops) {

package/src/inference/pipelines/diffusion/text-encoder-gpu.d.ts

@@ -80,3 +80,8 @@ export declare function projectContext(
 ): Promise<Tensor>;
 
 export declare function assertClipHiddenActivationSupported(config: { hidden_act?: string }): void;
+
+export declare function resolveGemma2WeightRoot(
+  weights: Map<string, any>,
+  prefix?: string
+): string;

package/src/inference/pipelines/diffusion/text-encoder-gpu.js

@@ -723,8 +723,19 @@ function buildGemma2LayerTypes(layerCount, slidingWindow) {
   ));
 }
 
-function getGemma2LayerWeight(weights, prefix, layerIdx, suffix, required = true) {
-  const key = `${prefix}.model.layers.${layerIdx}.${suffix}`;
+export function resolveGemma2WeightRoot(weights, prefix = 'text_encoder') {
+  const nestedRoot = `${prefix}.model`;
+  if (weights?.has(`${nestedRoot}.embed_tokens.weight`)) {
+    return nestedRoot;
+  }
+  if (weights?.has(`${prefix}.embed_tokens.weight`)) {
+    return prefix;
+  }
+  return nestedRoot;
+}
+
+function getGemma2LayerWeight(weights, weightRoot, layerIdx, suffix, required = true) {
+  const key = `${weightRoot}.layers.${layerIdx}.${suffix}`;
   const weight = weights.get(key) || null;
   if (!weight && required) {
     throw new Error(`Missing Gemma2 diffusion weight "${key}".`);
@@ -805,8 +816,9 @@ async function runGemma2TextEncoder(tokens, weightsEntry, config, runtime, optio
   const tokenIds = normalizeTokens(tokens, options.maxLength ?? resolved.maxPositionEmbeddings, padTokenId);
   const numTokens = tokenIds.length;
   const tokenBuffer = createDiffusionIndexBuffer(device, tokenIds, `${prefix}_tokens`);
+  const weightRoot = resolveGemma2WeightRoot(weights, prefix);
 
-  const embedKey = `${prefix}.model.embed_tokens.weight`;
+  const embedKey = `${weightRoot}.embed_tokens.weight`;
   const embedWeight = expectDiffusionWeight(
     weights.get(embedKey),
     embedKey
@@ -837,16 +849,16 @@ async function runGemma2TextEncoder(tokens, weightsEntry, config, runtime, optio
   const layerWeights = new Map();
   for (let layerIdx = 0; layerIdx < resolved.numLayers; layerIdx++) {
     layerWeights.set(`layer_${layerIdx}`, {
-      inputNorm: getGemma2LayerWeight(weights, prefix, layerIdx, 'input_layernorm.weight'),
-      qProj: getGemma2LayerWeight(weights, prefix, layerIdx, 'self_attn.q_proj.weight'),
-      kProj: getGemma2LayerWeight(weights, prefix, layerIdx, 'self_attn.k_proj.weight'),
-      vProj: getGemma2LayerWeight(weights, prefix, layerIdx, 'self_attn.v_proj.weight'),
-      oProj: getGemma2LayerWeight(weights, prefix, layerIdx, 'self_attn.o_proj.weight'),
-      postAttentionNorm: getGemma2LayerWeight(weights, prefix, layerIdx, 'post_attention_layernorm.weight'),
-      preFeedforwardNorm: getGemma2LayerWeight(weights, prefix, layerIdx, 'pre_feedforward_layernorm.weight'),
-      gate: getGemma2LayerWeight(weights, prefix, layerIdx, 'mlp.gate_proj.weight'),
-      up: getGemma2LayerWeight(weights, prefix, layerIdx, 'mlp.up_proj.weight'),
-      down: getGemma2LayerWeight(weights, prefix, layerIdx, 'mlp.down_proj.weight'),
+      inputNorm: getGemma2LayerWeight(weights, weightRoot, layerIdx, 'input_layernorm.weight'),
+      qProj: getGemma2LayerWeight(weights, weightRoot, layerIdx, 'self_attn.q_proj.weight'),
+      kProj: getGemma2LayerWeight(weights, weightRoot, layerIdx, 'self_attn.k_proj.weight'),
+      vProj: getGemma2LayerWeight(weights, weightRoot, layerIdx, 'self_attn.v_proj.weight'),
+      oProj: getGemma2LayerWeight(weights, weightRoot, layerIdx, 'self_attn.o_proj.weight'),
+      postAttentionNorm: getGemma2LayerWeight(weights, weightRoot, layerIdx, 'post_attention_layernorm.weight'),
+      preFeedforwardNorm: getGemma2LayerWeight(weights, weightRoot, layerIdx, 'pre_feedforward_layernorm.weight'),
+      gate: getGemma2LayerWeight(weights, weightRoot, layerIdx, 'mlp.gate_proj.weight'),
+      up: getGemma2LayerWeight(weights, weightRoot, layerIdx, 'mlp.up_proj.weight'),
+      down: getGemma2LayerWeight(weights, weightRoot, layerIdx, 'mlp.down_proj.weight'),
     });
   }
 
@@ -910,10 +922,10 @@ async function runGemma2TextEncoder(tokens, weightsEntry, config, runtime, optio
       numTokens * resolved.hiddenSize,
       context
     );
-    hidden = createTensor(output.buffer, output.dtype, [numTokens, resolved.hiddenSize], `gemma2_layer_${layerIdx}`);
+    hidden = createTensor(output, activationDtype, [numTokens, resolved.hiddenSize], `gemma2_layer_${layerIdx}`);
   }
 
-  const finalNormKey = `${prefix}.model.norm.weight`;
+  const finalNormKey = `${weightRoot}.norm.weight`;
   const finalNorm = expectDiffusionWeight(weights.get(finalNormKey), finalNormKey);
   const final = await ops.rmsNorm(hidden, getBuffer(finalNorm), resolved.rmsNormEps, {
     batchSize: numTokens,

package/src/inference/pipelines/diffusion/vae.js

@@ -118,13 +118,9 @@ function resolveAttentionHeadShape(channels, config) {
       headDim: channels / configuredNumHeads,
     };
   }
-
-  const fallbackHeadDims = [64, 40, 32, 24, 20, 16, 12, 10, 8, 6, 5, 4, 3, 2, 1];
-  const headDim = fallbackHeadDims.find((candidate) => candidate <= channels && channels % candidate === 0) || 1;
-  return {
-    numHeads: Math.max(1, channels / headDim),
-    headDim,
-  };
+  throw new Error(
+    `VAE attention requires explicit compatible attention_head_dim or num_attention_heads for channels=${channels}.`
+  );
 }
 
 function createBiasTensor(weight, label, fallbackDtype = 'f16') {

package/src/inference/pipelines/energy/pipeline.js

@@ -16,10 +16,10 @@ import { log, trace } from '../../../debug/index.js';
 import { DEFAULT_ENERGY_CONFIG } from '../../../config/schema/energy.schema.js';
 import { f32ToF16Array, f16ToF32Array } from '../../kv-cache/types.js';
 import { registerPipeline } from '../registry.js';
-import { applyPipelineContexts } from '../context.js';
+import { applyPipelineContexts, restorePipelineContexts } from '../context.js';
 import { createInitializedPipeline } from '../factory.js';
 import { createRng, sampleNormal } from '../rng.js';
-import { mergeQuintelConfig, runQuintelEnergyLoop } from './quintel.js';
+import { buildQuintelKernelFlags, mergeQuintelConfig, runQuintelEnergyLoop } from './quintel.js';
 
 
 function generateRandomArray(count, mode, seed, scale) {
@@ -140,24 +140,28 @@ async function createEnergyTensor(device, data, dtype, shape, label) {
   const byteLength = data.byteLength;
   const alignedSize = Math.ceil(byteLength / 4) * 4;
   const buffer = acquireBuffer(alignedSize, undefined, label);
+  try {
+    let payload = data;
+    if (alignedSize !== byteLength) {
+      const padded = new Uint8Array(alignedSize);
+      const view = data instanceof ArrayBuffer
+        ? new Uint8Array(data)
+        : new Uint8Array(data.buffer, data.byteOffset, data.byteLength);
+      padded.set(view);
+      payload = padded;
+    }
 
-  let payload = data;
-  if (alignedSize !== byteLength) {
-    const padded = new Uint8Array(alignedSize);
-    const view = data instanceof ArrayBuffer
-      ? new Uint8Array(data)
-      : new Uint8Array(data.buffer, data.byteOffset, data.byteLength);
-    padded.set(view);
-    payload = padded;
-  }
-
-  device.queue.writeBuffer(buffer, 0, payload);
-  const tensor = createTensor(buffer, dtype, shape, label);
-  const expectedBytes = tensorBytes(shape, dtype);
-  if (expectedBytes !== byteLength) {
-    log.warn('Energy', `${label} byte length mismatch: expected ${expectedBytes}, got ${byteLength}`);
+    device.queue.writeBuffer(buffer, 0, payload);
+    const tensor = createTensor(buffer, dtype, shape, label);
+    const expectedBytes = tensorBytes(shape, dtype);
+    if (expectedBytes !== byteLength) {
+      log.warn('Energy', `${label} byte length mismatch: expected ${expectedBytes}, got ${byteLength}`);
+    }
+    return tensor;
+  } catch (error) {
+    releaseBuffer(buffer);
+    throw error;
   }
-  return tensor;
 }
 
 async function readTensorToFloat32(tensor) {
@@ -202,6 +206,7 @@ export class EnergyPipeline {
 
   async unload() {
     this.manifest = null;
+    restorePipelineContexts(this);
   }
 
   async generate(request = {}) {
@@ -336,6 +341,7 @@ export class EnergyPipeline {
         const centerWeight = Number.isFinite(weights.center) ? weights.center : 1.0;
         const binarizeWeight = Number.isFinite(weights.binarize) ? weights.binarize : 0.0;
         const centerTarget = Number.isFinite(quintelConfig.centerTarget) ? quintelConfig.centerTarget : 1.0;
+        const flags = buildQuintelKernelFlags(rules, binarizeWeight);
         const energyHistory = [];
         const stepTimesMs = [];
         let lastEnergy = null;
@@ -387,11 +393,11 @@ export class EnergyPipeline {
             await runEnergyQuintelReduce(stateTensor, {
               count: elementCount,
               size,
+              flags,
               symmetryWeight,
               centerWeight,
               binarizeWeight,
               centerTarget,
-              rules,
               outputBuffer: reduceBuffer,
             });
 
@@ -447,13 +453,13 @@ export class EnergyPipeline {
             await runEnergyQuintelGrad(stateTensor, {
               count: elementCount,
               size,
+              flags,
               countDiff: safeCountDiff,
               symmetryWeight,
               countWeight,
               centerWeight,
               binarizeWeight,
               centerTarget,
-              rules,
               outputBuffer: gradBuffer,
             });
 
@@ -471,6 +477,7 @@ export class EnergyPipeline {
             await runEnergyQuintelUpdate(stateTensor, {
               count: elementCount,
               size,
+              flags,
               stepSize,
               gradientScale,
               countDiff: safeCountDiff,
@@ -481,7 +488,6 @@ export class EnergyPipeline {
               centerTarget,
               clampMin,
               clampMax,
-              rules,
             });
           }
 

package/src/inference/pipelines/energy/quintel.d.ts

@@ -84,4 +84,9 @@ export function mergeQuintelConfig(
   override?: Partial<QuintelEnergyConfig> | null
 ): QuintelEnergyConfig;
 
+export function buildQuintelKernelFlags(
+  rules: Partial<QuintelRuleConfig> | null | undefined,
+  binarizeWeight?: number
+): number;
+
 export function runQuintelEnergyLoop(options: QuintelEnergyLoopOptions): QuintelEnergyLoopResult;

package/src/inference/pipelines/energy/quintel.js

@@ -22,6 +22,17 @@ export function mergeQuintelConfig(base, override) {
   };
 }
 
+export function buildQuintelKernelFlags(rules, binarizeWeight) {
+  let flags = 0;
+  if (rules?.mirrorX) flags |= 1;
+  if (rules?.mirrorY) flags |= 2;
+  if (rules?.diagonal) flags |= 4;
+  if (rules?.count) flags |= 8;
+  if (rules?.center) flags |= 16;
+  if (Number.isFinite(binarizeWeight) && binarizeWeight !== 0) flags |= 32;
+  return flags >>> 0;
+}
+
 function applyPairEnergy(state, gradients, indexA, indexB, weight) {
   const diff = state[indexA] - state[indexB];
   const energy = weight * diff * diff;

package/src/inference/pipelines/energy-head/row-head-pipeline.js

@@ -5,7 +5,7 @@ import { runEnergyEval, runEnergyUpdate } from '../../../gpu/kernels/index.js';
 import { log } from '../../../debug/index.js';
 import { f16ToF32Array, f32ToF16Array } from '../../kv-cache/types.js';
 import { registerPipeline } from '../registry.js';
-import { applyPipelineContexts } from '../context.js';
+import { applyPipelineContexts, restorePipelineContexts } from '../context.js';
 import { createInitializedPipeline } from '../factory.js';
 import { selectRuleValue } from '../../../rules/rule-registry.js';
 
@@ -165,19 +165,22 @@ async function createFeatureTensor(device, values, dtype, label) {
   const byteLength = payload.byteLength;
   const alignedSize = Math.ceil(byteLength / 4) * 4;
   const buffer = acquireBuffer(alignedSize, undefined, label);
-
-  if (alignedSize === byteLength) {
-    device.queue.writeBuffer(buffer, 0, payload);
-  } else {
-    const bytes = payload instanceof Uint16Array
-      ? new Uint8Array(payload.buffer, payload.byteOffset, payload.byteLength)
-      : new Uint8Array(payload.buffer, payload.byteOffset, payload.byteLength);
-    const padded = new Uint8Array(alignedSize);
-    padded.set(bytes);
-    device.queue.writeBuffer(buffer, 0, padded);
+  try {
+    if (alignedSize === byteLength) {
+      device.queue.writeBuffer(buffer, 0, payload);
+    } else {
+      const bytes = payload instanceof Uint16Array
+        ? new Uint8Array(payload.buffer, payload.byteOffset, payload.byteLength)
+        : new Uint8Array(payload.buffer, payload.byteOffset, payload.byteLength);
+      const padded = new Uint8Array(alignedSize);
+      padded.set(bytes);
+      device.queue.writeBuffer(buffer, 0, padded);
+    }
+    return createTensor(buffer, dtype, [values.length], label);
+  } catch (error) {
+    releaseBuffer(buffer);
+    throw error;
   }
-
-  return createTensor(buffer, dtype, [values.length], label);
 }
 
 async function readTensorF32(tensor) {
@@ -307,6 +310,7 @@ export class EnergyRowHeadPipeline {
     this.manifest = null;
     this.model = null;
     this.stats = {};
+    restorePipelineContexts(this);
   }
 
   async scoreRows(request = {}) {