mlx-bun 0.0.8 → 0.0.10

package/src/spec/drafter.ts

@@ -4,14 +4,23 @@
 //
 // Q-only 4-layer transformer: K/V come from the TARGET's last sliding-
 // and full-attention donor caches at call time; pre/post projections
-// bridge target hidden space (2560) <-> drafter space (256); output via
-// the centroid-clustered head (2048 centroids, top-32 × 128 tokens).
+// bridge target hidden space (backbone) <-> drafter space.
 //
-// Implementation deviation (argmax-equivalent): the reference scatters
-// candidate scores into full-vocab logits initialised to -1e30 and
-// argmaxes; we argmax over the 4096 candidate scores and map through
-// candidate ids directly — same winner, no 262k materialisation. The
-// drafter's numeric details only influence acceptance RATE, never
+// Two output-head variants ship across the Gemma-4 drafter family — we
+// pick by TENSOR PRESENCE, not config (the larger artifacts declare
+// num_centroids in config.json but don't ship the centroid tensors —
+// the field is vestigial):
+//   - CENTROID head (small E2B/E4B drafters, `gemma4_assistant`):
+//     2048 centroids, top-32 × 128 tokens via a precomputed ordering.
+//   - TIED-EMBEDDING head (larger 12B/26B `gemma4_unified_assistant`):
+//     no centroids, no separate lm_head — logits = embed_tokens · h over
+//     the full vocab (the dense form the centroid head approximates).
+//
+// Implementation deviation (centroid path, argmax-equivalent): the
+// reference scatters candidate scores into full-vocab logits initialised
+// to -1e30 and argmaxes; we argmax over the 4096 candidate scores and map
+// through candidate ids directly — same winner, no 262k materialisation.
+// The drafter's numeric details only influence acceptance RATE, never
 // output correctness (the target's verify decides every token).
 
 import type { ModelConfig } from "../config";
@@ -56,10 +65,11 @@ export class GemmaAssistantDrafter {
   readonly vocabPerCentroid: number;
   #w: Weights;
   #blocks: Block[] = [];
+  readonly useCentroids: boolean;
   #prePT: MlxArray; // pre_projection transposed [in, out]
   #postPT: MlxArray;
-  #centroidsT: MlxArray; // [hidden, numCentroids]
-  #ordering: MlxArray;   // [numCentroids, vocabPerCentroid] int32
+  #centroidsT: MlxArray | null = null; // [hidden, numCentroids]
+  #ordering: MlxArray | null = null;   // [numCentroids, vocabPerCentroid] int32
   #embed: MlxArray;      // drafter token embeddings [vocab, hidden]
   #norm: MlxArray;
 
@@ -84,14 +94,20 @@ export class GemmaAssistantDrafter {
     };
     this.#prePT = transposed("pre_projection.weight");
     this.#postPT = transposed("post_projection.weight");
-    this.#centroidsT = transposed("masked_embedding.centroids.weight");
     this.#embed = T("model.embed_tokens.weight");
     this.#norm = T("model.norm.weight");
-    const ord64 = T("masked_embedding.token_ordering");
-    const ord32 = ord64.astype(Dtype.int32);
-    this.#ordering = ops.reshape(ord32, [this.numCentroids, this.vocabPerCentroid]);
-    this.#ordering.eval();
-    ord32.dispose();
+    // Head variant: only the small drafters ship the centroid tensors.
+    // The 12B/26B artifacts ship a tied embed_tokens head instead (their
+    // config still declares num_centroids — ignore it, trust the tensors).
+    this.useCentroids = w.has("masked_embedding.centroids.weight");
+    if (this.useCentroids) {
+      this.#centroidsT = transposed("masked_embedding.centroids.weight");
+      const ord64 = T("masked_embedding.token_ordering");
+      const ord32 = ord64.astype(Dtype.int32);
+      this.#ordering = ops.reshape(ord32, [this.numCentroids, this.vocabPerCentroid]);
+      this.#ordering.eval();
+      ord32.dispose();
+    }
 
     const ropeP = t.rope_parameters ?? {};
     for (let i = 0; i < t.num_hidden_layers; i++) {
@@ -215,41 +231,57 @@ export class GemmaAssistantDrafter {
 
     h = disposing(h, this.#rms(h, this.#norm));
 
-    // centroid decode → drafted token (argmax-equivalent shortcut)
-    const cScores = ops.matmul(h, this.#centroidsT);          // (1,1,2048)
-    const flat = ops.reshape(cScores, [1, this.numCentroids]);
-    cScores.dispose();
-    const neg = ops.neg(flat);
-    flat.dispose();
-    const part = ops.argpartitionAxis(neg, this.topK - 1, -1);
-    neg.dispose();
-    const topIdx = part.slice([0, 0], [1, this.topK]);        // (1, topK)
-    part.dispose();
-    const topFlat = ops.reshape(topIdx, [this.topK]);
-    topIdx.dispose();
-    const candIds = ops.takeAxis(this.#ordering, topFlat, 0); // (topK, vpc) int32
-    topFlat.dispose();
-    const candFlat = ops.reshape(candIds, [this.topK * this.vocabPerCentroid]);
-    candIds.dispose();
-    const candEmb = ops.takeAxis(this.#embed, candFlat, 0);   // (topK*vpc, hidden)
-    const hVec = ops.reshape(h, [this.hidden, 1]);
-    const scores = ops.matmul(candEmb, hVec);                  // (topK*vpc, 1)
-    candEmb.dispose();
-    hVec.dispose();
-    const scoresFlat = ops.reshape(scores, [1, this.topK * this.vocabPerCentroid]);
-    scores.dispose();
-    const am = ops.argmaxAxis(scoresFlat, -1);
-    scoresFlat.dispose();
-    const winner = ops.itemUint32(am);
-    am.dispose();
-    const tokArr = candFlat.slice([winner], [winner + 1]);
-    candFlat.dispose();
-    const tokU32 = tokArr.astype(Dtype.uint32);
-    tokArr.dispose();
-    const token = ops.itemUint32(tokU32);
-    tokU32.dispose();
+    let token: number;
+    if (this.useCentroids) {
+      // centroid decode → drafted token (argmax-equivalent shortcut)
+      const cScores = ops.matmul(h, this.#centroidsT!);         // (1,1,2048)
+      const flat = ops.reshape(cScores, [1, this.numCentroids]);
+      cScores.dispose();
+      const neg = ops.neg(flat);
+      flat.dispose();
+      const part = ops.argpartitionAxis(neg, this.topK - 1, -1);
+      neg.dispose();
+      const topIdx = part.slice([0, 0], [1, this.topK]);        // (1, topK)
+      part.dispose();
+      const topFlat = ops.reshape(topIdx, [this.topK]);
+      topIdx.dispose();
+      const candIds = ops.takeAxis(this.#ordering!, topFlat, 0); // (topK, vpc) int32
+      topFlat.dispose();
+      const candFlat = ops.reshape(candIds, [this.topK * this.vocabPerCentroid]);
+      candIds.dispose();
+      const candEmb = ops.takeAxis(this.#embed, candFlat, 0);   // (topK*vpc, hidden)
+      const hVec = ops.reshape(h, [this.hidden, 1]);
+      const scores = ops.matmul(candEmb, hVec);                  // (topK*vpc, 1)
+      candEmb.dispose();
+      hVec.dispose();
+      const scoresFlat = ops.reshape(scores, [1, this.topK * this.vocabPerCentroid]);
+      scores.dispose();
+      const am = ops.argmaxAxis(scoresFlat, -1);
+      scoresFlat.dispose();
+      const winner = ops.itemUint32(am);
+      am.dispose();
+      const tokArr = candFlat.slice([winner], [winner + 1]);
+      candFlat.dispose();
+      const tokU32 = tokArr.astype(Dtype.uint32);
+      tokArr.dispose();
+      token = ops.itemUint32(tokU32);
+      tokU32.dispose();
+    } else {
+      // tied-embedding head: token = argmax_v (embed[v] · h) over the
+      // full vocab — the dense form the centroid head approximates.
+      const V = this.#embed.shape[0]!;
+      const hVec = ops.reshape(h, [this.hidden, 1]);
+      const scores = ops.matmul(this.#embed, hVec);             // (vocab, 1)
+      hVec.dispose();
+      const scoresFlat = ops.reshape(scores, [1, V]);
+      scores.dispose();
+      const am = ops.argmaxAxis(scoresFlat, -1);
+      scoresFlat.dispose();
+      token = ops.itemUint32(am);
+      am.dispose();
+    }
 
-    const nextHidden = ops.matmul(h, this.#postPT);           // (1,1,2560)
+    const nextHidden = ops.matmul(h, this.#postPT);           // (1,1,backbone)
     h.dispose();
     return { token, nextHidden };
   }

package/src/spec/dspark/data-dflash.ts

@@ -0,0 +1,156 @@
+// Data path for faithful DFlash (multi-layer H_ctx). Parallel to data.ts.
+//
+// Shard stores, per packed sequence, the target's tapped hiddens at m layers
+// concatenated on the feature axis: `hidden_ml` [Ltot, m*H] bf16 (the m layers
+// in tapLayers order, the LAST being the post-finalNorm hidden used for p^t),
+// plus `ids` [Ltot], `seq_lens` [nSeq], `resp_starts` [nSeq].
+//
+// A training anchor at response position t attends the FULL prefix context
+// [0..t] (paper Eq 2/3), so sampleBatch slices a per-anchor variable-length
+// prefix and LEFT-PADS the batch to a fixed maxCtx with a key-padding mask.
+
+import { MlxArray } from "../../mlx/array";
+import { Dtype } from "../../mlx/ffi";
+import * as ops from "../../mlx/ops";
+import { loadAdapterTensors } from "../../lora";
+import { writeShardedSafetensors, type NamedTensor } from "../../quantize/safetensors-writer";
+import { mkdirSync, writeFileSync, readFileSync, existsSync, readdirSync } from "node:fs";
+import { join } from "node:path";
+
+/** One re-forwarded sequence: ids, the response-region start, and the m-layer
+ *  hiddens over the FULL sequence (row-major [L, m*H] bf16 bytes). */
+export interface DflashRecord {
+  ids: number[];
+  respStart: number;
+  hiddenMlBf16: Uint8Array;
+}
+
+export interface DflashShardMeta {
+  nSeq: number; nTokens: number; hiddenSize: number; m: number; tapLayers: number[];
+}
+
+export function writeDflashShard(outDir: string, shardIdx: number, records: DflashRecord[], hiddenSize: number, m: number, tapLayers: number[]): DflashShardMeta {
+  const dir = join(outDir, `shard_${String(shardIdx).padStart(5, "0")}`);
+  mkdirSync(dir, { recursive: true });
+  const mH = m * hiddenSize;
+  const seqLens: number[] = [], respStarts: number[] = [], allIds: number[] = [], parts: Uint8Array[] = [];
+  for (const r of records) {
+    seqLens.push(r.ids.length); respStarts.push(r.respStart);
+    for (const id of r.ids) allIds.push(id);
+    parts.push(r.hiddenMlBf16);
+  }
+  const nTokens = allIds.length;
+  let total = 0; for (const p of parts) total += p.length;
+  const bytes = new Uint8Array(total); let off = 0; for (const p of parts) { bytes.set(p, off); off += p.length; }
+  const hidden = MlxArray.fromBytesCopy(bytes, [nTokens, mH], Dtype.bfloat16);
+  const ids = MlxArray.fromInt32(new Int32Array(allIds), [nTokens]);
+  const lens = MlxArray.fromInt32(new Int32Array(seqLens), [seqLens.length]);
+  const rs = MlxArray.fromInt32(new Int32Array(respStarts), [respStarts.length]);
+  const tensors: NamedTensor[] = [
+    { name: "hidden_ml", array: hidden }, { name: "ids", array: ids },
+    { name: "seq_lens", array: lens }, { name: "resp_starts", array: rs },
+  ];
+  writeShardedSafetensors(dir, tensors);
+  for (const t of tensors) t.array.dispose();
+  const meta: DflashShardMeta = { nSeq: records.length, nTokens, hiddenSize, m, tapLayers };
+  writeFileSync(join(dir, "shard.json"), JSON.stringify(meta, null, 2));
+  return meta;
+}
+
+export class DflashShard {
+  readonly hiddenMl: MlxArray; // [Ltot, m*H] bf16
+  readonly ids: Int32Array;
+  readonly seqStart: number[]; readonly seqLen: number[]; readonly respStart: number[];
+  readonly hiddenSize: number; readonly m: number;
+  constructor(hiddenMl: MlxArray, ids: Int32Array, seqLens: Int32Array, respStarts: Int32Array, hiddenSize: number, m: number) {
+    this.hiddenMl = hiddenMl; this.ids = ids; this.hiddenSize = hiddenSize; this.m = m;
+    this.seqStart = []; this.seqLen = []; this.respStart = [];
+    let off = 0;
+    for (let i = 0; i < seqLens.length; i++) { this.seqStart.push(off); this.seqLen.push(seqLens[i]!); this.respStart.push(respStarts[i]!); off += seqLens[i]!; }
+  }
+  static load(dir: string): DflashShard {
+    const meta = JSON.parse(readFileSync(join(dir, "shard.json"), "utf8")) as DflashShardMeta;
+    const t = loadAdapterTensors(join(dir, "model.safetensors"));
+    try {
+      const hidden = t.get("hidden_ml"), idsA = t.get("ids"), lensA = t.get("seq_lens"), rsA = t.get("resp_starts");
+      if (!hidden || !idsA || !lensA || !rsA) throw new Error(`shard ${dir} missing tensors`);
+      const ids = Int32Array.from(idsA.toFloat32(), (v) => Math.round(v));
+      const lens = Int32Array.from(lensA.toFloat32(), (v) => Math.round(v));
+      const rs = Int32Array.from(rsA.toFloat32(), (v) => Math.round(v));
+      idsA.dispose(); lensA.dispose(); rsA.dispose();
+      return new DflashShard(hidden, ids, lens, rs, meta.hiddenSize, meta.m);
+    } catch (e) { for (const [, a] of t) a.dispose(); throw e; }
+  }
+  dispose(): void { this.hiddenMl.dispose(); }
+}
+
+export function listDflashShards(root: string): string[] {
+  if (!existsSync(root)) return [];
+  return readdirSync(root).filter((n) => n.startsWith("shard_")).sort().map((n) => join(root, n));
+}
+
+export interface DflashBatch {
+  hCtx: MlxArray;         // [A, maxCtx, m*H] bf16, LEFT-padded
+  ctxMask: MlxArray;      // [A, maxCtx] f32 (1=real, 0=pad)
+  targetHidden: MlxArray; // [A, γ, H] bf16 (final-layer hiddens at block positions → p^t)
+  anchorToks: number[];
+  blockToks: number[][];
+  size: number;
+}
+
+/**
+ * Sample A anchors and gather them into a faithful DFlash batch. Each anchor at
+ * response position t contributes its full prefix context [0..t] (capped at
+ * maxCtx, keeping the most recent), left-padded. The final tapped layer (last
+ * in tapLayers = post-finalNorm) supplies the block's p^t hiddens.
+ */
+export function sampleDflashBatch(shard: DflashShard, A: number, gamma: number, maxCtx: number, rng: () => number): DflashBatch | null {
+  const H = shard.hiddenSize, m = shard.m, mH = m * H;
+  const ctxIdx: number[] = [];   // A*maxCtx global indices (dummy 0 for pads)
+  const maskRows: number[] = []; // A*maxCtx
+  const blockIdx: number[] = []; // A*γ global indices
+  const anchorToks: number[] = [];
+  const blockToks: number[][] = [];
+
+  let got = 0, guard = 0;
+  while (got < A && guard < A * 64) {
+    guard++;
+    const s = Math.floor(rng() * shard.seqLen.length);
+    const len = shard.seqLen[s]!, rStart = shard.respStart[s]!, start = shard.seqStart[s]!;
+    // anchor local t in [rStart, len-γ-1]; need γ tokens after and t in response.
+    const lo = rStart, hi = len - gamma - 1;
+    if (hi < lo) continue;
+    const t = lo + Math.floor(rng() * (hi - lo + 1));
+    const g = start + t;
+    // prefix context [0..t], capped to the most-recent maxCtx positions.
+    const La = Math.min(t + 1, maxCtx);
+    const ctxStart = g - La + 1;
+    const pad = maxCtx - La;
+    for (let j = 0; j < pad; j++) { ctxIdx.push(0); maskRows.push(0); }
+    for (let j = 0; j < La; j++) { ctxIdx.push(ctxStart + j); maskRows.push(1); }
+    const bt: number[] = [];
+    for (let k = 1; k <= gamma; k++) bt.push(shard.ids[g + k]!); // CE/markov targets x_{t+1..t+γ}
+    // p^t hiddens: block position k predicts x_{t+1+k}, whose target dist is
+    // softmax(LM_head(h_{t+k})) — the hidden that PREDICTS it. So gather
+    // positions [t..t+γ-1] = [g..g+γ-1], NOT [g+1..g+γ] (off-by-one fix).
+    for (let k = 0; k < gamma; k++) blockIdx.push(g + k);
+    anchorToks.push(shard.ids[g]!); blockToks.push(bt); got++;
+  }
+  if (got === 0) return null;
+
+  const cIdx = MlxArray.fromInt32(new Int32Array(ctxIdx), [got * maxCtx]);
+  const gathered = ops.takeAxis(shard.hiddenMl, cIdx, 0); // [A*maxCtx, m*H]
+  cIdx.dispose();
+  const hCtx = ops.reshape(gathered, [got, maxCtx, mH]); gathered.dispose();
+  const ctxMask = MlxArray.fromFloat32(new Float32Array(maskRows), [got, maxCtx]);
+
+  const bIdx = MlxArray.fromInt32(new Int32Array(blockIdx), [got * gamma]);
+  const bGathered = ops.takeAxis(shard.hiddenMl, bIdx, 0); // [A*γ, m*H]
+  bIdx.dispose();
+  // final tapped layer = last H columns → the post-finalNorm hidden for p^t
+  const finalSlice = bGathered.slice([0, (m - 1) * H], [got * gamma, m * H]); // [A*γ, H]
+  bGathered.dispose();
+  const targetHidden = ops.reshape(finalSlice, [got, gamma, H]); finalSlice.dispose();
+
+  return { hCtx, ctxMask, targetHidden, anchorToks, blockToks, size: got };
+}

package/src/spec/dspark/data.ts

@@ -0,0 +1,201 @@
+// DSpark training data — the regen target (paper §3.3 + §5.1 "hidden state
+// communication"). Training the drafter requires, per block position, the
+// frozen target's TRUE distribution p^t and the tapped hidden state. Rather
+// than store full-vocab logits (V≈262k → enormous), we store only the target's
+// per-position final hidden states (the tap, H≈2816) plus the token ids, and
+// recompute p^t on the fly during training via the SHARED frozen LM head —
+// O(d) storage instead of O(V). (paper §5.1.)
+//
+// On-disk layout: a directory of shards, each `shard_NNNNN/model.safetensors`
+// holding three tensors over a pack of sequences:
+//   - "hidden"   [Ltot, H] bf16  — concatenated per-position final hiddens
+//   - "ids"      [Ltot]    int32 — concatenated token ids
+//   - "seq_lens" [nSeq]    int32 — per-sequence lengths (to split Ltot)
+// plus a sidecar `shard.json` with counts. The regen script (Josh runs it on
+// GPU) writes these; the trainer reads them and samples anchor blocks.
+//
+// An anchor at sequence position t yields a training example:
+//   hCtx        = hidden[t]            (the injected context, paper H_ctx)
+//   anchorTok   = ids[t]              (the bonus token x0)
+//   blockToks   = ids[t+1 .. t+γ]     (ground truth x*_1..x*_γ)
+//   targetHidden= hidden[t+1 .. t+γ]  (→ p^t_k via the frozen head at train)
+
+import { MlxArray } from "../../mlx/array";
+import { Dtype } from "../../mlx/ffi";
+import * as ops from "../../mlx/ops";
+import { loadAdapterTensors } from "../../lora";
+import { writeShardedSafetensors, type NamedTensor } from "../../quantize/safetensors-writer";
+import { mkdirSync, writeFileSync, readFileSync, existsSync, readdirSync } from "node:fs";
+import { join } from "node:path";
+
+/** One re-forwarded sequence: ids + per-position final hidden (bf16 bytes). */
+export interface DSparkRecord {
+  ids: number[];
+  /** Row-major [L, H] bf16 bytes (from MlxArray.rawBytes of the tap). */
+  hiddenBf16: Uint8Array;
+}
+
+export interface DSparkShardMeta {
+  nSeq: number;
+  nTokens: number;
+  hiddenSize: number;
+}
+
+/** Write a shard of re-forwarded sequences to `dir/shard_NNNNN/`. */
+export function writeDSparkShard(
+  outDir: string,
+  shardIdx: number,
+  records: DSparkRecord[],
+  hiddenSize: number,
+): DSparkShardMeta {
+  const dir = join(outDir, `shard_${String(shardIdx).padStart(5, "0")}`);
+  mkdirSync(dir, { recursive: true });
+
+  const seqLens: number[] = [];
+  const allIds: number[] = [];
+  const hiddenParts: Uint8Array[] = [];
+  for (const rec of records) {
+    seqLens.push(rec.ids.length);
+    for (const id of rec.ids) allIds.push(id);
+    hiddenParts.push(rec.hiddenBf16);
+  }
+  const nTokens = allIds.length;
+
+  // concat hidden bytes
+  let total = 0;
+  for (const p of hiddenParts) total += p.length;
+  const hiddenBytes = new Uint8Array(total);
+  let off = 0;
+  for (const p of hiddenParts) { hiddenBytes.set(p, off); off += p.length; }
+
+  const hidden = MlxArray.fromBytesCopy(hiddenBytes, [nTokens, hiddenSize], Dtype.bfloat16);
+  const ids = MlxArray.fromInt32(new Int32Array(allIds), [nTokens]);
+  const lens = MlxArray.fromInt32(new Int32Array(seqLens), [seqLens.length]);
+  const tensors: NamedTensor[] = [
+    { name: "hidden", array: hidden },
+    { name: "ids", array: ids },
+    { name: "seq_lens", array: lens },
+  ];
+  writeShardedSafetensors(dir, tensors);
+  hidden.dispose(); ids.dispose(); lens.dispose();
+
+  const meta: DSparkShardMeta = { nSeq: records.length, nTokens, hiddenSize };
+  writeFileSync(join(dir, "shard.json"), JSON.stringify(meta, null, 2));
+  return meta;
+}
+
+/** A loaded shard, ready for anchor sampling. Owns `hidden`; call dispose(). */
+export class DSparkShard {
+  readonly hidden: MlxArray;   // [Ltot, H] bf16
+  readonly ids: Int32Array;    // [Ltot]
+  readonly seqStart: number[]; // per-seq start offset into Ltot
+  readonly seqLen: number[];   // per-seq length
+  readonly hiddenSize: number;
+
+  constructor(hidden: MlxArray, ids: Int32Array, seqLens: Int32Array, hiddenSize: number) {
+    this.hidden = hidden;
+    this.ids = ids;
+    this.hiddenSize = hiddenSize;
+    this.seqStart = [];
+    this.seqLen = [];
+    let off = 0;
+    for (const l of seqLens) {
+      this.seqStart.push(off);
+      this.seqLen.push(l);
+      off += l;
+    }
+  }
+
+  static load(dir: string): DSparkShard {
+    const metaPath = join(dir, "shard.json");
+    if (!existsSync(metaPath)) throw new Error(`no shard.json in ${dir}`);
+    const meta = JSON.parse(readFileSync(metaPath, "utf8")) as DSparkShardMeta;
+    const t = loadAdapterTensors(join(dir, "model.safetensors"));
+    try {
+      const hidden = t.get("hidden");
+      const idsArr = t.get("ids");
+      const lensArr = t.get("seq_lens");
+      if (!hidden || !idsArr || !lensArr) throw new Error(`shard ${dir} missing tensors`);
+      // int32 ids round-trip exactly through f32 (vocab < 2^24).
+      const ids = Int32Array.from(idsArr.toFloat32(), (v) => Math.round(v));
+      const lens = Int32Array.from(lensArr.toFloat32(), (v) => Math.round(v));
+      idsArr.dispose(); lensArr.dispose();
+      return new DSparkShard(hidden, ids, lens, meta.hiddenSize);
+    } catch (e) {
+      for (const [, arr] of t) arr.dispose();
+      throw e;
+    }
+  }
+
+  dispose(): void {
+    this.hidden.dispose();
+  }
+}
+
+/** List shard directories under a data root, sorted. */
+export function listShards(root: string): string[] {
+  if (!existsSync(root)) return [];
+  return readdirSync(root)
+    .filter((n) => n.startsWith("shard_"))
+    .sort()
+    .map((n) => join(root, n));
+}
+
+export interface DSparkBatch {
+  /** [A, H] bf16 — tapped target hidden at each anchor. */
+  hCtx: MlxArray;
+  /** [A, γ, H] bf16 — target final hiddens at the γ block positions. */
+  targetHidden: MlxArray;
+  /** anchor token x0 per row. */
+  anchorToks: number[];
+  /** ground-truth block tokens x*_1..x*_γ per row [A][γ]. */
+  blockToks: number[][];
+  /** number of anchors A. */
+  size: number;
+}
+
+/**
+ * Sample A anchors from a loaded shard and gather them into a training batch.
+ * A valid anchor at global index g (within a sequence) needs γ continuation
+ * tokens in the SAME sequence. `rng` returns [0,1).
+ */
+export function sampleBatch(shard: DSparkShard, A: number, gamma: number, rng: () => number): DSparkBatch | null {
+  const anchorGlobal: number[] = [];
+  const anchorToks: number[] = [];
+  const blockToks: number[][] = [];
+  const blockGlobal: number[] = []; // flattened γ indices per anchor
+
+  let guard = 0;
+  while (anchorGlobal.length < A && guard < A * 64) {
+    guard++;
+    const s = Math.floor(rng() * shard.seqLen.length);
+    const len = shard.seqLen[s]!;
+    if (len < gamma + 1) continue;
+    const start = shard.seqStart[s]!;
+    // local anchor position t in [0, len-γ-1]
+    const t = Math.floor(rng() * (len - gamma));
+    const g = start + t;
+    anchorGlobal.push(g);
+    anchorToks.push(shard.ids[g]!);
+    const bt: number[] = [];
+    for (let k = 1; k <= gamma; k++) {
+      bt.push(shard.ids[g + k]!);
+      blockGlobal.push(g + k);
+    }
+    blockToks.push(bt);
+  }
+  if (anchorGlobal.length === 0) return null;
+  const got = anchorGlobal.length;
+
+  const anchorIdx = MlxArray.fromInt32(new Int32Array(anchorGlobal), [got]);
+  const hCtx = ops.takeAxis(shard.hidden, anchorIdx, 0); // [A,H]
+  anchorIdx.dispose();
+
+  const blockIdx = MlxArray.fromInt32(new Int32Array(blockGlobal), [got * gamma]);
+  const tgtFlat = ops.takeAxis(shard.hidden, blockIdx, 0); // [A*γ, H]
+  blockIdx.dispose();
+  const targetHidden = ops.reshape(tgtFlat, [got, gamma, shard.hiddenSize]);
+  tgtFlat.dispose();
+
+  return { hCtx, targetHidden, anchorToks, blockToks, size: got };
+}