diff --git a/.gitignore b/.gitignore index 45445fc89..3fe8d5eeb 100644 --- a/.gitignore +++ b/.gitignore @@ -137,3 +137,4 @@ dmypy.json # .DS_Store files .DS_Store +mlx-lm-full-repo.md diff --git a/mlx_lm/__init__.py b/mlx_lm/__init__.py index 7fb419423..6fcd66afd 100644 --- a/mlx_lm/__init__.py +++ b/mlx_lm/__init__.py @@ -10,6 +10,11 @@ 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", @@ -17,4 +22,5 @@ "generate", "stream_generate", "load", -] + "load_streaming", +] \ No newline at end of file diff --git a/mlx_lm/convert.py b/mlx_lm/convert.py index f28b0fb45..6ed577f12 100644 --- a/mlx_lm/convert.py +++ b/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 diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index 38792a160..ab352236b 100644 --- a/mlx_lm/generate.py +++ b/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,8 +1752,16 @@ 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 [ ( @@ -1664,7 +1769,7 @@ def _make_new_cache(self): 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, ) diff --git a/mlx_lm/models/base.py b/mlx_lm/models/base.py index d7c3efb28..c3a5e6ac9 100644 --- a/mlx_lm/models/base.py +++ b/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.") diff --git a/mlx_lm/models/cache.py b/mlx_lm/models/cache.py index b84c9d650..fe699b952 100644 --- a/mlx_lm/models/cache.py +++ b/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 diff --git a/mlx_lm/server.py b/mlx_lm/server.py index c000f4c10..ac81e48ce 100644 --- a/mlx_lm/server.py +++ b/mlx_lm/server.py @@ -3,6 +3,7 @@ import argparse import json import logging +import os import pickle import platform import socket @@ -344,6 +345,23 @@ def _load(self, model_path, adapter_path=None, draft_model_path=None): tokenizer_config=self._tokenizer_config, trust_remote_code=self.cli_args.trust_remote_code, ) + elif getattr(self.cli_args, "streaming", False): + if adapter_path is not None or draft_model_path is not None: + raise ValueError( + "Layer streaming does not support adapters or draft models" + ) + from mlx_lm.streaming import StreamingConfig, load_streaming + + streaming_config = StreamingConfig( + max_memory_gb=self.cli_args.max_memory_gb, + window_size=getattr(self.cli_args, "streaming_window", None), + verbose=self.cli_args.log_level == "DEBUG", + ) + model, tokenizer, _ = load_streaming( + model_path, + streaming_config=streaming_config, + tokenizer_config=self._tokenizer_config, + ) else: model, tokenizer = load( model_path, @@ -370,7 +388,16 @@ def _load(self, model_path, adapter_path=None, draft_model_path=None): # Compute batchability is_batchable = draft_model is None is_batchable = is_batchable and all( - hasattr(c, "merge") for c in make_prompt_cache(model) + hasattr(c, "merge") + for c in make_prompt_cache( + model, + kv_cache_mode=self.cli_args.kv_cache_mode, + tq_k_bits=self.cli_args.tq_k_bits, + tq_v_bits=self.cli_args.tq_v_bits, + tq_fp16_layers=self.cli_args.tq_fp16_layers, + tq_head_dim=self.cli_args.tq_head_dim, + tq_seed=self.cli_args.tq_seed, + ) ) # Update the member variables @@ -437,6 +464,17 @@ def _format_top_logprobs(logprobs, top_n, tokenizer) -> Tuple[Dict[str, Any]]: ) +def _kv_cache_kwargs(cli_args: argparse.Namespace) -> dict: + return dict( + kv_cache_mode=cli_args.kv_cache_mode, + tq_k_bits=cli_args.tq_k_bits, + tq_v_bits=cli_args.tq_v_bits, + tq_fp16_layers=cli_args.tq_fp16_layers, + tq_head_dim=cli_args.tq_head_dim, + tq_seed=cli_args.tq_seed, + ) + + class ResponseGenerator: def __init__(self, model_provider: ModelProvider, prompt_cache: LRUPromptCache): self.model_provider = model_provider @@ -824,6 +862,7 @@ def get_next_request(timeout=None): prefill_batch_size=self.cli_args.prompt_concurrency, prefill_step_size=self.cli_args.prefill_step_size, stream=generation_stream, + **_kv_cache_kwargs(self.cli_args), ) unprocessed_requests.append((rqueue, request, args)) continue @@ -967,10 +1006,13 @@ def progress(tokens_processed, tokens_total): ) ctx.prompt_cache_count = len(prompt) - len(rest) cache_key = prompt[:] + kv_kwargs = _kv_cache_kwargs(self.cli_args) if cache is None: - cache = make_prompt_cache(self.model_provider.model) + cache = make_prompt_cache(self.model_provider.model, **kv_kwargs) if self.model_provider.draft_model is not None: - cache += make_prompt_cache(self.model_provider.draft_model) + cache += make_prompt_cache( + self.model_provider.draft_model, **kv_kwargs + ) # Process the prompt and generate tokens for gen in stream_generate( @@ -985,6 +1027,7 @@ def progress(tokens_processed, tokens_total): num_draft_tokens=args.num_draft_tokens, prompt_progress_callback=progress, prefill_step_size=self.cli_args.prefill_step_size, + **kv_kwargs, ): finish_reason = gen.finish_reason sm_state, match_sequence, current_state = sm.match(sm_state, gen.token) @@ -1879,11 +1922,65 @@ def main(): type=_parse_size, help="Maximum size in bytes of the KV caches", ) + parser.add_argument( + "--kv-cache-mode", + type=str, + choices=["fp16", "tq_asymmetric"], + default=os.environ.get("BONSAI_KV_CACHE_MODE", "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 when using TurboQuant.", + ) + parser.add_argument( + "--tq-seed", + type=int, + default=42, + help="Base RNG seed for per-layer TurboQuant rotations.", + ) parser.add_argument( "--pipeline", action="store_true", help="Use pipelining instead of tensor parallelism", ) + 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).", + ) args = parser.parse_args() if mx.metal.is_available(): wired_limit = mx.device_info()["max_recommended_working_set_size"] diff --git a/mlx_lm/streaming/__init__.py b/mlx_lm/streaming/__init__.py new file mode 100644 index 000000000..190e59d7e --- /dev/null +++ b/mlx_lm/streaming/__init__.py @@ -0,0 +1,18 @@ +# Copyright © 2024 Apple Inc. + +"""Layer-streaming inference for models larger than available unified memory.""" + +from .config import StreamingConfig +from .layer_loader import RollingWindowLoader +from .load import load_streaming +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", +] \ No newline at end of file diff --git a/mlx_lm/streaming/benchmark.py b/mlx_lm/streaming/benchmark.py new file mode 100644 index 000000000..136a4c27d --- /dev/null +++ b/mlx_lm/streaming/benchmark.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() \ No newline at end of file diff --git a/mlx_lm/streaming/config.py b/mlx_lm/streaming/config.py new file mode 100644 index 000000000..440171d5f --- /dev/null +++ b/mlx_lm/streaming/config.py @@ -0,0 +1,67 @@ +# Copyright © 2024 Apple Inc. + +"""Streaming inference configuration and memory estimation.""" + +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class StreamingConfig: + """Runtime settings for layer-streaming inference.""" + + max_memory_gb: float = 20.0 + window_size: Optional[int] = None + safety_margin: float = 0.8 + prefetch_layers: int = 1 + clear_cache_every_n_layers: int = 5 + max_seq_len: int = 4096 + verbose: bool = False + + def estimate_layer_memory( + self, + hidden_size: int, + intermediate_size: int, + num_attention_heads: int, + num_key_value_heads: int, + dtype_bytes: int = 2, + ) -> int: + """Estimate bytes for one transformer layer.""" + head_dim = hidden_size // num_attention_heads + kv_dim = head_dim * num_key_value_heads + + attn_params = ( + hidden_size * hidden_size + + hidden_size * kv_dim * 2 + + hidden_size * hidden_size + ) + ffn_params = ( + hidden_size * intermediate_size * 2 + + intermediate_size * hidden_size + ) + norm_params = hidden_size * 2 + total_params = attn_params + ffn_params + norm_params + return int(total_params * dtype_bytes) + + def compute_optimal_window( + self, + layer_size_bytes: int, + kv_cache_bytes: int, + fixed_memory_bytes: int = 2_000_000_000, + ) -> int: + """Compute how many layers fit in the rolling window.""" + available = self.max_memory_gb * 1_000_000_000 * self.safety_margin + activation_overhead = layer_size_bytes * 2 + usable = available - kv_cache_bytes - activation_overhead - fixed_memory_bytes + window = max(1, int(usable / layer_size_bytes)) + + if self.verbose: + print("Streaming memory breakdown:") + print(f" Available: {available / 1e9:.2f} GB") + print(f" KV cache: {kv_cache_bytes / 1e9:.2f} GB") + print(f" Activations: {activation_overhead / 1e9:.2f} GB") + print(f" Fixed overhead: {fixed_memory_bytes / 1e9:.2f} GB") + print(f" Layer size: {layer_size_bytes / 1e6:.2f} MB") + print(f" Window size: {window} layers") + + return window \ No newline at end of file diff --git a/mlx_lm/streaming/layer_loader.py b/mlx_lm/streaming/layer_loader.py new file mode 100644 index 000000000..d82e4c023 --- /dev/null +++ b/mlx_lm/streaming/layer_loader.py @@ -0,0 +1,175 @@ +# Copyright © 2024 Apple Inc. + +"""Rolling-window layer weight loader with async prefetch.""" + +import mlx.core as mx +from collections import deque +from pathlib import Path +from typing import Dict, Optional + +from .config import StreamingConfig + + +class LayerWeights: + """Weights for a single transformer layer.""" + + def __init__(self, weights: Dict[str, mx.array], layer_idx: int): + self.weights = weights + self.layer_idx = layer_idx + self._memory_bytes: Optional[int] = None + + @property + def memory_bytes(self) -> int: + if self._memory_bytes is None: + self._memory_bytes = sum(w.nbytes for w in self.weights.values()) + return self._memory_bytes + + +class RollingWindowLoader: + """ + Sliding window of layer safetensors with eviction and prefetch. + + Expects per-layer files produced by ``split_model_by_layers``: + ``layer_{i}.safetensors`` plus optional ``fixed_weights.safetensors``. + """ + + def __init__( + self, + model_path: Path, + num_layers: int, + layer_size_bytes: int, + streaming_config: StreamingConfig, + kv_cache_bytes: int = 0, + layer_key_prefix: str = "model.layers", + ): + self.model_path = Path(model_path) + self.num_layers = num_layers + self.streaming_config = streaming_config + self.layer_key_prefix = layer_key_prefix + + if streaming_config.window_size is None: + self.window_size = streaming_config.compute_optimal_window( + layer_size_bytes, kv_cache_bytes + ) + else: + self.window_size = streaming_config.window_size + + if streaming_config.verbose: + print(f"RollingWindowLoader: path={model_path}, window={self.window_size}") + + self.layer_queue: deque = deque(maxlen=self.window_size) + self.loaded_layers: set = set() + self._layer_files = self._discover_layer_files() + + def _discover_layer_files(self) -> Dict[int, Optional[Path]]: + patterns = [ + "layer_{}.safetensors", + "model.layers.{}.safetensors", + "layers.{}.safetensors", + ] + layer_files = {} + for layer_idx in range(self.num_layers): + found = None + for pattern in patterns: + path = self.model_path / pattern.format(layer_idx) + if path.exists(): + found = path + break + layer_files[layer_idx] = found + return layer_files + + def _load_from_monolithic(self, layer_idx: int) -> Dict[str, mx.array]: + model_file = self.model_path / "model.safetensors" + if not model_file.exists(): + weight_files = list(self.model_path.glob("model*.safetensors")) + if not weight_files: + raise FileNotFoundError( + f"No weights for layer {layer_idx} in {self.model_path}" + ) + model_file = weight_files[0] + + all_weights = mx.load(str(model_file)) + patterns = [ + f"{self.layer_key_prefix}.{layer_idx}.", + f"layers.{layer_idx}.", + f"transformer.h.{layer_idx}.", + ] + weights = {} + for key, tensor in all_weights.items(): + for pattern in patterns: + if key.startswith(pattern): + weights[key[len(pattern):]] = tensor + break + if not weights: + raise FileNotFoundError( + f"Layer {layer_idx} not found in {model_file}" + ) + return weights + + def load_layer(self, layer_idx: int, prefetch: bool = False) -> LayerWeights: + layer_file = self._layer_files.get(layer_idx) + if layer_file is None: + weights = self._load_from_monolithic(layer_idx) + else: + weights = mx.load(str(layer_file)) + + layer_weights = LayerWeights(weights, layer_idx) + if prefetch and weights: + mx.async_eval(list(weights.values())) + return layer_weights + + def preload_window(self, start_idx: int = 0): + end_idx = min(start_idx + self.window_size, self.num_layers) + for i in range(start_idx, end_idx): + lw = self.load_layer(i, prefetch=True) + self.layer_queue.append(lw) + self.loaded_layers.add(i) + all_arrays = [] + for layer in self.layer_queue: + all_arrays.extend(layer.weights.values()) + if all_arrays: + mx.eval(all_arrays) + + def get_layer(self, layer_idx: int) -> LayerWeights: + for layer in self.layer_queue: + if layer.layer_idx == layer_idx: + return layer + + if len(self.layer_queue) >= self.window_size: + evicted = self.layer_queue.popleft() + self.loaded_layers.discard(evicted.layer_idx) + del evicted + if ( + layer_idx % self.streaming_config.clear_cache_every_n_layers == 0 + ): + mx.clear_cache() + + layer_weights = self.load_layer(layer_idx, prefetch=False) + mx.eval(list(layer_weights.weights.values())) + self.layer_queue.append(layer_weights) + self.loaded_layers.add(layer_idx) + + for offset in range(1, self.streaming_config.prefetch_layers + 1): + next_idx = layer_idx + offset + if ( + next_idx < self.num_layers + and next_idx not in self.loaded_layers + ): + self.load_layer(next_idx, prefetch=True) + + 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.clear_cache() \ No newline at end of file diff --git a/mlx_lm/streaming/load.py b/mlx_lm/streaming/load.py new file mode 100644 index 000000000..c74874a13 --- /dev/null +++ b/mlx_lm/streaming/load.py @@ -0,0 +1,110 @@ +# Copyright © 2024 Apple Inc. + +"""Load mlx-lm models for layer-streaming inference.""" + +from pathlib import Path +from typing import Dict, Optional, Tuple, Union + +import mlx.core as mx +import mlx.nn as nn + +from mlx_lm.tokenizer_utils import TokenizerWrapper +from mlx_lm.utils import _download, _get_classes, load_config, load_tokenizer + +from .config import StreamingConfig +from .layer_loader import RollingWindowLoader +from .split_model import ensure_streaming_layout +from .wrapper import StreamingModelWrapper + + +def _dtype_bytes(config: dict) -> int: + q = config.get("quantization") or config.get("quantization_config") + if q: + bits = q.get("bits", 4) if isinstance(q, dict) else 4 + return max(1, bits // 8) + torch_dtype = config.get("torch_dtype", "bfloat16") + if "float32" in str(torch_dtype): + return 4 + return 2 + + +def _estimate_kv_cache_bytes(config: dict, streaming_config: StreamingConfig) -> int: + n_layers = config.get("num_hidden_layers", 32) + n_kv = config.get("num_key_value_heads") or config.get("num_attention_heads", 32) + hidden = config.get("hidden_size", 4096) + n_heads = config.get("num_attention_heads", 32) + head_dim = hidden // n_heads + bpe = 2 + return int( + 2 * n_layers * streaming_config.max_seq_len * n_kv * head_dim * bpe + ) + + +def load_streaming( + path_or_hf_repo: str, + streaming_config: Optional[StreamingConfig] = None, + tokenizer_config: Optional[dict] = None, + revision: Optional[str] = None, + load_tokenizer: bool = True, +) -> Tuple[StreamingModelWrapper, Optional[TokenizerWrapper], Dict]: + """ + Load a model for layer-streaming inference. + + The model directory should contain per-layer safetensors (from + ``mlx_lm.streaming.split_model``) or a monolithic ``model*.safetensors`` + fallback. + + Args: + path_or_hf_repo: Local path or HuggingFace repo id. + streaming_config: Memory/window settings. + tokenizer_config: Passed to ``load_tokenizer``. + revision: HF revision id. + + Returns: + (StreamingModelWrapper, tokenizer, config dict) + """ + streaming_config = streaming_config or StreamingConfig() + tokenizer_config = tokenizer_config or {} + + 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" + 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"], + intermediate_size=config["intermediate_size"], + num_attention_heads=config["num_attention_heads"], + num_key_value_heads=config.get("num_key_value_heads") + or config["num_attention_heads"], + dtype_bytes=_dtype_bytes(config), + ) + + kv_bytes = _estimate_kv_cache_bytes(config, streaming_config) + loader = RollingWindowLoader( + model_path=model_path, + num_layers=config["num_hidden_layers"], + layer_size_bytes=layer_size, + streaming_config=streaming_config, + kv_cache_bytes=kv_bytes, + ) + loader.preload_window() + + wrapped = StreamingModelWrapper(model, loader) + wrapped.eval() + + tokenizer = ( + load_tokenizer(model_path, tokenizer_config) if load_tokenizer else None + ) + return wrapped, tokenizer, config \ No newline at end of file diff --git a/mlx_lm/streaming/split_model.py b/mlx_lm/streaming/split_model.py new file mode 100644 index 000000000..2efea10cb --- /dev/null +++ b/mlx_lm/streaming/split_model.py @@ -0,0 +1,103 @@ +# Copyright © 2024 Apple Inc. + +"""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, Optional + + +def split_model_by_layers( + model_path: Path, + output_dir: Path, + layer_prefix: str = "model.layers", +) -> None: + """ + Split a monolithic safetensors model into per-layer files. + + Writes ``layer_{i}.safetensors`` with keys stripped of the layer prefix, + and ``fixed_weights.safetensors`` for embeddings, norm, and lm_head. + """ + output_dir.mkdir(parents=True, exist_ok=True) + print(f"Splitting {model_path} -> {output_dir}") + + all_weights = mx.load(str(model_path)) + layer_keys: Dict[int, List[str]] = {} + fixed_keys: List[str] = [] + + for key in all_weights.keys(): + if layer_prefix in key: + parts = key.split(".") + try: + layer_idx_pos = parts.index("layers") + 1 + layer_idx = int(parts[layer_idx_pos]) + layer_keys.setdefault(layer_idx, []).append(key) + except (ValueError, IndexError): + fixed_keys.append(key) + else: + fixed_keys.append(key) + + print(f"Layers: {len(layer_keys)}, fixed tensors: {len(fixed_keys)}") + + for layer_idx in sorted(layer_keys.keys()): + layer_dict = {} + for key in layer_keys[layer_idx]: + local_key = key.replace(f"{layer_prefix}.{layer_idx}.", "") + layer_dict[local_key] = all_weights[key] + out = output_dir / f"layer_{layer_idx}.safetensors" + mx.save_safetensors(str(out), layer_dict) + + if fixed_keys: + fixed_dict = {key: all_weights[key] for key in fixed_keys} + mx.save_safetensors(str(output_dir / "fixed_weights.safetensors"), fixed_dict) + + 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" + ) + parser.add_argument("model_path", type=Path, help="Path to model.safetensors") + parser.add_argument("output_dir", type=Path, help="Output directory") + parser.add_argument( + "--layer-prefix", + default="model.layers", + help="Layer key prefix in the checkpoint", + ) + args = parser.parse_args() + split_model_by_layers(args.model_path, args.output_dir, args.layer_prefix) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/mlx_lm/streaming/wrapper.py b/mlx_lm/streaming/wrapper.py new file mode 100644 index 000000000..58c848a18 --- /dev/null +++ b/mlx_lm/streaming/wrapper.py @@ -0,0 +1,120 @@ +# Copyright © 2024 Apple Inc. + +"""Model wrapper that streams transformer layers from disk.""" + +from typing import List, Optional + +import mlx.core as mx +import mlx.nn as nn + +from mlx_lm.models.base import create_attention_mask + +from .layer_loader import RollingWindowLoader + + +class StreamingModelWrapper: + """ + Wraps an mlx-lm model and loads layer weights on demand. + + Compatible with ``mlx_lm.generate`` — exposes the same ``__call__`` signature + and delegates KV cache handling to the underlying architecture. + + The wrapped model is stored privately (not as an nn.Module child) to avoid + conflicting with MLX parameter tracking during per-layer load_weights. + """ + + def __init__(self, model: nn.Module, loader: RollingWindowLoader): + self._model = model + self.loader = loader + self._streaming_enabled = self._detect_streaming_support() + + def _detect_streaming_support(self) -> bool: + inner = getattr(self._model, "model", None) + return inner is not None and hasattr(inner, "layers") + + @property + def args(self): + return self._model.args + + def make_cache(self): + if hasattr(self._model, "make_cache"): + return self._model.make_cache() + inner = self._model.model + return [None] * len(inner.layers) + + def sanitize(self, weights): + if hasattr(self._model, "sanitize"): + return self._model.sanitize(weights) + return weights + + def set_dtype(self, dtype): + if hasattr(self._model, "set_dtype"): + self._model.set_dtype(dtype) + return self + + def __getattr__(self, name): + if name in ("_model", "loader", "_streaming_enabled"): + raise AttributeError(name) + return getattr(self._model, name) + + def get_stats(self) -> dict: + """Memory and window statistics for streaming inference.""" + usage = self.loader.get_memory_usage() + return { + "streaming": { + "window_size": usage["window_size"], + "loaded_layers": usage["loaded_layers"], + "layer_indices": usage["layer_indices"], + "total_mb": usage["total_mb"], + }, + } + + def _forward_streaming( + self, + inputs: mx.array, + cache: Optional[List], + input_embeddings: Optional[mx.array], + ) -> mx.array: + m = self._model.model + if input_embeddings is not None: + h = input_embeddings + else: + h = m.embed_tokens(inputs) + + if cache is None: + cache = [None] * len(m.layers) + + fa_idx = getattr(m, "fa_idx", 0) + swa_idx = getattr(m, "swa_idx", None) + sliding_window = getattr(m, "sliding_window", None) + + fa_mask = create_attention_mask(h, cache[fa_idx]) + swa_mask = None + if swa_idx is not None: + swa_mask = create_attention_mask( + h, cache[swa_idx], window_size=sliding_window + ) + + for i, (layer, layer_cache) in enumerate(zip(m.layers, cache)): + layer_weights = self.loader.get_layer(i) + layer.load_weights(list(layer_weights.weights.items()), strict=False) + mx.eval(layer.parameters()) + + mask = swa_mask if getattr(layer, "use_sliding", False) else fa_mask + h = layer(h, mask, cache=layer_cache) + + h = m.norm(h) + + if self._model.args.tie_word_embeddings: + return m.embed_tokens.as_linear(h) + return self._model.lm_head(h) + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ): + if self._streaming_enabled: + return self._forward_streaming(inputs, cache, input_embeddings) + return self._model(inputs, cache=cache, input_embeddings=input_embeddings) \ No newline at end of file diff --git a/mlx_lm/turboquant/__init__.py b/mlx_lm/turboquant/__init__.py new file mode 100644 index 000000000..34515a558 --- /dev/null +++ b/mlx_lm/turboquant/__init__.py @@ -0,0 +1,5 @@ +# Copyright © 2025 Bonsai Demo contributors. + +from mlx_lm.turboquant.factory import make_turboquant_cache + +__all__ = ["make_turboquant_cache"] \ No newline at end of file diff --git a/mlx_lm/turboquant/attention.py b/mlx_lm/turboquant/attention.py new file mode 100644 index 000000000..022e2efcd --- /dev/null +++ b/mlx_lm/turboquant/attention.py @@ -0,0 +1,187 @@ +# Copyright © 2025 Bonsai Demo contributors. + +"""Fused TurboQuant attention — packed KV, no persistent dequant buffers.""" + +from __future__ import annotations + +import math +from typing import Optional, Tuple + +import mlx.core as mx + +from mlx_lm.turboquant.codebooks import get_codebook +from mlx_lm.turboquant.kernels import ( + decode_mse_metal, + metal_available, + tq_sdpa_metal, +) +from mlx_lm.turboquant.packing import unpack_indices +from mlx_lm.turboquant.quantize import decode_mse, decode_prod + +_QJL_SCALE = math.sqrt(math.pi / 2.0) + + +def qk_scores_vectorized( + queries: mx.array, + k_packed: mx.array, + k_norms: mx.array, + k_qjl_signs: mx.array, + k_gamma: mx.array, + rotation: mx.array, + s_matrix: mx.array, + k_bits: int, + dim: int, + scale: float, +) -> mx.array: + """Compute QK logits from packed keys without materializing K (for tests).""" + mse_bits = k_bits - 1 + B, n_q_heads, L, D = queries.shape + n_kv_heads = k_packed.shape[1] + S = k_packed.shape[2] + n_repeats = n_q_heads // n_kv_heads + + q_rot = queries @ rotation.T + q_s = queries @ s_matrix.T + + if n_repeats > 1: + q_rot = q_rot.reshape(B, n_kv_heads, n_repeats, L, D) + q_s = q_s.reshape(B, n_kv_heads, n_repeats, L, D) + + indices = unpack_indices(k_packed, mse_bits, dim) + centroids, _ = get_codebook(mse_bits) + rot_c = centroids[indices.astype(mx.int32)] + + mse_dot = mx.sum( + q_rot[..., None, :] * rot_c[:, :, None, None, :, :], axis=-1 + ) + signs = unpack_indices(k_qjl_signs, 1, dim).astype(mx.float32) * 2.0 - 1.0 + qjl_dot = mx.sum( + q_s[..., None, :] * signs[:, :, None, None, :, :], axis=-1 + ) + + kn = k_norms[..., 0][:, :, None, None, :] + kg = k_gamma[..., 0][:, :, None, None, :] + scores = kn * (mse_dot + (_QJL_SCALE / dim) * kg * qjl_dot) + + if n_repeats > 1: + scores = scores.reshape(B, n_q_heads, L, S) + return scores * scale + + +def av_weighted_sum_vectorized( + attn: mx.array, + v_packed: mx.array, + v_norms: mx.array, + rotation: mx.array, + v_bits: int, + dim: int, +) -> mx.array: + """Reference: attn @ dequant(V) without materializing full V.""" + B, n_q_heads, L, S = attn.shape + n_kv_heads = v_packed.shape[1] + n_repeats = n_q_heads // n_kv_heads + + values_deq = decode_mse(v_packed, v_norms, rotation, v_bits, dim) + if n_repeats > 1: + values_deq = mx.repeat(values_deq, n_repeats, axis=1) + return mx.matmul(attn, values_deq) + + +def _qk_scores_reference( + queries: mx.array, + k_packed: mx.array, + k_norms: mx.array, + k_qjl_signs: mx.array, + k_gamma: mx.array, + rotation: mx.array, + s_matrix: mx.array, + k_bits: int, + dim: int, + scale: float, +) -> mx.array: + keys = decode_prod( + k_packed, k_norms, k_qjl_signs, k_gamma, rotation, s_matrix, k_bits, dim + ) + n_repeats = queries.shape[1] // keys.shape[1] + if n_repeats > 1: + keys = mx.repeat(keys, n_repeats, axis=1) + return mx.matmul(queries, keys.transpose(0, 1, 3, 2)) * scale + + +def _decode_values( + v_packed: mx.array, + v_norms: mx.array, + rotation: mx.array, + v_bits: int, + dim: int, + dtype: mx.dtype, +) -> mx.array: + if metal_available(): + return decode_mse_metal(v_packed, v_norms, rotation, v_bits, dim, dtype=dtype) + return decode_mse(v_packed, v_norms, rotation, v_bits, dim).astype(dtype) + + +def _sdpa_decode_fallback( + queries: mx.array, + keys: Tuple[mx.array, mx.array, mx.array, mx.array], + values: Tuple[mx.array, mx.array], + cache, + scale: float, + mask: Optional[mx.array], +) -> mx.array: + """MLX fallback: ephemeral K/V decode + fast SDPA.""" + k_packed, k_norms, k_qjl_signs, k_gamma = keys + v_packed, v_norms = values + dtype = queries.dtype + dim = queries.shape[-1] + n_repeats = queries.shape[1] // k_packed.shape[1] + + keys_deq = decode_prod( + k_packed, k_norms, k_qjl_signs, k_gamma, + cache._k_rotation, cache._k_qjl, cache.k_bits, dim, + ).astype(dtype) + values_deq = _decode_values( + v_packed, v_norms, cache._v_rotation, cache.v_bits, dim, dtype, + ) + if n_repeats > 1: + keys_deq = mx.repeat(keys_deq, n_repeats, axis=1) + values_deq = mx.repeat(values_deq, n_repeats, axis=1) + + return mx.fast.scaled_dot_product_attention( + queries, keys_deq, values_deq, scale=scale, mask=mask + ) + + +def turboquant_scaled_dot_product_attention( + queries: mx.array, + keys: Tuple[mx.array, mx.array, mx.array, mx.array], + values: Tuple[mx.array, mx.array], + cache, + scale: float, + mask: Optional[mx.array], +) -> mx.array: + """SDPA from packed TurboQuant KV without persistent dequant buffers.""" + if metal_available(): + k_packed, k_norms, k_qjl_signs, k_gamma = keys + v_packed, v_norms = values + dim = queries.shape[-1] + L = queries.shape[2] + do_causal = mask == "causal" and L > 1 + return tq_sdpa_metal( + queries, + k_packed, + k_norms, + k_qjl_signs, + k_gamma, + v_packed, + v_norms, + cache._k_rotation, + cache._v_rotation, + cache._k_qjl, + cache.k_bits, + cache.v_bits, + dim, + scale, + do_causal, + ) + return _sdpa_decode_fallback(queries, keys, values, cache, scale, mask) \ No newline at end of file diff --git a/mlx_lm/turboquant/cache.py b/mlx_lm/turboquant/cache.py new file mode 100644 index 000000000..5ecc54097 --- /dev/null +++ b/mlx_lm/turboquant/cache.py @@ -0,0 +1,620 @@ +# Copyright © 2025 Bonsai Demo contributors. + +"""TurboQuant KV cache implementations for mlx-lm.""" + +from __future__ import annotations + +import copy +from typing import List + +import mlx.core as mx + +from mlx_lm.models.cache import ( + _BaseCache, + create_attention_mask, + create_causal_mask, + dynamic_roll, +) +from mlx_lm.turboquant.codebooks import HEAD_DIM +from mlx_lm.turboquant.packing import packed_dim +from mlx_lm.turboquant.qjl import make_qjl_matrix, qjl_packed_dim +from mlx_lm.turboquant.quantize import encode_kv +from mlx_lm.turboquant.rotation import make_rotation_matrix + + +class AsymmetricTurboQuantCache(_BaseCache): + """KV cache with TurboQuant_prod keys and TurboQuant_mse values. + + Keys use an inner-product-preserving quantizer (Lloyd-Max + QJL residual). + Values use an MSE-optimal quantizer. Packed tensors are returned from + ``update_and_fetch`` for fused attention and fused KV encode. + """ + + step = 256 + turboquant = True + + def __init__( + self, + head_dim: int = HEAD_DIM, + k_bits: int = 4, + v_bits: int = 3, + seed: int = 42, + ): + if k_bits < 2: + raise ValueError("k_bits must be >= 2 for TurboQuant_prod") + if v_bits not in (2, 3, 4): + raise ValueError("v_bits must be 2, 3, or 4") + self.head_dim = head_dim + self.k_bits = k_bits + self.v_bits = v_bits + self.seed = seed + self.offset = 0 + + self._k_rotation = make_rotation_matrix(head_dim, seed) + self._v_rotation = make_rotation_matrix(head_dim, seed + 97) + self._k_qjl = make_qjl_matrix(head_dim, seed + 193) + + self._k_packed = None + self._k_norms = None + self._k_qjl_signs = None + self._k_gamma = None + + self._v_packed = None + self._v_norms = None + + self._dtype = mx.float16 + + @property + def k_pdim(self) -> int: + return packed_dim(self.head_dim, self.k_bits - 1) + + @property + def v_pdim(self) -> int: + return packed_dim(self.head_dim, self.v_bits) + + @property + def qjl_pdim(self) -> int: + return qjl_packed_dim(self.head_dim) + + def _grow(self, b: int, h: int, needed: int): + prev = self.offset + if self._k_packed is not None and needed <= self._k_packed.shape[2]: + return + n = ((needed + self.step - 1) // self.step) * self.step + k_shape = (b, h, n, self.k_pdim) + v_shape = (b, h, n, self.v_pdim) + if self._k_packed is None: + self._k_packed = mx.zeros(k_shape, dtype=mx.uint32) + self._k_norms = mx.zeros((b, h, n, 1), dtype=mx.float32) + self._k_qjl_signs = mx.zeros((b, h, n, self.qjl_pdim), dtype=mx.uint32) + self._k_gamma = mx.zeros((b, h, n, 1), dtype=mx.float32) + self._v_packed = mx.zeros(v_shape, dtype=mx.uint32) + self._v_norms = mx.zeros((b, h, n, 1), dtype=mx.float32) + return + + def _extend(tensor, shape): + new = mx.zeros(shape, dtype=tensor.dtype) + new[:, :, :prev, :] = tensor[:, :, :prev, :] + return new + + self._k_packed = _extend(self._k_packed, k_shape) + self._k_norms = _extend(self._k_norms, (b, h, n, 1)) + self._k_qjl_signs = _extend(self._k_qjl_signs, (b, h, n, self.qjl_pdim)) + self._k_gamma = _extend(self._k_gamma, (b, h, n, 1)) + self._v_packed = _extend(self._v_packed, v_shape) + self._v_norms = _extend(self._v_norms, (b, h, n, 1)) + + def _fetch_packed(self): + t = self.offset + keys = ( + self._k_packed[..., :t, :], + self._k_norms[..., :t, :], + self._k_qjl_signs[..., :t, :], + self._k_gamma[..., :t, :], + ) + values = (self._v_packed[..., :t, :], self._v_norms[..., :t, :]) + return keys, values + + def update_and_fetch(self, keys, values): + b, h, s, d = keys.shape + if d != self.head_dim: + raise ValueError(f"Expected head_dim={self.head_dim}, got {d}") + self._dtype = keys.dtype + prev = self.offset + self._grow(b, h, prev + s) + + k_pack, k_norm, k_sign, k_gamma, v_pack, v_norm = encode_kv( + keys, + values, + self._k_rotation, + self._v_rotation, + self._k_qjl, + self.k_bits, + self.v_bits, + ) + + self._k_packed[..., prev : prev + s, :] = k_pack + self._k_norms[..., prev : prev + s, :] = k_norm + self._k_qjl_signs[..., prev : prev + s, :] = k_sign + self._k_gamma[..., prev : prev + s, :] = k_gamma + self._v_packed[..., prev : prev + s, :] = v_pack + self._v_norms[..., prev : prev + s, :] = v_norm + self.offset += s + + return self._fetch_packed() + + def make_mask(self, *args, **kwargs): + return create_attention_mask(*args, offset=self.offset, **kwargs) + + def is_trimmable(self): + return True + + def trim(self, n): + n = min(self.offset, n) + self.offset -= n + return n + + def size(self): + return self.offset + + def empty(self): + return self._k_packed is None + + @property + def nbytes(self): + if self._k_packed is None: + return 0 + t = self.offset + total = 0 + total += self._k_packed[..., :t, :].nbytes + total += self._k_norms[..., :t, :].nbytes + total += self._k_qjl_signs[..., :t, :].nbytes + total += self._k_gamma[..., :t, :].nbytes + total += self._v_packed[..., :t, :].nbytes + total += self._v_norms[..., :t, :].nbytes + return total + + @property + def state(self): + if self.empty(): + return [] + t = self.offset + return [ + self._k_packed[..., :t, :], + self._k_norms[..., :t, :], + self._k_qjl_signs[..., :t, :], + self._k_gamma[..., :t, :], + self._v_packed[..., :t, :], + self._v_norms[..., :t, :], + ] + + @state.setter + def state(self, value): + if not value: + return + ( + self._k_packed, + self._k_norms, + self._k_qjl_signs, + self._k_gamma, + self._v_packed, + self._v_norms, + ) = value + self.offset = self._k_packed.shape[2] + + @property + def meta_state(self): + return f"{self.offset},{self.k_bits},{self.v_bits},{self.seed},{self.head_dim}" + + @meta_state.setter + def meta_state(self, meta): + parts = meta.split(",") + self.offset = int(parts[0]) + self.k_bits = int(parts[1]) + self.v_bits = int(parts[2]) + self.seed = int(parts[3]) + self.head_dim = int(parts[4]) + self._k_rotation = make_rotation_matrix(self.head_dim, self.seed) + self._v_rotation = make_rotation_matrix(self.head_dim, self.seed + 97) + self._k_qjl = make_qjl_matrix(self.head_dim, self.seed + 193) + + @classmethod + def from_state(cls, state, meta_state): + obj = cls.__new__(cls) + obj._k_packed = None + obj._k_norms = None + obj._k_qjl_signs = None + obj._k_gamma = None + obj._v_packed = None + obj._v_norms = None + obj._dtype = mx.float16 + obj.meta_state = meta_state + obj.state = state + return obj + + @classmethod + def merge(cls, caches: List["AsymmetricTurboQuantCache"]): + return BatchAsymmetricTurboQuantCache.merge(caches) + + def __deepcopy__(self, memo): + cls = self.__class__ + obj = cls.__new__(cls) + memo[id(self)] = obj + for key, value in self.__dict__.items(): + if key == "_dtype": + setattr(obj, key, value) + else: + setattr(obj, key, copy.deepcopy(value, memo)) + return obj + + +class BatchAsymmetricTurboQuantCache(_BaseCache): + """Batched TurboQuant KV cache for concurrent server requests.""" + + step = 256 + turboquant = True + + def __init__( + self, + left_padding: List[int], + head_dim: int = HEAD_DIM, + k_bits: int = 4, + v_bits: int = 3, + seed: int = 42, + ): + if k_bits < 2: + raise ValueError("k_bits must be >= 2 for TurboQuant_prod") + if v_bits not in (2, 3, 4): + raise ValueError("v_bits must be 2, 3, or 4") + self.head_dim = head_dim + self.k_bits = k_bits + self.v_bits = v_bits + self.seed = seed + self._k_rotation = make_rotation_matrix(head_dim, seed) + self._v_rotation = make_rotation_matrix(head_dim, seed + 97) + self._k_qjl = make_qjl_matrix(head_dim, seed + 193) + + self._k_packed = None + self._k_norms = None + self._k_qjl_signs = None + self._k_gamma = None + self._v_packed = None + self._v_norms = None + + self.left_padding = mx.array(left_padding) + self.offset = mx.array([-l for l in left_padding]) + self._idx = 0 + self._right_padding = None + self._dtype = mx.float16 + + @property + def k_pdim(self) -> int: + return packed_dim(self.head_dim, self.k_bits - 1) + + @property + def v_pdim(self) -> int: + return packed_dim(self.head_dim, self.v_bits) + + @property + def qjl_pdim(self) -> int: + return qjl_packed_dim(self.head_dim) + + def _tensor_fields(self): + return ( + self._k_packed, + self._k_norms, + self._k_qjl_signs, + self._k_gamma, + self._v_packed, + self._v_norms, + ) + + def _set_tensor_fields(self, fields): + ( + self._k_packed, + self._k_norms, + self._k_qjl_signs, + self._k_gamma, + self._v_packed, + self._v_norms, + ) = fields + + def _grow(self, b: int, h: int, needed: int): + prev = self._idx + if self._k_packed is not None and needed <= self._k_packed.shape[2]: + return + n = ((needed + self.step - 1) // self.step) * self.step + shapes = ( + (b, h, n, self.k_pdim), + (b, h, n, 1), + (b, h, n, self.qjl_pdim), + (b, h, n, 1), + (b, h, n, self.v_pdim), + (b, h, n, 1), + ) + dtypes = (mx.uint32, mx.float32, mx.uint32, mx.float32, mx.uint32, mx.float32) + if self._k_packed is None: + self._set_tensor_fields( + tuple(mx.zeros(shape, dtype=dtype) for shape, dtype in zip(shapes, dtypes)) + ) + return + + def _extend(tensor, shape, dtype): + new = mx.zeros(shape, dtype=dtype) + new[:, :, :prev, :] = tensor[:, :, :prev, :] + return new + + self._set_tensor_fields( + tuple( + _extend(tensor, shape, dtype) + for tensor, shape, dtype in zip(self._tensor_fields(), shapes, dtypes) + ) + ) + + def _fetch_packed(self): + t = self._idx + keys = ( + self._k_packed[..., :t, :], + self._k_norms[..., :t, :], + self._k_qjl_signs[..., :t, :], + self._k_gamma[..., :t, :], + ) + values = (self._v_packed[..., :t, :], self._v_norms[..., :t, :]) + return keys, values + + def update_and_fetch(self, keys, values): + b, h, s, d = keys.shape + if d != self.head_dim: + raise ValueError(f"Expected head_dim={self.head_dim}, got {d}") + self._dtype = keys.dtype + prev = self._idx + self._grow(b, h, prev + s) + + k_pack, k_norm, k_sign, k_gamma, v_pack, v_norm = encode_kv( + keys, + values, + self._k_rotation, + self._v_rotation, + self._k_qjl, + self.k_bits, + self.v_bits, + ) + + self._k_packed[..., prev : prev + s, :] = k_pack + self._k_norms[..., prev : prev + s, :] = k_norm + self._k_qjl_signs[..., prev : prev + s, :] = k_sign + self._k_gamma[..., prev : prev + s, :] = k_gamma + self._v_packed[..., prev : prev + s, :] = v_pack + self._v_norms[..., prev : prev + s, :] = v_norm + + self.offset += s + self._idx += s + return self._fetch_packed() + + def prepare(self, *, left_padding=None, lengths=None, right_padding=None): + if left_padding is not None: + if self._k_packed is not None: + raise ValueError( + "Left padding can only be added to an empty BatchTurboQuantCache" + ) + left_padding = mx.array(left_padding) + self.left_padding += left_padding + self.offset -= left_padding + + if right_padding is not None and max(right_padding) > 0: + self._right_padding = mx.array(right_padding) + + def finalize(self): + if self._right_padding is not None: + padding = self._right_padding + self._set_tensor_fields( + tuple( + dynamic_roll(tensor, padding[:, None], axis=2) + for tensor in self._tensor_fields() + ) + ) + self.offset -= padding + self.left_padding += padding + self._right_padding = None + + def make_mask(self, N: int, return_array: bool = False, **kwargs): + return create_causal_mask( + N, offset=self._idx, left_padding=self.left_padding, **kwargs + ) + + def is_trimmable(self): + return True + + def trim(self, n): + n = min(self._idx, n) + self._idx -= n + self.offset -= n + return n + + def filter(self, batch_indices): + if self._k_packed is not None: + self._set_tensor_fields( + tuple(tensor[batch_indices] for tensor in self._tensor_fields()) + ) + self.offset = self.offset[batch_indices] + self.left_padding = self.left_padding[batch_indices] + + min_left_pad = self.left_padding.min().item() + if min_left_pad > 0: + if self._k_packed is not None: + self._set_tensor_fields( + tuple( + tensor[..., min_left_pad:, :] + for tensor in self._tensor_fields() + ) + ) + self._idx -= min_left_pad + self.left_padding -= min_left_pad + + def extend(self, other: "BatchAsymmetricTurboQuantCache"): + if self._k_packed is None and other._k_packed is None: + self.left_padding = mx.concatenate([self.left_padding, other.left_padding]) + self.offset = mx.concatenate([self.offset, other.offset]) + return + + max_idx = max(self._idx, other._idx) + h = None + if self._k_packed is not None: + _, h, _, _ = self._k_packed.shape + elif other._k_packed is not None: + _, h, _, _ = other._k_packed.shape + + max_size = 0 + if self._k_packed is not None: + max_size = max(max_size, self._k_packed.shape[2]) + if other._k_packed is not None: + max_size = max(max_size, other._k_packed.shape[2]) + + def pad(cache): + tensors = cache._tensor_fields() + if cache._k_packed is None: + bc = cache.offset.shape[0] + tensors = ( + mx.array([]).reshape(bc, h, 0, cache.k_pdim), + mx.array([]).reshape(bc, h, 0, 1), + mx.array([]).reshape(bc, h, 0, cache.qjl_pdim), + mx.array([]).reshape(bc, h, 0, 1), + mx.array([]).reshape(bc, h, 0, cache.v_pdim), + mx.array([]).reshape(bc, h, 0, 1), + ) + left = max_idx - cache._idx + right = max_size - tensors[0].shape[2] - left + if right < 0: + tensors = tuple(t[..., :right, :] for t in tensors) + right = 0 + if left != 0 or right != 0: + pad_spec = [(0, 0), (0, 0), (left, right), (0, 0)] + tensors = tuple(mx.pad(t, pad_spec) for t in tensors) + left_padding = cache.left_padding + left + return tensors, cache.offset, left_padding + + self_fields = pad(self) + other_fields = pad(other) + + self._set_tensor_fields( + tuple( + mx.concatenate([a, b]) + for a, b in zip(self_fields[0], other_fields[0]) + ) + ) + self.offset, self.left_padding = map( + mx.concatenate, zip(self_fields[1:], other_fields[1:]) + ) + self._idx = max_idx + + def extract(self, idx: int) -> AsymmetricTurboQuantCache: + cache = AsymmetricTurboQuantCache( + head_dim=self.head_dim, + k_bits=self.k_bits, + v_bits=self.v_bits, + seed=self.seed, + ) + if self._k_packed is None: + return cache + padding = self.left_padding[idx].item() + end = self._idx + cache._k_packed = mx.contiguous(self._k_packed[idx : idx + 1, :, padding:end, :]) + cache._k_norms = mx.contiguous(self._k_norms[idx : idx + 1, :, padding:end, :]) + cache._k_qjl_signs = mx.contiguous( + self._k_qjl_signs[idx : idx + 1, :, padding:end, :] + ) + cache._k_gamma = mx.contiguous(self._k_gamma[idx : idx + 1, :, padding:end, :]) + cache._v_packed = mx.contiguous(self._v_packed[idx : idx + 1, :, padding:end, :]) + cache._v_norms = mx.contiguous(self._v_norms[idx : idx + 1, :, padding:end, :]) + cache.offset = cache._k_packed.shape[2] + cache._dtype = self._dtype + return cache + + @classmethod + def merge(cls, caches: List[AsymmetricTurboQuantCache]): + if not caches: + raise ValueError("Cannot merge an empty cache list") + + first = caches[0] + lengths = [c.size() for c in caches] + max_length = max(lengths) + if max_length == 0: + return cls( + [0] * len(caches), + head_dim=first.head_dim, + k_bits=first.k_bits, + v_bits=first.v_bits, + seed=first.seed, + ) + + padding = [max_length - length for length in lengths] + b = len(caches) + h = next(c._k_packed.shape[1] for c in caches if not c.empty()) + k_pdim = first.k_pdim + v_pdim = first.v_pdim + qjl_pdim = first.qjl_pdim + + k_packed = mx.zeros((b, h, max_length, k_pdim), dtype=mx.uint32) + k_norms = mx.zeros((b, h, max_length, 1), dtype=mx.float32) + k_signs = mx.zeros((b, h, max_length, qjl_pdim), dtype=mx.uint32) + k_gamma = mx.zeros((b, h, max_length, 1), dtype=mx.float32) + v_packed = mx.zeros((b, h, max_length, v_pdim), dtype=mx.uint32) + v_norms = mx.zeros((b, h, max_length, 1), dtype=mx.float32) + + for i, (pad, cache) in enumerate(zip(padding, caches)): + if cache.empty(): + continue + t = cache.offset + sl = slice(pad, pad + t) + k_packed[i : i + 1, :, sl, :] = cache._k_packed[..., :t, :] + k_norms[i : i + 1, :, sl, :] = cache._k_norms[..., :t, :] + k_signs[i : i + 1, :, sl, :] = cache._k_qjl_signs[..., :t, :] + k_gamma[i : i + 1, :, sl, :] = cache._k_gamma[..., :t, :] + v_packed[i : i + 1, :, sl, :] = cache._v_packed[..., :t, :] + v_norms[i : i + 1, :, sl, :] = cache._v_norms[..., :t, :] + + batch = cls( + padding, + head_dim=first.head_dim, + k_bits=first.k_bits, + v_bits=first.v_bits, + seed=first.seed, + ) + batch._k_packed = k_packed + batch._k_norms = k_norms + batch._k_qjl_signs = k_signs + batch._k_gamma = k_gamma + batch._v_packed = v_packed + batch._v_norms = v_norms + batch._idx = max_length + batch.offset += max_length + return batch + + def size(self): + return self._idx + + def empty(self): + return self._k_packed is None + + @property + def nbytes(self): + if self._k_packed is None: + return 0 + t = self._idx + total = 0 + for tensor in self._tensor_fields(): + total += tensor[..., :t, :].nbytes + return total + + def __deepcopy__(self, memo): + cls = self.__class__ + obj = cls.__new__(cls) + memo[id(self)] = obj + for key, value in self.__dict__.items(): + if key == "_dtype": + setattr(obj, key, value) + elif key in ("_k_rotation", "_v_rotation", "_k_qjl"): + setattr(obj, key, value) + else: + setattr(obj, key, copy.deepcopy(value, memo)) + return obj \ No newline at end of file diff --git a/mlx_lm/turboquant/codebooks.py b/mlx_lm/turboquant/codebooks.py new file mode 100644 index 000000000..dd721df06 --- /dev/null +++ b/mlx_lm/turboquant/codebooks.py @@ -0,0 +1,109 @@ +# Copyright © 2025 Bonsai Demo contributors. + +"""Lloyd-Max scalar codebooks for TurboQuant on the unit-sphere Beta distribution. + +Centroids are precomputed for head_dim=128 (Bonsai GQA) via the Beta coordinate +distribution from TurboQuant (arXiv:2504.19874), Lemma 1. +""" + +from __future__ import annotations + +import mlx.core as mx + +HEAD_DIM = 128 + +# Lloyd-Max centroids for d=128, computed with MLX grid integration. +_CENTROIDS = { + 2: [ + -0.133163, + -0.039921, + 0.040372, + 0.13355, + ], + 3: [ + -0.189914, + -0.119752, + -0.067866, + -0.021982, + 0.022232, + 0.067866, + 0.119752, + 0.189914, + ], + 4: [ + -0.277515, + -0.225104, + -0.185344, + -0.149934, + -0.116278, + -0.082991, + -0.0497, + -0.016406, + 0.016889, + 0.050183, + 0.083261, + 0.116478, + 0.150416, + 0.185657, + 0.225104, + 0.277516, + ], +} + +_BOUNDARIES = { + 2: [-0.566582, -0.086542, 0.000226, 0.086961, 0.566775], + 3: [ + -0.594957, + -0.154833, + -0.093809, + -0.044924, + 0.000125, + 0.045049, + 0.093809, + 0.154833, + 0.594957, + ], + 4: [ + -0.638758, + -0.25131, + -0.205224, + -0.167639, + -0.133106, + -0.099634, + -0.066346, + -0.033053, + 0.000241, + 0.033536, + 0.066722, + 0.09987, + 0.133447, + 0.168036, + 0.20538, + 0.25131, + 0.638758, + ], +} + + +def get_codebook(bits: int) -> tuple[mx.array, mx.array]: + if bits not in _CENTROIDS: + raise ValueError( + f"Unsupported TurboQuant bit width {bits}. Supported: {sorted(_CENTROIDS)}" + ) + centroids = mx.array(_CENTROIDS[bits], dtype=mx.float32) + boundaries = mx.array(_BOUNDARIES[bits], dtype=mx.float32) + return centroids, boundaries + + +def quantize_coords(values: mx.array, bits: int) -> mx.array: + """Nearest-centroid quantization along the last axis.""" + centroids, _ = get_codebook(bits) + # values: (..., d), centroids: (2**bits,) + expanded = values[..., None] + dist = mx.abs(expanded - centroids) + return mx.argmin(dist, axis=-1).astype(mx.uint8) + + +def dequantize_coords(indices: mx.array, bits: int) -> mx.array: + centroids, _ = get_codebook(bits) + return centroids[indices.astype(mx.int32)] \ No newline at end of file diff --git a/mlx_lm/turboquant/factory.py b/mlx_lm/turboquant/factory.py new file mode 100644 index 000000000..e987f4e23 --- /dev/null +++ b/mlx_lm/turboquant/factory.py @@ -0,0 +1,62 @@ +# Copyright © 2025 Bonsai Demo contributors. + +from __future__ import annotations + +from typing import List, Optional + +import mlx.nn as nn + +from mlx_lm.turboquant.cache import AsymmetricTurboQuantCache + + +def _layer_count(model: nn.Module) -> int: + if hasattr(model, "layers"): + return len(model.layers) + if hasattr(model, "model") and hasattr(model.model, "layers"): + return len(model.model.layers) + raise ValueError("Could not determine transformer layer count for TurboQuant cache") + + +def make_turboquant_cache( + model: nn.Module, + *, + max_kv_size: Optional[int] = None, + k_bits: int = 4, + v_bits: int = 3, + fp16_layers: int = 4, + head_dim: int = 128, + seed: int = 42, +) -> List: + """Build a per-layer cache list with TurboQuant on middle layers. + + First and last ``fp16_layers`` use standard FP16 ``KVCache`` (or + ``RotatingKVCache`` when ``max_kv_size`` is set). Middle layers use + ``AsymmetricTurboQuantCache``. + """ + from mlx_lm.models.cache import KVCache, RotatingKVCache + + num_layers = _layer_count(model) + if fp16_layers < 0: + raise ValueError("fp16_layers must be >= 0") + if 2 * fp16_layers >= num_layers: + raise ValueError( + f"fp16_layers={fp16_layers} leaves no middle layers for TurboQuant " + f"(model has {num_layers} layers)" + ) + caches = [] + for i in range(num_layers): + if i < fp16_layers or i >= num_layers - fp16_layers: + if max_kv_size is not None: + caches.append(RotatingKVCache(max_size=max_kv_size, keep=4)) + else: + caches.append(KVCache()) + else: + caches.append( + AsymmetricTurboQuantCache( + head_dim=head_dim, + k_bits=k_bits, + v_bits=v_bits, + seed=seed + i, + ) + ) + return caches \ No newline at end of file diff --git a/mlx_lm/turboquant/kernels.py b/mlx_lm/turboquant/kernels.py new file mode 100644 index 000000000..9cdc0d94c --- /dev/null +++ b/mlx_lm/turboquant/kernels.py @@ -0,0 +1,2128 @@ +# Copyright © 2025 Bonsai Demo contributors. + +"""Metal kernels for TurboQuant encode/decode and fused attention.""" + +from __future__ import annotations + +import math +from functools import lru_cache +from typing import Optional + +import mlx.core as mx + +from mlx_lm.turboquant.codebooks import get_codebook +from mlx_lm.turboquant.packing import packed_dim +from mlx_lm.turboquant.qjl import qjl_packed_dim + +_QJL_SCALE = math.sqrt(math.pi / 2.0) + + +def metal_available() -> bool: + return mx.metal.is_available() + + +@lru_cache(maxsize=None) +def _decode_mse_kernel(bits: int, dim: int): + pdim = packed_dim(dim, bits) + n_cent = 2**bits + source = f""" + constexpr int DIM = {dim}; + constexpr int BITS = {bits}; + constexpr int PDIM = {pdim}; + constexpr int N_CENT = {n_cent}; + constexpr uint MASK = (1u << BITS) - 1u; + constexpr int BITS_PER_WORD = 32; + + uint n = thread_position_in_grid.x; + float norm = norms[n]; + uint packed_off = n * PDIM; + + float rotated[DIM]; + for (int i = 0; i < DIM; i++) {{ + int word = (i * BITS) / BITS_PER_WORD; + int shift = (i * BITS) % BITS_PER_WORD; + uint w = packed[packed_off + word]; + uint val = (w >> shift) & MASK; + if (shift + BITS > BITS_PER_WORD) {{ + int spill = BITS - (BITS_PER_WORD - shift); + uint w2 = packed[packed_off + word + 1]; + val = val | ((w2 & ((1u << spill) - 1u)) << (BITS_PER_WORD - shift)); + }} + rotated[i] = centroids[val]; + }} + + for (int j = 0; j < DIM; j++) {{ + float sum = 0.0f; + for (int i = 0; i < DIM; i++) {{ + sum += rotated[i] * rotation[i * DIM + j]; + }} + out[n * DIM + j] = static_cast(norm * sum); + }} + """ + return mx.fast.metal_kernel( + name=f"tq_decode_mse_{bits}b_{dim}d", + input_names=["packed", "norms", "rotation", "centroids"], + output_names=["out"], + source=source, + ) + + +@lru_cache(maxsize=None) +def _decode_prod_kernel(mse_bits: int, dim: int): + pdim = packed_dim(dim, mse_bits) + qjl_pdim = qjl_packed_dim(dim) + n_cent = 2**mse_bits + scale = _QJL_SCALE / dim + source = f""" + constexpr int DIM = {dim}; + constexpr int MSE_BITS = {mse_bits}; + constexpr int PDIM = {pdim}; + constexpr int QJL_PDIM = {qjl_pdim}; + constexpr int N_CENT = {n_cent}; + constexpr uint MSE_MASK = (1u << MSE_BITS) - 1u; + constexpr int BITS_PER_WORD = 32; + constexpr float QJL_SCALE = {scale}f; + + uint n = thread_position_in_grid.x; + float norm = norms[n]; + float gamma = qjl_gamma[n]; + uint packed_off = n * PDIM; + uint signs_off = n * QJL_PDIM; + + float mse_unit[DIM]; + for (int i = 0; i < DIM; i++) {{ + int word = (i * MSE_BITS) / BITS_PER_WORD; + int shift = (i * MSE_BITS) % BITS_PER_WORD; + uint w = mse_packed[packed_off + word]; + uint val = (w >> shift) & MSE_MASK; + if (shift + MSE_BITS > BITS_PER_WORD) {{ + int spill = MSE_BITS - (BITS_PER_WORD - shift); + uint w2 = mse_packed[packed_off + word + 1]; + val = val | ((w2 & ((1u << spill) - 1u)) << (BITS_PER_WORD - shift)); + }} + mse_unit[i] = centroids[val]; + }} + + float unit[DIM]; + for (int j = 0; j < DIM; j++) {{ + float sum = 0.0f; + for (int i = 0; i < DIM; i++) {{ + sum += mse_unit[i] * rotation[i * DIM + j]; + }} + unit[j] = sum; + }} + + float qjl[DIM]; + for (int j = 0; j < DIM; j++) {{ + float sum = 0.0f; + for (int i = 0; i < DIM; i++) {{ + int word = i / BITS_PER_WORD; + int shift = i % BITS_PER_WORD; + uint bit = (qjl_signs[signs_off + word] >> shift) & 1u; + float sign = bit ? 1.0f : -1.0f; + sum += sign * s_matrix[i * DIM + j]; + }} + qjl[j] = QJL_SCALE * gamma * sum; + }} + + for (int j = 0; j < DIM; j++) {{ + out[n * DIM + j] = static_cast(norm * (unit[j] + qjl[j])); + }} + """ + return mx.fast.metal_kernel( + name=f"tq_decode_prod_{mse_bits}b_{dim}d", + input_names=[ + "mse_packed", + "norms", + "qjl_signs", + "qjl_gamma", + "rotation", + "centroids", + "s_matrix", + ], + output_names=["out"], + source=source, + ) + + +def _launch_decode_mse( + packed: mx.array, + norms: mx.array, + rotation: mx.array, + bits: int, + dim: int, + dtype: mx.dtype, +) -> mx.array: + shape = packed.shape[:-1] + n = int(math.prod(shape)) if shape else packed.shape[0] + flat_packed = packed.reshape(n, packed.shape[-1]) + flat_norms = norms.reshape(n, 1).astype(mx.float32) + centroids, _ = get_codebook(bits) + kernel = _decode_mse_kernel(bits, dim) + out = kernel( + inputs=[flat_packed, flat_norms, rotation, centroids], + template=[("T", dtype)], + grid=(n, 1, 1), + threadgroup=(min(256, n), 1, 1), + output_shapes=[(n, dim)], + output_dtypes=[dtype], + stream=mx.gpu, + )[0] + return out.reshape(shape + (dim,)) + + +def _launch_decode_prod( + mse_packed: mx.array, + norms: mx.array, + qjl_signs: mx.array, + gamma: mx.array, + rotation: mx.array, + s_matrix: mx.array, + bits: int, + dim: int, + dtype: mx.dtype, +) -> mx.array: + mse_bits = bits - 1 + shape = mse_packed.shape[:-1] + n = int(math.prod(shape)) if shape else mse_packed.shape[0] + flat_mse = mse_packed.reshape(n, mse_packed.shape[-1]) + flat_signs = qjl_signs.reshape(n, qjl_signs.shape[-1]) + flat_norms = norms.reshape(n, 1).astype(mx.float32) + flat_gamma = gamma.reshape(n, 1).astype(mx.float32) + centroids, _ = get_codebook(mse_bits) + kernel = _decode_prod_kernel(mse_bits, dim) + out = kernel( + inputs=[ + flat_mse, + flat_norms, + flat_signs, + flat_gamma, + rotation, + centroids, + s_matrix, + ], + template=[("T", dtype)], + grid=(n, 1, 1), + threadgroup=(min(256, n), 1, 1), + output_shapes=[(n, dim)], + output_dtypes=[dtype], + stream=mx.gpu, + )[0] + return out.reshape(shape + (dim,)) + + +def decode_mse_metal( + packed: mx.array, + norms: mx.array, + rotation: mx.array, + bits: int, + dim: int, + dtype: Optional[mx.dtype] = None, +) -> mx.array: + if not metal_available(): + raise RuntimeError("Metal is not available") + if dtype is None: + dtype = mx.float32 + return _launch_decode_mse(packed, norms, rotation, bits, dim, dtype) + + +def decode_prod_metal( + mse_packed: mx.array, + norms: mx.array, + qjl_signs: mx.array, + gamma: mx.array, + rotation: mx.array, + s_matrix: mx.array, + bits: int, + dim: int, + dtype: Optional[mx.dtype] = None, +) -> mx.array: + if not metal_available(): + raise RuntimeError("Metal is not available") + if dtype is None: + dtype = mx.float32 + return _launch_decode_prod( + mse_packed, norms, qjl_signs, gamma, rotation, s_matrix, bits, dim, dtype + ) + + +def _pack_loop_source( + var: str, + bits_name: str, + *, + mask: str = "MASK", + packed_buf: str = "packed", + offset: str = "packed_off", +) -> str: + """Metal snippet: pack per-coordinate indices into uint32 words.""" + return f""" + for (int i = 0; i < DIM; i++) {{ + int word = (i * {bits_name}) / BITS_PER_WORD; + int shift = (i * {bits_name}) % BITS_PER_WORD; + uint val = {var}[i]; + uint chunk = (val & {mask}) << shift; + {packed_buf}[{offset} + word] |= chunk; + if (shift + {bits_name} > BITS_PER_WORD) {{ + int spill = {bits_name} - (BITS_PER_WORD - shift); + {packed_buf}[{offset} + word + 1] |= (val & {mask}) >> (BITS_PER_WORD - shift); + }} + }} + """ + + +@lru_cache(maxsize=None) +def _encode_mse_kernel(bits: int, dim: int): + pdim = packed_dim(dim, bits) + n_cent = 2**bits + pack_loop = _pack_loop_source("indices", "BITS") + source = f""" + constexpr int DIM = {dim}; + constexpr int BITS = {bits}; + constexpr int PDIM = {pdim}; + constexpr int N_CENT = {n_cent}; + constexpr uint MASK = (1u << BITS) - 1u; + constexpr int BITS_PER_WORD = 32; + + uint n = thread_position_in_grid.x; + uint packed_off = n * PDIM; + + for (int w = 0; w < PDIM; w++) {{ + packed[packed_off + w] = 0u; + }} + + float norm_sq = 0.0f; + for (int j = 0; j < DIM; j++) {{ + float v = static_cast(vectors[n * DIM + j]); + norm_sq += v * v; + }} + float norm = sqrt(norm_sq); + norms[n] = norm; + float inv_norm = 1.0f / max(norm, 1e-8f); + + float unit[DIM]; + for (int j = 0; j < DIM; j++) {{ + unit[j] = static_cast(vectors[n * DIM + j]) * inv_norm; + }} + + uint indices[DIM]; + for (int i = 0; i < DIM; i++) {{ + float rot = 0.0f; + for (int j = 0; j < DIM; j++) {{ + rot += unit[j] * rotation[i * DIM + j]; + }} + uint best = 0u; + float best_d = abs(rot - centroids[0]); + for (int c = 1; c < N_CENT; c++) {{ + float d = abs(rot - centroids[c]); + if (d < best_d) {{ + best_d = d; + best = uint(c); + }} + }} + indices[i] = best; + }} + + {pack_loop} + """ + return mx.fast.metal_kernel( + name=f"tq_encode_mse_{bits}b_{dim}d", + input_names=["vectors", "rotation", "centroids"], + output_names=["packed", "norms"], + source=source, + ) + + +@lru_cache(maxsize=None) +def _encode_prod_kernel(bits: int, dim: int): + mse_bits = bits - 1 + pdim = packed_dim(dim, mse_bits) + qjl_pdim = qjl_packed_dim(dim) + n_cent = 2**mse_bits + mse_pack = _pack_loop_source( + "mse_indices", + "MSE_BITS", + mask="MSE_MASK", + packed_buf="mse_packed", + offset="packed_off", + ) + source = f""" + constexpr int DIM = {dim}; + constexpr int MSE_BITS = {mse_bits}; + constexpr int PDIM = {pdim}; + constexpr int QJL_PDIM = {qjl_pdim}; + constexpr int N_CENT = {n_cent}; + constexpr uint MSE_MASK = (1u << MSE_BITS) - 1u; + constexpr int BITS_PER_WORD = 32; + + uint n = thread_position_in_grid.x; + uint packed_off = n * PDIM; + uint signs_off = n * QJL_PDIM; + + for (int w = 0; w < PDIM; w++) {{ + mse_packed[packed_off + w] = 0u; + }} + for (int w = 0; w < QJL_PDIM; w++) {{ + qjl_signs[signs_off + w] = 0u; + }} + + float norm_sq = 0.0f; + for (int j = 0; j < DIM; j++) {{ + float v = static_cast(vectors[n * DIM + j]); + norm_sq += v * v; + }} + float norm = sqrt(norm_sq); + norms[n] = norm; + float inv_norm = 1.0f / max(norm, 1e-8f); + + float unit[DIM]; + for (int j = 0; j < DIM; j++) {{ + unit[j] = static_cast(vectors[n * DIM + j]) * inv_norm; + }} + + uint mse_indices[DIM]; + float mse_rot[DIM]; + for (int i = 0; i < DIM; i++) {{ + float rot = 0.0f; + for (int j = 0; j < DIM; j++) {{ + rot += unit[j] * rotation[i * DIM + j]; + }} + uint best = 0u; + float best_d = abs(rot - centroids[0]); + for (int c = 1; c < N_CENT; c++) {{ + float d = abs(rot - centroids[c]); + if (d < best_d) {{ + best_d = d; + best = uint(c); + }} + }} + mse_indices[i] = best; + mse_rot[i] = centroids[best]; + }} + + {mse_pack} + + float mse_unit[DIM]; + for (int j = 0; j < DIM; j++) {{ + float sum = 0.0f; + for (int i = 0; i < DIM; i++) {{ + sum += mse_rot[i] * rotation[i * DIM + j]; + }} + mse_unit[j] = sum; + }} + + float residual[DIM]; + float gamma_sq = 0.0f; + for (int j = 0; j < DIM; j++) {{ + residual[j] = unit[j] - mse_unit[j]; + gamma_sq += residual[j] * residual[j]; + }} + qjl_gamma[n] = sqrt(gamma_sq); + + for (int i = 0; i < DIM; i++) {{ + float proj = 0.0f; + for (int j = 0; j < DIM; j++) {{ + proj += residual[j] * s_matrix[i * DIM + j]; + }} + uint bit = proj >= 0.0f ? 1u : 0u; + int word = i / BITS_PER_WORD; + int shift = i % BITS_PER_WORD; + qjl_signs[signs_off + word] |= bit << shift; + }} + """ + return mx.fast.metal_kernel( + name=f"tq_encode_prod_{bits}b_{dim}d", + input_names=["vectors", "rotation", "centroids", "s_matrix"], + output_names=["mse_packed", "norms", "qjl_signs", "qjl_gamma"], + source=source, + ) + + +_ENCODE_SIMD_THREADS = 32 + + +@lru_cache(maxsize=None) +def _encode_mse_simd_kernel(bits: int, dim: int): + pdim = packed_dim(dim, bits) + n_cent = 2**bits + source = f""" + constexpr int DIM = {dim}; + constexpr int BITS = {bits}; + constexpr int PDIM = {pdim}; + constexpr int N_CENT = {n_cent}; + constexpr uint MASK = (1u << BITS) - 1u; + constexpr int BITS_PER_WORD = 32; + constexpr int BD = 32; + constexpr int per_thread = DIM / BD; + + uint tg = threadgroup_position_in_grid.x; + uint simd_lid = thread_index_in_simdgroup; + + threadgroup float unit_shmem[DIM]; + threadgroup uint indices_shmem[DIM]; + threadgroup float norm_shmem; + + float norm_sq = 0.0f; + for (int t = 0; t < per_thread; t++) {{ + int j = simd_lid * per_thread + t; + float v = vectors[tg * DIM + j]; + norm_sq += v * v; + }} + norm_sq = simd_sum(norm_sq); + if (simd_lid == 0) {{ + norm_shmem = sqrt(norm_sq); + norms[tg] = norm_shmem; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + float inv_norm = 1.0f / max(norm_shmem, 1e-8f); + for (int t = 0; t < per_thread; t++) {{ + int j = simd_lid * per_thread + t; + unit_shmem[j] = vectors[tg * DIM + j] * inv_norm; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int t = 0; t < per_thread; t++) {{ + int i = simd_lid * per_thread + t; + float rot = 0.0f; + for (int j = 0; j < DIM; j++) {{ + rot += unit_shmem[j] * rotation[i * DIM + j]; + }} + uint best = 0u; + float best_d = abs(rot - centroids[0]); + for (int c = 1; c < N_CENT; c++) {{ + float d = abs(rot - centroids[c]); + if (d < best_d) {{ + best_d = d; + best = uint(c); + }} + }} + indices_shmem[i] = best; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (simd_lid == 0) {{ + uint packed_off = tg * PDIM; + for (int w = 0; w < PDIM; w++) {{ + packed[packed_off + w] = 0u; + }} + for (int i = 0; i < DIM; i++) {{ + int word = (i * BITS) / BITS_PER_WORD; + int shift = (i * BITS) % BITS_PER_WORD; + uint val = indices_shmem[i]; + packed[packed_off + word] |= (val & MASK) << shift; + if (shift + BITS > BITS_PER_WORD) {{ + int spill = BITS - (BITS_PER_WORD - shift); + packed[packed_off + word + 1] |= (val & MASK) >> (BITS_PER_WORD - shift); + }} + }} + }} + """ + return mx.fast.metal_kernel( + name=f"tq_encode_mse_simd_{bits}b_{dim}d", + input_names=["vectors", "rotation", "centroids"], + output_names=["packed", "norms"], + source=source, + ) + + +@lru_cache(maxsize=None) +def _encode_prod_simd_kernel(bits: int, dim: int): + mse_bits = bits - 1 + pdim = packed_dim(dim, mse_bits) + qjl_pdim = qjl_packed_dim(dim) + n_cent = 2**mse_bits + source = f""" + constexpr int DIM = {dim}; + constexpr int MSE_BITS = {mse_bits}; + constexpr int PDIM = {pdim}; + constexpr int QJL_PDIM = {qjl_pdim}; + constexpr int N_CENT = {n_cent}; + constexpr uint MSE_MASK = (1u << MSE_BITS) - 1u; + constexpr int BITS_PER_WORD = 32; + constexpr int BD = 32; + constexpr int per_thread = DIM / BD; + + uint tg = threadgroup_position_in_grid.x; + uint simd_lid = thread_index_in_simdgroup; + + threadgroup float unit_shmem[DIM]; + threadgroup float mse_rot_shmem[DIM]; + threadgroup float residual_shmem[DIM]; + threadgroup uint mse_indices_shmem[DIM]; + threadgroup uint qjl_bits_shmem[DIM]; + threadgroup float norm_shmem; + + float norm_sq = 0.0f; + for (int t = 0; t < per_thread; t++) {{ + int j = simd_lid * per_thread + t; + float v = vectors[tg * DIM + j]; + norm_sq += v * v; + }} + norm_sq = simd_sum(norm_sq); + if (simd_lid == 0) {{ + norm_shmem = sqrt(norm_sq); + norms[tg] = norm_shmem; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + float inv_norm = 1.0f / max(norm_shmem, 1e-8f); + for (int t = 0; t < per_thread; t++) {{ + int j = simd_lid * per_thread + t; + unit_shmem[j] = vectors[tg * DIM + j] * inv_norm; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int t = 0; t < per_thread; t++) {{ + int i = simd_lid * per_thread + t; + float rot = 0.0f; + for (int j = 0; j < DIM; j++) {{ + rot += unit_shmem[j] * rotation[i * DIM + j]; + }} + uint best = 0u; + float best_d = abs(rot - centroids[0]); + for (int c = 1; c < N_CENT; c++) {{ + float d = abs(rot - centroids[c]); + if (d < best_d) {{ + best_d = d; + best = uint(c); + }} + }} + mse_indices_shmem[i] = best; + mse_rot_shmem[i] = centroids[best]; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + float res_part = 0.0f; + for (int t = 0; t < per_thread; t++) {{ + int j = simd_lid * per_thread + t; + float mse_u = 0.0f; + for (int i = 0; i < DIM; i++) {{ + mse_u += mse_rot_shmem[i] * rotation[i * DIM + j]; + }} + float r = unit_shmem[j] - mse_u; + residual_shmem[j] = r; + res_part += r * r; + }} + float gamma_sq = simd_sum(res_part); + if (simd_lid == 0) {{ + qjl_gamma[tg] = sqrt(gamma_sq); + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int t = 0; t < per_thread; t++) {{ + int i = simd_lid * per_thread + t; + float proj = 0.0f; + for (int j = 0; j < DIM; j++) {{ + proj += residual_shmem[j] * s_matrix[i * DIM + j]; + }} + qjl_bits_shmem[i] = proj >= 0.0f ? 1u : 0u; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (simd_lid == 0) {{ + uint packed_off = tg * PDIM; + uint signs_off = tg * QJL_PDIM; + for (int w = 0; w < PDIM; w++) {{ + mse_packed[packed_off + w] = 0u; + }} + for (int w = 0; w < QJL_PDIM; w++) {{ + qjl_signs[signs_off + w] = 0u; + }} + for (int i = 0; i < DIM; i++) {{ + int word = (i * MSE_BITS) / BITS_PER_WORD; + int shift = (i * MSE_BITS) % BITS_PER_WORD; + uint val = mse_indices_shmem[i]; + mse_packed[packed_off + word] |= (val & MSE_MASK) << shift; + if (shift + MSE_BITS > BITS_PER_WORD) {{ + mse_packed[packed_off + word + 1] |= (val & MSE_MASK) >> (BITS_PER_WORD - shift); + }} + int s_word = i / BITS_PER_WORD; + int s_shift = i % BITS_PER_WORD; + qjl_signs[signs_off + s_word] |= qjl_bits_shmem[i] << s_shift; + }} + }} + """ + return mx.fast.metal_kernel( + name=f"tq_encode_prod_simd_{bits}b_{dim}d", + input_names=["vectors", "rotation", "centroids", "s_matrix"], + output_names=["mse_packed", "norms", "qjl_signs", "qjl_gamma"], + source=source, + ) + + +@lru_cache(maxsize=None) +def _encode_kv_fused_kernel(k_bits: int, v_bits: int, dim: int): + """Fused K (prod) + V (mse) encode in one dispatch per token.""" + mse_bits = k_bits - 1 + k_pdim = packed_dim(dim, mse_bits) + v_pdim = packed_dim(dim, v_bits) + qjl_pdim = qjl_packed_dim(dim) + k_n_cent = 2**mse_bits + v_n_cent = 2**v_bits + source = f""" + constexpr int DIM = {dim}; + constexpr int MSE_BITS = {mse_bits}; + constexpr int V_BITS = {v_bits}; + constexpr int K_PDIM = {k_pdim}; + constexpr int V_PDIM = {v_pdim}; + constexpr int QJL_PDIM = {qjl_pdim}; + constexpr int K_N_CENT = {k_n_cent}; + constexpr int V_N_CENT = {v_n_cent}; + constexpr uint MSE_MASK = (1u << MSE_BITS) - 1u; + constexpr uint V_MASK = (1u << V_BITS) - 1u; + constexpr int BITS_PER_WORD = 32; + constexpr int BD = 32; + constexpr int per_thread = DIM / BD; + + uint tg = threadgroup_position_in_grid.x; + uint simd_lid = thread_index_in_simdgroup; + + threadgroup float k_unit_shmem[DIM]; + threadgroup float v_unit_shmem[DIM]; + threadgroup float k_mse_rot_shmem[DIM]; + threadgroup float k_residual_shmem[DIM]; + threadgroup uint k_mse_indices_shmem[DIM]; + threadgroup uint k_qjl_bits_shmem[DIM]; + threadgroup uint v_indices_shmem[DIM]; + threadgroup float k_norm_shmem; + threadgroup float v_norm_shmem; + + float k_norm_sq = 0.0f; + float v_norm_sq = 0.0f; + for (int t = 0; t < per_thread; t++) {{ + int j = simd_lid * per_thread + t; + float kv = keys[tg * DIM + j]; + float vv = values[tg * DIM + j]; + k_norm_sq += kv * kv; + v_norm_sq += vv * vv; + }} + k_norm_sq = simd_sum(k_norm_sq); + v_norm_sq = simd_sum(v_norm_sq); + if (simd_lid == 0) {{ + k_norm_shmem = sqrt(k_norm_sq); + v_norm_shmem = sqrt(v_norm_sq); + k_norms[tg] = k_norm_shmem; + v_norms[tg] = v_norm_shmem; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + float k_inv = 1.0f / max(k_norm_shmem, 1e-8f); + float v_inv = 1.0f / max(v_norm_shmem, 1e-8f); + for (int t = 0; t < per_thread; t++) {{ + int j = simd_lid * per_thread + t; + k_unit_shmem[j] = keys[tg * DIM + j] * k_inv; + v_unit_shmem[j] = values[tg * DIM + j] * v_inv; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int t = 0; t < per_thread; t++) {{ + int i = simd_lid * per_thread + t; + float k_rot = 0.0f; + for (int j = 0; j < DIM; j++) {{ + k_rot += k_unit_shmem[j] * k_rotation[i * DIM + j]; + }} + uint k_best = 0u; + float k_best_d = abs(k_rot - k_centroids[0]); + for (int c = 1; c < K_N_CENT; c++) {{ + float d = abs(k_rot - k_centroids[c]); + if (d < k_best_d) {{ + k_best_d = d; + k_best = uint(c); + }} + }} + k_mse_indices_shmem[i] = k_best; + k_mse_rot_shmem[i] = k_centroids[k_best]; + + float v_rot = 0.0f; + for (int j = 0; j < DIM; j++) {{ + v_rot += v_unit_shmem[j] * v_rotation[i * DIM + j]; + }} + uint v_best = 0u; + float v_best_d = abs(v_rot - v_centroids[0]); + for (int c = 1; c < V_N_CENT; c++) {{ + float d = abs(v_rot - v_centroids[c]); + if (d < v_best_d) {{ + v_best_d = d; + v_best = uint(c); + }} + }} + v_indices_shmem[i] = v_best; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + float res_part = 0.0f; + for (int t = 0; t < per_thread; t++) {{ + int j = simd_lid * per_thread + t; + float mse_u = 0.0f; + for (int i = 0; i < DIM; i++) {{ + mse_u += k_mse_rot_shmem[i] * k_rotation[i * DIM + j]; + }} + float r = k_unit_shmem[j] - mse_u; + k_residual_shmem[j] = r; + res_part += r * r; + }} + float gamma_sq = simd_sum(res_part); + if (simd_lid == 0) {{ + k_gamma[tg] = sqrt(gamma_sq); + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int t = 0; t < per_thread; t++) {{ + int i = simd_lid * per_thread + t; + float proj = 0.0f; + for (int j = 0; j < DIM; j++) {{ + proj += k_residual_shmem[j] * k_qjl[i * DIM + j]; + }} + k_qjl_bits_shmem[i] = proj >= 0.0f ? 1u : 0u; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (simd_lid == 0) {{ + uint k_off = tg * K_PDIM; + uint v_off = tg * V_PDIM; + uint signs_off = tg * QJL_PDIM; + for (int w = 0; w < K_PDIM; w++) {{ + k_packed[k_off + w] = 0u; + }} + for (int w = 0; w < V_PDIM; w++) {{ + v_packed[v_off + w] = 0u; + }} + for (int w = 0; w < QJL_PDIM; w++) {{ + k_qjl_signs[signs_off + w] = 0u; + }} + for (int i = 0; i < DIM; i++) {{ + int k_word = (i * MSE_BITS) / BITS_PER_WORD; + int k_shift = (i * MSE_BITS) % BITS_PER_WORD; + uint k_val = k_mse_indices_shmem[i]; + k_packed[k_off + k_word] |= (k_val & MSE_MASK) << k_shift; + if (k_shift + MSE_BITS > BITS_PER_WORD) {{ + k_packed[k_off + k_word + 1] |= (k_val & MSE_MASK) >> (BITS_PER_WORD - k_shift); + }} + int v_word = (i * V_BITS) / BITS_PER_WORD; + int v_shift = (i * V_BITS) % BITS_PER_WORD; + uint v_val = v_indices_shmem[i]; + v_packed[v_off + v_word] |= (v_val & V_MASK) << v_shift; + if (v_shift + V_BITS > BITS_PER_WORD) {{ + v_packed[v_off + v_word + 1] |= (v_val & V_MASK) >> (BITS_PER_WORD - v_shift); + }} + int s_word = i / BITS_PER_WORD; + int s_shift = i % BITS_PER_WORD; + k_qjl_signs[signs_off + s_word] |= k_qjl_bits_shmem[i] << s_shift; + }} + }} + """ + return mx.fast.metal_kernel( + name=f"tq_encode_kv_fused_{k_bits}k_{v_bits}v_{dim}d", + input_names=[ + "keys", + "values", + "k_rotation", + "v_rotation", + "k_centroids", + "v_centroids", + "k_qjl", + ], + output_names=[ + "k_packed", + "k_norms", + "k_qjl_signs", + "k_gamma", + "v_packed", + "v_norms", + ], + source=source, + ) + + +def _launch_encode_mse( + vectors: mx.array, + rotation: mx.array, + bits: int, + dim: int, +) -> tuple[mx.array, mx.array]: + shape = vectors.shape[:-1] + n = int(math.prod(shape)) if shape else vectors.shape[0] + pdim = packed_dim(dim, bits) + flat = vectors.astype(mx.float32).reshape(n, dim) + centroids, _ = get_codebook(bits) + kernel = _encode_mse_simd_kernel(bits, dim) + packed, norms = kernel( + inputs=[flat, rotation, centroids], + template=[("T", mx.float32)], + grid=(n * _ENCODE_SIMD_THREADS, 1, 1), + threadgroup=(_ENCODE_SIMD_THREADS, 1, 1), + output_shapes=[(n, pdim), (n, 1)], + output_dtypes=[mx.uint32, mx.float32], + stream=mx.gpu, + ) + return packed.reshape(shape + (pdim,)), norms.reshape(shape + (1,)) + + +def _launch_encode_prod( + vectors: mx.array, + rotation: mx.array, + s_matrix: mx.array, + bits: int, + dim: int, +) -> tuple[mx.array, mx.array, mx.array, mx.array]: + mse_bits = bits - 1 + shape = vectors.shape[:-1] + n = int(math.prod(shape)) if shape else vectors.shape[0] + pdim = packed_dim(dim, mse_bits) + qpdim = qjl_packed_dim(dim) + flat = vectors.astype(mx.float32).reshape(n, dim) + centroids, _ = get_codebook(mse_bits) + kernel = _encode_prod_simd_kernel(bits, dim) + mse_packed, norms, qjl_signs, qjl_gamma = kernel( + inputs=[flat, rotation, centroids, s_matrix], + template=[("T", mx.float32)], + grid=(n * _ENCODE_SIMD_THREADS, 1, 1), + threadgroup=(_ENCODE_SIMD_THREADS, 1, 1), + output_shapes=[(n, pdim), (n, 1), (n, qpdim), (n, 1)], + output_dtypes=[mx.uint32, mx.float32, mx.uint32, mx.float32], + stream=mx.gpu, + ) + return ( + mse_packed.reshape(shape + (pdim,)), + norms.reshape(shape + (1,)), + qjl_signs.reshape(shape + (qpdim,)), + qjl_gamma.reshape(shape + (1,)), + ) + + +def _launch_encode_kv( + keys: mx.array, + values: mx.array, + k_rotation: mx.array, + v_rotation: mx.array, + k_qjl: mx.array, + k_bits: int, + v_bits: int, + dim: int, +) -> tuple[mx.array, mx.array, mx.array, mx.array, mx.array, mx.array]: + shape = keys.shape[:-1] + n = int(math.prod(shape)) if shape else keys.shape[0] + mse_bits = k_bits - 1 + k_pdim = packed_dim(dim, mse_bits) + v_pdim = packed_dim(dim, v_bits) + qpdim = qjl_packed_dim(dim) + flat_k = keys.astype(mx.float32).reshape(n, dim) + flat_v = values.astype(mx.float32).reshape(n, dim) + k_centroids, _ = get_codebook(mse_bits) + v_centroids, _ = get_codebook(v_bits) + kernel = _encode_kv_fused_kernel(k_bits, v_bits, dim) + k_packed, k_norms, k_signs, k_gamma, v_packed, v_norms = kernel( + inputs=[ + flat_k, + flat_v, + k_rotation, + v_rotation, + k_centroids, + v_centroids, + k_qjl, + ], + template=[("T", mx.float32)], + grid=(n * _ENCODE_SIMD_THREADS, 1, 1), + threadgroup=(_ENCODE_SIMD_THREADS, 1, 1), + output_shapes=[ + (n, k_pdim), + (n, 1), + (n, qpdim), + (n, 1), + (n, v_pdim), + (n, 1), + ], + output_dtypes=[ + mx.uint32, + mx.float32, + mx.uint32, + mx.float32, + mx.uint32, + mx.float32, + ], + stream=mx.gpu, + ) + return ( + k_packed.reshape(shape + (k_pdim,)), + k_norms.reshape(shape + (1,)), + k_signs.reshape(shape + (qpdim,)), + k_gamma.reshape(shape + (1,)), + v_packed.reshape(shape + (v_pdim,)), + v_norms.reshape(shape + (1,)), + ) + + +def encode_mse_metal( + vectors: mx.array, + rotation: mx.array, + bits: int, + dim: int, +) -> tuple[mx.array, mx.array]: + if not metal_available(): + raise RuntimeError("Metal is not available") + return _launch_encode_mse(vectors, rotation, bits, dim) + + +def encode_prod_metal( + vectors: mx.array, + rotation: mx.array, + s_matrix: mx.array, + bits: int, + dim: int, +) -> tuple[mx.array, mx.array, mx.array, mx.array]: + if not metal_available(): + raise RuntimeError("Metal is not available") + return _launch_encode_prod(vectors, rotation, s_matrix, bits, dim) + + +def encode_kv_metal( + keys: mx.array, + values: mx.array, + k_rotation: mx.array, + v_rotation: mx.array, + k_qjl: mx.array, + k_bits: int, + v_bits: int, + dim: int, +) -> tuple[mx.array, mx.array, mx.array, mx.array, mx.array, mx.array]: + """Fused K+V encode in one Metal dispatch.""" + if not metal_available(): + raise RuntimeError("Metal is not available") + return _launch_encode_kv( + keys, values, k_rotation, v_rotation, k_qjl, k_bits, v_bits, dim + ) + + +@lru_cache(maxsize=None) +def _qk_scores_kernel(k_bits: int, dim: int): + """Fused Q·K for TurboQuant_prod keys. + + Math (row-vector convention, ``R`` from ``make_rotation_matrix``): + unit_mse = c @ R, c_i = Lloyd-Max centroid of rotated coord i + q·unit_mse = (q @ R.T)·c = sum_i c_i * (q @ R.T)_i + + qjl = (sqrt(pi/2)/d) * gamma * (signs @ S), signs in {{±1}} + q·qjl = (sqrt(pi/2)/d) * gamma * sum_i sign_i * (q @ S.T)_i + + score = norm * (q·unit_mse + q·qjl) * attn_scale + """ + mse_bits = k_bits - 1 + pdim = packed_dim(dim, mse_bits) + qjl_pdim = qjl_packed_dim(dim) + n_cent = 2**mse_bits + qjl_scale = _QJL_SCALE / dim + source = f""" + constexpr int DIM = {dim}; + constexpr int MSE_BITS = {mse_bits}; + constexpr int PDIM = {pdim}; + constexpr int QJL_PDIM = {qjl_pdim}; + constexpr int N_CENT = {n_cent}; + constexpr uint MSE_MASK = (1u << MSE_BITS) - 1u; + constexpr int BITS_PER_WORD = 32; + constexpr float QJL_SCALE = {qjl_scale}f; + + uint gid = thread_position_in_grid.x; + + uint tmp = gid; + uint s = tmp % seq_len; + tmp /= seq_len; + uint l = tmp % query_len; + tmp /= query_len; + uint h_q = tmp % n_q_heads; + uint b = tmp / n_q_heads; + uint h_kv = h_q / n_repeats; + + uint q_base = ((b * n_q_heads + h_q) * query_len + l) * DIM; + uint kv_slot = (b * n_kv_heads + h_kv) * seq_len + s; + uint k_packed_off = kv_slot * PDIM; + uint k_signs_off = kv_slot * QJL_PDIM; + + float kn = k_norms[kv_slot]; + float gamma = k_gamma[kv_slot]; + + float qv[DIM]; + for (int j = 0; j < DIM; j++) {{ + qv[j] = static_cast(queries[q_base + j]); + }} + + float q_rot[DIM]; + float q_s[DIM]; + for (int i = 0; i < DIM; i++) {{ + float rot = 0.0f; + float ps = 0.0f; + for (int j = 0; j < DIM; j++) {{ + rot += qv[j] * rotation[i * DIM + j]; + ps += qv[j] * s_matrix[i * DIM + j]; + }} + q_rot[i] = rot; + q_s[i] = ps; + }} + + float mse_dot = 0.0f; + for (int i = 0; i < DIM; i++) {{ + int word = (i * MSE_BITS) / BITS_PER_WORD; + int shift = (i * MSE_BITS) % BITS_PER_WORD; + uint w = mse_packed[k_packed_off + word]; + uint val = (w >> shift) & MSE_MASK; + if (shift + MSE_BITS > BITS_PER_WORD) {{ + int spill = MSE_BITS - (BITS_PER_WORD - shift); + uint w2 = mse_packed[k_packed_off + word + 1]; + val = val | ((w2 & ((1u << spill) - 1u)) << (BITS_PER_WORD - shift)); + }} + mse_dot += centroids[val] * q_rot[i]; + }} + + float qjl_dot = 0.0f; + for (int i = 0; i < DIM; i++) {{ + int s_word = i / BITS_PER_WORD; + int s_shift = i % BITS_PER_WORD; + uint bit = (qjl_signs[k_signs_off + s_word] >> s_shift) & 1u; + float sign = bit ? 1.0f : -1.0f; + qjl_dot += sign * q_s[i]; + }} + + scores[gid] = static_cast(kn * (mse_dot + QJL_SCALE * gamma * qjl_dot)); + """ + return mx.fast.metal_kernel( + name=f"tq_qk_scores_{k_bits}b_{dim}d", + input_names=[ + "queries", + "mse_packed", + "k_norms", + "qjl_signs", + "k_gamma", + "rotation", + "centroids", + "s_matrix", + ], + output_names=["scores"], + source=source, + ) + + +def qk_scores_metal( + queries: mx.array, + mse_packed: mx.array, + k_norms: mx.array, + qjl_signs: mx.array, + k_gamma: mx.array, + rotation: mx.array, + s_matrix: mx.array, + k_bits: int, + dim: int, + scale: float, +) -> mx.array: + if not metal_available(): + raise RuntimeError("Metal is not available") + B, n_q_heads, L, _ = queries.shape + S = mse_packed.shape[2] + n_kv_heads = mse_packed.shape[1] + n_repeats = n_q_heads // n_kv_heads + mse_bits = k_bits - 1 + pdim = packed_dim(dim, mse_bits) + qpdim = qjl_packed_dim(dim) + centroids, _ = get_codebook(mse_bits) + # Row-contiguous flatten (bitlinear / QuantizedKVCache pattern). + flat_q = mx.contiguous(queries.reshape(-1, dim)) + flat_packed = mx.contiguous(mse_packed.reshape(-1, pdim)) + flat_signs = mx.contiguous(qjl_signs.reshape(-1, qpdim)) + flat_norms = mx.contiguous(k_norms.reshape(-1).astype(mx.float32)) + flat_gamma = mx.contiguous(k_gamma.reshape(-1).astype(mx.float32)) + n_scores = B * n_q_heads * L * S + kernel = _qk_scores_kernel(k_bits, dim) + scores = kernel( + inputs=[ + flat_q, + flat_packed, + flat_norms, + flat_signs, + flat_gamma, + rotation, + centroids, + s_matrix, + ], + template=[ + ("T", mx.float32), + ("n_q_heads", n_q_heads), + ("n_kv_heads", n_kv_heads), + ("n_repeats", n_repeats), + ("query_len", L), + ("seq_len", S), + ], + grid=(n_scores, 1, 1), + threadgroup=(1, 1, 1), + output_shapes=[(n_scores,)], + output_dtypes=[mx.float32], + stream=mx.gpu, + )[0] + return (scores * scale).reshape(B, n_q_heads, L, S).astype(queries.dtype) + + +@lru_cache(maxsize=None) +def _av_weighted_sum_kernel(v_bits: int, dim: int): + """Fused attn @ V for TurboQuant_mse values without materializing V. + + Per output coord (b, h_q, l, d): + out = sum_s attn[b,h_q,l,s] * norm[s] * sum_i c_i * R[i,d] + where c_i is the Lloyd-Max centroid from packed indices at coord i. + """ + pdim = packed_dim(dim, v_bits) + n_cent = 2**v_bits + source = f""" + constexpr int DIM = {dim}; + constexpr int BITS = {v_bits}; + constexpr int PDIM = {pdim}; + constexpr int N_CENT = {n_cent}; + constexpr uint MASK = (1u << BITS) - 1u; + constexpr int BITS_PER_WORD = 32; + + uint gid = thread_position_in_grid.x; + + uint tmp = gid; + uint d = tmp % DIM; + tmp /= DIM; + uint l = tmp % query_len; + tmp /= query_len; + uint h_q = tmp % n_q_heads; + uint b = tmp / n_q_heads; + uint h_kv = h_q / n_repeats; + + uint attn_base = ((b * n_q_heads + h_q) * query_len + l) * seq_len; + uint kv_base = (b * n_kv_heads + h_kv) * seq_len; + + float acc = 0.0f; + for (uint s = 0; s < seq_len; s++) {{ + float w = static_cast(attn[attn_base + s]); + uint kv_slot = kv_base + s; + float vn = v_norms[kv_slot]; + uint packed_off = kv_slot * PDIM; + + float vd = 0.0f; + for (int i = 0; i < DIM; i++) {{ + int word = (i * BITS) / BITS_PER_WORD; + int shift = (i * BITS) % BITS_PER_WORD; + uint w_packed = v_packed[packed_off + word]; + uint val = (w_packed >> shift) & MASK; + if (shift + BITS > BITS_PER_WORD) {{ + int spill = BITS - (BITS_PER_WORD - shift); + uint w2 = v_packed[packed_off + word + 1]; + val = val | ((w2 & ((1u << spill) - 1u)) << (BITS_PER_WORD - shift)); + }} + vd += centroids[val] * rotation[i * DIM + d]; + }} + acc += w * vn * vd; + }} + + out[gid] = static_cast(acc); + """ + return mx.fast.metal_kernel( + name=f"tq_av_weighted_sum_{v_bits}b_{dim}d", + input_names=["attn", "v_packed", "v_norms", "rotation", "centroids"], + output_names=["out"], + source=source, + ) + + +def av_weighted_sum_metal( + attn: mx.array, + v_packed: mx.array, + v_norms: mx.array, + rotation: mx.array, + v_bits: int, + dim: int, + dtype: Optional[mx.dtype] = None, +) -> mx.array: + """Weighted sum of dequantized TurboQuant_mse values without full V decode.""" + if not metal_available(): + raise RuntimeError("Metal is not available") + if dtype is None: + dtype = attn.dtype + B, n_q_heads, L, S = attn.shape + n_kv_heads = v_packed.shape[1] + n_repeats = n_q_heads // n_kv_heads + pdim = packed_dim(dim, v_bits) + centroids, _ = get_codebook(v_bits) + flat_attn = mx.contiguous(attn.reshape(-1, S)) + flat_packed = mx.contiguous(v_packed.reshape(-1, pdim)) + flat_norms = mx.contiguous(v_norms.reshape(-1).astype(mx.float32)) + n_out = B * n_q_heads * L * dim + kernel = _av_weighted_sum_kernel(v_bits, dim) + out = kernel( + inputs=[flat_attn, flat_packed, flat_norms, rotation, centroids], + template=[ + ("T", dtype), + ("n_q_heads", n_q_heads), + ("n_kv_heads", n_kv_heads), + ("n_repeats", n_repeats), + ("query_len", L), + ("seq_len", S), + ], + grid=(n_out, 1, 1), + threadgroup=(1, 1, 1), + output_shapes=[(n_out,)], + output_dtypes=[dtype], + stream=mx.gpu, + )[0] + return out.reshape(B, n_q_heads, L, dim) + +def _metal_unpack_inline(bits: str, mask: str, packed: str, offset: str, idx: str) -> str: + return f""" + int word = ({idx} * {bits}) / BITS_PER_WORD; + int shift = ({idx} * {bits}) % BITS_PER_WORD; + uint w = {packed}[{offset} + word]; + uint val = (w >> shift) & {mask}; + if (shift + {bits} > BITS_PER_WORD) {{ + int spill = {bits} - (BITS_PER_WORD - shift); + uint w2 = {packed}[{offset} + word + 1]; + val = val | ((w2 & ((1u << spill) - 1u)) << (BITS_PER_WORD - shift)); + }} + """ + + +@lru_cache(maxsize=None) +def _tq_sdpa_vector_kernel(k_bits: int, v_bits: int, dim: int, scale: float): + """Fused vector SDPA for packed TurboQuant KV (L <= 8).""" + mse_bits = k_bits - 1 + k_pdim = packed_dim(dim, mse_bits) + v_pdim = packed_dim(dim, v_bits) + qjl_pdim = qjl_packed_dim(dim) + qjl_scale = _QJL_SCALE / dim + unpack_k_inline = _metal_unpack_inline("MSE_BITS", "MSE_MASK", "k_packed", "k_packed_off", "i") + unpack_v_inline = _metal_unpack_inline("V_BITS", "V_MASK", "v_packed", "v_packed_off", "i") + source = f""" + constexpr int DIM = {dim}; + constexpr int MSE_BITS = {mse_bits}; + constexpr int V_BITS = {v_bits}; + constexpr int K_PDIM = {k_pdim}; + constexpr int V_PDIM = {v_pdim}; + constexpr int QJL_PDIM = {qjl_pdim}; + constexpr uint MSE_MASK = (1u << MSE_BITS) - 1u; + constexpr uint V_MASK = (1u << V_BITS) - 1u; + constexpr int BITS_PER_WORD = 32; + constexpr int BD = 32; + constexpr int qk_per_thread = DIM / BD; + constexpr int v_per_thread = DIM / BD; + constexpr float QJL_SCALE = {qjl_scale}f; + constexpr float ATTN_SCALE = {scale}f; + + uint tg = threadgroup_position_in_grid.x; + uint q_batch_head_idx = tg / query_len; + uint q_seq_idx = tg % query_len; + uint simd_lid = thread_index_in_simdgroup; + + uint b = q_batch_head_idx / n_q_heads; + uint h_q = q_batch_head_idx % n_q_heads; + uint h_kv = h_q / n_repeats; + uint kv_base = (b * n_kv_heads + h_kv) * seq_len; + uint q_offset = q_batch_head_idx * query_len + q_seq_idx; + + threadgroup float q_shmem[DIM]; + threadgroup float q_rot_shmem[DIM]; + threadgroup float q_s_shmem[DIM]; + threadgroup float tg_max_score; + threadgroup float tg_sum_exp; + threadgroup float tg_factor; + threadgroup float tg_exp_score; + + thread float o[v_per_thread]; + for (int i = 0; i < v_per_thread; i++) {{ + o[i] = 0.0f; + }} + if (simd_lid == 0) {{ + tg_max_score = -1e30f; + tg_sum_exp = 0.0f; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint q_base = q_offset * DIM; + for (int t = 0; t < qk_per_thread; t++) {{ + int idx = simd_lid * qk_per_thread + t; + q_shmem[idx] = static_cast(queries[q_base + idx]) * ATTN_SCALE; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int idx = simd_lid; idx < DIM; idx += BD) {{ + float rot = 0.0f; + float ps = 0.0f; + for (int j = 0; j < DIM; j++) {{ + rot += q_shmem[j] * k_rotation[idx * DIM + j]; + ps += q_shmem[j] * k_qjl[idx * DIM + j]; + }} + q_rot_shmem[idx] = rot; + q_s_shmem[idx] = ps; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int ki = 0; ki < seq_len; ki++) {{ + if (do_causal && ki > int(seq_len - query_len + q_seq_idx)) {{ + continue; + }} + + uint kv_slot = kv_base + ki; + uint k_packed_off = kv_slot * K_PDIM; + uint k_signs_off = kv_slot * QJL_PDIM; + float kn = k_norms[kv_slot]; + float gamma = k_gamma[kv_slot]; + + float mse_part = 0.0f; + float qjl_part = 0.0f; + for (int t = 0; t < qk_per_thread; t++) {{ + int i = simd_lid * qk_per_thread + t; + {unpack_k_inline} + mse_part += centroids_k[val] * q_rot_shmem[i]; + int s_word = i / BITS_PER_WORD; + int s_shift = i % BITS_PER_WORD; + uint bit = (qjl_signs[k_signs_off + s_word] >> s_shift) & 1u; + float sign = bit ? 1.0f : -1.0f; + qjl_part += sign * q_s_shmem[i]; + }} + float mse_dot = simd_sum(mse_part); + float qjl_dot = simd_sum(qjl_part); + float score = kn * (mse_dot + QJL_SCALE * gamma * qjl_dot); + + if (simd_lid == 0) {{ + float new_max = max(tg_max_score, score); + tg_factor = metal::fast::exp(tg_max_score - new_max); + tg_exp_score = metal::fast::exp(score - new_max); + tg_max_score = new_max; + tg_sum_exp = tg_sum_exp * tg_factor + tg_exp_score; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint v_packed_off = kv_slot * V_PDIM; + float vn = v_norms[kv_slot]; + for (int t = 0; t < v_per_thread; t++) {{ + int d = simd_lid * v_per_thread + t; + float vd = 0.0f; + for (int i = 0; i < DIM; i++) {{ + {unpack_v_inline} + vd += centroids_v[val] * v_rotation[i * DIM + d]; + }} + o[t] = o[t] * tg_factor + tg_exp_score * vn * vd; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + }} + + float inv_sum = tg_sum_exp == 0.0f ? 1.0f : (1.0f / tg_sum_exp); + uint out_base = q_offset * DIM + simd_lid * v_per_thread; + for (int t = 0; t < v_per_thread; t++) {{ + out[out_base + t] = static_cast(o[t] * inv_sum); + }} + """ + return mx.fast.metal_kernel( + name=f"tq_sdpa_vector_v3_{k_bits}k_{v_bits}v_{dim}d", + input_names=[ + "queries", + "k_packed", + "k_norms", + "qjl_signs", + "k_gamma", + "v_packed", + "v_norms", + "k_rotation", + "v_rotation", + "centroids_k", + "centroids_v", + "k_qjl", + ], + output_names=["out"], + source=source, + ) + + +def tq_sdpa_vector_metal( + queries: mx.array, + k_packed: mx.array, + k_norms: mx.array, + qjl_signs: mx.array, + k_gamma: mx.array, + v_packed: mx.array, + v_norms: mx.array, + k_rotation: mx.array, + v_rotation: mx.array, + k_qjl: mx.array, + k_bits: int, + v_bits: int, + dim: int, + scale: float, + do_causal: bool, +) -> mx.array: + """Fused vector SDPA from packed TurboQuant KV (no dense K/V/scores).""" + if not metal_available(): + raise RuntimeError("Metal is not available") + B, n_q_heads, L, _ = queries.shape + S = k_packed.shape[2] + n_kv_heads = k_packed.shape[1] + n_repeats = n_q_heads // n_kv_heads + mse_bits = k_bits - 1 + k_pdim = packed_dim(dim, mse_bits) + v_pdim = packed_dim(dim, v_bits) + centroids_k, _ = get_codebook(mse_bits) + centroids_v, _ = get_codebook(v_bits) + flat_q = mx.contiguous(queries.reshape(-1, dim)) + flat_k = mx.contiguous(k_packed.reshape(-1, k_pdim)) + flat_signs = mx.contiguous(qjl_signs.reshape(-1, qjl_packed_dim(dim))) + flat_v = mx.contiguous(v_packed.reshape(-1, v_pdim)) + flat_kn = mx.contiguous(k_norms.reshape(-1).astype(mx.float32)) + flat_kg = mx.contiguous(k_gamma.reshape(-1).astype(mx.float32)) + flat_vn = mx.contiguous(v_norms.reshape(-1).astype(mx.float32)) + kernel = _tq_sdpa_vector_kernel(k_bits, v_bits, dim, scale) + out = kernel( + inputs=[ + flat_q, + flat_k, + flat_kn, + flat_signs, + flat_kg, + flat_v, + flat_vn, + k_rotation, + v_rotation, + centroids_k, + centroids_v, + k_qjl, + ], + template=[ + ("T", queries.dtype), + ("n_q_heads", n_q_heads), + ("n_kv_heads", n_kv_heads), + ("n_repeats", n_repeats), + ("query_len", L), + ("seq_len", S), + ("do_causal", do_causal), + ], + grid=(B * n_q_heads * L * 32, 1, 1), + threadgroup=(32, 1, 1), + output_shapes=[(B * n_q_heads * L, dim)], + output_dtypes=[queries.dtype], + stream=mx.gpu, + )[0] + return out.reshape(B, n_q_heads, L, dim) + + +@lru_cache(maxsize=None) +def _tq_sdpa_tiled_kernel(k_bits: int, v_bits: int, dim: int, scale: float): + """Fused tiled SDPA: 32 simdgroups stride over KV (L > 8 prefill).""" + mse_bits = k_bits - 1 + k_pdim = packed_dim(dim, mse_bits) + v_pdim = packed_dim(dim, v_bits) + qjl_pdim = qjl_packed_dim(dim) + qjl_scale = _QJL_SCALE / dim + unpack_k_inline = _metal_unpack_inline("MSE_BITS", "MSE_MASK", "k_packed", "k_packed_off", "i") + unpack_v_inline = _metal_unpack_inline("V_BITS", "V_MASK", "v_packed", "v_packed_off", "i") + source = f""" + constexpr int DIM = {dim}; + constexpr int MSE_BITS = {mse_bits}; + constexpr int V_BITS = {v_bits}; + constexpr int K_PDIM = {k_pdim}; + constexpr int V_PDIM = {v_pdim}; + constexpr int QJL_PDIM = {qjl_pdim}; + constexpr uint MSE_MASK = (1u << MSE_BITS) - 1u; + constexpr uint V_MASK = (1u << V_BITS) - 1u; + constexpr int BITS_PER_WORD = 32; + constexpr int BN = 32; + constexpr int BD = 32; + constexpr int qk_per_thread = DIM / BD; + constexpr int v_per_thread = DIM / BD; + constexpr float QJL_SCALE = {qjl_scale}f; + constexpr float ATTN_SCALE = {scale}f; + + uint tg = threadgroup_position_in_grid.x; + uint q_batch_head_idx = tg / query_len; + uint q_seq_idx = tg % query_len; + uint simd_gid = simdgroup_index_in_threadgroup; + uint simd_lid = thread_index_in_simdgroup; + + uint b = q_batch_head_idx / n_q_heads; + uint h_q = q_batch_head_idx % n_q_heads; + uint h_kv = h_q / n_repeats; + uint kv_base = (b * n_kv_heads + h_kv) * seq_len; + uint q_offset = q_batch_head_idx * query_len + q_seq_idx; + + threadgroup float q_shmem[DIM]; + threadgroup float q_rot_shmem[DIM]; + threadgroup float q_s_shmem[DIM]; + threadgroup float max_scores[BN]; + threadgroup float sum_exp_scores[BN]; + threadgroup float outputs[BN * DIM]; + + thread float o[v_per_thread]; + for (int i = 0; i < v_per_thread; i++) {{ + o[i] = 0.0f; + }} + + uint q_base = q_offset * DIM; + for (int t = 0; t < qk_per_thread; t++) {{ + int idx = simd_lid * qk_per_thread + t; + q_shmem[idx] = static_cast(queries[q_base + idx]) * ATTN_SCALE; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int idx = simd_lid; idx < DIM; idx += BD) {{ + float rot = 0.0f; + float ps = 0.0f; + for (int j = 0; j < DIM; j++) {{ + rot += q_shmem[j] * k_rotation[idx * DIM + j]; + ps += q_shmem[j] * k_qjl[idx * DIM + j]; + }} + q_rot_shmem[idx] = rot; + q_s_shmem[idx] = ps; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + float max_score = -1e30f; + float sum_exp_score = 0.0f; + + for (int ki = simd_gid; ki < seq_len; ki += BN) {{ + if (do_causal && ki > int(seq_len - query_len + q_seq_idx)) {{ + continue; + }} + + uint kv_slot = kv_base + ki; + uint k_packed_off = kv_slot * K_PDIM; + uint k_signs_off = kv_slot * QJL_PDIM; + float kn = k_norms[kv_slot]; + float gamma = k_gamma[kv_slot]; + + float mse_part = 0.0f; + float qjl_part = 0.0f; + for (int t = 0; t < qk_per_thread; t++) {{ + int i = simd_lid * qk_per_thread + t; + {unpack_k_inline} + mse_part += centroids_k[val] * q_rot_shmem[i]; + int s_word = i / BITS_PER_WORD; + int s_shift = i % BITS_PER_WORD; + uint bit = (qjl_signs[k_signs_off + s_word] >> s_shift) & 1u; + float sign = bit ? 1.0f : -1.0f; + qjl_part += sign * q_s_shmem[i]; + }} + float mse_dot = simd_sum(mse_part); + float qjl_dot = simd_sum(qjl_part); + float score = kn * (mse_dot + QJL_SCALE * gamma * qjl_dot); + + float new_max = max(max_score, score); + float factor = metal::fast::exp(max_score - new_max); + float exp_score = metal::fast::exp(score - new_max); + max_score = new_max; + sum_exp_score = sum_exp_score * factor + exp_score; + + uint v_packed_off = kv_slot * V_PDIM; + float vn = v_norms[kv_slot]; + for (int t = 0; t < v_per_thread; t++) {{ + int d = simd_lid * v_per_thread + t; + float vd = 0.0f; + for (int i = 0; i < DIM; i++) {{ + {unpack_v_inline} + vd += centroids_v[val] * v_rotation[i * DIM + d]; + }} + o[t] = o[t] * factor + exp_score * vn * vd; + }} + }} + + if (simd_lid == 0) {{ + max_scores[simd_gid] = max_score; + sum_exp_scores[simd_gid] = sum_exp_score; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + max_score = max_scores[simd_lid]; + float global_max = simd_max(max_score); + float global_factor = metal::fast::exp(max_score - global_max); + sum_exp_score = simd_sum(sum_exp_scores[simd_lid] * global_factor); + + for (int t = 0; t < v_per_thread; t++) {{ + outputs[simd_lid * DIM + simd_gid] = o[t]; + threadgroup_barrier(mem_flags::mem_threadgroup); + o[t] = simd_sum(outputs[simd_gid * DIM + simd_lid] * global_factor); + o[t] = sum_exp_score == 0.0f ? o[t] : (o[t] / sum_exp_score); + threadgroup_barrier(mem_flags::mem_threadgroup); + }} + + if (simd_lid == 0) {{ + uint out_base = q_offset * DIM + simd_gid * v_per_thread; + for (int t = 0; t < v_per_thread; t++) {{ + out[out_base + t] = static_cast(o[t]); + }} + }} + """ + return mx.fast.metal_kernel( + name=f"tq_sdpa_tiled_{k_bits}k_{v_bits}v_{dim}d", + input_names=[ + "queries", + "k_packed", + "k_norms", + "qjl_signs", + "k_gamma", + "v_packed", + "v_norms", + "k_rotation", + "v_rotation", + "centroids_k", + "centroids_v", + "k_qjl", + ], + output_names=["out"], + source=source, + ) + + +_TQ_TILE_THREADS = 1024 + + +def tq_sdpa_tiled_metal( + queries: mx.array, + k_packed: mx.array, + k_norms: mx.array, + qjl_signs: mx.array, + k_gamma: mx.array, + v_packed: mx.array, + v_norms: mx.array, + k_rotation: mx.array, + v_rotation: mx.array, + k_qjl: mx.array, + k_bits: int, + v_bits: int, + dim: int, + scale: float, + do_causal: bool, +) -> mx.array: + """Fused tiled SDPA from packed TurboQuant KV (simdgroup-parallel over S).""" + if not metal_available(): + raise RuntimeError("Metal is not available") + B, n_q_heads, L, _ = queries.shape + S = k_packed.shape[2] + n_kv_heads = k_packed.shape[1] + n_repeats = n_q_heads // n_kv_heads + mse_bits = k_bits - 1 + k_pdim = packed_dim(dim, mse_bits) + v_pdim = packed_dim(dim, v_bits) + centroids_k, _ = get_codebook(mse_bits) + centroids_v, _ = get_codebook(v_bits) + flat_q = mx.contiguous(queries.reshape(-1, dim)) + flat_k = mx.contiguous(k_packed.reshape(-1, k_pdim)) + flat_signs = mx.contiguous(qjl_signs.reshape(-1, qjl_packed_dim(dim))) + flat_v = mx.contiguous(v_packed.reshape(-1, v_pdim)) + flat_kn = mx.contiguous(k_norms.reshape(-1).astype(mx.float32)) + flat_kg = mx.contiguous(k_gamma.reshape(-1).astype(mx.float32)) + flat_vn = mx.contiguous(v_norms.reshape(-1).astype(mx.float32)) + n_tg = B * n_q_heads * L + kernel = _tq_sdpa_tiled_kernel(k_bits, v_bits, dim, scale) + out = kernel( + inputs=[ + flat_q, + flat_k, + flat_kn, + flat_signs, + flat_kg, + flat_v, + flat_vn, + k_rotation, + v_rotation, + centroids_k, + centroids_v, + k_qjl, + ], + template=[ + ("T", queries.dtype), + ("n_q_heads", n_q_heads), + ("n_kv_heads", n_kv_heads), + ("n_repeats", n_repeats), + ("query_len", L), + ("seq_len", S), + ("do_causal", do_causal), + ], + grid=(n_tg * _TQ_TILE_THREADS, 1, 1), + threadgroup=(_TQ_TILE_THREADS, 1, 1), + output_shapes=[(B * n_q_heads * L, dim)], + output_dtypes=[queries.dtype], + stream=mx.gpu, + )[0] + return out.reshape(B, n_q_heads, L, dim) + +TQ_2PASS_THRESHOLD = 1024 +TQ_2PASS_BLOCKS = 64 + + +@lru_cache(maxsize=None) +def _tq_sdpa_2pass1_kernel(k_bits: int, v_bits: int, dim: int, scale: float): + """Pass 1: per-block partial softmax stats and unnormalized output.""" + mse_bits = k_bits - 1 + k_pdim = packed_dim(dim, mse_bits) + v_pdim = packed_dim(dim, v_bits) + qjl_pdim = qjl_packed_dim(dim) + qjl_scale = _QJL_SCALE / dim + unpack_k_inline = _metal_unpack_inline("MSE_BITS", "MSE_MASK", "k_packed", "k_packed_off", "i") + unpack_v_inline = _metal_unpack_inline("V_BITS", "V_MASK", "v_packed", "v_packed_off", "i") + source = f""" + constexpr int DIM = {dim}; + constexpr int MSE_BITS = {mse_bits}; + constexpr int V_BITS = {v_bits}; + constexpr int K_PDIM = {k_pdim}; + constexpr int V_PDIM = {v_pdim}; + constexpr int QJL_PDIM = {qjl_pdim}; + constexpr uint MSE_MASK = (1u << MSE_BITS) - 1u; + constexpr uint V_MASK = (1u << V_BITS) - 1u; + constexpr int BITS_PER_WORD = 32; + constexpr int BD = 32; + constexpr int qk_per_thread = DIM / BD; + constexpr int v_per_thread = DIM / BD; + constexpr float QJL_SCALE = {qjl_scale}f; + constexpr float ATTN_SCALE = {scale}f; + + uint tg = threadgroup_position_in_grid.x; + uint block_idx = tg % n_blocks; + uint q_offset = tg / n_blocks; + uint q_batch_head_idx = q_offset / query_len; + uint q_seq_idx = q_offset % query_len; + uint simd_lid = thread_index_in_simdgroup; + + uint b = q_batch_head_idx / n_q_heads; + uint h_q = q_batch_head_idx % n_q_heads; + uint h_kv = h_q / n_repeats; + uint kv_base = (b * n_kv_heads + h_kv) * seq_len; + + threadgroup float q_shmem[DIM]; + threadgroup float q_rot_shmem[DIM]; + threadgroup float q_s_shmem[DIM]; + threadgroup float tg_max_score; + threadgroup float tg_sum_exp; + threadgroup float tg_factor; + threadgroup float tg_exp_score; + + thread float o[v_per_thread]; + for (int i = 0; i < v_per_thread; i++) {{ + o[i] = 0.0f; + }} + if (simd_lid == 0) {{ + tg_max_score = -1e30f; + tg_sum_exp = 0.0f; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint q_base = q_offset * DIM; + for (int t = 0; t < qk_per_thread; t++) {{ + int idx = simd_lid * qk_per_thread + t; + q_shmem[idx] = static_cast(queries[q_base + idx]) * ATTN_SCALE; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int idx = simd_lid; idx < DIM; idx += BD) {{ + float rot = 0.0f; + float ps = 0.0f; + for (int j = 0; j < DIM; j++) {{ + rot += q_shmem[j] * k_rotation[idx * DIM + j]; + ps += q_shmem[j] * k_qjl[idx * DIM + j]; + }} + q_rot_shmem[idx] = rot; + q_s_shmem[idx] = ps; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int ki = block_idx; ki < seq_len; ki += n_blocks) {{ + if (do_causal && ki > int(seq_len - query_len + q_seq_idx)) {{ + continue; + }} + + uint kv_slot = kv_base + ki; + uint k_packed_off = kv_slot * K_PDIM; + uint k_signs_off = kv_slot * QJL_PDIM; + float kn = k_norms[kv_slot]; + float gamma = k_gamma[kv_slot]; + + float mse_part = 0.0f; + float qjl_part = 0.0f; + for (int t = 0; t < qk_per_thread; t++) {{ + int i = simd_lid * qk_per_thread + t; + {unpack_k_inline} + mse_part += centroids_k[val] * q_rot_shmem[i]; + int s_word = i / BITS_PER_WORD; + int s_shift = i % BITS_PER_WORD; + uint bit = (qjl_signs[k_signs_off + s_word] >> s_shift) & 1u; + float sign = bit ? 1.0f : -1.0f; + qjl_part += sign * q_s_shmem[i]; + }} + float mse_dot = simd_sum(mse_part); + float qjl_dot = simd_sum(qjl_part); + float score = kn * (mse_dot + QJL_SCALE * gamma * qjl_dot); + + if (simd_lid == 0) {{ + float new_max = max(tg_max_score, score); + tg_factor = metal::fast::exp(tg_max_score - new_max); + tg_exp_score = metal::fast::exp(score - new_max); + tg_max_score = new_max; + tg_sum_exp = tg_sum_exp * tg_factor + tg_exp_score; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint v_packed_off = kv_slot * V_PDIM; + float vn = v_norms[kv_slot]; + for (int t = 0; t < v_per_thread; t++) {{ + int d = simd_lid * v_per_thread + t; + float vd = 0.0f; + for (int i = 0; i < DIM; i++) {{ + {unpack_v_inline} + vd += centroids_v[val] * v_rotation[i * DIM + d]; + }} + o[t] = o[t] * tg_factor + tg_exp_score * vn * vd; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + }} + + uint out_slot = q_offset * n_blocks + block_idx; + if (simd_lid == 0) {{ + sums[out_slot] = tg_sum_exp; + maxs[out_slot] = tg_max_score; + }} + uint partial_base = out_slot * DIM + simd_lid * v_per_thread; + for (int t = 0; t < v_per_thread; t++) {{ + partials[partial_base + t] = static_cast(o[t]); + }} + """ + return mx.fast.metal_kernel( + name=f"tq_sdpa_2pass1_{k_bits}k_{v_bits}v_{dim}d", + input_names=[ + "queries", + "k_packed", + "k_norms", + "qjl_signs", + "k_gamma", + "v_packed", + "v_norms", + "k_rotation", + "v_rotation", + "centroids_k", + "centroids_v", + "k_qjl", + ], + output_names=["partials", "sums", "maxs"], + source=source, + ) + + +@lru_cache(maxsize=None) +def _tq_sdpa_2pass2_kernel(dim: int): + """Pass 2: merge per-block partials with global online softmax.""" + source = f""" + constexpr int DIM = {dim}; + constexpr int BN = 32; + constexpr int BD = 32; + constexpr int v_per_thread = DIM / BD; + + uint tg = threadgroup_position_in_grid.x; + uint q_offset = tg; + uint simd_gid = simdgroup_index_in_threadgroup; + uint simd_lid = thread_index_in_simdgroup; + + uint block_base = q_offset * n_blocks; + partials += block_base * DIM + simd_gid * DIM + simd_lid * v_per_thread; + sums += block_base; + maxs += block_base; + + thread float o[v_per_thread]; + for (int i = 0; i < v_per_thread; i++) {{ + o[i] = 0.0f; + }} + + float max_score = -1e30f; + for (int b = 0; b < n_blocks / BN; ++b) {{ + max_score = max(max_score, maxs[simd_lid + BN * b]); + }} + float global_max = simd_max(max_score); + + float sum_exp_score = 0.0f; + for (int b = 0; b < n_blocks / BN; ++b) {{ + float factor = metal::fast::exp(maxs[simd_lid + BN * b] - global_max); + sum_exp_score += factor * sums[simd_lid + BN * b]; + }} + sum_exp_score = simd_sum(sum_exp_score); + + for (int b = 0; b < n_blocks / BN; ++b) {{ + float factor = metal::fast::exp(maxs[simd_gid] - global_max); + for (int t = 0; t < v_per_thread; t++) {{ + o[t] += factor * static_cast(partials[t]); + }} + maxs += BN; + sums += BN; + partials += BN * DIM; + }} + + threadgroup float outputs[BN * DIM]; + for (int t = 0; t < v_per_thread; t++) {{ + outputs[simd_lid * DIM + simd_gid] = o[t]; + threadgroup_barrier(mem_flags::mem_threadgroup); + o[t] = simd_sum(outputs[simd_gid * DIM + simd_lid]); + o[t] = sum_exp_score == 0.0f ? o[t] : (o[t] / sum_exp_score); + threadgroup_barrier(mem_flags::mem_threadgroup); + }} + + if (simd_lid == 0) {{ + uint out_base = q_offset * DIM + simd_gid * v_per_thread; + for (int t = 0; t < v_per_thread; t++) {{ + out[out_base + t] = static_cast(o[t]); + }} + }} + """ + return mx.fast.metal_kernel( + name=f"tq_sdpa_2pass2_{dim}d", + input_names=["partials", "sums", "maxs"], + output_names=["out"], + source=source, + ) + + +def tq_sdpa_2pass_metal( + queries: mx.array, + k_packed: mx.array, + k_norms: mx.array, + qjl_signs: mx.array, + k_gamma: mx.array, + v_packed: mx.array, + v_norms: mx.array, + k_rotation: mx.array, + v_rotation: mx.array, + k_qjl: mx.array, + k_bits: int, + v_bits: int, + dim: int, + scale: float, + do_causal: bool, + n_blocks: int = TQ_2PASS_BLOCKS, +) -> mx.array: + """Two-pass fused SDPA for long sequences (S >= 1024).""" + if not metal_available(): + raise RuntimeError("Metal is not available") + B, n_q_heads, L, _ = queries.shape + S = k_packed.shape[2] + n_kv_heads = k_packed.shape[1] + n_repeats = n_q_heads // n_kv_heads + mse_bits = k_bits - 1 + k_pdim = packed_dim(dim, mse_bits) + v_pdim = packed_dim(dim, v_bits) + centroids_k, _ = get_codebook(mse_bits) + centroids_v, _ = get_codebook(v_bits) + flat_q = mx.contiguous(queries.reshape(-1, dim)) + flat_k = mx.contiguous(k_packed.reshape(-1, k_pdim)) + flat_signs = mx.contiguous(qjl_signs.reshape(-1, qjl_packed_dim(dim))) + flat_v = mx.contiguous(v_packed.reshape(-1, v_pdim)) + flat_kn = mx.contiguous(k_norms.reshape(-1).astype(mx.float32)) + flat_kg = mx.contiguous(k_gamma.reshape(-1).astype(mx.float32)) + flat_vn = mx.contiguous(v_norms.reshape(-1).astype(mx.float32)) + n_q = B * n_q_heads * L + kernel1 = _tq_sdpa_2pass1_kernel(k_bits, v_bits, dim, scale) + partials, sums, maxs = kernel1( + inputs=[ + flat_q, + flat_k, + flat_kn, + flat_signs, + flat_kg, + flat_v, + flat_vn, + k_rotation, + v_rotation, + centroids_k, + centroids_v, + k_qjl, + ], + template=[ + ("T", queries.dtype), + ("n_q_heads", n_q_heads), + ("n_kv_heads", n_kv_heads), + ("n_repeats", n_repeats), + ("query_len", L), + ("seq_len", S), + ("do_causal", do_causal), + ("n_blocks", n_blocks), + ], + grid=(n_q * n_blocks * 32, 1, 1), + threadgroup=(32, 1, 1), + output_shapes=[(n_q * n_blocks, dim), (n_q * n_blocks,), (n_q * n_blocks,)], + output_dtypes=[queries.dtype, mx.float32, mx.float32], + stream=mx.gpu, + ) + kernel2 = _tq_sdpa_2pass2_kernel(dim) + out = kernel2( + inputs=[partials, sums, maxs], + template=[("T", queries.dtype), ("n_blocks", n_blocks)], + grid=(n_q * _TQ_TILE_THREADS, 1, 1), + threadgroup=(_TQ_TILE_THREADS, 1, 1), + output_shapes=[(n_q, dim)], + output_dtypes=[queries.dtype], + stream=mx.gpu, + )[0] + return out.reshape(B, n_q_heads, L, dim) + + +def tq_sdpa_metal( + queries: mx.array, + k_packed: mx.array, + k_norms: mx.array, + qjl_signs: mx.array, + k_gamma: mx.array, + v_packed: mx.array, + v_norms: mx.array, + k_rotation: mx.array, + v_rotation: mx.array, + k_qjl: mx.array, + k_bits: int, + v_bits: int, + dim: int, + scale: float, + do_causal: bool, +) -> mx.array: + """Dispatch fused SDPA: 2-pass (S>=1024), vector (L<=8), else tiled.""" + S = k_packed.shape[2] + L = queries.shape[2] + if S >= TQ_2PASS_THRESHOLD: + return tq_sdpa_2pass_metal( + queries, + k_packed, + k_norms, + qjl_signs, + k_gamma, + v_packed, + v_norms, + k_rotation, + v_rotation, + k_qjl, + k_bits, + v_bits, + dim, + scale, + do_causal, + ) + if L <= 8: + return tq_sdpa_vector_metal( + queries, + k_packed, + k_norms, + qjl_signs, + k_gamma, + v_packed, + v_norms, + k_rotation, + v_rotation, + k_qjl, + k_bits, + v_bits, + dim, + scale, + do_causal, + ) + return tq_sdpa_tiled_metal( + queries, + k_packed, + k_norms, + qjl_signs, + k_gamma, + v_packed, + v_norms, + k_rotation, + v_rotation, + k_qjl, + k_bits, + v_bits, + dim, + scale, + do_causal, + ) \ No newline at end of file diff --git a/mlx_lm/turboquant/packing.py b/mlx_lm/turboquant/packing.py new file mode 100644 index 000000000..e52767786 --- /dev/null +++ b/mlx_lm/turboquant/packing.py @@ -0,0 +1,55 @@ +# Copyright © 2025 Bonsai Demo contributors. + +"""Bit packing for low-bit TurboQuant index streams.""" + +from __future__ import annotations + +import math + +import mlx.core as mx + +_BITS_PER_WORD = 32 + + +def packed_dim(dim: int, bits: int) -> int: + return math.ceil(dim * bits / _BITS_PER_WORD) + + +def pack_indices(indices: mx.array, bits: int) -> mx.array: + """Pack unsigned indices along the last axis into uint32 words.""" + if bits not in (1, 2, 3, 4): + raise ValueError(f"pack_indices supports 1-4 bits, got {bits}") + *batch, dim = indices.shape + pdim = packed_dim(dim, bits) + flat = indices.reshape(-1, dim).astype(mx.uint32) + n = flat.shape[0] + words = mx.zeros((n, pdim), dtype=mx.uint32) + mask = (1 << bits) - 1 + for i in range(dim): + word = i * bits // _BITS_PER_WORD + shift = (i * bits) % _BITS_PER_WORD + chunk = (flat[:, i] & mask) << shift + words[:, word] = words[:, word] | chunk + if shift + bits > _BITS_PER_WORD: + spill = bits - (_BITS_PER_WORD - shift) + words[:, word + 1] = words[:, word + 1] | ((flat[:, i] & mask) >> (_BITS_PER_WORD - shift)) + return words.reshape(*batch, pdim) + + +def unpack_indices(packed: mx.array, bits: int, dim: int) -> mx.array: + if bits not in (1, 2, 3, 4): + raise ValueError(f"unpack_indices supports 1-4 bits, got {bits}") + *batch, pdim = packed.shape + flat = packed.reshape(-1, pdim).astype(mx.uint32) + n = flat.shape[0] + out = mx.zeros((n, dim), dtype=mx.uint32) + mask = (1 << bits) - 1 + for i in range(dim): + word = i * bits // _BITS_PER_WORD + shift = (i * bits) % _BITS_PER_WORD + val = (flat[:, word] >> shift) & mask + if shift + bits > _BITS_PER_WORD: + spill = bits - (_BITS_PER_WORD - shift) + val = val | ((flat[:, word + 1] & ((1 << spill) - 1)) << (_BITS_PER_WORD - shift)) + out[:, i] = val + return out.reshape(*batch, dim).astype(mx.uint8) \ No newline at end of file diff --git a/mlx_lm/turboquant/qjl.py b/mlx_lm/turboquant/qjl.py new file mode 100644 index 000000000..454e4c044 --- /dev/null +++ b/mlx_lm/turboquant/qjl.py @@ -0,0 +1,45 @@ +# Copyright © 2025 Bonsai Demo contributors. + +"""Quantized Johnson-Lindenstrauss (QJL) 1-bit residual quantizer.""" + +from __future__ import annotations + +import math + +import mlx.core as mx + +from mlx_lm.turboquant.packing import pack_indices, packed_dim, unpack_indices + + +def make_qjl_matrix(dim: int, seed: int) -> mx.array: + mx.random.seed(seed + 1_000_003) + with mx.stream(mx.cpu): + s = mx.random.normal(shape=(dim, dim)).astype(mx.float32) + mx.eval(s) + return s + + +def qjl_packed_dim(dim: int) -> int: + """Packed uint32 words for 1-bit QJL signs along ``dim``.""" + return packed_dim(dim, 1) + + +def quantize_qjl(residual: mx.array, s_matrix: mx.array) -> mx.array: + """Pack 1-bit QJL signs; last axis is ``qjl_packed_dim(dim)`` uint32 words.""" + dim = residual.shape[-1] + projected = residual @ s_matrix.T + bits = (projected >= 0).astype(mx.uint8) + return pack_indices(bits, bits=1) + + +def dequantize_qjl( + packed_signs: mx.array, gamma: mx.array, s_matrix: mx.array, dim: int +) -> mx.array: + """Unbiased QJL reconstruction scaled by residual norm gamma.""" + scale = math.sqrt(math.pi / 2.0) / dim + out_shape = packed_signs.shape[:-1] + (dim,) + bits = unpack_indices(packed_signs, bits=1, dim=dim) + flat_signs = bits.reshape(-1, dim).astype(mx.float32) * 2.0 - 1.0 + flat_gamma = gamma.reshape(-1, 1) + out = scale * flat_gamma * (flat_signs @ s_matrix) + return out.reshape(out_shape) \ No newline at end of file diff --git a/mlx_lm/turboquant/quantize.py b/mlx_lm/turboquant/quantize.py new file mode 100644 index 000000000..23a494c8a --- /dev/null +++ b/mlx_lm/turboquant/quantize.py @@ -0,0 +1,180 @@ +# Copyright © 2025 Bonsai Demo contributors. + +"""TurboQuant encode/decode primitives.""" + +from __future__ import annotations + +import mlx.core as mx + +from mlx_lm.turboquant.codebooks import dequantize_coords, quantize_coords +from mlx_lm.turboquant.kernels import ( + decode_mse_metal, + decode_prod_metal, + encode_kv_metal, + encode_mse_metal, + encode_prod_metal, + metal_available, +) +from mlx_lm.turboquant.packing import pack_indices, unpack_indices +from mlx_lm.turboquant.qjl import dequantize_qjl, qjl_packed_dim, quantize_qjl + +_USE_METAL = metal_available() + + +def _decode_mse_impl( + packed: mx.array, + norms: mx.array, + rotation: mx.array, + bits: int, + dim: int, + dtype: mx.dtype, +) -> mx.array: + if _USE_METAL: + return decode_mse_metal(packed, norms, rotation, bits, dim, dtype) + out_shape = packed.shape[:-1] + (dim,) + indices = unpack_indices(packed, bits, dim) + flat = indices.reshape(-1, dim) + rotated = dequantize_coords(flat, bits) + unit = _inverse_rotate(rotated, rotation) + scaled = unit * norms.reshape(-1, 1) + return scaled.reshape(out_shape) + + +def _decode_prod_impl( + mse_packed: mx.array, + norms: mx.array, + qjl_signs: mx.array, + gamma: mx.array, + rotation: mx.array, + s_matrix: mx.array, + bits: int, + dim: int, + dtype: mx.dtype, +) -> mx.array: + if _USE_METAL: + return decode_prod_metal( + mse_packed, norms, qjl_signs, gamma, rotation, s_matrix, bits, dim, dtype + ) + mse_bits = bits - 1 + ones = mx.ones(norms.shape, dtype=norms.dtype) + mse_unit = _decode_mse_impl(mse_packed, ones, rotation, mse_bits, dim, mx.float32) + qjl_part = dequantize_qjl(qjl_signs, gamma, s_matrix, dim) + return (norms * (mse_unit + qjl_part)).astype(dtype) + + +def _rotate(vectors: mx.array, rotation: mx.array) -> mx.array: + # vectors: (N, d), rotation: (d, d) — row-vector convention: y = x @ R.T + return vectors @ rotation.T + + +def _inverse_rotate(vectors: mx.array, rotation: mx.array) -> mx.array: + return vectors @ rotation + + +def _encode_mse_ref(vectors: mx.array, rotation: mx.array, bits: int): + shape = vectors.shape + dim = shape[-1] + flat = vectors.astype(mx.float32).reshape(-1, dim) + norms = mx.linalg.norm(flat, axis=-1, keepdims=True) + unit = flat / mx.maximum(norms, 1e-8) + rotated = _rotate(unit, rotation) + indices = quantize_coords(rotated, bits).reshape(shape[:-1] + (dim,)) + packed = pack_indices(indices, bits) + return packed, norms.reshape(shape[:-1] + (1,)) + + +def encode_mse(vectors: mx.array, rotation: mx.array, bits: int): + """Quantize vectors with TurboQuant_mse. Returns packed indices and norms.""" + shape = vectors.shape + dim = shape[-1] + if _USE_METAL: + return encode_mse_metal(vectors, rotation, bits, dim) + return _encode_mse_ref(vectors, rotation, bits) + + +def decode_mse( + packed: mx.array, norms: mx.array, rotation: mx.array, bits: int, dim: int +) -> mx.array: + return _decode_mse_impl(packed, norms, rotation, bits, dim, mx.float32) + + +def _encode_prod_ref(vectors: mx.array, rotation: mx.array, s_matrix: mx.array, bits: int): + if bits < 2: + raise ValueError("TurboQuant_prod requires bits >= 2") + mse_bits = bits - 1 + shape = vectors.shape + dim = shape[-1] + flat = vectors.astype(mx.float32).reshape(-1, dim) + norms = mx.linalg.norm(flat, axis=-1, keepdims=True) + unit = flat / mx.maximum(norms, 1e-8) + unit_batched = unit.reshape(shape) + + mse_packed, _ = _encode_mse_ref(unit_batched, rotation, mse_bits) + mse_unit = _decode_mse_impl( + mse_packed, + mx.ones_like(norms.reshape(shape[:-1] + (1,))), + rotation, + mse_bits, + dim, + mx.float32, + ).reshape(-1, dim) + + residual = unit - mse_unit + gamma = mx.linalg.norm(residual, axis=-1, keepdims=True) + qjl_signs = quantize_qjl( + residual.reshape(shape), s_matrix + ).reshape(shape[:-1] + (qjl_packed_dim(dim),)) + + return ( + mse_packed, + norms.reshape(shape[:-1] + (1,)), + qjl_signs, + gamma.reshape(shape[:-1] + (1,)), + ) + + +def encode_prod(vectors: mx.array, rotation: mx.array, s_matrix: mx.array, bits: int): + """TurboQuant_prod: (bits-1) Lloyd-Max on the unit vector + 1-bit QJL residual.""" + if bits < 2: + raise ValueError("TurboQuant_prod requires bits >= 2") + dim = vectors.shape[-1] + if _USE_METAL: + return encode_prod_metal(vectors, rotation, s_matrix, bits, dim) + return _encode_prod_ref(vectors, rotation, s_matrix, bits) + + +def encode_kv( + keys: mx.array, + values: mx.array, + k_rotation: mx.array, + v_rotation: mx.array, + k_qjl: mx.array, + k_bits: int, + v_bits: int, +): + """Fused TurboQuant_prod(K) + TurboQuant_mse(V) in one dispatch.""" + if k_bits < 2: + raise ValueError("TurboQuant_prod requires k_bits >= 2") + dim = keys.shape[-1] + if _USE_METAL: + return encode_kv_metal( + keys, values, k_rotation, v_rotation, k_qjl, k_bits, v_bits, dim + ) + k_out = encode_prod(keys, k_rotation, k_qjl, k_bits) + v_out = encode_mse(values, v_rotation, v_bits) + return (*k_out, *v_out) + + +def decode_prod( + mse_packed: mx.array, + norms: mx.array, + qjl_signs: mx.array, + gamma: mx.array, + rotation: mx.array, + s_matrix: mx.array, + bits: int, + dim: int, +) -> mx.array: + return _decode_prod_impl( + mse_packed, norms, qjl_signs, gamma, rotation, s_matrix, bits, dim, mx.float32 + ) \ No newline at end of file diff --git a/mlx_lm/turboquant/rotation.py b/mlx_lm/turboquant/rotation.py new file mode 100644 index 000000000..2c6d47d1e --- /dev/null +++ b/mlx_lm/turboquant/rotation.py @@ -0,0 +1,20 @@ +# Copyright © 2025 Bonsai Demo contributors. + +"""Random orthogonal rotations for TurboQuant (QR on a Gaussian draw).""" + +from __future__ import annotations + +import mlx.core as mx + + +def make_rotation_matrix(dim: int, seed: int) -> mx.array: + """Return a row-orthonormal rotation matrix of shape (dim, dim).""" + if dim <= 0: + raise ValueError("dim must be positive") + cpu = mx.cpu + mx.random.seed(seed) + with mx.stream(cpu): + gaussian = mx.random.normal(shape=(dim, dim)) + q, _ = mx.linalg.qr(gaussian) + mx.eval(q) + return q.astype(mx.float32) \ No newline at end of file diff --git a/scripts/verify_streaming_consumer.py b/scripts/verify_streaming_consumer.py new file mode 100644 index 000000000..919396d43 --- /dev/null +++ b/scripts/verify_streaming_consumer.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python3 +# Copyright © 2024 Apple Inc. + +"""Consumer verification for mlx-lm layer-streaming (MLX-first tensor paths).""" + +import json +import sys +import tempfile +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +import mlx.core as mx + +from mlx_lm.generate import generate_step +from mlx_lm.models.llama import Model, ModelArgs +from mlx_lm.streaming import StreamingConfig, load_streaming +from mlx_lm.streaming.split_model import ensure_streaming_layout, split_model_by_layers +from mlx_lm.utils import save_model + + +def _tiny_config(): + return { + "model_type": "llama", + "vocab_size": 128, + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "rms_norm_eps": 1e-5, + "tie_word_embeddings": True, + } + + +def _build_model_dir(tmp_path: Path) -> Path: + config = _tiny_config() + with open(tmp_path / "config.json", "w") as f: + json.dump(config, f) + args = ModelArgs.from_dict(config) + model = Model(args) + mx.eval(model.parameters()) + save_model(tmp_path, model) + ensure_streaming_layout(tmp_path, verbose=False) + return tmp_path + + +def main() -> int: + with tempfile.TemporaryDirectory() as tmp: + model_dir = _build_model_dir(Path(tmp)) + + # Split I/O uses mx.load / mx.save_safetensors + weights = mx.load(str(model_dir / "layer_0.safetensors")) + assert all(isinstance(v, mx.array) for v in weights.values()) + print("SPLIT OK", isinstance(weights["self_attn.q_proj.weight"], mx.array)) + + model, _, _ = load_streaming( + str(model_dir), + StreamingConfig(window_size=1, verbose=False), + load_tokenizer=False, + ) + print("LOAD OK", type(model).__name__) + + cache = model.make_cache() + logits = model(mx.array([[1, 2, 3]]), cache=cache) + mx.eval(logits) + assert isinstance(logits, mx.array) + print("FORWARD OK", isinstance(logits, mx.array), logits.shape) + + prompt = mx.array([1, 2, 3]) + for token, logprobs in generate_step(prompt, model, max_tokens=2): + assert isinstance(logprobs, mx.array) + assert isinstance(token, (int, mx.integer_types if hasattr(mx, "integer_types") else int)) + print("SAMPLE OK", isinstance(logprobs, mx.array)) + + # Standalone split entry path + mono = tmp + "/mono.safetensors" + mx.save_safetensors( + mono, + { + "model.layers.0.self_attn.q_proj.weight": mx.random.normal((8, 8)), + "model.embed_tokens.weight": mx.random.normal((100, 8)), + }, + ) + out = Path(tmp) / "split2" + split_model_by_layers(Path(mono), out) + assert (out / "layer_0.safetensors").exists() + print("SPLIT_CLI OK", True) + + print("ALL CONSUMER TESTS PASSED") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) \ No newline at end of file diff --git a/setup.py b/setup.py index f8e547b95..87a3e74ab 100644 --- a/setup.py +++ b/setup.py @@ -39,6 +39,8 @@ "mlx_lm.tuner", "mlx_lm.tool_parsers", "mlx_lm.chat_templates", + "mlx_lm.turboquant", + "mlx_lm.streaming", ], python_requires=">=3.8", extras_require={ @@ -69,6 +71,8 @@ "mlx_lm.share = mlx_lm.share:main", "mlx_lm.manage = mlx_lm.manage:main", "mlx_lm.upload = mlx_lm.upload:main", + "mlx_lm.split_model = mlx_lm.streaming.split_model:main", + "mlx_lm.streaming_benchmark = mlx_lm.streaming.benchmark:main", ] }, ) diff --git a/tests/test_streaming.py b/tests/test_streaming.py new file mode 100644 index 000000000..355dbb7ef --- /dev/null +++ b/tests/test_streaming.py @@ -0,0 +1,50 @@ +# Copyright © 2024 Apple Inc. + +import tempfile +from pathlib import Path + +import mlx.core as mx +import pytest + +from mlx_lm.streaming.config import StreamingConfig +from mlx_lm.streaming.layer_loader import RollingWindowLoader +from mlx_lm.streaming.split_model import split_model_by_layers + + +def test_streaming_config_window_size(): + cfg = StreamingConfig(max_memory_gb=16.0, safety_margin=0.8) + layer_size = 200_000_000 + window = cfg.compute_optimal_window(layer_size, kv_cache_bytes=1_000_000_000) + assert window >= 1 + + +def test_split_and_load_layers(): + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + weights = { + "model.layers.0.self_attn.q_proj.weight": mx.random.normal((8, 8)), + "model.layers.1.self_attn.q_proj.weight": mx.random.normal((8, 8)), + "model.embed_tokens.weight": mx.random.normal((100, 8)), + "model.norm.weight": mx.ones((8,)), + } + model_path = tmp_path / "model.safetensors" + mx.save_safetensors(str(model_path), weights) + + out_dir = tmp_path / "split" + split_model_by_layers(model_path, out_dir) + + assert (out_dir / "layer_0.safetensors").exists() + assert (out_dir / "fixed_weights.safetensors").exists() + + cfg = StreamingConfig(window_size=2, verbose=False) + loader = RollingWindowLoader( + model_path=out_dir, + num_layers=2, + layer_size_bytes=1024, + streaming_config=cfg, + ) + lw = loader.get_layer(0) + assert isinstance(lw.weights, dict) + assert all(isinstance(v, mx.array) for v in lw.weights.values()) + assert "self_attn.q_proj.weight" in lw.weights + assert isinstance(lw.weights["self_attn.q_proj.weight"], mx.array) \ No newline at end of file diff --git a/tests/test_streaming_generate.py b/tests/test_streaming_generate.py new file mode 100644 index 000000000..2314585cc --- /dev/null +++ b/tests/test_streaming_generate.py @@ -0,0 +1,117 @@ +# Copyright © 2024 Apple Inc. + +"""End-to-end tests for layer-streaming through mlx_lm.generate.""" + +import json +import tempfile +from pathlib import Path + +import mlx.core as mx +import pytest + +from mlx_lm.generate import generate_step +from mlx_lm.models.llama import Model, ModelArgs +from mlx_lm.streaming import StreamingConfig, load_streaming +from mlx_lm.streaming.split_model import ensure_streaming_layout +from mlx_lm.utils import save_model + + +def _tiny_llama_config(): + return { + "model_type": "llama", + "vocab_size": 128, + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "rms_norm_eps": 1e-5, + "tie_word_embeddings": True, + } + + +def _build_tiny_model_dir(tmp_path: Path) -> Path: + config = _tiny_llama_config() + with open(tmp_path / "config.json", "w") as f: + json.dump(config, f) + + args = ModelArgs.from_dict(config) + model = Model(args) + mx.eval(model.parameters()) + save_model(tmp_path, model) + ensure_streaming_layout(tmp_path, verbose=False) + return tmp_path + + +def test_load_streaming_auto_split(): + with tempfile.TemporaryDirectory() as tmp: + model_dir = _build_tiny_model_dir(Path(tmp)) + assert (model_dir / "fixed_weights.safetensors").exists() + assert (model_dir / "layer_0.safetensors").exists() + + model, tokenizer, config = load_streaming( + str(model_dir), + StreamingConfig(window_size=1, verbose=False), + load_tokenizer=False, + ) + assert tokenizer is None + stats = model.get_stats() + assert stats["streaming"]["window_size"] >= 1 + assert config["num_hidden_layers"] == 2 + + +def test_streaming_forward_returns_mx_array(): + """StreamingModelWrapper forward must return mx.array logits.""" + with tempfile.TemporaryDirectory() as tmp: + model_dir = _build_tiny_model_dir(Path(tmp)) + model, _, _ = load_streaming( + str(model_dir), + StreamingConfig(window_size=1, verbose=False), + load_tokenizer=False, + ) + + cache = model.make_cache() + inputs = mx.array([[1, 2, 3]]) + logits = model(inputs, cache=cache) + mx.eval(logits) + + assert isinstance(logits, mx.array) + assert logits.ndim == 3 + assert logits.shape[-1] == _tiny_llama_config()["vocab_size"] + + +def test_generate_step_on_streaming_model(): + with tempfile.TemporaryDirectory() as tmp: + model_dir = _build_tiny_model_dir(Path(tmp)) + model, _, _ = load_streaming( + str(model_dir), + StreamingConfig(window_size=1, verbose=False), + load_tokenizer=False, + ) + + prompt = mx.array([1, 2, 3]) + tokens = [] + for token, _logprobs in generate_step(prompt, model, max_tokens=3): + if isinstance(token, mx.array): + tokens.append(int(token.item())) + else: + tokens.append(int(token)) + + assert len(tokens) == 3 + + +def test_generate_step_sampling_returns_mx_array_logprobs(): + """generate_step logprobs must be mx.array through the streaming path.""" + with tempfile.TemporaryDirectory() as tmp: + model_dir = _build_tiny_model_dir(Path(tmp)) + model, _, _ = load_streaming( + str(model_dir), + StreamingConfig(window_size=1, verbose=False), + load_tokenizer=False, + ) + + prompt = mx.array([1, 2, 3]) + for _token, logprobs in generate_step(prompt, model, max_tokens=2): + assert isinstance(logprobs, mx.array) + assert logprobs.ndim == 1 + assert logprobs.shape[0] == _tiny_llama_config()["vocab_size"] \ No newline at end of file diff --git a/tests/test_turboquant.py b/tests/test_turboquant.py new file mode 100644 index 000000000..41358083c --- /dev/null +++ b/tests/test_turboquant.py @@ -0,0 +1,432 @@ +# Copyright © 2025 Bonsai Demo contributors. + +import unittest + +import mlx.core as mx + +from mlx_lm.models.cache import KVCache, save_prompt_cache, load_prompt_cache +from mlx_lm.turboquant.cache import AsymmetricTurboQuantCache +from mlx_lm.turboquant.factory import make_turboquant_cache +from mlx_lm.turboquant.packing import pack_indices, unpack_indices +from mlx_lm.turboquant.quantize import ( + decode_mse, + decode_prod, + encode_kv, + encode_mse, + encode_prod, +) +from mlx_lm.turboquant.rotation import make_rotation_matrix +from mlx_lm.turboquant.qjl import make_qjl_matrix, qjl_packed_dim +from mlx_lm.turboquant.attention import ( + av_weighted_sum_vectorized, + qk_scores_vectorized, + turboquant_scaled_dot_product_attention, + _qk_scores_reference, + _sdpa_decode_fallback, +) +from mlx_lm.turboquant.kernels import ( + av_weighted_sum_metal, + decode_mse_metal, + decode_prod_metal, + encode_mse_metal, + encode_prod_metal, + metal_available, + qk_scores_metal, +) +from mlx_lm.turboquant.quantize import _encode_mse_ref, _encode_prod_ref + + +class TurboQuantTests(unittest.TestCase): + def test_pack_roundtrip(self): + for bits in (2, 3, 4): + dim = 128 + indices = mx.random.randint(0, 2**bits, (4, 8, dim)).astype(mx.uint8) + packed = pack_indices(indices, bits) + restored = unpack_indices(packed, bits, dim) + self.assertTrue(mx.array_equal(indices, restored)) + + def test_mse_roundtrip_error(self): + dim = 128 + rotation = make_rotation_matrix(dim, 7) + vec = mx.random.normal(shape=(2, 3, dim)).astype(mx.float32) + packed, norms = encode_mse(vec, rotation, bits=3) + restored = decode_mse(packed, norms, rotation, bits=3, dim=dim) + err = mx.mean((vec - restored) ** 2).item() + self.assertLess(err, 0.05) + + def test_prod_reconstruction(self): + dim = 128 + rotation = make_rotation_matrix(dim, 11) + s_matrix = make_qjl_matrix(dim, 11) + vec = mx.random.normal(shape=(1, 2, dim)).astype(mx.float32) + packed, norms, signs, gamma = encode_prod(vec, rotation, s_matrix, bits=4) + restored = decode_prod( + packed, norms, signs, gamma, rotation, s_matrix, bits=4, dim=dim + ) + num = mx.linalg.norm(vec - restored) + den = mx.linalg.norm(vec) + self.assertLess((num / den).item(), 0.45) + + def test_qjl_signs_packed(self): + dim = 128 + s_matrix = make_qjl_matrix(dim, 3) + vec = mx.random.normal(shape=(2, dim)).astype(mx.float32) + rotation = make_rotation_matrix(dim, 5) + packed, _, signs, gamma = encode_prod(vec, rotation, s_matrix, bits=4) + self.assertEqual(signs.shape[-1], qjl_packed_dim(dim)) + self.assertEqual(signs.dtype, mx.uint32) + restored = decode_prod( + packed, mx.ones((2, 1)), signs, gamma, rotation, s_matrix, bits=4, dim=dim + ) + self.assertEqual(restored.shape, vec.shape) + + def test_cache_update_shapes(self): + cache = AsymmetricTurboQuantCache(head_dim=128, k_bits=4, v_bits=3, seed=1) + keys = mx.random.normal(shape=(1, 8, 16, 128)).astype(mx.float16) + values = mx.random.normal(shape=(1, 8, 16, 128)).astype(mx.float16) + k_out, v_out = cache.update_and_fetch(keys, values) + self.assertEqual(len(k_out), 4) + self.assertEqual(len(v_out), 2) + self.assertEqual(k_out[0].shape[:3], (1, 8, 16)) + self.assertEqual(v_out[0].shape[:3], (1, 8, 16)) + self.assertEqual(cache.offset, 16) + self.assertGreater(cache.nbytes, 0) + self.assertEqual(cache._k_qjl_signs.shape[-1], qjl_packed_dim(128)) + self.assertTrue(cache.turboquant) + + k2, v2 = cache.update_and_fetch( + mx.random.normal(shape=(1, 8, 1, 128)).astype(mx.float16), + mx.random.normal(shape=(1, 8, 1, 128)).astype(mx.float16), + ) + self.assertEqual(k2[0].shape[2], 17) + self.assertEqual(v2[0].shape[2], 17) + + def test_factory_layer_adaptive(self): + class Dummy: + layers = [object() for _ in range(6)] + + caches = make_turboquant_cache( + Dummy(), k_bits=4, v_bits=3, fp16_layers=1, head_dim=128, seed=0 + ) + self.assertEqual(len(caches), 6) + self.assertIsInstance(caches[0], type(caches[5])) + from mlx_lm.models.cache import KVCache + + self.assertIsInstance(caches[0], KVCache) + self.assertIsInstance(caches[1], AsymmetricTurboQuantCache) + self.assertIsInstance(caches[4], AsymmetricTurboQuantCache) + self.assertIsInstance(caches[5], KVCache) + + def test_cache_state_roundtrip(self): + import tempfile + import os + + cache = AsymmetricTurboQuantCache(head_dim=128, k_bits=4, v_bits=3, seed=9) + keys = mx.random.normal(shape=(1, 8, 8, 128)).astype(mx.float16) + values = mx.random.normal(shape=(1, 8, 8, 128)).astype(mx.float16) + cache.update_and_fetch(keys, values) + nbytes_before = cache.nbytes + + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "tq_cache.safetensors") + save_prompt_cache(path, [cache]) + loaded = load_prompt_cache(path)[0] + + self.assertIsInstance(loaded, AsymmetricTurboQuantCache) + self.assertEqual(loaded.offset, 8) + self.assertEqual(loaded.nbytes, nbytes_before) + _, _ = loaded.update_and_fetch( + mx.random.normal(shape=(1, 8, 1, 128)).astype(mx.float16), + mx.random.normal(shape=(1, 8, 1, 128)).astype(mx.float16), + ) + self.assertEqual(loaded.offset, 9) + + @unittest.skipUnless(metal_available(), "Metal not available") + def test_metal_encode_kv_parity(self): + dim = 128 + rotation_k = make_rotation_matrix(dim, 19) + rotation_v = make_rotation_matrix(dim, 23) + s_matrix = make_qjl_matrix(dim, 23) + vec = mx.random.normal(shape=(1, 8, 64, dim)).astype(mx.float16) + + ref_k = _encode_prod_ref(vec, rotation_k, s_matrix, bits=4) + ref_v = _encode_mse_ref(vec, rotation_v, bits=3) + met = encode_kv(vec, vec, rotation_k, rotation_v, s_matrix, 4, 3) + + self.assertTrue(mx.array_equal(ref_k[0], met[0])) + self.assertTrue(mx.array_equal(ref_k[2], met[2])) + self.assertTrue(mx.array_equal(ref_v[0], met[4])) + self.assertLess(mx.max(mx.abs(ref_k[1] - met[1])).item(), 1e-5) + self.assertLess(mx.max(mx.abs(ref_k[3] - met[3])).item(), 1e-5) + self.assertLess(mx.max(mx.abs(ref_v[1] - met[5])).item(), 1e-5) + + @unittest.skipUnless(metal_available(), "Metal not available") + def test_metal_encode_parity(self): + dim = 128 + rotation = make_rotation_matrix(dim, 19) + s_matrix = make_qjl_matrix(dim, 23) + vec = mx.random.normal(shape=(1, 8, 64, dim)).astype(mx.float16) + + ref_vp, ref_vn = _encode_mse_ref(vec, rotation, bits=3) + met_vp, met_vn = encode_mse_metal(vec, rotation, bits=3, dim=dim) + self.assertTrue(mx.array_equal(ref_vp, met_vp)) + self.assertLess(mx.max(mx.abs(ref_vn - met_vn)).item(), 1e-5) + + ref_k = _encode_prod_ref(vec, rotation, s_matrix, bits=4) + met_k = encode_prod_metal(vec, rotation, s_matrix, bits=4, dim=dim) + self.assertTrue(mx.array_equal(ref_k[0], met_k[0])) + self.assertTrue(mx.array_equal(ref_k[2], met_k[2])) + self.assertLess(mx.max(mx.abs(ref_k[1] - met_k[1])).item(), 1e-5) + self.assertLess(mx.max(mx.abs(ref_k[3] - met_k[3])).item(), 1e-5) + + @unittest.skipUnless(metal_available(), "Metal not available") + def test_metal_decode_parity(self): + dim = 128 + rotation = make_rotation_matrix(dim, 13) + s_matrix = make_qjl_matrix(dim, 17) + vec = mx.random.normal(shape=(1, 4, 32, dim)).astype(mx.float16) + v_pack, v_norm = encode_mse(vec, rotation, bits=3) + k_pack, k_norm, k_sign, k_gamma = encode_prod(vec, rotation, s_matrix, bits=4) + + ref_v = decode_mse(v_pack, v_norm, rotation, bits=3, dim=dim) + met_v = decode_mse_metal(v_pack, v_norm, rotation, bits=3, dim=dim, dtype=mx.float16) + self.assertLess(mx.max(mx.abs(ref_v.astype(mx.float16) - met_v)).item(), 1e-3) + + ref_k = decode_prod( + k_pack, k_norm, k_sign, k_gamma, rotation, s_matrix, bits=4, dim=dim + ) + met_k = decode_prod_metal( + k_pack, k_norm, k_sign, k_gamma, rotation, s_matrix, bits=4, dim=dim, dtype=mx.float16 + ) + self.assertLess(mx.max(mx.abs(ref_k.astype(mx.float16) - met_k)).item(), 1e-3) + + @unittest.skipUnless(metal_available(), "Metal not available") + def test_metal_av_parity(self): + dim = 128 + cache = AsymmetricTurboQuantCache(head_dim=dim, k_bits=4, v_bits=3, seed=11) + keys = mx.random.normal(shape=(1, 8, 32, dim)).astype(mx.float16) + values = mx.random.normal(shape=(1, 8, 32, dim)).astype(mx.float16) + queries = mx.random.normal(shape=(1, 32, 4, dim)).astype(mx.float16) + _, vpack = cache.update_and_fetch(keys, values) + attn = mx.softmax( + mx.random.normal(shape=(1, 32, 4, 32)).astype(mx.float32), + axis=-1, + precise=True, + ).astype(mx.float16) + met = av_weighted_sum_metal( + attn, *vpack, cache._v_rotation, cache.v_bits, dim, dtype=mx.float16 + ) + ref = av_weighted_sum_vectorized( + attn, *vpack, cache._v_rotation, cache.v_bits, dim + ) + self.assertLess( + mx.max(mx.abs(met.astype(mx.float32) - ref.astype(mx.float32))).item(), 1e-3 + ) + + @unittest.skipUnless(metal_available(), "Metal not available") + def test_metal_qk_scores_parity(self): + dim = 128 + cache = AsymmetricTurboQuantCache(head_dim=dim, k_bits=4, v_bits=3, seed=7) + keys = mx.random.normal(shape=(1, 8, 32, dim)).astype(mx.float16) + values = mx.random.normal(shape=(1, 8, 32, dim)).astype(mx.float16) + queries = mx.random.normal(shape=(1, 32, 4, dim)).astype(mx.float16) + kpack, _ = cache.update_and_fetch(keys, values) + scale = dim**-0.5 + met = qk_scores_metal( + queries, *kpack, cache._k_rotation, cache._k_qjl, 4, dim, scale + ) + ref = qk_scores_vectorized( + queries, *kpack, cache._k_rotation, cache._k_qjl, 4, dim, scale + ) + self.assertLess( + mx.max(mx.abs(met.astype(mx.float32) - ref.astype(mx.float32))).item(), 1e-3 + ) + + def test_fused_qk_scores(self): + dim = 128 + cache = AsymmetricTurboQuantCache(head_dim=dim, k_bits=4, v_bits=3, seed=3) + keys = mx.random.normal(shape=(1, 8, 24, dim)).astype(mx.float16) + values = mx.random.normal(shape=(1, 8, 24, dim)).astype(mx.float16) + queries = mx.random.normal(shape=(1, 32, 4, dim)).astype(mx.float16) + kpack, _ = cache.update_and_fetch(keys, values) + scale = dim**-0.5 + fused = qk_scores_vectorized( + queries, *kpack, cache._k_rotation, cache._k_qjl, 4, dim, scale + ) + ref = _qk_scores_reference( + queries, *kpack, cache._k_rotation, cache._k_qjl, 4, dim, scale + ) + self.assertLess( + mx.max(mx.abs(fused.astype(mx.float32) - ref.astype(mx.float32))).item(), 1e-4 + ) + + def test_fused_sdpa(self): + dim = 128 + cache = AsymmetricTurboQuantCache(head_dim=dim, k_bits=4, v_bits=3, seed=5) + keys = mx.random.normal(shape=(1, 8, 16, dim)).astype(mx.float16) + values = mx.random.normal(shape=(1, 8, 16, dim)).astype(mx.float16) + queries = mx.random.normal(shape=(1, 32, 4, dim)).astype(mx.float16) + kpack, vpack = cache.update_and_fetch(keys, values) + scale = dim**-0.5 + out = turboquant_scaled_dot_product_attention( + queries, kpack, vpack, cache, scale, "causal" + ) + self.assertEqual(out.shape, queries.shape) + + @unittest.skipUnless(metal_available(), "Metal not available") + def test_tq_sdpa_2pass_parity(self): + dim = 128 + for L, S in ((1, 2048), (1, 4096), (32, 2048)): + cache = AsymmetricTurboQuantCache(head_dim=dim, k_bits=4, v_bits=3, seed=23) + keys = mx.random.normal(shape=(1, 8, S, dim)).astype(mx.float16) + values = mx.random.normal(shape=(1, 8, S, dim)).astype(mx.float16) + queries = mx.random.normal(shape=(1, 32, L, dim)).astype(mx.float16) + kpack, vpack = cache.update_and_fetch(keys, values) + scale = dim**-0.5 + met = turboquant_scaled_dot_product_attention( + queries, kpack, vpack, cache, scale, "causal" + ) + ref = _sdpa_decode_fallback( + queries, kpack, vpack, cache, scale, "causal" + ) + self.assertLess( + mx.max(mx.abs(met.astype(mx.float32) - ref.astype(mx.float32))).item(), + 1e-2, + ) + + @unittest.skipUnless(metal_available(), "Metal not available") + def test_tq_sdpa_tiled_parity(self): + dim = 128 + for L, S in ((32, 128), (64, 256), (128, 512)): + cache = AsymmetricTurboQuantCache(head_dim=dim, k_bits=4, v_bits=3, seed=19) + keys = mx.random.normal(shape=(1, 8, S, dim)).astype(mx.float16) + values = mx.random.normal(shape=(1, 8, S, dim)).astype(mx.float16) + queries = mx.random.normal(shape=(1, 32, L, dim)).astype(mx.float16) + kpack, vpack = cache.update_and_fetch(keys, values) + scale = dim**-0.5 + met = turboquant_scaled_dot_product_attention( + queries, kpack, vpack, cache, scale, "causal" + ) + ref = _sdpa_decode_fallback( + queries, kpack, vpack, cache, scale, "causal" + ) + self.assertLess( + mx.max(mx.abs(met.astype(mx.float32) - ref.astype(mx.float32))).item(), + 1e-2, + ) + + @unittest.skipUnless(metal_available(), "Metal not available") + def test_tq_sdpa_vector_parity(self): + dim = 128 + for L, S in ((1, 64), (1, 512), (4, 128), (8, 256)): + cache = AsymmetricTurboQuantCache(head_dim=dim, k_bits=4, v_bits=3, seed=17) + keys = mx.random.normal(shape=(1, 8, S, dim)).astype(mx.float16) + values = mx.random.normal(shape=(1, 8, S, dim)).astype(mx.float16) + queries = mx.random.normal(shape=(1, 32, L, dim)).astype(mx.float16) + kpack, vpack = cache.update_and_fetch(keys, values) + scale = dim**-0.5 + met = turboquant_scaled_dot_product_attention( + queries, kpack, vpack, cache, scale, "causal" + ) + ref = _sdpa_decode_fallback( + queries, kpack, vpack, cache, scale, "causal" + ) + self.assertLess( + mx.max(mx.abs(met.astype(mx.float32) - ref.astype(mx.float32))).item(), + 1e-2, + ) + + @unittest.skipUnless(metal_available(), "Metal not available") + def test_metal_sdpa_parity(self): + dim = 128 + cache = AsymmetricTurboQuantCache(head_dim=dim, k_bits=4, v_bits=3, seed=13) + keys = mx.random.normal(shape=(1, 8, 24, dim)).astype(mx.float16) + values = mx.random.normal(shape=(1, 8, 24, dim)).astype(mx.float16) + queries = mx.random.normal(shape=(1, 32, 4, dim)).astype(mx.float16) + kpack, vpack = cache.update_and_fetch(keys, values) + scale = dim**-0.5 + met = turboquant_scaled_dot_product_attention( + queries, kpack, vpack, cache, scale, "causal" + ) + ref = _sdpa_decode_fallback( + queries, kpack, vpack, cache, scale, "causal" + ) + self.assertLess( + mx.max(mx.abs(met.astype(mx.float32) - ref.astype(mx.float32))).item(), 1e-2 + ) + + def test_tq_smaller_than_fp16(self): + dim = 128 + seq = 256 + tq = AsymmetricTurboQuantCache(head_dim=dim, k_bits=4, v_bits=3, seed=0) + fp16 = KVCache() + keys = mx.zeros((1, 8, seq, dim), dtype=mx.float16) + values = mx.zeros((1, 8, seq, dim), dtype=mx.float16) + tq.update_and_fetch(keys, values) + fp16.update_and_fetch(keys, values) + mx.eval(tq.nbytes, fp16.nbytes) + self.assertLess(tq.nbytes, fp16.nbytes) + + def test_merge_supports_batching(self): + from mlx_lm.turboquant.cache import BatchAsymmetricTurboQuantCache + from mlx_lm.generate import _merge_caches + + caches = [] + for length in (5, 6, 7): + cache = AsymmetricTurboQuantCache(head_dim=128, k_bits=4, v_bits=3, seed=7) + keys = mx.random.normal(shape=(1, 8, length, 128)).astype(mx.float16) + values = mx.random.normal(shape=(1, 8, length, 128)).astype(mx.float16) + cache.update_and_fetch(keys, values) + caches.append([cache]) + + merged = _merge_caches(caches) + self.assertEqual(len(merged), 1) + self.assertIsInstance(merged[0], BatchAsymmetricTurboQuantCache) + self.assertEqual(merged[0]._k_packed.shape[0], 3) + self.assertEqual(merged[0].size(), 7) + + def test_batch_prefill_matches_single(self): + from mlx_lm.turboquant.cache import BatchAsymmetricTurboQuantCache + + dim = 128 + keys_a = mx.random.normal(shape=(1, 8, 5, dim)).astype(mx.float16) + vals_a = mx.random.normal(shape=(1, 8, 5, dim)).astype(mx.float16) + keys_b = mx.random.normal(shape=(1, 8, 3, dim)).astype(mx.float16) + vals_b = mx.random.normal(shape=(1, 8, 3, dim)).astype(mx.float16) + + single_a = AsymmetricTurboQuantCache(head_dim=dim, k_bits=4, v_bits=3, seed=7) + single_b = AsymmetricTurboQuantCache(head_dim=dim, k_bits=4, v_bits=3, seed=7) + single_a.update_and_fetch(keys_a, vals_a) + single_b.update_and_fetch(keys_b, vals_b) + + batch = BatchAsymmetricTurboQuantCache.merge([single_a, single_b]) + extracted_a = batch.extract(0) + extracted_b = batch.extract(1) + mx.eval( + single_a._k_packed, + extracted_a._k_packed, + single_b._k_packed, + extracted_b._k_packed, + ) + self.assertEqual(extracted_a.offset, single_a.offset) + self.assertEqual(extracted_b.offset, single_b.offset) + self.assertTrue( + mx.allclose( + extracted_a._k_norms[..., : extracted_a.offset, :].astype(mx.float32), + single_a._k_norms[..., : single_a.offset, :].astype(mx.float32), + rtol=1e-5, + atol=1e-5, + ) + ) + self.assertTrue( + mx.allclose( + extracted_b._k_norms[..., : extracted_b.offset, :].astype(mx.float32), + single_b._k_norms[..., : single_b.offset, :].astype(mx.float32), + rtol=1e-5, + atol=1e-5, + ) + ) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file