@mlx-bun/quantize 0.0.0-stage → 1.0.0

package/src/calibration.ts

@@ -0,0 +1,182 @@
+// Calibration data loading for sensitivity analysis.
+//
+// Port of optiq/calibration/datasets.py `load_llm_calibration` /
+// `_load_optiq_calibration`: load a multi-domain JSONL mix, tokenize, build one
+// long token stream, and chunk it into fixed-length sequences. The
+// `calibrationFn` contract is the TS analogue of OptIQ's
+// `list[(args_tuple, kwargs_dict)]`: here, just a list of `number[]` token-id
+// prompts (each a (1, seq_len) row the model forwards).
+//
+// Two deviations from the Python, both forced by our pure-JS tokenizer surface
+// (packages/inference/src/input/tokenizer.ts exposes plain `encode` only, no
+// `apply_chat_template`):
+//   - `messages` samples use the role-prefixed dumb-concatenation fallback —
+//     exactly OptIQ's `no chat template` branch.
+//   - `add_special_tokens=False` is the OptIQ encode flag; our `encode` already
+//     omits special tokens for raw substrings.
+//
+// Source mix: the bundled `optiq.jsonl` if locatable on disk (the oracle venv's
+// optiq/calibration/data/optiq.jsonl), else a small built-in fallback corpus so
+// the path works with no external data.
+
+import { existsSync } from "node:fs";
+import type { LoadedTokenizer } from "@mlx-bun/inference/input";
+import { runtimeValue } from "@mlx-bun/inference/runtime/config";
+
+/** One JSONL calibration sample (subset of the OptIQ schema we can consume). */
+interface CalibrationSample {
+  domain?: string;
+  text?: string;
+  messages?: Array<{ role: string; content: string }>;
+  tools?: unknown;
+}
+
+/** Candidate on-disk locations for the bundled optiq.jsonl mix. */
+function bundledMixCandidates(): string[] {
+  const home = process.env.HOME ?? "";
+  const rel = "lib/python3.14/site-packages/optiq/calibration/data/optiq.jsonl";
+  return [
+    runtimeValue("MLX_BUN_CALIBRATION_JSONL"),
+    `${home}/Code/mlx-lm-example/.venv/${rel}`,
+    `${home}/Code/mlx-lm/.venv/${rel}`,
+  ].filter((p): p is string => !!p);
+}
+
+/** A tiny built-in fallback corpus (used only when no JSONL mix is found).
+ *  Multi-domain prose so groups have meaningful activation diversity. */
+const FALLBACK_CORPUS: string[] = [
+  "The history of computing stretches back to mechanical calculators and the abacus. Charles Babbage designed the Analytical Engine, a general-purpose mechanical computer, in the nineteenth century. Ada Lovelace wrote what is often considered the first algorithm intended for such a machine. The theoretical foundations were later laid by Alan Turing, whose model of computation defined the limits of what machines can compute.",
+  "In physics, the conservation of energy states that the total energy of an isolated system remains constant over time. Energy can neither be created nor destroyed; rather, it transforms from one form to another. Kinetic energy becomes potential energy as a ball rises, and back again as it falls. Thermodynamics extends these ideas to heat, work, and entropy in large systems.",
+  "Functions in programming are reusable blocks of code that perform a specific task. They take inputs, called arguments, and may return a value. Pure functions have no side effects and always return the same output for the same input, which makes them easier to test and reason about. Recursion is a technique where a function calls itself to solve smaller subproblems.",
+  "The water cycle describes the continuous movement of water on, above, and below the surface of the Earth. Water evaporates from oceans and lakes, condenses into clouds, and precipitates as rain or snow. Rivers carry water back to the sea, and groundwater slowly percolates through soil and rock. Plants release water vapor through transpiration, contributing to atmospheric moisture.",
+  "Economic markets coordinate the production and exchange of goods through prices. When demand for a product rises and supply stays fixed, prices tend to increase, signaling producers to make more. Competition drives firms to lower costs and innovate. Externalities, such as pollution, are costs not reflected in market prices and often require regulation to address.",
+  "A neural network is composed of layers of interconnected nodes, each applying a weighted sum followed by a nonlinear activation. During training, the network adjusts its weights to minimize a loss function via gradient descent. Backpropagation efficiently computes gradients by applying the chain rule layer by layer. Deeper networks can represent more complex functions but are harder to train.",
+];
+
+/**
+ * Build a calibration callable from a tokenizer.
+ *
+ * @returns a zero-arg function producing `number[][]` — N token-id rows, each
+ *          of length `seqLen`, ready to forward through a model.
+ */
+export function loadLlmCalibration(
+  tokenizer: LoadedTokenizer,
+  opts: {
+    nSamples?: number;
+    seqLen?: number;
+    seed?: number;
+    mix?: string;
+  } = {},
+): () => number[][] {
+  const nSamples = opts.nSamples ?? 32;
+  const seqLen = opts.seqLen ?? 1024;
+  const seed = opts.seed ?? 42;
+  const mix = opts.mix ?? "optiq";
+
+  const samples = loadSamples(mix);
+
+  // Concatenate every sample's tokens into one long stream.
+  const allTokens: number[] = [];
+  for (const s of samples) {
+    let text: string | undefined;
+    if (s.messages) {
+      // No chat template available in our JS tokenizer — role-prefixed
+      // concatenation (OptIQ's no-chat-template fallback branch).
+      text = s.messages.map((m) => `[${m.role}]: ${m.content}`).join("\n\n");
+    } else if (s.text !== undefined) {
+      text = s.text;
+    } else {
+      continue;
+    }
+    const toks = tokenizer.encode(text);
+    for (const t of toks) allTokens.push(t);
+  }
+
+  if (allTokens.length < seqLen) {
+    throw new Error(
+      `Calibration mix produced only ${allTokens.length} tokens (< seqLen=${seqLen}). ` +
+        `Try a smaller seqLen or a larger mix.`,
+    );
+  }
+
+  // Chunk start offsets, shuffled deterministically (seeded), like OptIQ.
+  const starts: number[] = [];
+  for (let s = 0; s < allTokens.length - seqLen; s += seqLen) starts.push(s);
+  shuffleInPlace(starts, seed);
+
+  const chunks: number[][] = [];
+  for (const start of starts.slice(0, nSamples)) {
+    chunks.push(allTokens.slice(start, start + seqLen));
+  }
+
+  if (chunks.length === 0) {
+    throw new Error("Could not create any calibration sequences.");
+  }
+
+  return () => chunks;
+}
+
+/** Load raw samples from the mix spec ("optiq" → bundled, else a JSONL path,
+ *  else the built-in fallback corpus). */
+function loadSamples(mix: string): CalibrationSample[] {
+  let jsonlPath: string | null = null;
+  if (mix === "optiq") {
+    jsonlPath = bundledMixCandidates().find((p) => existsSync(p)) ?? null;
+  } else if (existsSync(mix)) {
+    jsonlPath = mix;
+  } else {
+    throw new Error(
+      `calibration mix not found at ${mix}. Use mix='optiq' (default) or a ` +
+        `path to a JSONL file.`,
+    );
+  }
+
+  if (jsonlPath) return parseJsonl(jsonlPath);
+
+  // Built-in fallback: synthesize one prose sample per corpus entry.
+  return FALLBACK_CORPUS.map((text) => ({ domain: "prose", text }));
+}
+
+/** Parse a JSONL file into CalibrationSample[] (one object per non-blank line). */
+function parseJsonl(path: string): CalibrationSample[] {
+  // Synchronous read so the caller can stay synchronous; the file is small.
+  const text = readFileSyncUtf8(path);
+  const out: CalibrationSample[] = [];
+  for (const line of text.split("\n")) {
+    const trimmed = line.trim();
+    if (trimmed) out.push(JSON.parse(trimmed) as CalibrationSample);
+  }
+  return out;
+}
+
+function readFileSyncUtf8(path: string): string {
+  // Bun's readFileSync via node:fs is available in this runtime.
+  // (Kept local so the rest of the module is import-light.)
+  const { readFileSync } = require("node:fs") as typeof import("node:fs");
+  return readFileSync(path, "utf8");
+}
+
+/** Deterministic Fisher–Yates shuffle (mulberry32 PRNG seeded by `seed`).
+ *  NOTE: this does NOT reproduce numpy's RandomState shuffle order — it only
+ *  needs to be deterministic for our own reproducibility; the exact sample
+ *  ordering does not affect the allocator's 1:1 parity (that test feeds
+ *  synthetic SensitivityResults directly). */
+function shuffleInPlace<T>(arr: T[], seed: number): void {
+  const rand = mulberry32(seed >>> 0);
+  for (let i = arr.length - 1; i > 0; i--) {
+    const j = Math.floor(rand() * (i + 1));
+    const tmp = arr[i]!;
+    arr[i] = arr[j]!;
+    arr[j] = tmp;
+  }
+}
+
+function mulberry32(a: number): () => number {
+  return () => {
+    a |= 0;
+    a = (a + 0x6d2b79f5) | 0;
+    let t = Math.imul(a ^ (a >>> 15), 1 | a);
+    t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
+    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
+  };
+}

package/src/config-writer.ts

@@ -0,0 +1,109 @@
+// Config + sidecar writer for a quantized model directory.
+//
+// Produces the `config.json` block the loader (src/config.ts → parseQuantization)
+// reads back: a top-level `quantization` object AND a mirror `quantization_config`
+// (HF compatibility — config.ts reads `raw.quantization ?? raw.quantization_config`).
+// The block carries default {group_size, bits, mode} plus per-module overrides
+// keyed by module path (e.g. "model.embed_tokens") — exactly the OptiQ layout.
+//
+// Aux files (tokenizer, chat template, generation config, …) are copied through
+// verbatim so the output directory is a complete, loadable model snapshot.
+
+import { join } from "node:path";
+import { copyAuxFiles } from "@mlx-bun/inference/artifacts/auxiliary-files";
+
+/** Compatibility re-export; `@mlx-bun/inference/artifacts/auxiliary-files` owns it. */
+export { copyAuxFiles };
+
+/** A per-module quantization override. `false` means "left unquantized". */
+export type PerLayerEntry =
+  | { bits: number; groupSize: number; mode?: string; trellis?: { L: number; code: string; axis: 0 | 1 } }
+  | false;
+
+/** The default quantization definition for the model. */
+export interface QuantDef {
+  bits: number;
+  groupSize: number;
+  mode?: string;
+}
+
+/** The serialized quantization block, ready to drop into config.json. */
+export interface QuantizationBlock {
+  group_size: number;
+  bits: number;
+  mode: string;
+  // plus per-module entries: { [modulePath]: {bits, group_size} | false }
+  [modulePath: string]: number | string | { bits: number; group_size: number; mode?: string; trellis?: { L: number; code: string; axis: 0 | 1 } } | false;
+}
+
+/**
+ * Build the `quantization` config object: a `{group_size, bits, mode}` default
+ * plus one entry per module in `perLayer`. A module mapped to `false` is
+ * serialized as `false` (mlx's "not quantized" convention); otherwise it
+ * becomes `{bits, group_size}`.
+ */
+export function buildQuantizationBlock(
+  def: QuantDef,
+  perLayer: Map<string, PerLayerEntry>,
+): QuantizationBlock {
+  const block: QuantizationBlock = {
+    group_size: def.groupSize,
+    bits: def.bits,
+    mode: def.mode ?? "affine",
+  };
+  for (const [path, entry] of perLayer) {
+    if (entry === false) { block[path] = false; continue; }
+    const v: { bits: number; group_size: number; mode?: string; trellis?: { L: number; code: string; axis: 0 | 1 } } =
+      { bits: entry.bits, group_size: entry.groupSize };
+    if (entry.mode) v.mode = entry.mode;
+    if (entry.trellis) v.trellis = entry.trellis;
+    block[path] = v;
+  }
+  return block;
+}
+
+/** Metadata for the OptiQ-style sidecar describing the achieved quantization. */
+export interface OptiqMetadata {
+  method: string;
+  base_model: string;
+  bits: number;
+  group_size: number;
+  achieved_bpw: number;
+  per_layer_count: number;
+  weight_transforms?: Array<{
+    id: string;
+    seed: number;
+    family: string;
+    deviations: string[];
+  }>;
+}
+
+/**
+ * Write `config.json` (deep copy of `srcConfigRaw` with both `quantization`
+ * and `quantization_config` set to `block`) into `outDir`, copy through any
+ * present aux files, and write `optiq_metadata.json` describing the run.
+ *
+ * `*.model` tokenizer files beyond the known names are also swept in so SPM /
+ * tiktoken sidecars survive the round-trip.
+ */
+export async function writeQuantizedConfig(
+  srcConfigRaw: Record<string, unknown>,
+  outDir: string,
+  block: QuantizationBlock,
+  opts: { srcDir?: string; optiq?: OptiqMetadata } = {},
+): Promise<void> {
+  // Deep copy so we never mutate the caller's config object.
+  const config = JSON.parse(JSON.stringify(srcConfigRaw)) as Record<string, unknown>;
+  config.quantization = block;
+  config.quantization_config = block;
+  await Bun.write(join(outDir, "config.json"), JSON.stringify(config, null, 2));
+
+  if (opts.srcDir) await copyAuxFiles(opts.srcDir, outDir);
+
+  if (opts.optiq) {
+    await Bun.write(
+      join(outDir, "optiq_metadata.json"),
+      JSON.stringify(opts.optiq, null, 2),
+    );
+  }
+}

package/src/convert.ts

@@ -0,0 +1,79 @@
+// Non-quantizing conversion (mlx_lm.convert without -q): rewrite a checkpoint
+// with its floating tensors cast to one dtype and/or its quantized modules
+// dequantized to dense weights. Tensor names, shard layout rules, and aux files
+// follow the quantizer's writer; the CPU stream discipline is the same.
+
+import { mkdirSync } from "node:fs";
+import { join } from "node:path";
+import { clearCache } from "@mlx-bun/mlx/ffi";
+import { cpuStream } from "@mlx-bun/mlx/array";
+import { dequantize } from "@mlx-bun/mlx/ops";
+import { Weights, loadModelConfig, quantFor, writeShardedSafetensors, type NamedTensor, type WriteResult } from "@mlx-bun/inference/artifacts";
+import { copyAuxFiles } from "@mlx-bun/inference/artifacts/auxiliary-files";
+import { writeAtomicDirectory } from "./atomic-output";
+import { castTensor, type ConvertDtype, type ProgressEvent } from "./quantizer";
+
+export interface ConvertOptions {
+  /** Cast every floating tensor (quantization scales/biases included) to this dtype. */
+  dtype?: ConvertDtype;
+  /** Replace each quantized module with its dense weight and drop the config's quantization block. */
+  dequantize?: boolean;
+}
+
+export interface ConvertResult {
+  outDir: string;
+  /** Modules that were dequantized to dense weights. */
+  nDequantized: number;
+  write: WriteResult;
+}
+
+/** Rewrite `srcDir` into `outDir` (published atomically). Without options the
+ * tensors are copied unchanged, as `mlx_lm.convert` without -q/--dtype does. */
+export async function convertModelDir(
+  srcDir: string,
+  outDir: string,
+  opts: ConvertOptions = {},
+  onProgress?: (e: ProgressEvent) => void,
+): Promise<ConvertResult> {
+  return writeAtomicDirectory(outDir, async (staging) => {
+    const progress = (stage: string, message: string, fraction: number) => onProgress?.({ stage, message, progress: fraction });
+    progress("loading", `Reading ${srcDir}`, 0);
+    const config = await loadModelConfig(srcDir);
+    const weights = await Weights.open(srcDir);
+    const out: NamedTensor[] = [];
+    try {
+      const names = weights.tensorNames;
+      const present = new Set(names), scaled = new Set<string>();
+      for (const name of names) if (name.endsWith(".scales")) scaled.add(name.slice(0, -".scales".length));
+      // Modules to dequantize: a `.scales` sibling and a quantization spec for the module.
+      const dense = new Set([...scaled].filter(base => opts.dequantize && quantFor(config.quantization, base)));
+      let processed = 0;
+      for (const name of names) {
+        const base = /^(.*)\.(weight|scales|biases)$/.exec(name)?.[1];
+        if (base && dense.has(base)) {
+          if (!name.endsWith(".weight")) continue; // scales/biases are folded into the dense weight
+          const spec = quantFor(config.quantization, base)!;
+          const packed = weights.tensor(name), scales = weights.tensor(`${base}.scales`);
+          const biases = present.has(`${base}.biases`) ? weights.tensor(`${base}.biases`) : null;
+          const weight = dequantize(packed, scales, biases, spec, cpuStream);
+          out.push({ name, array: castTensor(name, weight, opts.dtype) });
+        } else out.push({ name, array: castTensor(name, weights.tensor(name), opts.dtype) });
+        if (++processed % 16 === 0) clearCache();
+        progress("converting", `Tensor ${processed}/${names.length}`, processed / names.length);
+      }
+      progress("writing", `Writing ${out.length} tensors to ${outDir}`, 1);
+      mkdirSync(staging, { recursive: true });
+      const write = writeShardedSafetensors(staging, out);
+      const raw = JSON.parse(JSON.stringify(config.raw)) as Record<string, unknown>;
+      if (opts.dequantize) { delete raw.quantization; delete raw.quantization_config; }
+      await Bun.write(join(staging, "config.json"), JSON.stringify(raw, null, 2));
+      await copyAuxFiles(srcDir, staging);
+      progress("done", `Converted ${out.length} tensors → ${outDir}`, 1);
+      return { outDir, nDequantized: dense.size, write };
+    } finally {
+      for (const tensor of out) { try { tensor.array.dispose(); } catch { /* already released */ } }
+      weights.dispose();
+      clearCache();
+    }
+  });
+}

package/src/drafter.ts

@@ -0,0 +1,94 @@
+import { rowOperationFor } from "@mlx-bun/inference/models/graph";
+// Quantize a DeepSpec-format drafter checkpoint (DeepSeek's released
+// `Gemma4DSparkModel`; a locally trained dspark.json drafter is small enough to
+// serve as trained).
+//
+// Thin policy layer over this package's quantizer: the generic
+// shape-eligibility rule already captures every 2-D matmul weight — the
+// layer projections, fc, lm_head, markov_w2, and the two gather tables
+// (embed_tokens, markov_w1: quantized-embedding gather-dequant at load,
+// mlx QuantizedEmbedding semantics). The drafter-specific overlay:
+//
+//   KEEP bf16: confidence_head.* — threshold-sensitive (the truncation
+//   comparison runs on its sigmoid; borderline flips change proposal
+//   lengths) and tiny (~4k params), so quantizing it buys nothing.
+//   Norms / layer_scalar / biases are 1-D and never shape-eligible.
+//
+// Drafter numerics only move ACCEPTANCE, never correctness (the target
+// verifies every draft) — the quality gate for any scheme swept through
+// here is the drafter acceptance A/B (`scripts/drafter-ab.ts`), not a KL battery.
+
+import { quantizeModelDir, type QuantizeResult } from "./quantizer";
+import { isDeepspecDrafterConfig } from "@mlx-bun/inference/models/drafters";
+import { loadDeepspecDrafter } from "@mlx-bun/inference/models/drafter-loaders";
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+
+/** Module bases kept full-precision regardless of shape eligibility. */
+const KEEP_BF16 = [/^confidence_head\./];
+
+/** The drafter per-tensor policy: quantize everything shape-eligible
+ *  except the KEEP_BF16 set. */
+export function drafterQuantizePredicate(base: string): boolean {
+  return !KEEP_BF16.some((re) => re.test(base));
+}
+
+export interface QuantizeDrafterOptions {
+  bits: 4 | 8;
+  groupSize: 32 | 64;
+  /** Skip the post-write load smoke (tests exercise loading separately). */
+  skipLoadSmoke?: boolean;
+  onProgress?: (stage: string, message: string, progress: number) => void;
+}
+
+/**
+ * Quantize the DeepSpec drafter at `srcDir` into `outDir` and smoke the
+ * result: reload through the DeepSpec drafter loader (quantized detection + graph
+ * build) and eval one projectContext row through the quantized fc path.
+ */
+export async function quantizeDrafterDir(
+  srcDir: string,
+  outDir: string,
+  opts: QuantizeDrafterOptions,
+): Promise<QuantizeResult> {
+  const raw = (await Bun.file(`${srcDir}/config.json`).json()) as Record<string, any>;
+  if (!isDeepspecDrafterConfig(raw))
+    throw new Error(
+      `quantizeDrafterDir: expected a DeepSpec checkpoint (architectures[0]==="Gemma4DSparkModel"), got ${JSON.stringify(raw.architectures)} — for regular models use \`mlx-bun convert -q\``,
+    );
+  if (raw.quantization ?? raw.quantization_config)
+    throw new Error("quantizeDrafterDir: source is already quantized — start from the bf16 checkpoint");
+
+  const result = await quantizeModelDir(
+    srcDir,
+    outDir,
+    {
+      bits: opts.bits,
+      groupSize: opts.groupSize,
+      quantizePredicate: drafterQuantizePredicate,
+    },
+    opts.onProgress ? (e) => opts.onProgress!(e.stage, e.message, e.progress) : undefined,
+  );
+
+  if (!opts.skipLoadSmoke) {
+    // Load smoke: quantized detection + one real quantized matmul (fc) so a
+    // wrong spec/layout fails HERE, not at serve time.
+    const d = await loadDeepspecDrafter(outDir);
+    try {
+      const width = d.tapLayers.length * d.hidden;
+      const zeros = MlxArray.fromFloat32(new Float32Array(width), [1, 1, width]);
+      const x = zeros.astype(Dtype.bfloat16);
+      zeros.dispose();
+      const proj = rowOperationFor(d.projectContext,1)(x); // one context row
+      x.dispose();
+      const got = proj.shape;
+      proj.dispose();
+      if (got.length !== 3 || got[2] !== d.hidden)
+        throw new Error(`load smoke: projectContext shape ${JSON.stringify(got)}, expected [1,1,${d.hidden}]`);
+    } finally {
+      d.dispose();
+    }
+  }
+
+  return result;
+}

package/src/index.ts

@@ -0,0 +1,86 @@
+// Public checkpoint-quantization API: allocation, weight transforms, sharded
+// output, and quantization metadata. Model discovery belongs to @mlx-bun/hub.
+
+export { quantizeModelDir, isQuantizable, withPreparedProbe, CONVERT_DTYPES } from "./quantizer";
+export { convertModelDir } from "./convert";
+export { quantizeDrafterDir, drafterQuantizePredicate } from "./drafter";
+export type { QuantizeDrafterOptions } from "./drafter";
+export type { ConvertOptions, ConvertResult } from "./convert";
+export type {
+  ConvertDtype,
+  PreparedProbe,
+  ProbeSource,
+  QuantizeOptions,
+  QuantizeResult,
+  ProgressEvent,
+} from "./quantizer";
+
+export {
+  planQwen35Fold,
+  planQwenMtpFold,
+} from "./rotate";
+export type {
+  FoldOptions,
+  QwenFoldOptions,
+} from "./rotate";
+
+export {
+  automaticRotationWeightTransform,
+  llamaWeightTransform,
+  planLlamaWeightTransform,
+  qwen35WeightTransform,
+  qwenMtpWeightTransform,
+} from "./weight-transform";
+export type {
+  LlamaWeightTransformPlan,
+  QwenWeightTransformPlan,
+  WeightTransform,
+  WeightTransformContext,
+  WeightTransformMetadata,
+  WeightTransformPlan,
+} from "./weight-transform";
+
+export { quantizeTrellisModelDir } from "./trellis-quantizer";
+export type { TrellisQuantizeOptions, TrellisQuantizeResult } from "./trellis-quantizer";
+export { TrellisEncoder, loadHessianFactor } from "./trellis-encoder";
+export type { PackedTrellisTensor, LdlqEncoding, TrellisEncoderOptions } from "./trellis-encoder";
+export { parseKMap, treatmentFor, trellisBpw, mlxLmPath } from "./trellis-allocation";
+export type { DownAxis, Treatment, TreatmentPolicy, TrellisKMap } from "./trellis-allocation";
+export { planQwen35HfFold } from "./qwen35-hf-fold";
+export type { Qwen35HfFoldPlan, Qwen35HfFoldSource } from "./qwen35-hf-fold";
+
+export { writeShardedSafetensors, DEFAULT_SHARD_BYTES } from "@mlx-bun/inference/artifacts";
+export type { NamedTensor, ShardInfo, SafetensorsIndex, WriteResult, WriteOpts } from "@mlx-bun/inference/artifacts";
+
+export {
+  buildQuantizationBlock,
+  writeQuantizedConfig,
+} from "./config-writer";
+export type {
+  QuantDef,
+  PerLayerEntry,
+  QuantizationBlock,
+  OptiqMetadata,
+} from "./config-writer";
+
+// Mixed-precision (OptiQ sensitivity + knapsack) surface.
+export {
+  optimizeMixedPrecision,
+  computeBpw,
+  klReduction,
+  UpgradeHeap,
+} from "./allocator";
+export type {
+  LayerQuantConfig,
+  OptimizationResult,
+  OptimizeMixedPrecisionOptions,
+  UpgradeEntry,
+} from "./allocator";
+
+export { analyzeSensitivityExact, klFromRef } from "./sensitivity";
+export type {
+  SensitivityResult,
+  AnalyzeSensitivityOptions,
+} from "./sensitivity";
+
+export { loadLlmCalibration } from "./calibration";