mlx-bun 0.5.0 → 1.0.0

package/src/spec/dspark/deepspec-module.ts

@@ -1,776 +0,0 @@
-// 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 `projectContext` /
-// `draftBlock` below for the seam.
-
-import { Weights } from "../../weights";
-import { MlxArray } from "../../mlx/array";
-import { Dtype } from "../../mlx/ffi";
-import * as ops from "../../mlx/ops";
-import { parseQuantization, quantFor, type QuantizationConfig } from "../../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
-// 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 {
-  const arch = raw.architectures?.[0];
-  if (arch !== "Gemma4DSparkModel")
-    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]
-}
-
-/** Per-layer context K/V rows, already projected + normed + roped —
- *  cacheable across draft rounds (see projectContext doc). */
-export interface ContextKV {
-  k: MlxArray; // [1, nKvHeads, ctxLen, headDim]
-  v: MlxArray; // [1, nKvHeads, ctxLen, headDim] (v ≡ k when attention_k_eq_v)
-}
-
-/** Storage owns context validity; the graph supplies only Q and new block KV. */
-export interface DeepspecContextAttention {
-  attend(layer: number, query: MlxArray, blockKeys: MlxArray, blockValues: MlxArray): MlxArray;
-}
-export function plainDeepspecContext(context: readonly ContextKV[]): DeepspecContextAttention {
-  return { attend(layer, query, blockKeys, blockValues) {
-    using keys = ops.concatAxis([context[layer]!.k, blockKeys], 2);
-    using values = ops.concatAxis([context[layer]!.v, blockValues], 2);
-    return ops.sdpa(query, keys, values, 1, "", null);
-  } };
-}
-export interface DraftRowsResult {
-  tokens: number[][];
-  conf: number[][];
-  baseLogits?: MlxArray;
-}
-
-export interface DraftBlockResult {
-  /** Sequentially-sampled draft tokens, length 0..gamma (0 iff confidence
-   *  truncation fires at the very first position — EMPTY proposal). */
-  tokens: number[];
-  /** Per-position sigmoid confidence, aligned with `tokens` (same length;
-   *  empty when the confidence head is disabled). */
-  conf: number[];
-  /** Raw base_logits (post target-forward, PRE-Markov-bias, post-softcap)
-   *  for the full block [1, gamma, vocab] (float32 COPY; sampling runs in
-   *  model dtype separately) — caller disposes. NOTE: the reference
-   *  verifier's temp>0 draft_probs come from the Markov-CORRECTED logits
-   *  (draft_ops.py:140-143), NOT these — at temp 0 (our oracle regime) the
-   *  distinction is moot (verify degenerates to argmax token-match), but a
-   *  future temp>0 verify must add the per-position Markov bias first. */
-  baseLogits: MlxArray;
-}
-
-export class DeepspecDrafter {
-  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
-  #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 src/spec/drafter.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
-    } else {
-      this.#markovW1 = null;
-      this.#markovW2 = null;
-    }
-
-    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 (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
-   * projectContextKV).
-   */
-  projectContext(targetHiddens: MlxArray): MlxArray {
-    const proj = matmulW(targetHiddens, this.#fc); // [1,L,hidden]
-    return disposing(proj, this.#rms(proj, this.#hiddenNorm));
-  }
-
-  /**
-   * 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).
-   */
-  projectContextKV(contextHidden: MlxArray, positions: number[]): ContextKV[] {
-    if (positions.length !== contextHidden.shape[1])
-      throw new Error(`projectContextKV: ${positions.length} positions for ${contextHidden.shape[1]} context rows`);
-    assertContiguous(positions);
-    return this.projectContextKVRows(contextHidden, positions[0]!);
-  }
-
-  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 projectContextKV'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)
-   *
-   * ctxKV: this round's cached context K/V (see projectContextKV/crop notes
-   * on the class doc — caller owns cache lifetime and crop-after-round).
-   * anchorTok: the just-verified/bonus token starting this block.
-   * anchorPos: its absolute sequence position (blockPositions[0]).
-   */
-  draftBlock(ctxKV: ContextKV[], anchorTok: number, anchorPos: number): DraftBlockResult {
-    const result = this.draftRows(plainDeepspecContext(ctxKV), [anchorTok], anchorPos, true);
-    return { tokens: result.tokens[0]!, conf: result.conf[0]!, baseLogits: result.baseLogits! };
-  }
-
-  /** The graph and sequential proposal head run at the active row count.
-   * Confidence selects each row's length after the full device chain. */
-  draftRows(context: DeepspecContextAttention, anchors: readonly number[], position: number | MlxArray,
-    collectLogits = false): DraftRowsResult {
-    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]);
-        using bias = this.#markovW1 ? this.#markovStepBias(previous) : null;
-        using corrected = bias ? ops.add(logits, bias) : null;
-        const token = ops.argmaxAxis(corrected ?? logits, -1);
-        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();
-  }
-}
-
-function assertContiguous(positions: number[]): void {
-  for (let i = 1; i < positions.length; i++) {
-    if (positions[i] !== positions[i - 1]! + 1)
-      throw new Error(`DeepspecDrafter: expected contiguous positions, got ${JSON.stringify(positions)}`);
-  }
-}