VeloxQuant-MLX 0.2.0__py3-none-any.whl

mlx_kv_quant/__init__.py

@@ -0,0 +1,41 @@
+"""mlx_kv_quant — KV cache quantization for Apple Silicon MLX.
+
+Implements TurboQuant, PolarQuant, and QJL for production LLM inference.
+"""
+from __future__ import annotations
+
+from mlx_kv_quant.cache.base import KVCacheBuilder, KVCacheConfig, KVCacheFactory
+from mlx_kv_quant.core.abstractions import (
+    ArtifactStore,
+    KVCache,
+    Quantizer,
+    QuantizationObserver,
+)
+from mlx_kv_quant.core.context import EncodedVector, QuantizationContext, TransformResult
+from mlx_kv_quant.core.exceptions import (
+    ArtifactNotFoundError,
+    CodebookDimensionMismatch,
+    CyclicPipelineError,
+    QuantizerConfigError,
+)
+from mlx_kv_quant.quantizers.base import QuantizerFactory
+
+__all__ = [
+    "KVCacheBuilder",
+    "KVCacheConfig",
+    "KVCacheFactory",
+    "ArtifactStore",
+    "KVCache",
+    "Quantizer",
+    "QuantizationObserver",
+    "EncodedVector",
+    "QuantizationContext",
+    "TransformResult",
+    "ArtifactNotFoundError",
+    "CodebookDimensionMismatch",
+    "CyclicPipelineError",
+    "QuantizerConfigError",
+    "QuantizerFactory",
+]
+
+__version__ = "0.2.0"

mlx_kv_quant/__main__.py

@@ -0,0 +1,28 @@
+"""Entry point for `python -m mlx_kv_quant <command>`."""
+from __future__ import annotations
+
+import sys
+
+
+def main() -> None:
+    if len(sys.argv) < 2:
+        print("Usage: veloxquant {precompute|benchmark}")
+        sys.exit(1)
+
+    command = sys.argv[1]
+    # Remove the subcommand so sub-parsers see argv correctly
+    sys.argv = [f"mlx_kv_quant {command}"] + sys.argv[2:]
+
+    if command == "precompute":
+        from mlx_kv_quant.cli.precompute import main as _main
+        _main()
+    elif command == "benchmark":
+        from mlx_kv_quant.cli.benchmark import main as _main
+        _main()
+    else:
+        print(f"Unknown command: {command!r}. Choices: precompute, benchmark")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

mlx_kv_quant/artifacts/__init__.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from mlx_kv_quant.artifacts.memory_store import InMemoryArtifactStore
+from mlx_kv_quant.artifacts.npy_store import NpyArtifactStore
+
+__all__ = ["InMemoryArtifactStore", "NpyArtifactStore"]

mlx_kv_quant/artifacts/base.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from mlx_kv_quant.core.abstractions import ArtifactStore
+
+__all__ = ["ArtifactStore"]

mlx_kv_quant/artifacts/memory_store.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Tuple
+
+import numpy as np
+
+from mlx_kv_quant.core.abstractions import ArtifactStore
+from mlx_kv_quant.core.exceptions import ArtifactNotFoundError
+
+
+class InMemoryArtifactStore(ArtifactStore):
+    """In-memory artifact store for testing — performs no disk I/O.
+
+    All artifacts are stored in plain Python dicts keyed by descriptor tuples.
+    Arrays are stored as float16 numpy arrays and wrapped in MLX on load.
+
+    Args:
+        None.
+    """
+
+    def __init__(self) -> None:
+        self._rotations: Dict[Tuple, np.ndarray] = {}
+        self._codebooks: Dict[Tuple, np.ndarray] = {}
+        self._jls: Dict[Tuple, np.ndarray] = {}
+
+    # ------------------------------------------------------------------
+    # Rotation matrix
+    # ------------------------------------------------------------------
+
+    def load_rotation_matrix(self, d: int, seed: int) -> Any:
+        key = (d, seed)
+        if key not in self._rotations:
+            raise ArtifactNotFoundError(
+                f"InMemoryArtifactStore: rotation d={d} seed={seed} not found."
+            )
+        import mlx.core as mx
+        return mx.array(self._rotations[key])
+
+    def save_rotation_matrix(self, Pi: Any, d: int, seed: int) -> None:
+        self._rotations[(d, seed)] = np.array(Pi, dtype=np.float16)
+
+    # ------------------------------------------------------------------
+    # Codebook
+    # ------------------------------------------------------------------
+
+    def load_codebook(self, distribution: str, b: int, d: int) -> Any:
+        key = (distribution, b, d)
+        if key not in self._codebooks:
+            raise ArtifactNotFoundError(
+                f"InMemoryArtifactStore: codebook dist={distribution} b={b} d={d} not found."
+            )
+        import mlx.core as mx
+        return mx.array(self._codebooks[key])
+
+    def save_codebook(self, cb: Any, distribution: str, b: int, d: int) -> None:
+        self._codebooks[(distribution, b, d)] = np.array(cb, dtype=np.float16)
+
+    # ------------------------------------------------------------------
+    # JL matrix
+    # ------------------------------------------------------------------
+
+    def load_jl_matrix(self, d: int, m: int, seed: int) -> Any:
+        key = (d, m, seed)
+        if key not in self._jls:
+            raise ArtifactNotFoundError(
+                f"InMemoryArtifactStore: JL d={d} m={m} seed={seed} not found."
+            )
+        import mlx.core as mx
+        return mx.array(self._jls[key])
+
+    def save_jl_matrix(self, S: Any, d: int, m: int, seed: int) -> None:
+        self._jls[(d, m, seed)] = np.array(S, dtype=np.float16)
+
+    # ------------------------------------------------------------------
+    # Existence check
+    # ------------------------------------------------------------------
+
+    def exists(self, artifact_type: str, **kwargs: Any) -> bool:
+        if artifact_type == "rotation":
+            return (kwargs["d"], kwargs["seed"]) in self._rotations
+        if artifact_type == "codebook":
+            return (kwargs["distribution"], kwargs["b"], kwargs["d"]) in self._codebooks
+        if artifact_type == "jl":
+            return (kwargs["d"], kwargs["m"], kwargs["seed"]) in self._jls
+        return False
+
+    def __repr__(self) -> str:
+        return (
+            f"InMemoryArtifactStore("
+            f"rotations={len(self._rotations)}, "
+            f"codebooks={len(self._codebooks)}, "
+            f"jls={len(self._jls)})"
+        )

mlx_kv_quant/artifacts/npy_store.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from mlx_kv_quant.core.abstractions import ArtifactStore
+from mlx_kv_quant.core.exceptions import ArtifactNotFoundError
+
+
+class NpyArtifactStore(ArtifactStore):
+    """Artifact store that reads and writes ``.npy`` files from a local directory.
+
+    File naming conventions:
+        rotation_d{d}_seed{seed}.npy
+        codebook_{distribution}_b{b}_d{d}.npy
+        jl_d{d}_m{m}_seed{seed}.npy
+
+    Args:
+        root_dir: Path to the directory where artifacts are stored.
+            Created automatically on first save if absent.
+    """
+
+    def __init__(self, root_dir: str | Path) -> None:
+        self._root = Path(root_dir)
+        self._root.mkdir(parents=True, exist_ok=True)
+
+    # ------------------------------------------------------------------
+    # Rotation matrix
+    # ------------------------------------------------------------------
+
+    def _rotation_path(self, d: int, seed: int) -> Path:
+        return self._root / f"rotation_d{d}_seed{seed}.npy"
+
+    def load_rotation_matrix(self, d: int, seed: int) -> Any:
+        path = self._rotation_path(d, seed)
+        if not path.exists():
+            raise ArtifactNotFoundError(
+                f"Rotation matrix not found at {path}. "
+                f"Run `python -m mlx_kv_quant precompute --head_dim {d}` first."
+            )
+        import mlx.core as mx
+        return mx.array(np.load(path).astype(np.float16))
+
+    def save_rotation_matrix(self, Pi: Any, d: int, seed: int) -> None:
+        path = self._rotation_path(d, seed)
+        arr = np.array(Pi, dtype=np.float16)
+        np.save(path, arr)
+
+    # ------------------------------------------------------------------
+    # Codebook
+    # ------------------------------------------------------------------
+
+    def _codebook_path(self, distribution: str, b: int, d: int) -> Path:
+        return self._root / f"codebook_{distribution}_b{b}_d{d}.npy"
+
+    def load_codebook(self, distribution: str, b: int, d: int) -> Any:
+        path = self._codebook_path(distribution, b, d)
+        if not path.exists():
+            raise ArtifactNotFoundError(
+                f"Codebook not found at {path}. "
+                f"Run `python -m mlx_kv_quant precompute --head_dim {d} --bits {b}` first."
+            )
+        import mlx.core as mx
+        return mx.array(np.load(path).astype(np.float16))
+
+    def save_codebook(self, cb: Any, distribution: str, b: int, d: int) -> None:
+        path = self._codebook_path(distribution, b, d)
+        arr = np.array(cb, dtype=np.float16)
+        np.save(path, arr)
+
+    # ------------------------------------------------------------------
+    # JL matrix
+    # ------------------------------------------------------------------
+
+    def _jl_path(self, d: int, m: int, seed: int) -> Path:
+        return self._root / f"jl_d{d}_m{m}_seed{seed}.npy"
+
+    def load_jl_matrix(self, d: int, m: int, seed: int) -> Any:
+        path = self._jl_path(d, m, seed)
+        if not path.exists():
+            raise ArtifactNotFoundError(
+                f"JL matrix not found at {path}. "
+                f"Run `python -m mlx_kv_quant precompute --head_dim {d} --jl_dim {m}` first."
+            )
+        import mlx.core as mx
+        return mx.array(np.load(path).astype(np.float16))
+
+    def save_jl_matrix(self, S: Any, d: int, m: int, seed: int) -> None:
+        path = self._jl_path(d, m, seed)
+        arr = np.array(S, dtype=np.float16)
+        np.save(path, arr)
+
+    # ------------------------------------------------------------------
+    # Existence check
+    # ------------------------------------------------------------------
+
+    def exists(self, artifact_type: str, **kwargs: Any) -> bool:
+        if artifact_type == "rotation":
+            return self._rotation_path(kwargs["d"], kwargs["seed"]).exists()
+        if artifact_type == "codebook":
+            return self._codebook_path(
+                kwargs["distribution"], kwargs["b"], kwargs["d"]
+            ).exists()
+        if artifact_type == "jl":
+            return self._jl_path(
+                kwargs["d"], kwargs["m"], kwargs["seed"]
+            ).exists()
+        return False
+
+    def __repr__(self) -> str:
+        return f"NpyArtifactStore(root={self._root!r})"

mlx_kv_quant/benchmarks/attend_benchmark.py

@@ -0,0 +1,184 @@
+"""Attend latency and memory benchmark across sequence lengths.
+
+Compares four configurations:
+  baseline   — no optimizations
+  vectorized — enable_vectorized_attend=True
+  fused      — enable_vectorized_attend=True + enable_fused_query_dot=True
+  all        — fused + enable_outlier_two_stream=True
+
+Usage::
+
+    python -m mlx_kv_quant.benchmarks.attend_benchmark
+    python -m mlx_kv_quant.benchmarks.attend_benchmark --method turboquant_mse --bits 2
+    python -m mlx_kv_quant.benchmarks.attend_benchmark --seq_lens 64 256 1024 4096
+"""
+from __future__ import annotations
+
+import argparse
+import time
+from typing import List
+
+import numpy as np
+
+
+_SEQ_LENS: List[int] = [128, 512, 1000, 2048]
+_N_ATTEND_CALLS = 20          # per measurement
+_N_OUTLIER_CHANNELS = 4
+_N_CALIB_TOKENS = 50          # short so calibration completes inside every seq_len
+
+
+def _build(
+    method: str,
+    d: int,
+    bits: int,
+    jl_dim: int,
+    seed: int,
+    *,
+    vectorized: bool = False,
+    fused: bool = False,
+    outlier: bool = False,
+):
+    from mlx_kv_quant.cache.base import KVCacheBuilder
+
+    return (
+        KVCacheBuilder()
+        .with_method(method)
+        .with_head_dim(d)
+        .with_bit_width(inlier=bits)
+        .with_jl_dim(jl_dim)
+        .with_seed(seed)
+        .with_vectorized_attend(vectorized)
+        .with_fused_query_dot(fused)
+        .with_outlier_two_stream(outlier)
+        .with_n_outlier_channels(_N_OUTLIER_CHANNELS)
+        .with_n_calib_tokens(_N_CALIB_TOKENS)
+        .build()
+    )
+
+
+def _fill(cache, keys, vals) -> None:
+    for i in range(len(keys)):
+        cache.append(keys[i], vals[i])
+
+
+def _measure_attend_ms(cache, q, n_calls: int) -> float:
+    import mlx.core as mx
+
+    # Warm-up
+    mx.eval(cache.attend(q))
+
+    t0 = time.perf_counter()
+    for _ in range(n_calls):
+        mx.eval(cache.attend(q))
+    return (time.perf_counter() - t0) * 1_000.0 / n_calls
+
+
+def _correctness_check(cache_a, cache_b, q, label_a: str, label_b: str) -> None:
+    """Assert two caches produce numerically close attend outputs."""
+    import mlx.core as mx
+
+    out_a = np.array(cache_a.attend(q))
+    out_b = np.array(cache_b.attend(q))
+    mx.eval()
+    try:
+        np.testing.assert_allclose(out_a, out_b, rtol=5e-3, atol=5e-3)
+        print(f"  Correctness {label_a} vs {label_b}: OK (max_diff={np.max(np.abs(out_a-out_b)):.5f})")
+    except AssertionError as e:
+        print(f"  Correctness {label_a} vs {label_b}: FAIL — {e}")
+
+
+def run(
+    method: str,
+    d: int,
+    bits: int,
+    jl_dim: int,
+    seed: int,
+    seq_lens: List[int],
+    n_calls: int,
+    correctness: bool,
+) -> None:
+    import mlx.core as mx
+
+    rng = np.random.default_rng(seed)
+
+    configs = {
+        "baseline":   dict(vectorized=False, fused=False, outlier=False),
+        "vectorized": dict(vectorized=True,  fused=False, outlier=False),
+        "fused":      dict(vectorized=True,  fused=True,  outlier=False),
+        "all_opts":   dict(vectorized=True,  fused=True,  outlier=True),
+    }
+    # turboquant_mse doesn't have a fused path; skip 'fused'/'all_opts' for it.
+    if method == "turboquant_mse":
+        configs = {
+            "baseline":   dict(vectorized=False, fused=False, outlier=False),
+            "vectorized": dict(vectorized=True,  fused=False, outlier=False),
+            "all_opts":   dict(vectorized=True,  fused=False, outlier=True),
+        }
+
+    col_w = 14
+    header = f"{'seq_len':>8}  " + "  ".join(f"{k:>{col_w}}" for k in configs)
+    print(f"\n=== attend latency (ms/call) — method={method}, d={d}, bits={bits} ===")
+    print(header)
+    print("-" * len(header))
+
+    for seq_len in seq_lens:
+        keys = [mx.array(rng.standard_normal(d).astype(np.float16)) for _ in range(seq_len)]
+        vals = [mx.array(rng.standard_normal(d).astype(np.float16)) for _ in range(seq_len)]
+        q    = mx.array(rng.standard_normal(d).astype(np.float16))
+
+        caches = {}
+        latencies = {}
+        for name, flags in configs.items():
+            c = _build(method, d, bits, jl_dim, seed, **flags)
+            _fill(c, keys, vals)
+            caches[name] = c
+            latencies[name] = _measure_attend_ms(c, q, n_calls)
+
+        row = f"{seq_len:>8}  " + "  ".join(f"{latencies[k]:>{col_w}.3f}" for k in configs)
+        print(row)
+
+        # Speedup summary
+        base_ms = latencies["baseline"]
+        for name, ms in latencies.items():
+            if name != "baseline":
+                print(f"          {name:>12}: {base_ms/max(ms, 1e-9):.2f}× speedup vs baseline")
+
+        # Optional correctness check
+        if correctness and len(caches) >= 2:
+            names = list(caches.keys())
+            _correctness_check(caches[names[0]], caches[names[1]], q, names[0], names[1])
+
+        # Memory footprint
+        print(f"          memory (bytes): "
+              + ", ".join(f"{n}={c.memory_bytes()}" for n, c in caches.items()))
+        print()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="TurboQuant attend latency sweep")
+    parser.add_argument("--method", default="turboquant_prod",
+                        choices=["turboquant_prod", "turboquant_mse"])
+    parser.add_argument("--head_dim", type=int, default=128)
+    parser.add_argument("--bits", type=int, default=3)
+    parser.add_argument("--jl_dim", type=int, default=128)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--seq_lens", type=int, nargs="*", default=_SEQ_LENS)
+    parser.add_argument("--n_calls", type=int, default=_N_ATTEND_CALLS)
+    parser.add_argument("--correctness", action="store_true",
+                        help="Run cross-config correctness checks at each seq_len")
+    args = parser.parse_args()
+
+    run(
+        method=args.method,
+        d=args.head_dim,
+        bits=args.bits,
+        jl_dim=args.jl_dim,
+        seed=args.seed,
+        seq_lens=args.seq_lens,
+        n_calls=args.n_calls,
+        correctness=args.correctness,
+    )
+
+
+if __name__ == "__main__":
+    main()

mlx_kv_quant/cache/__init__.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from mlx_kv_quant.cache.base import KVCacheBuilder, KVCacheConfig, KVCacheFactory
+from mlx_kv_quant.cache.polar_cache import PolarQuantKVCache
+from mlx_kv_quant.cache.qjl_cache import QJLKVCache
+from mlx_kv_quant.cache.sliding_window_cache import SlidingWindowKVCache
+from mlx_kv_quant.cache.turboquant_cache import TurboQuantKVCache
+
+__all__ = [
+    "KVCacheBuilder",
+    "KVCacheConfig",
+    "KVCacheFactory",
+    "PolarQuantKVCache",
+    "QJLKVCache",
+    "SlidingWindowKVCache",
+    "TurboQuantKVCache",
+]