v0.2.0: Add Qwen3.5-35B-A3B support (5.78 tok/s, 19.5 GB on 16 GB RAM)

README.md

@@ -1,17 +1,16 @@
 # CLI Agent — `mlx-expert-sniper`
 
 Pip-installable CLI that wraps the Expert Sniper research into a production tool.
+Run MoE models larger than your RAM on Apple Silicon.
 
 ## 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 |
+| Model | Size | Standard mlx_lm | Sniper tok/s | RAM pinned | Cache hit |
+|-------|------|-----------------|--------------|------------|-----------|
+| Qwen3-30B-A3B | 17.2 GB | OOM | **4.33 tok/s** | 0.87 GB | 88.5% |
+| Qwen3.5-35B-A3B | 19.5 GB | OOM | **5.78 tok/s** | 1.38 GB | 75% |
+
+Both models exceed 16 GB RAM. Both produce coherent multi-paragraph output.
 
 ## Install
 

src/mlx_expert_sniper/config.py

@@ -52,19 +52,21 @@ class SniperConfig:
         with open(config_path) as f:
             model_config = json.load(f)
 
-        quant = model_config.get("quantization", {})
+        # Handle nested configs (qwen3_5_moe has text_config)
+        quant = model_config.get("quantization", model_config.get("quantization_config", {}))
+        text_cfg = model_config.get("text_config", model_config)
 
         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),
+            num_hidden_layers=text_cfg.get("num_hidden_layers", 0),
+            num_experts=text_cfg.get("num_experts", 0),
+            num_experts_per_tok=text_cfg.get("num_experts_per_tok", 0),
+            hidden_size=text_cfg.get("hidden_size", 0),
+            moe_intermediate_size=text_cfg.get("moe_intermediate_size", 0),
+            vocab_size=text_cfg.get("vocab_size", 0),
+            norm_topk_prob=text_cfg.get("norm_topk_prob", True),
             model_type=model_config.get("model_type", ""),
             tokenizer_name=model_config.get("_name_or_path", ""),
         )

src/mlx_expert_sniper/preprocess.py

@@ -163,6 +163,16 @@ def preprocess_model(model_dir: str, output_dir: str, verbose: bool = True):
                 layer = int(key.split(".layers.")[1].split(".")[0])
                 short = key.split(".switch_mlp.")[1]
                 layer_experts.setdefault(layer, {})[short] = tensor
+            elif "experts.gate_up_proj" in key:
+                # qwen3_5_moe: fused gate+up proj — split into separate tensors
+                layer = int(key.split(".layers.")[1].split(".")[0])
+                gate_up = tensor
+                mid = gate_up.shape[-2] // 2
+                layer_experts.setdefault(layer, {})["gate_proj.weight"] = gate_up[..., :mid, :]
+                layer_experts.setdefault(layer, {})["up_proj.weight"] = gate_up[..., mid:, :]
+            elif "experts.down_proj" in key:
+                layer = int(key.split(".layers.")[1].split(".")[0])
+                layer_experts.setdefault(layer, {})["down_proj.weight"] = tensor
             else:
                 pinned[key] = tensor
 

src/mlx_expert_sniper/sniper.py

@@ -108,8 +108,11 @@ class SniperEngine:
         if mt in ("qwen3_moe", "qwen2_moe"):
             from mlx_lm.models.qwen3_moe import Model, ModelArgs
             return Model, ModelArgs
+        if mt in ("qwen3_5_moe",):
+            from mlx_lm.models.qwen3_5_moe import Model, ModelArgs
+            return Model, ModelArgs
         raise ValueError(f"Unsupported model_type: {mt}. "
-                         f"Currently supported: qwen3_moe, qwen2_moe")
+                         f"Currently supported: qwen3_moe, qwen3_5_moe")
 
     def _quantize_model(self, model_config: dict, quant: dict):
         """Apply quantization matching the stored format."""
@@ -129,35 +132,87 @@ class SniperEngine:
             class_predicate=class_predicate,
         )
 
+    def _get_layers_and_head(self):
+        """Get (layers, embed_tokens, norm, lm_head) for the model architecture."""
+        mt = self.config.model_type
+        if mt in ("qwen3_5_moe",):
+            lm = self.model.language_model
+            return lm.model.layers, lm.model.embed_tokens, lm.model.norm, lm.lm_head
+        return self.model.model.layers, self.model.model.embed_tokens, self.model.model.norm, self.model.lm_head
+
+    def _is_moe_layer(self, layer) -> bool:
+        """Check if a layer has a MoE MLP (vs dense MLP)."""
+        return hasattr(layer.mlp, "gate") and hasattr(layer.mlp, "switch_mlp")
+
+    def _has_shared_expert(self, layer) -> bool:
+        """Check if the MoE block has a shared expert."""
+        return hasattr(layer.mlp, "shared_expert")
+
     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)
+        if hasattr(self.model, "make_cache"):
+            self.kv_cache = self.model.make_cache()
+        else:
+            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.
+        Supports both qwen3_moe (standard attention) and qwen3_5_moe
+        (hybrid linear/full attention with shared experts).
         """
         from mlx_lm.models.base import create_attention_mask
 
         cfg = self.config
         bits = cfg.bits
         group_size = cfg.group_size
+        is_hybrid = cfg.model_type in ("qwen3_5_moe",)
+
+        layers, embed_tokens, norm, lm_head = self._get_layers_and_head()
+        h = embed_tokens(input_ids)
+
+        # Create masks
+        if is_hybrid:
+            from mlx_lm.models.base import create_ssm_mask
+            # Find first full-attention and first linear-attention layer cache
+            fa_cache = None
+            ssm_cache = None
+            for li, layer in enumerate(layers):
+                if hasattr(layer, "is_linear"):
+                    if layer.is_linear and ssm_cache is None:
+                        ssm_cache = self.kv_cache[li]
+                    elif not layer.is_linear and fa_cache is None:
+                        fa_cache = self.kv_cache[li]
+            fa_mask = create_attention_mask(h, fa_cache) if fa_cache is not None else None
+            ssm_mask = create_ssm_mask(h, ssm_cache) if ssm_cache is not None else None
+        else:
+            mask = create_attention_mask(h, self.kv_cache[0])
 
-        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):
+        for i, layer in enumerate(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])
+
+            if is_hybrid and hasattr(layer, "is_linear") and layer.is_linear:
+                attn_out = layer.linear_attn(normed, mask=ssm_mask, cache=self.kv_cache[i])
+            elif is_hybrid and hasattr(layer, "self_attn"):
+                attn_out = layer.self_attn(normed, mask=fa_mask, cache=self.kv_cache[i])
+            else:
+                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 ──
+            # ── Post-attention norm ──
             normed = layer.post_attention_layernorm(h)
+
+            # ── Check if this is an MoE layer ──
+            if not self._is_moe_layer(layer):
+                # Dense MLP — just run it (pinned weights)
+                h = h + layer.mlp(normed)
+                continue
+
+            # ── Router: compute expert scores ──
             gates = layer.mlp.gate(normed)
             gates = mx.softmax(gates, axis=-1, precise=True)
             k = layer.mlp.top_k
@@ -226,11 +281,19 @@ class SniperEngine:
                 do = do.squeeze(-2)
 
             # Weighted sum of expert outputs
-            h = h + (do * scores[..., None]).sum(axis=-2)
+            moe_out = (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)
+            # ── Shared expert (pinned, always runs) ──
+            if self._has_shared_expert(layer):
+                shared_out = layer.mlp.shared_expert(normed)
+                shared_out = mx.sigmoid(layer.mlp.shared_expert_gate(normed)) * shared_out
+                h = h + moe_out + shared_out
+            else:
+                h = h + moe_out
+
+        h = norm(h)
+        return lm_head(h)
 
     def generate(self, prompt, max_tokens=None, temperature=None,
                  chat_messages=None):
@@ -252,15 +315,27 @@ class SniperEngine:
         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)
+        messages = chat_messages if chat_messages else [{"role": "user", "content": prompt}]
+        if self.tokenizer.chat_template:
+            try:
+                text = self.tokenizer.apply_chat_template(
+                    messages, tokenize=False,
+                    add_generation_prompt=True, enable_thinking=False)
+            except TypeError:
+                text = self.tokenizer.apply_chat_template(
+                    messages, tokenize=False,
+                    add_generation_prompt=True)
         else:
-            text = self.tokenizer.apply_chat_template(
-                [{"role": "user", "content": prompt}],
-                tokenize=False, add_generation_prompt=True,
-                enable_thinking=False)
+            # Fallback: Qwen/ChatML format
+            parts = []
+            for m in messages:
+                parts.append(f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>")
+            parts.append("<|im_start|>assistant\n")
+            text = "\n".join(parts)
+
+        # Strip thinking from output by tracking state
+        self._in_thinking = False
+        self._thinking_buffer = ""
 
         tokens = self.tokenizer.encode(text)
         input_ids = mx.array([tokens])
@@ -275,14 +350,39 @@ class SniperEngine:
         # Sample first token
         next_token = self._sample(logits[:, -1, :], temperature)
 
+        # Build EOS set from tokenizer config
+        eos_ids = set()
+        if hasattr(self.tokenizer, "eos_token_id"):
+            eid = self.tokenizer.eos_token_id
+            if isinstance(eid, list):
+                eos_ids.update(eid)
+            elif eid is not None:
+                eos_ids.add(eid)
+        eos_ids.update({151643, 151645, 248044, 248046})  # Qwen3 + Qwen3.5 EOS
+
         # Autoregressive generation
         for _ in range(max_tokens):
-            if next_token in {151643, 151645}:  # EOS tokens
+            if next_token in eos_ids:
                 break
             word = self.tokenizer.decode([next_token])
             if "<|im_end|>" in word or "<|endoftext|>" in word:
                 break
-            yield word
+
+            # Filter out <think>...</think> blocks
+            self._thinking_buffer += word
+            if "<think>" in self._thinking_buffer:
+                self._in_thinking = True
+                self._thinking_buffer = ""
+            elif "</think>" in self._thinking_buffer:
+                self._in_thinking = False
+                # Yield anything after </think>
+                after = self._thinking_buffer.split("</think>", 1)[-1].lstrip("\n")
+                self._thinking_buffer = ""
+                if after:
+                    yield after
+            elif not self._in_thinking:
+                yield self._thinking_buffer
+                self._thinking_buffer = ""
 
             input_ids = mx.array([[next_token]])
             logits = self.forward_token(input_ids)

stream_preprocess_35b.py

@@ -0,0 +1,228 @@
+#!/usr/bin/env python3
+"""Stream-preprocess Qwen3.5-35B-A3B-4bit: download one shard, process, delete."""
+
+import os, sys, json, time, gc, shutil, glob
+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.5-35B-A3B-4bit"
+OUTPUT_DIR = os.path.expanduser("~/models/qwen35-35b")
+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:
+        if name not in layer_data:
+            continue
+        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:
+                if name not in layer_data:
+                    continue
+                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
+
+    print("=" * 55)
+    print("  Stream Preprocess Qwen3.5-35B-A3B-4bit")
+    print(f"  Output: {OUTPUT_DIR}")
+    print("=" * 55)
+
+    os.makedirs(os.path.join(OUTPUT_DIR, "bin"), exist_ok=True)
+
+    # Download config + tokenizer
+    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_35b")
+            shutil.copy(path, os.path.join(OUTPUT_DIR, fname))
+            print(f"  Downloaded {fname}")
+        except Exception as e:
+            print(f"  Skipped {fname}: {e}")
+
+    with open(os.path.join(OUTPUT_DIR, "config.json")) as f:
+        config = json.load(f)
+    text_cfg = config.get("text_config", config)
+    num_layers = text_cfg.get("num_hidden_layers", 40)
+    print(f"  Layers: {num_layers}, Experts: {text_cfg.get('num_experts', 0)}")
+
+    idx_path = hf_hub_download(REPO, "model.safetensors.index.json",
+                                local_dir="/tmp/sniper_dl_35b")
+    with open(idx_path) as f:
+        idx = json.load(f)
+    shards = sorted(set(idx["weight_map"].values()))
+    print(f"  {len(shards)} shards")
+
+    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: {sorted(existing)}")
+
+    pinned = {}
+    layers_done = set(existing)
+    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_35b")
+        dl_time = time.time() - t0
+        shard_size = os.path.getsize(shard_path) / 1e9
+        print(f"    {shard_size:.1f} GB in {dl_time:.0f}s")
+
+        data = mx.load(shard_path)
+        print(f"    {len(data)} tensors")
+
+        layer_experts = {}
+        for key, tensor in data.items():
+            # Skip vision tower
+            if "vision_tower" in key or "model.visual" in key:
+                continue
+
+            if "switch_mlp" in key and ".layers." in key:
+                layer = int(key.split(".layers.")[1].split(".")[0])
+                short = key.split(".switch_mlp.")[1]
+                layer_experts.setdefault(layer, {})[short] = tensor
+            elif "experts.gate_up_proj" in key and ".layers." in key:
+                # Fused gate+up — split
+                layer = int(key.split(".layers.")[1].split(".")[0])
+                gate_up = tensor
+                mid = gate_up.shape[-2] // 2
+                ld = layer_experts.setdefault(layer, {})
+                ld["gate_proj.weight"] = gate_up[..., :mid, :]
+                ld["up_proj.weight"] = gate_up[..., mid:, :]
+            elif "experts.down_proj" in key and ".layers." in key:
+                layer = int(key.split(".layers.")[1].split(".")[0])
+                layer_experts.setdefault(layer, {})["down_proj.weight"] = tensor
+            else:
+                pinned[key] = tensor
+
+        for layer_idx, tensors in layer_experts.items():
+            if layer_idx in layers_done:
+                continue
+            if layer_idx in partial_layers:
+                partial_layers[layer_idx].update(tensors)
+                tensors = partial_layers[layer_idx]
+
+            # Check how many tensor groups we have
+            # For quantized: need weight + scales + biases for each of gate/up/down = 9
+            # For non-quantized: just weight for gate/up/down = 3
+            n_keys = len(tensors)
+            has_all = all(f"{p}.weight" in tensors for p in ["gate_proj", "up_proj", "down_proj"])
+
+            if not has_all:
+                partial_layers[layer_idx] = tensors
+                print(f"    Layer {layer_idx}: partial ({n_keys} tensors)")
+                continue
+
+            num_experts = tensors["gate_proj.weight"].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 ({num_experts} experts)")
+
+        del data, layer_experts
+        gc.collect()
+        mx.clear_cache()
+
+        try:
+            os.remove(shard_path)
+            print(f"    Deleted shard ({shard_size:.1f} GB freed)")
+        except:
+            pass
+
+    # Handle remaining partials
+    for layer_idx, tensors in partial_layers.items():
+        if layer_idx in layers_done:
+            continue
+        has_all = all(f"{p}.weight" in tensors for p in ["gate_proj", "up_proj", "down_proj"])
+        if has_all:
+            num_experts = tensors["gate_proj.weight"].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 pinned:
+        print(f"\n  Saving pinned ({len(pinned)} tensors)...")
+        mx.save_safetensors(os.path.join(OUTPUT_DIR, "pinned.safetensors"), pinned)
+        psz = os.path.getsize(os.path.join(OUTPUT_DIR, "pinned.safetensors")) / 1e9
+        print(f"    Pinned: {psz:.2f} GB")
+    else:
+        psz = 0
+
+    shutil.rmtree("/tmp/sniper_dl_35b", ignore_errors=True)
+
+    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)
+    missing = set(range(num_layers)) - layers_done
+    print(f"\n  Expert layers: {len(bin_files)}/{num_layers}")
+    print(f"  Expert total:  {total/1e9:.2f} GB")
+    print(f"  Pinned:        {psz:.2f} GB")
+    if missing:
+        print(f"  WARNING: Missing layers: {sorted(missing)}")
+    else:
+        print(f"\n  All {num_layers} layers converted!")
+        print(f"  Test: mlx-sniper run {OUTPUT_DIR} -p 'What is 2+2?' -v")
+
+
+if __name__ == "__main__":
+    main()