forward.js

/* forward.js — a real llama-family forward pass in the browser, over a GGUF's real weight tensors.
 *
 * VENDORED from `templates/static/hdc/forward.js` (same repository, same author) for the DMCRA Lab
 * project, with exactly ONE change: the quantiser is now a callback (`opts.quantize`) instead of a
 * hardcoded ternary branch. The hook is marked "THE DMCRA HOOK" below. With no `quantize` callback
 * and no `variant`, behaviour is unchanged, so the vendored maths is still the maths the HDC Space
 * measures. The HDC Space itself is frozen and untouched.
 *
 * It was measured before it was bundle-loadable: next-token top-1 of 16.9% against a uniform chance of
 * 0.20% over the 512-symbol stories260K vocabulary, and 5.6% with ternarised weights. This module
 * exists so the PAGE can run that pass rather than only the test suite. It is a port, not a second
 * implementation — the layer order, the caching and the accounting below are the measured ones, and
 * `scripts/hdc_tests/forward.test.mjs` loads THIS file (after weights.js, off the global, exactly as
 * index.html does) and asserts it still reproduces those numbers, so the shipped code and the
 * measured code cannot drift apart.
 *
 * Two things it needs that a test-only port did not:
 *   `requirements(header)` is the gate. It reads the tensor INDEX only (no weight bytes) and reports
 *   whether a llama pass is runnable on this file, which required tensors are missing, and which are
 *   present but have no codec — so the page can refuse with a reason instead of throwing mid-pass.
 *   `topK`/`softmax` turn the logits into the ranked probabilities a page has to display.
 *
 * CONVENTIONS KEPT FROM THE REFERENCE (stated because they are choices, not facts):
 *   RoPE     rotate-half (NeoX): the pair rotated together is (x[i], x[i + head_dim/2]), with
 *            theta_i = pos * freq_base^(-2i/head_dim) for i in [0, head_dim/2), at absolute
 *            positions 0..n-1. llama.cpp selects this for LLM_ARCH_LLAMA. The other common
 *            convention pairs (x[2i], x[2i+1]) and is deliberately NOT used.
 *   Scale    1/sqrt(head_dim) on the Q·K scores, per head.
 *   GQA      query head h reads KV head floor(h / (n_head / n_head_kv)).
 *   Mask     causal and inclusive: position t attends over j <= t.
 *   Layout   GGUF stores [in, out] with `in` contiguous, so y[o] = sum_i W[o*in + i] * x[i], which is
 *            exactly `HDCWeights.exactLinear(x, values, inDim)` — every projection goes through it,
 *            so this file cannot disagree with the layer the HDC suite measures.
 *   Ternary  (arXiv 2402.17764, BitNet b1.58) applies to 2-D weight matrices only, through the
 *            project's own `toTernary`/`fromTernary`. Norms and embeddings stay float: a ternary norm
 *            would be a different model rather than a quantised one.
 *
 * KNOWN LIMIT, carried over rather than papered over: the pass adds NO attention biases. The reference
 * ignores them, so on a file that carries them (Qwen2 does) the output is not that model's exact
 * output. `requirements()` reports that condition and sets `exact: false` for it.
 *
 * Dependencies are taken off the global object (like every other bundle here): `HDCWeights` for
 * `readTensor`/`exactLinear`/`toTernary`/`fromTernary`/`tensorByteLength`, and `HDCGGUF` only for type
 * names. No GGUF parsing happens in this file.
 */
(function (root) {
  "use strict";

  var W = root.HDCWeights, G = root.HDCGGUF;
  var NEED_W = "forward.js requires HDCWeights (weights.js) to be loaded first";

  /* ---------------------------------------------------------------- metadata */

  /* `parse` returns `kv` as an ARRAY of {key, type, value}, not an object — so index it explicitly. */
  function kvMapFromHeader(header) {
    var d = {}, kv = (header && header.kv) || [], i;
    for (i = 0; i < kv.length; i++) d[kv[i].key] = kv[i].value;
    return d;
  }

  function num(d, k, def) { return typeof d[k] === "number" ? d[k] : def; }
  function str(d, k, def) { return typeof d[k] === "string" ? d[k] : def; }

  function normalise(m, defaults) {
    var headDim = Math.floor(m.nEmbd / m.nHead);
    var out = {}, k;
    for (k in m) if (Object.prototype.hasOwnProperty.call(m, k)) out[k] = m[k];
    out.headDim = headDim;
    out.gqaGroups = Math.max(1, Math.floor(m.nHead / m.nHeadKv));
    out.defaulted = (m.defaulted || []).concat(defaults || []);
    return out;
  }

  /* Accepts a header from HDCGGUF.parse(), a raw KV map, or an already-resolved meta object.
   * Only what the file genuinely omits is defaulted, and each default is recorded in `defaulted`. */
  function resolveMeta(input) {
    var d = {}, i;
    if (input && input.meta && typeof input.meta === "object" && !Array.isArray(input.meta)) {
      if (typeof input.meta.nEmbd === "number") return normalise(input.meta, []);
      d = input.meta;
    } else if (input && (input.header || input.kv)) {
      d = (input.header && input.header.kv) ? kvMapFromHeader(input.header) : input.kv;
      if (Array.isArray(d)) d = kvMapFromHeader({ kv: d });
    }
    var has = function (k) { return d[k] !== undefined; };
    if (!Object.keys(d).length && input && typeof input === "object" && !Array.isArray(input) &&
        (has("general.architecture") || has("llama.embedding_length") || has("general.name"))) {
      d = input;                                  // a plain KV map given directly
    }

    var defaults = [];
    var nEmbd = num(d, "llama.embedding_length", 64);
    var eps = num(d, "llama.attention.layer_norm_rms_epsilon", 1e-5);
    if (d["llama.attention.layer_norm_rms_epsilon"] === undefined) defaults.push("llama.attention.layer_norm_rms_epsilon=1e-5");
    var ropeBase = num(d, "llama.rope.freq_base", 10000);
    if (d["llama.rope.freq_base"] === undefined) defaults.push("llama.rope.freq_base=10000");
    var nHead = num(d, "llama.attention.head_count", nEmbd / 16);
    if (d["llama.attention.head_count"] === undefined) defaults.push("llama.attention.head_count=" + nHead);
    var nHeadKv = num(d, "llama.attention.head_count_kv", nHead);
    if (d["llama.attention.head_count_kv"] === undefined) defaults.push("llama.attention.head_count_kv=n_head");
    var nLayer = num(d, "llama.block_count", 5);
    if (d["llama.block_count"] === undefined) defaults.push("llama.block_count=5");
    var nFfn = num(d, "llama.feed_forward_length", 172);
    if (d["llama.feed_forward_length"] === undefined) defaults.push("llama.feed_forward_length=172");
    return normalise({
      nLayer: nLayer, nEmbd: nEmbd, nHead: nHead, nHeadKv: nHeadKv, nFfn: nFfn, eps: eps,
      ropeBase: ropeBase, vocab: num(d, "llama.vocab_size", 0),
      arch: str(d, "general.architecture", "llama")
    }, defaults);
  }

  /* The tensors a llama pass reads, in the order it reads them. */
  function requiredTensors(nLayer) {
    var req = ["token_embd.weight", "output_norm.weight", "output.weight"], L, p;
    for (L = 0; L < nLayer; L++) {
      p = "blk." + L + ".";
      req.push(p + "attn_q.weight", p + "attn_k.weight", p + "attn_v.weight", p + "attn_output.weight",
               p + "attn_norm.weight", p + "ffn_gate.weight", p + "ffn_up.weight", p + "ffn_down.weight");
    }
    return req;
  }

  /* The page's gate: can this file be run, and if not, why not? Reads the tensor index only. */
  function requirements(header) {
    var meta = resolveMeta(header);
    var tensors = (header && header.tensors) || [];
    var index = {}, i, nm;
    for (i = 0; i < tensors.length; i++) index[tensors[i].name] = tensors[i];

    var req = requiredTensors(meta.nLayer);
    var missing = [], unsupported = [], params = 0, t;
    for (i = 0; i < req.length; i++) {
      t = index[req[i]];
      if (!t) { missing.push(req[i]); continue; }
      params += Number(t.params) || 0;
      if (W && W.tensorByteLength(t) === null) unsupported.push(req[i]);
    }

    var biases = [];
    for (i = 0; i < tensors.length; i++) {
      nm = tensors[i].name;
      if (/^blk\.\d+\.attn_[qkv]\.bias$/.test(nm)) biases.push(nm);
    }

    var vocab = 0;
    t = index["token_embd.weight"];
    if (t && t.dims && t.dims.length > 1) vocab = Number(t.dims[1]);

    var reason = null;
    if (!tensors.length) {
      reason = "this header has no tensor index — it was parsed with metadataOnly, so the weights " +
               "are not described; re-read the file without stopping at the metadata block";
    } else if (missing.length) {
      reason = missing.length + " required tensor(s) absent, e.g. " + missing.slice(0, 4).join(", ");
    } else if (unsupported.length) {
      reason = "no codec for " + unsupported.length + " required tensor(s), e.g. " +
               unsupported.slice(0, 3).join(", ");
    } else if (!(vocab > 0)) {
      reason = "token_embd.weight carries no vocabulary dimension";
    }

    return {
      feasible: reason === null,
      reason: reason,
      exact: biases.length === 0,            // false when attention biases exist and are not added
      missing: missing,
      tensors: missing,                      // same list, under the name the brief specified
      unsupported: unsupported,
      biasesIgnored: biases,
      note: biases.length
        ? biases.length + " attention bias tensor(s) are present and are NOT added by this pass, so " +
          "the output is not this model's exact output (the reference ignores them too)"
        : null,
      params: params,
      vocab: vocab,
      layers: meta.nLayer,
      meta: meta
    };
  }

  /* ---------------------------------------------------------------- primitives */

  function rmsNorm(x, weight, eps, out) {
    var ss = 0, i;
    for (i = 0; i < x.length; i++) ss += x[i] * x[i];
    var inv = 1 / Math.sqrt(ss / x.length + eps);
    for (i = 0; i < x.length; i++) out[i] = x[i] * inv * weight[i];
    return out;
  }

  /* Rotate-half RoPE, in place, on one head's slice. */
  function ropeHead(vec, base, headDim, pos, freqBase) {
    var half = headDim >> 1, i;
    for (i = 0; i < half; i++) {
      var theta = pos * Math.pow(freqBase, -2 * i / headDim);
      var c = Math.cos(theta), s = Math.sin(theta);
      var a = vec[base + i], b = vec[base + i + half];
      vec[base + i] = a * c - b * s;
      vec[base + i + half] = a * s + b * c;
    }
  }

  function silu(v) { return v / (1 + Math.exp(-v)); }

  /* Uses the project's own toTernary/fromTernary so this file cannot disagree with the rest of the
   * measurement suite about what "ternary" means: per-block scale = mean(|w|), levels {-1,0,+1}. */
  function ternarise(values, blockSize) {
    var t = W.toTernary(values, { blockSize: blockSize });
    return { values: W.fromTernary(t, blockSize), bits: t.perValueBits, scale: t.scale,
             blocks: t.scales.length };
  }

  /* ---- THE DMCRA HOOK (the only edit to this vendored file).
   * The quantiser becomes a callback so the elastic level-set ladder of quant.js can run through
   * this same already-verified pass. The old behaviour is preserved exactly: no callback and no
   * variant quantises nothing, and `variant:"ternary"` still calls ternarise() above. The callback
   * receives the caller's `(values, tensorName)` and returns the reconstruction, so a page can pick
   * a different grid per matrix. It may also return its own bit account; if it does not, the pass
   * falls back to the old 2-bit-planes-plus-scale arithmetic so a caller is never silently
   * over-charged for a code it did not describe. */
  function makeQuantiser(opts) {
    if (typeof opts.quantize === "function") {
      return {
        name: opts.quantizeName || "custom",
        fn: function (values, name) {
          var r = opts.quantize(values, name);
          var bits = null;
          if (r && r.bits && r.bits.effBits !== undefined) bits = r.bits.effBits;
          else if (r && typeof r.bits === "number") bits = r.bits;
          var bytes = null;
          if (r && r.bits && r.bits.storedBytes !== undefined) bytes = r.bits.storedBytes;
          return { values: r.values, bits: bits, bytes: bytes };
        }
      };
    }
    if (opts.variant === "ternary") {
      return { name: "ternary", fn: function (values) { return ternarise(values, opts.blockSize | 0); } };
    }
    return null;
  }

  /* ---------------------------------------------------------------- display helpers */

  function softmax(logits) {
    var n = logits.length, out = new Float32Array(n), i, max = -Infinity, sum = 0, e;
    if (!n) return out;
    for (i = 0; i < n; i++) if (logits[i] > max) max = logits[i];
    for (i = 0; i < n; i++) { e = Math.exp(logits[i] - max); out[i] = e; sum += e; }
    if (sum > 0) for (i = 0; i < n; i++) out[i] /= sum;
    return out;
  }

  /* Ranked next-token candidates: {id, logit, p}. */
  function topK(logits, k) {
    var n = logits.length, idx = [], i, limit;
    for (i = 0; i < n; i++) idx.push(i);
    idx.sort(function (a, b) { return logits[b] - logits[a]; });
    var p = softmax(logits);
    limit = Math.min(k || 8, idx.length);
    var out = [];
    for (i = 0; i < limit; i++) out.push({ id: idx[i], logit: logits[idx[i]], p: p[idx[i]] });
    return out;
  }

  /* Display-only surface form for the llama/SentencePiece encoding, matching the reference's helper.
   * The page has its own decodeToken for the byte-level (gpt2) encoding; this is here so a caller
   * without a vocabulary can still label a prediction. */
  function surface(tok) {
    if (typeof tok !== "string") return String(tok);
    var m = /^<0x([0-9A-Fa-f]{2})>$/.exec(tok);
    if (m) return String.fromCharCode(parseInt(m[1], 16));
    return tok.replace(/\u2581/g, " ");
  }

  /* ---------------------------------------------------------------- the forward pass */

  function makeRead(buffer, header) {
    return function (name) {
      var t = W.readTensor(buffer, header, name);
      if (!t) return null;
      return { values: t.values, dims: t.dims, typeName: t.typeName,
               supported: t.supported, error: t.error };
    };
  }

  function run(buffer, header, tokenIds, opts) {
    if (!W) throw new Error(NEED_W);
    opts = opts || {};
    var variant = opts.variant === "ternary" ? "ternary" : "exact";
    var blockSize = opts.blockSize | 0;
    var read = makeRead(buffer, header);
    var meta = resolveMeta(header);
    var nLayer = meta.nLayer, nEmbd = meta.nEmbd, nHead = meta.nHead, nHeadKv = meta.nHeadKv;
    var nFfn = meta.nFfn, eps = meta.eps, ropeBase = meta.ropeBase;
    var headDim = meta.headDim, gqaGroups = meta.gqaGroups;
    var nTok = tokenIds.length;
    if (!nTok) throw new Error("tokenIds must not be empty");

    var quantiser = makeQuantiser(opts);
    /* One cache per call: a tensor is quantised at most once, and both variants share the reading. */
    var cache = new Map();
    var quantisedMatrices = 0, totalMatrices = 0, quantBytes = 0, floatBytes = 0;
    var quantisedBits = 0, quantisedWeights = 0;
    function mat(name) {
      if (cache.has(name)) return cache.get(name);
      var t = read(name);
      if (!t || !t.values) throw new Error("tensor unavailable: " + name +
        (t && t.error ? " (" + t.error + ")" : ""));
      var v = t.values;
      var isMatrix = t.dims && t.dims.length === 2;
      var out = v;
      if (quantiser && isMatrix) {
        var q = quantiser.fn(v, name);
        out = q.values;
        quantisedMatrices++; totalMatrices++;
        floatBytes += v.length * 4;
        quantisedWeights += v.length;
        if (q.bits !== null) quantisedBits += q.bits * v.length;
        quantBytes += q.bytes !== null ? q.bytes
          : Math.ceil(v.length / 8) * 2 +
            Math.ceil(v.length / (blockSize > 0 ? blockSize : v.length)) * 4;
      } else if (isMatrix) { totalMatrices++; floatBytes += v.length * 4; }
      cache.set(name, out);
      return out;
    }
    function vec(name) {
      if (cache.has("v:" + name)) return cache.get("v:" + name);
      var t = read(name);
      if (!t || !t.values) throw new Error("tensor unavailable: " + name +
        (t && t.error ? " (" + t.error + ")" : ""));
      cache.set("v:" + name, t.values);
      return t.values;
    }
    /* 1-D tensors (norms) are read through vec() and never ternarised: a ternary norm would be a
     * different model rather than a quantised one. */

    // ---- embeddings
    var emb = vec("token_embd.weight");
    var embDims = read("token_embd.weight").dims;      // [nEmbd, vocab]
    var nVocab = embDims && embDims.length > 1 ? Number(embDims[1]) : 0;
    var H = new Float32Array(nTok * nEmbd);
    var t0, i0;
    for (t0 = 0; t0 < nTok; t0++) {
      var id = tokenIds[t0];
      if (!(id >= 0 && id < nVocab)) throw new Error("token id " + id + " outside vocabulary of " + nVocab);
      for (i0 = 0; i0 < nEmbd; i0++) H[t0 * nEmbd + i0] = emb[id * nEmbd + i0];
    }

    // ---- layers
    for (var L = 0; L < nLayer; L++) {
      var p = "blk." + L + ".";
      var wq = mat(p + "attn_q.weight");
      var wk = mat(p + "attn_k.weight");
      var wv = mat(p + "attn_v.weight");
      var wo = mat(p + "attn_output.weight");
      var wg = mat(p + "ffn_gate.weight");
      var wu = mat(p + "ffn_up.weight");
      var wd = mat(p + "ffn_down.weight");
      var an = vec(p + "attn_norm.weight");
      var fn = vec(p + "ffn_norm.weight");

      // norm, then the three projections
      var normed = [], t;
      for (t = 0; t < nTok; t++) normed.push(new Float32Array(nEmbd));
      for (t = 0; t < nTok; t++) rmsNorm(H.subarray(t * nEmbd, (t + 1) * nEmbd), an, eps, normed[t]);

      var q = new Float32Array(nTok * nHead * headDim);
      var k = new Float32Array(nTok * nHeadKv * headDim);
      var v = new Float32Array(nTok * nHeadKv * headDim);
      for (t = 0; t < nTok; t++) {
        var qq = W.exactLinear(normed[t], wq, nEmbd);
        var kk = W.exactLinear(normed[t], wk, nEmbd);
        var vv = W.exactLinear(normed[t], wv, nEmbd);
        q.set(qq, t * nHead * headDim);
        k.set(kk, t * nHeadKv * headDim);
        v.set(vv, t * nHeadKv * headDim);
      }

      // positions 0..n-1; llama applies RoPE at the absolute sequence position
      for (t = 0; t < nTok; t++) {
        var h;
        for (h = 0; h < nHead; h++) ropeHead(q, t * nHead * headDim + h * headDim, headDim, t, ropeBase);
        for (h = 0; h < nHeadKv; h++) ropeHead(k, t * nHeadKv * headDim + h * headDim, headDim, t, ropeBase);
      }

      // causal attention, per query head, with GQA
      var attn = new Float32Array(nTok * nEmbd);
      var scale = 1 / Math.sqrt(headDim);
      var scores = new Float32Array(nTok);
      for (t = 0; t < nTok; t++) {
        for (var h2 = 0; h2 < nHead; h2++) {
          var kv = Math.floor(h2 / gqaGroups);
          var qOff = t * nHead * headDim + h2 * headDim;
          var best = -Infinity, j, d, s;
          for (j = 0; j <= t; j++) {
            var kOff = j * nHeadKv * headDim + kv * headDim;
            s = 0;
            for (d = 0; d < headDim; d++) s += q[qOff + d] * k[kOff + d];
            s *= scale;
            scores[j] = s;
            if (s > best) best = s;
          }
          var sum = 0;
          for (j = 0; j <= t; j++) { var e = Math.exp(scores[j] - best); scores[j] = e; sum += e; }
          var oOff = t * nEmbd + h2 * headDim;
          for (j = 0; j <= t; j++) {
            var w = scores[j] / sum;
            if (w === 0) continue;
            var vOff = j * nHeadKv * headDim + kv * headDim;
            for (d = 0; d < headDim; d++) attn[oOff + d] += w * v[vOff + d];
          }
        }
      }

      // attention output projection + residual
      for (t = 0; t < nTok; t++) {
        var proj = W.exactLinear(attn.subarray(t * nEmbd, (t + 1) * nEmbd), wo, nEmbd);
        for (var i = 0; i < nEmbd; i++) H[t * nEmbd + i] += proj[i];
      }

      // SwiGLU MLP + residual
      for (t = 0; t < nTok; t++) {
        var hn = new Float32Array(nEmbd);
        rmsNorm(H.subarray(t * nEmbd, (t + 1) * nEmbd), fn, eps, hn);
        var gate = W.exactLinear(hn, wg, nEmbd);
        var up = W.exactLinear(hn, wu, nEmbd);
        var act = new Float32Array(nFfn);
        for (var a = 0; a < nFfn; a++) act[a] = silu(gate[a]) * up[a];
        var down = W.exactLinear(act, wd, nFfn);
        for (var o = 0; o < nEmbd; o++) H[t * nEmbd + o] += down[o];
      }
    }

    // ---- final norm and unembedding, from the LAST position
    var last = H.subarray((nTok - 1) * nEmbd, nTok * nEmbd);
    var on = vec("output_norm.weight");
    var hidden = new Float32Array(nEmbd);
    rmsNorm(last, on, eps, hidden);
    var wOut = mat("output.weight");     // tied to token_embd in some fixtures; that is expected
    var logits = W.exactLinear(hidden, wOut, nEmbd);

    meta.vocab = nVocab || meta.vocab;
    meta.variant = variant;
    meta.blockSize = blockSize;
    meta.quantiser = quantiser ? quantiser.name : "exact";
    meta.quantisedWeights = quantisedWeights;
    meta.meanBitsPerWeight = quantisedWeights ? quantisedBits / quantisedWeights : null;

    return {
      logits: logits, hidden: hidden, meta: meta, kv: meta,
      accounting: { quantisedMatrices: quantisedMatrices, totalMatrices: totalMatrices,
                    floatBytes: floatBytes, quantBytes: quantBytes,
                    ratio: quantBytes > 0 ? floatBytes / quantBytes : 0 }
    };
  }

  var api = {
    run: run, requirements: requirements, resolveMeta: resolveMeta,
    softmax: softmax, topK: topK, surface: surface,
    requiredTensors: requiredTensors, TYPE_NAMES: G ? G.GGML_TYPES : null
  };

  if (typeof module !== "undefined" && module.exports) module.exports = api;
  root.HDCFORWARD = api;
})(typeof globalThis !== "undefined" ? globalThis : this);