v0.1.0: MoE expert sniping for MLX — run models larger than your RAM

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.egg-info/
+dist/
+build/
+*.pyc
+.eggs/

README.md

@@ -0,0 +1,81 @@
+# CLI Agent — `mlx-expert-sniper`
+
+Pip-installable CLI that wraps the Expert Sniper research into a production tool.
+
+## Verified Results (M4 Mac Mini, 16 GB)
+
+| Metric | Value |
+|--------|-------|
+| Model | Qwen3-30B-A3B, 17.2 GB at 4-bit |
+| Standard mlx_lm | OOM |
+| **Sniper steady-state** | **4.22–4.68 tok/s** |
+| Cache hit rate | 85% (cold start) → 88.5% (warm) |
+| RAM used | 0.87 GB pinned |
+| Output | Coherent code, math, essays |
+
+## Install
+
+```bash
+cd research/expert-sniper/cli-agent
+pip install -e .
+```
+
+## Usage
+
+```bash
+# Preprocess model (one-time, ~17 GB on disk)
+mlx-sniper preprocess <hf-model-dir> -o ~/models/qwen3-30b
+
+# Or use the streaming preprocessor (downloads one shard at a time):
+python3 stream_preprocess.py
+
+# Generate
+mlx-sniper run ~/models/qwen3-30b -p "What is 2+2?" -v
+
+# Interactive chat
+mlx-sniper chat ~/models/qwen3-30b
+
+# OpenAI-compatible server
+mlx-sniper server ~/models/qwen3-30b --port 8899
+
+# Profile performance
+mlx-sniper profile ~/models/qwen3-30b --tokens 100
+
+# Show model info
+mlx-sniper info ~/models/qwen3-30b
+```
+
+## Files
+
+| File | Purpose |
+|------|---------|
+| `src/mlx_expert_sniper/sniper.py` | Core engine — the proven forward pass with expert sniping |
+| `src/mlx_expert_sniper/cache.py` | Per-expert LRU cache + pread/F_NOCACHE binary reader |
+| `src/mlx_expert_sniper/config.py` | SniperConfig dataclass |
+| `src/mlx_expert_sniper/generate.py` | High-level generate / stream_generate |
+| `src/mlx_expert_sniper/preprocess.py` | Convert HuggingFace model → sniper binary format |
+| `src/mlx_expert_sniper/server.py` | OpenAI-compatible HTTP server |
+| `src/mlx_expert_sniper/profile.py` | Per-token profiling tools |
+| `src/mlx_expert_sniper/cli.py` | `mlx-sniper` CLI entry point |
+| `stream_preprocess.py` | Streaming preprocessor (downloads one shard at a time) |
+
+## How It Relates to the Research
+
+This is a packaged version of the same forward pass in `../qwen3_agent.py`:
+- `cache.py` = extracted from `../expert_io.py`
+- `sniper.py:forward_token()` = same loop as `Qwen3SniperEngine.forward_token()`
+- `preprocess.py` = extracted from `../convert_qwen3_30b.py`
+- `server.py` = extracted from `../sniper_server.py`
+
+No algorithmic changes — same pread + F_NOCACHE + per-expert LRU + gather_qmm.
+
+## Python API
+
+```python
+from mlx_expert_sniper import SniperEngine
+
+engine = SniperEngine.from_dir("~/models/qwen3-30b")
+
+for token in engine.generate("Write a haiku about AI"):
+    print(token, end="", flush=True)
+```

pyproject.toml

@@ -0,0 +1,39 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mlx-expert-sniper"
+version = "0.1.0"
+description = "Run MoE models larger than RAM on Apple Silicon via expert sniping"
+readme = "README.md"
+license = "Apache-2.0"
+requires-python = ">=3.10"
+authors = [
+    {name = "walter-grace"},
+]
+keywords = ["mlx", "moe", "apple-silicon", "expert-sniping", "inference"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Operating System :: MacOS",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "mlx>=0.22.0",
+    "mlx-lm>=0.22.0",
+    "transformers>=4.40.0",
+    "numpy",
+    "huggingface-hub",
+]
+
+[project.optional-dependencies]
+dev = ["pytest", "rich"]
+server = ["rich"]
+
+[project.scripts]
+mlx-sniper = "mlx_expert_sniper.cli:main"
+
+[tool.setuptools.packages.find]
+where = ["src"]

src/mlx_expert_sniper/__init__.py

@@ -0,0 +1,18 @@
+"""mlx-expert-sniper: Run MoE models larger than RAM on Apple Silicon."""
+
+from .config import SniperConfig
+from .cache import PerExpertCache, ExpertReader
+from .sniper import SniperEngine
+from .generate import generate, stream_generate
+from .preprocess import preprocess_model
+
+__version__ = "0.1.0"
+__all__ = [
+    "SniperConfig",
+    "PerExpertCache",
+    "ExpertReader",
+    "SniperEngine",
+    "generate",
+    "stream_generate",
+    "preprocess_model",
+]

src/mlx_expert_sniper/cache.py

@@ -0,0 +1,206 @@
+"""
+Per-expert LRU cache + binary expert reader via pread + F_NOCACHE.
+
+Proven: 88.5% hit rate with 2,000-expert LRU on Qwen3-30B-A3B.
+Key insight: consecutive tokens share ~87.5% of experts (temporal locality),
+so per-expert LRU is the correct eviction policy.
+"""
+
+import os
+import json
+import fcntl
+import time
+import numpy as np
+from collections import OrderedDict
+from concurrent.futures import ThreadPoolExecutor
+
+import mlx.core as mx
+
+F_NOCACHE = 48  # macOS fcntl constant for bypassing page cache
+PAGE_SIZE = 16384  # 16KB — matches DART IOMMU page size
+
+MLX_DTYPES = {
+    "uint32": mx.uint32,
+    "float16": mx.float16,
+    "bfloat16": mx.float16,  # bf16 converted to f16 in bin files
+    "float32": mx.float32,
+}
+
+
+class PerExpertCache:
+    """LRU cache keyed by (layer_idx, expert_id).
+
+    NOT per-expert-set (which gets 0% hit rate because 1 different
+    expert out of 8 causes a full miss). Per-expert caching gets
+    88.5% because 7 of 8 experts are reused between consecutive tokens.
+    """
+
+    def __init__(self, max_experts: int = 2000):
+        self.cache: OrderedDict = OrderedDict()
+        self.max_experts = max_experts
+        self.hits = 0
+        self.misses = 0
+
+    def get(self, layer_idx: int, expert_id: int):
+        """Look up a single expert. Returns dict of tensors or None."""
+        key = (layer_idx, expert_id)
+        if key in self.cache:
+            self.hits += 1
+            self.cache.move_to_end(key)
+            return self.cache[key]
+        self.misses += 1
+        return None
+
+    def put(self, layer_idx: int, expert_id: int, data: dict):
+        """Insert expert data, evicting LRU entries if full."""
+        key = (layer_idx, expert_id)
+        self.cache[key] = data
+        self.cache.move_to_end(key)
+        while len(self.cache) > self.max_experts:
+            self.cache.popitem(last=False)
+
+    @property
+    def hit_rate(self) -> float:
+        total = self.hits + self.misses
+        return self.hits / max(total, 1)
+
+    def stats(self) -> str:
+        total = self.hits + self.misses
+        rate = self.hits / max(total, 1) * 100
+        return (f"{rate:.0f}% hit ({self.hits}h/{self.misses}m), "
+                f"{len(self.cache)} cached")
+
+    def clear(self):
+        self.cache.clear()
+        self.hits = 0
+        self.misses = 0
+
+
+class ExpertReader:
+    """Reads specific experts from binary layer files via pread + F_NOCACHE.
+
+    Binary format per layer file:
+        Header (PAGE_SIZE bytes): JSON with layout info
+        Data: expert_0 | expert_1 | ... | expert_N
+        Each expert: gate_proj.weight | gate_proj.scales | gate_proj.biases |
+                     up_proj.weight | up_proj.scales | up_proj.biases |
+                     down_proj.weight | down_proj.scales | down_proj.biases
+    """
+
+    def __init__(self, bin_dir: str, num_workers: int = 4):
+        self.bin_dir = bin_dir
+        self.executor = ThreadPoolExecutor(max_workers=num_workers)
+
+        # Parse headers for all MoE layer files
+        self.headers = {}
+        self.fds = {}
+        for f in os.listdir(bin_dir):
+            if f.startswith("moe_layer_") and f.endswith(".bin"):
+                layer_idx = int(f.split("_")[2].split(".")[0])
+                path = os.path.join(bin_dir, f)
+                with open(path, "rb") as fh:
+                    raw = fh.read(PAGE_SIZE)
+                self.headers[layer_idx] = json.loads(raw.rstrip(b"\x00"))
+
+        if not self.headers:
+            raise FileNotFoundError(f"No moe_layer_*.bin files in {bin_dir}")
+
+        # Layout from first layer (all layers share the same layout)
+        first = next(iter(self.headers.values()))["layout"]
+        self.expert_block_size = first["expert_block_size"]
+        self.data_start = first["data_start"]
+        self.tensor_layout = first["tensors"]
+
+        # Stats
+        self.read_time = 0.0
+        self.reads = 0
+        self.bytes_read = 0
+
+    @property
+    def num_layers(self) -> int:
+        return len(self.headers)
+
+    def _get_fd(self, layer_idx: int) -> int:
+        """Get (or open) a file descriptor with F_NOCACHE for a layer."""
+        if layer_idx not in self.fds:
+            path = os.path.join(self.bin_dir, f"moe_layer_{layer_idx:02d}.bin")
+            fd = os.open(path, os.O_RDONLY)
+            fcntl.fcntl(fd, F_NOCACHE, 1)  # Bypass macOS Unified Buffer Cache
+            self.fds[layer_idx] = fd
+        return self.fds[layer_idx]
+
+    def _read_expert(self, layer_idx: int, expert_id: int) -> bytes:
+        """Read one expert's raw bytes via pread (no seek, thread-safe)."""
+        fd = self._get_fd(layer_idx)
+        offset = self.data_start + expert_id * self.expert_block_size
+        return os.pread(fd, self.expert_block_size, offset)
+
+    def _parse_expert(self, raw_bytes: bytes) -> dict:
+        """Parse raw bytes into dict of MLX arrays."""
+        result = {}
+        for name, info in self.tensor_layout.items():
+            off = info["inner_offset"]
+            nbytes = info["nbytes"]
+            shape = info["shape_per_expert"]
+            dtype = MLX_DTYPES.get(info["dtype"], mx.float16)
+
+            arr_bytes = raw_bytes[off:off + nbytes]
+            if dtype == mx.uint32:
+                np_arr = np.frombuffer(arr_bytes, dtype=np.uint32).reshape(shape)
+            elif dtype == mx.float16:
+                np_arr = np.frombuffer(arr_bytes, dtype=np.float16).reshape(shape)
+            elif dtype == mx.float32:
+                np_arr = np.frombuffer(arr_bytes, dtype=np.float32).reshape(shape)
+            else:
+                np_arr = np.frombuffer(arr_bytes, dtype=np.float16).reshape(shape)
+
+            result[name] = mx.array(np_arr)
+        return result
+
+    def get_experts(self, layer_idx: int, expert_ids: list) -> dict:
+        """Read and parse active experts for a layer (parallel pread).
+
+        Args:
+            layer_idx: MoE layer index
+            expert_ids: list of expert IDs to load
+
+        Returns:
+            dict mapping expert_id -> dict of tensor_name -> mx.array
+        """
+        t0 = time.time()
+
+        futures = {
+            eid: self.executor.submit(self._read_expert, layer_idx, eid)
+            for eid in expert_ids
+        }
+
+        experts = {}
+        for eid, future in futures.items():
+            raw = future.result()
+            experts[eid] = self._parse_expert(raw)
+            self.bytes_read += len(raw)
+
+        self.read_time += time.time() - t0
+        self.reads += len(expert_ids)
+        return experts
+
+    def stats(self) -> str:
+        if self.reads == 0:
+            return "No reads"
+        avg_ms = self.read_time / self.reads * 1000
+        throughput = self.bytes_read / max(self.read_time, 0.001) / 1e9
+        return (f"reads={self.reads}, avg={avg_ms:.1f}ms/expert, "
+                f"throughput={throughput:.1f} GB/s, "
+                f"total={self.bytes_read / 1e9:.2f} GB")
+
+    def close(self):
+        for fd in self.fds.values():
+            os.close(fd)
+        self.executor.shutdown(wait=False)
+        self.fds.clear()
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass

src/mlx_expert_sniper/cli.py

@@ -0,0 +1,331 @@
+"""
+CLI for mlx-expert-sniper.
+
+Usage:
+    mlx-sniper preprocess <model_dir> -o <output_dir>
+    mlx-sniper run <sniper_dir> --prompt "Hello" --max-tokens 100
+    mlx-sniper chat <sniper_dir>
+    mlx-sniper profile <sniper_dir> --tokens 50
+    mlx-sniper server <sniper_dir> --port 8899
+    mlx-sniper info <sniper_dir>
+"""
+
+import argparse
+import json
+import sys
+import time
+import os
+
+
+def cmd_preprocess(args):
+    """Preprocess a HuggingFace model into sniper format."""
+    from .preprocess import preprocess_model
+
+    print("=" * 50)
+    print("  mlx-expert-sniper: preprocess")
+    print("=" * 50)
+
+    result = preprocess_model(
+        model_dir=args.model_dir,
+        output_dir=args.output,
+        verbose=True,
+    )
+    print(f"\n  Ready: {result['output_dir']}")
+    print(f"  Run:   mlx-sniper run {result['output_dir']}")
+
+
+def cmd_run(args):
+    """Generate text from a prompt."""
+    from .sniper import SniperEngine
+
+    print("  Loading...", end="", flush=True)
+    engine = SniperEngine.from_dir(
+        args.sniper_dir,
+        max_cached_experts=args.cache_size,
+    )
+    print(" ready.")
+
+    t0 = time.time()
+    tokens = 0
+
+    if args.stream:
+        for token in engine.generate(
+            args.prompt,
+            max_tokens=args.max_tokens,
+            temperature=args.temperature,
+        ):
+            print(token, end="", flush=True)
+            tokens += 1
+        print()
+    else:
+        result = []
+        for token in engine.generate(
+            args.prompt,
+            max_tokens=args.max_tokens,
+            temperature=args.temperature,
+        ):
+            result.append(token)
+            tokens += 1
+        print("".join(result))
+
+    elapsed = time.time() - t0
+    tps = tokens / elapsed if elapsed > 0 else 0
+
+    if args.verbose:
+        print(f"\n  {tokens} tokens in {elapsed:.1f}s ({tps:.2f} tok/s)")
+        print(f"  Cache: {engine.expert_cache.stats()}")
+        print(f"  I/O:   {engine.reader.stats()}")
+
+    engine.close()
+
+
+def cmd_chat(args):
+    """Interactive chat session."""
+    from .sniper import SniperEngine
+
+    print("=" * 50)
+    print("  mlx-expert-sniper: chat")
+    print("=" * 50)
+    print("  Loading...", end="", flush=True)
+
+    engine = SniperEngine.from_dir(
+        args.sniper_dir,
+        max_cached_experts=args.cache_size,
+    )
+    print(" ready.")
+    print(f"  Model: {engine.config.model_type}")
+    print(f"  Cache: {engine.config.max_cached_experts} experts")
+    print(f"  Commands: /stats, /clear, /quit")
+    print()
+
+    history = []
+    session_tokens = 0
+    session_time = 0.0
+
+    while True:
+        try:
+            user_input = input("  > ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\n  goodbye.")
+            break
+
+        if not user_input:
+            continue
+
+        if user_input in ("/quit", "/exit", "/q"):
+            break
+
+        if user_input == "/stats":
+            avg = session_tokens / session_time if session_time > 0 else 0
+            print(f"  Tokens: {session_tokens:,}")
+            print(f"  Time:   {session_time:.1f}s")
+            print(f"  Speed:  {avg:.2f} tok/s")
+            print(f"  Cache:  {engine.expert_cache.stats()}")
+            print(f"  I/O:    {engine.reader.stats()}")
+            print()
+            continue
+
+        if user_input == "/clear":
+            history.clear()
+            engine.expert_cache.clear()
+            print("  (cleared)")
+            continue
+
+        history.append({"role": "user", "content": user_input})
+        messages = history[-20:]  # Keep last 10 turns
+
+        t0 = time.time()
+        response = ""
+        token_count = 0
+
+        print("  ", end="", flush=True)
+        for token in engine.generate(
+            prompt=None,
+            chat_messages=messages,
+            max_tokens=args.max_tokens,
+            temperature=args.temperature,
+        ):
+            print(token, end="", flush=True)
+            response += token
+            token_count += 1
+        print()
+
+        elapsed = time.time() - t0
+        tps = token_count / elapsed if elapsed > 0 else 0
+        session_tokens += token_count
+        session_time += elapsed
+        print(f"  [{token_count} tokens, {elapsed:.1f}s, {tps:.2f} tok/s]")
+        print()
+
+        history.append({"role": "assistant", "content": response})
+
+    engine.close()
+
+
+def cmd_profile(args):
+    """Profile inference performance."""
+    from .sniper import SniperEngine
+    from .profile import ProfileSession, TokenProfile
+
+    print("  Loading...", end="", flush=True)
+    engine = SniperEngine.from_dir(
+        args.sniper_dir,
+        max_cached_experts=args.cache_size,
+    )
+    print(" ready.")
+
+    prompt = args.prompt or "Explain the theory of relativity in detail."
+    session = ProfileSession()
+
+    print(f"  Generating {args.tokens} tokens...")
+    print(f"  Prompt: {prompt[:60]}...")
+    print()
+
+    t_prev = time.time()
+    token_idx = 0
+    prev_hits = 0
+    prev_misses = 0
+
+    for token in engine.generate(
+        prompt,
+        max_tokens=args.tokens,
+        temperature=0.0,  # Deterministic for profiling
+    ):
+        t_now = time.time()
+        cur_hits = engine.expert_cache.hits
+        cur_misses = engine.expert_cache.misses
+
+        profile = TokenProfile(
+            token_idx=token_idx,
+            total_ms=(t_now - t_prev) * 1000,
+            cache_hits=cur_hits - prev_hits,
+            cache_misses=cur_misses - prev_misses,
+        )
+        session.add(profile)
+
+        if token_idx % 20 == 0:
+            print(f"  [{token_idx:3d}] {repr(token):20s}  {profile.total_ms:.0f}ms  "
+                  f"hits={profile.cache_hits} miss={profile.cache_misses}", flush=True)
+
+        prev_hits = cur_hits
+        prev_misses = cur_misses
+        token_idx += 1
+        t_prev = t_now
+
+    print()
+    print(session.summary())
+
+    if args.csv:
+        session.to_csv(args.csv)
+        print(f"\n  CSV: {args.csv}")
+
+    print(f"\n  Cache: {engine.expert_cache.stats()}")
+    print(f"  I/O:   {engine.reader.stats()}")
+    engine.close()
+
+
+def cmd_server(args):
+    """Start an OpenAI-compatible HTTP server."""
+    from .server import run_server
+
+    run_server(
+        sniper_dir=args.sniper_dir,
+        host=args.host,
+        port=args.port,
+        cache_size=args.cache_size,
+    )
+
+
+def cmd_info(args):
+    """Show info about a sniper directory."""
+    from .config import SniperConfig
+    from pathlib import Path
+    import glob
+
+    cfg = SniperConfig.from_dir(args.sniper_dir)
+
+    pinned_size = os.path.getsize(cfg.pinned_path) / 1e9 if os.path.exists(cfg.pinned_path) else 0
+    bin_files = sorted(glob.glob(os.path.join(cfg.bin_dir, "moe_layer_*.bin")))
+    expert_total = sum(os.path.getsize(f) for f in bin_files) / 1e9
+
+    print(f"  Directory:  {cfg.sniper_dir}")
+    print(f"  Model type: {cfg.model_type}")
+    print(f"  Tokenizer:  {cfg.tokenizer_name}")
+    print(f"  Layers:     {cfg.num_hidden_layers}")
+    print(f"  Experts:    {cfg.num_experts} per layer, top-{cfg.num_experts_per_tok}")
+    print(f"  Hidden:     {cfg.hidden_size}")
+    print(f"  MoE inter:  {cfg.moe_intermediate_size}")
+    print(f"  Quant:      {cfg.bits}-bit, group_size={cfg.group_size}")
+    print(f"  Pinned:     {pinned_size:.2f} GB")
+    print(f"  Experts:    {len(bin_files)} layers, {expert_total:.2f} GB")
+    print(f"  Total:      {pinned_size + expert_total:.2f} GB")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="mlx-sniper",
+        description="Run MoE models larger than RAM on Apple Silicon",
+    )
+    sub = parser.add_subparsers(dest="command")
+
+    # preprocess
+    p = sub.add_parser("preprocess", help="Convert HF model to sniper format")
+    p.add_argument("model_dir", help="Path to HuggingFace model directory")
+    p.add_argument("-o", "--output", required=True, help="Output directory")
+
+    # run
+    p = sub.add_parser("run", help="Generate text from a prompt")
+    p.add_argument("sniper_dir", help="Path to sniper directory")
+    p.add_argument("--prompt", "-p", required=True, help="Text prompt")
+    p.add_argument("--max-tokens", type=int, default=200)
+    p.add_argument("--temperature", type=float, default=0.7)
+    p.add_argument("--stream", action="store_true", default=True)
+    p.add_argument("--no-stream", dest="stream", action="store_false")
+    p.add_argument("--cache-size", type=int, default=2000)
+    p.add_argument("--verbose", "-v", action="store_true")
+
+    # chat
+    p = sub.add_parser("chat", help="Interactive chat")
+    p.add_argument("sniper_dir", help="Path to sniper directory")
+    p.add_argument("--max-tokens", type=int, default=500)
+    p.add_argument("--temperature", type=float, default=0.7)
+    p.add_argument("--cache-size", type=int, default=2000)
+
+    # profile
+    p = sub.add_parser("profile", help="Profile inference performance")
+    p.add_argument("sniper_dir", help="Path to sniper directory")
+    p.add_argument("--tokens", type=int, default=50)
+    p.add_argument("--prompt", default=None)
+    p.add_argument("--csv", default=None, help="Output CSV path")
+    p.add_argument("--cache-size", type=int, default=2000)
+
+    # server
+    p = sub.add_parser("server", help="OpenAI-compatible HTTP server")
+    p.add_argument("sniper_dir", help="Path to sniper directory")
+    p.add_argument("--host", default="0.0.0.0")
+    p.add_argument("--port", type=int, default=8899)
+    p.add_argument("--cache-size", type=int, default=2000)
+
+    # info
+    p = sub.add_parser("info", help="Show sniper directory info")
+    p.add_argument("sniper_dir", help="Path to sniper directory")
+
+    args = parser.parse_args()
+
+    if args.command is None:
+        parser.print_help()
+        sys.exit(1)
+
+    commands = {
+        "preprocess": cmd_preprocess,
+        "run": cmd_run,
+        "chat": cmd_chat,
+        "profile": cmd_profile,
+        "server": cmd_server,
+        "info": cmd_info,
+    }
+    commands[args.command](args)
+
+
+if __name__ == "__main__":
+    main()

src/mlx_expert_sniper/config.py

@@ -0,0 +1,88 @@
+"""Configuration for the expert sniper."""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+import json
+
+
+@dataclass
+class SniperConfig:
+    """Configuration for expert sniping inference."""
+
+    # Model paths
+    sniper_dir: str = ""  # Directory with pinned.safetensors + bin/
+
+    # Cache settings
+    max_cached_experts: int = 2000  # LRU cache capacity
+    num_io_workers: int = 4  # Parallel pread threads
+
+    # Quantization (read from config.json)
+    bits: int = 4
+    group_size: int = 64
+
+    # MLX memory limits
+    memory_limit_gb: float = 12.0
+    cache_limit_mb: int = 256
+
+    # Generation defaults
+    max_tokens: int = 500
+    temperature: float = 0.7
+    repetition_penalty: float = 1.0
+
+    # Model architecture (populated from config.json)
+    num_hidden_layers: int = 0
+    num_experts: int = 0
+    num_experts_per_tok: int = 0
+    hidden_size: int = 0
+    moe_intermediate_size: int = 0
+    vocab_size: int = 0
+    norm_topk_prob: bool = True
+    model_type: str = ""
+    tokenizer_name: str = ""
+
+    @classmethod
+    def from_dir(cls, sniper_dir: str, **overrides) -> "SniperConfig":
+        """Load config from a sniper directory (must have config.json)."""
+        sniper_dir = str(Path(sniper_dir).expanduser().resolve())
+        config_path = Path(sniper_dir) / "config.json"
+        if not config_path.exists():
+            raise FileNotFoundError(f"No config.json in {sniper_dir}")
+
+        with open(config_path) as f:
+            model_config = json.load(f)
+
+        quant = model_config.get("quantization", {})
+
+        cfg = cls(
+            sniper_dir=sniper_dir,
+            bits=quant.get("bits", 4),
+            group_size=quant.get("group_size", 64),
+            num_hidden_layers=model_config.get("num_hidden_layers", 0),
+            num_experts=model_config.get("num_experts", 0),
+            num_experts_per_tok=model_config.get("num_experts_per_tok", 0),
+            hidden_size=model_config.get("hidden_size", 0),
+            moe_intermediate_size=model_config.get("moe_intermediate_size", 0),
+            vocab_size=model_config.get("vocab_size", 0),
+            norm_topk_prob=model_config.get("norm_topk_prob", True),
+            model_type=model_config.get("model_type", ""),
+            tokenizer_name=model_config.get("_name_or_path", ""),
+        )
+
+        for k, v in overrides.items():
+            if hasattr(cfg, k):
+                setattr(cfg, k, v)
+
+        return cfg
+
+    @property
+    def pinned_path(self) -> str:
+        return str(Path(self.sniper_dir) / "pinned.safetensors")
+
+    @property
+    def bin_dir(self) -> str:
+        return str(Path(self.sniper_dir) / "bin")
+
+    @property
+    def config_json_path(self) -> str:
+        return str(Path(self.sniper_dir) / "config.json")

src/mlx_expert_sniper/generate.py

@@ -0,0 +1,80 @@
+"""
+High-level generation functions.
+
+These wrap SniperEngine for simple one-shot usage.
+"""
+
+import time
+from typing import Optional
+
+from .config import SniperConfig
+from .sniper import SniperEngine
+
+
+def generate(
+    sniper_dir: str,
+    prompt: str,
+    max_tokens: int = 200,
+    temperature: float = 0.7,
+    verbose: bool = False,
+    **config_overrides,
+) -> str:
+    """Generate text from a prompt. Returns the full response.
+
+    Args:
+        sniper_dir: Path to preprocessed sniper directory
+        prompt: Text prompt
+        max_tokens: Maximum tokens to generate
+        temperature: Sampling temperature (0 = greedy)
+        verbose: Print timing info
+        **config_overrides: Override SniperConfig fields
+
+    Returns:
+        Generated text as a string
+    """
+    engine = SniperEngine.from_dir(sniper_dir, **config_overrides)
+    try:
+        result = []
+        t0 = time.time()
+        for token in engine.generate(prompt, max_tokens=max_tokens,
+                                     temperature=temperature):
+            result.append(token)
+        elapsed = time.time() - t0
+
+        if verbose:
+            n = len(result)
+            tps = n / elapsed if elapsed > 0 else 0
+            print(f"\n[{n} tokens in {elapsed:.1f}s — {tps:.2f} tok/s]")
+            print(f"[Cache: {engine.expert_cache.stats()}]")
+            print(f"[I/O: {engine.reader.stats()}]")
+
+        return "".join(result)
+    finally:
+        engine.close()
+
+
+def stream_generate(
+    sniper_dir: str,
+    prompt: str,
+    max_tokens: int = 200,
+    temperature: float = 0.7,
+    **config_overrides,
+):
+    """Stream tokens from a prompt. Yields each token as it's generated.
+
+    Args:
+        sniper_dir: Path to preprocessed sniper directory
+        prompt: Text prompt
+        max_tokens: Maximum tokens to generate
+        temperature: Sampling temperature (0 = greedy)
+        **config_overrides: Override SniperConfig fields
+
+    Yields:
+        str: Each generated token
+    """
+    engine = SniperEngine.from_dir(sniper_dir, **config_overrides)
+    try:
+        yield from engine.generate(prompt, max_tokens=max_tokens,
+                                   temperature=temperature)
+    finally:
+        engine.close()

src/mlx_expert_sniper/preprocess.py

@@ -0,0 +1,234 @@
+"""
+Preprocess a HuggingFace MoE model into sniper format.
+
+Converts from:
+  ~/.cache/huggingface/hub/models--<repo>/snapshots/<hash>/
+    model-00001-of-00007.safetensors
+    ...
+    config.json, tokenizer.json, etc.
+
+To:
+  <output_dir>/
+    config.json
+    tokenizer.json
+    tokenizer_config.json
+    pinned.safetensors      (~0.87 GB for Qwen3-30B)
+    bin/
+      moe_layer_00.bin      (~340 MB per layer)
+      moe_layer_01.bin
+      ...
+"""
+
+import os
+import sys
+import json
+import time
+import gc
+import glob
+import shutil
+from pathlib import Path
+
+import numpy as np
+import mlx.core as mx
+
+PAGE_SIZE = 16384
+
+TENSOR_NAMES = [
+    "gate_proj.weight", "gate_proj.scales", "gate_proj.biases",
+    "up_proj.weight", "up_proj.scales", "up_proj.biases",
+    "down_proj.weight", "down_proj.scales", "down_proj.biases",
+]
+
+
+def _convert_layer_to_bin(layer_data, layer_idx, num_experts, output_dir):
+    """Convert one layer's expert tensors to binary format."""
+    tensor_info = {}
+    expert_block_size = 0
+
+    for name in TENSOR_NAMES:
+        t = layer_data[name]
+        per_expert_shape = list(t.shape[1:])
+        if t.dtype == mx.uint32:
+            elem_size = 4
+        elif t.dtype in (mx.bfloat16, mx.float16):
+            elem_size = 2
+        else:
+            elem_size = 4
+        nbytes = 1
+        for s in per_expert_shape:
+            nbytes *= s
+        nbytes *= elem_size
+        tensor_info[name] = {
+            "shape_per_expert": per_expert_shape,
+            "dtype": str(t.dtype).replace("mlx.core.", ""),
+            "nbytes": nbytes,
+            "inner_offset": expert_block_size,
+        }
+        expert_block_size += nbytes
+
+    header = {
+        "layer_idx": layer_idx,
+        "num_experts": num_experts,
+        "layout": {
+            "expert_block_size": expert_block_size,
+            "data_start": PAGE_SIZE,
+            "tensors": tensor_info,
+        }
+    }
+    header_bytes = json.dumps(header, indent=2).encode()
+    assert len(header_bytes) < PAGE_SIZE
+    header_bytes += b"\x00" * (PAGE_SIZE - len(header_bytes))
+
+    out_path = os.path.join(output_dir, "bin", f"moe_layer_{layer_idx:02d}.bin")
+    with open(out_path, "wb") as f:
+        f.write(header_bytes)
+        for expert_id in range(num_experts):
+            for name in TENSOR_NAMES:
+                t = layer_data[name][expert_id]
+                if t.dtype == mx.bfloat16:
+                    raw = np.array(t.astype(mx.float16)).astype(np.float16).tobytes()
+                elif t.dtype == mx.uint32:
+                    raw = np.array(t).astype(np.uint32).tobytes()
+                else:
+                    raw = np.array(t).tobytes()
+                f.write(raw)
+
+    return os.path.getsize(out_path), expert_block_size
+
+
+def preprocess_model(model_dir: str, output_dir: str, verbose: bool = True):
+    """Convert a HuggingFace MoE model to sniper binary format.
+
+    Args:
+        model_dir: Path to HuggingFace model directory (with safetensors shards)
+        output_dir: Where to write the sniper format
+        verbose: Print progress
+    """
+    model_dir = str(Path(model_dir).expanduser().resolve())
+    output_dir = str(Path(output_dir).expanduser().resolve())
+
+    # Find model directory (could be a HF cache path or direct path)
+    if not os.path.exists(os.path.join(model_dir, "config.json")):
+        # Try HF cache format
+        snap_dirs = glob.glob(os.path.join(model_dir, "snapshots", "*/"))
+        if snap_dirs:
+            model_dir = snap_dirs[0]
+
+    config_path = os.path.join(model_dir, "config.json")
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(f"No config.json in {model_dir}")
+
+    os.makedirs(os.path.join(output_dir, "bin"), exist_ok=True)
+
+    # Copy config + tokenizer files
+    for fname in ["config.json", "tokenizer.json", "tokenizer_config.json",
+                   "special_tokens_map.json", "vocab.json", "merges.txt"]:
+        src = os.path.join(model_dir, fname)
+        if os.path.exists(src):
+            shutil.copy(src, os.path.join(output_dir, fname))
+
+    # Find shard index
+    index_path = os.path.join(model_dir, "model.safetensors.index.json")
+    if os.path.exists(index_path):
+        with open(index_path) as f:
+            idx = json.load(f)
+        shards = sorted(set(idx["weight_map"].values()))
+    else:
+        # Single shard
+        shards = [f for f in os.listdir(model_dir) if f.endswith(".safetensors")]
+
+    if verbose:
+        print(f"  Source: {model_dir}")
+        print(f"  Output: {output_dir}")
+        print(f"  Shards: {len(shards)}")
+
+    pinned = {}
+    layers_done = set()
+
+    with open(os.path.join(output_dir, "config.json")) as f:
+        config = json.load(f)
+    num_layers = config.get("num_hidden_layers", 48)
+
+    for si, shard_name in enumerate(shards):
+        if verbose:
+            print(f"\n  [{si + 1}/{len(shards)}] {shard_name}...")
+
+        t0 = time.time()
+        data = mx.load(os.path.join(model_dir, shard_name))
+
+        # Classify tensors as expert vs pinned
+        layer_experts = {}
+        for key, tensor in data.items():
+            if "switch_mlp" in key:
+                layer = int(key.split(".layers.")[1].split(".")[0])
+                short = key.split(".switch_mlp.")[1]
+                layer_experts.setdefault(layer, {})[short] = tensor
+            else:
+                pinned[key] = tensor
+
+        # Convert complete expert layers
+        for layer_idx, tensors in layer_experts.items():
+            if len(tensors) < 9:  # Partial layer (spans shards)
+                continue
+            if layer_idx in layers_done:
+                continue
+
+            num_experts = tensors[list(tensors.keys())[0]].shape[0]
+            sz, ebs = _convert_layer_to_bin(tensors, layer_idx, num_experts,
+                                            output_dir)
+            layers_done.add(layer_idx)
+            if verbose:
+                print(f"    Layer {layer_idx}: {sz / 1e6:.0f} MB "
+                      f"({ebs / 1e6:.2f} MB/expert)")
+
+        del data, layer_experts
+        gc.collect()
+        mx.clear_cache()
+        if verbose:
+            print(f"    {time.time() - t0:.1f}s")
+
+    # Handle partial layers (spanning multiple shards)
+    incomplete = set(range(num_layers)) - layers_done
+    if incomplete:
+        if verbose:
+            print(f"\n  Merging {len(incomplete)} partial layers...")
+        for layer_idx in sorted(incomplete):
+            tensors = {}
+            for shard_name in shards:
+                data = mx.load(os.path.join(model_dir, shard_name))
+                for key, tensor in data.items():
+                    if f".layers.{layer_idx}." in key and "switch_mlp" in key:
+                        short = key.split(".switch_mlp.")[1]
+                        tensors[short] = tensor
+                del data
+            if len(tensors) >= 9:
+                num_experts = tensors[list(tensors.keys())[0]].shape[0]
+                sz, ebs = _convert_layer_to_bin(tensors, layer_idx, num_experts,
+                                                output_dir)
+                layers_done.add(layer_idx)
+                if verbose:
+                    print(f"    Layer {layer_idx}: {sz / 1e6:.0f} MB (merged)")
+
+    # Save pinned weights
+    if verbose:
+        print(f"\n  Saving pinned weights ({len(pinned)} tensors)...")
+    pinned_path = os.path.join(output_dir, "pinned.safetensors")
+    mx.save_safetensors(pinned_path, pinned)
+    psz = os.path.getsize(pinned_path) / 1e9
+    del pinned
+
+    # Summary
+    bin_files = sorted(glob.glob(os.path.join(output_dir, "bin", "moe_layer_*.bin")))
+    total = sum(os.path.getsize(f) for f in bin_files)
+    if verbose:
+        print(f"\n  Expert layers: {len(bin_files)}/{num_layers}")
+        print(f"  Expert total:  {total / 1e9:.2f} GB")
+        print(f"  Pinned:        {psz:.2f} GB")
+        print(f"  Done!")
+
+    return {
+        "expert_layers": len(bin_files),
+        "expert_total_gb": total / 1e9,
+        "pinned_gb": psz,
+        "output_dir": output_dir,
+    }

src/mlx_expert_sniper/profile.py

@@ -0,0 +1,93 @@
+"""
+Profiling tools for expert sniping inference.
+
+Measures per-token breakdown: attention, router, I/O, compute, cache stats.
+"""
+
+import time
+from dataclasses import dataclass, field
+
+
+@dataclass
+class TokenProfile:
+    """Timing breakdown for a single token."""
+    token_idx: int = 0
+    attention_ms: float = 0.0
+    router_ms: float = 0.0
+    io_ms: float = 0.0
+    compute_ms: float = 0.0
+    total_ms: float = 0.0
+    cache_hits: int = 0
+    cache_misses: int = 0
+    experts_loaded: int = 0
+
+
+@dataclass
+class ProfileSession:
+    """Accumulated profiling data across multiple tokens."""
+    tokens: list = field(default_factory=list)
+
+    def add(self, profile: TokenProfile):
+        self.tokens.append(profile)
+
+    @property
+    def num_tokens(self) -> int:
+        return len(self.tokens)
+
+    @property
+    def avg_tok_per_sec(self) -> float:
+        if not self.tokens:
+            return 0.0
+        total_ms = sum(t.total_ms for t in self.tokens)
+        return len(self.tokens) / (total_ms / 1000) if total_ms > 0 else 0
+
+    @property
+    def total_cache_hits(self) -> int:
+        return sum(t.cache_hits for t in self.tokens)
+
+    @property
+    def total_cache_misses(self) -> int:
+        return sum(t.cache_misses for t in self.tokens)
+
+    @property
+    def cache_hit_rate(self) -> float:
+        total = self.total_cache_hits + self.total_cache_misses
+        return self.total_cache_hits / max(total, 1)
+
+    def summary(self) -> str:
+        if not self.tokens:
+            return "No tokens profiled"
+
+        n = len(self.tokens)
+        total_ms = sum(t.total_ms for t in self.tokens)
+        avg_ms = total_ms / n
+        tps = n / (total_ms / 1000) if total_ms > 0 else 0
+
+        # Skip first token (prefill) for steady-state stats
+        steady = self.tokens[1:] if n > 1 else self.tokens
+        if steady:
+            ss_ms = sum(t.total_ms for t in steady) / len(steady)
+            ss_tps = len(steady) / (sum(t.total_ms for t in steady) / 1000)
+        else:
+            ss_ms = avg_ms
+            ss_tps = tps
+
+        lines = [
+            f"  Tokens:     {n}",
+            f"  Total:      {total_ms / 1000:.2f}s",
+            f"  Avg:        {avg_ms:.1f} ms/token ({tps:.2f} tok/s)",
+            f"  Steady:     {ss_ms:.1f} ms/token ({ss_tps:.2f} tok/s)",
+            f"  Cache:      {self.cache_hit_rate * 100:.1f}% hit rate "
+            f"({self.total_cache_hits}h/{self.total_cache_misses}m)",
+        ]
+        return "\n".join(lines)
+
+    def to_csv(self, path: str):
+        """Write per-token profiles to CSV."""
+        with open(path, "w") as f:
+            f.write("token,total_ms,attention_ms,router_ms,io_ms,compute_ms,"
+                    "cache_hits,cache_misses\n")
+            for t in self.tokens:
+                f.write(f"{t.token_idx},{t.total_ms:.2f},{t.attention_ms:.2f},"
+                        f"{t.router_ms:.2f},{t.io_ms:.2f},{t.compute_ms:.2f},"
+                        f"{t.cache_hits},{t.cache_misses}\n")

src/mlx_expert_sniper/server.py

@@ -0,0 +1,167 @@
+"""
+OpenAI-compatible HTTP server wrapping the sniper engine.
+
+Endpoints:
+    GET  /v1/models              — list available models
+    POST /v1/chat/completions    — chat (streaming + non-streaming)
+    GET  /health                 — health check
+"""
+
+import json
+import time
+import uuid
+from http.server import HTTPServer, BaseHTTPRequestHandler
+
+from .sniper import SniperEngine
+
+_engine = None
+_model_name = "mlx-expert-sniper"
+
+
+class _Handler(BaseHTTPRequestHandler):
+    def log_message(self, format, *args):
+        pass
+
+    def _send_json(self, data, status=200):
+        body = json.dumps(data).encode()
+        self.send_response(status)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", len(body))
+        self.send_header("Access-Control-Allow-Origin", "*")
+        self.end_headers()
+        self.wfile.write(body)
+
+    def _send_sse(self, data):
+        self.wfile.write(f"data: {json.dumps(data)}\n\n".encode())
+        self.wfile.flush()
+
+    def do_OPTIONS(self):
+        self.send_response(200)
+        self.send_header("Access-Control-Allow-Origin", "*")
+        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
+        self.send_header("Access-Control-Allow-Headers", "Content-Type, Authorization")
+        self.end_headers()
+
+    def do_GET(self):
+        if self.path == "/health":
+            self._send_json({"status": "ok", "model": _model_name})
+        elif self.path == "/v1/models":
+            self._send_json({
+                "object": "list",
+                "data": [{
+                    "id": _model_name,
+                    "object": "model",
+                    "created": int(time.time()),
+                    "owned_by": "mlx-expert-sniper",
+                }]
+            })
+        else:
+            self._send_json({"error": "not found"}, 404)
+
+    def do_POST(self):
+        if self.path != "/v1/chat/completions":
+            self._send_json({"error": "not found"}, 404)
+            return
+
+        content_length = int(self.headers.get("Content-Length", 0))
+        body = json.loads(self.rfile.read(content_length))
+
+        messages = body.get("messages", [])
+        max_tokens = body.get("max_tokens", 500)
+        temperature = body.get("temperature", 0.7)
+        stream = body.get("stream", False)
+        request_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
+
+        if stream:
+            self.send_response(200)
+            self.send_header("Content-Type", "text/event-stream")
+            self.send_header("Cache-Control", "no-cache")
+            self.send_header("Access-Control-Allow-Origin", "*")
+            self.end_headers()
+
+            self._send_sse({
+                "id": request_id, "object": "chat.completion.chunk",
+                "model": _model_name, "created": int(time.time()),
+                "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}]
+            })
+
+            total_tokens = 0
+            for token_text in _engine.generate(
+                prompt=None, chat_messages=messages,
+                max_tokens=max_tokens, temperature=temperature,
+            ):
+                total_tokens += 1
+                self._send_sse({
+                    "id": request_id, "object": "chat.completion.chunk",
+                    "model": _model_name, "created": int(time.time()),
+                    "choices": [{"index": 0, "delta": {"content": token_text}, "finish_reason": None}]
+                })
+
+            self._send_sse({
+                "id": request_id, "object": "chat.completion.chunk",
+                "model": _model_name, "created": int(time.time()),
+                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]
+            })
+            self.wfile.write(b"data: [DONE]\n\n")
+            self.wfile.flush()
+
+        else:
+            t0 = time.time()
+            full_text = ""
+            token_count = 0
+            for token_text in _engine.generate(
+                prompt=None, chat_messages=messages,
+                max_tokens=max_tokens, temperature=temperature,
+            ):
+                full_text += token_text
+                token_count += 1
+            elapsed = time.time() - t0
+
+            self._send_json({
+                "id": request_id, "object": "chat.completion",
+                "model": _model_name, "created": int(time.time()),
+                "choices": [{
+                    "index": 0,
+                    "message": {"role": "assistant", "content": full_text},
+                    "finish_reason": "stop",
+                }],
+                "usage": {
+                    "prompt_tokens": 0,
+                    "completion_tokens": token_count,
+                    "total_tokens": token_count,
+                },
+                "timings": {
+                    "total_ms": elapsed * 1000,
+                    "tokens_per_second": token_count / elapsed if elapsed > 0 else 0,
+                }
+            })
+
+
+def run_server(sniper_dir: str, host: str = "0.0.0.0", port: int = 8899,
+               cache_size: int = 2000):
+    """Start the OpenAI-compatible server."""
+    global _engine, _model_name
+
+    print("=" * 50)
+    print("  mlx-expert-sniper: server")
+    print("=" * 50)
+    print("\n  Loading...", end="", flush=True)
+
+    _engine = SniperEngine.from_dir(sniper_dir, max_cached_experts=cache_size)
+    _model_name = f"sniper-{_engine.config.model_type}"
+    print(" ready.")
+
+    print(f"\n  Server:  http://{host}:{port}")
+    print(f"  API:     http://localhost:{port}/v1/chat/completions")
+    print(f"  Models:  http://localhost:{port}/v1/models")
+    print(f"  Health:  http://localhost:{port}/health")
+    print(f"\n{'=' * 50}")
+    print("  Listening...\n")
+
+    server = HTTPServer((host, port), _Handler)
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        print("\n  Shutting down...")
+        server.shutdown()
+        _engine.close()

src/mlx_expert_sniper/sniper.py

@@ -0,0 +1,305 @@
+"""
+SniperEngine: the core inference engine for expert sniping.
+
+Implements the proven forward pass from research/expert-sniper/:
+  1. Attention (pinned weights, always in RAM)
+  2. Router → mx.eval() → top-k expert indices
+  3. Cache lookup → pread misses from SSD
+  4. gather_qmm for fused expert computation
+  5. Weighted sum of expert outputs
+
+Measured: 4.33 tok/s on Qwen3-30B-A3B (17.2 GB) with 0.87 GB RAM.
+"""
+
+import gc
+import json
+import os
+import time
+from pathlib import Path
+
+import numpy as np
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten
+
+from .config import SniperConfig
+from .cache import PerExpertCache, ExpertReader
+
+
+class SniperEngine:
+    """MoE expert sniping inference engine for MLX.
+
+    Usage:
+        engine = SniperEngine.from_dir("~/models/qwen3-30b")
+        for token in engine.generate("What is 2+2?"):
+            print(token, end="", flush=True)
+    """
+
+    def __init__(self, config: SniperConfig):
+        self.config = config
+        self.model = None
+        self.tokenizer = None
+        self.reader = None
+        self.expert_cache = None
+        self.kv_cache = None
+        self._loaded = False
+
+    @classmethod
+    def from_dir(cls, sniper_dir: str, **overrides) -> "SniperEngine":
+        """Create engine from a sniper directory."""
+        config = SniperConfig.from_dir(sniper_dir, **overrides)
+        engine = cls(config)
+        engine.load()
+        return engine
+
+    def load(self) -> float:
+        """Load model skeleton + pinned weights. Returns pinned size in GB."""
+        cfg = self.config
+
+        # Import model class based on model_type
+        model_cls, args_cls = self._get_model_classes()
+
+        with open(cfg.config_json_path) as f:
+            model_config = json.load(f)
+
+        # Build empty model skeleton
+        self.model = model_cls(args_cls.from_dict(model_config))
+
+        # Apply quantization matching stored format
+        quant = model_config.get("quantization", {})
+        self._quantize_model(model_config, quant)
+
+        # Load pinned (non-expert) weights
+        pinned = mx.load(cfg.pinned_path)
+        if hasattr(self.model, "sanitize"):
+            pinned = self.model.sanitize(pinned)
+        self.model.load_weights(list(pinned.items()), strict=False)
+
+        # Evaluate ONLY pinned params (skip switch_mlp expert weights)
+        params = [p for name, p in tree_flatten(self.model.parameters())
+                  if "switch_mlp" not in name]
+        mx.eval(*params)
+        pinned_gb = sum(p.nbytes for p in params) / 1e9
+        del pinned
+        gc.collect()
+
+        # Initialize I/O and cache
+        self.reader = ExpertReader(cfg.bin_dir, num_workers=cfg.num_io_workers)
+        self.expert_cache = PerExpertCache(max_experts=cfg.max_cached_experts)
+
+        # Load tokenizer — try sniper dir first, then HF name
+        from transformers import AutoTokenizer
+        tokenizer_path = cfg.sniper_dir
+        if not os.path.exists(os.path.join(tokenizer_path, "tokenizer.json")):
+            tokenizer_path = cfg.tokenizer_name or cfg.sniper_dir
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_path, trust_remote_code=True)
+
+        # Set MLX memory limits
+        mx.set_memory_limit(int(cfg.memory_limit_gb * 1024**3))
+        mx.set_cache_limit(cfg.cache_limit_mb * 1024**2)
+
+        self._loaded = True
+        return pinned_gb
+
+    def _get_model_classes(self):
+        """Get model and args classes for the model type."""
+        mt = self.config.model_type
+        if mt in ("qwen3_moe", "qwen2_moe"):
+            from mlx_lm.models.qwen3_moe import Model, ModelArgs
+            return Model, ModelArgs
+        raise ValueError(f"Unsupported model_type: {mt}. "
+                         f"Currently supported: qwen3_moe, qwen2_moe")
+
+    def _quantize_model(self, model_config: dict, quant: dict):
+        """Apply quantization matching the stored format."""
+        def class_predicate(path, module):
+            # Respect per-path quantization overrides from config
+            if path in model_config.get("quantization", {}):
+                return model_config["quantization"][path]
+            if not hasattr(module, "to_quantized"):
+                return False
+            return True
+
+        nn.quantize(
+            self.model,
+            group_size=quant.get("group_size", 64),
+            bits=quant.get("bits", 4),
+            mode=quant.get("mode", "affine"),
+            class_predicate=class_predicate,
+        )
+
+    def reset_kv_cache(self):
+        """Reset the KV cache for a new conversation."""
+        from mlx_lm.models.cache import make_prompt_cache
+        self.kv_cache = make_prompt_cache(self.model)
+
+    def forward_token(self, input_ids: mx.array) -> mx.array:
+        """Run one forward pass with expert sniping.
+
+        This is the proven forward pass: attention (pinned) → router →
+        mx.eval(indices) → cache/pread experts → gather_qmm → combine.
+        """
+        from mlx_lm.models.base import create_attention_mask
+
+        cfg = self.config
+        bits = cfg.bits
+        group_size = cfg.group_size
+
+        h = self.model.model.embed_tokens(input_ids)
+        mask = create_attention_mask(h, self.kv_cache[0])
+
+        for i, layer in enumerate(self.model.model.layers):
+            # ── Attention (pinned weights, always in RAM) ──
+            normed = layer.input_layernorm(h)
+            attn_out = layer.self_attn(normed, mask=mask, cache=self.kv_cache[i])
+            h = h + attn_out
+            mx.eval(h)  # Must eval before router (data-dependent)
+
+            # ── Router: compute expert scores ──
+            normed = layer.post_attention_layernorm(h)
+            gates = layer.mlp.gate(normed)
+            gates = mx.softmax(gates, axis=-1, precise=True)
+            k = layer.mlp.top_k
+            inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
+            scores = mx.take_along_axis(gates, inds, axis=-1)
+            if layer.mlp.norm_topk_prob:
+                scores = scores / scores.sum(axis=-1, keepdims=True)
+
+            # ── CRITICAL: eval indices before loading experts ──
+            mx.eval(inds, scores)
+
+            # ── Cache lookup + SSD read for misses ──
+            active = sorted(set(int(e) for e in np.array(inds).flatten()))
+            cached_experts = {}
+            misses = []
+            for eid in active:
+                c = self.expert_cache.get(i, eid)
+                if c is not None:
+                    cached_experts[eid] = c
+                else:
+                    misses.append(eid)
+
+            if misses:
+                loaded = self.reader.get_experts(i, misses)
+                for eid in misses:
+                    self.expert_cache.put(i, eid, loaded[eid])
+                    cached_experts[eid] = loaded[eid]
+                del loaded
+
+            # ── Remap global expert IDs → local [0..K-1] for gather_qmm ──
+            id_to_local = {eid: j for j, eid in enumerate(active)}
+            local_inds = mx.array(
+                np.vectorize(lambda x: id_to_local.get(int(x), 0))(np.array(inds)))
+
+            # ── Stack active experts into (K, ...) tensors ──
+            gw = mx.stack([cached_experts[e]["gate_proj.weight"] for e in active])
+            gs = mx.stack([cached_experts[e]["gate_proj.scales"] for e in active])
+            gb = mx.stack([cached_experts[e]["gate_proj.biases"] for e in active])
+            uw = mx.stack([cached_experts[e]["up_proj.weight"] for e in active])
+            us = mx.stack([cached_experts[e]["up_proj.scales"] for e in active])
+            ub = mx.stack([cached_experts[e]["up_proj.biases"] for e in active])
+            dw = mx.stack([cached_experts[e]["down_proj.weight"] for e in active])
+            ds = mx.stack([cached_experts[e]["down_proj.scales"] for e in active])
+            db = mx.stack([cached_experts[e]["down_proj.biases"] for e in active])
+            del cached_experts
+
+            # ── gather_qmm: fused quantized expert computation ──
+            x_exp = mx.expand_dims(normed, (-2, -3))
+
+            # SwiGLU: silu(gate_proj(x)) * up_proj(x)
+            go = mx.gather_qmm(x_exp, gw, scales=gs, biases=gb,
+                                rhs_indices=local_inds, transpose=True,
+                                group_size=group_size, bits=bits)
+            uo = mx.gather_qmm(x_exp, uw, scales=us, biases=ub,
+                                rhs_indices=local_inds, transpose=True,
+                                group_size=group_size, bits=bits)
+            hid = nn.silu(go) * uo
+
+            # down_proj
+            do = mx.gather_qmm(hid, dw, scales=ds, biases=db,
+                                rhs_indices=local_inds, transpose=True,
+                                group_size=group_size, bits=bits)
+
+            # Squeeze extra dims from gather_qmm output
+            while do.ndim > 4:
+                do = do.squeeze(-2)
+
+            # Weighted sum of expert outputs
+            h = h + (do * scores[..., None]).sum(axis=-2)
+            del gw, gs, gb, uw, us, ub, dw, ds, db
+
+        h = self.model.model.norm(h)
+        return self.model.lm_head(h)
+
+    def generate(self, prompt, max_tokens=None, temperature=None,
+                 chat_messages=None):
+        """Generate tokens, yielding each as a string.
+
+        Args:
+            prompt: Text prompt (used if chat_messages is None)
+            max_tokens: Maximum tokens to generate
+            temperature: Sampling temperature (0 = greedy)
+            chat_messages: List of {"role": ..., "content": ...} dicts
+
+        Yields:
+            str: Generated token text
+        """
+        if not self._loaded:
+            raise RuntimeError("Engine not loaded. Call load() first.")
+
+        max_tokens = max_tokens or self.config.max_tokens
+        temperature = temperature if temperature is not None else self.config.temperature
+
+        # Tokenize
+        if chat_messages:
+            text = self.tokenizer.apply_chat_template(
+                chat_messages, tokenize=False,
+                add_generation_prompt=True, enable_thinking=False)
+        else:
+            text = self.tokenizer.apply_chat_template(
+                [{"role": "user", "content": prompt}],
+                tokenize=False, add_generation_prompt=True,
+                enable_thinking=False)
+
+        tokens = self.tokenizer.encode(text)
+        input_ids = mx.array([tokens])
+
+        # Reset KV cache
+        self.reset_kv_cache()
+
+        # Prefill
+        logits = self.forward_token(input_ids)
+        mx.eval(logits)
+
+        # Sample first token
+        next_token = self._sample(logits[:, -1, :], temperature)
+
+        # Autoregressive generation
+        for _ in range(max_tokens):
+            if next_token in {151643, 151645}:  # EOS tokens
+                break
+            word = self.tokenizer.decode([next_token])
+            if "<|im_end|>" in word or "<|endoftext|>" in word:
+                break
+            yield word
+
+            input_ids = mx.array([[next_token]])
+            logits = self.forward_token(input_ids)
+            mx.eval(logits)
+            next_token = self._sample(logits[:, -1, :], temperature)
+
+    def _sample(self, logits: mx.array, temperature: float) -> int:
+        """Sample next token from logits."""
+        if temperature <= 0.01:
+            return int(mx.argmax(logits, axis=-1).item())
+        probs = mx.softmax(logits / temperature, axis=-1)
+        return int(mx.random.categorical(mx.log(probs + 1e-10)).item())
+
+    def close(self):
+        """Clean up resources."""
+        if self.reader:
+            self.reader.close()
+        self.expert_cache = None
+        self.model = None
+        self.tokenizer = None

stream_preprocess.py

@@ -0,0 +1,236 @@
+#!/usr/bin/env python3
+"""
+Stream-preprocess: download one shard at a time, process it, delete it.
+Avoids needing 17 GB free for the full model download.
+
+Usage:
+    python3 stream_preprocess.py
+"""
+
+import os
+import sys
+import json
+import time
+import gc
+import shutil
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))
+
+import numpy as np
+import mlx.core as mx
+
+REPO = "mlx-community/Qwen3-30B-A3B-4bit"
+OUTPUT_DIR = os.path.expanduser("~/models/qwen3-30b")
+PAGE_SIZE = 16384
+
+TENSOR_NAMES = [
+    "gate_proj.weight", "gate_proj.scales", "gate_proj.biases",
+    "up_proj.weight", "up_proj.scales", "up_proj.biases",
+    "down_proj.weight", "down_proj.scales", "down_proj.biases",
+]
+
+
+def convert_layer_to_bin(layer_data, layer_idx, num_experts, output_dir):
+    tensor_info = {}
+    expert_block_size = 0
+    for name in TENSOR_NAMES:
+        t = layer_data[name]
+        per_expert_shape = list(t.shape[1:])
+        if t.dtype == mx.uint32:
+            elem_size = 4
+        elif t.dtype in (mx.bfloat16, mx.float16):
+            elem_size = 2
+        else:
+            elem_size = 4
+        nbytes = 1
+        for s in per_expert_shape:
+            nbytes *= s
+        nbytes *= elem_size
+        tensor_info[name] = {
+            "shape_per_expert": per_expert_shape,
+            "dtype": str(t.dtype).replace("mlx.core.", ""),
+            "nbytes": nbytes,
+            "inner_offset": expert_block_size,
+        }
+        expert_block_size += nbytes
+
+    header = {
+        "layer_idx": layer_idx,
+        "num_experts": num_experts,
+        "layout": {
+            "expert_block_size": expert_block_size,
+            "data_start": PAGE_SIZE,
+            "tensors": tensor_info,
+        }
+    }
+    header_bytes = json.dumps(header, indent=2).encode()
+    assert len(header_bytes) < PAGE_SIZE
+    header_bytes += b"\x00" * (PAGE_SIZE - len(header_bytes))
+
+    out_path = os.path.join(output_dir, "bin", f"moe_layer_{layer_idx:02d}.bin")
+    with open(out_path, "wb") as f:
+        f.write(header_bytes)
+        for expert_id in range(num_experts):
+            for name in TENSOR_NAMES:
+                t = layer_data[name][expert_id]
+                if t.dtype == mx.bfloat16:
+                    raw = np.array(t.astype(mx.float16)).astype(np.float16).tobytes()
+                elif t.dtype == mx.uint32:
+                    raw = np.array(t).astype(np.uint32).tobytes()
+                else:
+                    raw = np.array(t).tobytes()
+                f.write(raw)
+
+    return os.path.getsize(out_path)
+
+
+def main():
+    from huggingface_hub import hf_hub_download, HfApi
+
+    print("=" * 55)
+    print("  Stream Preprocess — one shard at a time")
+    print(f"  Model:  {REPO}")
+    print(f"  Output: {OUTPUT_DIR}")
+    print("=" * 55)
+
+    os.makedirs(os.path.join(OUTPUT_DIR, "bin"), exist_ok=True)
+
+    # Download config + tokenizer files (small)
+    for fname in ["config.json", "tokenizer.json", "tokenizer_config.json",
+                   "special_tokens_map.json"]:
+        try:
+            path = hf_hub_download(REPO, fname, local_dir="/tmp/sniper_dl")
+            shutil.copy(path, os.path.join(OUTPUT_DIR, fname))
+            print(f"  Downloaded {fname}")
+        except Exception as e:
+            print(f"  Skipped {fname}: {e}")
+
+    # Get shard list
+    with open(os.path.join(OUTPUT_DIR, "config.json")) as f:
+        config = json.load(f)
+    num_layers = config.get("num_hidden_layers", 48)
+
+    # Download the index to find shard names
+    idx_path = hf_hub_download(REPO, "model.safetensors.index.json",
+                                local_dir="/tmp/sniper_dl")
+    with open(idx_path) as f:
+        idx = json.load(f)
+    shards = sorted(set(idx["weight_map"].values()))
+    print(f"\n  {len(shards)} shards to process")
+
+    # Check which layers already done
+    existing = set()
+    for f in os.listdir(os.path.join(OUTPUT_DIR, "bin")):
+        if f.startswith("moe_layer_") and f.endswith(".bin"):
+            existing.add(int(f.split("_")[2].split(".")[0]))
+    if existing:
+        print(f"  Already done: layers {sorted(existing)}")
+
+    pinned = {}
+    layers_done = set(existing)
+
+    # Track partial layers that span shards
+    partial_layers = {}
+
+    for si, shard_name in enumerate(shards):
+        print(f"\n  [{si+1}/{len(shards)}] Downloading {shard_name}...")
+        t0 = time.time()
+
+        shard_path = hf_hub_download(REPO, shard_name, local_dir="/tmp/sniper_dl")
+        dl_time = time.time() - t0
+        shard_size = os.path.getsize(shard_path) / 1e9
+        print(f"    Downloaded {shard_size:.1f} GB in {dl_time:.0f}s")
+
+        print(f"    Loading tensors...")
+        data = mx.load(shard_path)
+        print(f"    {len(data)} tensors")
+
+        # Classify
+        layer_experts = {}
+        for key, tensor in data.items():
+            if "switch_mlp" in key:
+                layer = int(key.split(".layers.")[1].split(".")[0])
+                short = key.split(".switch_mlp.")[1]
+                layer_experts.setdefault(layer, {})[short] = tensor
+            else:
+                pinned[key] = tensor
+
+        # Convert complete expert layers
+        for layer_idx, tensors in layer_experts.items():
+            if layer_idx in layers_done:
+                continue
+
+            # Merge with partial data from previous shards
+            if layer_idx in partial_layers:
+                partial_layers[layer_idx].update(tensors)
+                tensors = partial_layers[layer_idx]
+
+            if len(tensors) < 9:
+                # Partial — save for later
+                partial_layers[layer_idx] = tensors
+                print(f"    Layer {layer_idx}: partial ({len(tensors)}/9 tensors)")
+                continue
+
+            num_experts = tensors[list(tensors.keys())[0]].shape[0]
+            sz = convert_layer_to_bin(tensors, layer_idx, num_experts, OUTPUT_DIR)
+            layers_done.add(layer_idx)
+            if layer_idx in partial_layers:
+                del partial_layers[layer_idx]
+            print(f"    Layer {layer_idx}: {sz/1e6:.0f} MB")
+
+        del data, layer_experts
+        gc.collect()
+        mx.clear_cache()
+
+        # Delete the downloaded shard to free disk
+        try:
+            os.remove(shard_path)
+            print(f"    Deleted shard ({shard_size:.1f} GB freed)")
+        except:
+            pass
+
+    # Handle remaining partial layers
+    if partial_layers:
+        print(f"\n  {len(partial_layers)} partial layers remain — re-downloading...")
+        for layer_idx, tensors in partial_layers.items():
+            if layer_idx in layers_done:
+                continue
+            if len(tensors) >= 9:
+                num_experts = tensors[list(tensors.keys())[0]].shape[0]
+                sz = convert_layer_to_bin(tensors, layer_idx, num_experts, OUTPUT_DIR)
+                layers_done.add(layer_idx)
+                print(f"    Layer {layer_idx}: {sz/1e6:.0f} MB (merged)")
+
+    # Save pinned if we don't already have it or if it's stale
+    pinned_path = os.path.join(OUTPUT_DIR, "pinned.safetensors")
+    if pinned:
+        print(f"\n  Saving pinned ({len(pinned)} tensors)...")
+        mx.save_safetensors(pinned_path, pinned)
+        psz = os.path.getsize(pinned_path) / 1e9
+        print(f"    Pinned: {psz:.2f} GB")
+    else:
+        psz = os.path.getsize(pinned_path) / 1e9 if os.path.exists(pinned_path) else 0
+
+    # Clean up temp downloads
+    shutil.rmtree("/tmp/sniper_dl", ignore_errors=True)
+
+    # Summary
+    import glob
+    bin_files = sorted(glob.glob(os.path.join(OUTPUT_DIR, "bin", "moe_layer_*.bin")))
+    total = sum(os.path.getsize(f) for f in bin_files)
+    print(f"\n  Expert layers: {len(bin_files)}/{num_layers}")
+    print(f"  Expert total:  {total/1e9:.2f} GB")
+    print(f"  Pinned:        {psz:.2f} GB")
+    print(f"  Total:         {(total/1e9 + psz):.2f} GB")
+
+    missing = set(range(num_layers)) - layers_done
+    if missing:
+        print(f"\n  WARNING: Missing layers: {sorted(missing)}")
+    else:
+        print(f"\n  All {num_layers} layers converted!")
+        print(f"\n  Test with:")
+        print(f"    mlx-sniper run {OUTPUT_DIR} -p 'What is 2+2?' -v")
+
+
+if __name__ == "__main__":
+    main()