mlx-bun 0.4.0 → 1.0.0

package/src/config.ts

@@ -1,399 +0,0 @@
-// Model config + quantization metadata reader for the load path.
-// Parses HF config.json (Gemma 4 unified layout: text_config nested),
-// the OptiQ per-layer quantization map, and kv_config.json when present.
-
-import {
-  parseGlm52Config,
-  type Glm52Config,
-} from "./model/glm52-config";
-
-export interface RopeParams {
-  ropeTheta: number;
-  ropeType: string;
-  partialRotaryFactor: number;
-  /** ProportionalRoPE frequency multiplier (rope_utils: factor, default 1). */
-  factor: number;
-}
-
-export interface TextConfig {
-  hiddenSize: number;
-  numHiddenLayers: number;
-  numAttentionHeads: number;
-  numKeyValueHeads: number;
-  headDim: number;
-  /** Global (full-attention) layers can use different KV geometry. */
-  numGlobalKeyValueHeads: number;
-  globalHeadDim: number;
-  attentionKEqV: boolean;
-  intermediateSize: number;
-  hiddenActivation: string;
-  rmsNormEps: number;
-  vocabSize: number;
-  maxPositionEmbeddings: number;
-  slidingWindow: number;
-  /** "sliding_attention" | "full_attention" per layer. */
-  layerTypes: string[];
-  /** e2b/e4b: per-layer input embedding width (0 = disabled). */
-  hiddenSizePerLayerInput: number;
-  vocabSizePerLayerInput: number;
-  /** Trailing layers that share KV with earlier same-type layers. */
-  numKvSharedLayers: number;
-  /** 26B-A4B MoE block: dense MLP + routed experts per layer. */
-  enableMoeBlock: boolean;
-  numExperts: number;
-  topKExperts: number;
-  moeIntermediateSize: number;
-  /** Qwen3-MoE sparse-block gating (qwen3_moe): a layer is MoE unless it is
-   *  in `mlpOnlyLayers` and `(idx+1) % decoderSparseStep == 0` selects it.
-   *  `normTopkProb` renormalizes the top-k gate weights. Defaults are inert
-   *  for non-qwen3_moe families. */
-  decoderSparseStep: number;
-  mlpOnlyLayers: number[];
-  normTopkProb: boolean;
-  /** Qwen3.5 hybrid (gated-DeltaNet) geometry; 0/false for other families. */
-  linearNumValueHeads: number;
-  linearNumKeyHeads: number;
-  linearKeyHeadDim: number;
-  linearValueHeadDim: number;
-  linearConvKernelDim: number;
-  /** Every `fullAttentionInterval`-th layer is full-attention (Qwen3.5). */
-  fullAttentionInterval: number;
-  /** Qwen3.5 full-attention output gate: q_proj emits 2× and gates the output. */
-  attnOutputGate: boolean;
-  /** Rotary fraction of head_dim (Qwen3.5 full attention: 0.25 → 64 of 256). */
-  partialRotaryFactor: number;
-  ropeParameters: Record<string, RopeParams>;
-  finalLogitSoftcapping: number | null;
-  tieWordEmbeddings: boolean;
-  bosTokenId: number;
-  eosTokenId: number | number[];
-  /** DiffusionGemma: fixed denoising-canvas length (256). Undefined elsewhere. */
-  canvasLength?: number;
-}
-
-export interface QuantSpec {
-  bits: number;
-  groupSize: number;
-  mode: string;
-  /** `mode: "trellis"` (packed trellis-coded weights, src/model/trellis-linear.ts):
-   *  bits = k per weight, groupSize = block T, plus the codec geometry. */
-  trellis?: TrellisSpec;
-}
-
-export interface TrellisSpec {
-  L: number;
-  code: string;
-  /** 1 = coded along the input dim (stored [out, in·k/32]); 0 = along the
-   *  output dim (stored [in, out·k/32], i.e. Wᵀ coded along its last axis). */
-  axis: 0 | 1;
-}
-
-export interface QuantizationConfig {
-  default: QuantSpec;
-  /** Per-module overrides keyed by module path (e.g. "language_model.model.embed_tokens"). */
-  perLayer: Map<string, QuantSpec>;
-}
-
-export interface KvQuantSpec {
-  layerIdx: number;
-  bits: number;
-  groupSize: number;
-}
-
-/** TurboQuant scheme (docs/design/turboquant.md): rotation-based KV
- *  quantization, a distinct axis from the uniform/per-layer affine kvQuant
- *  above (composes with it in principle; v1 ships as a standalone CLI-only
- *  runtime lever, same class as uniform kvBits — see
- *  src/generate.ts maybeQuantizeKv). Valid kBits: {2,4,5,8}; vBits:
- *  {2,3,4,5,8} (docs/design/turboquant.md's supported Lloyd-Max tables). */
-export interface TurboQuantScheme {
-  kBits: number;
-  vBits: number;
-}
-
-export const TURBOQUANT_VALID_KBITS = [2, 4, 5, 8] as const;
-export const TURBOQUANT_VALID_VBITS = [2, 3, 4, 5, 8] as const;
-
-/** Parse `turbo` / `turbo:k<bits>v<bits>` (default k8v3). Returns null for
- *  non-turbo strings (caller falls through to the uniform/config parser);
- *  throws Error with a user-facing message for a malformed turbo spec. */
-export function parseTurboQuantScheme(raw: string): TurboQuantScheme | null {
-  if (raw !== "turbo" && !raw.startsWith("turbo:")) return null;
-  if (raw === "turbo") return { kBits: 8, vBits: 3 };
-  const m = /^turbo:k(\d+)v(\d+)$/.exec(raw);
-  if (!m) throw new Error(`--kv-quant turbo spec must look like "turbo:k<bits>v<bits>" (got "${raw}")`);
-  const kBits = Number(m[1]);
-  const vBits = Number(m[2]);
-  if (!(TURBOQUANT_VALID_KBITS as readonly number[]).includes(kBits))
-    throw new Error(`--kv-quant turbo: kBits must be one of ${TURBOQUANT_VALID_KBITS.join(",")} (got ${kBits})`);
-  if (!(TURBOQUANT_VALID_VBITS as readonly number[]).includes(vBits))
-    throw new Error(`--kv-quant turbo: vBits must be one of ${TURBOQUANT_VALID_VBITS.join(",")} (got ${vBits})`);
-  return { kBits, vBits };
-}
-
-export interface ModelConfig {
-  modelDir: string;
-  modelType: string;
-  architectures: string[];
-  dtype: string;
-  text: TextConfig;
-  quantization: QuantizationConfig | null;
-  kvQuant: KvQuantSpec[] | null;
-  hasVisionSidecar: boolean;
-  eosTokenIds: number[];
-  raw: Record<string, unknown>;
-}
-
-function parseRope(raw: Record<string, any>): Record<string, RopeParams> {
-  const out: Record<string, RopeParams> = {};
-  for (const [kind, p] of Object.entries(raw ?? {})) {
-    // transformers ≥5.10 configs mix flat scalars into the per-attention-type
-    // map (`rope_parameters: { full_attention: {...}, rope_theta: null,
-    // rope_type: "default", ... }` — the DeepSpec drafter ships this shape);
-    // only object entries are attention-type params.
-    if (!p || typeof p !== "object") continue;
-    out[kind] = {
-      ropeTheta: p.rope_theta,
-      ropeType: p.rope_type ?? "default",
-      partialRotaryFactor: p.partial_rotary_factor ?? 1.0,
-      factor: p.factor ?? 1.0,
-    };
-  }
-  return out;
-}
-
-export function parseQuantization(raw: Record<string, any> | undefined): QuantizationConfig | null {
-  if (!raw) return null;
-  const def: QuantSpec = {
-    bits: raw.bits,
-    groupSize: raw.group_size,
-    mode: raw.mode ?? "affine",
-  };
-  const perLayer = new Map<string, QuantSpec>();
-  for (const [key, value] of Object.entries(raw)) {
-    if (key === "bits" || key === "group_size" || key === "mode") continue;
-    if (typeof value === "object" && value !== null) {
-      const spec: QuantSpec = {
-        bits: (value as any).bits ?? def.bits,
-        groupSize: (value as any).group_size ?? def.groupSize,
-        mode: (value as any).mode ?? def.mode,
-      };
-      const tr = (value as any).trellis;
-      if (spec.mode === "trellis") {
-        if (!tr || typeof tr !== "object") throw new Error(`quantization.${key}: mode "trellis" needs a trellis {L, code, axis} block`);
-        spec.trellis = { L: Number(tr.L), code: String(tr.code ?? "1mad"), axis: tr.axis === 0 ? 0 : 1 };
-      }
-      perLayer.set(key, spec);
-    } else if (value === false) {
-      // mlx convention: `"layer": false` means not quantized
-      perLayer.set(key, { bits: 0, groupSize: 0, mode: "none" });
-    }
-  }
-  return { default: def, perLayer };
-}
-
-/** Quant spec for a module path, or null if the module is unquantized. */
-export function quantFor(q: QuantizationConfig | null, modulePath: string): QuantSpec | null {
-  if (!q) return null;
-  const spec = q.perLayer.get(modulePath) ?? q.default;
-  return spec.mode === "none" ? null : spec;
-}
-
-/** DiffusionGemma's config.json ships only token ids + canvas_length + the
- *  quant map — the architecture dims live in optiq's `config.py` TextConfig
- *  defaults (and layer_types / rope_parameters are computed in __post_init__).
- *  We reproduce those defaults in snake_case so the generic parser below picks
- *  them up; any field actually present in config.json still overrides. Source:
- *  optiq/vlm/_mlxvlm/models/diffusion_gemma/config.py. */
-function diffusionGemmaRawDefaults(): Record<string, any> {
-  const numLayers = 30;
-  const pattern = ["sliding_attention", "sliding_attention", "sliding_attention",
-    "sliding_attention", "sliding_attention", "full_attention"];
-  const layer_types = Array.from({ length: numLayers }, (_, i) => pattern[i % pattern.length]);
-  layer_types[numLayers - 1] = "full_attention"; // last forced full
-  return {
-    hidden_size: 2816,
-    num_hidden_layers: numLayers,
-    num_attention_heads: 16,
-    num_key_value_heads: 8,
-    num_global_key_value_heads: 2,
-    head_dim: 256,
-    global_head_dim: 512,
-    intermediate_size: 2112,
-    moe_intermediate_size: 704,
-    hidden_activation: "gelu_pytorch_tanh",
-    rms_norm_eps: 1e-6,
-    vocab_size: 262144,
-    max_position_embeddings: 262144,
-    sliding_window: 1024,
-    layer_types,
-    enable_moe_block: true,
-    num_experts: 128,
-    top_k_experts: 8,
-    final_logit_softcapping: 30.0,
-    tie_word_embeddings: true,
-    bos_token_id: 2,
-    eos_token_id: 1,
-    rope_parameters: {
-      sliding_attention: { rope_type: "default", rope_theta: 10000.0 },
-      full_attention: { rope_type: "proportional", partial_rotary_factor: 0.25, rope_theta: 1000000.0 },
-    },
-  };
-}
-
-export async function loadModelConfig(modelDir: string): Promise<ModelConfig> {
-  const raw = (await Bun.file(`${modelDir}/config.json`).json()) as Record<string, any>;
-  const modelType = raw.model_type;
-  const isDiffusion = modelType === "diffusion_gemma";
-  const isGlm52 = modelType === "glm_moe_dsa";
-  // The converted Colibri artifact keeps stop metadata split across the two
-  // HF files. Preserve the old config-only behavior for every other family;
-  // GLM-5.2 explicitly unions both lists through its dedicated parser.
-  let glm52: Glm52Config | null = null;
-  if (isGlm52) {
-    const generationFile = Bun.file(`${modelDir}/generation_config.json`);
-    const generation = await generationFile.exists()
-      ? await generationFile.json() as Record<string, any>
-      : null;
-    glm52 = parseGlm52Config(modelDir, raw, generation);
-  }
-  // Gemma 4 unified nests the LM config; plain text models keep it at top level.
-  // DiffusionGemma keeps it flat too, but omits all dims — backfill from defaults.
-  const baseT = (raw.text_config ?? raw) as Record<string, any>;
-  const t = isDiffusion ? { ...diffusionGemmaRawDefaults(), ...baseT } : baseT;
-  const isLlama = modelType === "llama";
-  const isQwen35 = typeof modelType === "string" && modelType.startsWith("qwen3_5");
-  // Plain Qwen3 (Qwen3ForCausalLM, e.g. Qwen3-Embedding): a flat HF config with
-  // a scalar rope_theta and no rope_parameters map — handled like llama below.
-  const isQwen3 = modelType === "qwen3";
-  // Qwen3.5 rope_parameters is a flat dict ({type, rope_theta, mrope_section,
-  // partial_rotary_factor}), not the gemma per-attention-type map — and
-  // type "default" means plain partial nn.RoPE (mrope_section ignored for text).
-  const qwenRope = (t.rope_parameters ?? {}) as Record<string, any>;
-  // Qwen3.8 adds output_gate_type ("swish"), which every implementation
-  // (transformers ground truth, mlx-lm pinned + main) currently ignores —
-  // the attention output gate is hardcoded o_proj(out·σ(gate)). Only accept
-  // values verified to mean that; an unknown value must fail at load, not
-  // silently mis-gate a future checkpoint where the field starts mattering.
-  if (isQwen35 && t.output_gate_type != null && t.output_gate_type !== "swish")
-    throw new Error(
-      `qwen3_5: unverified output_gate_type "${t.output_gate_type}" (only "swish" ` +
-        `is verified to mean the standard sigmoid output gate)`,
-    );
-
-  const text: TextConfig = {
-    hiddenSize: glm52?.hiddenSize ?? t.hidden_size,
-    numHiddenLayers: glm52?.numHiddenLayers ?? t.num_hidden_layers,
-    numAttentionHeads: glm52?.numAttentionHeads ?? t.num_attention_heads,
-    numKeyValueHeads: glm52?.numKeyValueHeads ?? t.num_key_value_heads,
-    headDim: glm52?.qkHeadDim ?? t.head_dim ?? Math.floor(t.hidden_size / t.num_attention_heads),
-    numGlobalKeyValueHeads:
-      glm52?.numKeyValueHeads ?? t.num_global_key_value_heads ?? t.num_key_value_heads,
-    globalHeadDim:
-      glm52?.qkHeadDim ??
-      t.global_head_dim ??
-      t.head_dim ??
-      Math.floor(t.hidden_size / t.num_attention_heads),
-    attentionKEqV: t.attention_k_eq_v ?? false,
-    intermediateSize: glm52?.intermediateSize ?? t.intermediate_size,
-    hiddenActivation: t.hidden_activation ?? t.hidden_act ?? "gelu_pytorch_tanh",
-    rmsNormEps: glm52?.rmsNormEps ?? t.rms_norm_eps,
-    vocabSize: glm52?.vocabSize ?? t.vocab_size,
-    maxPositionEmbeddings:
-      glm52?.maxPositionEmbeddings ?? t.max_position_embeddings,
-    slidingWindow: t.sliding_window ?? 0,
-    layerTypes: glm52
-      ? Array(glm52.numHiddenLayers).fill("full_attention")
-      : t.layer_types ?? (isLlama ? Array(t.num_hidden_layers).fill("full_attention") : []),
-    hiddenSizePerLayerInput: t.hidden_size_per_layer_input ?? 0,
-    vocabSizePerLayerInput:
-      t.vocab_size_per_layer_input ?? glm52?.vocabSize ?? t.vocab_size,
-    numKvSharedLayers: t.num_kv_shared_layers ?? 0,
-    enableMoeBlock: glm52 ? true : t.enable_moe_block ?? false,
-    numExperts: glm52?.numRoutedExperts ?? t.num_experts ?? 0,
-    // gemma4 uses `top_k_experts`; qwen3_moe uses `num_experts_per_tok`.
-    topKExperts:
-      glm52?.numExpertsPerToken ?? t.top_k_experts ?? t.num_experts_per_tok ?? 0,
-    moeIntermediateSize: glm52?.moeIntermediateSize ?? t.moe_intermediate_size ?? 0,
-    decoderSparseStep: glm52 ? 1 : t.decoder_sparse_step ?? 1,
-    mlpOnlyLayers: glm52
-      ? Array.from({ length: glm52.firstKDenseReplace }, (_, layer) => layer)
-      : t.mlp_only_layers ?? [],
-    normTopkProb: glm52?.normTopkProb ?? t.norm_topk_prob ?? false,
-    linearNumValueHeads: t.linear_num_value_heads ?? 0,
-    linearNumKeyHeads: t.linear_num_key_heads ?? 0,
-    linearKeyHeadDim: t.linear_key_head_dim ?? 0,
-    linearValueHeadDim: t.linear_value_head_dim ?? 0,
-    linearConvKernelDim: t.linear_conv_kernel_dim ?? 0,
-    fullAttentionInterval: t.full_attention_interval ?? 0,
-    attnOutputGate: t.attn_output_gate ?? false,
-    partialRotaryFactor: glm52
-      ? glm52.qkRopeHeadDim / glm52.qkHeadDim
-      : qwenRope.partial_rotary_factor ?? t.partial_rotary_factor ?? 1.0,
-    ropeParameters: glm52
-      ? {
-          full_attention: {
-            ropeTheta: glm52.ropeTheta,
-            ropeType: "default",
-            partialRotaryFactor: glm52.qkRopeHeadDim / glm52.qkHeadDim,
-            factor: 1.0,
-          },
-        }
-      : isQwen35
-      ? {
-          full_attention: {
-            ropeTheta: qwenRope.rope_theta ?? t.rope_theta ?? 10000,
-            ropeType: "default",
-            partialRotaryFactor: qwenRope.partial_rotary_factor ?? 1.0,
-            factor: 1.0,
-          },
-        }
-      : t.rope_parameters
-        ? parseRope(t.rope_parameters)
-        : isLlama || isQwen3
-          ? {
-              full_attention: {
-                ropeTheta: t.rope_theta ?? 10000,
-                ropeType: "default",
-                partialRotaryFactor: 1.0,
-                factor: 1.0,
-              },
-            }
-          : {},
-    finalLogitSoftcapping: t.final_logit_softcapping ?? null,
-    tieWordEmbeddings: t.tie_word_embeddings ?? raw.tie_word_embeddings ?? false,
-    bosTokenId: t.bos_token_id ?? raw.bos_token_id,
-    eosTokenId: glm52?.eosTokenIds ?? t.eos_token_id ?? raw.eos_token_id,
-    canvasLength: raw.canvas_length ?? t.canvas_length,
-  };
-
-  let kvQuant: KvQuantSpec[] | null = null;
-  if (await Bun.file(`${modelDir}/kv_config.json`).exists()) {
-    const kv = (await Bun.file(`${modelDir}/kv_config.json`).json()) as any[];
-    kvQuant = kv.map((e) => ({
-      layerIdx: e.layer_idx,
-      bits: e.bits,
-      groupSize: e.group_size,
-    }));
-  }
-
-  const eos = glm52?.eosTokenIds ?? raw.eos_token_id ?? text.eosTokenId;
-  return {
-    modelDir,
-    modelType,
-    architectures: raw.architectures ?? [],
-    dtype: raw.dtype ?? "bfloat16",
-    text,
-    // Colibri's config retains the SOURCE FP8 quantization_config, while the
-    // serving shards are converted U8 + `.qs` tensors whose int4/int8/group
-    // geometry is inferred and validated by Glm52Container. Advertising the
-    // source FP8 block as mlx affine quantization would be actively wrong.
-    quantization: glm52
-      ? null
-      : parseQuantization(raw.quantization ?? raw.quantization_config),
-    kvQuant,
-    hasVisionSidecar: await Bun.file(`${modelDir}/optiq_vision.safetensors`).exists(),
-    eosTokenIds: Array.isArray(eos) ? eos : [eos],
-    raw,
-  };
-}

package/src/contracts/completion.ts

@@ -1,16 +0,0 @@
-import type { Cancellation } from "./generation";
-
-/** Application-facing calls. The implementation owns tokenization, numerical
- * method and transport; callers provide request data and consume results. */
-export interface CompletionClient<Request, Result> {
-  complete(request: Request): Promise<Result>;
-}
-
-export interface BatchCompletionClient<Request, Result> extends CompletionClient<Request, Result> {
-  /** Results preserve input order; the implementation chooses execution groups. */
-  completeBatch(requests: readonly Request[]): Promise<Result[]>;
-}
-
-export interface TaskClient<Request, Progress, Result> {
-  run(request: Request, report: (progress: Progress) => void, cancellation?: Cancellation): Promise<Result>;
-}

package/src/contracts/execution.ts

@@ -1,60 +0,0 @@
-/** Request facts only. Native resources remain with the preparation owner. */
-export interface ExecutionRequirements {
-  readonly hasVision: boolean;
-  readonly hasAdapters: boolean;
-  readonly hasRepetitionPenalty: boolean;
-  readonly userSeed: boolean;
-  readonly kvQuant: boolean;
-  readonly turboQuant: boolean;
-  readonly hasLogitsExtras: boolean;
-  readonly hasGrammar: boolean;
-  readonly wantsLogprobs: boolean;
-  readonly hasDraft: boolean;
-}
-
-export interface ExecutionCapabilities {
-  readonly method: "autoregressive" | "denoising";
-  readonly continuous: boolean;
-  readonly quantizedBatch: boolean;
-  readonly grammarBatch: boolean;
-  readonly adapterBatch?: boolean;
-  readonly pagedBatch?: boolean;
-  readonly checkpoints: boolean;
-  /** Bound ordinary driver can restore/capture this request configuration. */
-  readonly sharedCheckpoints?: boolean;
-  /** Model-qualified speculative execution can retain this request's KV codec. */
-  readonly speculativeKvQuant?: boolean;
-  readonly speculativeTurboQuant?: boolean;
-  readonly turboQuantBatch?: boolean;
-  readonly speculativeLogprobs?: boolean;
-  /** The grouped provider supports the target's mounted adapter context. */
-  readonly sharedSpeculativeAdapters?: boolean;
-  /** Methods supplied by the model's shared execution binding. */
-  readonly groupedMethods?: readonly string[];
-  /** A graph-owned compiled step exists; cache geometry can still decline it. */
-  readonly compiledDecode?: boolean;
-}
-
-export interface ExecutionFeatures {
-  readonly pagedKv: boolean;
-  readonly fill: boolean;
-  readonly compiledDecode?: boolean;
-  readonly grammarJump?: boolean;
-}
-
-/** Selected once; consumers execute/report these values without reselecting. */
-export interface ResolvedExecution {
-  /** Implementation-owned method ID. The built-in planner retains its known
-   * methods; another model may register a different set without editing this contract. */
-  readonly method: string;
-  readonly mechanism: "serial" | "continuous";
-  readonly pagedKv: boolean;
-  readonly promptCache: boolean;
-  readonly checkpoint: boolean;
-  readonly fill: boolean;
-  /** Permission to attempt the bound compiled step, never a promise that
-   * every dynamic row/state shape supports replay. */
-  readonly compiledDecode: boolean;
-  readonly grammarJump: boolean;
-  readonly reasons: readonly string[];
-}

package/src/contracts/generation.ts

@@ -1,103 +0,0 @@
-/** Portable execution contracts. No transport, model, tensor, or runtime imports. */
-export type CancelReason = "requested" | "consumer_closed" | "consumer_idle" | "engine_closed";
-
-export interface Cancellation {
-  readonly reason: CancelReason | undefined;
-  /** Invoked immediately when already cancelled. Returns an idempotent unsubscribe. */
-  subscribe(listener: (reason: CancelReason) => void): () => void;
-}
-
-export interface TokenLogprobs {
-  /** Selected-token logprob, before sampler truncation. */
-  logprob?: number;
-  /** Top-k pairs, sorted by descending logprob. */
-  top?: { id: number; logprob: number }[];
-}
-
-export type GenerationEvent =
-  | { readonly type: "committed"; readonly sequence: number;
-      readonly tokenIds: readonly number[]; readonly logprobs?: readonly (TokenLogprobs | undefined)[] }
-  | { readonly type: "progress"; readonly completed: number; readonly total?: number };
-
-/** Methods publish in order and await each publication before publishing again.
- * Only committed tokens cross this boundary. Tentative canvases/drafts stay private. */
-export interface GenerationOutput {
-  /** False requests a normal stop. The method must stop publishing and settle
-   * its metrics after releasing its private execution state. */
-  commit(tokenIds: readonly number[], logprobs?: readonly (TokenLogprobs | undefined)[]): Promise<void | false>;
-  progress(completed: number, total?: number): Promise<void>;
-}
-
-export interface MethodResult<Metrics> {
-  readonly finishReason: "stop" | "length";
-  readonly metrics: Metrics;
-}
-
-/** One attempt's private state, including any pending device work.
- * close must await safe resource release, including on execute failure. */
-export interface MethodRun<Metrics> {
-  execute(output: GenerationOutput, cancellation: Cancellation): Promise<MethodResult<Metrics>>;
-  close(): Promise<void>;
-}
-
-export interface InferenceMethod<Metrics> {
-  readonly id: string;
-  /** A rejected creation must release its own partial allocations. */
-  createRun(cancellation: Cancellation): Promise<MethodRun<Metrics>>;
-}
-
-export interface ExecutionPlan<Metrics> {
-  readonly id: string;
-  readonly outputTokenLimit: number;
-  readonly method: InferenceMethod<Metrics>;
-}
-
-export interface ExecutionPlanner<Request, Metrics> {
-  /** Planning selects a bound implementation; it does not allocate device state. */
-  plan(request: Request, cancellation: Cancellation): Promise<ExecutionPlan<Metrics>>;
-}
-
-export type GenerationOutcome<Metrics> =
-  | { readonly status: "completed"; readonly committedTokens: number;
-      readonly output?: Uint32Array; readonly result: MethodResult<Metrics> }
-  | { readonly status: "cancelled"; readonly committedTokens: number;
-      readonly output?: Uint32Array; readonly reason: CancelReason }
-  | { readonly status: "failed"; readonly committedTokens: number;
-      readonly output?: Uint32Array; readonly error: { readonly code: string; readonly message: string;
-        readonly cleanupError?: string } };
-
-export type SessionState = "created" | "preparing" | "running" | "settling" | "terminal";
-
-export type RunControl = {
-  /** collect returns token IDs; per-token logprobs require stream consumption. */
-  readonly output: "stream" | "collect";
-  readonly cancellation?: Cancellation;
-} | {
-  /** Direct delivery for consumers that already own bounded output storage.
-   * The callback is awaited before execution continues; false is a normal
-   * token stop. Borrowed arrays are valid until the callback settles. */
-  readonly output: "callback";
-  readonly onTokens: GenerationOutput["commit"];
-  readonly cancellation?: Cancellation;
-};
-
-export interface GenerationSession<Metrics> {
-  readonly id: string;
-  readonly state: SessionState;
-  /** Single consumer. Stream mode starts on first next(); collect mode has no events. */
-  readonly events: AsyncIterable<GenerationEvent>;
-  /** Always resolves with one terminal outcome, after native resources are released. */
-  readonly result: Promise<GenerationOutcome<Metrics>>;
-  cancel(reason?: CancelReason): Promise<void>;
-  close(): Promise<void>;
-}
-
-export interface InferenceEngine<Request, Metrics> {
-  open(request: Request, control: RunControl): Promise<GenerationSession<Metrics>>;
-  close(): Promise<void>;
-}
-
-/** Injected host timer. The portable engine has no DOM/Node/Bun dependency. */
-export interface Timer {
-  after(milliseconds: number, callback: () => void): () => void;
-}

package/src/contracts/host.ts

@@ -1,7 +0,0 @@
-/** Transport-specific request/response values stay outside portable policy. */
-export interface EngineHost<Request, Response> {
-  readonly ready: Promise<void>;
-  forward(request: Request): Promise<Response>;
-  /** Stop accepting work and wait until the owned worker has exited. */
-  close(): Promise<void>;
-}

package/src/contracts/pi-provider.ts

@@ -1,30 +0,0 @@
-export interface PiModelOptions {
-  name?: string;
-  contextWindow?: number;
-  maxTokens?: number;
-  reasoning?: boolean;
-  vision?: boolean;
-}
-
-/** Self-contained so the standalone Pi extension can embed the same builder. */
-export function piModelDefinition(options: PiModelOptions = {}) {
-  return {
-    id: "local",
-    name: options.name ?? "mlx-bun (local)",
-    api: "openai-completions" as const,
-    reasoning: options.reasoning ?? false,
-    compat: { supportsDeveloperRole: false,
-      ...(options.reasoning ? { thinkingFormat: "qwen-chat-template" as const } : {}) },
-    input: (options.vision ? ["text", "image"] : ["text"]) as ("text" | "image")[],
-    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-    contextWindow: options.contextWindow ?? 32768,
-    maxTokens: options.maxTokens ?? 8192,
-  };
-}
-
-export const PI_PROVIDER_ID = "mlx-bun";
-export const PI_API_KEY = "sk-mlx-bun-local";
-export const PI_LOCAL_MODEL_ID = piModelDefinition().id;
-export const PI_API = piModelDefinition().api;
-export const DEFAULT_CONTEXT_WINDOW = piModelDefinition().contextWindow;
-export const DEFAULT_MAX_TOKENS = piModelDefinition().maxTokens;