Feature/layer streaming

.gitignore

@@ -137,3 +137,4 @@ dmypy.json
 
 # .DS_Store files
 .DS_Store
+mlx-lm-full-repo.md

mlx_lm/__init__.py

@@ -10,11 +10,17 @@
 from .generate import batch_generate, generate, stream_generate
 from .utils import load
 
+try:
+    from .streaming import load_streaming
+except ImportError:
+    load_streaming = None
+
 __all__ = [
     "__version__",
     "convert",
     "batch_generate",
     "generate",
     "stream_generate",
     "load",
-]
+    "load_streaming",
+]

mlx_lm/convert.py

@@ -97,6 +97,7 @@ def convert(
         Union[Callable[[str, nn.Module, dict], Union[bool, dict]], str]
     ] = None,
     trust_remote_code: bool = False,
+    split_for_streaming: bool = False,
 ):
     # Check the save path is empty
     if isinstance(mlx_path, str):
@@ -172,6 +173,12 @@ def set_dtype(k, v):
         config,
     )
 
+    if split_for_streaming:
+        from mlx_lm.streaming.split_model import ensure_streaming_layout
+
+        print("[INFO] Splitting weights for layer-streaming inference")
+        ensure_streaming_layout(mlx_path, verbose=True)
+
     if upload_repo is not None:
         upload_to_hub(mlx_path, upload_repo)
 
@@ -251,6 +258,12 @@ def configure_parser() -> argparse.ArgumentParser:
         action="store_true",
         default=False,
     )
+    parser.add_argument(
+        "--split-for-streaming",
+        help="Split saved weights into per-layer files for streaming inference.",
+        action="store_true",
+        default=False,
+    )
     return parser
 
 

mlx_lm/generate.py

@@ -219,11 +219,71 @@ def setup_arg_parser():
         help="Number of tokens to draft when using speculative decoding.",
         default=3,
     )
+    parser.add_argument(
+        "--kv-cache-mode",
+        type=str,
+        choices=["fp16", "tq_asymmetric"],
+        default="fp16",
+        help="KV cache storage mode. tq_asymmetric uses TurboQuant on middle layers.",
+    )
+    parser.add_argument(
+        "--tq-k-bits",
+        type=int,
+        default=4,
+        help="TurboQuant_prod bit width for keys when --kv-cache-mode=tq_asymmetric.",
+    )
+    parser.add_argument(
+        "--tq-v-bits",
+        type=int,
+        default=3,
+        help="TurboQuant_mse bit width for values when --kv-cache-mode=tq_asymmetric.",
+    )
+    parser.add_argument(
+        "--tq-fp16-layers",
+        type=int,
+        default=4,
+        help="FP16 anchor layers at the start and end when using TurboQuant.",
+    )
+    parser.add_argument(
+        "--tq-head-dim",
+        type=int,
+        default=128,
+        help="Attention head dimension for TurboQuant caches.",
+    )
+    parser.add_argument(
+        "--tq-seed",
+        type=int,
+        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
 
 
 # A stream on the default device just for generation
-generation_stream = mx.new_thread_local_stream(mx.default_device())
+def _make_generation_stream():
+    if hasattr(mx, "new_thread_local_stream"):
+        return mx.new_thread_local_stream(mx.default_device())
+    return mx.default_stream(mx.default_device())
+
+
+generation_stream = _make_generation_stream()
 
 
 @contextlib.contextmanager
@@ -317,6 +377,12 @@ def generate_step(
     kv_bits: Optional[int] = None,
     kv_group_size: int = 64,
     quantized_kv_start: int = 0,
+    kv_cache_mode: str = "fp16",
+    tq_k_bits: int = 4,
+    tq_v_bits: int = 3,
+    tq_fp16_layers: int = 4,
+    tq_head_dim: int = 128,
+    tq_seed: int = 42,
     prompt_progress_callback: Optional[Callable[[int, int], None]] = None,
     input_embeddings: Optional[mx.array] = None,
 ) -> Generator[Tuple[mx.array, mx.array], None, None]:
@@ -372,6 +438,12 @@ def generate_step(
         prompt_cache = cache.make_prompt_cache(
             model,
             max_kv_size=max_kv_size,
+            kv_cache_mode=kv_cache_mode,
+            tq_k_bits=tq_k_bits,
+            tq_v_bits=tq_v_bits,
+            tq_fp16_layers=tq_fp16_layers,
+            tq_head_dim=tq_head_dim,
+            tq_seed=tq_seed,
         )
 
     prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
@@ -854,6 +926,19 @@ def to_batch_cache(c):
         elif isinstance(c, CacheList):
             return CacheList(*(to_batch_cache(sub_c) for sub_c in c.caches))
         else:
+            from mlx_lm.turboquant.cache import (
+                AsymmetricTurboQuantCache,
+                BatchAsymmetricTurboQuantCache,
+            )
+
+            if isinstance(c, AsymmetricTurboQuantCache):
+                return BatchAsymmetricTurboQuantCache(
+                    left_padding,
+                    head_dim=c.head_dim,
+                    k_bits=c.k_bits,
+                    v_bits=c.v_bits,
+                    seed=c.seed,
+                )
             raise ValueError(f"{type(c)} does not yet support batching")
 
     if hasattr(model, "make_cache"):
@@ -1508,6 +1593,12 @@ def __init__(
         prefill_batch_size: int = 8,
         prefill_step_size: int = 2048,
         max_kv_size: Optional[int] = None,
+        kv_cache_mode: str = "fp16",
+        tq_k_bits: int = 4,
+        tq_v_bits: int = 3,
+        tq_fp16_layers: int = 4,
+        tq_head_dim: int = 128,
+        tq_seed: int = 42,
         stream=None,
     ):
         self.model = model
@@ -1519,6 +1610,12 @@ def __init__(
         self.prefill_batch_size = prefill_batch_size
         self.completion_batch_size = max(completion_batch_size, prefill_batch_size)
         self.max_kv_size = max_kv_size
+        self.kv_cache_mode = kv_cache_mode
+        self.tq_k_bits = tq_k_bits
+        self.tq_v_bits = tq_v_bits
+        self.tq_fp16_layers = tq_fp16_layers
+        self.tq_head_dim = tq_head_dim
+        self.tq_seed = tq_seed
 
         self._stream = stream or generation_stream
 
@@ -1655,16 +1752,24 @@ def insert_segments(
         return uids
 
     def _make_new_cache(self):
+        cache_kwargs = dict(
+            kv_cache_mode=self.kv_cache_mode,
+            tq_k_bits=self.tq_k_bits,
+            tq_v_bits=self.tq_v_bits,
+            tq_fp16_layers=self.tq_fp16_layers,
+            tq_head_dim=self.tq_head_dim,
+            tq_seed=self.tq_seed,
+        )
         if self.max_kv_size is None:
-            return cache.make_prompt_cache(self.model)
+            return cache.make_prompt_cache(self.model, **cache_kwargs)
 
         return [
             (
                 RotatingKVCache(max_size=self.max_kv_size)
                 if isinstance(ci, KVCache)
                 else ci
             )
-            for ci in cache.make_prompt_cache(self.model)
+            for ci in cache.make_prompt_cache(self.model, **cache_kwargs)
         ]
 
     def _find_uids(self, uids):
@@ -2026,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)
 
@@ -2103,6 +2224,12 @@ def main():
         kv_bits=args.kv_bits,
         kv_group_size=args.kv_group_size,
         quantized_kv_start=args.quantized_kv_start,
+        kv_cache_mode=args.kv_cache_mode,
+        tq_k_bits=args.tq_k_bits,
+        tq_v_bits=args.tq_v_bits,
+        tq_fp16_layers=args.tq_fp16_layers,
+        tq_head_dim=args.tq_head_dim,
+        tq_seed=args.tq_seed,
         draft_model=draft_model,
         num_draft_tokens=args.num_draft_tokens,
     )

mlx_lm/models/base.py

@@ -114,6 +114,16 @@ def scaled_dot_product_attention(
     mask: Optional[mx.array],
     sinks: Optional[mx.array] = None,
 ) -> mx.array:
+    if getattr(cache, "turboquant", False):
+        if sinks is not None:
+            raise ValueError("TurboQuant SDPA does not support attention sinks.")
+        from mlx_lm.turboquant.attention import (
+            turboquant_scaled_dot_product_attention,
+        )
+
+        return turboquant_scaled_dot_product_attention(
+            queries, keys, values, cache, scale=scale, mask=mask
+        )
     if hasattr(cache, "bits"):
         if sinks is not None:
             raise ValueError("Quantized SDPA does not support attention sinks.")

mlx_lm/models/cache.py

@@ -15,6 +15,12 @@
 def make_prompt_cache(
     model: nn.Module,
     max_kv_size: Optional[int] = None,
+    kv_cache_mode: str = "fp16",
+    tq_k_bits: int = 4,
+    tq_v_bits: int = 3,
+    tq_fp16_layers: int = 4,
+    tq_head_dim: int = 128,
+    tq_seed: int = 42,
 ) -> List[Any]:
     """
     Construct the model's cache for use in generation.
@@ -27,7 +33,32 @@ def make_prompt_cache(
         max_kv_size (Optional[int]): If provided and the model does not have a
             ``make_cache`` method, a ``RotatingKVCache`` is used with a maximum
             size of ``max_kv_size``
+        kv_cache_mode (str): ``fp16`` (default) or ``tq_asymmetric`` for
+            TurboQuant KV compression on middle layers.
+        tq_k_bits (int): TurboQuant_prod bit width for keys (>= 2).
+        tq_v_bits (int): TurboQuant_mse bit width for values (2, 3, or 4).
+        tq_fp16_layers (int): FP16 anchor layers at the start and end.
+        tq_head_dim (int): Attention head dimension.
+        tq_seed (int): Base RNG seed for per-layer rotations.
     """
+    if kv_cache_mode == "tq_asymmetric":
+        from mlx_lm.turboquant.factory import make_turboquant_cache
+
+        return make_turboquant_cache(
+            model,
+            max_kv_size=max_kv_size,
+            k_bits=tq_k_bits,
+            v_bits=tq_v_bits,
+            fp16_layers=tq_fp16_layers,
+            head_dim=tq_head_dim,
+            seed=tq_seed,
+        )
+    if kv_cache_mode != "fp16":
+        raise ValueError(
+            f"Unknown kv_cache_mode={kv_cache_mode!r}. "
+            "Supported: 'fp16', 'tq_asymmetric'."
+        )
+
     if hasattr(model, "make_cache"):
         return model.make_cache()
 
@@ -1761,3 +1792,7 @@ def stats_by_type(self):
                 "n_bytes": self._n_bytes_by_type[cache_type],
             }
         return result
+
+
+# Register TurboQuant cache for load_prompt_cache() class lookup.
+from mlx_lm.turboquant.cache import AsymmetricTurboQuantCache  # noqa: E402,F401