Continue AirMLX port: generate integration and e2e tests

- Auto-split via ensure_streaming_layout in load_streaming
- --streaming flags on mlx_lm.generate CLI
- Fix StreamingModelWrapper recursion and layer weight keys
- Add benchmark CLI, get_stats, generate_step e2e tests

Author: SashimiSaketoro

mlx-lm-full-repo.md

@@ -256,6 +256,23 @@ def setup_arg_parser():
         default=42,
         help="Base seed for TurboQuant per-layer rotations.",
     )
+    parser.add_argument(
+        "--streaming",
+        action="store_true",
+        help="Enable layer-streaming inference for models larger than RAM.",
+    )
+    parser.add_argument(
+        "--max-memory-gb",
+        type=float,
+        default=20.0,
+        help="Memory budget (GB) for layer window when --streaming is set.",
+    )
+    parser.add_argument(
+        "--streaming-window",
+        type=int,
+        default=None,
+        help="Fixed layer window size (auto-computed if omitted).",
+    )
     return parser
 
 
@@ -2114,13 +2131,29 @@ def main():
             )
     model_path = model_path or DEFAULT_MODEL
 
-    model, tokenizer = load(
-        model_path,
-        adapter_path=args.adapter_path,
-        tokenizer_config=tokenizer_config,
-        model_config={"quantize_activations": args.quantize_activations},
-        trust_remote_code=args.trust_remote_code,
-    )
+    if args.streaming:
+        from mlx_lm.streaming import StreamingConfig, load_streaming
+
+        streaming_config = StreamingConfig(
+            max_memory_gb=args.max_memory_gb,
+            window_size=args.streaming_window,
+            verbose=args.verbose,
+        )
+        model, tokenizer, _ = load_streaming(
+            model_path,
+            streaming_config=streaming_config,
+            tokenizer_config=tokenizer_config,
+        )
+        if args.verbose:
+            print(model.get_stats())
+    else:
+        model, tokenizer = load(
+            model_path,
+            adapter_path=args.adapter_path,
+            tokenizer_config=tokenizer_config,
+            model_config={"quantize_activations": args.quantize_activations},
+            trust_remote_code=args.trust_remote_code,
+        )
     for eos_token in args.extra_eos_token:
         tokenizer.add_eos_token(eos_token)
 

mlx_lm/generate.py

@@ -5,13 +5,14 @@
 from .config import StreamingConfig
 from .layer_loader import RollingWindowLoader
 from .load import load_streaming
-from .split_model import split_model_by_layers
+from .split_model import ensure_streaming_layout, split_model_by_layers
 from .wrapper import StreamingModelWrapper
 
 __all__ = [
     "StreamingConfig",
     "RollingWindowLoader",
     "load_streaming",
     "split_model_by_layers",
+    "ensure_streaming_layout",
     "StreamingModelWrapper",
 ]

mlx_lm/streaming/__init__.py

@@ -0,0 +1,68 @@
+# Copyright © 2024 Apple Inc.
+
+"""Benchmark layer-streaming generation throughput."""
+
+import argparse
+import time
+
+from mlx_lm.generate import generate_step
+from mlx_lm.sample_utils import make_sampler
+
+from .config import StreamingConfig
+from .load import load_streaming
+
+
+def benchmark_streaming(
+    model_path: str,
+    prompt_tokens: list,
+    max_tokens: int = 32,
+    max_memory_gb: float = 20.0,
+    verbose: bool = True,
+) -> dict:
+    """Run a streaming benchmark and return timing stats."""
+    import mlx.core as mx
+
+    config = StreamingConfig(max_memory_gb=max_memory_gb, verbose=verbose)
+    model, _tokenizer, _cfg = load_streaming(model_path, streaming_config=config)
+
+    prompt = mx.array([prompt_tokens])
+    sampler = make_sampler(temp=0.0)
+
+    start = time.perf_counter()
+    count = 0
+    for _ in generate_step(prompt, model, max_tokens=max_tokens, sampler=sampler):
+        count += 1
+    elapsed = time.perf_counter() - start
+
+    stats = {
+        "tokens": count,
+        "elapsed_s": elapsed,
+        "tok_per_s": count / elapsed if elapsed > 0 else 0.0,
+        "ms_per_token": (elapsed / count * 1000) if count > 0 else 0.0,
+        "streaming": model.get_stats(),
+    }
+    if verbose:
+        print(f"Tokens: {stats['tokens']}")
+        print(f"Throughput: {stats['tok_per_s']:.2f} tok/s")
+        print(f"Window: {stats['streaming']['streaming']['window_size']} layers")
+    return stats
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Benchmark layer-streaming inference")
+    parser.add_argument("--model", required=True, help="Model path or HF repo")
+    parser.add_argument("--max-tokens", type=int, default=32)
+    parser.add_argument("--max-memory-gb", type=float, default=20.0)
+    parser.add_argument("--prompt-ids", default="1,2,3,4", help="Comma-separated token ids")
+    args = parser.parse_args()
+    prompt_tokens = [int(x) for x in args.prompt_ids.split(",")]
+    benchmark_streaming(
+        args.model,
+        prompt_tokens,
+        max_tokens=args.max_tokens,
+        max_memory_gb=args.max_memory_gb,
+    )
+
+
+if __name__ == "__main__":
+    main()

mlx_lm/streaming/benchmark.py

@@ -109,18 +109,9 @@ def _load_from_monolithic(self, layer_idx: int) -> Dict[str, mx.array]:
     def load_layer(self, layer_idx: int, prefetch: bool = False) -> LayerWeights:
         layer_file = self._layer_files.get(layer_idx)
         if layer_file is None:
-            raw = self._load_from_monolithic(layer_idx)
-            # Prefix keys for mlx-lm load_weights on TransformerBlock
-            weights = {
-                f"{self.layer_key_prefix}.{layer_idx}.{k}": v
-                for k, v in raw.items()
-            }
+            weights = self._load_from_monolithic(layer_idx)
         else:
-            local = mx.load(str(layer_file))
-            weights = {
-                f"{self.layer_key_prefix}.{layer_idx}.{k}": v
-                for k, v in local.items()
-            }
+            weights = mx.load(str(layer_file))
 
         layer_weights = LayerWeights(weights, layer_idx)
         if prefetch and weights:
@@ -151,7 +142,7 @@ def get_layer(self, layer_idx: int) -> LayerWeights:
             if (
                 layer_idx % self.streaming_config.clear_cache_every_n_layers == 0
             ):
-                mx.metal.clear_cache()
+                mx.clear_cache()
 
         layer_weights = self.load_layer(layer_idx, prefetch=False)
         mx.eval(list(layer_weights.weights.values()))
@@ -168,7 +159,17 @@ def get_layer(self, layer_idx: int) -> LayerWeights:
 
         return layer_weights
 
+    def get_memory_usage(self) -> dict:
+        total_bytes = sum(layer.memory_bytes for layer in self.layer_queue)
+        return {
+            "loaded_layers": len(self.loaded_layers),
+            "layer_indices": sorted(list(self.loaded_layers)),
+            "total_mb": total_bytes / 1e6,
+            "total_gb": total_bytes / 1e9,
+            "window_size": self.window_size,
+        }
+
     def clear(self):
         self.layer_queue.clear()
         self.loaded_layers.clear()
-        mx.metal.clear_cache()
+        mx.clear_cache()

mlx_lm/streaming/layer_loader.py

@@ -13,6 +13,7 @@
 
 from .config import StreamingConfig
 from .layer_loader import RollingWindowLoader
+from .split_model import ensure_streaming_layout
 from .wrapper import StreamingModelWrapper
 
 
@@ -44,7 +45,8 @@ def load_streaming(
     streaming_config: Optional[StreamingConfig] = None,
     tokenizer_config: Optional[dict] = None,
     revision: Optional[str] = None,
-) -> Tuple[StreamingModelWrapper, TokenizerWrapper, Dict]:
+    load_tokenizer: bool = True,
+) -> Tuple[StreamingModelWrapper, Optional[TokenizerWrapper], Dict]:
     """
     Load a model for layer-streaming inference.
 
@@ -67,22 +69,18 @@ def load_streaming(
     model_path = _download(path_or_hf_repo, revision=revision)
     config = load_config(model_path)
 
+    ensure_streaming_layout(model_path, verbose=streaming_config.verbose)
+
     model_class, model_args_class = _get_classes(config)
     model_args = model_args_class.from_dict(config)
     model = model_class(model_args)
 
     fixed_file = model_path / "fixed_weights.safetensors"
-    if fixed_file.exists():
-        fixed_weights = mx.load(str(fixed_file))
-        if hasattr(model, "sanitize"):
-            fixed_weights = model.sanitize(fixed_weights)
-        model.load_weights(list(fixed_weights.items()), strict=False)
-        mx.eval(model.parameters())
-    else:
-        raise FileNotFoundError(
-            f"{fixed_file} not found. Run: python -m mlx_lm.streaming.split_model "
-            f"<model.safetensors> {model_path}"
-        )
+    fixed_weights = mx.load(str(fixed_file))
+    if hasattr(model, "sanitize"):
+        fixed_weights = model.sanitize(fixed_weights)
+    model.load_weights(list(fixed_weights.items()), strict=False)
+    mx.eval(model.parameters())
 
     layer_size = streaming_config.estimate_layer_memory(
         hidden_size=config["hidden_size"],
@@ -106,5 +104,7 @@ def load_streaming(
     wrapped = StreamingModelWrapper(model, loader)
     wrapped.eval()
 
-    tokenizer = load_tokenizer(model_path, tokenizer_config)
+    tokenizer = (
+        load_tokenizer(model_path, tokenizer_config) if load_tokenizer else None
+    )
     return wrapped, tokenizer, config

mlx_lm/streaming/load.py

@@ -3,9 +3,10 @@
 """Split monolithic safetensors into per-layer files for streaming inference."""
 
 import argparse
+import glob
 import mlx.core as mx
 from pathlib import Path
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 
 def split_model_by_layers(
@@ -55,6 +56,34 @@ def split_model_by_layers(
     print("Split complete.")
 
 
+def ensure_streaming_layout(
+    model_dir: Path,
+    layer_prefix: str = "model.layers",
+    verbose: bool = False,
+) -> bool:
+    """
+    Ensure a model directory has per-layer safetensors for streaming.
+
+    If ``fixed_weights.safetensors`` is missing, split the first ``model*.safetensors``
+    file in place. Returns True if a split was performed.
+    """
+    model_dir = Path(model_dir)
+    if (model_dir / "fixed_weights.safetensors").exists():
+        return False
+
+    weight_files = sorted(glob.glob(str(model_dir / "model*.safetensors")))
+    if not weight_files:
+        raise FileNotFoundError(
+            f"No model*.safetensors in {model_dir}. "
+            "Convert or download a model first, then enable streaming."
+        )
+
+    if verbose:
+        print(f"Preparing streaming layout in {model_dir}")
+    split_model_by_layers(Path(weight_files[0]), model_dir, layer_prefix)
+    return True
+
+
 def main():
     parser = argparse.ArgumentParser(
         description="Split a model into per-layer safetensors for streaming inference"