Add first-class quantized inference for v0.3.0

CONTRIBUTING.md

@@ -21,7 +21,7 @@ HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 \
   python -m unittest discover -s tests -v
 ```
 
-The expected baseline is 98 discovered tests: 76 pass and 22 Hub tests skip.
+The expected baseline is 114 discovered tests: 92 pass and 22 Hub tests skip.
 The exact count should change only when tests are intentionally added or
 removed.
 

README.md

@@ -56,18 +56,74 @@ inputs = {k: mx.array(v) for k, v in inputs.items()}
 outputs = model(**inputs)
 ```
 
-Quantized loading:
+## Quantized Inference
+
+MLX Transformers can auto-detect an MLX pre-quantized checkpoint and run it
+without extra loader flags:
+
+```python
+from mlx_transformers import generate_text, load_causal_model
+
+loaded = load_causal_model(
+    "mlx-community/Phi-3-mini-4k-instruct-4bit",
+)
+result = generate_text(
+    loaded.model,
+    loaded.tokenizer,
+    "Explain weight quantization in one sentence.",
+    max_new_tokens=64,
+)
+print(loaded.quantization)
+print(result.text)
+```
+
+To quantize a regular safetensors checkpoint in memory after loading:
 
 ```python
-model.from_pretrained(
+import mlx.core as mx
+
+from mlx_transformers import QuantizationConfig, load_causal_model
+
+loaded = load_causal_model(
     model_name,
-    quantize=True,
-    group_size=64,
-    bits=4,
-    mode="affine",
+    dtype=mx.float16,
+    quantization=QuantizationConfig(
+        group_size=64,
+        bits=4,
+        mode="affine",
+    ),
 )
 ```
 
+The existing model-specific loader flags remain supported:
+
+```python
+model.from_pretrained(model_name, quantize=True, group_size=64, bits=4)
+```
+
+The installed CLI offers the same two paths:
+
+```bash
+# Auto-detect a pre-quantized MLX checkpoint.
+mlx-transformers-generate \
+  --model mlx-community/Phi-3-mini-4k-instruct-4bit \
+  --prompt "Explain attention masking." \
+  --max-new-tokens 64
+
+# Quantize a regular checkpoint after loading it.
+mlx-transformers-generate \
+  --model meta-llama/Llama-3.2-1B-Instruct \
+  --prompt "Explain attention masking." \
+  --quantize --group-size 64 --bits 4 \
+  --max-new-tokens 64
+```
+
+On-load quantization temporarily materializes the regular checkpoint before
+replacing supported layers, so peak memory is higher than loading an already
+quantized checkpoint. Prefer a reviewed pre-quantized MLX checkpoint for large
+models. See [docs/load_model.md](docs/load_model.md) for supported modes,
+offline use, metadata inspection, and safety constraints.
+
 Generation is finite and uses Hugging Face-style `max_new_tokens`:
 
 ```python
@@ -184,5 +240,5 @@ Tests that use external checkpoints are skipped unless
 download sizes before opting in. Some checkpoints are gated; set `HF_TOKEN`
 only for an explicitly reviewed integration run.
 
-The verified 2026-07-26 Apple-silicon baseline is 98 discovered tests:
-76 pass and 22 Hub integration tests skip.
+The verified 2026-07-26 Apple-silicon baseline is 114 discovered tests:
+92 pass and 22 Hub integration tests skip.

SUPPORT.md

@@ -75,9 +75,19 @@ Gemma 3 4B multimodal inference produced non-finite values in float16 during
 validation. Use float32 or `mx.bfloat16` for that path. Float16 is not a
 verified Gemma 3 4B multimodal configuration.
 
-Quantization structure and error paths have bounded tests. Real quantized
-generation is currently verified only for the tested 4-bit Phi-3 checkpoint
-with group size 64.
+Quantized inference is verified through two bounded paths:
+
+- a regular synthetic Llama checkpoint loaded and quantized at runtime with
+  affine 4-bit weights and group size 32;
+- saved/reloaded pre-quantized synthetic Llama weights plus the cached
+  `mlx-community/Phi-3-mini-4k-instruct-4bit` checkpoint with affine 4-bit
+  weights and group size 64.
+
+Both paths pass finite forward execution and bounded cached generation. The
+public CLI was run offline against the cached Phi-3 checkpoint on Apple
+silicon. Affine, MXFP4, MXFP8, and NVFP4 settings are validated against the
+MLX 0.31/0.32 contract, but real-checkpoint evidence is currently limited to
+affine 4-bit weights.
 
 ## Generation contract
 
@@ -102,7 +112,14 @@ Transformers stopping-criteria API are not implemented.
 `from_pretrained` supports local checkpoint directories and Hugging Face Hub
 model IDs. It validates safetensors shards and indexes, rejects missing or
 unexpected model parameters, and supports explicit offline, cache, revision,
-token, worker, and dtype controls.
+token, worker, dtype, and quantization controls. Pre-quantized MLX checkpoints
+are detected from checkpoint metadata and scale tensors. Regular checkpoints
+can be quantized after loading through `QuantizationConfig` or the compatible
+legacy flags.
+
+Runtime quantization has a higher peak-memory requirement because the regular
+checkpoint is materialized before its supported modules are replaced. It does
+not convert and publish a reusable checkpoint.
 
 PyTorch `.bin` checkpoints and remote model code are unsupported.
 `trust_remote_code=True` is ignored with a warning rather than executing code

docs/load_model.md

@@ -0,0 +1,119 @@
+# Loading and running quantized models
+
+MLX Transformers 0.3.0 supports two bounded quantized-inference workflows on
+Apple silicon:
+
+1. Load an MLX checkpoint that already contains quantized weights.
+2. Load a regular safetensors checkpoint and quantize supported modules in
+   memory before inference.
+
+## Pre-quantized checkpoints
+
+Use `load_causal_model` for a supported text-generation family. Quantization
+metadata and `.scales` tensors are detected automatically:
+
+```python
+from mlx_transformers import generate_text, load_causal_model
+
+loaded = load_causal_model(
+    "mlx-community/Phi-3-mini-4k-instruct-4bit",
+)
+
+print(loaded.quantization)
+result = generate_text(
+    loaded.model,
+    loaded.tokenizer,
+    "Reply with one word: ready",
+    max_new_tokens=8,
+    temperature=0.0,
+)
+print(result.text)
+```
+
+For a local or cached checkpoint with networking disabled:
+
+```python
+loaded = load_causal_model(
+    "/path/to/checkpoint",
+    local_files_only=True,
+)
+```
+
+The loader preserves the stored quantized tensor dtypes and exposes a
+`QuantizationInfo` object whose `source` is `"checkpoint"`.
+
+## Quantize after loading
+
+Pass a validated `QuantizationConfig` for a regular checkpoint:
+
+```python
+import mlx.core as mx
+
+from mlx_transformers import QuantizationConfig, load_causal_model
+
+loaded = load_causal_model(
+    "meta-llama/Llama-3.2-1B-Instruct",
+    dtype=mx.float16,
+    quantization=QuantizationConfig(
+        group_size=64,
+        bits=4,
+        mode="affine",
+    ),
+)
+```
+
+The resulting `QuantizationInfo.source` is `"runtime"`. This path temporarily
+holds the regular checkpoint in memory, so its peak memory can be materially
+higher than direct pre-quantized loading. It changes the in-memory model only;
+it does not write or publish a converted checkpoint.
+
+The legacy `from_pretrained(..., quantize=True, group_size=..., bits=...)`
+arguments remain supported for compatibility. Do not combine them with
+`quantization=QuantizationConfig(...)`.
+
+## Supported settings
+
+Settings are validated before checkpoint loading:
+
+| Mode | Group size | Bits | Quantized inputs |
+| --- | --- | --- | --- |
+| `affine` | 32, 64, or 128 | 2, 3, 4, 5, 6, or 8 | No |
+| `mxfp4` | 32 | 4 | No |
+| `mxfp8` | 32 | 8 | Optional |
+| `nvfp4` | 16 | 4 | Optional |
+
+The dimension quantized by MLX must be divisible by the chosen group size.
+Runtime support also depends on the installed MLX version and Apple hardware.
+The package bounds its supported MLX versions in `pyproject.toml`.
+
+## Command-line generation
+
+Pre-quantized checkpoint:
+
+```bash
+mlx-transformers-generate \
+  --model mlx-community/Phi-3-mini-4k-instruct-4bit \
+  --prompt "Reply with one word: ready" \
+  --max-new-tokens 8
+```
+
+Runtime quantization:
+
+```bash
+mlx-transformers-generate \
+  --model meta-llama/Llama-3.2-1B-Instruct \
+  --prompt "Reply with one word: ready" \
+  --quantize --mode affine --group-size 64 --bits 4 \
+  --max-new-tokens 8
+```
+
+The CLI requires both `--model` and `--prompt`, always uses a finite token
+bound, never enables remote model code, and prints the effective quantization
+metadata to standard error. Add `--local-files-only` to prohibit Hub network
+resolution. Use `--raw-prompt` when a tokenizer chat template should not be
+applied.
+
+Supported causal model types are Gemma 3 text, Llama, OpenELM, Persimmon, Phi,
+Phi-3, and Qwen3. OpenELM requires an explicit `--tokenizer` because its model
+checkpoint does not bundle one. OpenELM and Persimmon are maintenance-only;
+new verification work prioritizes active families listed in `SUPPORT.md`.

pyproject.toml

@@ -34,6 +34,9 @@ dependencies = [
   "transformers>=4.57.6,<6",
 ]
 
+[project.scripts]
+mlx-transformers-generate = "mlx_transformers.cli:main"
+
 [project.optional-dependencies]
 tokenizers = [
   "sentencepiece>=0.2.0",

src/mlx_transformers/__init__.py

@@ -1,8 +1,47 @@
+from importlib import import_module
 from importlib.metadata import PackageNotFoundError, version
 
 try:
     __version__ = version("mlx-transformers")
 except PackageNotFoundError:
     __version__ = "0+unknown"
 
-__all__ = ["__version__"]
+__all__ = [
+    "GenerationResult",
+    "LoadedCausalModel",
+    "QuantizationConfig",
+    "QuantizationInfo",
+    "__version__",
+    "generate_text",
+    "get_quantization_info",
+    "load_causal_model",
+    "resolve_causal_model_class",
+]
+
+_LAZY_IMPORTS = {
+    "GenerationResult": (".inference", "GenerationResult"),
+    "LoadedCausalModel": (".inference", "LoadedCausalModel"),
+    "QuantizationConfig": (".quantization", "QuantizationConfig"),
+    "QuantizationInfo": (".quantization", "QuantizationInfo"),
+    "generate_text": (".inference", "generate_text"),
+    "get_quantization_info": (".quantization", "get_quantization_info"),
+    "load_causal_model": (".inference", "load_causal_model"),
+    "resolve_causal_model_class": (".inference", "resolve_causal_model_class"),
+}
+
+
+def __getattr__(name):
+    try:
+        module_name, attribute_name = _LAZY_IMPORTS[name]
+    except KeyError as error:
+        raise AttributeError(
+            f"module {__name__!r} has no attribute {name!r}"
+        ) from error
+
+    value = getattr(import_module(module_name, __name__), attribute_name)
+    globals()[name] = value
+    return value
+
+
+def __dir__():
+    return sorted(set(globals()) | set(__all__))