@mlx-bun/quantize 0.0.0-stage → 1.0.0

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Josh Rossi
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/README.md

@@ -1,3 +1,57 @@
-# Temporary Holding Version
+# @mlx-bun/quantize
+
+Produces the quantized checkpoints that [`@mlx-bun/inference`](../inference/README.md)
+loads and builds its graphs for. Inference never quantizes; this package never serves.
+
+| Module | Owns |
+| --- | --- |
+| `quantizer` | Quantize a checkpoint directory: per-module bit selection, dequantize/requantize through MLX, sharded output, config rewrite |
+| `convert` | Non-quantizing rewrite of a checkpoint: cast floating tensors to one dtype and/or dequantize quantized modules to dense weights (`mlx_lm.convert` without `-q`) |
+| `allocator` | Mixed-precision allocation from per-layer sensitivity under a bits-per-weight budget |
+| `sensitivity` | Exact per-layer KL sensitivity on calibration text |
+| `calibration` | Calibration sample loading and tokenization |
+| `drafter` | Quantize a released DeepSpec (`Gemma4DSparkModel`) drafter: every matmul weight and gather table, `confidence_head` kept bf16, then a load smoke through the drafter loader |
+| `rotate`, `weight-transform` | Rotation and fold plans for Llama and Qwen families before quantization |
+| `trellis` | Trellis state packing, interleaving, and the host decoder. The 1MAD codebook and word geometry are imported from the inference kernels, which own the packed format |
+| `trellis-quantizer` | `quantizeTrellisModelDir`: fold an HF-layout Qwen3.5-family checkpoint (γ + R1), code every MLP tensor with the Trellis (QTIP TCQ, L=12, 1MAD, T=256, tail-biting) and write the packed format the inference kernels serve; the rest keeps the shipped-compact affine tiers. `turboquant_fold.json` records the R1 seed and the final-norm gain the fold moved into the head, the residual basis a drafter reads through the graph |
+| `trellis-encoder` | The Viterbi encode of one folded tensor into packed codes and fp16 row scales, unweighted or with BlockLDLQ error feedback against a Hessian factor |
+| `trellis-allocation`, `qwen35-hf-fold` | Pure planning: which tensor gets which treatment (k-map, axis, layer limit) and the fold plan for the `model.language_model.*` layout with stored γ−1 norms |
+| `config-writer`, `atomic-output` | Quantization metadata and atomic directory publication |
+
+Sharded safetensors writing lives in `@mlx-bun/inference/artifacts` and is
+re-exported here for convenience. Auxiliary checkpoint file copying lives in
+`@mlx-bun/inference/artifacts/auxiliary-files`; `config-writer` re-exports
+`copyAuxFiles` for compatibility.
+
+Executed usage lives in the package tests: [weight-transform-plan](tests/weight-transform-plan.test.ts)
+plans a fold, [weight-transform-numerics](tests/weight-transform-numerics.test.ts) applies one, and
+[trellis-roundtrip](tests/trellis-roundtrip.test.ts) packs states here and expands them with the
+inference kernels. `quantizeModelDir(sourceDirectory, outputDirectory, { bits, groupSize })`
+is the entry point for a whole checkpoint; it needs a real one.
+
+`quantizeDrafterDir(sourceDirectory, outputDirectory, { bits, groupSize })` applies the drafter policy
+(the confidence head's sigmoid is compared against thresholds, so it stays bf16); drafter numerics move
+acceptance, never correctness, so gate the result with `scripts/drafter-ab.ts`. A [synthetic DeepSpec
+checkpoint test](tests/drafter.test.ts) covers the policy and refusals; a released drafter is not run here.
+
+`quantizeTrellisModelDir(sourceDirectory, outputDirectory, options)` produces a packed
+Trellis artifact from a full-precision `model.language_model.*` checkpoint (Qwen3.8-27B
+layout; a quantized source is refused). `bits` sets k for every coded tensor; `kMap`
+(an allocation file's `budgets[budget]`) sets it per tensor and may raise affine tiers;
+`downAxis: "in"` codes down_proj along the un-rotated dim; `layers` codes only the first
+N layers; `interleave` reorders eligible k=3 axis-0 codes for the scatter kernel; `reuse`
+copies tensors of unchanged geometry from earlier packed artifacts, so a new allocation
+encodes only the tensor/k pairs no source holds; `ldlq` names a directory of per-layer
+Hessian factors (`layer-NNN-mlp.safetensors` for gate/up, `layer-NNN-down.safetensors`,
+each holding one tensor `L`, the block-LDL factor of that site's input Hessian in the
+folded basis) and switches the objective to BlockLDLQ with a guard that falls back to the
+unweighted codes for any tensor whose feedback grows past 4x the unweighted magnitude.
+The Hessian factors and the k-map are inputs: collecting activation Hessians runs a
+calibration forward pass and is not part of this package. The encoder is numerically
+the frozen producer's (`02d723a:scripts/turboquant/tq-quantize-trellis-packed.ts`);
+[trellis-quantizer](tests/trellis-quantizer.test.ts) runs the whole path on a synthetic
+checkpoint, [trellis-encoder](tests/trellis-encoder.test.ts) checks the codec and the
+LDLQ arm, and [trellis-allocation](tests/trellis-allocation.test.ts) the planning.
+
+Job orchestration and the CLI verbs belong to the app's modules: `mlx-bun convert --q-mode trellis` to [`@mlx-bun/module-quantize`](../module-quantize/README.md), `mlx-bun draft quantize` to [`@mlx-bun/module-train`](../module-train/README.md).
 
-This version is a temporary placeholder for this package. An operational version to replace this has been submitted for review and is awaiting a staged release.

package/package.json

@@ -1,6 +1,46 @@
 {
   "name": "@mlx-bun/quantize",
-  "version": "0.0.0-stage",
-  "stub": true,
-  "description": "Temporary package placeholder for staged publishing"
-}
+  "version": "1.0.0",
+  "description": "Quantization for MLX checkpoints on Apple Silicon: calibration, sensitivity, mixed-precision allocation, rotation, Trellis packing, and safetensors output",
+  "type": "module",
+  "license": "MIT",
+  "engines": {
+    "bun": ">=1.4.2"
+  },
+  "exports": {
+    ".": "./src/index.ts",
+    "./quantizer": "./src/quantizer.ts",
+    "./convert": "./src/convert.ts",
+    "./drafter": "./src/drafter.ts",
+    "./allocator": "./src/allocator.ts",
+    "./sensitivity": "./src/sensitivity.ts",
+    "./calibration": "./src/calibration.ts",
+    "./rotate": "./src/rotate.ts",
+    "./weight-transform": "./src/weight-transform.ts",
+    "./trellis": "./src/trellis.ts",
+    "./trellis-quantizer": "./src/trellis-quantizer.ts",
+    "./trellis-allocation": "./src/trellis-allocation.ts",
+    "./trellis-encoder": "./src/trellis-encoder.ts",
+    "./qwen35-hf-fold": "./src/qwen35-hf-fold.ts",
+    "./config-writer": "./src/config-writer.ts",
+    "./atomic-output": "./src/atomic-output.ts"
+  },
+  "files": [
+    "src",
+    "LICENSE"
+  ],
+  "scripts": {
+    "typecheck": "bun ../../scripts/typecheck.ts --noEmit",
+    "test": "bun test"
+  },
+  "dependencies": {
+    "@mlx-bun/inference": "1.0.0",
+    "@mlx-bun/mlx": "1.0.0"
+  },
+  "os": [
+    "darwin"
+  ],
+  "cpu": [
+    "arm64"
+  ]
+}