@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/generation/speculative/source.ts

@@ -0,0 +1,144 @@
+// The draft contract: a provider owns the draft machinery for the server's
+// lifetime and opens row groups that the speculative group (the one served
+// lane, `execution/speculative-group.ts`) drafts, verifies and commits through.
+// A single request is a group of one row. Every drafter shares that one
+// verify/accept executor (designed in `02d723a:docs/design/speculative-decoding.md`
+// §1); the providers differ only in what fills the draft:
+//   - TwoModelProvider — mlx-lm parity (L1 oracle: mlx_lm.server --draft-model),
+//     a full second model; reads no target port.
+//   - AssistantProvider — the optiq KV-borrowing Gemma drafter
+//     (models/gemma4/assistant.ts; L2 oracle: optiq spec_generate): reads the
+//     target's donor K/V and the true hidden at each row's anchor.
+//   - DflashProvider, Dflash2Provider, DeepspecProvider — block drafters over
+//     the target's tapped hidden layers, projected once into context rows.
+//   - QwenMtpProvider, Glm52NativeMtpProvider — multi-token-prediction heads.
+//   - NgramProvider — model-free prompt lookup.
+// KV-borrowing providers read target state through the `target` view their
+// rows open with (see [[dspark-seam-kv-borrowing]]).
+
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { PreparedStateChange } from "../../contracts/portable/resources";
+import type { MlxDraftRows } from "./round";
+
+export { targetLacks, type AssistantRowsTarget, type HiddenLayerTapsTarget, type RecurrentMtpTarget, type TargetPort, type TargetView } from "../../contracts/mlx/draft-target";
+import type { TargetView } from "../../contracts/mlx/draft-target";
+
+/** How a draft group's rows choose their tokens: each row's draw, and the
+ *  selection the group resolved when it bound its requests. A drafter that
+ *  selects over its target's head binds one selection operation from
+ *  `selection` when its rows open; its draft loop calls that operation. */
+export interface DraftRowSampling {
+  /** Borrow [B,V] log-probabilities; return owned [B] IDs on device: each row's
+   *  own draw (argmax under the greedy draft policy). */
+  sample(logprobs: MlxArray, steps: readonly number[]): MlxArray;
+  readonly selection: DraftSelection;
+}
+
+/** The selection every row of one draft group shares, resolved once per request
+ *  from its sampling and the vocabulary its drafter declares.
+ *  - `sampled`: each row's `sample` over the full head's log-probabilities.
+ *  - `greedy`: every row drafts by argmax (the composition's `greedyDraft`), so a
+ *    drafter may select with the fused normalized argmax over the full head's
+ *    logits instead of calling `sample`; the selected IDs are identical.
+ *  - `listed`: the frequency-ranked draft vocabulary its drafter declared
+ *    (ascending token ids) and each row's coupled draw over logits [B, N]
+ *    restricted to them, returning full-vocabulary ids [B]. Bound only for
+ *    requests whose draft sampler reproduces on the subset. */
+export type DraftSelection =
+  | { readonly kind: "sampled" }
+  | { readonly kind: "greedy" }
+  | { readonly kind: "listed"; readonly ids: readonly number[];
+      sample(logits: MlxArray, head: import("./draft-vocab").DraftVocabularyHead, steps: readonly number[]): MlxArray };
+
+/** Request-owned constraints can propose known continuations without changing
+ * their committed state. The target verifier still decides every output. */
+export interface DraftRowConstraints {
+  propose(row: number, maxTokens: number): Promise<number[]>;
+}
+
+/** Method-owned state membership, independent of scheduling and cache tiers.
+ * Open/append borrow checkpoints; capture returns owned immutable state.
+ * Membership changes and capture occur only between committed rounds.
+ * The caller retains the provider's residency lease for the group's lifetime. */
+export interface DraftRowGroup extends MlxDraftRows {
+  /** Consume externally chosen target inputs at a committed round boundary.
+   * Context and tokens are borrowed rectangular windows; lengths retain each
+   * row's consumed prefix, excluding right padding and the pending output. */
+  consume?(tokens: MlxArray, context: MlxArray, lengths: readonly number[]): void | Promise<void>;
+  /** Optional: the proposals `draftDevice` would make from `pending` [B, depth], owned and left on the device,
+   *  without changing the rows' state (a recurrent head drafts under a round of its own state that it rolls back), so
+   *  a method can compare them with the target's own next samples instead of verifying them. A drafter may propose
+   *  fewer positions than `depth` (the first ones): a recurrent head proposes one. */
+  preview?(pending: readonly number[], depth: number, steps: readonly number[]): MlxArray | null;
+  readonly namespace: string;
+  readonly tapLayers: readonly number[];
+  readonly rowCount: number;
+  append(checkpoints: readonly DraftRowCheckpoint[]): void;
+  prepareAppend(checkpoints: readonly DraftRowCheckpoint[]): PreparedStateChange;
+  filterRows(keep: readonly number[]): void;
+  capture(row: number): DraftRowCheckpoint;
+  dispose(): void;
+}
+
+/** Draft preparation consumes each target chunk at the same batch geometry.
+ * Inputs are borrowed and contain only new tokens/context, with equal chunk
+ * lengths across rows. A null checkpoint appends a cold request. The provider
+ * owns alignment and retained hidden state; scheduling and persistence do not.
+ * Membership and capture occur between prefill calls. Capture returns owned
+ * state after at least one target token has been consumed for that row. */
+export interface DraftPrefillGroup {
+  readonly namespace: string;
+  /** Target prefill shape required by this provider's oracle. Most mlx-lm
+   *  providers leave the final prompt token pending (`tail-split`); the MTP
+   *  heads start from a full-prompt target forward. */
+  readonly prefillMode: "tail-split" | "full";
+  readonly tapLayers: readonly number[];
+  readonly rowCount: number;
+  append(checkpoints: readonly (DraftRowCheckpoint | null)[]): void;
+  prepareAppend(checkpoints: readonly (DraftRowCheckpoint | null)[]): PreparedStateChange;
+  prefill(tokens: MlxArray, context?: MlxArray): void | Promise<void>;
+  /** Resolve borrowed backing before its external restore lease is released. */
+  materialize(): void;
+  filterRows(keep: readonly number[]): void;
+  capture(row: number): DraftRowCheckpoint;
+  dispose(): void;
+}
+
+export interface GroupedDraftProvider {
+  readonly supportsExternalTokens?: boolean;
+  /** The frequency-ranked draft vocabulary (ascending token ids) this provider's
+   *  rows can select over for this target, or null. The binding resolves it once
+   *  for the bound target; a request whose draft sampler reproduces on it drafts
+   *  over it (`DraftSelection` `listed`). Absent: none. */
+  draftVocabulary?(target: TargetView): readonly number[] | null;
+  /** Resolve persistence identity without allocating a draft row. */
+  checkpointNamespace?(): string;
+  /** Draft state remains valid when target forwards run under a mounted
+   * adapter context. Unqualified learned draft graphs leave this absent. */
+  readonly supportsTargetAdapters?: boolean;
+  /** Target hidden layers every row of this provider taps for this target, in
+   * capture order. A provider whose rows tap must implement it; rows opened for
+   * the target tap exactly this list. The binding resolves it once and admits the
+   * provider only when the target forward can capture these layers. Absent: no taps. */
+  targetTapLayers?(target: TargetView): readonly number[];
+  openPrefill(options: { target: TargetView;
+    checkpoints: readonly (DraftRowCheckpoint | null)[] }): DraftPrefillGroup;
+  open(options: { target: TargetView; sampling: DraftRowSampling; constraints?: DraftRowConstraints;
+    checkpoints: readonly DraftRowCheckpoint[] }): DraftRowGroup;
+}
+
+/** Server-lifetime owner of the draft machinery (the loaded draft model). Its
+ * grouped rows are how every request drafts, one request or several. */
+export interface DraftProvider {
+  /** Human-readable id (registry id / path tail) for logs + cache namespacing. */
+  readonly id: string;
+  readonly weightsBytes: number;
+  readonly grouped: GroupedDraftProvider;
+  dispose(): void;
+}
+
+import { DraftRowCheckpoint } from "../../contracts/mlx/draft-checkpoint";
+export { type DraftRowCheckpoint } from "../../contracts/mlx/draft-checkpoint";
+export { type DraftProjection } from "../../contracts/mlx/draft-projection";
+
+import { DraftProjection } from "../../contracts/mlx/draft-projection";

package/src/generation/speculative/sources/assistant-source.ts

@@ -0,0 +1,64 @@
+import { artifactIdentity } from "../../../artifacts/identity";
+import { assistantGroups } from "../bindings/assistant-rows";
+// AssistantProvider — the optiq KV-borrowing Gemma "-assistant" drafter
+// (packages/inference/src/models/gemma4/assistant.ts) behind the grouped draft
+// rows. L2 oracle: optiq spec_generate. Its rows plug into the shared
+// speculative group, so they compose with grammar, logits processors,
+// admission accounting, and stats like any other provider's.
+//
+// KV-borrowing (why the rows read a TargetView): the drafter has NO cache of
+// its own — each draft step reads the target's donor K/V (last sliding + last
+// full layer, in chronological order) and, for the block's first step, the
+// target's hidden at each row's anchor position, which the rows keep from the
+// target's prefill and verify windows. See [[dspark-seam-kv-borrowing]].
+//
+// Drafting is GREEDY (the drafter's own argmax head), independent of the
+// request sampler — drafts are only proposals; the target's verify (which does
+// honor the sampler) decides every emitted token, so correctness holds at any
+// temperature and only the acceptance RATE moves.
+
+import { readdirSync, statSync } from "node:fs";
+import { join } from "node:path";
+import type { AssistantDrafterModel } from "../../../contracts/mlx/drafter";
+import { DRAFT_CHECKPOINT_SCHEMA } from "../../../contracts/mlx/draft-checkpoint";
+import { loadAssistantDrafter } from "../../../models/drafter-loaders";
+import type { DraftProvider } from "../source";
+
+/** Sum of the drafter's on-disk safetensors — admission accounting only. */
+function safetensorsBytes(dir: string): number {
+  let total = 0;
+  for (const f of readdirSync(dir))
+    if (f.endsWith(".safetensors")) total += statSync(join(dir, f)).size;
+  return total;
+}
+
+export class AssistantProvider implements DraftProvider {
+  readonly id: string;
+  readonly weightsBytes: number;
+  readonly grouped: import("../source").GroupedDraftProvider;
+
+  private constructor(
+    private readonly drafter: AssistantDrafterModel,
+    id: string,
+    weightsBytes: number, namespace: string,
+  ) {
+    this.id = id;
+    this.weightsBytes = weightsBytes;
+    this.grouped = assistantGroups(drafter,namespace);
+  }
+
+  static async load(modelDir: string): Promise<AssistantProvider> {
+    const drafter = await loadAssistantDrafter(modelDir);
+    const id = modelDir.split("/").filter(Boolean).at(-1)!;
+    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 AssistantProvider(drafter, id, safetensorsBytes(modelDir), `${DRAFT_CHECKPOINT_SCHEMA.assistant}:${identity}`);
+    } catch (error) { drafter.dispose(); throw error; }
+  }
+
+  dispose(): void {
+    this.drafter.dispose();
+  }
+}

package/src/generation/speculative/sources/deepspec-source.ts

@@ -0,0 +1,70 @@
+// DeepspecProvider — DeepSeek's released DSpark drafter checkpoints
+// (Gemma4DSparkModel, e.g. deepseek-ai/dspark_gemma4_12b_block7) behind the
+// grouped draft rows. The reference implementation is DeepSpec
+// (github.com/deepseek-ai/DeepSpec, MIT); the module port with transcription
+// notes lives in packages/inference/src/models/speculative/deepspec.ts. Oracle: DeepSpec's eval
+// at temperature 0 (RNG-free — verify degenerates to argmax token-match),
+// `02d723a:scripts/oracle/oracle-dspark-deepspec.py` → dspark-deepspec-compare.ts (deleted 2026-08-23; git history).
+//
+// Row mapping (why this fits the shared speculative group): the reference
+// recycles target hidden states from each verify pass into the drafter's
+// growing context — exactly the tapped context the rows consume:
+//   prefill                  → project the prompt's tapped hiddens into the
+//                              per-layer context K/V rows (their DynamicCache
+//                              of PROJECTED context rows, positions 0..Lp-1);
+//   draft                    → one block forward per row against its rows
+//                              (anchor = the row's pending token at its
+//                              context length); confidence truncation may
+//                              return 0..depth tokens (ℓ=0 = plain step);
+//   commit                   → project + append ONLY the accepted rows
+//                              [anchor + accepted drafts] (their crop()
+//                              semantics: rejected rows never enter the
+//                              context).
+// Context rows are cached post-norm-post-rope — bit-equivalent to their
+// in-round concat (per-row RMSNorm/RoPE; argued in projectContextKVRows's doc).
+
+import { artifactIdentity } from "../../../artifacts/identity";
+import { deepspecGroups } from "../bindings/deepspec-rows";
+import { readdirSync, statSync } from "node:fs";
+import { join } from "node:path";
+import type { DeepspecDrafterModel } from "../../../contracts/mlx/drafter";
+import { loadDeepspecDrafter } from "../../../models/drafter-loaders";
+import type { DraftProvider } from "../source";
+
+function safetensorsBytes(dir: string): number {
+  let total = 0;
+  for (const f of readdirSync(dir))
+    if (f.endsWith(".safetensors")) total += statSync(join(dir, f)).size;
+  return total;
+}
+
+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: DeepspecDrafterModel, 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 loadDeepspecDrafter(modelDir);
+    const id = modelDir.split("/").filter(Boolean).at(-1)!;
+    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; }
+  }
+
+  dispose(): void {
+    this.drafter.dispose(); // owned MlxArrays (transpose views + mmap'd weights)
+  }
+}

package/src/generation/speculative/sources/dflash-source.ts

@@ -0,0 +1,84 @@
+// DflashProvider — DSpark (faithful DFlash + Markov + confidence,
+// packages/inference/src/models/speculative/dflash.ts) behind the grouped draft
+// rows. L3 (KL/quality-gated). Its rows plug into the shared speculative group.
+//
+// H_ctx (why the rows tap): DSpark drafts by attending to a GROWING
+// multi-layer context — the target's tapped hiddens over the accepted stream
+// (paper Eq 2-3). The provider declares its tap layers, so the group taps the
+// target's prefill (→ seeds each row's context) and every verify forward (→
+// grown at commit by the accepted window; rejected tips are never appended,
+// mirroring the target-cache trim). See [[dspark-seam-kv-borrowing]].
+//
+// v1 acceptance: the group verifies with TOKEN-MATCH acceptance
+// (mlx-lm/optiq style — lossless at any temperature), NOT the paper's
+// distribution-level rejection-sampling verify. Greedy is identical; temp>0 is
+// still lossless but lower-acceptance. Drafting here is greedy (proposals; the
+// target verify decides). Confidence-scheduled pruning (Alg 1) may return fewer
+// than the requested tokens: active when the checkpoint carries STS thresholds
+// (cfg.sts, §3.2.1) or via the MLX_BUN_DSPARK_MINCONF override; uncalibrated
+// checkpoints draft fixed-length.
+
+import { artifactIdentity } from "../../../artifacts/identity";
+import { projectedDraftGroups } from "../bindings/projected-draft-rows";
+import { readdirSync, statSync } from "node:fs";
+import { join } from "node:path";
+import type { DsparkDrafterModel } from "../../../contracts/mlx/drafter";
+import { loadDsparkDrafter } from "../../../models/drafter-loaders";
+import type { DraftProvider } from "../source";
+import { targetLacks } from "../source";
+import { runtimeValue } from "../../../runtime/config";
+
+function safetensorsBytes(dir: string): number {
+  let total = 0;
+  for (const f of readdirSync(dir))
+    if (f.endsWith(".safetensors")) total += statSync(join(dir, f)).size;
+  return total;
+}
+
+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
+   *  group never asks for more positions than the block was trained for. */
+  readonly gamma: number;
+
+  private constructor(private readonly drafter: DsparkDrafterModel, id: string, weightsBytes: number, namespace: string) {
+    this.id = id;
+    this.weightsBytes = weightsBytes;
+    this.gamma = drafter.cfg.gamma;
+    this.grouped = projectedDraftGroups(namespace, drafter.cfg.tapLayers, target => {
+      if (!target.hiddenLayerTaps) throw targetLacks("hiddenLayerTaps");
+      const projection = target.hiddenLayerTaps.projection;
+      const minConf = runtimeValue("MLX_BUN_DSPARK_MINCONF");
+      return {
+        namespace, schema: "dflash-context-v1",
+        layers: drafter.cfg.nLayers,
+        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 = await loadDsparkDrafter(modelDir); // variant dispatch (dspark|legacy dflash)
+    const id = modelDir.split("/").filter(Boolean).at(-1)!;
+    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; }
+  }
+
+  dispose(): void {
+    // Provider-owned weights are released when the mounted draft is closed.
+    this.drafter.dispose();
+  }
+}

package/src/generation/speculative/sources/dflash2-source.ts

@@ -0,0 +1,86 @@
+import { rowOperationFor } from "../../../models/graph";
+// Dflash2Provider — the DFlash 2 block drafter (models/speculative/dflash2.ts)
+// behind the grouped draft seam. Context rows are the target's tapped layer
+// outputs over the accepted stream (prefill seeds them; each commit appends
+// the anchor plus accepted rows), projected once into per-layer K/V.
+// Drafting is greedy; the shared verify decides acceptance.
+//
+// The drafter loads in the basis it was trained on, its projections quantized
+// to the width its composer passes (the draft registry passes 4: 1.0 GB
+// resident instead of the 3.85 GB BF16 checkpoint). A target whose residual
+// stream carries a TurboQuant R1 fold declares it on its draft target
+// (`residualBasis`), and one whose composition measured them the kernels of
+// those projections (`draftProjections`); binding a draft group applies both
+// before the first draft.
+
+import { artifactIdentity } from "../../../artifacts/identity";
+import { projectedDraftGroups, type ProjectedDraftMethod } from "../bindings/projected-draft-rows";
+
+import { readdirSync, statSync } from "node:fs";
+import { join } from "node:path";
+import type { DraftProvider, GroupedDraftProvider } from "../source";
+import { targetLacks } from "../source";
+import type { Dflash2DrafterModel } from "../../../contracts/mlx/drafter";
+import { loadDflash2Drafter } from "../../../models/drafter-loaders";
+
+export class Dflash2Provider implements DraftProvider {
+  readonly id: string;
+  readonly weightsBytes: number;
+  readonly grouped: GroupedDraftProvider;
+  readonly gamma: number;
+
+  private constructor(private readonly drafter: Dflash2DrafterModel, id: string, weightsBytes: number, namespace: string) {
+    this.id = id;
+    this.weightsBytes = weightsBytes;
+    this.gamma = drafter.gamma;
+    this.grouped = projectedDraftGroups(namespace, drafter.tapLayers, target => {
+      if (!target.hiddenLayerTaps) throw targetLacks("hiddenLayerTaps");
+      const layers = drafter.cfg.numTargetLayers;
+      if (target.hiddenLayerTaps.layerCount !== layers) throw new Error(
+        `DFlash 2 drafter was trained for a ${layers}-layer target; this model has ` +
+        `${target.hiddenLayerTaps.layerCount} layers — wrong (target, drafter) pairing`);
+      drafter.bindTarget(target.residualBasis ?? null, target.draftProjections ?? null);
+      const projection = target.hiddenLayerTaps.projection;
+      // The block for the rows drafted this round: anchors × columns rows index the drafter's draft table.
+      const draftDevice: NonNullable<ProjectedDraftMethod["draftDevice"]> = (context, pending, positions, depth) => {
+        const columns = drafter.blockColumns(depth);
+        return rowOperationFor(drafter.draft,(pending.length * columns))(context, projection, pending, positions, columns);
+      };
+      return {
+        namespace, schema: "dflash2-context-v1", layers: drafter.cfg.layers,
+        // The B × L context rows a prefill or a commit appends index the drafter's context tables.
+        project(hidden, positions) {
+          const width = (hidden.shape[0]! * hidden.shape[1]!);
+          using context = rowOperationFor(drafter.projectContext,width)(hidden);
+          return rowOperationFor(drafter.projectContextKVRows,width)(context, positions);
+        },
+        draft(context, pending, positions, depth) {
+          const steps = drafter.blockColumns(depth) - 1;
+          if (!steps) return pending.map(() => []);
+          using tokens = draftDevice(context, pending, positions, depth);
+          const flat = tokens.toIntTokens();
+          return pending.map((_, row) => flat.slice(row * steps, (row + 1) * steps));
+        },
+        draftDevice,
+      };
+    });
+  }
+
+  /** `bits`: the drafter's projection width (0 keeps BF16). Persisted context is
+   *  keyed by the drafter and that width; the target's own identity scopes it
+   *  to one residual basis. */
+  static async load(modelDir: string, options: { bits: number }): Promise<Dflash2Provider> {
+    const { bits } = options;
+    const drafter = await loadDflash2Drafter(modelDir, { bits });
+    const id = modelDir.split("/").filter(Boolean).at(-1)!;
+    try {
+      const files = readdirSync(modelDir).filter(file => file.endsWith(".safetensors"));
+      const identity = await artifactIdentity(await Bun.file(`${modelDir}/config.json`).text(),
+        files.map(name => ({ name, path: join(modelDir, name) })));
+      const bytes = files.reduce((total, file) => total + statSync(join(modelDir, file)).size, 0);
+      return new Dflash2Provider(drafter, id, bytes, `dflash2-context-v1:${identity}:b${bits}`);
+    } catch (error) { drafter.dispose(); throw error; }
+  }
+
+  dispose(): void { this.drafter.dispose(); }
+}

package/src/generation/speculative/sources/glm52-mtp-source.ts

@@ -0,0 +1,66 @@
+// Native GLM-5.2 MTP draft provider.
+//
+// The target and MTP row share embeddings, the output head, and the resident
+// dense target weights. Only the MTP layer's own dense tensors and int8 routed
+// experts add model weight residency. The shared speculative group owns
+// verification, grammar composition, target-cache rollback, sampling, and
+// telemetry. Prompt lookup remains the model-free alternative DraftProvider;
+// selecting one provider per request prevents two independent draft histories
+// from advancing the target or grammar state twice.
+
+import type { NativeMtpHead } from "../../../contracts/mlx/drafter";
+import { nativeMtpHeadOf } from "../../../models/drafter-loaders";
+import type { DraftProvider, GroupedDraftProvider, TargetView, DraftRowSampling, DraftRowGroup, DraftPrefillGroup } from "../source";
+import { DRAFT_CHECKPOINT_SCHEMA, type DraftRowCheckpoint } from "../../../contracts/mlx/draft-checkpoint";
+import { Glm52MtpRows } from "../bindings/glm52-mtp-rows";
+import { type Glm52MtpRowState } from "../../../state/speculative/glm52-mtp-state";
+import { captureGlm52MtpState, restoreGlm52MtpState } from "../../../state/speculative/glm52-mtp-state";
+import { applyStateChanges, disposeResources } from "../../../runtime/resources";
+
+export class Glm52NativeMtpProvider implements DraftProvider {
+  readonly grouped: GroupedDraftProvider = {
+    checkpointNamespace: () => DRAFT_CHECKPOINT_SCHEMA.glm52Mtp,
+    open: options => this.#openRows(options.target, options.sampling, options.checkpoints),
+    openPrefill: options => this.#openRows(options.target, null, options.checkpoints),
+  };
+  readonly id = "glm52-native-mtp";
+  readonly weightsBytes = 0;
+  readonly #head: NativeMtpHead;
+
+  /** Binds the native MTP head `model` declares (`GraphCapabilities.nativeDraft`); it owns every shared and MTP weight. */
+  constructor(readonly model: object) {
+    this.#head = nativeMtpHeadOf(model);
+  }
+
+  #openRows(target: TargetView, sampling: DraftRowSampling | null,
+    checkpoints: readonly (DraftRowCheckpoint | null)[]): DraftRowGroup & DraftPrefillGroup {
+    if (target.identity !== this.model) throw new Error("native MTP provider was opened for a different target");
+    const rows = new Glm52MtpRows(this.#head, sampling);
+    const prepareAppend = (checkpoints: readonly (DraftRowCheckpoint | null)[]) => {
+      const states: (Glm52MtpRowState | null)[] = [];
+      try {
+        for (const checkpoint of checkpoints) states.push(checkpoint ? restoreGlm52MtpState(checkpoint, () => rows.makeCache()) : null);
+        return rows.prepareAppend(states);
+      } finally { disposeResources(states.flatMap(state => state ? [state.cache, state.hidden] : [])); }
+    };
+    const append = (states: readonly (DraftRowCheckpoint | null)[]) => applyStateChanges([() => prepareAppend(states)]);
+    try { append(checkpoints); } catch (error) { rows.dispose(); throw error; }
+    return {
+      namespace: DRAFT_CHECKPOINT_SCHEMA.glm52Mtp, prefillMode: "full", tapLayers: [],
+      get rowCount() { return rows.rowCount; },
+      append, prepareAppend, prefill: (tokens, context) => rows.prefill(tokens, context!),
+      materialize: rows.materialize.bind(rows), filterRows: rows.filterRows.bind(rows),
+      draft: rows.draft.bind(rows), commit: rows.commit.bind(rows),
+      capture(row) {
+        const state = rows.extractRow(row);
+        try { return captureGlm52MtpState(state); }
+        finally { state.cache.dispose(); state.hidden.dispose(); }
+      },
+      dispose: rows.dispose.bind(rows),
+    };
+  }
+
+  dispose(): void {
+    // The target model owns every shared/MTP weight and the expert tier.
+  }
+}