engine/weights.js

// Weights loader: validates manifest.json, fetches weights.bin, uploads the
// whole thing into ONE storage GPUBuffer. Tensors are addressed by 256-aligned
// offsets into that buffer (bindingFor). On devices without shader-f16, pass
// targetDtype: 'f32' to expand all f16 tensors on the CPU before upload.

import { alignUp } from './shapes.js';
import { expandF16, f16ToF32 } from './f16.js';
import { applyModelConfig, parseModelConfig } from './constants.js';

const BYTES = { f16: 2, f32: 4 };

// --- pure parsing/validation ---

export function parseManifest(json) {
  if (!json || typeof json !== 'object') throw new Error('manifest: not an object');
  if (json.version !== 1) throw new Error(`manifest: expected version 1, got ${json.version}`);
  if (!json.model || typeof json.model !== 'object') throw new Error('manifest: missing model');
  if (!Array.isArray(json.tensors) || json.tensors.length === 0) {
    throw new Error('manifest: tensors must be a non-empty array');
  }

  const tensors = new Map();
  let prevEnd = 0;
  let prevName = null;
  for (const t of json.tensors) {
    const { name, dtype, shape, byteOffset, byteLength } = t ?? {};
    if (typeof name !== 'string' || !name) throw new Error('manifest: tensor without a name');
    if (dtype !== 'f16' && dtype !== 'f32') {
      throw new Error(`manifest: tensor "${name}" has unsupported dtype "${dtype}"`);
    }
    if (!Array.isArray(shape) || shape.length === 0 ||
        !shape.every((d) => Number.isInteger(d) && d > 0)) {
      throw new Error(`manifest: tensor "${name}" has invalid shape ${JSON.stringify(shape)}`);
    }
    if (!Number.isInteger(byteOffset) || byteOffset < 0 ||
        !Number.isInteger(byteLength) || byteLength <= 0) {
      throw new Error(`manifest: tensor "${name}" has invalid byteOffset/byteLength`);
    }
    const elems = shape.reduce((a, d) => a * d, 1);
    const expectBytes = elems * BYTES[dtype];
    if (byteLength !== expectBytes) {
      throw new Error(
        `manifest: tensor "${name}" byteLength ${byteLength} != shape·dtype ${expectBytes}`);
    }
    if (byteOffset % 256 !== 0) {
      throw new Error(`manifest: tensor "${name}" byteOffset ${byteOffset} is not 256-aligned`);
    }
    if (tensors.has(name)) throw new Error(`manifest: duplicate tensor name "${name}"`);
    if (byteOffset < prevEnd) {
      throw new Error(
        `manifest: tensor "${name}" (offset ${byteOffset}) overlaps or is not ascending ` +
        `after "${prevName}" (ends at ${prevEnd})`);
    }
    tensors.set(name, { dtype, shape, elems, byteOffset, byteLength });
    prevEnd = byteOffset + byteLength;
    prevName = name;
  }

  if (json.bins !== undefined) {
    if (!Array.isArray(json.bins) || json.bins.length === 0) {
      throw new Error('manifest: bins must be a non-empty array');
    }
    for (const b of json.bins) {
      if (typeof b?.file !== 'string' || !b.file) {
        throw new Error('manifest: bins[].file must be a non-empty string');
      }
      if (!Number.isInteger(b.byteLength) || b.byteLength <= 0) {
        throw new Error(`manifest: bins "${b.file}" byteLength must be a positive integer`);
      }
      // Optional (older deploys lack it): full SHA-256 of the part, checked
      // after every landed part in fetchWeightsBin.
      if (b.sha256 !== undefined &&
          (typeof b.sha256 !== 'string' || !/^[0-9a-f]{64}$/.test(b.sha256))) {
        throw new Error(`manifest: bins "${b.file}" sha256 must be 64 lowercase hex chars`);
      }
    }
  }

  return { model: json.model, tensors };
}

export function expectInventory(tensorsMap, { encLayers = 8, decLayers = 2 } = {}) {
  const expected = new Set(['shared.weight', 'pos_embed', 'final_logits_bias']);
  for (let l = 0; l < encLayers; l++) {
    for (const mod of ['qkv', 'out', 'fc1', 'fc2', 'ln1', 'ln2']) {
      expected.add(`enc.${l}.${mod}.weight`);
      expected.add(`enc.${l}.${mod}.bias`);
    }
  }
  for (let l = 0; l < decLayers; l++) {
    for (const mod of ['self_qkv', 'self_out', 'cross_q', 'cross_kv', 'cross_out',
      'fc1', 'fc2', 'ln1', 'ln2', 'ln3']) {
      expected.add(`dec.${l}.${mod}.weight`);
      expected.add(`dec.${l}.${mod}.bias`);
    }
  }

  const missing = [...expected].filter((n) => !tensorsMap.has(n));
  const unexpected = [...tensorsMap.keys()].filter((n) => !expected.has(n));
  if (missing.length || unexpected.length) {
    const parts = [];
    if (missing.length) parts.push(`missing: ${missing.join(', ')}`);
    if (unexpected.length) parts.push(`unexpected: ${unexpected.join(', ')}`);
    throw new Error(`weights inventory mismatch — ${parts.join('; ')}`);
  }
}

// --- CPU-side fp32 expansion (for devices without shader-f16) ---

function totalBytes(tensors) {
  let end = 0;
  for (const t of tensors.values()) end = Math.max(end, t.byteOffset + t.byteLength);
  return end;
}

// Rewrites the bin so every f16 tensor becomes f32. New offsets are assigned by
// walking tensors in ascending original order, 256-aligning each. f32 tensors
// (final_logits_bias) are copied unchanged.
function expandBinToF32(tensors, binBytes) {
  // Compute new layout first (Map iteration order == ascending, enforced by parseManifest).
  let cursor = 0;
  const newTensors = new Map();
  for (const [name, t] of tensors) {
    const byteOffset = alignUp(cursor, 256);
    const byteLength = t.elems * 4;
    newTensors.set(name, { dtype: 'f32', shape: t.shape, elems: t.elems, byteOffset, byteLength });
    cursor = byteOffset + byteLength;
  }
  const out = new Uint8Array(alignUp(cursor, 4));
  for (const [name, t] of tensors) {
    const nt = newTensors.get(name);
    if (t.dtype === 'f16') {
      // binBytes is freshly allocated, so t.byteOffset is even within its buffer.
      const src = new Uint16Array(binBytes.buffer, binBytes.byteOffset + t.byteOffset, t.elems);
      const dst = new Float32Array(out.buffer, nt.byteOffset, t.elems);
      dst.set(expandF16(src));
    } else {
      out.set(binBytes.subarray(t.byteOffset, t.byteOffset + t.byteLength), nt.byteOffset);
    }
  }
  return { tensors: newTensors, bytes: out };
}

// --- INT8 lm_head quantization (W8A16) ---

// Symmetric per-row int8 quantization, 4 values packed per u32 along the
// column axis (little-endian lane order — lane l holds column c4+l, matching
// gemm_tiled2's WQ8 shift-unpack). Rows here = vocab entries of
// shared.weight [24000, 448]; the per-row scale factors out of the lm_head
// dot product and is applied in the kernel epilogue. Clamped to ±127 (never
// -128) so dequantization is exactly q·scale.
function quantizeQ8Row(data, dataBase, cols, packed, packedBase, scales, scaleIndex) {
  let m = 0;
  for (let c = 0; c < cols; c++) m = Math.max(m, Math.abs(data[dataBase + c]));
  const scale = m > 0 ? m / 127 : 1; // all-zero row: q stays 0, any scale works
  scales[scaleIndex] = scale;
  for (let c4 = 0; c4 < cols; c4 += 4) {
    let word = 0;
    for (let l = 0; l < 4; l++) {
      const q = Math.max(-127, Math.min(127, Math.round(data[dataBase + c4 + l] / scale)));
      word |= (q & 0xff) << (8 * l);
    }
    packed[packedBase + c4 / 4] = word >>> 0;
  }
}

export function quantizeQ8Rows(data, rows, cols) {
  if (cols % 4 !== 0) throw new Error(`quantizeQ8Rows: cols=${cols} not a multiple of 4`);
  const packed = new Uint32Array((rows * cols) / 4);
  const scales = new Float32Array(rows);
  const wordsPerRow = cols / 4;
  for (let r = 0; r < rows; r++) {
    quantizeQ8Row(data, r * cols, cols, packed, r * wordsPerRow, scales, r);
  }
  return { packed, scales };
}

// Row-at-a-time f16 decode for Q8 tensors. Unlike expandF16()+gather, this
// keeps only one decoded row alive while writing directly into the final Q8
// arrays. The shared quantizeQ8Row helper makes scale rounding and byte
// packing exactly identical to the legacy Float32 path.
export function quantizeQ8RowsFromF16(src, srcRows, cols, rowIds = null) {
  if (cols % 4 !== 0) {
    throw new Error(`quantizeQ8RowsFromF16: cols=${cols} not a multiple of 4`);
  }
  if (!Number.isInteger(srcRows) || srcRows < 0 || !Number.isInteger(cols) || cols <= 0) {
    throw new Error(`quantizeQ8RowsFromF16: invalid geometry rows=${srcRows} cols=${cols}`);
  }
  const needed = srcRows * cols;
  if (src.length < needed) {
    throw new Error(`quantizeQ8RowsFromF16: source too short: ${src.length} < ${needed}`);
  }
  const rows = rowIds ? rowIds.length : srcRows;
  const wordsPerRow = cols / 4;
  const packed = new Uint32Array(rows * wordsPerRow);
  const scales = new Float32Array(rows);
  const decoded = new Float32Array(cols);
  for (let r = 0; r < rows; r++) {
    const srcRow = rowIds ? rowIds[r] : r;
    if (!Number.isInteger(srcRow) || srcRow < 0 || srcRow >= srcRows) {
      throw new Error(`quantizeQ8RowsFromF16: row id ${srcRow} out of range ${srcRows}`);
    }
    const srcBase = srcRow * cols;
    for (let c = 0; c < cols; c++) decoded[c] = f16ToF32(src[srcBase + c]);
    quantizeQ8Row(decoded, 0, cols, packed, r * wordsPerRow, scales, r);
  }
  return { packed, scales, temporaryBytes: decoded.byteLength };
}

// --- GPU upload ---

// Core, fetch-free: takes the manifest JSON object and the raw bin bytes.
// Exported separately so tests can drive it with synthetic data.
// lmHeadQ8: additionally quantize shared.weight to per-row-scaled int8
// (quantizeQ8Rows) and append the packed words + f32 scales to the SAME
// weights buffer as synthetic tensors 'lm_head.q8' / 'lm_head.scales' —
// bindingFor and the buffer-destroy story stay unchanged.
export function uploadParsed(device, manifestJson, binBytes, { targetDtype = 'f16', lmHeadQ8 = false, ffnQ8 = false, ffnWT = false, projWT = false, lmHeadIds = null, directQ8 = true } = {}) {
  if (targetDtype !== 'f16' && targetDtype !== 'f32') {
    throw new Error(`unsupported targetDtype "${targetDtype}"`);
  }
  if (lmHeadIds && !lmHeadQ8) {
    throw new Error('lmHeadIds needs lmHeadQ8: true (the shortlist repacks the q8 lm_head only)');
  }
  const { model, tensors: parsed } = parseManifest(manifestJson);
  // Make this manifest's geometry the ACTIVE engine config (dims, caps, token
  // ids, derived scales) — see constants.js. Validates before any GPU work.
  const cfg = applyModelConfig(model);
  const need = totalBytes(parsed);
  if (binBytes.byteLength < need) {
    throw new Error(`weights.bin too short: ${binBytes.byteLength} < ${need}`);
  }

  let tensors = parsed;
  let bytes = binBytes;
  if (targetDtype === 'f32') {
    ({ tensors, bytes } = expandBinToF32(parsed, binBytes));
  }

  // Quantization specs: shared.weight is already [N,K] (vocab-major); the
  // FFN weights are stored [K_in, N_out] and get a CPU transpose so all q8
  // tensors share the [N,K] row-packed layout of the WQ8 kernel paths.
  const q8Specs = [];
  if (lmHeadQ8) q8Specs.push({ base: 'lm_head', src: 'shared.weight', transpose: false });
  if (ffnQ8) {
    for (let l = 0; l < cfg.decLayers; l++) {
      for (const mod of ['fc1', 'fc2']) {
        q8Specs.push({ base: `dec.${l}.${mod}`, src: `dec.${l}.${mod}.weight`, transpose: true });
      }
    }
  }
  // Shortlist: pad the emittable-id list to a whole number of BN=64 column
  // tiles so the tiled lm_head never sees a ragged N. Pad slots reuse the eos
  // row's weights and get a floor bias below, so they can never win argmax.
  let lmShort = null;
  if (lmHeadIds) {
    const seen = new Set();
    for (const id of lmHeadIds) {
      if (!Number.isInteger(id) || id < 0 || id >= cfg.vocab) {
        throw new Error(`lmHeadIds: id ${id} out of vocab ${cfg.vocab}`);
      }
      if (seen.has(id)) throw new Error(`lmHeadIds: duplicate id ${id}`);
      seen.add(id);
    }
    for (const [f, v] of [['eos', cfg.eos], ['pad', cfg.pad], ['decoderStart', cfg.decoderStart]]) {
      if (!seen.has(v)) throw new Error(`lmHeadIds must contain ${f} (${v})`);
    }
    const real = lmHeadIds.length;
    const padded = new Uint32Array(Math.ceil(real / 64) * 64).fill(cfg.eos);
    padded.set(lmHeadIds);
    lmShort = { ids: padded, real };
  }

  const q8Blobs = [];
  const q8Stats = {
    directEnabled: directQ8,
    directTensors: 0,
    sourceBytes: 0,
    packedBytes: 0,
    scaleBytes: 0,
    peakTemporaryBytes: 0,
    legacyPeakTemporaryBytes: 0,
    temporaryBytesSaved: 0,
    quantizeMs: 0,
    tensors: [],
  };
  for (const spec of q8Specs) {
    const t = tensors.get(spec.src);
    if (!t) throw new Error(`q8: no ${spec.src} tensor`);
    const started = globalThis.performance?.now?.() ?? Date.now();
    const sourceBytes = t.byteLength;
    let temporaryBytes = 0;
    let legacyTemporaryBytes = 0;
    let direct = false;
    let quantized;
    let [rows, cols] = t.shape;

    // Production lm_head: decode one requested f16 row at a time, including
    // repeated EOS padding rows, and write straight into the final Q8 arrays.
    // Transposed FFN Q8 and f32 sources retain the established path.
    if (directQ8 && t.dtype === 'f16' && !spec.transpose) {
      const src = new Uint16Array(bytes.buffer, bytes.byteOffset + t.byteOffset, t.elems);
      const rowIds = spec.base === 'lm_head' && lmShort ? lmShort.ids : null;
      quantized = quantizeQ8RowsFromF16(src, rows, cols, rowIds);
      if (rowIds) rows = rowIds.length;
      direct = true;
      temporaryBytes = quantized.temporaryBytes;
      // What the compatibility path holds simultaneously: the whole decoded
      // tensor plus, for a shortlist, its gathered Float32 row matrix.
      legacyTemporaryBytes = t.elems * Float32Array.BYTES_PER_ELEMENT
        + (rowIds ? rows * cols * Float32Array.BYTES_PER_ELEMENT : 0);
    } else {
      const f = t.dtype === 'f16'
        ? expandF16(new Uint16Array(bytes.buffer, bytes.byteOffset + t.byteOffset, t.elems))
        : new Float32Array(bytes.buffer, bytes.byteOffset + t.byteOffset, t.elems);
      temporaryBytes += t.dtype === 'f16' ? f.byteLength : 0;
      let data = f;
      if (spec.transpose) {
        const [K, N] = t.shape;
        data = new Float32Array(t.elems);
        temporaryBytes += data.byteLength;
        for (let k = 0; k < K; k++) {
          for (let n = 0; n < N; n++) data[n * K + k] = f[k * N + n];
        }
        rows = N;
        cols = K;
      }
      if (spec.base === 'lm_head' && lmShort) {
        const g = new Float32Array(lmShort.ids.length * cols);
        temporaryBytes += g.byteLength;
        for (let i = 0; i < lmShort.ids.length; i++) {
          g.set(data.subarray(lmShort.ids[i] * cols, (lmShort.ids[i] + 1) * cols), i * cols);
        }
        data = g;
        rows = lmShort.ids.length;
      }
      quantized = quantizeQ8Rows(data, rows, cols);
      legacyTemporaryBytes = temporaryBytes;
    }

    const quantizeMs = (globalThis.performance?.now?.() ?? Date.now()) - started;
    const tensorStats = {
      name: spec.base,
      rows,
      cols,
      direct,
      sourceBytes,
      temporaryBytes,
      legacyTemporaryBytes,
      temporaryBytesSaved: Math.max(0, legacyTemporaryBytes - temporaryBytes),
      packedBytes: quantized.packed.byteLength,
      scaleBytes: quantized.scales.byteLength,
      quantizeMs,
    };
    q8Stats.directTensors += direct ? 1 : 0;
    q8Stats.sourceBytes += sourceBytes;
    q8Stats.packedBytes += tensorStats.packedBytes;
    q8Stats.scaleBytes += tensorStats.scaleBytes;
    q8Stats.peakTemporaryBytes = Math.max(q8Stats.peakTemporaryBytes, temporaryBytes);
    q8Stats.legacyPeakTemporaryBytes = Math.max(
      q8Stats.legacyPeakTemporaryBytes, legacyTemporaryBytes,
    );
    q8Stats.quantizeMs += quantizeMs;
    q8Stats.tensors.push(tensorStats);
    q8Blobs.push({ spec, rows, cols, packed: quantized.packed, scales: quantized.scales });
  }
  q8Stats.temporaryBytesSaved = Math.max(
    0, q8Stats.legacyPeakTemporaryBytes - q8Stats.peakTemporaryBytes,
  );

  // Shortlist companions: the local→global id map, and the logits bias
  // gathered into list order (same dtype as final_logits_bias so the argmax
  // epilogue's binding just points at a different tensor). Pad slots get a
  // finite floor (-65504 f16 / -1e30 f32) — never argmax, no inf arithmetic.
  const lmBlobs = [];
  if (lmShort) {
    const bt = tensors.get('final_logits_bias');
    if (!bt) throw new Error('lmHeadIds: no final_logits_bias tensor');
    const View = bt.dtype === 'f16' ? Uint16Array : Float32Array;
    const src = new View(bytes.buffer, bytes.byteOffset + bt.byteOffset, bt.elems);
    const sb = new View(lmShort.ids.length);
    const FLOOR = bt.dtype === 'f16' ? 0xFBFF : -1e30; // f16 bits for -65504
    for (let i = 0; i < lmShort.ids.length; i++) {
      sb[i] = i < lmShort.real ? src[lmShort.ids[i]] : FLOOR;
    }
    lmBlobs.push({ name: 'lm_head.idmap', dtype: 'u32', shape: [lmShort.ids.length], data: lmShort.ids });
    lmBlobs.push({ name: 'lm_head.sbias', dtype: bt.dtype, shape: [lmShort.ids.length], data: sb });
  }

  // ffnWT / projWT: value-preserving TRANSPOSED copies of decode weights
  // ([K,N] → [N,K], same element type) so those sites can run the GEMV WT /
  // tiled wt paths (per-workgroup W-tile reuse for MT rows) with full f16
  // precision. ffnWT covers fc1/fc2 (+3.2MB); projWT the four attention-side
  // projections — self_qkv, self_out, cross_q, cross_out (+4.8MB), whose NWT
  // per-row W re-reads were 46.6% of the b128 step (prod_profile 2026-07-06).
  const wtBlobs = [];
  const wtMods = [
    ...(ffnWT ? ['fc1', 'fc2'] : []),
    ...(projWT ? ['self_qkv', 'self_out', 'cross_q', 'cross_out'] : []),
  ];
  if (wtMods.length) {
    const eb = BYTES[targetDtype];
    const View = targetDtype === 'f16' ? Uint16Array : Float32Array;
    for (let l = 0; l < cfg.decLayers; l++) {
      for (const mod of wtMods) {
        const name = `dec.${l}.${mod}.weight`;
        const t = tensors.get(name);
        if (!t) throw new Error(`wt: no ${name} tensor`);
        const [K, N] = t.shape;
        const src = new View(bytes.buffer, bytes.byteOffset + t.byteOffset, t.elems);
        const out = new View(t.elems);
        for (let k = 0; k < K; k++) {
          for (let n = 0; n < N; n++) out[n * K + k] = src[k * N + n];
        }
        wtBlobs.push({ name: `dec.${l}.${mod}.wt`, dtype: targetDtype, shape: [N, K], data: out, eb });
      }
    }
  }

  // writeBuffer size must be a multiple of 4; pad if the bin isn't.
  let size = alignUp(totalBytes(tensors), 4);
  if (bytes.byteLength < size) {
    const padded = new Uint8Array(size);
    padded.set(bytes.subarray(0, totalBytes(tensors)));
    bytes = padded;
  }
  const binSize = size;
  for (const blob of q8Blobs) {
    blob.qOff = alignUp(size, 256);
    blob.sOff = alignUp(blob.qOff + blob.packed.byteLength, 256);
    size = blob.sOff + blob.scales.byteLength;
    tensors.set(`${blob.spec.base}.q8`, {
      dtype: 'u32', shape: [blob.rows, blob.cols], elems: blob.packed.length,
      byteOffset: blob.qOff, byteLength: blob.packed.byteLength,
    });
    tensors.set(`${blob.spec.base}.scales`, {
      dtype: 'f32', shape: [blob.rows], elems: blob.scales.length,
      byteOffset: blob.sOff, byteLength: blob.scales.byteLength,
    });
  }
  for (const blob of [...wtBlobs, ...lmBlobs]) {
    blob.off = alignUp(size, 256);
    size = blob.off + blob.data.byteLength;
    tensors.set(blob.name, {
      dtype: blob.dtype, shape: blob.shape, elems: blob.data.length,
      byteOffset: blob.off, byteLength: blob.data.byteLength,
    });
  }
  size = alignUp(size, 4);
  const buffer = device.createBuffer({
    size,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
  });
  device.queue.writeBuffer(buffer, 0, bytes, 0, binSize);
  for (const blob of q8Blobs) {
    device.queue.writeBuffer(buffer, blob.qOff, blob.packed);
    device.queue.writeBuffer(buffer, blob.sOff, blob.scales);
  }
  for (const blob of [...wtBlobs, ...lmBlobs]) device.queue.writeBuffer(buffer, blob.off, blob.data);

  return {
    model,
    dtype: targetDtype,
    buffer,
    tensors,
    q8Stats,
    // CPU copy of the shortlist (padded ids + real length): decode state
    // needs local indices (DECODER_START pre-set) without a GPU readback.
    lmHeadShort: lmShort,
    bindingFor(name) {
      const t = tensors.get(name);
      if (!t) throw new Error(`unknown tensor "${name}"`);
      return { buffer, offset: t.byteOffset, size: t.byteLength };
    },
  };
}

// Fetch the weight blob(s) described by a manifest into one aligned buffer.
// Without manifest.bins this is the original single weights.bin streaming
// path; with bins (sharded deploys — Cloudflare Pages caps files at 25MiB)
// parts are fetched sequentially into their offsets. onProgress(loaded,
// total) is cumulative across parts.
// Parts are fetched over this many connections at once: a single connection
// nowhere near saturates the CDNs we serve from (HF Xet ~3MB/s, Cloudflare
// ~10MB/s measured), so the parallelism roughly halves model load time.
const CONCURRENT_PARTS = 4;

// Cross-visit shard cache. Shard filenames are content-addressed (sha in the
// name), so entries can never go stale — a redeploy changes the names and the
// old ones are pruned below. This matters on hosts whose CDN serves weights
// via short-lived signed redirect URLs the HTTP cache can't reuse (HF spaces
// re-download the full model on every visit without it). The single-bin dev
// path is NOT content-addressed and is never cached.
const WEIGHTS_CACHE = 'moxhi-weights-v1';

async function openWeightsCache() {
  try {
    if (typeof caches === 'undefined') return null;
    return await caches.open(WEIGHTS_CACHE);
  } catch {
    return null; // private mode / storage denied — plain network fetch
  }
}

// Content check for a landed part. The byteLength check cannot catch an entry
// whose bytes rotted at the declared size, and a wrong shard is silent garbage
// inference. Manifests without per-part sha256 (older deploys, dev) and pages
// without crypto.subtle (insecure context) skip this, keeping size-only.
async function assertPartSha(part, view) {
  const subtle = globalThis.crypto?.subtle;
  if (!part.sha256 || !subtle) return;
  const digest = new Uint8Array(await subtle.digest('SHA-256', view));
  let hex = '';
  for (const b of digest) hex += b.toString(16).padStart(2, '0');
  if (hex !== part.sha256) throw new Error(`${part.file}: SHA-256 mismatch`);
}

export async function fetchWeightsBin(baseUrl, manifestJson, { onProgress } = {}) {
  const { tensors } = parseManifest(manifestJson);
  const need = totalBytes(tensors);
  const parts = manifestJson.bins ?? null;

  // Parts land out of order, so progress is an aggregate byte counter —
  // monotonic, but not "prefix of the file complete".
  let loaded = 0;
  const streamInto = async (res, bin, offset, budget, total) => {
    let got = 0;
    const reader = res.body.getReader();
    for (;;) {
      const { done, value } = await reader.read();
      if (done) break;
      if (got + value.byteLength > budget) {
        throw new Error(`weights part larger than expected ${budget} bytes`);
      }
      bin.set(value, offset + got);
      got += value.byteLength;
      loaded += value.byteLength;
      if (onProgress) onProgress(loaded, total);
    }
    return got;
  };

  if (!parts) {
    const res = await fetch(`${baseUrl}/weights.bin`);
    if (!res.ok) throw new Error(`fetch weights.bin: HTTP ${res.status}`);
    const total = Number(res.headers.get('content-length')) || need;
    const bin = new Uint8Array(alignUp(Math.max(total, need), 4));
    const got = await streamInto(res, bin, 0, bin.byteLength, total);
    if (got < need) throw new Error(`weights.bin truncated: got ${got} of ${need} bytes`);
    return bin;
  }

  const total = parts.reduce((s, p) => s + p.byteLength, 0);
  if (total < need) throw new Error(`manifest.bins total ${total} < tensors need ${need} bytes`);
  const bin = new Uint8Array(alignUp(total, 4));
  const offsets = [];
  for (let off = 0, i = 0; i < parts.length; off += parts[i++].byteLength) offsets.push(off);
  const cache = await openWeightsCache();
  const loadPart = async (i) => {
    const part = parts[i];
    const url = `${baseUrl}/${part.file}`;
    for (let attempt = 0; ; attempt++) {
      let res = null;
      if (cache && attempt === 0) {
        try { res = (await cache.match(url)) ?? null; } catch { res = null; }
      }
      const fromCache = res !== null;
      if (!res) {
        res = await fetch(url);
        if (!res.ok) throw new Error(`fetch ${part.file}: HTTP ${res.status}`);
      }
      let got = 0;
      try {
        got = await streamInto(res, bin, offsets[i], part.byteLength, total);
        if (got !== part.byteLength) {
          throw new Error(`${part.file}: got ${got} of declared ${part.byteLength} bytes`);
        }
        await assertPartSha(part, bin.subarray(offsets[i], offsets[i] + got));
      } catch (err) {
        loaded -= got; // undo this attempt's progress
        if (fromCache) {
          try { await cache.delete(url); } catch { /* ignore */ }
          continue; // corrupt/truncated cache entry — refetch from network
        }
        throw err;
      }
      if (cache && !fromCache) {
        try {
          await cache.put(url, new Response(bin.subarray(offsets[i], offsets[i] + got)));
        } catch { /* quota exceeded — keep serving from network */ }
      }
      return;
    }
  };
  let next = 0;
  const worker = async () => {
    for (;;) {
      const i = next++;
      if (i >= parts.length) return;
      await loadPart(i);
    }
  };
  await Promise.all(Array.from({ length: Math.min(CONCURRENT_PARTS, parts.length) }, worker));
  if (cache) await pruneWeightsCache(cache, baseUrl, parts);
  return bin;
}

// Drop same-directory entries that are not in the current manifest (old
// shas after a redeploy). Other models live in subdirectories — untouched.
async function pruneWeightsCache(cache, baseUrl, parts) {
  try {
    const names = new Set(parts.map((p) => p.file));
    const prefix = `${baseUrl}/`;
    for (const req of await cache.keys()) {
      const path = new URL(req.url).pathname;
      if (!path.startsWith(prefix)) continue;
      const rest = path.slice(prefix.length);
      if (rest.includes('/') || names.has(rest)) continue;
      await cache.delete(req);
    }
  } catch { /* best-effort */ }
}

export async function loadWeights(device, baseUrl = '/weights', { targetDtype = 'f16', lmHeadQ8 = false, ffnQ8 = false, ffnWT = false, projWT = false, lmHeadIds = null, directQ8 = true, onProgress } = {}) {
  const manifestRes = await fetch(`${baseUrl}/manifest.json`);
  if (!manifestRes.ok) throw new Error(`fetch manifest.json: HTTP ${manifestRes.status}`);
  const manifestJson = await manifestRes.json();
  const { model, tensors } = parseManifest(manifestJson);
  expectInventory(tensors, parseModelConfig(model));
  const bin = await fetchWeightsBin(baseUrl, manifestJson, { onProgress });
  const up = uploadParsed(device, manifestJson, bin, {
    targetDtype, lmHeadQ8, ffnQ8, ffnWT, projWT, lmHeadIds, directQ8,
  });
  // Content-version tag for caches keyed on "these exact weights" (the
  // app's row translation memory): sharded deploys carry a content hash in
  // every bin name; the unsharded dev path has none — 'dev' entries may go
  // stale across a local weights swap, production entries cannot.
  up.weightsTag = manifestJson.bins?.[0]?.file ?? 'dev';
  return up;
}