@simulatte/doppler 0.1.6 → 0.1.7

package/src/tooling/source-runtime-bundle.js

@@ -14,11 +14,13 @@ import {
   parseGroupLayerIndex,
   sortGroupIds,
 } from '../formats/rdrr/index.js';
+import { computeHash } from '../storage/shard-manager.js';
 
-const PLACEHOLDER_HASH = '0'.repeat(64);
 export const DIRECT_SOURCE_RUNTIME_MODE = 'direct-source';
 export const DIRECT_SOURCE_RUNTIME_SCHEMA_VERSION = 1;
 export const DIRECT_SOURCE_RUNTIME_SCHEMA = `direct-source/v${DIRECT_SOURCE_RUNTIME_SCHEMA_VERSION}`;
+export const DIRECT_SOURCE_PATH_RUNTIME_LOCAL = 'runtime-local';
+export const DIRECT_SOURCE_PATH_ARTIFACT_RELATIVE = 'artifact-relative';
 
 function toPathKey(value) {
   return String(value || '').trim().replace(/\\/g, '/');
@@ -38,6 +40,31 @@ function toUint8Chunk(value, label) {
   return value instanceof Uint8Array ? value : new Uint8Array(toArrayBuffer(value, label));
 }
 
+function encodeUtf8(value) {
+  return new TextEncoder().encode(String(value ?? ''));
+}
+
+function normalizeHashAlgorithm(value) {
+  const normalized = String(value || '').trim().toLowerCase();
+  return normalized === 'blake3' ? 'blake3' : 'sha256';
+}
+
+function normalizeHashString(value, label) {
+  if (value == null) return null;
+  const normalized = String(value).trim().toLowerCase();
+  if (!normalized) return null;
+  if (!/^[a-f0-9]{64}$/.test(normalized)) {
+    throw new Error(`${label} must be a 64-character lowercase hex digest.`);
+  }
+  return normalized;
+}
+
+function normalizeAssetKind(value) {
+  const normalized = String(value || '').trim().toLowerCase();
+  if (!normalized) return 'unknown';
+  return normalized;
+}
+
 function normalizePositiveInteger(value, label) {
   const parsed = Number(value);
   if (!Number.isFinite(parsed) || parsed < 0) {
@@ -66,7 +93,12 @@ async function resolveSourceFiles(tensors, sourceFiles, resolveSourceSize) {
     const path = toPathKey(entry?.path);
     if (!path) continue;
     const size = normalizePositiveInteger(entry?.size, `source file size (${path})`);
-    fileMap.set(path, { path, size });
+    fileMap.set(path, {
+      path,
+      size,
+      hash: normalizeHashString(entry?.hash, `source file hash (${path})`),
+      hashAlgorithm: normalizeHashAlgorithm(entry?.hashAlgorithm),
+    });
   }
 
   for (const tensor of tensors) {
@@ -106,7 +138,7 @@ function buildSourceShards(sourceFiles, hashAlgorithm) {
       index,
       filename,
       size: file.size,
-      hash: PLACEHOLDER_HASH,
+      hash: file.hash ?? '',
       hashAlgorithm,
       offset,
     });
@@ -115,6 +147,8 @@ function buildSourceShards(sourceFiles, hashAlgorithm) {
       path: file.path,
       filename,
       size: file.size,
+      hash: file.hash ?? '',
+      hashAlgorithm,
     });
     offset += file.size;
   }
@@ -203,7 +237,7 @@ function buildSourceGroups(tensorLocations, modelType) {
       version: '1.0.0',
       shards: Array.from(entry.shards).sort((left, right) => left - right),
       tensors: [...entry.tensors].sort((left, right) => left.localeCompare(right)),
-      hash: PLACEHOLDER_HASH,
+      hash: '',
       ...(Number.isInteger(layerIndex) ? { layerIndex } : {}),
       ...(Number.isInteger(expertIndex) ? { expertIndex } : {}),
     };
@@ -212,6 +246,170 @@ function buildSourceGroups(tensorLocations, modelType) {
   return groups;
 }
 
+async function assignGroupHashes(groups, tensorLocations, hashAlgorithm) {
+  const groupIds = sortGroupIds(Object.keys(groups ?? {}));
+  for (const groupId of groupIds) {
+    const group = groups[groupId];
+    if (!group) continue;
+    const tensors = Array.isArray(group.tensors) ? group.tensors : [];
+    const payload = {
+      groupId,
+      type: group.type ?? null,
+      version: group.version ?? null,
+      layerIndex: Number.isInteger(group.layerIndex) ? group.layerIndex : null,
+      expertIndex: Number.isInteger(group.expertIndex) ? group.expertIndex : null,
+      tensors: tensors.map((tensorName) => {
+        const location = tensorLocations?.[tensorName] ?? null;
+        return {
+          name: tensorName,
+          shard: location?.shard ?? null,
+          offset: location?.offset ?? null,
+          size: location?.size ?? null,
+          dtype: location?.dtype ?? null,
+          shape: Array.isArray(location?.shape) ? location.shape : null,
+          layout: location?.layout ?? null,
+        };
+      }),
+    };
+    group.hash = await computeHash(encodeUtf8(JSON.stringify(payload)), hashAlgorithm);
+  }
+}
+
+function normalizeAuxiliaryFileEntry(entry, defaultHashAlgorithm) {
+  const path = toPathKey(entry?.path);
+  if (!path) return null;
+  return {
+    path,
+    size: normalizePositiveInteger(entry?.size, `source auxiliary file size (${path})`),
+    hash: normalizeHashString(entry?.hash, `source auxiliary file hash (${path})`),
+    hashAlgorithm: normalizeHashAlgorithm(entry?.hashAlgorithm ?? defaultHashAlgorithm),
+    kind: normalizeAssetKind(entry?.kind),
+  };
+}
+
+function normalizeAuxiliaryFiles(auxiliaryFiles, defaultHashAlgorithm) {
+  const normalized = [];
+  for (const entry of Array.isArray(auxiliaryFiles) ? auxiliaryFiles : []) {
+    const resolved = normalizeAuxiliaryFileEntry(entry, defaultHashAlgorithm);
+    if (resolved) normalized.push(resolved);
+  }
+  normalized.sort((left, right) => left.path.localeCompare(right.path));
+  return normalized;
+}
+
+function buildSourceRuntimeMetadata(options, manifest, shardSources, auxiliaryFiles, hashAlgorithm) {
+  const tokenizerJsonPath = typeof options.tokenizerJsonPath === 'string' && options.tokenizerJsonPath.trim()
+    ? toPathKey(options.tokenizerJsonPath)
+    : null;
+  const tokenizerConfigPath = typeof options.tokenizerConfigPath === 'string' && options.tokenizerConfigPath.trim()
+    ? toPathKey(options.tokenizerConfigPath)
+    : null;
+  const tokenizerModelPath = typeof options.tokenizerModelPath === 'string' && options.tokenizerModelPath.trim()
+    ? toPathKey(options.tokenizerModelPath)
+    : null;
+  const hasFullSourceDigests = shardSources.every((entry) => typeof entry.hash === 'string' && entry.hash.length > 0);
+  const hasFullAuxDigests = auxiliaryFiles.every((entry) => typeof entry.hash === 'string' && entry.hash.length > 0);
+
+  return {
+    mode: DIRECT_SOURCE_RUNTIME_MODE,
+    schema: DIRECT_SOURCE_RUNTIME_SCHEMA,
+    schemaVersion: DIRECT_SOURCE_RUNTIME_SCHEMA_VERSION,
+    sourceKind: typeof options.sourceKind === 'string' && options.sourceKind.trim()
+      ? String(options.sourceKind).trim().toLowerCase()
+      : null,
+    hashAlgorithm,
+    pathSemantics: DIRECT_SOURCE_PATH_RUNTIME_LOCAL,
+    sourceFileCount: shardSources.length,
+    auxiliaryFileCount: auxiliaryFiles.length,
+    sourceFiles: shardSources.map((entry) => ({
+      index: entry.index,
+      path: entry.path,
+      filename: entry.filename,
+      size: entry.size,
+      hash: entry.hash,
+      hashAlgorithm: entry.hashAlgorithm,
+    })),
+    auxiliaryFiles,
+    tokenizer: {
+      jsonPath: tokenizerJsonPath,
+      configPath: tokenizerConfigPath,
+      modelPath: tokenizerModelPath,
+    },
+    invariants: {
+      tensorIdentity: 'tensor.name',
+      shardIdentity: 'sourceFiles[index].path',
+      byteOffsets: 'shard-relative bytes',
+      hashSemantics: hasFullSourceDigests && hasFullAuxDigests
+        ? 'sourceFiles[*].hash digests raw source files; auxiliaryFiles[*].hash digests config/index/tokenizer assets'
+        : 'source digests are incomplete; persist a materialized direct-source manifest before release claims',
+      cacheKeying: hasFullSourceDigests ? 'path:size:hash' : 'path:size',
+      tokenizerAssetsCovered: tokenizerJsonPath != null || tokenizerModelPath != null,
+      manifestFamily: manifest?.modelType ?? null,
+    },
+  };
+}
+
+export function getSourceRuntimeMetadata(manifest) {
+  const metadata = manifest?.metadata?.sourceRuntime;
+  if (!metadata || typeof metadata !== 'object') {
+    return null;
+  }
+  if (metadata.mode !== DIRECT_SOURCE_RUNTIME_MODE) {
+    return null;
+  }
+
+  const hashAlgorithm = normalizeHashAlgorithm(metadata.hashAlgorithm);
+  const sourceFiles = Array.isArray(metadata.sourceFiles)
+    ? metadata.sourceFiles
+      .map((entry) => {
+        const path = toPathKey(entry?.path);
+        if (!path) return null;
+        return {
+          index: normalizePositiveInteger(entry?.index ?? 0, `source runtime sourceFiles index (${path})`),
+          path,
+          filename: typeof entry?.filename === 'string' && entry.filename.trim()
+            ? entry.filename.trim()
+            : null,
+          size: normalizePositiveInteger(entry?.size, `source runtime sourceFiles size (${path})`),
+          hash: normalizeHashString(entry?.hash, `source runtime sourceFiles hash (${path})`),
+          hashAlgorithm: normalizeHashAlgorithm(entry?.hashAlgorithm ?? hashAlgorithm),
+        };
+      })
+      .filter(Boolean)
+      .sort((left, right) => left.index - right.index)
+    : [];
+  const auxiliaryFiles = normalizeAuxiliaryFiles(metadata.auxiliaryFiles, hashAlgorithm);
+  const tokenizer = metadata.tokenizer && typeof metadata.tokenizer === 'object'
+    ? {
+      jsonPath: typeof metadata.tokenizer.jsonPath === 'string' && metadata.tokenizer.jsonPath.trim()
+        ? toPathKey(metadata.tokenizer.jsonPath)
+        : null,
+      configPath: typeof metadata.tokenizer.configPath === 'string' && metadata.tokenizer.configPath.trim()
+        ? toPathKey(metadata.tokenizer.configPath)
+        : null,
+      modelPath: typeof metadata.tokenizer.modelPath === 'string' && metadata.tokenizer.modelPath.trim()
+        ? toPathKey(metadata.tokenizer.modelPath)
+        : null,
+    }
+    : { jsonPath: null, configPath: null, modelPath: null };
+
+  return {
+    mode: DIRECT_SOURCE_RUNTIME_MODE,
+    schema: DIRECT_SOURCE_RUNTIME_SCHEMA,
+    schemaVersion: DIRECT_SOURCE_RUNTIME_SCHEMA_VERSION,
+    sourceKind: typeof metadata.sourceKind === 'string' && metadata.sourceKind.trim()
+      ? String(metadata.sourceKind).trim().toLowerCase()
+      : null,
+    hashAlgorithm,
+    pathSemantics: metadata.pathSemantics === DIRECT_SOURCE_PATH_ARTIFACT_RELATIVE
+      ? DIRECT_SOURCE_PATH_ARTIFACT_RELATIVE
+      : DIRECT_SOURCE_PATH_RUNTIME_LOCAL,
+    sourceFiles,
+    auxiliaryFiles,
+    tokenizer,
+  };
+}
+
 function resolveModelQuantization(options, tensorLocations) {
   const sourceQuantization = options.sourceQuantization
     ? normalizeQuantTag(options.sourceQuantization)
@@ -270,14 +468,15 @@ export async function buildSourceRuntimeBundle(options = {}) {
     throw new Error('source runtime bundle: tensors[] is required.');
   }
 
-  const hashAlgorithmRaw = String(options.hashAlgorithm || '').trim().toLowerCase();
-  const hashAlgorithm = hashAlgorithmRaw === 'blake3' ? 'blake3' : 'sha256';
+  const hashAlgorithm = normalizeHashAlgorithm(options.hashAlgorithm);
   const sourceFiles = await resolveSourceFiles(tensors, options.sourceFiles, options.resolveSourceSize);
   const { shards, shardSources } = buildSourceShards(sourceFiles, hashAlgorithm);
   const shardIndexByPath = new Map(shardSources.map((entry) => [entry.path, entry.index]));
   const tensorLocations = buildSourceTensorLocations(tensors, shardIndexByPath, modelType);
   const groups = buildSourceGroups(tensorLocations, modelType);
+  await assignGroupHashes(groups, tensorLocations, hashAlgorithm);
   const { quantizationInfo, manifestQuantization } = resolveModelQuantization(options, tensorLocations);
+  const auxiliaryFiles = normalizeAuxiliaryFiles(options.auxiliaryFiles, hashAlgorithm);
 
   const model = {
     name: options.modelName || modelId,
@@ -316,19 +515,13 @@ export async function buildSourceRuntimeBundle(options = {}) {
   if (!manifest.metadata || typeof manifest.metadata !== 'object') {
     manifest.metadata = {};
   }
-  manifest.metadata.sourceRuntime = {
-    mode: DIRECT_SOURCE_RUNTIME_MODE,
-    schema: DIRECT_SOURCE_RUNTIME_SCHEMA,
-    schemaVersion: DIRECT_SOURCE_RUNTIME_SCHEMA_VERSION,
-    sourceFileCount: shardSources.length,
-    invariants: {
-      tensorIdentity: 'tensor.name',
-      shardIdentity: 'sourcePath -> shard index',
-      byteOffsets: 'shard-relative bytes',
-      hashSemantics: 'placeholder shard/group hashes; verifyHashes must be false',
-      cacheKeying: 'sourcePath:size',
-    },
-  };
+  manifest.metadata.sourceRuntime = buildSourceRuntimeMetadata(
+    options,
+    manifest,
+    shardSources,
+    auxiliaryFiles,
+    hashAlgorithm
+  );
 
   return {
     manifest,
@@ -357,7 +550,10 @@ export function createSourceStorageContext(options = {}) {
     throw new Error('source storage context: manifest is required.');
   }
 
-  const shardSources = Array.isArray(options.shardSources) ? options.shardSources : null;
+  const sourceRuntime = getSourceRuntimeMetadata(manifest);
+  const shardSources = Array.isArray(options.shardSources) && options.shardSources.length > 0
+    ? options.shardSources
+    : (sourceRuntime?.sourceFiles ?? null);
   if (!shardSources || shardSources.length === 0) {
     throw new Error('source storage context: shardSources[] is required.');
   }
@@ -376,17 +572,15 @@ export function createSourceStorageContext(options = {}) {
   const readBinary = typeof options.readBinary === 'function'
     ? options.readBinary
     : null;
-  const tokenizerJsonPath = options.tokenizerJsonPath ?? null;
-  const tokenizerModelPath = options.tokenizerModelPath ?? null;
+  const auxiliaryFileMap = new Map(
+    (sourceRuntime?.auxiliaryFiles ?? []).map((entry) => [entry.path, entry])
+  );
+  const tokenizerJsonPath = options.tokenizerJsonPath ?? sourceRuntime?.tokenizer?.jsonPath ?? null;
+  const tokenizerModelPath = options.tokenizerModelPath ?? sourceRuntime?.tokenizer?.modelPath ?? null;
   const verifyHashes = options.verifyHashes === true;
-  if (verifyHashes) {
-    throw new Error(
-      'source storage context: verifyHashes=true is not supported for direct-source manifests. ' +
-      'Convert to persisted RDRR shards first when hash verification is required.'
-    );
-  }
+  const allowRangeFastPath = verifyHashes !== true;
 
-  const loadShardRange = async (index, offset = 0, length = null) => {
+  const loadShardRange = allowRangeFastPath ? async (index, offset = 0, length = null) => {
     const { sourcePath, shardSize } = resolveSourceEntry(index, manifest, shardSources);
     const start = normalizePositiveInteger(offset, `shard offset (${index})`);
     const maxLength = Math.max(0, shardSize - start);
@@ -398,14 +592,19 @@ export function createSourceStorageContext(options = {}) {
     }
     const payload = await readRange(sourcePath, start, requested);
     return toArrayBuffer(payload, `readRange(${sourcePath})`);
-  };
+  } : null;
 
   const loadShard = async (index) => {
     const { shardSize } = resolveSourceEntry(index, manifest, shardSources);
-    return loadShardRange(index, 0, shardSize);
+    if (loadShardRange) {
+      return loadShardRange(index, 0, shardSize);
+    }
+    const { sourcePath } = resolveSourceEntry(index, manifest, shardSources);
+    const payload = await readRange(sourcePath, 0, shardSize);
+    return toArrayBuffer(payload, `readRange(${sourcePath})`);
   };
 
-  const streamShardRange = async function* (index, offset = 0, length = null, streamOptions = {}) {
+  const streamShardRange = allowRangeFastPath ? async function* (index, offset = 0, length = null, streamOptions = {}) {
     const { sourcePath, shardSize } = resolveSourceEntry(index, manifest, shardSources);
     const start = normalizePositiveInteger(offset, `shard stream offset (${index})`);
     const maxLength = Math.max(0, shardSize - start);
@@ -441,18 +640,35 @@ export function createSourceStorageContext(options = {}) {
         break;
       }
     }
-  };
+  } : null;
 
   const loadTokenizerJson = readText && tokenizerJsonPath
     ? async () => {
       const raw = await readText(tokenizerJsonPath);
       if (typeof raw === 'string') {
+        if (verifyHashes) {
+          const descriptor = auxiliaryFileMap.get(tokenizerJsonPath);
+          if (descriptor?.hash) {
+            const computedHash = await computeHash(encodeUtf8(raw), descriptor.hashAlgorithm);
+            if (computedHash !== descriptor.hash) {
+              throw new Error(
+                `Tokenizer asset hash mismatch for ${tokenizerJsonPath}. ` +
+                `Expected ${descriptor.hash}, got ${computedHash}.`
+              );
+            }
+          }
+        }
         return JSON.parse(raw);
       }
+      if (verifyHashes && raw && typeof raw === 'object') {
+        throw new Error(
+          `readText(${tokenizerJsonPath}) must return the original JSON string when verifyHashes=true.`
+        );
+      }
       if (raw && typeof raw === 'object') {
         return raw;
       }
-      return null;
+      throw new Error(`readText(${tokenizerJsonPath}) did not return tokenizer JSON data.`);
     }
     : null;
 
@@ -466,6 +682,17 @@ export function createSourceStorageContext(options = {}) {
       }
       const raw = await readBinary(targetPath);
       const buffer = toArrayBuffer(raw, `readBinary(${targetPath})`);
+      if (verifyHashes) {
+        const descriptor = auxiliaryFileMap.get(targetPath);
+        if (descriptor?.hash) {
+          const computedHash = await computeHash(new Uint8Array(buffer), descriptor.hashAlgorithm);
+          if (computedHash !== descriptor.hash) {
+            throw new Error(
+              `Binary asset hash mismatch for ${targetPath}. Expected ${descriptor.hash}, got ${computedHash}.`
+            );
+          }
+        }
+      }
       if (buffer.byteLength <= 0) {
         throw new Error(`readBinary(${targetPath}) returned an empty tokenizer model payload.`);
       }

package/src/tooling/source-runtime-materializer.d.ts

@@ -0,0 +1,6 @@
+import type { RDRRManifest } from '../formats/rdrr/index.js';
+
+export declare function materializeSourceRuntimeManifest(
+  manifest: RDRRManifest,
+  artifactDir: string
+): RDRRManifest;

package/src/tooling/source-runtime-materializer.js

@@ -0,0 +1,93 @@
+import path from 'node:path';
+
+import {
+  DIRECT_SOURCE_PATH_ARTIFACT_RELATIVE,
+  DIRECT_SOURCE_RUNTIME_MODE,
+  DIRECT_SOURCE_RUNTIME_SCHEMA,
+  DIRECT_SOURCE_RUNTIME_SCHEMA_VERSION,
+  getSourceRuntimeMetadata,
+} from './source-runtime-bundle.js';
+
+function cloneJsonValue(value) {
+  if (typeof structuredClone === 'function') {
+    return structuredClone(value);
+  }
+  return JSON.parse(JSON.stringify(value));
+}
+
+function toRelativeArtifactPath(value, artifactDir, label) {
+  const raw = String(value || '').trim();
+  if (!raw) {
+    throw new Error(`${label} path is required.`);
+  }
+  const resolvedArtifactDir = path.resolve(artifactDir);
+  const resolvedTarget = path.resolve(raw);
+  const relativePath = path.relative(resolvedArtifactDir, resolvedTarget).replace(/\\/g, '/');
+  if (!relativePath || relativePath.startsWith('../') || relativePath === '..') {
+    throw new Error(
+      `${label} "${raw}" must live inside artifactDir "${resolvedArtifactDir}" for a persisted direct-source manifest.`
+    );
+  }
+  return relativePath;
+}
+
+export function materializeSourceRuntimeManifest(manifest, artifactDir) {
+  const sourceRuntime = getSourceRuntimeMetadata(manifest);
+  if (!sourceRuntime) {
+    throw new Error('materializeSourceRuntimeManifest requires manifest.metadata.sourceRuntime.');
+  }
+  const resolvedArtifactDir = String(artifactDir || '').trim();
+  if (!resolvedArtifactDir) {
+    throw new Error('materializeSourceRuntimeManifest requires artifactDir.');
+  }
+
+  const nextManifest = cloneJsonValue(manifest);
+  if (!nextManifest.metadata || typeof nextManifest.metadata !== 'object') {
+    nextManifest.metadata = {};
+  }
+  const sourceMetadata = nextManifest.metadata.sourceRuntime && typeof nextManifest.metadata.sourceRuntime === 'object'
+    ? cloneJsonValue(nextManifest.metadata.sourceRuntime)
+    : {};
+
+  sourceMetadata.mode = DIRECT_SOURCE_RUNTIME_MODE;
+  sourceMetadata.schema = DIRECT_SOURCE_RUNTIME_SCHEMA;
+  sourceMetadata.schemaVersion = DIRECT_SOURCE_RUNTIME_SCHEMA_VERSION;
+  sourceMetadata.hashAlgorithm = sourceRuntime.hashAlgorithm;
+  sourceMetadata.pathSemantics = DIRECT_SOURCE_PATH_ARTIFACT_RELATIVE;
+  sourceMetadata.sourceFiles = sourceRuntime.sourceFiles.map((entry) => ({
+    index: entry.index,
+    filename: entry.filename ?? null,
+    path: toRelativeArtifactPath(
+      entry.path,
+      resolvedArtifactDir,
+      `source runtime source file ${entry.index}`
+    ),
+    size: entry.size,
+    hash: entry.hash,
+    hashAlgorithm: entry.hashAlgorithm,
+  }));
+  sourceMetadata.auxiliaryFiles = sourceRuntime.auxiliaryFiles.map((entry) => ({
+    path: toRelativeArtifactPath(
+      entry.path,
+      resolvedArtifactDir,
+      `source runtime auxiliary file ${entry.kind}`
+    ),
+    size: entry.size,
+    hash: entry.hash,
+    hashAlgorithm: entry.hashAlgorithm,
+    kind: entry.kind,
+  }));
+  sourceMetadata.tokenizer = {
+    jsonPath: sourceRuntime.tokenizer.jsonPath
+      ? toRelativeArtifactPath(sourceRuntime.tokenizer.jsonPath, resolvedArtifactDir, 'source runtime tokenizer json')
+      : null,
+    configPath: sourceRuntime.tokenizer.configPath
+      ? toRelativeArtifactPath(sourceRuntime.tokenizer.configPath, resolvedArtifactDir, 'source runtime tokenizer config')
+      : null,
+    modelPath: sourceRuntime.tokenizer.modelPath
+      ? toRelativeArtifactPath(sourceRuntime.tokenizer.modelPath, resolvedArtifactDir, 'source runtime tokenizer model')
+      : null,
+  };
+  nextManifest.metadata.sourceRuntime = sourceMetadata;
+  return nextManifest;
+}

package/src/training/attention-backward.js

@@ -1,6 +1,7 @@
-import { acquireBuffer, uploadData, readBuffer } from '../memory/buffer-pool.js';
+import { acquireBuffer, uploadData, readBuffer, releaseBuffer } from '../memory/buffer-pool.js';
 import { createTensor, tensorBytes } from '../gpu/tensor.js';
 import { f16ToF32Array } from '../inference/kv-cache/types.js';
+import { createUploadedTensor } from './tensor-factory.js';
 
 function toFloat32(buffer, dtype) {
   if (dtype === 'f16') {
@@ -67,9 +68,7 @@ export async function buildAttentionSoftmaxCache(q, k, options) {
   const kData = toFloat32(kBuf, k.dtype);
   const sData = computeSoftmax(qData, kData, options);
   const { seqLen, numHeads } = options;
-  const outBuf = acquireBuffer(tensorBytes([numHeads, seqLen, seqLen], 'f32'), undefined, 'attn_softmax_cache');
-  uploadData(outBuf, sData);
-  return createTensor(outBuf, 'f32', [numHeads, seqLen, seqLen], 'attn_softmax_cache');
+  return createUploadedTensor(sData, 'f32', [numHeads, seqLen, seqLen], 'attn_softmax_cache');
 }
 
 export async function attentionBackwardCpu(
@@ -201,17 +200,33 @@ export async function attentionBackwardCpu(
     }
   }
 
-  const qBufOut = acquireBuffer(tensorBytes([seqLen, numHeads, headDim], 'f32'), undefined, 'attn_backward_q');
-  const kBufOut = acquireBuffer(tensorBytes([seqLen, numHeads, headDim], 'f32'), undefined, 'attn_backward_k');
-  const vBufOut = acquireBuffer(tensorBytes([seqLen, numHeads, headDim], 'f32'), undefined, 'attn_backward_v');
-
-  uploadData(qBufOut, dQ);
-  uploadData(kBufOut, dK);
-  uploadData(vBufOut, dV);
-
-  return {
-    gradQ: createTensor(qBufOut, 'f32', [seqLen, numHeads, headDim], 'attn_grad_q'),
-    gradK: createTensor(kBufOut, 'f32', [seqLen, numHeads, headDim], 'attn_grad_k'),
-    gradV: createTensor(vBufOut, 'f32', [seqLen, numHeads, headDim], 'attn_grad_v'),
-  };
+  let qBufOut = null;
+  let kBufOut = null;
+  let vBufOut = null;
+  try {
+    qBufOut = acquireBuffer(tensorBytes([seqLen, numHeads, headDim], 'f32'), undefined, 'attn_backward_q');
+    kBufOut = acquireBuffer(tensorBytes([seqLen, numHeads, headDim], 'f32'), undefined, 'attn_backward_k');
+    vBufOut = acquireBuffer(tensorBytes([seqLen, numHeads, headDim], 'f32'), undefined, 'attn_backward_v');
+
+    uploadData(qBufOut, dQ);
+    uploadData(kBufOut, dK);
+    uploadData(vBufOut, dV);
+
+    return {
+      gradQ: createTensor(qBufOut, 'f32', [seqLen, numHeads, headDim], 'attn_grad_q'),
+      gradK: createTensor(kBufOut, 'f32', [seqLen, numHeads, headDim], 'attn_grad_k'),
+      gradV: createTensor(vBufOut, 'f32', [seqLen, numHeads, headDim], 'attn_grad_v'),
+    };
+  } catch (error) {
+    if (qBufOut) {
+      releaseBuffer(qBufOut);
+    }
+    if (kBufOut) {
+      releaseBuffer(kBufOut);
+    }
+    if (vBufOut) {
+      releaseBuffer(vBufOut);
+    }
+    throw error;
+  }
 }