1. Expert Double-Buffered(TPS x2.3) 2. Remove KV cache 3. Support Local Model Path

.python-version

@@ -1 +1 @@
-3.13
+3.12

AGENTS.md

@@ -214,3 +214,86 @@ Limits:
 - No defensive checks for impossible states.
 - No silent fallbacks that hide failures.
 - Remove dead code paths cleanly; no compatibility shims for deleted behavior.
+
+### Python Style Guidelines
+
+**Imports** (ordered by PEP 8 with custom groupings):
+1. Standard library (`asyncio`, `json`, `pathlib`, `re`, `time`, `uuid`)
+2. Third-party packages (`mlx`, `mlx.nn`, `mlx_lm`, `numpy`, `pytest`, `starlette`, `uvicorn`)
+3. Local imports (`from mlx_moe.lazy_experts import ...`)
+
+Use absolute imports for the package (`from mlx_moe.lazy_experts.modules import ...`), not relative (`from .modules import ...`).
+
+**Formatting**:
+- Maximum line length: 100 characters
+- Use 4 spaces for indentation (no tabs)
+- Use hanging indents for long function signatures
+- Put imports in a single block (no multiple `import` statements scattered)
+
+**Types**:
+- Use Python 3.12+ type hints: `def foo(x: int) -> str:`
+- Use `X | None` instead of `Optional[X]`
+- Use `dict[str, int]` instead of `Dict[str, int]`
+- For internal/private functions, type hints are optional but encouraged
+- Use `Path` from `pathlib` for file paths
+
+**Naming Conventions**:
+- `snake_case` for functions, methods, and variables
+- `SCREAMING_SNAKE_CASE` for constants
+- `PascalCase` for classes
+- Leading underscore (`_func`) for private functions
+- Double leading underscore (`__method`) for name mangling (use sparingly)
+- Suffix `_cb` for callback functions
+- Suffix `_map` for dict-based mappings
+- Prefix `num_` or `n_` for counts
+
+**Error Handling**:
+- Raise specific exceptions with clear messages
+- Do not catch bare `Exception` unless re-raising or logging
+- Use `assert` for internal invariants, not for runtime validation
+- Fail fast with descriptive errors, not silent fallbacks
+
+**MLX-Specific**:
+- Never call `mx.eval()` in the forward pass of predictive modules (critical for performance)
+- Use `mx.stop_gradient()` where needed to prevent unwanted gradient flow
+- Prefer in-place operations when possible to reduce memory allocation
+- Remember MLX GPU eval is not thread-safe; serialization in server is intentional
+
+## Testing
+
+**Run all tests**:
+```bash
+uv run pytest
+```
+
+**Run a single test file**:
+```bash
+uv run pytest tests/test_unit_core.py
+```
+
+**Run a single test class**:
+```bash
+uv run pytest tests/test_unit_core.py::TestExpertCache
+```
+
+**Run a single test function**:
+```bash
+uv run pytest tests/test_unit_core.py::TestExpertCache::test_put_and_lookup -v
+```
+
+**Run tests matching a pattern**:
+```bash
+uv run pytest -k "test_lcp"
+```
+
+**Run with output capture disabled** (see print statements):
+```bash
+uv run pytest -s
+```
+
+**Run integration tests only**:
+```bash
+uv run pytest tests/test_integration.py
+```
+
+Note: Some tests in `test_unit_core.py` use synthetic mocks and don't require model files. Integration tests may require a real model to be available.

README.md

@@ -227,25 +227,45 @@ session.close()                          # release model
 
 For full control over the loading pipeline:
 
-```python
+```bash
+uv run python -c "
 import mlx.core as mx
 import mlx_lm
-from mlx_lm.utils import hf_repo_to_path
-from mlx_moe.lazy_experts import (
-    enable_lazy_experts, upgrade_to_predictive, get_fallback_stats
-)
-
-model, tokenizer = mlx_lm.load("mlx-community/Qwen3-Coder-Next-4bit", lazy=True)
-model_path = hf_repo_to_path("mlx-community/Qwen3-Coder-Next-4bit")
-
+from mlx_moe.lazy_experts import enable_lazy_experts, upgrade_to_predictive, get_fallback_stats
+from mlx_moe.lazy_experts.loading import _find_switch_mlp
+from mlx_moe.lazy_experts.discovery import router_only_discovery
+from mlx_moe.lazy_experts.core import dynamic_cache_update, enable_skip_fallback
+model_path = '/Users/steven/.lmstudio/models/lmstudio-community/Qwen3-Coder-Next-MLX-4bit'
+print('Loading model...')
+model, tokenizer = mlx_lm.load(model_path, lazy=True)
+print('Model loaded')
+print('Enable cached capacity=208...')
 enable_lazy_experts(model, model_path, cache_capacity_per_layer=208, predictive=True)
 mx.eval(model.parameters())
-
-mlx_lm.generate(model, tokenizer, prompt="Hello", max_tokens=10, verbose=False)
+switch, _ = _find_switch_mlp(model.layers[0], 0)
+print('Module type:', type(switch.gate_proj).__name__)
+print('=== Router-only discovery (does NOT load experts) ===')
+router_only_discovery(model, tokenizer, 'git pull origin main', max_tokens=10)
+print('Discovery done')
+print('Upgrade to predictive...')
 upgrade_to_predictive(model, model_path, 208)
-
-output = mlx_lm.generate(model, tokenizer, prompt="Write a Flask server",
-                          max_tokens=200, verbose=False)
+switch, _ = _find_switch_mlp(model.layers[0], 0)
+print('After upgrade:', type(switch.gate_proj).__name__)
+print('Enable skip fallback...')
+enable_skip_fallback(model)
+print('=== Hybrid warmup (10 tokens) ===')
+for resp in mlx_lm.stream_generate(model, tokenizer, prompt='git pull origin main', max_tokens=10):
+    pass
+dynamic_cache_update(model, max_layer_updates=48)
+print('Hybrid done')
+print()
+print('=== First generate ===')
+output = mlx_lm.generate(model, tokenizer, prompt='git pull origin main', max_tokens=50, verbose=False)
+print('Output:', repr(output[:150]))
+stats = get_fallback_stats(model)
+rate = stats['fallback_rate'] * 100
+print('Fallback rate:', rate, '%')
+"
 ```
 
 ## Testing

docs/server.md

@@ -50,7 +50,7 @@ flowchart TD
     G --> F
     F --> H["stream generation"]
     H --> I{"completed successfully?"}
-    I -- yes --> J["store (prompt_tokens + generated, prompt_cache) in keyed LRU"]
+    I -- yes --> J["store (prompt_tokens, prompt_cache) in keyed LRU"]
     I -- no --> K["do not restore cache entry"]
 ```
 

mlx_moe/cli.py

@@ -21,8 +21,6 @@ def main():
                        help="Max input tokens — rejects requests over this (default: 16384)")
     serve.add_argument("--kv-bits", type=int, default=None,
                        help="Quantize KV cache to N bits (8 recommended). Saves ~45%% KV memory.")
-    serve.add_argument("--kv-cache-slots", type=int, default=1,
-                       help="Number of keyed KV cache entries to keep (default: 1).")
     serve.add_argument("--shutdown-timeout", type=int, default=5,
                        help="Graceful shutdown timeout in seconds before cancelling active requests (default: 5).")
     serve.add_argument("--warmup", choices=["hybrid", "full", "none"], default="hybrid",
@@ -44,7 +42,6 @@ def main():
             max_input_tokens=args.max_input_tokens,
             kv_bits=args.kv_bits,
             warmup=args.warmup,
-            kv_cache_slots=args.kv_cache_slots,
             shutdown_timeout=args.shutdown_timeout,
         )
     else:

mlx_moe/lazy_experts/core.py

@@ -637,6 +637,30 @@ def dynamic_cache_update(model, max_layer_updates: int = 12) -> list[dict]:
 
         layer_stats["layer"] = i
         stats.append(layer_stats)
+
+    # Prefetch: trigger next layer prefetch based on current layer requests
+    # This runs after current layer updates to predict next layer needs
+    for i in range(len(model.layers) - 1):
+        curr_proj = getattr(model.layers[i].mlp.switch_mlp, "up_proj", None)
+        if not curr_proj or not isinstance(curr_proj, PredictiveCachedSwitchLinear):
+            continue
+
+        next_layer = model.layers[i + 1]
+        next_switch, _ = _find_switch_mlp(next_layer, i + 1)
+        if next_switch is None:
+            continue
+
+        next_proj = getattr(next_switch, "up_proj", None)
+        if not next_proj or not isinstance(next_proj, PredictiveCachedSwitchLinear):
+            continue
+
+        # Use current layer requests to predict next layer
+        predicted_experts = curr_proj._cache.predict_next_experts()
+        if predicted_experts:
+            next_proj._cache.prefetch_async(predicted_experts)
+            # Check if we can swap buffers for next layer
+            next_proj._cache.check_prefetch_and_swap()
+
     return stats
 
 

mlx_moe/lazy_experts/generate.py

@@ -31,6 +31,29 @@
 _WARMUP_CACHE = 256 * 1024 * 1024
 
 
+def _fix_group_size(model):
+    """Fix group_size mismatch by inferring from weight/scales shapes.
+
+    Some models (e.g., Qwen3-Coder-Next-MLX-8bit) have incorrect group_size
+    in config.json. The actual group_size must be computed from tensor shapes.
+    """
+    for layer in model.layers:
+        switch, _ = _find_switch_mlp(layer)
+        if switch is None:
+            continue
+        for name in ("gate_proj", "up_proj", "down_proj"):
+            proj = getattr(switch, name, None)
+            if proj is None or not hasattr(proj, "weight") or not hasattr(proj, "scales"):
+                continue
+            # Infer actual group_size from shapes
+            # weight shape: (num_experts, output_dim, input_dim)
+            # scales shape: (num_experts, output_dim, num_groups)
+            # group_size = input_dim / num_groups
+            actual = proj.weight.shape[-1] // proj.scales.shape[-1]
+            if actual != proj.group_size:
+                proj.group_size = actual
+
+
 def _startup(
     model_name,
     prompt,
@@ -54,12 +77,66 @@ def _startup(
     from .core import upgrade_to_predictive
     import mlx_lm as _mlx_lm
     from mlx_lm.utils import hf_repo_to_path
+    from huggingface_hub.errors import HFValidationError
 
     t_total_start = time.perf_counter()
 
-    model_path = hf_repo_to_path(model_name)
+    # First check if the path exists at all
+    if os.path.exists(model_name):
+        if not os.path.isdir(model_name):
+            raise FileNotFoundError(
+                f"Model path exists but is not a directory: {model_name}"
+            )
+        model_path = Path(model_name)
+    elif "/" not in model_name:
+        # Local path without slash - try as relative path
+        raise FileNotFoundError(
+            f"Model path does not exist: {model_name}\n"
+            f"Use a valid local path or HuggingFace model ID (e.g., mlx-community/Qwen3-Coder-Next-4bit)"
+        )
+    else:
+        # HuggingFace model ID - verify it looks like one (has namespace/repo format)
+        # and doesn't look like a file path
+        if model_name.startswith("/") or model_name.startswith("."):
+            raise FileNotFoundError(
+                f"Model path does not exist: {model_name}\n"
+                f"Use a valid local path or HuggingFace model ID (e.g., mlx-community/Qwen3-Coder-Next-4bit)"
+            )
+        try:
+            model_path = hf_repo_to_path(model_name)
+        except HFValidationError:
+            raise FileNotFoundError(
+                f"Invalid model ID or model not found: {model_name}\n"
+                f"HuggingFace model IDs should be in the format 'namespace/repo_name'"
+            )
+        except Exception as e:
+            raise FileNotFoundError(
+                f"Failed to download model: {model_name}\n{str(e)}"
+            )
+
+    if not model_path.is_dir():
+        raise FileNotFoundError(f"Model path does not exist: {model_path}")
+
+    config_path = model_path / "config.json"
+    if not config_path.is_file():
+        safetensors_files = list(model_path.glob("*.safetensors"))
+        if not safetensors_files:
+            raise FileNotFoundError(
+                f"Model directory is empty or missing model files: {model_path}\n"
+                f"Please download the model first."
+            )
+        else:
+            raise FileNotFoundError(
+                f"config.json not found in: {model_path}\n"
+                f"This may not be a valid HuggingFace-compatible model directory."
+            )
+
     t0 = time.perf_counter()
-    model, tokenizer = _mlx_lm.load(model_name, lazy=True)
+    model, tokenizer = _mlx_lm.load(str(model_path), lazy=True)
+
+    # Fix group_size mismatch: config.json may declare wrong group_size
+    # Actual group_size must be inferred from weight/scales shapes
+    # _fix_group_size(model)
 
     # Detect MoE architecture
     num_moe_layers = 0
@@ -120,9 +197,17 @@ def _startup(
     if cache_dir is not None:
         cache_dir = os.path.expanduser(cache_dir)
         os.makedirs(cache_dir, exist_ok=True)
-        safe_name = model_name.replace("/", "--")
-        cache_path = os.path.join(cache_dir, f"{safe_name}.json")
+        model_path_str = str(model_path).rstrip("/")
+        safe_name_base = model_path_str.replace("/", "--").replace("\\", "--")
+
+        cache_path = os.path.join(cache_dir, f"{safe_name_base}.json")
         prepacked_path = cache_path.replace(".json", ".weights.safetensors")
+        meta_path = prepacked_path + ".meta.json"
+
+        if not os.path.exists(meta_path):
+            safe_name_alt = safe_name_base + "--"
+            cache_path = os.path.join(cache_dir, f"{safe_name_alt}.json")
+            prepacked_path = cache_path.replace(".json", ".weights.safetensors")
 
     used_saved_state = False
     mx.reset_peak_memory()