@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/models/speculative/deepspec.ts

@@ -0,0 +1,744 @@
+// DeepSpec's Gemma4DSparkModel — semi-autoregressive draft model, ported
+// verbatim (copy-verbatim methodology: transcribe THEIR forward, then port
+// op-for-op) from github.com/deepseek-ai/DeepSpec (MIT), files:
+//   deepspec/modeling/dspark/gemma4/modeling.py   (Gemma4DSparkModel)
+//   deepspec/modeling/dspark/common.py            (extract_context_feature)
+//   deepspec/modeling/dspark/markov_head.py       (VanillaMarkov)
+//   deepspec/eval/dspark/draft_ops.py             (inference-time draft loop)
+//   deepspec/utils/sampling.py                    (sample_tokens, temp=0 argmax)
+// Cross-checked against the installed HF reference (transformers
+// models/gemma4/modeling_gemma4.py) for Gemma4RMSNorm, apply_rotary_pos_emb,
+// Gemma4TextRotaryEmbedding (rope_type "proportional" via
+// modeling_rope_utils._compute_proportional_rope_parameters), and
+// Gemma4TextScaledWordEmbedding — since Gemma4DSparkModel subclasses/reuses
+// those directly rather than redefining them.
+//
+// This is a DRAFT model for target-Gemma4 speculative decoding: not
+// causal, not autoregressive over the whole sequence — it predicts a
+// γ-token block in one backbone pass (context-conditioned, bidirectional
+// within the block), then a cheap sequential Markov correction walks the
+// block left-to-right applying a low-rank transition bias per step.
+//
+// This module is STATELESS server-lifetime weights. Per-request context
+// K/V (the incremental cache DeepSpec threads through consecutive draft
+// rounds via a DynamicCache) is NOT stored here — see `projectContextKVRows` /
+// `draftRows` below for the seam.
+
+import { Weights } from "../../artifacts/weights";
+import type {
+  ContextKV, DeepspecContextAttention, DeepspecDraftRows, DeepspecDrafterModel, DrafterContextKV, DrafterContextProjection,
+} from "../../contracts/mlx/drafter";
+import type { MlxRowTable } from "../../contracts/mlx/graph";
+import { isDeepspecArchitecture } from "./deepspec-artifact";
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { parseQuantization, quantFor, type QuantizationConfig } from "../../artifacts/config";
+
+function disposing(old: MlxArray, next: MlxArray): MlxArray {
+  old.dispose();
+  return next;
+}
+
+// ---------------------------------------------------------------------------
+// Weight representation: every matmul weight and gather table loads as either
+// bf16 (the released checkpoint) or affine-quantized (a
+// `02d723a:scripts/dspark.ts` quantize sibling). Detection is the house
+// pattern — a `.scales` sibling key marks a quantized module; bits/group_size
+// come from the config's `quantization` block (per-module overrides + the
+// mlx `false` = kept-full-precision convention). Both formats flow through
+// the same forward: the quantized matmul path is mlx QuantizedLinear
+// semantics (x @ W^T via quantized_matmul transpose=true) and the quantized
+// gather path is mlx QuantizedEmbedding semantics (gather packed/scales/bias
+// rows, dequantize the gathered rows) — matching mlx-lm's numerics.
+// ---------------------------------------------------------------------------
+
+/** A Linear weight `W [out, in]`, used as `x @ W^T`. bf16 keeps the lazy
+ *  transpose VIEW (see the memory note below); quant keeps packed storage. */
+type MatW =
+  | { kind: "bf16"; t: MlxArray }
+  | { kind: "quant"; w: MlxArray; scales: MlxArray; biases: MlxArray | null; spec: ops.QuantSpec };
+
+/** An embedding table `[vocab, dim]`, used via row gather. */
+type TableW =
+  | { kind: "bf16"; w: MlxArray }
+  | { kind: "quant"; w: MlxArray; scales: MlxArray; biases: MlxArray | null; spec: ops.QuantSpec };
+
+/** x @ W^T for either format. Output dtype follows x (bf16 on the real
+ *  checkpoint) in BOTH branches — the model-dtype fidelity notes on the
+ *  markov/confidence paths hold unchanged for the quantized path. */
+function matmulW(x: MlxArray, m: MatW): MlxArray {
+  if (m.kind === "bf16") return ops.matmul(x, m.t);
+  return ops.quantizedMatmul(x, m.w, m.scales, m.biases, m.spec, true);
+}
+
+/** Row gather for either format (quant: QuantizedEmbedding.encode —
+ *  dequantize returns the scales' dtype, bf16, same as the bf16 branch). */
+function gatherW(t: TableW, idsU32: MlxArray): MlxArray {
+  if (t.kind === "bf16") return ops.takeAxis(t.w, idsU32, 0);
+  const rows = ops.takeAxis(t.w, idsU32, 0);
+  const scaleRows = ops.takeAxis(t.scales, idsU32, 0);
+  const biasRows = t.biases ? ops.takeAxis(t.biases, idsU32, 0) : null;
+  const out = ops.dequantize(rows, scaleRows, biasRows, t.spec);
+  rows.dispose();
+  scaleRows.dispose();
+  biasRows?.dispose();
+  return out;
+}
+
+function disposeW(m: MatW | TableW | null): void {
+  if (!m) return;
+  if (m.kind === "bf16") {
+    ("t" in m ? m.t : m.w).dispose();
+    return;
+  }
+  m.w.dispose();
+  m.scales.dispose();
+  m.biases?.dispose();
+}
+
+// ---------------------------------------------------------------------------
+// Config (their config.json field names, verbatim keys — no renaming so the
+// mapping to the checkpoint stays obvious).
+// ---------------------------------------------------------------------------
+
+export interface DeepspecConfig {
+  architectures: string[];
+  hidden_size: number;
+  num_hidden_layers: number;
+  num_attention_heads: number;
+  global_head_dim: number;
+  num_global_key_value_heads: number;
+  attention_k_eq_v: boolean;
+  intermediate_size: number;
+  hidden_activation: string;
+  rms_norm_eps: number;
+  final_logit_softcapping: number | null;
+  vocab_size: number;
+  block_size: number;
+  mask_token_id: number;
+  target_layer_ids: number[];
+  num_target_layers: number;
+  markov_rank: number;
+  markov_head_type: string;
+  enable_confidence_head: boolean;
+  confidence_head_with_markov: boolean;
+  rope_theta: number;
+  partial_rotary_factor: number;
+  /** Not in the checkpoint's config.json (draft_ops.py takes it as a call
+   *  arg to build_dspark_proposal); default 0 = confidence head disabled
+   *  functionally even when weights exist (threshold<=0 ⇒ full γ block). */
+  confidence_threshold: number;
+}
+
+function readConfig(raw: Record<string, any>): DeepspecConfig {
+  if (!isDeepspecArchitecture(raw.architectures))
+    throw new Error(`DeepspecDrafter: expected architectures[0]==="Gemma4DSparkModel", got ${JSON.stringify(raw.architectures)}`);
+
+  // RoPE: their config nests rope params under rope_parameters.full_attention
+  // (matches Gemma4TextConfig / Gemma4TextRotaryEmbedding's per-layer-type
+  // dict); ground truth also allows a flat rope_theta/partial_rotary_factor
+  // for simplified checkpoints — accept either shape.
+  const ropeFull = raw.rope_parameters?.full_attention ?? raw.rope_parameters ?? {};
+  const ropeTheta = ropeFull.rope_theta ?? raw.rope_theta ?? 1e6;
+  const partialRotary = ropeFull.partial_rotary_factor ?? raw.partial_rotary_factor ?? 1.0;
+
+  const targetLayerIds: number[] = raw.target_layer_ids;
+  if (!Array.isArray(targetLayerIds) || targetLayerIds.length === 0)
+    throw new Error("DeepspecDrafter: config.target_layer_ids must be a non-empty array");
+
+  // Silent-wrongness guards (2026-07-06 port review): each of these variants
+  // would LOAD cleanly and compute wrong math — this module implements only
+  // the released-checkpoint shape. Refuse loudly instead.
+  if (!raw.attention_k_eq_v)
+    throw new Error(
+      "DeepspecDrafter: attention_k_eq_v=false checkpoints need a real v_proj path — not implemented (this module wires V ≡ K)",
+    );
+  if ((raw.markov_rank ?? 0) > 0 && (raw.markov_head_type ?? "vanilla") !== "vanilla")
+    throw new Error(
+      `DeepspecDrafter: markov_head_type "${raw.markov_head_type}" not implemented (gated/rnn heads ship markov_w1/w2 too and would silently get vanilla math) — only "vanilla"`,
+    );
+  if ((raw.hidden_activation ?? "gelu_pytorch_tanh") !== "gelu_pytorch_tanh")
+    throw new Error(
+      `DeepspecDrafter: hidden_activation "${raw.hidden_activation}" not implemented — the MLP hardcodes gelu_pytorch_tanh`,
+    );
+
+  return {
+    architectures: raw.architectures,
+    hidden_size: raw.hidden_size,
+    num_hidden_layers: raw.num_hidden_layers,
+    num_attention_heads: raw.num_attention_heads,
+    global_head_dim: raw.global_head_dim,
+    num_global_key_value_heads: raw.num_global_key_value_heads,
+    attention_k_eq_v: !!raw.attention_k_eq_v,
+    intermediate_size: raw.intermediate_size,
+    hidden_activation: raw.hidden_activation ?? "gelu_pytorch_tanh",
+    rms_norm_eps: raw.rms_norm_eps ?? 1e-6,
+    final_logit_softcapping: raw.final_logit_softcapping ?? null,
+    vocab_size: raw.vocab_size,
+    block_size: raw.block_size,
+    mask_token_id: raw.mask_token_id,
+    target_layer_ids: targetLayerIds,
+    num_target_layers: raw.num_target_layers,
+    markov_rank: raw.markov_rank ?? 0,
+    markov_head_type: raw.markov_head_type ?? "vanilla",
+    enable_confidence_head: !!raw.enable_confidence_head,
+    confidence_head_with_markov: !!raw.confidence_head_with_markov,
+    rope_theta: ropeTheta,
+    partial_rotary_factor: partialRotary,
+    confidence_threshold: raw.confidence_threshold ?? 0,
+  };
+}
+
+// ---------------------------------------------------------------------------
+// Per-layer weights.
+// ---------------------------------------------------------------------------
+
+interface Layer {
+  qProj: MatW; // [nHeads*headDim, hidden] Linear storage, used as x@W^T
+  kProj: MatW; // [nKvHeads*headDim, hidden]
+  oProj: MatW; // [hidden, nHeads*headDim]
+  qNorm: MlxArray;  // [headDim]
+  kNorm: MlxArray;  // [headDim]
+  // v_norm has with_scale=False (Gemma4RMSNorm(..., with_scale=False)) — no
+  // weight tensor in the checkpoint; represented as `null` into ops.rmsNorm.
+  inputNorm: MlxArray;
+  postAttnNorm: MlxArray;
+  preFfNorm: MlxArray;
+  postFfNorm: MlxArray;
+  gateProj: MatW;
+  upProj: MatW;
+  downProj: MatW;
+  layerScalar: MlxArray; // [1]
+}
+
+export class DeepspecDrafter implements DeepspecDrafterModel {
+  readonly cfg: DeepspecConfig;
+  /** === config.target_layer_ids, alias for the ground-truth term. */
+  readonly tapLayers: number[];
+  /** === config.block_size (γ), alias for the ground-truth term. */
+  readonly gamma: number;
+  readonly hidden: number;
+  readonly nHeads: number;
+  readonly nKvHeads: number;
+  readonly headDim: number;
+  readonly eps: number;
+  readonly embedScale: number;
+  readonly ropeFreqs: MlxArray; // proportional-RoPE freqs, [headDim/2]
+
+  #w: Weights;
+  #embed: TableW;        // [vocab, hidden]
+  #fc: MatW;             // fc.weight [hidden, taps*hidden] Linear storage
+  #hiddenNorm: MlxArray; // [hidden]
+  #norm: MlxArray;       // [hidden]
+  #lmHead: MatW;         // lm_head.weight [vocab, hidden]
+  #layers: Layer[] = [];
+  #markovW1: TableW | null; // [vocab, markovRank] embedding
+  #markovW2: MatW | null;   // markov_w2.weight [vocab, rank] Linear storage
+  /** One block step's token per row, bound once: the argmax of the step's
+   *  logits [B, V], plus the Markov head's bias from the previous tokens [B]
+   *  when the checkpoint carries one. Borrows both; owns the ids. */
+  readonly #selectStep: (logits: MlxArray, previous: MlxArray) => MlxArray;
+  #confProj: MatW | null;   // proj.weight [1, confDim] Linear storage (bf16
+                            // keeps the once-built transpose view — per-call
+                            // transpose+cast leaked, 2026-07-06 review fix)
+  #confB: MlxArray | null;  // [1]
+
+  private constructor(w: Weights, cfg: DeepspecConfig, quant: QuantizationConfig | null) {
+    this.#w = w;
+    this.cfg = cfg;
+    this.tapLayers = cfg.target_layer_ids;
+    this.gamma = cfg.block_size;
+    this.hidden = cfg.hidden_size;
+    this.nHeads = cfg.num_attention_heads;
+    this.nKvHeads = cfg.num_global_key_value_heads; // k_eq_v enforced in readConfig (no v_proj path)
+    this.headDim = cfg.global_head_dim;
+    this.eps = cfg.rms_norm_eps;
+    this.embedScale = Math.sqrt(cfg.hidden_size);
+
+    const T = (name: string) => w.tensor(name);
+    // Lazy transpose VIEW — no contiguous copy. Materializing would store
+    // every matmul weight a
+    // second time: on the real 6.86 GB checkpoint that's ~+4.8 GB resident
+    // (lm_head alone ~2 GB), decisive on a 24 GB box. mlx's steel GEMM
+    // dispatches a transposed-B operand natively, so the view costs nothing
+    // per matmul. (2026-07-06 port review, MEDIUM-memory finding.)
+    const transposed = (name: string): MlxArray => ops.transposeAxes(T(name), [1, 0]);
+    // Quantized-module detection + spec (house pattern: `.scales` sibling
+    // key gates; bits/group_size from the config quantization block).
+    const qspec = (base: string): ops.QuantSpec => {
+      const spec = quantFor(quant, base);
+      if (!spec)
+        throw new Error(
+          `DeepspecDrafter: ${base} has quantized tensors (.scales sibling) but no quantization config entry`,
+        );
+      return { bits: spec.bits, groupSize: spec.groupSize, mode: spec.mode };
+    };
+    const quantParts = (base: string) => ({
+      w: T(`${base}.weight`),
+      scales: T(`${base}.scales`),
+      biases: w.has(`${base}.biases`) ? T(`${base}.biases`) : null,
+      spec: qspec(base),
+    });
+    const mat = (base: string): MatW =>
+      w.has(`${base}.scales`)
+        ? { kind: "quant", ...quantParts(base) }
+        : { kind: "bf16", t: transposed(`${base}.weight`) };
+    const table = (base: string): TableW =>
+      w.has(`${base}.scales`)
+        ? { kind: "quant", ...quantParts(base) }
+        : { kind: "bf16", w: T(`${base}.weight`) };
+
+    this.#embed = table("embed_tokens");
+    this.#fc = mat("fc");
+    this.#hiddenNorm = T("hidden_norm.weight");
+    this.#norm = T("norm.weight");
+    this.#lmHead = mat("lm_head");
+
+    // Proportional RoPE freqs (mlx_lm.models.rope_utils.ProportionalRoPE,
+    // transformers._compute_proportional_rope_parameters): rotated_dims =
+    // int(head_dim * partial_rotary_factor); exponents over [0, rotated_dims)
+    // step 2, divided by FULL head_dim (not rotated_dims); tail padded with
+    // +inf ⇒ angle pos/∞ = 0 ⇒ cos=1, sin=0 — the identity rotation, i.e.
+    // passthrough on those dims (freqs is the wavelength array fast::rope
+    // divides position by, same convention as packages/inference/src/models/gemma4/assistant.ts's
+    // ProportionalRoPE port; verified vs mlx_lm rope_utils.ProportionalRoPE).
+    const rotatedDims = Math.floor(this.headDim * cfg.partial_rotary_factor);
+    const n = this.headDim / 2;
+    const freqs = new Float32Array(n).fill(Infinity);
+    for (let k = 0; k < rotatedDims / 2; k++) {
+      freqs[k] = Math.pow(cfg.rope_theta, (2 * k) / this.headDim);
+    }
+    this.ropeFreqs = MlxArray.fromFloat32(freqs, [n]);
+
+    for (let i = 0; i < cfg.num_hidden_layers; i++) {
+      const p = `layers.${i}`;
+      this.#layers.push({
+        qProj: mat(`${p}.self_attn.q_proj`),
+        kProj: mat(`${p}.self_attn.k_proj`),
+        oProj: mat(`${p}.self_attn.o_proj`),
+        qNorm: T(`${p}.self_attn.q_norm.weight`),
+        kNorm: T(`${p}.self_attn.k_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: mat(`${p}.mlp.gate_proj`),
+        upProj: mat(`${p}.mlp.up_proj`),
+        downProj: mat(`${p}.mlp.down_proj`),
+        layerScalar: T(`${p}.layer_scalar`),
+      });
+    }
+
+    if (cfg.markov_rank > 0) {
+      this.#markovW1 = table("markov_head.markov_w1"); // [vocab, rank] embedding
+      this.#markovW2 = mat("markov_head.markov_w2");   // stored [vocab, rank] Linear
+      this.#selectStep = (logits, previous) => {
+        using bias = this.#markovStepBias(previous);
+        using corrected = ops.add(logits, bias);
+        return ops.argmaxAxis(corrected, -1);
+      };
+    } else {
+      this.#markovW1 = null;
+      this.#markovW2 = null;
+      this.#selectStep = logits => ops.argmaxAxis(logits, -1);
+    }
+
+    if (cfg.enable_confidence_head) {
+      this.#confProj = mat("confidence_head.proj"); // [1, confDim] Linear
+      this.#confB = T("confidence_head.proj.bias"); // [1]
+    } else {
+      this.#confProj = null;
+      this.#confB = null;
+    }
+  }
+
+  static async load(dir: string): Promise<DeepspecDrafter> {
+    const raw = (await Bun.file(`${dir}/config.json`).json()) as Record<string, any>;
+    const cfg = readConfig(raw);
+    // Quantized siblings (`02d723a:scripts/dspark.ts` quantize) carry the house
+    // quantization block; the released bf16 checkpoint has none → null.
+    const quant = parseQuantization(raw.quantization ?? raw.quantization_config);
+    const w = await Weights.open(dir);
+    return new DeepspecDrafter(w, cfg, quant);
+  }
+
+  /** f32-internal RMSNorm (Gemma4RMSNorm casts through float32 internally:
+   *  `hidden_states.float()` → norm → `* weight.float()` → `.type_as`).
+   *  weight=null reproduces `with_scale=False` (v_norm). */
+  #rms(x: MlxArray, weight: MlxArray | null): MlxArray {
+    const f = x.dtype === Dtype.float32 ? x : x.astype(Dtype.float32);
+    const n = ops.rmsNorm(f, weight, this.eps);
+    if (f !== x) f.dispose();
+    if (n.dtype === x.dtype) return n;
+    const back = n.astype(x.dtype);
+    n.dispose();
+    return back;
+  }
+
+  #geluTanhMlp(x: MlxArray, layer: Layer): MlxArray {
+    const g = matmulW(x, layer.gateProj);
+    const u = matmulW(x, layer.upProj);
+    const act = ops.geluApprox(g); // gelu_pytorch_tanh == config.hidden_activation
+    g.dispose();
+    let mlp = ops.mul(act, u);
+    act.dispose();
+    u.dispose();
+    mlp = disposing(mlp, matmulW(mlp, layer.downProj));
+    return mlp;
+  }
+
+  /**
+   * extract_context_feature (common.py):
+   *   torch.cat([hidden_states[0 if lid==-1 else lid+1] for lid in layer_ids], -1)
+   * i.e. tap i (target_layer_ids[i], never -1 for this checkpoint) reads the
+   * TARGET's post-layer-i residual output (list-of-all-hidden-states
+   * convention: index 0 = embedding, index L+1 = output of layer L). The
+   * caller is expected to hand us exactly that: targetHiddens is the
+   * concatenation over taps, [1, L, taps*hidden], already ordered to match
+   * config.target_layer_ids (ascending, enforced upstream by
+   * validate_target_layer_ids).
+   *
+   * context = hidden_norm(fc(target_hiddens))   — modeling.py _forward_backbone:
+   *   target_hidden_states = self.hidden_norm(self.fc(target_hidden_states))
+   *
+   * Returns the projected+normed context hidden [1, L, hidden] (still needs
+   * per-layer k_proj + k_norm + rope to become cacheable ContextKV — see
+   * projectContextKVRows).
+   */
+  #projectContext(targetHiddens: MlxArray): MlxArray {
+    const proj = matmulW(targetHiddens, this.#fc); // [1,L,hidden]
+    return disposing(proj, this.#rms(proj, this.#hiddenNorm));
+  }
+
+  /** The context projection for every positive B·L row count (`matmulW`:
+   *  MLX's quantized matmul, matmul for a BF16 checkpoint). */
+  readonly projectContext: MlxRowTable<DrafterContextProjection> =
+    Object.freeze({
+      default: (targetHiddens: MlxArray) => this.#projectContext(targetHiddens),
+    });
+
+  /** Each layer's context K/V: the same operation for every positive B·L row count. */
+  readonly projectContextKVRows: MlxRowTable<DrafterContextKV> =
+    Object.freeze({
+      default: (contextHidden: MlxArray, position: number | MlxArray) => this.#projectContextKVRows(contextHidden, position),
+    });
+
+  /**
+   * Per-layer context K/V rows, projected + normed + ROPE'd, ready to cache
+   * across draft rounds.
+   *
+   * Justification for caching at this stage (from modeling.py's
+   * Gemma4DSparkAttention.forward, transcribed above the port below):
+   *   k_ctx = self.k_proj(target_hidden_states)          # per-layer proj
+   *   k = torch.cat([k_ctx, k_noise], dim=1).view(...)    # concat BEFORE norm
+   *   k = self.k_norm(k).transpose(1, 2)                  # norm over the concat
+   *   k = apply_gemma4_rotary_pos_emb(k, cos, sin, ...)   # rope over the concat
+   *   if past_key_values is not None:
+   *       k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs)  # CACHE AFTER ROPE
+   *
+   * Gemma4RMSNorm normalizes over the LAST axis (head_dim) independently per
+   * row; concatenation happens along the sequence axis. So k_norm(cat(ctx,
+   * noise))[i] == k_norm(ctx)[i] for every ctx row i — norming the context
+   * rows alone and norming them as part of a larger concat batch produce
+   * IDENTICAL values (RMSNorm has no cross-row term). RoPE likewise applies
+   * per-position (absolute position index, stable per context row across
+   * rounds) independent of what else shares the batch. Caching post-norm,
+   * post-rope context K rows is therefore bit-exact equivalent to
+   * recomputing them fresh inside every round's concat — which is exactly
+   * what DeepSpec's own DynamicCache.update-after-rope does (it caches the
+   * already roped/normed tensor, then crop() drops only the noise suffix).
+   *
+   * v ≡ k when attention_k_eq_v (v_proj is absent; v_norm is the SAME
+   * concat-then-scale-less-norm, over k_proj's output — v gets NO rope).
+   * `position` is the first context row's absolute position (per row as an
+   * array when the rows start at different positions).
+   */
+  #projectContextKVRows(contextHidden: MlxArray, position: number | MlxArray): ContextKV[] {
+    const B = contextHidden.shape[0]!, L = contextHidden.shape[1]!;
+    const out: ContextKV[] = [];
+    for (const layer of this.#layers) {
+      const kFlat = matmulW(contextHidden, layer.kProj); // [1,L,nKvHeads*headDim]
+      const k4 = ops.reshape(kFlat, [B, L, this.nKvHeads, this.headDim]);
+      kFlat.dispose();
+      // v ≡ k pre-norm (attention_k_eq_v: v_ctx = k_ctx, same tensor before
+      // the two norms diverge — v_norm has no weight).
+      const vNormed = this.#rms(k4, null); // scale-less
+      const kNormed = this.#rms(k4, layer.kNorm);
+      k4.dispose();
+
+      const kT = ops.transposeAxes(kNormed, [0, 2, 1, 3]); // [1,nKvHeads,L,headDim]
+      kNormed.dispose();
+      const vT = ops.transposeAxes(vNormed, [0, 2, 1, 3]);
+      vNormed.dispose();
+
+      const kRoped = typeof position === "number"
+        ? ops.rope(kT, this.headDim, null, position, this.ropeFreqs)
+        : ops.ropeDynamic(kT, this.headDim, null, position, this.ropeFreqs);
+      kT.dispose();
+      out.push({ k: kRoped, v: vT });
+    }
+    return out;
+  }
+
+  /**
+   * One backbone forward over the noise block, given cached context K/V —
+   * modeling.py Gemma4DSparkModel._forward_backbone + Gemma4DSparkAttention,
+   * transcribed inline per-layer below. Non-causal SDPA (scale=1.0) over
+   * [ctx rows ++ block rows]; q uses only the block's own positions.
+   *
+   * blockIds: token ids for each noise block, shape [B, gamma] — [anchor, MASK,
+   * MASK, ..., MASK] (create_noise_embed: only position 0 gets the real
+   * anchor token when block_keep_mask; MASK elsewhere — at inference every
+   * anchor is "kept" so this is always [anchor, mask_token_id × (γ-1)]).
+   * position: each row's initial absolute position. RoPE advances contiguous
+   * positions within the block.
+   */
+  #forwardBackbone(context: DeepspecContextAttention, blockIds: MlxArray, position: number | MlxArray): MlxArray {
+    const B = blockIds.shape[0]!, G = blockIds.shape[1]!;
+    const idsU32 = blockIds.astype(Dtype.uint32);
+    const embRaw = gatherW(this.#embed, idsU32); // [1,G,hidden] (gather over vocab axis)
+    idsU32.dispose();
+    const embFlat = ops.reshape(embRaw, [B, G, this.hidden]);
+    embRaw.dispose();
+    // Gemma4TextScaledWordEmbedding.forward: embed(ids) * embed_scale
+    let h = ops.mulScalar(embFlat, this.embedScale);
+    embFlat.dispose();
+
+    for (let li = 0; li < this.#layers.length; li++) {
+      const layer = this.#layers[li]!;
+
+      const residual = h;
+      const x = this.#rms(h, layer.inputNorm);
+
+      // q = q_norm(q_proj(x).view(B,q_len,nHeads,headDim)).transpose(1,2)
+      const qFlat = matmulW(x, layer.qProj); // [1,G,nHeads*headDim]
+      const q4 = ops.reshape(qFlat, [B, G, this.nHeads, this.headDim]);
+      qFlat.dispose();
+      const qNormed = this.#rms(q4, layer.qNorm);
+      q4.dispose();
+      const qT = ops.transposeAxes(qNormed, [0, 2, 1, 3]); // [1,nHeads,G,headDim]
+      qNormed.dispose();
+      // rope on q uses the LAST q_len positions (cos[:, -q_len:, :]) — here
+      // q_len == G == the whole block, starting at each row's position.
+      const qRoped = typeof position === "number"
+        ? ops.rope(qT, this.headDim, null, position, this.ropeFreqs)
+        : ops.ropeDynamic(qT, this.headDim, null, position, this.ropeFreqs);
+      qT.dispose();
+
+      // k_noise = k_proj(x); norm+rope IDENTICALLY to the cached context
+      // rows (concat-then-norm == norm-then-concat per the RMSNorm argument
+      // in projectContextKVRows's doc comment).
+      const kNoiseFlat = matmulW(x, layer.kProj);
+      const kNoise4 = ops.reshape(kNoiseFlat, [B, G, this.nKvHeads, this.headDim]);
+      kNoiseFlat.dispose();
+      const vNoiseNormed = this.#rms(kNoise4, null);
+      const kNoiseNormed = this.#rms(kNoise4, layer.kNorm);
+      kNoise4.dispose();
+      const kNoiseT = ops.transposeAxes(kNoiseNormed, [0, 2, 1, 3]);
+      kNoiseNormed.dispose();
+      const vNoiseT = ops.transposeAxes(vNoiseNormed, [0, 2, 1, 3]);
+      vNoiseNormed.dispose();
+      const kNoiseRoped = typeof position === "number"
+        ? ops.rope(kNoiseT, this.headDim, null, position, this.ropeFreqs)
+        : ops.ropeDynamic(kNoiseT, this.headDim, null, position, this.ropeFreqs);
+      kNoiseT.dispose();
+      x.dispose();
+
+      const attnOut = context.attend(li, qRoped, kNoiseRoped, vNoiseT);
+      qRoped.dispose(); kNoiseRoped.dispose(); vNoiseT.dispose();
+
+      const attnT = ops.transposeAxes(attnOut, [0, 2, 1, 3]); // [1,G,nHeads,headDim]
+      attnOut.dispose();
+      const attnFlat = ops.reshape(attnT, [B, G, this.nHeads * this.headDim]);
+      attnT.dispose();
+      let attn = matmulW(attnFlat, layer.oProj);
+      attnFlat.dispose();
+      attn = disposing(attn, this.#rms(attn, layer.postAttnNorm));
+      h = ops.add(residual, attn);
+      residual.dispose();
+      attn.dispose();
+
+      const res2 = h;
+      const f = this.#rms(h, layer.preFfNorm);
+      let mlp = this.#geluTanhMlp(f, layer);
+      f.dispose();
+      mlp = disposing(mlp, this.#rms(mlp, layer.postFfNorm));
+      h = ops.add(res2, mlp);
+      res2.dispose();
+      mlp.dispose();
+
+      // Gemma4DSparkDecoderLayer.forward: `return hidden_states * self.layer_scalar`
+      // — the scalar multiplies the WHOLE layer output (residual-sum
+      // INCLUDED, not just the sublayer delta).
+      h = disposing(h, ops.mul(h, layer.layerScalar));
+    }
+
+    return disposing(h, this.#rms(h, this.#norm));
+  }
+
+  /** lm_head(h) then softcap: `tanh(logits/cap)*cap` (compute_logits). */
+  #computeLogits(h: MlxArray): MlxArray {
+    const logits = matmulW(h, this.#lmHead);
+    if (this.cfg.final_logit_softcapping === null) return logits;
+    const cap = this.cfg.final_logit_softcapping;
+    const capArr = ops.scalarLike(cap, logits);
+    const scaled = ops.div(logits, capArr);
+    logits.dispose();
+    const t = ops.tanh(scaled);
+    scaled.dispose();
+    const out = ops.mul(t, capArr);
+    t.dispose();
+    capArr.dispose();
+    return out;
+  }
+
+  /** B_k = markov_w2(markov_w1(prevTok)) — VanillaMarkov.compute_step_bias /
+   *  project_bias(get_prev_embeddings(token_ids)). Returns [1,vocab] in MODEL
+   *  dtype (bf16 on the real checkpoint) — the reference samples over the
+   *  bf16 `logits + bias` sum, and softcapped logits in [-30,30] have bf16
+   *  steps ~0.125-0.25, so an f32 sum here reorders near-ties vs torch and
+   *  derails the sequential block (2026-07-06 port review, fidelity fix). */
+  #markovStepBias(previous: MlxArray): MlxArray {
+    if (!this.#markovW1 || !this.#markovW2)
+      throw new Error("markovStepBias: markov_rank == 0 (no markov head)");
+    const idxU32 = previous.astype(Dtype.uint32);
+    const e1 = gatherW(this.#markovW1, idxU32); // [1,rank]
+    idxU32.dispose();
+    const bias = matmulW(e1, this.#markovW2); // [1,vocab], model dtype
+    e1.dispose();
+    return bias;
+  }
+
+  /** c_k = sigmoid(proj([h_k; markov_w1[prevTok]])). h1 is [1,hidden]
+   *  (already the right dtype/shape row for this position). AcceptRatePredictor
+   *  is nn.Linear(confDim,1) — proj.weight is [1,confDim]; we compute
+   *  features @ proj.weight^T + bias without transposing storage (row-vector
+   *  dot via matmul against a [confDim,1] view). */
+  #confidence(h1: MlxArray, previous: MlxArray): MlxArray {
+    if (!this.#confProj || !this.#confB)
+      throw new Error("confidence: confidence head disabled");
+    let feat = h1;
+    let ownFeat = false;
+    if (this.cfg.confidence_head_with_markov) {
+      if (!this.#markovW1) throw new Error("confidence_head_with_markov requires markov_rank > 0");
+      const idxU32 = previous.astype(Dtype.uint32);
+      const e1 = gatherW(this.#markovW1, idxU32); // [1,rank]
+      idxU32.dispose();
+      const e1f = e1.dtype === h1.dtype ? e1 : disposing(e1, e1.astype(h1.dtype));
+      feat = ops.concatAxis([h1, e1f], 1); // [1, hidden+rank]
+      e1f.dispose();
+      ownFeat = true;
+    }
+    // Reference precision (AcceptRatePredictor / predict_confidence_step):
+    // the nn.Linear runs in MODEL dtype (bf16 on the real checkpoint) and is
+    // .float()ed AFTER — computing it in f32 here flipped borderline
+    // threshold comparisons vs torch (2026-07-06 port review). So: matmul +
+    // bias in feat's dtype (matmulW: the once-built transpose view for bf16,
+    // qmm transpose=true for a quantized head), then cast.
+    let z = matmulW(feat, this.#confProj);
+    if (ownFeat) feat.dispose();
+    z = disposing(z, ops.add(z, this.#confB));
+    if (z.dtype !== Dtype.float32) z = disposing(z, z.astype(Dtype.float32));
+    z = disposing(z, ops.sigmoid(z));
+    return z;
+  }
+
+  /**
+   * The whole per-round draft. Two-pass structure, matching the source
+   * exactly (draft_ops.py build_dspark_proposal calls sample_draft_tokens
+   * for the FULL block first, THEN predicts confidence over the full
+   * sampled sequence, THEN truncates — confidence does not short-circuit
+   * sampling):
+   *
+   *   base_draft_logits = model.compute_logits(proposal_hidden_states)      # [1,γ,V], softcapped
+   *   sampled_tokens, draft_logits = model.sample_draft_tokens(             # full γ, sequential Markov
+   *       base_draft_logits, first_prev_token_ids=anchor, temperature=0)
+   *   confidence_logits = predict_confidence_step(                         # full γ, over the SAMPLED prevs
+   *       proposal_hidden_states,
+   *       prev_token_ids=cat([anchor, sampled_tokens[:, :-1]]))
+   *   proposal_draft_tokens = _confident_prefix_length(confidence_logits, threshold)
+   *   verify_input_ids = cat([anchor, sampled_tokens[:, :proposal_draft_tokens]])
+   *
+   * Transcribed sequential Markov sampling (modeling.py sample_draft_tokens
+   * → markov_head.py VanillaMarkov.sample_block_tokens, temperature=0 ⇒
+   * sample_tokens's argmax branch — RNG-free):
+   *   prev = first_prev_token_ids   # == anchorTok
+   *   for k in range(block_size):
+   *       step_logits = base_logits[:,k] + markov_w2(markov_w1(prev))
+   *       tok = argmax(step_logits)
+   *       prev = tok
+   *
+   * Transcribed confidence truncation (draft_ops.py _confident_prefix_length):
+   *   if threshold <= 0: return block_size          # disabled — full γ
+   *   below = sigmoid(confidence_logits) < threshold
+   *   if not below.any(): return block_size
+   *   return first index where below is True         # ℓ can be 0 (EMPTY proposal)
+   *
+   * Here the graph and the sequential proposal head run at the active row
+   * count, one block per anchor (the just-verified or bonus token starting
+   * it); confidence selects each row's length after the full device chain.
+   * `context` reads the rows' cached context K/V (projectContextKVRows; the
+   * caller owns its lifetime and the crop after the round); `position` is
+   * each block's first absolute position.
+   */
+  draftRows(context: DeepspecContextAttention, anchors: readonly number[], position: number | MlxArray,
+    collectLogits = false): DeepspecDraftRows {
+    const B = anchors.length, G = this.gamma, V = this.cfg.vocab_size;
+    using ids = ops.fromInt32(anchors.flatMap(anchor => [anchor, ...Array(G - 1).fill(this.cfg.mask_token_id)]), [B, G]);
+    using hidden = this.#forwardBackbone(context, ids, position);
+    using base = this.#computeLogits(hidden);
+    using anchorIds = ops.fromInt32([...anchors], [B]);
+    const sampled: MlxArray[] = [], confidence: MlxArray[] = [];
+    let previous = anchorIds;
+    try {
+      for (let step = 0; step < G; step++) {
+        using row = base.slice([0, step, 0], [B, step + 1, V]);
+        using logits = ops.reshape(row, [B, V]);
+        const token = this.#selectStep(logits, previous);
+        sampled.push(token); previous = token;
+      }
+      // Serving does not consume confidence when truncation is disabled.
+      if (this.#confProj && (collectLogits || this.cfg.confidence_threshold > 0)) for (let step = 0; step < G; step++) {
+        using row = hidden.slice([0, step, 0], [B, step + 1, this.hidden]);
+        using flat = ops.reshape(row, [B, this.hidden]);
+        confidence.push(this.#confidence(flat, step ? sampled[step - 1]! : anchorIds));
+      }
+      using packed = ops.concatAxis(sampled, 0);
+      const tokens = packed.toIntTokens();
+      using packedConfidence = confidence.length ? ops.concatAxis(confidence, 0) : null;
+      const conf = packedConfidence?.toFloat32();
+      const rows = anchors.map((_, row) => {
+        const values = conf ? Array.from({ length: G }, (_, step) => conf[step * B + row]!) : [];
+        const below = this.cfg.confidence_threshold > 0 ? values.findIndex(value => value < this.cfg.confidence_threshold) : -1;
+        const length = below < 0 ? G : below;
+        return { tokens: Array.from({ length }, (_, step) => tokens[step * B + row]!), conf: values.slice(0, length) };
+      });
+      return { tokens: rows.map(row => row.tokens), conf: rows.map(row => row.conf),
+        ...(collectLogits ? { baseLogits: base.astype(Dtype.float32) } : {}) };
+    } finally { for (const array of [...sampled, ...confidence]) array.dispose(); }
+  }
+
+  dispose(): void {
+    disposeW(this.#embed);
+    disposeW(this.#fc);
+    this.#hiddenNorm.dispose();
+    this.#norm.dispose();
+    disposeW(this.#lmHead);
+    this.ropeFreqs.dispose();
+    for (const l of this.#layers) {
+      disposeW(l.qProj); disposeW(l.kProj); disposeW(l.oProj);
+      l.qNorm.dispose(); l.kNorm.dispose();
+      l.inputNorm.dispose(); l.postAttnNorm.dispose();
+      l.preFfNorm.dispose(); l.postFfNorm.dispose();
+      disposeW(l.gateProj); disposeW(l.upProj); disposeW(l.downProj);
+      l.layerScalar.dispose();
+    }
+    disposeW(this.#markovW1);
+    disposeW(this.#markovW2);
+    disposeW(this.#confProj);
+    this.#confB?.dispose();
+    this.#w.dispose();
+  }
+}
+