mlx-bun 0.3.0 → 0.4.0

package/src/spec/deepspec-source.ts

@@ -22,6 +22,8 @@
 // Context rows are cached post-norm-post-rope — bit-equivalent to their
 // in-round concat (per-row RMSNorm/RoPE; argued in projectContextKV's doc).
 
+import { artifactIdentity } from "../model/artifact-identity";
+import { deepspecGroups, bindDeepspecTarget } from "./deepspec-rows";
 import { readdirSync, statSync } from "node:fs";
 import { join } from "node:path";
 import type { MlxArray } from "../mlx/array";
@@ -39,20 +41,27 @@ function safetensorsBytes(dir: string): number {
 export class DeepspecProvider implements DraftProvider {
   readonly id: string;
   readonly weightsBytes: number;
+  readonly grouped: import("./source").GroupedDraftProvider;
   /** Trained block width (config block_size, e.g. 7) — the server pins
    *  numDraftTokens to this. */
   readonly gamma: number;
 
-  private constructor(private readonly drafter: DeepspecDrafter, id: string, weightsBytes: number) {
+  private constructor(private readonly drafter: DeepspecDrafter, id: string, weightsBytes: number, namespace: string) {
     this.id = id;
     this.weightsBytes = weightsBytes;
     this.gamma = drafter.gamma;
+    this.grouped = deepspecGroups(drafter, namespace);
   }
 
   static async load(modelDir: string): Promise<DeepspecProvider> {
     const drafter = await DeepspecDrafter.load(modelDir);
     const id = modelDir.split("/").filter(Boolean).at(-1)!;
-    return new DeepspecProvider(drafter, id, safetensorsBytes(modelDir));
+    try {
+      const identity = await artifactIdentity(await Bun.file(`${modelDir}/config.json`).text(),
+        readdirSync(modelDir).filter(file => file.endsWith(".safetensors"))
+          .map(name => ({ name, path: join(modelDir, name) })));
+      return new DeepspecProvider(drafter, id, safetensorsBytes(modelDir), `deepspec-context-v1:${identity}`);
+    } catch (error) { drafter.dispose(); throw error; }
   }
 
   open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
@@ -72,14 +81,7 @@ export class DeepspecSource implements DraftSource {
 
   constructor(private readonly drafter: Pick<DeepspecDrafter,
     "cfg" | "tapLayers" | "projectContext" | "projectContextKV" | "draftBlock">, target: TargetView) {
-    if (!target.gemmaTaps)
-      throw new Error("DeepSpec drafter requires a Gemma4 target");
-    const nLayers = target.gemmaTaps.layerCount;
-    if (nLayers !== drafter.cfg.num_target_layers)
-      throw new Error(
-        `DeepSpec drafter was trained for a ${drafter.cfg.num_target_layers}-layer target; ` +
-          `this model has ${nLayers} layers — wrong (target, drafter) pairing`,
-      );
+    bindDeepspecTarget(target, drafter);
     this.tapLayers = drafter.tapLayers;
   }
 

package/src/spec/dflash-source.ts

@@ -18,6 +18,8 @@
 // target verify decides). The richer rejection-sampling verify stays in the
 // standalone dflashGenerate for the measure script.
 
+import { artifactIdentity } from "../model/artifact-identity";
+import { projectedDraftGroups } from "./projected-draft-rows";
 import { readdirSync, statSync } from "node:fs";
 import { join } from "node:path";
 import type { MlxArray } from "../mlx/array";
@@ -37,20 +39,43 @@ function safetensorsBytes(dir: string): number {
 export class DflashProvider implements DraftProvider {
   readonly id: string;
   readonly weightsBytes: number;
+  readonly grouped: import("./source").GroupedDraftProvider;
   /** The trained block width — the server pins numDraftTokens to this so the
    *  serve loop never asks for more positions than the block was trained for. */
   readonly gamma: number;
 
-  private constructor(private readonly drafter: DflashDrafter, id: string, weightsBytes: number) {
+  private constructor(private readonly drafter: DflashDrafter, id: string, weightsBytes: number, namespace: string) {
     this.id = id;
     this.weightsBytes = weightsBytes;
     this.gamma = drafter.cfg.gamma;
+    this.grouped = projectedDraftGroups(target => {
+      if (!target.gemmaTaps) throw new Error("DSpark drafter requires a Gemma4 target");
+      const projection = target.gemmaTaps.projection;
+      const minConf = runtimeValue("MLX_BUN_DSPARK_MINCONF");
+      return {
+        namespace, schema: "dflash-context-v1",
+        layers: drafter.cfg.nLayers, tapLayers: drafter.cfg.tapLayers,
+        project: hidden => drafter.projectContextRows(hidden),
+        draft(context, pending, _positions, depth) {
+          return drafter.forwardRows(projection, null, pending, depth, {
+            collectLogits: false, collectConfidence: false,
+            thresholds: drafter.cfg.sts?.thresholds,
+            ...(minConf ? { minConf: Number(minConf) } : {}),
+          }, context).tokens;
+        },
+      };
+    });
   }
 
   static async load(modelDir: string): Promise<DflashProvider> {
     const drafter = loadDsparkDrafter(modelDir); // variant dispatch (dspark|legacy dflash)
     const id = modelDir.split("/").filter(Boolean).at(-1)!;
-    return new DflashProvider(drafter, id, safetensorsBytes(modelDir));
+    try {
+      const identity = await artifactIdentity(await Bun.file(`${modelDir}/dspark.json`).text(),
+        readdirSync(modelDir).filter(file => file.endsWith(".safetensors"))
+          .map(name => ({ name, path: join(modelDir, name) })));
+      return new DflashProvider(drafter, id, safetensorsBytes(modelDir), `dflash-context-v1:${identity}`);
+    } catch (error) { drafter.dispose(); throw error; }
   }
 
   open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
@@ -58,8 +83,7 @@ export class DflashProvider implements DraftProvider {
   }
 
   dispose(): void {
-    // Unlike the assistant drafter (process-pinned mmaps), DflashDrafter
-    // materializes owned MlxArrays — free them on teardown.
+    // Provider-owned weights are released when the mounted draft is closed.
     this.drafter.dispose();
   }
 }

package/src/spec/drafter.ts

@@ -28,8 +28,32 @@ import { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
+import { disposeResources } from "../engine/resources";
 import type { SharedKv } from "../model/gemma4";
 
+/** The graph borrows donor attention; storage owns validity and encoding. */
+export interface AssistantAttention {
+  attend(query: MlxArray, scale: number, window: number | null): MlxArray;
+}
+export interface AssistantDonors { sliding: AssistantAttention; full: AssistantAttention; }
+
+/** Legacy plain donor adapter. The common graph never reads cache geometry. */
+export function plainAssistantDonors(shared: { sliding: [MlxArray, MlxArray]; full: [MlxArray, MlxArray] },
+  position: number): AssistantDonors {
+  const bind = ([keys, values]: [MlxArray, MlxArray]): AssistantAttention => ({
+    attend(query, scale, window) {
+      const length = keys.shape[2]!;
+      const from = window === null ? 0 : Math.max(0, Math.max(0, position - window + 1) - (position - length + 1));
+      using floats = from ? MlxArray.fromFloat32(Float32Array.from({ length }, (_, i) => i < from ? -1e9 : 0), [1,1,1,length]) : null;
+      using mask = floats?.astype(keys.dtype);
+      return ops.sdpa(query, keys, values, scale, mask ? "array" : "", mask ?? null);
+    },
+  });
+  return { sliding: bind(shared.sliding), full: bind(shared.full) };
+}
+
+export interface DrafterRowsStep { tokens: MlxArray; nextHidden: MlxArray; }
+
 export interface DrafterStep {
   /** Drafted token id. */
   token: number;
@@ -64,6 +88,7 @@ export class GemmaAssistantDrafter {
   readonly topK: number;
   readonly vocabPerCentroid: number;
   #w: Weights;
+  #owned: MlxArray[] = [];
   #blocks: Block[] = [];
   readonly useCentroids: boolean;
   #prePT: MlxArray; // pre_projection transposed [in, out]
@@ -84,73 +109,80 @@ export class GemmaAssistantDrafter {
     this.vocabPerCentroid = Math.floor(t.vocab_size / this.numCentroids);
     this.#w = w;
 
-    const T = (name: string) => w.tensor(name);
-    const transposed = (name: string): MlxArray => {
-      const a = ops.transposeAxes(T(name), [1, 0]);
-      const c = ops.contiguous(a);
-      a.dispose();
-      c.eval();
-      return c;
-    };
-    this.#prePT = transposed("pre_projection.weight");
-    this.#postPT = transposed("post_projection.weight");
-    this.#embed = T("model.embed_tokens.weight");
-    this.#norm = T("model.norm.weight");
-    // 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 own = (array: MlxArray) => { this.#owned.push(array); return array; };
+    try {
+      const T = (name: string) => w.tensor(name);
+      // Keep the oracle's strided weight.T: a contiguous transpose changes
+      // M4 matmul rounding and allocates a second copy of every linear weight.
+      const transposed = (name: string): MlxArray => {
+        const c = own(ops.transposeAxes(T(name), [1, 0]));
+        c.eval();
+        return c;
+      };
+      this.#prePT = transposed("pre_projection.weight");
+      this.#postPT = transposed("post_projection.weight");
+      this.#embed = T("model.embed_tokens.weight");
+      this.#norm = T("model.norm.weight");
+      // 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 = own(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++) {
-      const layerType = t.layer_types[i];
-      const isFull = layerType === "full_attention";
-      const headDim = isFull ? t.global_head_dim : t.head_dim;
-      const rp = ropeP[layerType] ?? {};
-      let ropeBase: number | null = rp.rope_theta ?? (isFull ? 1e6 : 1e4);
-      let ropeFreqs: MlxArray | null = null;
-      const partial = rp.partial_rotary_factor ?? 1.0;
-      if (isFull && partial < 1.0) {
-        // ProportionalRoPE freqs, matching the target's donor-K rope
-        const rotated = Math.floor(headDim * partial);
-        const n = headDim / 2;
-        const freqs = new Float32Array(n).fill(Infinity);
-        for (let k = 0; k < rotated / 2; k++)
-          freqs[k] = Math.pow(ropeBase!, (2 * k) / headDim);
-        ropeFreqs = MlxArray.fromFloat32(freqs, [n]);
-        ropeBase = null;
+      const ropeP = t.rope_parameters ?? {};
+      for (let i = 0; i < t.num_hidden_layers; i++) {
+        const layerType = t.layer_types[i];
+        const isFull = layerType === "full_attention";
+        const headDim = isFull ? t.global_head_dim : t.head_dim;
+        const rp = ropeP[layerType] ?? {};
+        let ropeBase: number | null = rp.rope_theta ?? (isFull ? 1e6 : 1e4);
+        let ropeFreqs: MlxArray | null = null;
+        const partial = rp.partial_rotary_factor ?? 1.0;
+        if (isFull && partial < 1.0) {
+          // ProportionalRoPE freqs, matching the target's donor-K rope
+          const rotated = Math.floor(headDim * partial);
+          const n = headDim / 2;
+          const freqs = new Float32Array(n).fill(Infinity);
+          for (let k = 0; k < rotated / 2; k++)
+            freqs[k] = Math.pow(ropeBase!, (2 * k) / headDim);
+          ropeFreqs = own(MlxArray.fromFloat32(freqs, [n]));
+          ropeBase = null;
+        }
+        const p = `model.layers.${i}`;
+        this.#blocks.push({
+          layerType, headDim, ropeBase, ropeFreqs,
+          qProj: transposed(`${p}.self_attn.q_proj.weight`),
+          oProj: transposed(`${p}.self_attn.o_proj.weight`),
+          qNorm: T(`${p}.self_attn.q_norm.weight`),
+          inputNorm: T(`${p}.input_layernorm.weight`),
+          postAttnNorm: T(`${p}.post_attention_layernorm.weight`),
+          preFfNorm: T(`${p}.pre_feedforward_layernorm.weight`),
+          postFfNorm: T(`${p}.post_feedforward_layernorm.weight`),
+          gateProj: transposed(`${p}.mlp.gate_proj.weight`),
+          upProj: transposed(`${p}.mlp.up_proj.weight`),
+          downProj: transposed(`${p}.mlp.down_proj.weight`),
+          layerScalar: T(`${p}.layer_scalar`),
+        });
       }
-      const p = `model.layers.${i}`;
-      this.#blocks.push({
-        layerType, headDim, ropeBase, ropeFreqs,
-        qProj: transposed(`${p}.self_attn.q_proj.weight`),
-        oProj: transposed(`${p}.self_attn.o_proj.weight`),
-        qNorm: T(`${p}.self_attn.q_norm.weight`),
-        inputNorm: T(`${p}.input_layernorm.weight`),
-        postAttnNorm: T(`${p}.post_attention_layernorm.weight`),
-        preFfNorm: T(`${p}.pre_feedforward_layernorm.weight`),
-        postFfNorm: T(`${p}.post_feedforward_layernorm.weight`),
-        gateProj: transposed(`${p}.mlp.gate_proj.weight`),
-        upProj: transposed(`${p}.mlp.up_proj.weight`),
-        downProj: transposed(`${p}.mlp.down_proj.weight`),
-        layerScalar: T(`${p}.layer_scalar`),
-      });
-    }
+    } catch (error) { disposeResources(this.#owned.splice(0)); throw error; }
   }
 
   static async load(modelDir: string): Promise<GemmaAssistantDrafter> {
     const cfg = (await Bun.file(`${modelDir}/config.json`).json()) as Record<string, any>;
-    return new GemmaAssistantDrafter(await Weights.open(modelDir), cfg);
+    const weights = await Weights.open(modelDir);
+    try { return new GemmaAssistantDrafter(weights, cfg); }
+    catch (error) { weights.dispose(); throw error; }
   }
 
+  dispose(): void { disposeResources([...this.#owned.splice(0), this.#w]); }
+
   /** f32-internal RMSNorm (reference _RMSNorm casts through float32). */
   #rms(x: MlxArray, weight: MlxArray): MlxArray {
     const f = x.astype(Dtype.float32);
@@ -161,49 +193,77 @@ export class GemmaAssistantDrafter {
     return back;
   }
 
-  /** One draft step. All inputs (1,1,·); sharedKv holds the target's
-   *  donor K/V in chronological order. position = target's last position. */
-  forward(
-    lastTokenEmb: MlxArray, targetHidden: MlxArray,
-    sharedKv: { sliding: [MlxArray, MlxArray]; full: [MlxArray, MlxArray] },
-    position: number,
-  ): DrafterStep {
-    let h = ops.concatAxis([lastTokenEmb, targetHidden], 2); // (1,1,2*backbone)
-    h = disposing(h, ops.matmul(h, this.#prePT));            // (1,1,256)
+  /** The request adapter consumes the same graph as a batch of one.
+   * Donors are chronological; position is the target's last position. */
+  forward(lastTokenEmb: MlxArray, targetHidden: MlxArray,
+    sharedKv: { sliding: [MlxArray, MlxArray]; full: [MlxArray, MlxArray] }, position: number): DrafterStep {
+    const result = this.forwardRows(lastTokenEmb, targetHidden, plainAssistantDonors(sharedKv, position), position);
+    try { return { token: ops.itemUint32(result.tokens), nextHidden: result.nextHidden }; }
+    catch (error) { result.nextHidden.dispose(); throw error; }
+    finally { result.tokens.dispose(); }
+  }
 
-    for (const blk of this.#blocks) {
-      const [k, v] = blk.layerType === "sliding_attention" ? sharedKv.sliding : sharedKv.full;
-      const kLen = k.shape[2]!;
+  #sampleRows(h: MlxArray, B: number): MlxArray {
+    if (this.useCentroids) {
+      using scores = ops.matmul(h, this.#centroidsT!);
+      using flat = ops.reshape(scores, [B, this.numCentroids]);
+      using neg = ops.neg(flat);
+      using partition = ops.argpartitionAxis(neg, this.topK - 1, -1);
+      using top = partition.slice([0, 0], [B, this.topK]);
+      using topFlat = ops.reshape(top, [B * this.topK]);
+      using candidates = ops.takeAxis(this.#ordering!, topFlat, 0);
+      const count = this.topK * this.vocabPerCentroid;
+      using ids = ops.reshape(candidates, [B, count]);
+      using flatIds = ops.reshape(ids, [B * count]);
+      using embeddings = ops.takeAxis(this.#embed, flatIds, 0);
+      using matrix = ops.reshape(embeddings, [B, count, this.hidden]);
+      using vector = ops.reshape(h, [B, 1, this.hidden]);
+      // The oracle rounds each bf16 product before reduction. Matmul's
+      // accumulation changes candidate rankings on M1, even at identical h.
+      using products = ops.mul(matrix, vector);
+      using candidateScores = ops.sumAxis(products, -1, false);
+      using maximum = ops.maxAxis(candidateScores, -1, true);
+      using winners = ops.equal(candidateScores, maximum);
+      // Sparse full-vocabulary argmax chooses the lowest token ID on a tie.
+      // Apply that rule directly to the candidates without a vocab-size scatter.
+      using sentinel = ops.fromInt32([this.#embed.shape[0]!], []);
+      using eligible = ops.where(winners, ids, sentinel);
+      using selected = ops.minAxis(eligible, -1, false);
+      return selected.astype(Dtype.uint32);
+    }
+    using rows = ops.reshape(h, [B, this.hidden]);
+    using vector = B === 1 ? ops.reshape(h, [this.hidden, 1]) : ops.transposeAxes(rows, [1, 0]);
+    using scores = ops.matmul(this.#embed, vector);
+    using flat = B === 1 ? ops.reshape(scores, [1, this.#embed.shape[0]!]) : ops.transposeAxes(scores, [1, 0]);
+    return ops.argmaxAxis(flat, -1);
+  }
 
-      // sliding-window additive mask over cached K positions
-      let mask: MlxArray | null = null;
-      if (blk.layerType === "sliding_attention") {
-        const firstCached = position - kLen + 1;
-        const windowStart = Math.max(0, position - this.slidingWindow + 1);
-        const allowedFrom = Math.max(0, windowStart - firstCached);
-        if (allowedFrom > 0) {
-          const m = new Float32Array(kLen);
-          for (let i = 0; i < allowedFrom; i++) m[i] = -1e9;
-          const mf = MlxArray.fromFloat32(m, [1, 1, 1, kLen]);
-          mask = mf.astype(k.dtype);
-          mf.dispose();
-        }
-      }
+  /** Numerical graph only. Callers supply donor validity masks for different
+   * row positions; token selection remains on device across the draft chain. */
+  forwardRows(lastTokenEmb: MlxArray, targetHidden: MlxArray,
+    donors: AssistantDonors, position: number | MlxArray): DrafterRowsStep {
+    const B = lastTokenEmb.shape[0]!;
+    let h = ops.concatAxis([lastTokenEmb, targetHidden], 2); // (B,1,2*backbone)
+    h = disposing(h, ops.matmul(h, this.#prePT));            // (B,1,256)
 
+    for (const blk of this.#blocks) {
+      const donor = blk.layerType === "sliding_attention" ? donors.sliding : donors.full;
       const residual = h;
       let x = this.#rms(h, blk.inputNorm);
       // Q-only attention
       let q = ops.matmul(x, blk.qProj);
       x.dispose();
-      q = disposing(q, ops.reshape(q, [1, 1, this.nHeads, blk.headDim]));
+      q = disposing(q, ops.reshape(q, [B, 1, this.nHeads, blk.headDim]));
       q = disposing(q, this.#rms(q, blk.qNorm));
       q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
-      q = disposing(q, ops.rope(q, blk.headDim, blk.ropeBase, position, blk.ropeFreqs));
-      let attn = ops.sdpa(q, k, v, Math.pow(blk.headDim, -0.5), mask ? "array" : "", mask);
+      q = disposing(q, typeof position === "number"
+        ? ops.rope(q, blk.headDim, blk.ropeBase, position, blk.ropeFreqs)
+        : ops.ropeDynamic(q, blk.headDim, blk.ropeBase, position, blk.ropeFreqs));
+      let attn = donor.attend(q, Math.pow(blk.headDim, -0.5),
+        blk.layerType === "sliding_attention" ? this.slidingWindow : null);
       q.dispose();
-      mask?.dispose();
       attn = disposing(attn, ops.transposeAxes(attn, [0, 2, 1, 3]));
-      attn = disposing(attn, ops.reshape(attn, [1, 1, this.nHeads * blk.headDim]));
+      attn = disposing(attn, ops.reshape(attn, [B, 1, this.nHeads * blk.headDim]));
       attn = disposing(attn, ops.matmul(attn, blk.oProj));
       attn = disposing(attn, this.#rms(attn, blk.postAttnNorm));
       h = ops.add(residual, attn);
@@ -231,59 +291,10 @@ export class GemmaAssistantDrafter {
 
     h = disposing(h, this.#rms(h, this.#norm));
 
-    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,backbone)
+    const tokens = this.#sampleRows(h, B);
+    const nextHidden = ops.matmul(h, this.#postPT);           // (B,1,backbone)
     h.dispose();
-    return { token, nextHidden };
+    return { tokens, nextHidden };
   }
 }