diff --git a/examples/generate/ssd_moe_qwen38/bench.py b/examples/generate/ssd_moe_qwen38/bench.py new file mode 100644 index 0000000000..6aa731eac2 --- /dev/null +++ b/examples/generate/ssd_moe_qwen38/bench.py @@ -0,0 +1,208 @@ +"""Qwen3.8-Flash-Next on a 32 GB Mac with Transformers expert streaming, following Strata's Mac protocol. + +Strata's Mac guide reports decode tokens/s on a 200-token story answer and prefill tokens/s on a 4K-token prompt, +with MTP off on 32 GB Macs. This script runs both on the checkpoint written by ``convert_fp8.py``: +- story: a short story request, exactly 200 new tokens (``min_new_tokens``), greedy; decode tok/s excludes the first token. + The story runs first on a cold cache, then three other prompts follow (``--runs`` repeats this list); the report's + ``decode_median`` is over all of them; +- prefill: a ~4,000-token prompt, time to first token and prefill tok/s. +Normal use (OS page cache on) unless ``--bypass-page-cache``. Writes one JSON report. + +Usage: bench.py CONVERTED_DIR --cache-gib 6 --output report.json [--bits 3] [--runs 1] +""" + +import argparse +import json +import platform +import resource +import threading +import time +from pathlib import Path + +import torch + +from transformers import AutoTokenizer, MetalConfig, Qwen4ExpForCausalLM + + +STORY = "Write a short story about a lighthouse keeper who finds a message in a bottle." +# Decode is measured on different prompts, each once: the expert cache carries over from one generation to the next, so +# repeating a prompt replays the same routing and overstates the hit rate (~91% instead of ~89% for the story). +DECODE_PROMPTS = [ + STORY, + "Write a Python function that parses an ISO 8601 date string without using datetime, with tests.", + "Explain the causes of the French Revolution to a high school student.", + "Translate into French and explain the grammar: 'The children would have gone if it had not rained.'", +] + + +class DevicePeak: + """Peak MPS driver memory, sampled every 20 ms (the allocator's own peak is not exposed on MPS).""" + + def __init__(self): + self.peak, self.running = 0, True + self.thread = threading.Thread(target=self.run, daemon=True) + self.thread.start() + + def run(self): + while self.running: + self.peak = max(self.peak, torch.mps.driver_allocated_memory()) + time.sleep(0.02) + + def stop(self): + self.running = False + self.thread.join() + return self.peak + + +def long_prompt(tokenizer, target=4000): + notes = [ + f"Note {i}: the harbour town kept records of ships, storms, cargo and the keepers who tended its light, " + f"and entry {i} describes one season in detail, with names, dates and the weather." + for i in range(1, 2000) + ] + text = [] + for note in notes: + text.append(note) + if len(tokenizer("\n".join(text)).input_ids) >= target: + break + return "Read these records.\n" + "\n".join(text) + "\nSummarize the history of the harbour town." + + +def generate(model, tokenizer, prompt, new_tokens, prefill_chunk_size=None): + messages = [{"role": "user", "content": prompt}] + try: + inputs = tokenizer.apply_chat_template( + messages, add_generation_prompt=True, return_tensors="pt", return_dict=True, enable_thinking=False + ) + except TypeError: + inputs = tokenizer.apply_chat_template( + messages, add_generation_prompt=True, return_tensors="pt", return_dict=True + ) + ids = inputs["input_ids"].to("mps") + first = {"calls": 0} + + class FirstToken: + """`generate` calls `put` once with the prompt, then once per new token: the second call is the first token.""" + + def put(self, value): + first["calls"] += 1 + if first["calls"] == 2: + first["t"] = time.perf_counter() + + def end(self): + pass + + stats_before = model.ssd_store.stats() + torch.mps.synchronize() + start = time.perf_counter() + with torch.no_grad(): + out = model.generate( + ids, + max_new_tokens=new_tokens, + min_new_tokens=new_tokens, + do_sample=False, + streamer=FirstToken(), + prefill_chunk_size=prefill_chunk_size, + ) + torch.mps.synchronize() + end = time.perf_counter() + stats = model.ssd_store.stats() + generated = out.shape[1] - ids.shape[1] + ttft = first.get("t", end) - start + delta = {k: v - stats_before[k] for k, v in stats.items() if isinstance(v, (int, float)) and k in stats_before} + return { + "prompt_tokens": ids.shape[1], + "new_tokens": generated, + "time_to_first_token_s": ttft, + "prefill_tokens_per_s": ids.shape[1] / ttft, + "decode_tokens_per_s": (generated - 1) / (end - start - ttft), + "total_s": end - start, + "store": delta, + "text": tokenizer.decode(out[0, ids.shape[1] :], skip_special_tokens=True), + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("converted", type=Path) + parser.add_argument("--cache-gib", type=float, required=True) + parser.add_argument("--bits", type=int, default=3) + parser.add_argument("--io-threads", type=int, default=4) + parser.add_argument("--runs", type=int, default=1) + parser.add_argument("--bypass-page-cache", action="store_true") + parser.add_argument("--skip-prefill", action="store_true") + parser.add_argument("--skip-story", action="store_true") + parser.add_argument("--cache-policy", default="lru", choices=("lru", "lfu")) + parser.add_argument("--prefill-chunk-size", type=int, default=None, help="Prefill long prompts in chunks") + parser.add_argument( + "--dense-bits", type=int, default=None, help="Quantize the dense nn.Linear weights (MetalConfig)" + ) + parser.add_argument( + "--no-overlap-shared-expert", action="store_true", help="A/B: shared expert before the routing" + ) + parser.add_argument("--no-mps-commit", action="store_true", help="A/B: no early MPS command-buffer commits") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + dense = args.converted / "dense" + tokenizer = AutoTokenizer.from_pretrained(dense) + peak = DevicePeak() + start = time.perf_counter() + model = Qwen4ExpForCausalLM.from_pretrained( + dense, + dtype=torch.bfloat16, + device_map="mps", + attn_implementation="sdpa", # the indexer's sparse mask rules out the flash kernels + quantization_config=( + MetalConfig(bits=args.dense_bits, group_size=64, modules_to_not_convert=["shared_expert_gate"]) + if args.dense_bits + else None + ), + experts_offload={ + "cache_bytes": int(args.cache_gib * 2**30), + "quantized_experts": args.converted / f"experts-q{args.bits}", + "disk_embeddings": {"model.layers.1.ple.ple_embedding.ngram_embedding": args.converted / "ngram"}, + "io_threads": args.io_threads, + "cache_policy": args.cache_policy, + "bypass_page_cache": args.bypass_page_cache, + "overlap_shared_expert": not args.no_overlap_shared_expert, + "mps_commit": not args.no_mps_commit, + }, + ) + load_s = time.perf_counter() - start + after_load = torch.mps.driver_allocated_memory() + report = { + "model": "Qwen3.8-Flash-Next (from Qwen/Qwen3.8-Flash-Next-FP8)", + "experts": f"{args.bits}-bit MLX affine, group 64", + "cache_gib": args.cache_gib, + "io_threads": args.io_threads, + "page_cache": "bypassed" if args.bypass_page_cache else "on", + "dense": f"{args.dense_bits}-bit MetalConfig" if args.dense_bits else "bf16", + "prefill_chunk_size": args.prefill_chunk_size, + "torch": torch.__version__, + "platform": platform.platform(), + "load_s": load_s, + "device_after_load_bytes": after_load, + "runs": {"story": [], "prefill_4k": []}, + } + for _ in range(0 if args.skip_story else args.runs): + for prompt in DECODE_PROMPTS: + report["runs"]["story"].append({"prompt": prompt, **generate(model, tokenizer, prompt, 200)}) + print(json.dumps({k: v for k, v in report["runs"]["story"][-1].items() if k != "text"}), flush=True) + if report["runs"]["story"]: + rates = sorted(run["decode_tokens_per_s"] for run in report["runs"]["story"]) + report["decode_median"] = (rates[(len(rates) - 1) // 2] + rates[len(rates) // 2]) / 2 + if not args.skip_prefill: + prompt = long_prompt(tokenizer) + for _ in range(args.runs): + report["runs"]["prefill_4k"].append(generate(model, tokenizer, prompt, 32, args.prefill_chunk_size)) + print(json.dumps({k: v for k, v in report["runs"]["prefill_4k"][-1].items() if k != "text"}), flush=True) + report["device_peak_sampled_bytes"] = peak.stop() + report["process_peak_rss_bytes"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + args.output.write_text(json.dumps(report, indent=2)) + if report["runs"]["story"]: + print(report["runs"]["story"][0]["text"][:400]) + print("decode median", report["decode_median"]) + + +if __name__ == "__main__": + main() diff --git a/examples/generate/ssd_moe_qwen38/convert_fp8.py b/examples/generate/ssd_moe_qwen38/convert_fp8.py new file mode 100644 index 0000000000..d2da883626 --- /dev/null +++ b/examples/generate/ssd_moe_qwen38/convert_fp8.py @@ -0,0 +1,226 @@ +"""Convert Qwen/Qwen3.8-Flash-Next-FP8 for SSD streaming on a 32 GB Mac, one shard at a time (peak disk: one shard). + +Output, under OUT_DIR: +- ``dense/``: the text model without experts and without the n-gram table, dequantized to BF16 (~9.9 GB), keys renamed + from ``model.language_model.*`` to ``model.*``, a text-only config, and the tokenizer files. Vision and MTP are skipped. +- ``experts-q{bits}/``: every routed expert quantized with MLX affine quantization (``quantize`` from + kernels-community/mlx-quantization-metal-kernels on MPS), one contiguous row per expert, in the layout + ``transformers.integrations.ssd_moe.quantize_experts`` writes (``model.layers.N.mlp.experts.experts`` [512, bytes]). +- ``ngram/``: the 51.2 GB hashed n-gram table, kept in FP8 E4M3 exactly as published, as one raw [rows, 160] file + plus its scalar scale, for a disk-backed embedding that reads only the rows each token needs. + +FP8 weights carry 128x128 block scales (``weight_scale_inv``): weight = fp8 * scale, per block. +Usage: convert_fp8.py OUT_DIR [--bits 3] [--group-size 64] [--revision REV] +""" + +import argparse +import json +import re +import shutil +import time +from collections import defaultdict +from pathlib import Path + +import torch +from huggingface_hub import hf_hub_download +from kernels import get_kernel +from safetensors import safe_open +from safetensors.torch import save_file + + +REPO = "Qwen/Qwen3.8-Flash-Next-FP8" +EXPERT = re.compile(r"^model\.language_model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$") +NGRAM = re.compile(r"^model\.language_model\.layers\.(\d+)\.ple\.ple_embedding\.ngram_embedding\.shard_(\d+)\.weight$") +TOKENIZER_FILES = ( + "tokenizer.json", + "tokenizer_config.json", + "chat_template.jinja", + "generation_config.json", + "vocab.json", + "merges.txt", + "special_tokens_map.json", +) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("out", type=Path) + parser.add_argument("--bits", type=int, default=3) + parser.add_argument("--group-size", type=int, default=64) + parser.add_argument("--revision", default="main") + args = parser.parse_args() + out, device = args.out, torch.device("mps") + staging, dense_dir = out / "staging", out / "dense" + experts_dir, ngram_dir = out / f"experts-q{args.bits}", out / "ngram" + for directory in (staging, dense_dir, experts_dir, ngram_dir): + directory.mkdir(parents=True, exist_ok=True) + kernel = get_kernel("kernels-community/mlx-quantization-metal-kernels", version=2) + # FP8 E4M3 -> float32 for all 256 bit patterns, so dequantization is a gather on the GPU. + fp8_values = torch.arange(256, dtype=torch.uint8).view(torch.float8_e4m3fn).float().to(device) + + def dequantize(raw, scale_inv): + weight = fp8_values[raw.to(device).view(torch.uint8).int()] + if scale_inv is not None: + rows, columns = weight.shape + scale = ( + scale_inv.to(device).float().repeat_interleave(128, 0)[:rows].repeat_interleave(128, 1)[:, :columns] + ) + weight = weight * scale + return weight + + index = json.loads( + Path(hf_hub_download(REPO, "model.safetensors.index.json", revision=args.revision)).read_text(encoding="utf-8") + ) + weight_map = index["weight_map"] + config = json.loads(Path(hf_hub_download(REPO, "config.json", revision=args.revision)).read_text(encoding="utf-8")) + text = dict(config["text_config"]) + text.pop("quantization_config", None) + text["architectures"] = ["Qwen4ExpForCausalLM"] + num_experts, num_layers = text["num_experts"], text["num_hidden_layers"] + (dense_dir / "config.json").write_text(json.dumps(text, indent=2)) + for name in TOKENIZER_FILES: + try: + shutil.copy(hf_hub_download(REPO, name, revision=args.revision), dense_dir / name) + except Exception as error: # not every repo has every tokenizer file + print(f"skipping {name}: {type(error).__name__}") + + # Resumable per file. Experts are buffered per layer until all of a layer's 512 are quantized (a layer spans two + # files), so a restart reprocesses every file holding experts of a layer that is not written yet. + progress_path = out / "progress.json" + progress = {"done": [], "layers_done": [], "dense_map": {}, "ngram_parts": []} + if progress_path.exists(): + progress.update(json.loads(progress_path.read_text(encoding="utf-8"))) + layers_done = set(progress["layers_done"]) + layer_files = defaultdict(set) + for name, filename in weight_map.items(): + if match := EXPERT.match(name): + layer_files[int(match[1])].add(filename) + redo = {f for layer, names in layer_files.items() if layer not in layers_done for f in names} + progress["done"] = [f for f in progress["done"] if f not in redo] + pending_experts = defaultdict(dict) # layer -> expert -> {projection: bf16 weight on device} + expert_rows = defaultdict(dict) # layer -> expert -> uint8 blob (CPU) + dense, dense_bytes, dense_map = {}, 0, progress["dense_map"] + dense_files = len(list(dense_dir.glob("dense-*.safetensors"))) + ngram_rows = text["split_ngram_parts"] * 2500012 # checked against the first part below + ngram_shape, ngram_file = (ngram_rows, 160), ngram_dir / "ngram.fp8" + + def flush_dense(final=False): + nonlocal dense, dense_bytes, dense_files + if dense and (final or dense_bytes > 2e9): + name = f"dense-{dense_files:05d}.safetensors" + save_file(dense, dense_dir / name, {"format": "pt"}) + dense_map.update(dict.fromkeys(dense, name)) + dense_files += 1 + dense, dense_bytes = {}, 0 + + def finish_expert(layer, expert): + parts = pending_experts[layer].pop(expert) + blob = [] + for weight in (torch.cat((parts["gate"], parts["up"])), parts["down"]): + packed, scales, biases = kernel.quantize( + weight.to(torch.bfloat16), group_size=args.group_size, bits=args.bits + ) + blob += [packed, scales, biases] + expert_rows[layer][expert] = torch.cat([t.cpu().contiguous().view(torch.uint8).reshape(-1) for t in blob]) + if len(expert_rows[layer]) == num_experts: + rows = expert_rows.pop(layer) + layers_done.add(layer) + name = f"experts-{layer:05d}.safetensors" + save_file( + {f"model.layers.{layer}.mlp.experts.experts": torch.stack([rows[e] for e in range(num_experts)])}, + experts_dir / name, + {"format": "pt"}, + ) + print(f" experts of layer {layer} written ({rows[0].numel() / 1e6:.2f} MB per expert)", flush=True) + + files = sorted(set(weight_map.values())) + started = time.perf_counter() + for count, filename in enumerate(files, 1): + if filename in progress["done"]: + continue + if shutil.disk_usage(out).free < 8e9: + raise RuntimeError("Less than 8 GB free; stopping before the next shard") + path = Path(hf_hub_download(REPO, filename, revision=args.revision, local_dir=staging)) + with safe_open(path, framework="pt") as handle: + names = list(handle.keys()) + for name in names: + if name.endswith("weight_scale_inv") or name.startswith(("mtp.", "model.visual.")): + continue + expert = EXPERT.match(name) + ngram = NGRAM.match(name) + if expert: + layer, index_, projection = int(expert[1]), int(expert[2]), expert[3] + if layer in layers_done: + continue + scale = handle.get_tensor(name + "_scale_inv") if name + "_scale_inv" in names else None + weight = dequantize(handle.get_tensor(name), scale).to(torch.bfloat16) + pending_experts[layer].setdefault(index_, {})[projection] = weight + if len(pending_experts[layer][index_]) == 3: + finish_expert(layer, index_) + elif ngram: + tensor = handle.get_tensor(name) + if tuple(tensor.shape) != (ngram_rows // text["split_ngram_parts"], ngram_shape[1]): + raise ValueError(f"Unexpected n-gram part shape {tuple(tensor.shape)}") + if not ngram_file.exists(): + with ngram_file.open("wb") as raw: + raw.truncate(ngram_shape[0] * ngram_shape[1]) + with ngram_file.open("r+b") as raw: + raw.seek(int(ngram[2]) * tensor.numel()) + raw.write(tensor.view(torch.uint8).numpy().tobytes()) + progress["ngram_parts"] = sorted(set(progress["ngram_parts"]) | {int(ngram[2])}) + elif name.endswith("ngram_embedding.weight_scale"): + (ngram_dir / "scale.json").write_text( + json.dumps({"scale": handle.get_tensor(name).float().item()}) + ) + else: + tensor = handle.get_tensor(name) + if tensor.dtype == torch.float8_e4m3fn: + scale = handle.get_tensor(name + "_scale_inv") if name + "_scale_inv" in names else None + tensor = dequantize(tensor, scale).to(torch.bfloat16).cpu() + key = name.replace("model.language_model.", "model.", 1) + dense[key] = tensor.contiguous() + dense_bytes += tensor.numel() * tensor.element_size() + flush_dense() + flush_dense(final=True) # dense tensors of a done file are on disk before it is marked done + path.unlink() + progress["done"].append(filename) + progress["layers_done"] = sorted(layers_done) + progress_path.write_text(json.dumps(progress)) + elapsed = time.perf_counter() - started + print(f"[{count}/{len(files)}] {filename} done, {elapsed / 60:.1f} min", flush=True) + if any(pending_experts.values()) or expert_rows or len(layers_done) != num_layers: + raise RuntimeError(f"Incomplete experts: {sorted(pending_experts)} {sorted(expert_rows)}") + if len(progress["ngram_parts"]) != text["split_ngram_parts"]: + raise RuntimeError(f"Missing n-gram parts: {len(progress['ngram_parts'])} of {text['split_ngram_parts']}") + (dense_dir / "model.safetensors.index.json").write_text(json.dumps({"metadata": {}, "weight_map": dense_map})) + expert_map = { + f"model.layers.{layer}.mlp.experts.experts": f"experts-{layer:05d}.safetensors" for layer in range(num_layers) + } + (experts_dir / "model.safetensors.index.json").write_text(json.dumps({"metadata": {}, "weight_map": expert_map})) + info = { + "format": "mlx-affine-expert-blob", + "bits": args.bits, + "group_size": args.group_size, + "quantizer": "mlx-quantization-metal-kernels", + "source": f"{REPO}@{args.revision}", + } + (experts_dir / "quantization.json").write_text(json.dumps(info, indent=2)) + scale = json.loads((ngram_dir / "scale.json").read_text(encoding="utf-8"))["scale"] + (ngram_dir / "ngram.json").write_text( + json.dumps( + { + "rows": ngram_shape[0], + "dim": ngram_shape[1], + "dtype": "float8_e4m3fn", + "scale": scale, + "file": "ngram.fp8", + }, + indent=2, + ) + ) + shutil.rmtree(staging, ignore_errors=True) + print("done") + + +if __name__ == "__main__": + main() diff --git a/examples/generate/ssd_moe_qwen38/run_guarded.sh b/examples/generate/ssd_moe_qwen38/run_guarded.sh new file mode 100755 index 0000000000..f5c01bb174 --- /dev/null +++ b/examples/generate/ssd_moe_qwen38/run_guarded.sh @@ -0,0 +1,20 @@ +#!/bin/bash +# run_guarded.sh NAME BENCH_ARGS... : bench.py under the same memory watchdog as run_local.sh +set -uo pipefail +OUT=/private/tmp/ssd-moe-local/qwen38-results +POC=/private/tmp/transformers-moe-ssd-streaming +export PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.9 PYTORCH_MPS_LOW_WATERMARK_RATIO=0.8 HF_HUB_OFFLINE=1 +name=$1; shift +swap_mb() { sysctl -n vm.swapusage | awk '{gsub("M","",$6); print int($6)}'; } +PYTHONPATH=$POC/src /private/tmp/ssd-moe-local/venv-kernels/bin/python "$POC/examples/generate/ssd_moe_qwen38/bench.py" \ + /private/tmp/ssd-moe-local/qwen38 --output "$OUT/$name.json" "$@" > "$OUT/$name.log" 2>&1 & +pid=$!; start=$(swap_mb) +while kill -0 "$pid" 2>/dev/null; do + level=$(sysctl -n kern.memorystatus_vm_pressure_level 2>/dev/null || echo 1) + if [ "$level" -ge 4 ] || [ $(( $(swap_mb) - start )) -gt 2048 ]; then + echo "WATCHDOG: pressure $level, swap +$(( $(swap_mb) - start )) MB: killing $pid" | tee -a "$OUT/$name.log" + kill -9 "$pid"; exit 1 + fi + sleep 1 +done +wait "$pid"; echo "exit $?" >> "$OUT/$name.log" diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 8b5072eead..3da6a4ad91 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -803,6 +803,7 @@ _HUB_KERNEL_MAPPING: dict[str, dict[str, str]] = { "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 2}, "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "revision": "ep-support"}, "nvfp4": {"repo_id": "kernels-community/nvfp4-gemm", "version": 1}, + "mlx-quantization": {"repo_id": "kernels-community/mlx-quantization-metal-kernels", "version": 2}, } _KERNEL_MODULE_MAPPING: dict[str, ModuleType | None] = {} @@ -987,6 +988,16 @@ def get_kernel( ) +def _builtin_kernel(func_name: str, args: tuple, kwargs: dict) -> Callable | None: + """A kernel shipped with Transformers for this call, used when the original package is not installed.""" + if func_name == "fused_recurrent_gated_delta_rule" and len(args) >= 3 and args[0].device.type == "mps": + from .metal_gated_delta_rule import gated_delta_rule_step, supports + + if supports(*args[:3], **kwargs): + return gated_delta_rule_step + return None + + def use_kernel_func_from_hub_with_fallback(func_name: str, package: str, internal_path: str | None = None): """ The same as `use_kernel_forward_from_hub` but with the optional fallback to an original package if it exists, e.g., @@ -1040,6 +1051,9 @@ def use_kernel_func_from_hub_with_fallback(func_name: str, package: str, interna ) if not is_new_implementation and not is_torchdynamo_compiling(): + builtin = _builtin_kernel(func_name, args, kwargs) + if builtin is not None: + return builtin(*args, **kwargs) distribution = _PACKAGE_TO_DISTRIBUTION.get(package, package) logger.warning_once( f"`{func_name}` is falling back to its reference PyTorch implementation because " diff --git a/src/transformers/integrations/metal_gated_delta_net.py b/src/transformers/integrations/metal_gated_delta_net.py new file mode 100644 index 0000000000..7dd50fcb69 --- /dev/null +++ b/src/transformers/integrations/metal_gated_delta_net.py @@ -0,0 +1,297 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""One decode step of a Gated DeltaNet layer (Qwen3-Next / Qwen3.5 / Qwen3.8 linear attention) on MPS, fused. + +The PyTorch path runs ~34 ops per layer and token. Here: + +1. the four input projections (qkv, z, b, a) are one matmul over their concatenated weights (the layers keep views + into the merged tensors, so no memory is duplicated); +2. one kernel updates the convolution state, applies the depthwise causal conv + SiLU, and computes the gates + beta = sigmoid(b) and g = -exp(A_log) * softplus(a + dt_bias); +3. one kernel runs the gated delta rule (q/k L2 norm, decay, delta update, output) on the float32 recurrent state in + place, then the gated RMSNorm of each head; +4. `out_proj`. + +Rounding follows the PyTorch path (bfloat16 between the same steps). Compiled at first use with +`torch.mps.compile_shader`. +""" + +import functools + +import torch +import torch.nn.functional as F + + +_SOURCE = """ +#include +using namespace metal; + +// Conv state (C, W) shifts left and takes x; out = silu(T(sum_j state[j] * w[j])). Threads c < H also compute the +// gates of value head c. params: C, W, H. +template +kernel void gdn_conv_gates( + device const T* x [[buffer(0)]], // projections: qkv (C) then z, b (H), a (H) + device T* conv_state [[buffer(1)]], // (C, W), updated in place + device const T* conv_weight [[buffer(2)]],// (C, W) + device const float* a_log [[buffer(3)]], // (H) + device const float* dt_bias [[buffer(4)]],// (H) + device T* qkv [[buffer(5)]], // (C) + device float* g [[buffer(6)]], // (H) + device T* beta [[buffer(7)]], // (H) + device const int* params [[buffer(8)]], // C, W, H, offset of b, offset of a + uint c [[thread_position_in_grid]]) { + const uint C = params[0], W = params[1], H = params[2]; + if (c < H) { + const float b = float(x[params[3] + c]); + beta[c] = T(1.0f / (1.0f + exp(-b))); + const float a = float(x[params[4] + c]) + dt_bias[c]; + const float softplus = a > 20.0f ? a : log(1.0f + exp(a)); + g[c] = -exp(a_log[c]) * softplus; + } + if (c >= C) return; + device T* state = conv_state + c * W; + float acc = 0.0f; + for (uint j = 0; j + 1 < W; ++j) { + state[j] = state[j + 1]; + acc += float(state[j]) * float(conv_weight[c * W + j]); + } + state[W - 1] = x[c]; + acc += float(x[c]) * float(conv_weight[c * W + W - 1]); + const float v = float(T(acc)); + qkv[c] = T(v / (1.0f + exp(-v))); +} + +// One threadgroup per value head, one thread per value column. dims: Hk, Hv, K, V, key_dim, gate_is_sigmoid, +// offset of z in `projections`. q = qkv[0:key_dim], k = qkv[key_dim:2 key_dim], v = qkv[2 key_dim:]; z: (Hv, V). +template +kernel void gdn_delta_norm( + device const T* qkv [[buffer(0)]], + device const float* g [[buffer(1)]], + device const T* beta [[buffer(2)]], + device float* state [[buffer(3)]], // (Hv, K, V) float32, in place + device const T* projections [[buffer(4)]], + device const T* norm_weight [[buffer(5)]],// (V) + device T* out [[buffer(6)]], // (Hv, V) + device const int* dims [[buffer(7)]], + device const float* eps [[buffer(8)]], + uint head [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]], + uint threads [[threads_per_threadgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint Hk = dims[0], Hv = dims[1], K = dims[2], V = dims[3], key_dim = dims[4]; + const bool sigmoid_gate = dims[5] != 0; + device const T* z = projections + dims[6]; + const uint k_head = head / (Hv / Hk); + device const T* q = qkv + k_head * K; + device const T* k = qkv + key_dim + k_head * K; + device const T* v = qkv + 2 * key_dim + head * V; + + threadgroup float qs[256]; + threadgroup float ks[256]; + threadgroup float partial[3][32]; + float q_sq = 0.0f, k_sq = 0.0f; + for (uint i = tid; i < K; i += threads) { + const float qi = float(q[i]), ki = float(k[i]); + qs[i] = qi; + ks[i] = ki; + q_sq += qi * qi; + k_sq += ki * ki; + } + q_sq = simd_sum(q_sq); + k_sq = simd_sum(k_sq); + if (lane == 0) { + partial[0][simd_group] = q_sq; + partial[1][simd_group] = k_sq; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + q_sq = 0.0f; + k_sq = 0.0f; + const uint groups = (threads + 31) / 32; + for (uint s = 0; s < groups; ++s) { + q_sq += partial[0][s]; + k_sq += partial[1][s]; + } + const float q_scale = rsqrt(q_sq + 1e-6f) * rsqrt(float(K)); + const float k_scale = rsqrt(k_sq + 1e-6f); + const float decay = exp(g[head]); + const float beta_t = float(beta[head]); + + // The kernel is launched with V threads: thread `tid` owns value column `tid`. + const uint j = tid; + device float* column = state + head * K * V + j; + float memory = 0.0f; + for (uint i = 0; i < K; ++i) { + memory += column[i * V] * decay * ks[i] * k_scale; + } + const float delta = (float(v[j]) - memory) * beta_t; + float o = 0.0f; + for (uint i = 0; i < K; ++i) { + const float updated = column[i * V] * decay + ks[i] * k_scale * delta; + column[i * V] = updated; + o += updated * qs[i] * q_scale; + } + // The reference rounds the attention output to T, then normalizes in float32. + o = float(T(o)); + float sq = simd_sum(o * o); + if (lane == 0) partial[2][simd_group] = sq; + threadgroup_barrier(mem_flags::mem_threadgroup); + sq = 0.0f; + for (uint s = 0; s < groups; ++s) sq += partial[2][s]; + const T normed = T(o * rsqrt(sq / float(V) + eps[0])); + const T weighted = T(float(norm_weight[j]) * float(normed)); + const float zf = float(z[head * V + j]); + const float gate = sigmoid_gate ? 1.0f / (1.0f + exp(-zf)) : zf / (1.0f + exp(-zf)); + out[head * V + j] = T(float(weighted) * gate); +} + +#define INSTANTIATE(T, NAME) \\ + template [[host_name("gdn_conv_gates_" NAME)]] kernel void gdn_conv_gates( \\ + device const T*, device T*, device const T*, device const float*, device const float*, device T*, \\ + device float*, device T*, device const int*, uint); \\ + template [[host_name("gdn_delta_norm_" NAME)]] kernel void gdn_delta_norm( \\ + device const T*, device const float*, device const T*, device float*, device const T*, device const T*, \\ + device T*, device const int*, device const float*, uint, uint, uint, uint, uint); +INSTANTIATE(bfloat, "bfloat") +INSTANTIATE(half, "half") +INSTANTIATE(float, "float") +""" + +_SUFFIX = {torch.bfloat16: "bfloat", torch.float16: "half", torch.float32: "float"} + + +@functools.cache +def _library(): + return torch.mps.compile_shader(_SOURCE) + + +def _is_metal_linear(linear): + return linear.weight.dtype == torch.uint32 and hasattr(linear, "scales") + + +def _merge_projections(module): + """Concatenate in_proj_qkv / z / b / a along the output dimension, once; the layers keep views into the result.""" + linears = [module.in_proj_qkv, module.in_proj_z, module.in_proj_b, module.in_proj_a] + names = ("weight", "scales", "qbiases") if _is_metal_linear(linears[0]) else ("weight",) + merged = {name: torch.cat([getattr(linear, name).data for linear in linears]) for name in names} + start = 0 + for linear in linears: + rows = linear.weight.shape[0] + for name in names: + setattr(linear, name, torch.nn.Parameter(merged[name][start : start + rows], requires_grad=False)) + start += rows + quantized = None + if len(names) == 3: + quantized = (linears[0].group_size, linears[0].bits) + module._metal_in_proj = (merged, quantized) + module._metal_gate_params = (module.A_log.detach().float(), module.dt_bias.detach().float()) + + +def supports(module, hidden_states, cache_params, attention_mask): + if hidden_states.device.type != "mps" or hidden_states.dtype not in _SUFFIX or torch.is_grad_enabled(): + return False + if hidden_states.shape[0] != 1 or hidden_states.shape[1] != 1 or attention_mask is not None: + return False + if cache_params is None or cache_params.is_recording_past(module.layer_idx): + return False + conv_state = cache_params.get_conv_state(module.layer_idx, state_idx=0) + state = cache_params.get_recurrent_state(module.layer_idx, state_idx=0) + if conv_state is None or state is None or state.dtype != torch.float32 or not state.is_contiguous(): + return False + linears = [module.in_proj_qkv, module.in_proj_z, module.in_proj_b, module.in_proj_a] + if len({_is_metal_linear(linear) for linear in linears}) != 1 or any( + linear.bias is not None for linear in linears + ): + return False + if _is_metal_linear(linears[0]) and len({(linear.group_size, linear.bits) for linear in linears}) != 1: + return False + return ( + conv_state.is_contiguous() + and conv_state.dtype == hidden_states.dtype + and module.conv1d.bias is None + and module.activation in ("silu", "swish") + and module.norm.activation in ("silu", "swish", "sigmoid") + and module.head_k_dim <= 256 + and module.head_v_dim <= 1024 + and module.num_v_heads % module.num_k_heads == 0 + ) + + +@functools.cache +def _int_params(values, device): + return torch.tensor(values, dtype=torch.int32, device=device) + + +@functools.cache +def _float_params(values, device): + return torch.tensor(values, dtype=torch.float32, device=device) + + +def gated_delta_net_step(module, hidden_states, cache_params): + """`GatedDeltaNet.forward` for one cached decode token; updates the cache states in place.""" + if not hasattr(module, "_metal_in_proj"): + _merge_projections(module) + merged, quantized = module._metal_in_proj + library, suffix = _library(), _SUFFIX[hidden_states.dtype] + device = hidden_states.device + x = hidden_states.reshape(1, -1) + if quantized is not None: + from .metal_quantization import load_metal_kernel + + group_size, bits = quantized + projections = load_metal_kernel().quantized_matmul( + x, merged["weight"], merged["scales"], merged["qbiases"], transpose=True, group_size=group_size, bits=bits + ) + else: + projections = F.linear(x, merged["weight"]) + projections = projections.reshape(-1) + conv_dim, value_dim, heads = module.conv_dim, module.value_dim, module.num_v_heads + conv_state = cache_params.get_conv_state(module.layer_idx, state_idx=0) + state = cache_params.get_recurrent_state(module.layer_idx, state_idx=0) + a_log, dt_bias = module._metal_gate_params + qkv = torch.empty(conv_dim, dtype=x.dtype, device=device) + g = torch.empty(heads, dtype=torch.float32, device=device) + beta = torch.empty(heads, dtype=x.dtype, device=device) + width = module.conv_kernel_size + getattr(library, f"gdn_conv_gates_{suffix}")( + projections, + conv_state, + module.conv1d.weight.reshape(conv_dim, width), + a_log, + dt_bias, + qkv, + g, + beta, + _int_params((conv_dim, width, heads, conv_dim + value_dim, conv_dim + value_dim + heads), device), + threads=conv_dim, + ) + out = torch.empty(value_dim, dtype=x.dtype, device=device) + sigmoid_gate = int(module.norm.activation == "sigmoid") + getattr(library, f"gdn_delta_norm_{suffix}")( + qkv, + g, + beta, + state, + projections, + module.norm.weight, + out, + _int_params( + (module.num_k_heads, heads, module.head_k_dim, module.head_v_dim, module.key_dim, sigmoid_gate, conv_dim), + device, + ), + _float_params((module.norm.variance_epsilon,), device), + threads=heads * module.head_v_dim, + group_size=module.head_v_dim, + ) + return module.out_proj(out.reshape(1, 1, value_dim)) diff --git a/src/transformers/integrations/metal_gated_delta_rule.py b/src/transformers/integrations/metal_gated_delta_rule.py new file mode 100644 index 0000000000..3ec4e2d96f --- /dev/null +++ b/src/transformers/integrations/metal_gated_delta_rule.py @@ -0,0 +1,171 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Metal kernel for one decode step of the gated delta rule (Gated DeltaNet linear attention) on MPS. + +One threadgroup per (batch, value head) and one thread per value column of the recurrent state: the query/key L2 norm, +the state decay, the delta update and the output are fused into a single dispatch, in float32 like the PyTorch +reference, with the float32 state updated in place. Compiled at first use with `torch.mps.compile_shader`; no build +step and no extra package. +""" + +import functools + +import torch + + +_SOURCE = """ +#include +using namespace metal; + +// dims: num_k_heads, num_v_heads, k_head_dim, v_head_dim. q, k: (B, Hk, K); v, out: (B, Hv, V); g (float32), beta: +// (B, Hv); state: (B, Hv, K, V) float32, updated in place. +template +kernel void gated_delta_step( + device const T* q [[buffer(0)]], + device const T* k [[buffer(1)]], + device const T* v [[buffer(2)]], + device const float* g [[buffer(3)]], + device const T* beta [[buffer(4)]], + device float* state [[buffer(5)]], + device T* out [[buffer(6)]], + device const int* dims [[buffer(7)]], + uint group [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]], + uint threads [[threads_per_threadgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint num_k_heads = dims[0], num_v_heads = dims[1], K = dims[2], V = dims[3]; + const uint batch = group / num_v_heads, head = group % num_v_heads; + const uint k_head = head / (num_v_heads / num_k_heads); + const uint qk_offset = (batch * num_k_heads + k_head) * K; + const uint bh = batch * num_v_heads + head; + + threadgroup float qs[256]; + threadgroup float ks[256]; + threadgroup float partial[2][32]; + float q_sq = 0.0f, k_sq = 0.0f; + for (uint i = tid; i < K; i += threads) { + const float qi = float(q[qk_offset + i]), ki = float(k[qk_offset + i]); + qs[i] = qi; + ks[i] = ki; + q_sq += qi * qi; + k_sq += ki * ki; + } + q_sq = simd_sum(q_sq); + k_sq = simd_sum(k_sq); + if (lane == 0) { + partial[0][simd_group] = q_sq; + partial[1][simd_group] = k_sq; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + q_sq = 0.0f; + k_sq = 0.0f; + for (uint s = 0; s < (threads + 31) / 32; ++s) { + q_sq += partial[0][s]; + k_sq += partial[1][s]; + } + // l2norm(x) = x * rsqrt(sum(x^2) + 1e-6); the query is also scaled by 1/sqrt(K). + const float q_scale = rsqrt(q_sq + 1e-6f) * rsqrt(float(K)); + const float k_scale = rsqrt(k_sq + 1e-6f); + const float decay = exp(g[bh]); + const float beta_t = float(beta[bh]); + + for (uint j = tid; j < V; j += threads) { + device float* column = state + bh * K * V + j; + float memory = 0.0f; + for (uint i = 0; i < K; ++i) { + memory += column[i * V] * decay * ks[i] * k_scale; + } + const float delta = (float(v[bh * V + j]) - memory) * beta_t; + float output = 0.0f; + for (uint i = 0; i < K; ++i) { + const float updated = column[i * V] * decay + ks[i] * k_scale * delta; + column[i * V] = updated; + output += updated * qs[i] * q_scale; + } + out[bh * V + j] = T(output); + } +} + +template [[host_name("gated_delta_step_bfloat")]] kernel void gated_delta_step( + device const bfloat*, device const bfloat*, device const bfloat*, device const float*, device const bfloat*, + device float*, device bfloat*, device const int*, uint, uint, uint, uint, uint); +template [[host_name("gated_delta_step_half")]] kernel void gated_delta_step( + device const half*, device const half*, device const half*, device const float*, device const half*, + device float*, device half*, device const int*, uint, uint, uint, uint, uint); +template [[host_name("gated_delta_step_float")]] kernel void gated_delta_step( + device const float*, device const float*, device const float*, device const float*, device const float*, + device float*, device float*, device const int*, uint, uint, uint, uint, uint); +""" + +_KERNELS = { + torch.bfloat16: "gated_delta_step_bfloat", + torch.float16: "gated_delta_step_half", + torch.float32: "gated_delta_step_float", +} + + +@functools.cache +def _library(): + return torch.mps.compile_shader(_SOURCE) + + +@functools.cache +def _dims(num_k_heads, num_v_heads, k_dim, v_dim, device): + return torch.tensor([num_k_heads, num_v_heads, k_dim, v_dim], dtype=torch.int32, device=device) + + +def supports(query, key, value, initial_state=None, use_qk_l2norm_in_kernel=False, **kwargs): + """Whether :func:`gated_delta_rule_step` handles this call: one token on MPS, a float32 state to update.""" + return ( + kwargs.get("cu_seqlens") is None + and kwargs.get("output_final_state", True) + and torch.backends.mps.is_available() + and query.device.type == "mps" + and query.shape[1] == 1 + and use_qk_l2norm_in_kernel + and initial_state is not None + and initial_state.dtype == torch.float32 + and initial_state.is_contiguous() + and query.dtype in _KERNELS + and key.dtype == value.dtype == query.dtype + and key.shape[-1] <= 256 + and value.shape[2] % key.shape[2] == 0 + ) + + +def gated_delta_rule_step(query, key, value, g, beta, initial_state, **kwargs): + """One decode step, same contract as the PyTorch recurrent reference with ``use_qk_l2norm_in_kernel=True``. + + query, key: (B, 1, Hk, K); value: (B, 1, Hv, V); g, beta: (B, 1, Hv); initial_state: (B, Hv, K, V) float32, + updated in place and returned. Returns ``(output (B, 1, Hv, V), state)``. + """ + batch, _, num_k_heads, k_dim = key.shape + num_v_heads, v_dim = value.shape[-2:] + out = torch.empty_like(value) + dims = _dims(num_k_heads, num_v_heads, k_dim, v_dim, query.device) + threads = min(v_dim, 1024) + getattr(_library(), _KERNELS[query.dtype])( + query.contiguous(), + key.contiguous(), + value.contiguous(), + g.float().contiguous(), + beta.to(query.dtype).contiguous(), + initial_state, + out, + dims, + threads=batch * num_v_heads * threads, + group_size=threads, + ) + return out, initial_state diff --git a/src/transformers/integrations/metal_hyper_connection.py b/src/transformers/integrations/metal_hyper_connection.py new file mode 100644 index 0000000000..86d8402693 --- /dev/null +++ b/src/transformers/integrations/metal_hyper_connection.py @@ -0,0 +1,198 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Metal kernels for one decode step of a gated hyper-connection (Qwen3.8's `GatedResidual`) on MPS. + +The reference runs ~25 small ops per call around its three projections; here the group RMSNorm, the SiLU of the +low-rank code and the stream mixing (with the injection weights) are one dispatch each, and the projections keep +their own (quantized) matmuls. Compiled at first use with `torch.mps.compile_shader`. +""" + +import functools + +import torch + + +_SOURCE = """ +#include +using namespace metal; + +// One threadgroup per stream: normed[s, h] = x[s, h] * rsqrt(mean(x[s, :]^2) + eps) * (1 + weight[s, h]). +template +kernel void hc_norm( + device const T* x [[buffer(0)]], + device const T* weight [[buffer(1)]], + device T* normed [[buffer(2)]], + device const float* params [[buffer(3)]], // hidden, eps + uint stream [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]], + uint threads [[threads_per_threadgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint hidden = uint(params[0]); + const uint base = stream * hidden; + threadgroup float partial[32]; + float sum = 0.0f; + for (uint h = tid; h < hidden; h += threads) { + const float v = float(x[base + h]); + sum += v * v; + } + sum = simd_sum(sum); + if (lane == 0) partial[simd_group] = sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + sum = 0.0f; + for (uint s = 0; s < (threads + 31) / 32; ++s) sum += partial[s]; + const float scale = rsqrt(sum / float(hidden) + params[1]); + for (uint h = tid; h < hidden; h += threads) { + normed[base + h] = T(float(x[base + h]) * scale * (1.0f + float(weight[base + h]))); + } +} + +// a = silu(T(d / count)), rounded like the reference (bf16 division, then SiLU). +template +kernel void hc_act( + device const T* d [[buffer(0)]], + device T* a [[buffer(1)]], + device const float* params [[buffer(2)]], // count + uint i [[thread_position_in_grid]]) { + const float v = float(T(float(d[i]) / params[0])); + a[i] = T(v / (1.0f + exp(-v))); +} + +// mixed[h] = mean_s sigmoid(u[s, h]) * normed[s, h]; injection[s] = 2 * sigmoid(inject[s] / count). +template +kernel void hc_mix( + device const T* u [[buffer(0)]], + device const T* normed [[buffer(1)]], + device const T* inject [[buffer(2)]], + device T* mixed [[buffer(3)]], + device T* injection [[buffer(4)]], + device const float* params [[buffer(5)]], // hidden, count + uint h [[thread_position_in_grid]]) { + const uint hidden = uint(params[0]), count = uint(params[1]); + if (h < count) { + const float z = float(T(float(inject[h]) / params[1])); + injection[h] = T(2.0f * float(T(1.0f / (1.0f + exp(-z))))); + } + if (h >= hidden) return; + float acc = 0.0f; + for (uint s = 0; s < count; ++s) { + const T gate = T(1.0f / (1.0f + exp(-float(u[s * hidden + h])))); + acc += float(T(float(gate) * float(normed[s * hidden + h]))); + } + mixed[h] = T(acc / float(count)); +} + +#define INSTANTIATE(T, NAME) \\ + template [[host_name("hc_norm_" NAME)]] kernel void hc_norm( \\ + device const T*, device const T*, device T*, device const float*, uint, uint, uint, uint, uint); \\ + template [[host_name("hc_act_" NAME)]] kernel void hc_act(device const T*, device T*, device const float*, \\ + uint); \\ + template [[host_name("hc_mix_" NAME)]] kernel void hc_mix(device const T*, device const T*, device const T*, \\ + device T*, device T*, device const float*, uint); +INSTANTIATE(bfloat, "bfloat") +INSTANTIATE(half, "half") +INSTANTIATE(float, "float") +""" + +_SUFFIX = {torch.bfloat16: "bfloat", torch.float16: "half", torch.float32: "float"} + + +@functools.cache +def _library(): + return torch.mps.compile_shader(_SOURCE) + + +@functools.cache +def _params(values, device): + return torch.tensor(values, dtype=torch.float32, device=device) + + +def supports(module, hyper_input): + """One token on MPS, with the combine projection (`use_combine`).""" + return ( + hyper_input.device.type == "mps" + and hyper_input.numel() == module.hc_count * module.hidden_size + and hyper_input.dtype in _SUFFIX + and module.block_inject_weight is not None + and module.hc_norm.group_size == module.hidden_size + and not torch.is_grad_enabled() + ) + + +def _fused_projections(module): + """``input_mix_weight_down`` and ``block_inject_weight`` as one quantized matmul: both read the normalized input, + and on its own the 4-row ``block_inject_weight`` costs twice the 320-row projection (49 vs 25 us on the M5), while + the fused 324-row call costs as much as the 320-row one. Each output row is quantized on its own, so the results are + bit-exact. The two layers' tensors become row views of the fused ones: no memory is added. + + Returns ``((weight, scales, biases), bits, group_size, rows of the first)``, or ``None`` when the two are not + MetalLinear layers with the same quantization.""" + if "_hc_fused" not in module.__dict__: + from .metal_quantization import MetalLinear + + down, inject, fused = module.input_mix_weight_down, module.block_inject_weight, None + if ( + isinstance(down, MetalLinear) + and isinstance(inject, MetalLinear) + and down.weight.dtype == inject.weight.dtype == torch.uint32 + and (down.bits, down.group_size) == (inject.bits, inject.group_size) + and down.bias is None + and inject.bias is None + ): + tensors = [] + for name in ("weight", "scales", "qbiases"): + first, second = getattr(down, name), getattr(inject, name) + both = torch.cat([first.data, second.data]) + first.data, second.data = both[: first.shape[0]], both[first.shape[0] :] + tensors.append(both) + fused = (tuple(tensors), down.bits, down.group_size, down.out_features) + module.__dict__["_hc_fused"] = fused + return module.__dict__["_hc_fused"] + + +def gated_residual_step(module, hyper_input): + """`GatedResidual.forward` for a single token: returns ``(mixed_input, hyper_input, injection_weights)``.""" + library, suffix = _library(), _SUFFIX[hyper_input.dtype] + count, hidden = module.hc_count, module.hidden_size + device = hyper_input.device + x = hyper_input.contiguous() + normed = torch.empty_like(x) + getattr(library, f"hc_norm_{suffix}")( + x, + module.hc_norm.weight.to(x.dtype), + normed, + _params((float(hidden), module.hc_norm.eps), device), + threads=count * 256, + group_size=256, + ) + fused = _fused_projections(module) + if fused is not None: + from .metal_quantization import load_metal_kernel + + (weight, scales, biases), bits, group_size, rows = fused + both = load_metal_kernel().quantized_matmul( + normed, weight, scales.to(x.dtype), biases.to(x.dtype), transpose=True, group_size=group_size, bits=bits + ) + down, inject = both[..., :rows].contiguous(), both[..., rows:].contiguous() + else: + down, inject = module.input_mix_weight_down(normed), module.block_inject_weight(normed) + code = torch.empty_like(down) + getattr(library, f"hc_act_{suffix}")(down, code, _params((float(count),), device), threads=down.numel()) + up = module.input_mix_weight_up(code) + mixed = torch.empty(*x.shape[:-1], hidden, dtype=x.dtype, device=device) + injection = torch.empty_like(inject) + getattr(library, f"hc_mix_{suffix}")( + up, normed, inject, mixed, injection, _params((float(hidden), float(count)), device), threads=hidden + ) + return mixed, hyper_input, injection diff --git a/src/transformers/integrations/metal_quantization.py b/src/transformers/integrations/metal_quantization.py index 6db0f9a65c..bda28a63df 100644 --- a/src/transformers/integrations/metal_quantization.py +++ b/src/transformers/integrations/metal_quantization.py @@ -23,19 +23,25 @@ This module provides: ``nn.Linear`` with ``MetalLinear``. - ``MetalQuantize`` / ``MetalDequantize``: weight conversion operations that participate in the new ``WeightConverter`` pipeline. + - ``affine_quantize`` / ``affine_dequantize``: MLX's ``mx.quantize`` / ``mx.dequantize`` on MPS, and a torch + reference (``_affine_quantize_tensor`` / ``_affine_dequantize_tensor``) elsewhere, as the GGUF integration does. -Weight layout (transposed, matching ``affine_qmm_t``): +Weight layout (MLX's affine layout, transposed like ``nn.Linear``): - ``weight``: ``[N, K_packed]`` (``uint32``) -- K is the packed dimension. - ``scales``: ``[N, K // group_size]`` (``float16 / bfloat16``) - ``qbiases``: ``[N, K // group_size]`` (same dtype as scales) -The kernel call is ``affine_qmm_t(x, weight, scales, qbiases, group_size, bits)`` -which computes ``y = x @ dequant(weight).T``, identical to ``nn.Linear``. +The kernel call is ``quantized_matmul(x, weight, scales, qbiases, transpose=True, group_size, bits)`` +(MLX's ``mx.quantized_matmul``), which computes ``y = x @ dequant(weight).T``, identical to ``nn.Linear``. """ +from collections.abc import Callable +from dataclasses import dataclass + from ..core_model_loading import ConversionOps, _IdentityOp from ..quantizers.quantizers_utils import should_convert_module from ..utils import is_torch_available, logging +from .hub_kernels import lazy_load_kernel if is_torch_available(): @@ -45,24 +51,42 @@ if is_torch_available(): logger = logging.get_logger(__name__) -_metal_kernel = None +@dataclass(frozen=True) +class MetalKernel: + """Entry points exposed by the `kernels-community/mlx-quantization-metal-kernels` kernel.""" + + quantized_matmul: Callable + quantize: Callable + dequantize: Callable + + +_METAL_KERNEL: MetalKernel | None = None + + +@torch._dynamo.allow_in_graph +def _load_metal_kernel() -> None: + """ + Load the MLX quantization kernel once into the `_METAL_KERNEL` module global. It must not return anything, as an + `@allow_in_graph` node's return must be proxyable: `load_metal_kernel` returns the global instead. + """ + global _METAL_KERNEL + if _METAL_KERNEL is not None: + return + kernel = lazy_load_kernel("mlx-quantization") + if kernel is None: + raise ImportError( + "Failed to load the `kernels-community/mlx-quantization-metal-kernels` kernel from the Hub. Make sure you " + "have `kernels` installed (`pip install kernels`) and are running on an Apple Silicon machine." + ) + _METAL_KERNEL = MetalKernel( + quantized_matmul=kernel.quantized_matmul, quantize=kernel.quantize, dequantize=kernel.dequantize + ) -def _get_metal_kernel(): - """Lazily load the quantization-mlx kernel from Hugging Face Hub.""" - global _metal_kernel - if _metal_kernel is None: - try: - from .hub_kernels import get_kernel - _metal_kernel = get_kernel("kernels-community/mlx-quantization-metal-kernels", version=1) - except Exception as e: - raise ImportError( - f"Failed to load the quantization-mlx kernel from the Hub: {e}. " - "Make sure you have `kernels` installed (`pip install kernels`) " - "and are running on an Apple Silicon machine." - ) from e - return _metal_kernel +def load_metal_kernel() -> MetalKernel: + _load_metal_kernel() + return _METAL_KERNEL # --------------------------------------------------------------------------- @@ -116,16 +140,21 @@ class MetalLinear(nn.Linear): if self.weight.dtype != torch.uint32: return nn.functional.linear(input, self.weight, self.bias) - kernel = _get_metal_kernel() + kernel = load_metal_kernel() - output = kernel.affine_qmm_t( - input, + # On M5 GPUs, mlx-quantization-metal-kernels v2's NAX matmul returns wrong results for bfloat16/float16 inputs + # once a call has enough rows (~30-128 depending on the shape; kernels-community issue #1332); float32 is + # correct. Single-row calls (decode) take the matrix-vector kernel, which is correct in every dtype. + compute_dtype = input.dtype if input.numel() == input.shape[-1] else torch.float32 + output = kernel.quantized_matmul( + input.to(compute_dtype), self.weight, - self.scales.to(input.dtype), - self.qbiases.to(input.dtype), - self.group_size, - self.bits, - ) + self.scales.to(compute_dtype), + self.qbiases.to(compute_dtype), + transpose=True, + group_size=self.group_size, + bits=self.bits, + ).to(input.dtype) if self.bias is not None: output = output + self.bias @@ -182,29 +211,55 @@ def replace_with_metal_linear( return model +def affine_quantize(weight: torch.Tensor, group_size: int, bits: int): + """``mx.quantize`` on MPS, the torch reference elsewhere. Scales and biases come back in ``weight``'s dtype.""" + if weight.device.type == "mps": + return load_metal_kernel().quantize(weight.contiguous(), group_size=group_size, bits=bits) + return _affine_quantize_tensor(weight, group_size, bits) + + +def affine_dequantize(w_packed: torch.Tensor, scales: torch.Tensor, biases: torch.Tensor, group_size: int, bits: int): + """``mx.dequantize`` on MPS, the torch reference elsewhere. Returns ``scales``' dtype.""" + if w_packed.device.type == "mps": + return load_metal_kernel().dequantize(w_packed, scales, biases, group_size=group_size, bits=bits) + return _affine_dequantize_tensor(w_packed, scales, biases, group_size, bits) + + def _affine_quantize_tensor(weight: torch.Tensor, group_size: int, bits: int): """ - Quantize a 2-D float weight ``[N, K]`` into packed uint32 + scales + biases. + Torch reference for ``mx.quantize``: quantize a 2-D float weight ``[N, K]`` into packed uint32 + scales + biases. Returns ``(w_packed, scales, biases)`` with: - ``w_packed``: ``[N, K // (32 // bits)]`` uint32 - - ``scales``: ``[N, K // group_size]`` float32/float16/bfloat16 - - ``biases``: ``[N, K // group_size]`` float32/float16/bfloat16 + - ``scales``: ``[N, K // group_size]`` in ``weight``'s dtype, as ``mx.quantize`` returns them + - ``biases``: ``[N, K // group_size]`` in ``weight``'s dtype """ N, K = weight.shape elems_per_int = 32 // bits max_val = (1 << bits) - 1 n_groups = K // group_size - w_grouped = weight.float().reshape(N, n_groups, group_size) - w_min = w_grouped.min(dim=-1).values # [N, n_groups] - w_max = w_grouped.max(dim=-1).values + # MLX's `affine_quantize` kernel, so a weight quantized here gets the codes `mx.quantize` gives it. + # Metal's `round` is half away from zero, unlike `torch.round`. + def round_half_away(t): + return t.sign() * (t.abs() + 0.5).floor() - scales = ((w_max - w_min) / max_val).clamp(min=1e-8) - biases = w_min + w_grouped = weight.float().reshape(N, n_groups, group_size) + w_min = w_grouped.amin(dim=-1) # [N, n_groups] + w_max = w_grouped.amax(dim=-1).clamp(min=0) # upstream starts its running max at 0 + + scales = ((w_max - w_min) / max_val).clamp(min=1e-7) + # Snap the larger-magnitude edge onto the grid, so it is represented exactly. + side = w_min.abs() > w_max.abs() + scales = torch.where(side, scales, -scales) + edge = torch.where(side, w_min, w_max) + q0 = round_half_away(edge / scales) + at_zero = q0 == 0 + scales = torch.where(at_zero, scales, edge / q0) + biases = torch.where(at_zero, torch.zeros_like(edge), edge) w_int = (w_grouped - biases.unsqueeze(-1)) / scales.unsqueeze(-1) - w_int = w_int.round().clamp(0, max_val).to(torch.int32).reshape(N, K) + w_int = round_half_away(w_int).clamp(0, max_val).to(torch.int32).reshape(N, K) # Pack into uint32 k_packed = K // elems_per_int @@ -212,16 +267,16 @@ def _affine_quantize_tensor(weight: torch.Tensor, group_size: int, bits: int): for i in range(elems_per_int): w_packed |= w_int[:, i::elems_per_int] << (bits * i) - return w_packed.to(torch.uint32), scales, biases + return w_packed.to(torch.uint32), scales.to(weight.dtype), biases.to(weight.dtype) def _affine_dequantize_tensor( w_packed: torch.Tensor, scales: torch.Tensor, biases: torch.Tensor, group_size: int, bits: int ): """ - Dequantize a packed uint32 weight ``[N, K_packed]`` back to float. + Torch reference for ``mx.dequantize``: dequantize a packed uint32 weight ``[N, K_packed]`` back to float. - Returns a ``[N, K]`` float32 tensor. + Returns a ``[N, K]`` tensor in ``scales``' dtype, as ``mx.dequantize`` does. """ N = w_packed.shape[0] elems_per_int = 32 // bits @@ -235,7 +290,7 @@ def _affine_dequantize_tensor( w_grouped = w_flat.reshape(N, -1, group_size) w_deq = w_grouped * scales.float().unsqueeze(-1) + biases.float().unsqueeze(-1) - return w_deq.reshape(N, K) + return w_deq.reshape(N, K).to(scales.dtype) class MetalQuantize(ConversionOps): @@ -256,18 +311,13 @@ class MetalQuantize(ConversionOps): bits = self.hf_quantizer.quantization_config.bits group_size = self.hf_quantizer.quantization_config.group_size - w_packed, scales, biases = _affine_quantize_tensor(value, group_size, bits) + w_packed, scales, biases = affine_quantize(value, group_size, bits) base = target_key.rsplit(".", 1)[0] if "." in target_key else "" scale_key = f"{base}.scales" if base else "scales" bias_key = f"{base}.qbiases" if base else "qbiases" - orig_dtype = value.dtype - return { - target_key: w_packed, - scale_key: scales.to(orig_dtype), - bias_key: biases.to(orig_dtype), - } + return {target_key: w_packed, scale_key: scales, bias_key: biases} class MetalDequantize(ConversionOps): @@ -296,8 +346,7 @@ class MetalDequantize(ConversionOps): continue scales, qbiases = input_dict[prefix + "scales"], input_dict[prefix + "qbiases"] result[key] = [ - _affine_dequantize_tensor(w, s, b, group_size, bits).to(s.dtype) - for w, s, b in zip(weights, scales, qbiases) + affine_dequantize(w, s, b, group_size, bits).to(s.dtype) for w, s, b in zip(weights, scales, qbiases) ] if "weight" in result: # the general converter: one module return {full_layer_name: result["weight"][0]} diff --git a/src/transformers/integrations/ssd_moe.py b/src/transformers/integrations/ssd_moe.py new file mode 100644 index 0000000000..1a14a0fd6d --- /dev/null +++ b/src/transformers/integrations/ssd_moe.py @@ -0,0 +1,1585 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Experimental, inference-only Qwen MoE expert streaming from SSD, without Accelerate or native I/O dependencies. + +Supported checkpoints: ``qwen3_moe``, ``olmoe`` and the text model of ``qwen3_5_moe`` (Qwen3.5/3.6 MoE), unquantized. +The experts can instead come from a llama.cpp GGUF of the same model (Q8_0, F16, BF16 or F32 expert tensors), so that +the streamed bytes match what llama.cpp-based engines read; dense weights still come from the safetensors checkpoint. + +How it works +============ + +1. The hook Transformers already has. + In Qwen3 MoE and Qwen3.5/3.6 MoE, each layer's experts live in one ``Experts`` module as two packed tensors, + ``gate_up_proj`` ``[experts, 2 * intermediate, hidden]`` and ``down_proj`` ``[experts, hidden, intermediate]``. + The router produces ``top_k_index`` / ``top_k_weights``, then the experts module calls whatever function is + registered under ``config._experts_implementation`` through ``ExpertsInterface`` (the same kind of registry as + attention). No model code is patched: this module registers ``"ssd_streaming"`` and points the config at it. + +2. Loading without the experts in memory (``load_moe_from_ssd``). + - Build the model on the ``meta`` device: nothing is allocated. + - For every experts module, delete the two big parameters (``register_parameter(..., None)``) and attach the + shared ``ExpertStore`` plus the module's checkpoint prefix (e.g. ``model.language_model.layers.12.mlp.experts``). + - Load everything else (attention, linear attention, router, shared expert, embeddings, norms) one tensor at a + time, straight to the device. + - Check that every checkpoint tensor was either loaded, deliberately left on disk (experts), or skipped (vision + tower, MTP head). + +3. Reading one expert (``SafetensorsReader`` / ``GGUFExpertReader``). + - safetensors: the header gives each tensor's byte offset; in a packed bank expert ``e`` is one contiguous slice at + ``offset + e * prod(shape[1:]) * itemsize``. One ``readinto`` fills a buffer, ``torch.frombuffer`` wraps it + without a copy, ``.to(device)`` uploads it. A whole bank or shard is never read. + - GGUF: a small header parser finds ``blk.N.ffn_{gate,up,down}_exps``. ggml stores those as + ``[columns, rows, experts]``, so one expert is again a contiguous ``rows x columns`` matrix in the PyTorch + ``(out_features, in_features)`` orientation. A Q8_0 row is blocks of 34 bytes (one float16 scale and 32 int8 + values), so a Qwen3.6 expert is 3.2 MiB instead of 6 MiB in BF16. + +4. The cache (``ExpertStore``). + One global LRU across all layers, keyed by ``(layer prefix, expert id)`` and budgeted in bytes. A miss evicts the + least recently used experts until the new one fits, reads it, uploads it and inserts it. In GGUF mode the cache + holds the GGUF's own bytes (e.g. Q8_0) and ``dequantize_gguf`` turns them into the execution dtype on the device + each time the expert runs, so the same budget holds about twice as many Q8_0 experts as BF16 ones. + +5. The forward pass (``streamed_experts_forward``):: + + routing = top_k_index.cpu() # one small device-to-host copy per layer + for expert in unique(routing): + tokens, slots = where(routing == expert) + gate_up, down = store.get(layer, expert) # cache hit, or read from the SSD + for each chunk of at most token_chunk_size tokens: # bounds prefill memory + gate, up = linear(hidden[tokens], gate_up).chunk(2) + out = linear(act(gate) * up, down) * top_k_weights[tokens, slots] + result.index_add_(0, tokens, out) + + This is the same math as the eager experts. The tests compare logits and greedy generations with the eager model + on tiny Qwen3 MoE and Qwen3.5 MoE checkpoints (packed and per-expert layouts, forced evictions, zero cache, and the + Q8_0 GGUF path). + +6. Where the time goes. + Measured on Qwen3.6-35B-A3B Q8_0 with an A10G limited to 12 GiB of GPU memory and 12 GiB of RAM: each decoded token + makes 40 layers x 8 = 320 expert lookups; with a 5.5 GiB cache and a 77% hit rate that is about 75 misses, about + 250 MB read per token. At 2.2 tok/s that is about 0.55 GB/s, far below what the disk delivers (5 GB/s direct). + The limits are elsewhere: reads are synchronous (the GPU waits on every miss, nothing is prefetched), there is a + Python loop over experts in every layer, and Q8_0 is dequantized on every use. The next steps would be reading + layer N+1's experts while layer N computes, a grouped MoE kernel instead of the loop, and a Q8_0 matmul. + +7. What it deliberately is not. + Inference only, one request at a time, a single device, not wired into ``from_pretrained``. By default reads use + ordinary buffered file I/O, so the OS page cache can help and ``read_bytes`` counts requested bytes, NOT physical + SSD traffic; ``bypass_page_cache=True`` keeps read bytes out of the page cache, so every miss reads the storage. Dense weights stay in the execution dtype even when the experts come from a GGUF. +""" + +import heapq +import json +import math +import os +import re +import struct +import threading +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import torch +from safetensors import safe_open + + +try: # POSIX only; used to keep streamed bytes out of the page cache on macOS (F_NOCACHE) + import fcntl +except ImportError: # Windows + fcntl = None + +from ..models.olmoe.configuration_olmoe import OlmoeConfig +from ..models.olmoe.modeling_olmoe import OlmoeExperts, OlmoeForCausalLM, OlmoeRotaryEmbedding +from ..models.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeTextConfig +from ..models.qwen3_5_moe.modeling_qwen3_5_moe import ( + Qwen3_5MoeExperts, + Qwen3_5MoeForCausalLM, + Qwen3_5MoeTextRotaryEmbedding, +) +from ..models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig +from ..models.qwen3_moe.modeling_qwen3_moe import ( + Qwen3MoeExperts, + Qwen3MoeForCausalLM, + Qwen3MoeRotaryEmbedding, +) +from ..utils import logging +from .moe import ExpertsInterface + + +logger = logging.get_logger(__name__) + + +_DTYPES = {"F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16} +# Quantized experts (see `quantize_experts`) also hold packed uint32 weights. +_PACKED_DTYPES = {**_DTYPES, "U32": torch.uint32, "U8": torch.uint8} +_QUANTIZED_FORMAT = "mlx-affine-expert-blob" + + +def _qwen3_5_moe_text_config(directory): + return Qwen3_5MoeConfig.from_pretrained(directory, local_files_only=True).text_config + + +# model_type -> (text config loader, causal LM class, experts class, rotary class, +# checkpoint prefix of the text model, checkpoint prefixes that are deliberately not loaded) +_ARCHITECTURES = { + "qwen3_moe": ( + lambda directory: Qwen3MoeConfig.from_pretrained(directory, local_files_only=True), + Qwen3MoeForCausalLM, + Qwen3MoeExperts, + Qwen3MoeRotaryEmbedding, + "model.", + (), + ), + # Multimodal Qwen3.5/3.6 MoE checkpoint: load only the text model; skip the vision tower and the MTP head. + "qwen3_5_moe": ( + _qwen3_5_moe_text_config, + Qwen3_5MoeForCausalLM, + Qwen3_5MoeExperts, + Qwen3_5MoeTextRotaryEmbedding, + "model.language_model.", + ("model.visual.", "mtp."), + ), + "olmoe": ( + lambda directory: OlmoeConfig.from_pretrained(directory, local_files_only=True), + OlmoeForCausalLM, + OlmoeExperts, + OlmoeRotaryEmbedding, + "model.", + (), + ), + "qwen3_5_moe_text": ( + lambda directory: Qwen3_5MoeTextConfig.from_pretrained(directory, local_files_only=True), + Qwen3_5MoeForCausalLM, + Qwen3_5MoeExperts, + Qwen3_5MoeTextRotaryEmbedding, + "model.", + ("model.visual.", "mtp."), + ), +} + + +# Reads run on I/O threads (os reads release the GIL); their byte counters are shared. +_COUNTER_LOCK = threading.Lock() + + +def _read_exact(path, offset, size, bypass_page_cache, what, out=None, fd=None): + """Read ``size`` bytes at ``offset``, into ``out`` (a writable buffer) when given. With ``bypass_page_cache``, the + bytes do not stay in the OS page cache (macOS: F_NOCACHE; elsewhere: drop the range after reading), so every miss + is a real storage read. ``fd``: an already open descriptor of ``path`` (positional reads are thread-safe).""" + data = bytearray(size) if out is None else out + own = fd is None + if own: + fd = os.open(path, os.O_RDONLY) + if bypass_page_cache and hasattr(fcntl, "F_NOCACHE"): + fcntl.fcntl(fd, fcntl.F_NOCACHE, 1) + try: + view = memoryview(data).cast("B") + done = 0 + while done < size: + count = os.preadv(fd, [view[done:size]], offset + done) + if not count: + raise OSError(f"Short checkpoint read: {path}, {what}") + done += count + if bypass_page_cache and hasattr(os, "posix_fadvise"): + os.posix_fadvise(fd, offset, size, os.POSIX_FADV_DONTNEED) + finally: + if own: + os.close(fd) + return data + + +class SafetensorsReader: + """Index validated safetensors headers, then read whole tensors or contiguous expert slices. + + File handles are scoped to individual reads. No whole-file mapping survives indexing and no full + packed expert tensor is loaded. Both sharded and single-file local checkpoints are supported. + """ + + def __init__(self, directory, bypass_page_cache=False, packed=False): + directory = Path(directory) + self.bypass_page_cache = bypass_page_cache + self.dtypes = _PACKED_DTYPES if packed else _DTYPES + index = directory / "model.safetensors.index.json" + if index.exists(): + weight_map = json.loads(index.read_text())["weight_map"] + filenames = sorted(set(weight_map.values())) + else: + weight_map = None + filenames = ["model.safetensors"] + self.tensors = {} + self.read_bytes = 0 + self.read_calls = 0 + self._fds = {} + for filename in filenames: + path = (directory / filename).resolve() + if not path.is_file(): + raise FileNotFoundError(f"Missing checkpoint shard: {path}") + # Delegate safetensors format validation (offsets, sizes, overlaps) to safetensors. + with safe_open(path, framework="pt", device="cpu") as handle: + keys = set(handle.keys()) + with path.open("rb") as handle: + header_size = int.from_bytes(handle.read(8), "little") + header = json.loads(handle.read(header_size)) + for name in keys: + info = header[name] + if name in self.tensors: + raise ValueError(f"Duplicate checkpoint tensor: {name}") + self.tensors[name] = (path, 8 + header_size + info["data_offsets"][0], info) + if weight_map is not None and weight_map.get(name) != filename: + raise ValueError(f"Checkpoint index disagrees with shard for {name}") + if weight_map is not None and set(weight_map) != set(self.tensors): + raise ValueError("Checkpoint index contains missing tensors") + + def shape(self, name): + return tuple(self.tensors[name][2]["shape"]) + + def read(self, name, expert=None, out=None): + path, offset, info = self.tensors[name] + shape = tuple(info["shape"]) + if info["dtype"] not in self.dtypes: + raise ValueError(f"Only F32/F16/BF16 (and packed) tensors can be read: {name} is {info['dtype']}") + dtype = self.dtypes[info["dtype"]] + item_size = torch.empty((), dtype=dtype).element_size() + if expert is not None: + if len(shape) < 2 or not 0 <= expert < shape[0]: + raise ValueError(f"Invalid expert slice {expert} for {name}: {shape}") + shape = shape[1:] + offset += expert * math.prod(shape) * item_size + size = math.prod(shape) * item_size + if out is not None: + if out.dtype != torch.uint8 or out.numel() != size or out.device.type != "cpu": + raise ValueError(f"`out` must be a CPU uint8 tensor of {size} bytes") + _read_exact(path, offset, size, self.bypass_page_cache, f"tensor {name}", out.numpy(), self._fd(path)) + data = out + else: + data = torch.frombuffer( + _read_exact(path, offset, size, self.bypass_page_cache, f"tensor {name}", fd=self._fd(path)), + dtype=torch.uint8, + ) + with _COUNTER_LOCK: + self.read_bytes += size + self.read_calls += 1 + return data.view(dtype).reshape(shape) + + def _fd(self, path): + """One read-only descriptor per shard, opened on first use and shared by the I/O threads.""" + fd = self._fds.get(path) + if fd is None: + with _COUNTER_LOCK: + fd = self._fds.get(path) + if fd is None: + fd = os.open(path, os.O_RDONLY) + if self.bypass_page_cache and hasattr(fcntl, "F_NOCACHE"): + fcntl.fcntl(fd, fcntl.F_NOCACHE, 1) + self._fds[path] = fd + return fd + + def __del__(self): + for fd in getattr(self, "_fds", {}).values(): + try: + os.close(fd) + except OSError: + pass + + +# GGUF metadata value types -> struct formats (8 = string and 9 = array are handled separately). +_GGUF_SCALARS = {0: " (elements per block, bytes per block). +_GGML_F32, _GGML_F16, _GGML_Q8_0, _GGML_BF16 = 0, 1, 8, 30 +_GGML_BLOCKS = {_GGML_F32: (1, 4), _GGML_F16: (1, 2), _GGML_Q8_0: (32, 34), _GGML_BF16: (1, 2)} +_GGUF_PROJECTIONS = {"gate": "ffn_gate_exps", "up": "ffn_up_exps", "down": "ffn_down_exps"} + + +def _read_gguf_index(handle): + """Parse a GGUF (v2/v3) header: {tensor name: (ggml dims, ggml type, absolute offset)}. Reads no tensor data.""" + + def read(fmt): + return struct.unpack(fmt, handle.read(struct.calcsize(fmt)))[0] + + def string(): + return handle.read(read(" a ``(rows, columns)`` tensor of ``dtype`` on that device.""" + if kind == _GGML_Q8_0: + blocks = raw.view(rows, columns // 32, 34) + scales = blocks[..., :2].contiguous().view(torch.float16).float() + values = blocks[..., 2:].contiguous().view(torch.int8).float() + return (values * scales).to(dtype).reshape(rows, columns) + source = {_GGML_F32: torch.float32, _GGML_F16: torch.float16, _GGML_BF16: torch.bfloat16}[kind] + return raw.view(source).reshape(rows, columns).to(dtype) + + +class ExpertStore: + """One global expert cache across layers, budgeted in bytes of what it keeps on the execution device: the execution + dtype, or the GGUF's own bytes (e.g. Q8_0) when the experts come from a GGUF, dequantized on each use. + + Policies: + - ``"lru"``: evict the least recently used expert. Decoding visits every layer in order, so when the experts used by + one token do not all fit, LRU evicts exactly the expert the next token needs first and the hit rate drops to 0. + - ``"lfu"``: keep a usage count per expert, halved every ``decay_every`` lookups so the set can follow the text. + When the cache is full, the least used resident expert is the victim, and a missing expert is admitted only when + it has been used more than that victim; otherwise it runs once without being cached. The cache settles on the + most used experts instead of cycling through all of them. + + A miss evicts before loading. One active expert plus conversion/staging memory may exceed the cache + budget. Eviction never overwrites a tensor used by queued work; PyTorch owns tensor lifetimes. + The store is synchronous and intended for one inference request at a time. + """ + + def __init__( + self, + reader, + cache_bytes, + device, + dtype, + gguf=None, + policy="lru", + decay_every=4096, + io_threads=0, + quantized=None, + ): + if cache_bytes < 0: + raise ValueError("cache_bytes must be nonnegative") + if policy not in ("lru", "lfu"): + raise ValueError(f"Unknown cache policy {policy!r}; use 'lru' or 'lfu'") + self.policy = policy + self.decay_every = decay_every + self.counts = {} + self.heap = [] # (count, sequence, key) for resident experts; stale entries are skipped lazily + self.sequence = 0 + self.lookups = 0 + self.not_admitted = 0 + self.reader = reader + self.gguf = gguf + self.cache_bytes = cache_bytes + self.device = torch.device(device) + self.dtype = dtype + self.cache = OrderedDict() + self.resident_bytes = 0 + self.peak_cache_bytes = 0 + self.hits = self.misses = self.evictions = 0 + self.layouts = {} + # Unquantized experts live in one preallocated device pool of `cache_bytes // expert_bytes` slots, so a + # layer's resident experts are slices of a single bank that Transformers' grouped_mm / batched_mm kernels + # can index directly. Each layer keeps a device table expert -> slot (`num_slots` when not resident). + self.gate_up_pool = self.down_pool = None + self.pools = None + # Quantized experts: (reader of the `quantize_experts` directory, bits, group_size), run by `gather_qmm`. + self.quantized = None + if quantized is not None: + info = json.loads((Path(quantized) / "quantization.json").read_text()) + if info.get("format") != _QUANTIZED_FORMAT: + raise ValueError(f"Unsupported quantized experts format {info.get('format')!r}") + self.quantized = ( + SafetensorsReader(quantized, reader.bypass_page_cache, packed=True), + info["bits"], + info["group_size"], + ) + self.gather_qmm = _load_gather_qmm(self.device) + self.num_slots = 0 + self.free_slots = [] + # expert -> slot per layer, kept on the host (installs and evictions are host-side bookkeeping) and copied to + # the device only when a device-side lookup needs it: one write per install would be one kernel launch. + self.slot_tables = {} + self.host_slot_tables = {} + self.stale_slot_tables = set() + self.total_experts = 0 + self.compute = "grouped_mm" + # Reads of experts that are not resident run on `io_threads` threads: a layer's misses in parallel, and the + # predicted experts of the next layer in the background. They land in host memory and only enter the cache + # when a forward call actually uses them, so a wrong prediction costs storage bandwidth, never a cached expert. + self.executor = ThreadPoolExecutor(io_threads, thread_name_prefix="ssd-moe-io") if io_threads else None + self.pending = {} + self.window = 2 * io_threads + self.prefetched = self.prefetch_used = self.installed_ahead = 0 + # Host buffers of non-blocking host-to-device copies still queued on the device; the next host + # synchronization (each layer's routing copy) proves they are done, see `release_host_buffers`. + self.inflight = [] + # Recycled host buffers of quantized experts (see `_host_buffer`). + self.free_host_buffers = [] + self.quantized_expert_bytes = 0 + # Pinned host buffers for the routing copies, by size; one layer's routing is consumed before the next copy. + self.routing_buffers = {} + self.commit = False # see `_commit_mps` + + def routing_buffer(self, numel, dtype): + buffer = self.routing_buffers.get((numel, dtype)) + if buffer is None: + buffer = self.routing_buffers[(numel, dtype)] = torch.empty(numel, dtype=dtype, pin_memory=True) + return buffer + + def register(self, prefix, module): + shapes = { + "gate_up_proj": (module.num_experts, 2 * module.intermediate_dim, module.hidden_dim), + "down_proj": (module.num_experts, module.hidden_dim, module.intermediate_dim), + } + if self.quantized is not None and f"{prefix}.experts" in self.quantized[0].tensors: + # The quantized experts are what runs; a BF16 copy in the checkpoint (if any) is not read but consumed. + candidates = [f"{prefix}.gate_up_proj", f"{prefix}.down_proj"] + [ + f"{prefix}.{expert}.{projection}.weight" + for expert in range(module.num_experts) + for projection in ("gate_proj", "up_proj", "down_proj") + ] + names = {name: None for name in candidates if name in self.reader.tensors} + packed = True + elif f"{prefix}.gate_up_proj" in self.reader.tensors: + names = {f"{prefix}.{suffix}": shape for suffix, shape in shapes.items()} + packed = True + else: + names = {} + for expert in range(module.num_experts): + for projection in ("gate_proj", "up_proj", "down_proj"): + shape = (module.intermediate_dim, module.hidden_dim) + if projection == "down_proj": + shape = shape[::-1] + names[f"{prefix}.{expert}.{projection}.weight"] = shape + packed = False + for name, shape in names.items(): + if shape is None: + continue + if name not in self.reader.tensors or self.reader.shape(name) != shape: + raise ValueError(f"Missing or incompatible expert tensor {name}; expected {shape}") + self.layouts[prefix] = (packed, 3 * module.intermediate_dim * module.hidden_dim) + if self.gguf is not None: + match = re.search(r"\.layers\.(\d+)\.", f".{prefix}.") + if match is None: + raise ValueError(f"Cannot find the layer index of {prefix}") + layer = int(match.group(1)) + expected = { + "gate": (module.intermediate_dim, module.hidden_dim), + "up": (module.intermediate_dim, module.hidden_dim), + "down": (module.hidden_dim, module.intermediate_dim), + } + layouts = {} + for projection, shape in expected.items(): + kind, rows, columns, experts, _, size = self.gguf.layout(layer, projection) + if (rows, columns) != shape or experts != module.num_experts: + raise ValueError( + f"GGUF layer {layer} {projection} is {experts}x{rows}x{columns}; expected {shape}" + ) + layouts[projection] = (kind, rows, columns, size) + # The safetensors expert banks are deliberately not read: they count as consumed. + self.layouts[prefix] = ("gguf", layer, layouts) + else: + if self.quantized is not None: + self._check_quantized(prefix, module) + self.layouts[prefix] = ("quantized",) + self._register_slots(prefix, module) + return set(names) + + def _quantized_shapes(self, module): + """Per-expert shapes and dtypes of the quantized tensors, in pool order.""" + _, bits, group_size = self.quantized + shapes = [] + for rows, columns in ( + (2 * module.intermediate_dim, module.hidden_dim), + (module.hidden_dim, module.intermediate_dim), + ): + if columns % group_size or (columns * bits) % 32: + raise ValueError(f"Input size {columns} does not fit {bits}-bit groups of {group_size}") + shapes += [ + ((rows, columns * bits // 32), torch.uint32), + ((rows, columns // group_size), self.dtype), + ((rows, columns // group_size), self.dtype), + ] + return shapes + + def _check_quantized(self, prefix, module): + reader, name = self.quantized[0], f"{prefix}.experts" + size = sum(math.prod(shape) * dtype.itemsize for shape, dtype in self._quantized_shapes(module)) + if name not in reader.tensors or reader.shape(name) != (module.num_experts, size): + raise ValueError( + f"Missing or incompatible quantized experts {name}; expected {module.num_experts} x {size} bytes " + "(quantized from a checkpoint in the execution dtype)" + ) + + def _register_slots(self, prefix, module): + if self.quantized is not None: + shapes = self._quantized_shapes(module) + else: + shapes = [ + ((2 * module.intermediate_dim, module.hidden_dim), self.dtype), + ((module.hidden_dim, module.intermediate_dim), self.dtype), + ] + if self.pools is None: + self._pool_shapes = shapes + self.slot_bytes = sum(math.prod(shape) * dtype.itemsize for shape, dtype in shapes) + self.num_slots = self.cache_bytes // self.slot_bytes + if self.quantized is not None: + # One byte buffer per slot holding the expert exactly as stored: a miss is one read and one copy. + # The kernel reads the six tensors through strided views of it. + self.slot_buffer = torch.empty( + (self.num_slots, self.slot_bytes), device=self.device, dtype=torch.uint8 + ) + self.pools = list(_split_expert(self.slot_buffer, shapes)) + self.quantized_expert_bytes = self.slot_bytes + else: + self.pools = [ + torch.empty((self.num_slots, *shape), device=self.device, dtype=dtype) for shape, dtype in shapes + ] + if self.quantized is None: + self.gate_up_pool, self.down_pool = self.pools + self.free_slots = list(range(self.num_slots - 1, -1, -1)) + elif [tuple(pool.shape[1:]) for pool in self.pools] != [shape for shape, _ in shapes]: + raise ValueError("All MoE layers must have the same expert shapes to share the expert pool") + self.slot_tables[prefix] = torch.full((module.num_experts,), self.num_slots, device=self.device) + self.host_slot_tables[prefix] = torch.full((module.num_experts,), self.num_slots) + self.total_experts += module.num_experts + + @property + def pooled(self): + return self.pools is not None + + def all_resident(self): + """True when every expert of the model is cached: nothing can be evicted, no routing needs to reach the host.""" + return self.pooled and len(self.cache) == self.total_experts + + def _count(self, key): + self.lookups += 1 + if self.policy != "lfu": + return + if self.decay_every and self.lookups % self.decay_every == 0: + self.counts = {k: v / 2 for k, v in self.counts.items() if v >= 0.25} + self.heap = [] + for resident in self.cache: + self._push(resident) + self.counts[key] = self.counts.get(key, 0) + 1 + if key in self.cache: + self._push(key) + + def _push(self, key): + self.sequence += 1 + heapq.heappush(self.heap, (self.counts.get(key, 0), self.sequence, key)) + + def _least_used(self): + while self.heap: + count, _, key = self.heap[0] + if key in self.cache and count == self.counts.get(key, 0): + return key, count + heapq.heappop(self.heap) + return None, None + + def _evict(self, victim): + entry = self.cache.pop(victim) + self.resident_bytes -= entry[-1] + self.evictions += 1 + if self.pooled: + self.free_slots.append(entry[0]) + self.host_slot_tables[victim[0]][victim[1]] = self.num_slots + self.stale_slot_tables.add(victim[0]) + + def _make_room(self, key, size, pinned=()): + """Evict until ``size`` fits; returns False when the policy declines to cache ``key``. + + ``pinned`` experts are used by the forward call in progress and are never evicted for it. + """ + if size > self.cache_bytes: + return False + if self.policy == "lfu": + while self.cache and self.resident_bytes + size > self.cache_bytes: + victim, count = self._least_used() + if victim is None or victim in pinned or self.counts.get(key, 0) <= count: + self.not_admitted += 1 + return False + heapq.heappop(self.heap) + self._evict(victim) + return True + while self.cache and self.resident_bytes + size > self.cache_bytes: + victim = next((k for k in self.cache if k not in pinned), None) + if victim is None: + self.not_admitted += 1 + return False + self._evict(victim) + return True + + def _expert_bytes(self, prefix): + layout = self.layouts[prefix] + if layout[0] == "quantized": + return self.slot_bytes + if layout[0] == "gguf": + return sum(entry[3] for entry in layout[2].values()) + return layout[1] * self.dtype.itemsize + + def _read_host(self, prefix, expert): + """Read one expert from storage into host memory (raw GGUF bytes, or gate_up and down in the execution dtype). + + Runs on I/O threads: storage and host work only, no device calls. + """ + layout = self.layouts[prefix] + if layout[0] == "quantized": + # One contiguous read per expert, straight into a recycled pinned buffer: copying pinned memory to the + # device costs the host ~5x less than pageable memory (0.03 vs 0.14 ms for a 2 MB expert on the M5). + return (self.quantized[0].read(f"{prefix}.experts", expert, out=self._host_buffer()),) + if layout[0] == "gguf": + layer = layout[1] + return tuple(self.gguf.read(layer, p, expert) for p in ("gate", "up")), self.gguf.read( + layer, "down", expert + ) + if layout[0]: + gate_up = self.reader.read(f"{prefix}.gate_up_proj", expert) + else: + gate_up = torch.cat([self.reader.read(f"{prefix}.{expert}.{p}_proj.weight") for p in ("gate", "up")]) + down = self.reader.read(f"{prefix}.down_proj", expert) if layout[0] else None + if down is None: + down = self.reader.read(f"{prefix}.{expert}.down_proj.weight") + return gate_up.to(self.dtype), down.to(self.dtype) + + def prefetch(self, prefix, experts): + """Start reading the experts of ``prefix`` that are neither resident nor already being read.""" + if self.executor is None: + return + for expert in experts: + key = (prefix, expert) + if key not in self.cache and key not in self.pending: + self.pending[key] = self.executor.submit(self._read_host, prefix, expert) + self.prefetched += 1 + + def discard(self, prefix): + """Drop the reads started for ``prefix`` that its forward call did not use (wrong predictions).""" + for key in [key for key in self.pending if key[0] == prefix]: + self.pending.pop(key).cancel() + + def _take_host(self, prefix, expert): + """One expert in host memory: from a read already started by :meth:`prefetch`, else read now.""" + future = self.pending.pop((prefix, expert), None) + if future is not None and not future.cancelled(): + self.prefetch_used += 1 + return future.result() + return self._read_host(prefix, expert) + + def _read(self, prefix, expert): + """One expert on the device, for a single use.""" + tensors = self._take_host(prefix, expert) + self.inflight.append(tensors) + if self.layouts[prefix][0] == "quantized": + return _split_expert(tensors[0].to(self.device, non_blocking=True), self._pool_shapes) + if isinstance(tensors[0], tuple): # GGUF: (gate, up) raw bytes, down raw bytes + return tuple(t.to(self.device, non_blocking=True) for t in tensors[0]), tensors[1].to( + self.device, non_blocking=True + ) + return tuple(t.to(self.device, non_blocking=True) for t in tensors) + + def _host_buffer(self): + """A host buffer of one quantized expert, pinned when the device can use it; called from I/O threads.""" + try: + return self.free_host_buffers.pop() + except IndexError: + return torch.empty(self.quantized_expert_bytes, dtype=torch.uint8, pin_memory=self.device.type != "cpu") + + def device_slot_table(self, prefix): + """The layer's expert -> slot table on the device, refreshed from the host copy when it changed.""" + if prefix in self.stale_slot_tables: + snapshot = self.host_slot_tables[prefix].clone() + self.inflight.append((snapshot,)) # read by a queued non-blocking copy + self.slot_tables[prefix].copy_(snapshot, non_blocking=True) + self.stale_slot_tables.discard(prefix) + return self.slot_tables[prefix] + + def device_slots(self, prefix, routing): + """Slots of host-side ``routing`` (expert IDs) as a device tensor: one small copy instead of a device lookup.""" + slots = self.host_slot_tables[prefix][routing] + self.inflight.append((slots,)) + return slots.to(self.device, non_blocking=True) + + def release_host_buffers(self): + """Called right after a host synchronization: every queued copy has read its host buffer.""" + if self.quantized is not None: + # list.append / pop are atomic, so I/O threads can take buffers back while this runs. + for tensors in self.inflight: + buffer = tensors[0] + if buffer.dtype == torch.uint8 and buffer.numel() == self.quantized_expert_bytes: + self.free_host_buffers.append(buffer) + self.inflight.clear() + + def synchronize(self): + """Wait for queued device work, then free the host buffers of its copies.""" + if self.device.type == "mps": + torch.mps.synchronize() + elif self.device.type == "cuda": + torch.cuda.synchronize(self.device) + self.release_host_buffers() + + def install_ahead(self, prefix, experts): + """Make guessed experts resident before their layer runs: called once this layer's work is queued, so waiting + for their reads and queuing their copies overlaps device work instead of following the next routing copy. + Not counted as lookups; the layer's own `acquire` then counts a hit (and `installed_ahead` says how many).""" + for expert in experts: + key = (prefix, expert) + # Only reads that already finished: waiting here would hold back the next layer's launches for guesses + # that may not be used. Unfinished reads stay pending for the next layer's own lookups. + if key in self.cache or key not in self.pending or not self.pending[key].done(): + continue + if self._install(key, self._expert_bytes(prefix)) is not None: + self.installed_ahead += 1 + + def _install(self, key, size, pinned=()): + if not self._make_room(key, size, pinned): + return None + slot = self.free_slots.pop() + # Non-blocking: a blocking copy waits for the whole device queue each time (1.6 ms vs 0.13 ms for a + # 3.5 MB expert on the M5). The host buffer stays referenced until the next host synchronization. + tensors = self._take_host(*key) + self.inflight.append(tensors) + if self.quantized is not None: + self.slot_buffer[slot].copy_(tensors[0], non_blocking=True) + else: + for pool, tensor in zip(self.pools, tensors): + pool[slot].copy_(tensor, non_blocking=True) + self.host_slot_tables[key[0]][key[1]] = slot + self.stale_slot_tables.add(key[0]) + self.cache[key] = (slot, size) + self.resident_bytes += size + self.peak_cache_bytes = max(self.peak_cache_bytes, self.resident_bytes) + if self.policy == "lfu": + self._push(key) + return slot + + def acquire(self, prefix, expert, pinned=()): + """Make an expert resident in the pool; returns its slot, or None when it runs once without being cached.""" + key = (prefix, expert) + self._count(key) + if key in self.cache: + self.hits += 1 + self.cache.move_to_end(key) + return self.cache[key][0] + self.misses += 1 + return self._install(key, self._expert_bytes(prefix), pinned) + + def get(self, prefix, expert): + """Return ``(gate_up, down, size)`` for one expert, cached or read for this use only.""" + if self.pooled: + slot = self.acquire(prefix, expert) + if slot is None: + return (*self._read(prefix, expert), self._expert_bytes(prefix)) + return (*(pool[slot] for pool in self.pools), self._expert_bytes(prefix)) + key = (prefix, expert) + self._count(key) + if key in self.cache: + self.hits += 1 + self.cache.move_to_end(key) + return self._materialize(prefix, self.cache[key]) + self.misses += 1 + size = self._expert_bytes(prefix) + admit = self._make_room(key, size) + # Keep the GGUF bytes on the device; dequantize each time the expert runs. + result = (*self._read(prefix, expert), size) + if admit: + self.cache[key] = result + self.resident_bytes += size + self.peak_cache_bytes = max(self.peak_cache_bytes, self.resident_bytes) + if self.policy == "lfu": + self._push(key) + return self._materialize(prefix, result) + + def _materialize(self, prefix, entry): + gate_up, down, size = entry + layout = self.layouts[prefix] + if layout[0] != "gguf": + return entry + layouts = layout[2] + gate_up = torch.cat( + [dequantize_gguf(raw, *layouts[p][:3], self.dtype) for raw, p in zip(gate_up, ("gate", "up"))] + ) + return gate_up, dequantize_gguf(down, *layouts["down"][:3], self.dtype), size + + def _readers(self): + return [self.reader] + ([self.gguf] if self.gguf else []) + ([self.quantized[0]] if self.quantized else []) + + def preload(self): + """Load every expert now (reads in parallel), when the cache holds them all: no storage reads afterwards.""" + if not self.pooled or self.num_slots < self.total_experts: + raise ValueError( + f"preload needs room for all {self.total_experts} experts; the cache has {self.num_slots} slots" + ) + for prefix, table in self.slot_tables.items(): + experts = list(range(table.numel())) + self.prefetch(prefix, experts) + for expert in experts: + self.acquire(prefix, expert) + self.synchronize() + self.hits = self.misses = 0 + + def stats(self): + return { + "hits": self.hits, + "misses": self.misses, + "evictions": self.evictions, + "not_admitted": self.not_admitted, + "policy": self.policy, + "slots": self.num_slots, + "prefetched": self.prefetched, + "prefetch_used": self.prefetch_used, + "installed_ahead": self.installed_ahead, + "resident_bytes": self.resident_bytes, + "peak_cache_bytes": self.peak_cache_bytes, + "read_bytes": sum(reader.read_bytes for reader in self._readers()), + "read_calls": sum(reader.read_calls for reader in self._readers()), + } + + +class _PoolExperts: + """The attributes Transformers' experts kernels read from an experts module, with the shared pool as its banks. + + Slot ``num_slots`` marks an expert that is not resident; the kernels' expert-parallel sentinel handling drops it. + """ + + def __init__(self, module, store): + self.module = module + self.num_experts = store.num_slots + self.gate_up_proj = store.gate_up_pool + self.down_proj = store.down_pool + self.act_fn = module.act_fn + self.has_gate = getattr(module, "has_gate", True) + self.has_bias = getattr(module, "has_bias", False) + self.is_transposed = getattr(module, "is_transposed", False) + self.has_post_expert_norm = getattr(module, "has_post_expert_norm", False) + self.post_expert_norm = getattr(module, "post_expert_norm", None) + self._is_expert_parallel = True + + def _apply_gate(self, gate_up_out): + return self.module._apply_gate(gate_up_out) + + +def _run_expert(module, result, hidden_states, top_k_weights, token, slot, gate_up, down): + # Bound the temporary activations during prefill, even when one expert serves many tokens. + for start in range(0, token.numel(), module.ssd_token_chunk_size): + indices = token[start : start + module.ssd_token_chunk_size] + slots = slot[start : start + module.ssd_token_chunk_size] + gate, up = torch.nn.functional.linear(hidden_states[indices], gate_up).chunk(2, dim=-1) + output = torch.nn.functional.linear(module.act_fn(gate) * up, down) + output = output * top_k_weights[indices, slots, None] + result.index_add_(0, indices, output.to(result.dtype)) + + +def _split_expert(blob, shapes): + """Views of one expert's bytes as its quantized tensors: gate_up (weight, scales, biases), then down.""" + tensors, offset = [], 0 + for shape, dtype in shapes: + size = math.prod(shape) * dtype.itemsize + if offset % dtype.itemsize: + raise ValueError("Misaligned quantized expert layout") + tensors.append(blob[..., offset : offset + size].view(dtype).view(*blob.shape[:-1], *shape)) + offset += size + return tuple(tensors) + + +_MLX_QUANTIZATION_KERNEL = ("kernels-community/mlx-quantization-metal-kernels", 2) + + +def _load_gather_qmm(device): + """MLX's quantized MoE matmul from the Hub (Metal) on MPS; elsewhere the slow PyTorch reference below.""" + if device.type == "mps": + try: + from .hub_kernels import get_kernel + + return get_kernel(_MLX_QUANTIZATION_KERNEL[0], version=_MLX_QUANTIZATION_KERNEL[1]).gather_qmm + except Exception as error: # the kernels package is optional + logger.warning(f"Using the PyTorch reference for quantized experts; the Metal kernel failed: {error}") + return _gather_qmm_reference + + +def _gather_qmm_reference(x, w, scales, biases, rhs_indices, group_size, bits, **_): + """`gather_qmm(x [S, 1, K], w [E, N, K * bits / 32], ..., rhs_indices [S], transpose=True)` -> [S, 1, N].""" + from .metal_quantization import _affine_dequantize_tensor + + output = x.new_empty(*x.shape[:-1], w.shape[1]) + for expert in torch.unique(rhs_indices).tolist(): + rows = rhs_indices == expert + weight = _affine_dequantize_tensor(w[expert], scales[expert], biases[expert], group_size, bits) + output[rows] = x[rows] @ weight.to(x.dtype).T + return output + + +def _quantized_mlp(module, store, x, banks, rhs): + """Gated expert MLP of rows ``x`` [S, hidden], row ``i`` through expert ``rhs[i]`` of ``banks``. + + ``sorted_indices`` stays off: in mlx-quantization-metal-kernels v2 (torch 2.15, M5) the sorted path returns wrong + results once a call has 4 or more rows per expert of the bank (prefill), and it is no faster below that. + """ + _, bits, group_size = store.quantized + + def matmul(inputs, weight, scales, biases): + return store.gather_qmm( + inputs[:, None, :], + weight, + scales, + biases, + rhs_indices=rhs, + transpose=True, + group_size=group_size, + bits=bits, + ).squeeze(1) + + return matmul(module._apply_gate(matmul(x, *banks[:3])), *banks[3:]) + + +# Experts per sorted `gather_qmm` call in prefill: bounds the transient bank (64 x 2.15 MB for 3-bit Qwen3.8 experts). +PREFILL_EXPERT_GROUP = 64 + + +def _sorted_kernel_applies(rows, experts): + """Whether `gather_qmm(sorted_indices=True)` takes MLX's batched `gather_qmm_rhs` path: M == 1, at least 16 rows + and 4 rows per expert of the bank. That path is several times faster for prefill.""" + return rows >= 16 and rows >= 4 * experts + + +def _quantized_sorted_forward(module, store, hidden_states, top_k_index, top_k_weights, experts, banks, result): + """Prefill: the rows routed to ``experts`` (sorted by expert) through the sorted `gather_qmm`, over ``banks`` holding + only ``experts``, added into ``result`` (float32). + + Runs in float32: in mlx-quantization-metal-kernels v2 the sorted (NAX) kernel returns wrong results for bfloat16 + and float16 inputs on the M5, but is correct, and still 4-6x faster than the unsorted path, in float32. Only this + call's scales and biases are converted. + """ + _, bits, group_size = store.quantized + top_k = top_k_index.shape[1] + lookup = torch.full((module.num_experts,), -1, dtype=torch.long, device=hidden_states.device) + lookup[torch.tensor(experts, device=hidden_states.device)] = torch.arange( + len(experts), device=hidden_states.device + ) + flat = lookup[top_k_index.reshape(-1)] + selected = torch.nonzero(flat >= 0).squeeze(1) + rhs, permutation = torch.sort(flat[selected], stable=True) + order = selected[permutation] + banks = [bank.float() if bank.is_floating_point() else bank for bank in banks] + + def matmul(inputs, weight, scales, biases): + return store.gather_qmm( + inputs[:, None, :], + weight, + scales, + biases, + rhs_indices=rhs, + transpose=True, + group_size=group_size, + bits=bits, + sorted_indices=True, + ).squeeze(1) + + x = hidden_states.float()[order // top_k] + output = matmul(module._apply_gate(matmul(x, *banks[:3])), *banks[3:]) + output = output * top_k_weights.reshape(-1)[order, None].float() + result.index_add_(0, order // top_k, output) + + +def _quantized_pool_forward(module, store, hidden_states, slots, top_k_weights): + """Resident quantized experts through `gather_qmm`; non-resident routes (slot ``num_slots``) contribute zero.""" + num_tokens, top_k = slots.shape + rhs = slots.reshape(-1) + valid = rhs < store.num_slots + weights = (top_k_weights.reshape(-1) * valid)[:, None] + output = _quantized_mlp(module, store, hidden_states.repeat_interleave(top_k, dim=0), store.pools, rhs * valid) + return (output * weights.to(output.dtype)).view(num_tokens, top_k, -1).sum(dim=1).to(hidden_states.dtype) + + +class _Routing: + """A layer's routing IDs on their way to the host, with the guessed experts of the next layers riding along. + + Routing IDs are small; they are copied once per layer, rather than synchronizing for every expert. The next MoE + layer's experts are guessed from this layer's hidden states with the next router (the residual stream changes + little from one layer to the next). With `ssd_after_next` (prefetch depth 2), the layer after the next one is + guessed the same way and its experts are only read into host memory, giving those reads two layers of compute to + finish; installs stay one layer ahead. + + On MPS/CUDA the copy goes to a pinned host buffer followed by an event, so device work queued after it (the shared + expert) keeps the device busy while the host waits for the routing and does the cache bookkeeping. + """ + + def __init__(self, module, hidden_states, top_k_index): + store = module.ssd_store + self.top_k_index = top_k_index + self.next_layer = getattr(module, "ssd_next", None) + self.predict = self.next_layer is not None and store.executor is not None + self.predict = self.predict and top_k_index.numel() <= module.ssd_predict_max_routes + self.after_next = getattr(module, "ssd_after_next", None) if self.predict else None + parts = [top_k_index.reshape(-1)] + if self.predict: + for _, router in [self.next_layer] + ([self.after_next] if self.after_next is not None else []): + parts.append( + torch.nn.functional.linear(hidden_states, router.weight).topk(router.top_k, dim=-1).indices + ) + routes = torch.cat([part.reshape(-1) for part in parts]) if len(parts) > 1 else parts[0] + self.event = None + if store.device.type in ("mps", "cuda"): + self.host = store.routing_buffer(routes.numel(), routes.dtype) + self.host.copy_(routes, non_blocking=True) + self.event = torch.mps.Event() if store.device.type == "mps" else torch.cuda.Event() + self.event.record() + else: + self.host = routes.cpu() + + def wait(self): + if self.event is not None: + self.event.synchronize() + return self.host + + +def _commit_mps(*_): + """Hand the queued MPS work to the GPU without waiting for it. + + PyTorch's MPS stream encodes small kernels into a command buffer that it only commits at the next synchronization + (or when memory runs low), so after each per-layer routing wait the GPU idles while the host encodes the next + layer. Recording an event commits the buffer (measured on an M5: 1,000 small matvecs launched then 106 ms of host + work take 213 ms without it, 109 ms with it).""" + torch.mps.Event().record() + + +def _overlap_shared_expert(block): + """Queue a shared-expert MoE block's shared expert after its routing copy instead of before it. + + ``Qwen*SparseMoeBlock.forward`` runs the shared expert, then the router, then the experts, whose routing copy waits + for all of it; here the routing copy is started first, so the shared expert runs on the device while the host + waits for the routing and does the cache bookkeeping. Same operations, same results. + """ + + def forward(hidden_states): + batch_size, sequence_length, hidden_dim = hidden_states.shape + flat = hidden_states.view(-1, hidden_dim) + _, routing_weights, selected_experts = block.gate(flat) + if not block.experts.ssd_store.all_resident(): + block.experts._ssd_routing = _Routing(block.experts, flat, selected_experts) + shared_expert_output = block.shared_expert(flat) + shared_expert_output = torch.nn.functional.sigmoid(block.shared_expert_gate(flat)) * shared_expert_output + expert_output = block.experts(flat, selected_experts, routing_weights) + shared_expert_output + return expert_output.reshape(batch_size, sequence_length, hidden_dim) + + block.forward = forward + + +def streamed_experts_forward(module, hidden_states, top_k_index, top_k_weights): + if module.training or torch.is_grad_enabled(): + raise RuntimeError("SSD expert streaming requires eval() and torch.no_grad() / inference_mode()") + store = module.ssd_store + if hidden_states.device != store.device or hidden_states.dtype != store.dtype: + raise ValueError( + "Moving/casting a streamed model after loading is unsupported; reload on the target device/dtype" + ) + prefix = module.ssd_prefix + if store.pooled: + kernel = ExpertsInterface._global_mapping[store.compute] + if store.all_resident(): + table = store.device_slot_table(prefix) + # Nothing to load or evict: stay on the device, no host synchronization (counted as one hit per route). + store.hits += top_k_index.numel() + if store.quantized is not None: + return _quantized_pool_forward(module, store, hidden_states, table[top_k_index], top_k_weights) + return kernel(_PoolExperts(module, store), hidden_states, table[top_k_index], top_k_weights) + # The routing copy may already be on its way (started by the MoE block before its shared expert, see + # `_overlap_shared_expert`); otherwise start it now. + pending = module.__dict__.pop("_ssd_routing", None) + if pending is None or pending.top_k_index is not top_k_index: + pending = _Routing(module, hidden_states, top_k_index) + routes = pending.wait() + predict, next_layer, after_next = pending.predict, pending.next_layer, pending.after_next + routing = routes[: top_k_index.numel()].view(top_k_index.shape) + if predict: + next_prefix = next_layer[0] + # That wait covered all device work queued before the routing copy, including earlier host-to-device expert copies. + store.release_host_buffers() + experts = torch.unique(routing).tolist() + if experts and not (experts[0] >= 0 and experts[-1] < module.num_experts): + raise ValueError("Expert parallelism and sentinel expert IDs are unsupported") + # This layer's misses are read in parallel (a window ahead of use, bounding host memory); the guessed next-layer + # reads queue behind them. + store.prefetch(prefix, experts[: store.window]) + width = top_k_index.shape[0] * (next_layer[1].top_k if predict else 0) + guessed = torch.unique(routes[top_k_index.numel() : top_k_index.numel() + width]).tolist() if predict else [] + if predict: + store.prefetch(next_prefix, guessed) + if after_next is not None: + store.prefetch(after_next[0], torch.unique(routes[top_k_index.numel() + width :]).tolist()) + try: + result = _streamed_experts(module, store, prefix, hidden_states, top_k_index, top_k_weights, routing, experts) + finally: + store.discard(prefix) + if store.commit: + _commit_mps() + if guessed and store.pooled: + # This layer's experts are queued on the device: install the next layer's guesses while it computes. + store.install_ahead(next_prefix, guessed) + return result + + +def _streamed_experts(module, store, prefix, hidden_states, top_k_index, top_k_weights, routing, experts): + if not store.pooled: + result = torch.zeros_like(hidden_states) + for position, expert in enumerate(experts): + store.prefetch(prefix, experts[position + 1 : position + 1 + store.window]) + token, slot = (t.to(hidden_states.device) for t in torch.where(routing == expert)) + gate_up, down, _ = store.get(prefix, expert) + _run_expert(module, result, hidden_states, top_k_weights, token, slot, gate_up, down) + return result + pinned, uncached = set(), [] + for position, expert in enumerate(experts): + store.prefetch(prefix, experts[position + 1 : position + 1 + store.window]) + if store.acquire(prefix, expert, pinned) is None: + uncached.append(expert) + else: + pinned.add((prefix, expert)) + if store.quantized is not None and _sorted_kernel_applies(top_k_index.numel(), len(experts)): + # Prefill-sized call: sorted gather_qmm over compact banks of this call's experts, cached or not, a group of + # experts at a time so a prefill that routes to every expert never holds a whole layer of them at once. + result = torch.zeros(hidden_states.shape, dtype=torch.float32, device=hidden_states.device) + for start in range(0, len(experts), PREFILL_EXPERT_GROUP): + group = experts[start : start + PREFILL_EXPERT_GROUP] + per_expert = [] + for position, expert in enumerate(group): + if (prefix, expert) in pinned: + slot = store.cache[(prefix, expert)][0] + per_expert.append([pool[slot] for pool in store.pools]) + else: + store.prefetch(prefix, experts[start + position + 1 : start + position + 1 + store.window]) + per_expert.append(store._read(prefix, expert)) + banks = [torch.stack(parts) for parts in zip(*per_expert)] + del per_expert + _quantized_sorted_forward(module, store, hidden_states, top_k_index, top_k_weights, group, banks, result) + if store.inflight: + store.synchronize() + return result.to(hidden_states.dtype) + if pinned and store.quantized is not None: + slots = store.device_slots(prefix, routing) + result = _quantized_pool_forward(module, store, hidden_states, slots, top_k_weights) + elif pinned: + kernel = ExpertsInterface._global_mapping[store.compute] + result = kernel(_PoolExperts(module, store), hidden_states, store.device_slots(prefix, routing), top_k_weights) + else: + result = torch.zeros_like(hidden_states) + # Experts the cache did not admit (budget smaller than one call's working set, or LFU declined) run once each. + for position, expert in enumerate(uncached): + store.prefetch(prefix, uncached[position + 1 : position + 1 + store.window]) + token, slot = (t.to(hidden_states.device) for t in torch.where(routing == expert)) + tensors = store._read(prefix, expert) + if store.quantized is not None: + banks = [tensor[None] for tensor in tensors] + rhs = torch.zeros_like(token) + output = _quantized_mlp(module, store, hidden_states[token], banks, rhs) * top_k_weights[token, slot, None] + result.index_add_(0, token, output.to(result.dtype)) + else: + _run_expert(module, result, hidden_states, top_k_weights, token, slot, *tensors) + return result + + +def load_moe_from_ssd( + directory, + *, + device="cpu", + dtype=torch.float32, + cache_bytes=0, + token_chunk_size=256, + experts_gguf=None, + bypass_page_cache=False, + cache_policy="lru", + experts_implementation="grouped_mm", + io_threads=4, + prefetch_next_layer=True, + predict_max_routes=64, + prefetch_depth=2, + quantized_experts=None, + preload_experts=False, +): + """Load a local, unquantized Qwen MoE checkpoint without allocating resident expert banks. + + With ``experts_gguf`` (a llama.cpp GGUF of the same model), the experts are streamed from that file instead, in its + own quantization (Q8_0 is dequantized on the device at each use); ``cache_bytes`` then counts the GGUF's bytes. + + Supports ``qwen3_moe`` (original per-expert or packed experts) and the text model of ``qwen3_5_moe``. + Does not require Accelerate, checkpoint conversion, kernels, or a second copy of the weights. + The returned model is inference-only: do not train, move/cast, save_pretrained, compile, or change + its experts implementation. Reload the original checkpoint to obtain an ordinary model. + """ + if dtype not in _DTYPES.values() or token_chunk_size <= 0: + raise ValueError("Use a floating dtype (float32/float16/bfloat16) and positive token_chunk_size") + raw_config = json.loads((Path(directory) / "config.json").read_text()) + model_type = raw_config.get("model_type") + if model_type not in _ARCHITECTURES or raw_config.get("quantization_config"): + raise ValueError(f"This prototype supports only unquantized {sorted(_ARCHITECTURES)} checkpoints") + load_config, model_class, experts_class, rotary_class, prefix, skipped = _ARCHITECTURES[model_type] + config = load_config(directory) + config._attn_implementation = "eager" + config._experts_implementation = "ssd_streaming" + ExpertsInterface.register("ssd_streaming", streamed_experts_forward) + # Resolve cuda -> cuda:0 and analogous device aliases for forward validation. + device = torch.empty(0, device=device).device + reader = SafetensorsReader(directory, bypass_page_cache) + gguf = GGUFExpertReader(experts_gguf, bypass_page_cache) if experts_gguf else None + store = ExpertStore( + reader, + cache_bytes, + device, + dtype, + gguf, + policy=cache_policy, + io_threads=io_threads, + quantized=quantized_experts, + ) + if experts_implementation not in ("grouped_mm", "batched_mm"): + raise ValueError("experts_implementation must be 'grouped_mm' or 'batched_mm'") + store.compute = experts_implementation + + def checkpoint_name(name): + return prefix + name.removeprefix("model.") if name.startswith("model.") else name + + with torch.device("meta"): + model = model_class(config) + experts_modules = [] + consumed = {name for name in reader.tensors if name.startswith(skipped)} if skipped else set() + for name, module in model.named_modules(): + if isinstance(module, experts_class): + consumed.update(store.register(checkpoint_name(name), module)) + # Keep no meta Parameters that an offload hook could materialize as whole banks. + module.register_parameter("gate_up_proj", None) + module.register_parameter("down_proj", None) + module.ssd_store = store + module.ssd_prefix = checkpoint_name(name) + module.ssd_token_chunk_size = token_chunk_size + module.ssd_predict_max_routes = predict_max_routes + experts_modules.append((name, module)) + # Load dense weights individually; peak host staging is one dense tensor, not one shard. + for name, parameter in list(model.named_parameters()): + source = checkpoint_name(name) + if config.tie_word_embeddings and name == "lm_head.weight" and name not in reader.tensors: + source = checkpoint_name("model.embed_tokens.weight") + if source not in reader.tensors or reader.shape(source) != tuple(parameter.shape): + raise ValueError(f"Missing or incompatible dense tensor: {source}") + value = reader.read(source).to(device=device, dtype=dtype) + parent, _, attr = name.rpartition(".") + model.get_submodule(parent).register_parameter(attr, torch.nn.Parameter(value, requires_grad=False)) + consumed.add(source) + # Each MoE layer knows the next one's router, to guess and prefetch its experts (decode-sized calls only). + for (name, module), (next_name, _) in zip(experts_modules, experts_modules[1:] + [(None, None)]): + router = model.get_submodule(next_name.rpartition(".")[0]).gate if next_name and prefetch_next_layer else None + module.ssd_next = (checkpoint_name(next_name), router) if router is not None else None + module.ssd_after_next = None + if prefetch_depth >= 2: + for (_, module), (_, following) in zip(experts_modules, experts_modules[1:]): + module.ssd_after_next = following.ssd_next + unexpected = set(reader.tensors) - consumed + if unexpected: + raise ValueError(f"Unsupported checkpoint tensors: {sorted(unexpected)[:10]}") + model.tie_weights() + # RoPE has nonpersistent buffers, so checkpoint loading cannot materialize them. + model.model.rotary_emb = rotary_class(config).to(device) + meta = [n for n, t in list(model.named_parameters()) + list(model.named_buffers()) if t.is_meta] + if meta: + raise RuntimeError(f"Unmaterialized model state after streaming load: {meta[:10]}") + model.eval() + model.ssd_store = store + if preload_experts: + store.preload() + # Exclude dense startup I/O from expert runtime counters. + reader.read_bytes = reader.read_calls = 0 + if store.gguf is not None: + store.gguf.read_bytes = store.gguf.read_calls = 0 + return model + + +# Backward-compatible name used by the first version of this prototype. +load_qwen3_moe_from_ssd = load_moe_from_ssd + + +def _checkpoint_prefix(reader, name, quantized_reader=None): + """The checkpoint key prefix of the experts module ``name``: the same path, possibly under another root (e.g. + ``model.language_model.``), in the packed (``gate_up_proj``) or per-expert (``0.gate_proj.weight``) layout, or + in the quantized experts (``experts``) when the checkpoint has no BF16 experts.""" + suffix = name.split(".", 1)[1] if name.startswith("model.") else name + found = set() + for key in reader.tensors: + if key.endswith(f"{suffix}.gate_up_proj"): + found.add(key.removesuffix(".gate_up_proj")) + elif key.endswith(f"{suffix}.0.gate_proj.weight"): + found.add(key.removesuffix(".0.gate_proj.weight")) + if not found and quantized_reader is not None: + found = {key.removesuffix(".experts") for key in quantized_reader.tensors if key.endswith(f"{suffix}.experts")} + if len(found) != 1: + raise ValueError(f"Cannot find the checkpoint experts of {name} (candidates: {sorted(found)})") + return found.pop() + + +class DiskEmbedding(torch.nn.Module): + """An embedding table that stays on disk; each lookup reads only the rows it needs (through the OS page cache). + + For tables far larger than memory that are read a few rows per token, such as Qwen3.8-Flash-Next's 51 GB hashed + n-gram table. ``directory`` holds ``ngram.json`` (``rows``, ``dim``, ``dtype``: ``float8_e4m3fn``, ``bfloat16`` or + ``float32``, ``scale``, ``file``) and the raw row-major file. Rows are dequantized as ``row * scale``. + """ + + _DTYPES = {"float8_e4m3fn": torch.float8_e4m3fn, "bfloat16": torch.bfloat16, "float32": torch.float32} + + def __init__(self, directory, device, dtype): + super().__init__() + import numpy as np + + info = json.loads((Path(directory) / "ngram.json").read_text(encoding="utf-8")) + self.rows, self.dim = info["rows"], info["dim"] + self.stored_dtype = self._DTYPES[info["dtype"]] + self.scale = float(info.get("scale", 1.0)) + row_bytes = self.dim * self.stored_dtype.itemsize + self.table = np.memmap(Path(directory) / info["file"], dtype=np.uint8, mode="r", shape=(self.rows, row_bytes)) + self.dtype = dtype + self._is_hf_initialized = True + # Callers read `.weight.device` to place their indices. + self.register_buffer("weight", torch.empty(0, device=device, dtype=dtype), persistent=False) + + def forward(self, ids): + flat = ids.reshape(-1).cpu().numpy() + rows = torch.from_numpy(self.table[flat]).view(self.stored_dtype).float() + if self.scale != 1.0: + rows = rows * self.scale + return rows.to(self.dtype).view(*ids.shape, self.dim).to(self.weight.device, non_blocking=True) + + +def prepare_experts_offload(model, checkpoint_files, *, dtype, device, options): + """Called by ``from_pretrained(..., experts_offload=...)`` on the meta-initialized model, before any weight loads. + + Replaces every experts bank with the streaming store: the bank parameters are removed (so the loader never + materializes them), the experts dispatch to ``ssd_streaming``, and the expert tensors are reported as consumed so + the loader skips them. ``options`` are the keyword arguments of :func:`load_moe_from_ssd` that concern the store: + ``cache_bytes`` (required), ``cache_policy``, ``io_threads``, ``prefetch_next_layer``, ``experts_implementation``, + ``bypass_page_cache``, ``token_chunk_size``, ``predict_max_routes``, ``overlap_shared_expert``, ``mps_commit``. + """ + options = dict(options) + unknown = set(options) - { + "cache_bytes", + "cache_policy", + "io_threads", + "prefetch_next_layer", + "experts_implementation", + "bypass_page_cache", + "token_chunk_size", + "predict_max_routes", + "quantized_experts", + "disk_embeddings", + "prefetch_depth", + "overlap_shared_expert", + "mps_commit", + } + if unknown or "cache_bytes" not in options: + raise ValueError(f"experts_offload needs `cache_bytes`; unknown options: {sorted(unknown)}") + if not checkpoint_files or not all(str(f).endswith(".safetensors") for f in checkpoint_files): + raise ValueError("experts_offload streams experts from safetensors checkpoints only") + if dtype not in _DTYPES.values(): + raise ValueError("experts_offload needs a floating dtype (float32/float16/bfloat16)") + ExpertsInterface.register("ssd_streaming", streamed_experts_forward) + device = torch.empty(0, device=device).device + reader = SafetensorsReader(Path(checkpoint_files[0]).parent, options.get("bypass_page_cache", False)) + store = ExpertStore( + reader, + options["cache_bytes"], + device, + dtype, + policy=options.get("cache_policy", "lru"), + io_threads=options.get("io_threads", 4), + quantized=options.get("quantized_experts"), + ) + store.compute = options.get("experts_implementation", "grouped_mm") + experts = [ + (name, module) + for name, module in model.named_modules() + if hasattr(module, "_apply_gate") and isinstance(getattr(module, "gate_up_proj", None), torch.nn.Parameter) + ] + if not experts: + raise ValueError(f"{type(model).__name__} has no experts module that can be offloaded") + consumed = set() + for name, module in experts: + if module.has_bias or module.is_transposed or not module.is_concatenated or module.has_post_expert_norm: + raise ValueError( + f"experts_offload supports gated, concatenated, untransposed experts without bias: {name}" + ) + prefix = _checkpoint_prefix(reader, name, store.quantized[0] if store.quantized else None) + consumed.update(store.register(prefix, module)) + module.register_parameter("gate_up_proj", None) + module.register_parameter("down_proj", None) + module._is_hf_initialized = True # nothing to initialize: the banks live in the checkpoint + module.config._experts_implementation = "ssd_streaming" + module.ssd_store = store + module.ssd_prefix = prefix + module.ssd_token_chunk_size = options.get("token_chunk_size", 256) + module.ssd_predict_max_routes = options.get("predict_max_routes", 64) + module.ssd_next = None + # On MPS, commit the queued work after each sub-block of a decoder layer and after the experts' kernels, so the + # GPU runs it while the host encodes what follows (see `_commit_mps`). + store.commit = torch.device(device).type == "mps" and options.get("mps_commit", True) + if store.commit: + for name, _ in experts: + layer = model.get_submodule(name.rpartition(".")[0].rpartition(".")[0]) + for child in layer.children(): + child.register_forward_hook(_commit_mps) + if options.get("overlap_shared_expert", True): + for name, _ in experts: + block = model.get_submodule(name.rpartition(".")[0]) + if type(block).__name__.endswith("SparseMoeBlock") and all( + hasattr(block, attr) for attr in ("gate", "shared_expert", "shared_expert_gate") + ): + _overlap_shared_expert(block) + if options.get("prefetch_next_layer", True): + routers = [] + for name, module in experts: + router = getattr(model.get_submodule(name.rpartition(".")[0]), "gate", None) + usable = router is not None and hasattr(router, "weight") and hasattr(router, "top_k") + routers.append((module.ssd_prefix, router) if usable else None) + depth = options.get("prefetch_depth", 2) + for index, (_, module) in enumerate(experts): + module.ssd_next = routers[index + 1] if index + 1 < len(routers) else None + module.ssd_after_next = routers[index + 2] if depth >= 2 and index + 2 < len(routers) else None + for name, directory in (options.get("disk_embeddings") or {}).items(): + parent, _, child = name.rpartition(".") + embedding = DiskEmbedding(directory, device, dtype) + expected = getattr(model.get_submodule(name), "weight", None) + if expected is not None and tuple(expected.shape) != (embedding.rows, embedding.dim): + raise ValueError(f"{name} is {tuple(expected.shape)}; {directory} holds {(embedding.rows, embedding.dim)}") + setattr(model.get_submodule(parent), child, embedding) + consumed.update( + key for key in reader.tensors if key.startswith(name + ".") or key.endswith("." + name + ".weight") + ) + model.ssd_store = store + return consumed + + +def quantize_experts(checkpoint, output, *, bits=4, group_size=64, device="cpu"): + """Write the experts of a safetensors MoE checkpoint, quantized for `ExpertStore(quantized=output)`. + + MLX's affine format, the one `gather_qmm` reads: per projection (``gate_up_proj``, then ``down_proj``), ``weight`` + [out, in * bits / 32] uint32, ``scales`` and ``biases`` [out, in / group_size] in the checkpoint dtype. The six + tensors of one expert are stored back to back as one uint8 row of ``{prefix}.experts`` [experts, bytes], so a miss + is a single contiguous read. One file per experts module plus an index. Quantizes with the Metal kernel on MPS, else with the PyTorch reference (same layout). + Peak memory is one experts module. + """ + from safetensors.torch import save_file + + from .metal_quantization import _affine_quantize_tensor + + device = torch.device(device) + quantize = lambda weight: _affine_quantize_tensor(weight, group_size, bits) # noqa: E731 + quantizer = "pytorch-reference" + if device.type == "mps": + from .hub_kernels import get_kernel + + kernel = get_kernel(_MLX_QUANTIZATION_KERNEL[0], version=_MLX_QUANTIZATION_KERNEL[1]) + quantize = lambda weight: kernel.quantize(weight, group_size=group_size, bits=bits) # noqa: E731 + quantizer = "mlx-quantization-metal-kernels" + reader = SafetensorsReader(checkpoint) + output = Path(output) + output.mkdir(parents=True, exist_ok=True) + packed = sorted(name.removesuffix(".gate_up_proj") for name in reader.tensors if name.endswith(".gate_up_proj")) + per_expert = sorted( + name.removesuffix(".0.gate_proj.weight") for name in reader.tensors if name.endswith(".0.gate_proj.weight") + ) + weight_map = {} + for index, prefix in enumerate(packed + per_expert): + if prefix in packed: + experts = reader.shape(f"{prefix}.gate_up_proj")[0] + + def read(expert): + return reader.read(f"{prefix}.gate_up_proj", expert), reader.read(f"{prefix}.down_proj", expert) + + else: + experts = sum( + 1 for name in reader.tensors if re.fullmatch(rf"{re.escape(prefix)}\.\d+\.gate_proj\.weight", name) + ) + + def read(expert): + gate_up = torch.cat([reader.read(f"{prefix}.{expert}.{p}_proj.weight") for p in ("gate", "up")]) + return gate_up, reader.read(f"{prefix}.{expert}.down_proj.weight") + + blobs = [] + for expert in range(experts): + parts = [] + for weight in read(expert): + weight_q, scales, biases = quantize(weight.to(device)) + parts += [weight_q, scales.to(weight.dtype), biases.to(weight.dtype)] + # Expert-major: one expert is one contiguous run of bytes, read in one call. + blobs.append(torch.cat([part.cpu().contiguous().view(torch.uint8).reshape(-1) for part in parts])) + filename = f"experts-{index:05d}.safetensors" + save_file({f"{prefix}.experts": torch.stack(blobs)}, output / filename, {"format": "pt"}) + weight_map[f"{prefix}.experts"] = filename + logger.info(f"quantized {prefix}: {experts} experts") + (output / "model.safetensors.index.json").write_text(json.dumps({"metadata": {}, "weight_map": weight_map})) + info = {"format": _QUANTIZED_FORMAT, "bits": bits, "group_size": group_size, "quantizer": quantizer} + (output / "quantization.json").write_text(json.dumps(info, indent=2)) + return output diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index ed410f9606..f659cdfa85 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -187,6 +187,8 @@ class LoadStateDictConfig: weights_only: bool = True weight_mapping: list[WeightConverter | WeightRenaming] | None = None disable_mmap: bool | None = None + # Checkpoint keys another component reads on its own (e.g. experts streamed by `experts_offload`). + skip_keys: set = field(default_factory=set) @property def is_quantized(self) -> bool: @@ -4166,6 +4168,7 @@ class PreTrainedModel( use_kernels = kwargs.pop("use_kernels", False) kernel_config = kwargs.pop("kernel_config", None) key_mapping = kwargs.pop("key_mapping", None) + experts_offload = kwargs.pop("experts_offload", None) # Not used anymore -- remove them from the kwargs for name in ["mirror", "_fast_init", "low_cpu_mem_usage", "from_tf", "from_flax", "offload_state_dict"]: @@ -4394,6 +4397,22 @@ class PreTrainedModel( if gguf_file: state_dict = hf_quantizer.get_state_dict(checkpoint_files[0], model) + skip_keys = set() + if experts_offload is not None: + from .integrations.ssd_moe import prepare_experts_offload + + devices = set(device_map.values()) if isinstance(device_map, dict) else {device_map or "cpu"} + # Metal quantization only converts nn.Linear modules, never the expert banks this removes. + dense_quantizer = hf_quantizer is not None and hf_quantizer.quantization_config.quant_method == "metal" + if len(devices) != 1 or gguf_file or (hf_quantizer is not None and not dense_quantizer) or state_dict: + raise ValueError( + "experts_offload needs a single-device device_map and a safetensors checkpoint, unquantized or " + "with MetalConfig for its dense layers" + ) + skip_keys = prepare_experts_offload( + model, checkpoint_files, dtype=dtype, device=devices.pop(), options=experts_offload + ) + # Create the dtype_plan to potentially use the `keep_in_fp32` flags (this needs to be called on the already # instantiated model, as the flags can be modified by instances sometimes) dtype_plan = model._get_dtype_plan(dtype) @@ -4425,6 +4444,7 @@ class PreTrainedModel( use_safetensors=use_safetensors, download_kwargs=download_kwargs, disable_mmap=disable_mmap, + skip_keys=skip_keys, ) if distributed_checkpoint_dir is not None: loading_info = model._load_and_finalize_distributed_pretrained_model( @@ -4570,7 +4590,8 @@ class PreTrainedModel( file_pointer = safe_open(file, framework="pt", device=device, backend=backend) all_pointer.add(file_pointer) for k in file_pointer.keys(): - merged_state_dict[k] = file_pointer.get_slice(k) # don't materialize yet + if k not in load_config.skip_keys: + merged_state_dict[k] = file_pointer.get_slice(k) # don't materialize yet # Checkpoints are .bin elif checkpoint_files is not None: merged_state_dict = {} diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 85d0c864d8..079556fc6a 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -549,6 +549,12 @@ class Qwen3_5GatedDeltaNet(nn.Module): attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ): + if hidden_states.device.type == "mps" and hidden_states.shape[1] == 1: + # Cached decode on Apple GPUs: one merged input projection and two Metal kernels instead of ~34 ops. + from ...integrations.metal_gated_delta_net import gated_delta_net_step, supports + + if supports(self, hidden_states, cache_params, attention_mask): + return gated_delta_net_step(self, hidden_states, cache_params) hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) # Set up dimensions for reshapes later diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py index 8c40031178..b8fb8a8401 100644 --- a/src/transformers/models/qwen3_5/modular_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py @@ -232,6 +232,12 @@ class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet): attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ): + if hidden_states.device.type == "mps" and hidden_states.shape[1] == 1: + # Cached decode on Apple GPUs: one merged input projection and two Metal kernels instead of ~34 ops. + from ...integrations.metal_gated_delta_net import gated_delta_net_step, supports + + if supports(self, hidden_states, cache_params, attention_mask): + return gated_delta_net_step(self, hidden_states, cache_params) hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) # Set up dimensions for reshapes later diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index fd16b2300e..0ebe19664b 100644 --- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -545,6 +545,12 @@ class Qwen3_5MoeGatedDeltaNet(nn.Module): attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ): + if hidden_states.device.type == "mps" and hidden_states.shape[1] == 1: + # Cached decode on Apple GPUs: one merged input projection and two Metal kernels instead of ~34 ops. + from ...integrations.metal_gated_delta_net import gated_delta_net_step, supports + + if supports(self, hidden_states, cache_params, attention_mask): + return gated_delta_net_step(self, hidden_states, cache_params) hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) # Set up dimensions for reshapes later diff --git a/src/transformers/models/qwen4_exp/modeling_qwen4_exp.py b/src/transformers/models/qwen4_exp/modeling_qwen4_exp.py index af706ee8e9..8101e0ed4b 100644 --- a/src/transformers/models/qwen4_exp/modeling_qwen4_exp.py +++ b/src/transformers/models/qwen4_exp/modeling_qwen4_exp.py @@ -515,6 +515,12 @@ class Qwen4ExpTextGatedDeltaNet(nn.Module): attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ): + if hidden_states.device.type == "mps" and hidden_states.shape[1] == 1: + # Cached decode on Apple GPUs: one merged input projection and two Metal kernels instead of ~34 ops. + from ...integrations.metal_gated_delta_net import gated_delta_net_step, supports + + if supports(self, hidden_states, cache_params, attention_mask): + return gated_delta_net_step(self, hidden_states, cache_params) hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) # Set up dimensions for reshapes later @@ -716,6 +722,10 @@ class Qwen4ExpTextQSAIndexer(nn.Module): # Note that the mask is never None here as we only allow eager and sdpa, and we do not allow sdpa's mask skip # It's always 4D with either bool (sdpa) or float (eager) and already gives us the valid indices visible_token_indices = attention_mask if attention_mask.dtype == torch.bool else attention_mask == 0 + if attention_mask.shape[-1] <= self.token_budget: + # At most `token_budget` keys: every complete block fits in `block_topk` and the tail is always kept, so + # every visible token is selected. Skips the per-query selection (and its host synchronizations). + return visible_token_indices selected_token_indices = torch.full( (batch_size, seq_length, self.token_budget + self.compress_ratio - 1), @@ -1012,6 +1022,12 @@ class Qwen4ExpTextGatedResidual(nn.Module): raise ValueError( f"Expected {self.hc_count * self.hidden_size} hyper-connection features, got {hyper_input.shape[-1]}." ) + if hyper_input.device.type == "mps": + # Decode on Apple GPUs: the norm, activation and stream mixing as three Metal kernels instead of ~25 ops. + from ...integrations.metal_hyper_connection import gated_residual_step, supports + + if supports(self, hyper_input): + return gated_residual_step(self, hyper_input) hyper_input_normed = self.hc_norm(hyper_input) input_mix_weight = F.silu(self.input_mix_weight_down(hyper_input_normed) / self.hc_count) input_mix_weight = torch.sigmoid(self.input_mix_weight_up(input_mix_weight)) diff --git a/src/transformers/models/qwen4_exp/modular_qwen4_exp.py b/src/transformers/models/qwen4_exp/modular_qwen4_exp.py index 82916d82e0..8c5a623936 100644 --- a/src/transformers/models/qwen4_exp/modular_qwen4_exp.py +++ b/src/transformers/models/qwen4_exp/modular_qwen4_exp.py @@ -413,6 +413,10 @@ class Qwen4ExpTextQSAIndexer(nn.Module): # Note that the mask is never None here as we only allow eager and sdpa, and we do not allow sdpa's mask skip # It's always 4D with either bool (sdpa) or float (eager) and already gives us the valid indices visible_token_indices = attention_mask if attention_mask.dtype == torch.bool else attention_mask == 0 + if attention_mask.shape[-1] <= self.token_budget: + # At most `token_budget` keys: every complete block fits in `block_topk` and the tail is always kept, so + # every visible token is selected. Skips the per-query selection (and its host synchronizations). + return visible_token_indices selected_token_indices = torch.full( (batch_size, seq_length, self.token_budget + self.compress_ratio - 1), @@ -542,6 +546,12 @@ class Qwen4ExpTextGatedResidual(nn.Module): raise ValueError( f"Expected {self.hc_count * self.hidden_size} hyper-connection features, got {hyper_input.shape[-1]}." ) + if hyper_input.device.type == "mps": + # Decode on Apple GPUs: the norm, activation and stream mixing as three Metal kernels instead of ~25 ops. + from ...integrations.metal_hyper_connection import gated_residual_step, supports + + if supports(self, hyper_input): + return gated_residual_step(self, hyper_input) hyper_input_normed = self.hc_norm(hyper_input) input_mix_weight = F.silu(self.input_mix_weight_down(hyper_input_normed) / self.hc_count) input_mix_weight = torch.sigmoid(self.input_mix_weight_up(input_mix_weight)) diff --git a/src/transformers/quantizers/quantizer_metal.py b/src/transformers/quantizers/quantizer_metal.py index 99ef8a9b6d..b809928994 100644 --- a/src/transformers/quantizers/quantizer_metal.py +++ b/src/transformers/quantizers/quantizer_metal.py @@ -15,7 +15,7 @@ from typing import TYPE_CHECKING, Any from ..utils import is_kernels_available, is_torch_available, logging from .base import HfQuantizer -from .quantizers_utils import get_module_from_name +from .quantizers_utils import get_module_from_name, is_attn_implementation_requested if is_torch_available(): @@ -100,6 +100,30 @@ class MetalHfQuantizer(HfQuantizer): pre_quantized=self.pre_quantized, ) + def update_attn_implementation(self, config): + """Remember whether an attention was asked for. The default itself is set once the weights are loaded.""" + self.attn_requested = is_attn_implementation_requested(config) + return config + + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + """ggml's layer kernels, as the GGUF path uses: the quantized matmuls are only part of a decode step, + and on MPS the layers around them (gated delta net, norms, router) otherwise run as many small ops. + Attention goes to MLX's attention kernels unless the caller chose one.""" + if self.quantization_config.dequantize: + return model + from ..integrations.gguf.kernels import kernelize_ggml_layers as apply_ggml_kernels + + apply_ggml_kernels(model) + if not getattr(self, "attn_requested", False): + try: + model.set_attn_implementation("kernels-community/metal-flash-sdpa") + except Exception as error: + logger.warning( + f"Could not use the `kernels-community/metal-flash-sdpa` attention kernel ({error}); keeping " + f"`{model.config._attn_implementation}`. Pass `attn_implementation=` to choose one explicitly." + ) + return model + def is_serializable(self): return True diff --git a/src/transformers/quantizers/quantizers_utils.py b/src/transformers/quantizers/quantizers_utils.py index c721695db6..b353493822 100644 --- a/src/transformers/quantizers/quantizers_utils.py +++ b/src/transformers/quantizers/quantizers_utils.py @@ -37,6 +37,18 @@ def should_convert_module(full_name, patterns: list[str] | None = None): return not any(re.match(rf"(?:{key}){boundary}", full_name) or full_name.endswith(key) for key in patterns) +def is_attn_implementation_requested(config) -> bool: + """Whether an attention implementation was requested for `config` or any of its sub-configs: a VLM or ALM can get + one for its `text_config` only, which leaves the top-level `_attn_implementation` unset.""" + if config._attn_implementation is not None: + return True + return any( + is_attn_implementation_requested(subconfig) + for key in config.sub_configs + if (subconfig := getattr(config, key, None)) is not None + ) + + @contextmanager def on_device(device): """Align the current accelerator device with a tensor or device-like object.""" diff --git a/tests/integrations/test_ssd_moe.py b/tests/integrations/test_ssd_moe.py new file mode 100644 index 0000000000..8124e7a2d3 --- /dev/null +++ b/tests/integrations/test_ssd_moe.py @@ -0,0 +1,728 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +import torch +from safetensors.torch import load_file, save_file + +from transformers import ( + OlmoeConfig, + OlmoeForCausalLM, + Qwen3_5MoeConfig, + Qwen3_5MoeForCausalLM, + Qwen3MoeConfig, + Qwen3MoeForCausalLM, +) +from transformers.integrations.ssd_moe import load_moe_from_ssd, load_qwen3_moe_from_ssd + + +class SSDMoETest(unittest.TestCase): + def setUp(self): + torch.manual_seed(17) + self.config = Qwen3MoeConfig( + vocab_size=64, + hidden_size=16, + intermediate_size=24, + moe_intermediate_size=8, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=8, + num_experts=4, + num_experts_per_tok=2, + ) + self.config._experts_implementation = "eager" + self.config._attn_implementation = "eager" + self.reference = Qwen3MoeForCausalLM(self.config).eval() + self.directory = tempfile.TemporaryDirectory() + self.addCleanup(self.directory.cleanup) + self.path = Path(self.directory.name) + self.config.save_pretrained(self.path) + # One expert is 3 * 8 * 16 * 4 = 1536 bytes in float32. + self.expert_bytes = 1536 + + def save_checkpoint(self, unpacked=False, sharded=False): + state = dict(self.reference.state_dict()) + if unpacked: + for name in list(state): + if ".experts." not in name: + continue + weight = state.pop(name) + prefix, _, projection = name.rpartition(".") + for expert in range(weight.shape[0]): + if projection == "gate_up_proj": + gate, up = weight[expert].chunk(2, dim=0) + state[f"{prefix}.{expert}.gate_proj.weight"] = gate.contiguous() + state[f"{prefix}.{expert}.up_proj.weight"] = up.contiguous() + else: + state[f"{prefix}.{expert}.down_proj.weight"] = weight[expert].contiguous() + if sharded: + names = list(state) + shards = [names[::2], names[1::2]] + weight_map = {} + for number, keys in enumerate(shards): + filename = f"model-{number}.safetensors" + save_file({key: state[key] for key in keys}, self.path / filename) + weight_map.update(dict.fromkeys(keys, filename)) + (self.path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + else: + save_file(state, self.path / "model.safetensors") + + def test_prefill_and_cached_decode_match_for_both_checkpoint_layouts(self): + ids = torch.tensor([[1, 2, 3, 4, 5], [5, 4, 3, 2, 1]]) + for unpacked in (False, True): + with self.subTest(unpacked=unpacked): + self.save_checkpoint(unpacked=unpacked, sharded=True) + streamed = load_qwen3_moe_from_ssd(self.path, cache_bytes=self.expert_bytes, token_chunk_size=2) + self.assertFalse(any(".experts." in name for name, _ in streamed.named_parameters())) + with torch.inference_mode(): + expected = self.reference(ids, use_cache=True) + actual = streamed(ids, use_cache=True) + torch.testing.assert_close(actual.logits, expected.logits, atol=1e-6, rtol=1e-5) + for _ in range(3): + next_id = expected.logits[:, -1].argmax(-1, keepdim=True) + expected = self.reference(next_id, past_key_values=expected.past_key_values, use_cache=True) + actual = streamed(next_id, past_key_values=actual.past_key_values, use_cache=True) + torch.testing.assert_close(actual.logits, expected.logits, atol=1e-6, rtol=1e-5) + self.assertGreater(streamed.ssd_store.evictions, 0) + self.assertLessEqual(streamed.ssd_store.peak_cache_bytes, self.expert_bytes) + + def test_sparse_reads_cache_hits_eviction_and_zero_budget(self): + self.save_checkpoint() + for budget in (0, self.expert_bytes): + with self.subTest(budget=budget): + # No I/O threads: no reads ahead, so the byte counts below are exact. + model = load_qwen3_moe_from_ssd(self.path, cache_bytes=budget, io_threads=0) + module = model.model.layers[0].mlp.experts + eager = self.reference.model.layers[0].mlp.experts + hidden = torch.randn(2, 16) + weights = torch.ones(2, 1) + ids = torch.zeros(2, 1, dtype=torch.long) + with torch.inference_mode(): + torch.testing.assert_close(module(hidden, ids, weights), eager(hidden, ids, weights)) + self.assertEqual(model.ssd_store.reader.read_bytes, self.expert_bytes) + module(hidden, ids, weights) + self.assertEqual(model.ssd_store.hits, int(budget > 0)) + module(hidden, ids + 1, weights) + torch.testing.assert_close(module(hidden, ids, weights), eager(hidden, ids, weights)) + self.assertLessEqual(model.ssd_store.resident_bytes, budget) + self.assertEqual(model.ssd_store.misses, 3 if budget else 4) + + def test_generation_and_inference_guard(self): + self.save_checkpoint() + model = load_qwen3_moe_from_ssd(self.path) + ids = torch.tensor([[1, 3, 5]]) + with torch.inference_mode(): + expected = self.reference.generate(ids, max_new_tokens=3, do_sample=False, pad_token_id=0) + actual = model.generate(ids, max_new_tokens=3, do_sample=False, pad_token_id=0) + torch.testing.assert_close(actual, expected) + with self.assertRaisesRegex(RuntimeError, "requires eval"): + model(ids) + + def test_missing_expert_fails_during_loading(self): + self.save_checkpoint(unpacked=True) + weights = load_file(self.path / "model.safetensors") + del weights["model.layers.0.mlp.experts.0.gate_proj.weight"] + save_file(weights, self.path / "model.safetensors") + with self.assertRaisesRegex(ValueError, "Missing or incompatible expert"): + load_qwen3_moe_from_ssd(self.path) + + def test_reduced_precision_and_chunked_prefill(self): + ids = torch.tensor([[1, 2, 3, 4, 5]]) + self.save_checkpoint() + for dtype in (torch.float16, torch.bfloat16): + with self.subTest(dtype=dtype): + eager = Qwen3MoeForCausalLM.from_pretrained( + self.path, dtype=dtype, experts_implementation="eager", attn_implementation="eager" + ).eval() + model = load_qwen3_moe_from_ssd(self.path, dtype=dtype, cache_bytes=self.expert_bytes // 2) + with torch.inference_mode(): + expected = eager(ids).logits[:, -1] + past = None + for chunk in ids.split(2, dim=1): + actual = model(chunk, past_key_values=past, use_cache=True) + past = actual.past_key_values + torch.testing.assert_close(actual.logits[:, -1], expected) + + def test_tied_embeddings(self): + self.config.tie_word_embeddings = True + eager = Qwen3MoeForCausalLM(self.config).eval() + eager.save_pretrained(self.path) + model = load_qwen3_moe_from_ssd(self.path) + self.assertIs(model.lm_head.weight, model.model.embed_tokens.weight) + with torch.inference_mode(): + ids = torch.tensor([[1, 2, 3]]) + torch.testing.assert_close(model(ids).logits, eager(ids).logits) + + def test_truncated_checkpoint_read_raises(self): + self.save_checkpoint() + model = load_qwen3_moe_from_ssd(self.path) + with (self.path / "model.safetensors").open("r+b") as handle: + handle.truncate(0) + with torch.inference_mode(), self.assertRaisesRegex(OSError, "Short checkpoint read"): + model(torch.tensor([[1, 2, 3]])) + + +class SSDQwen3_5MoeTest(unittest.TestCase): + """Qwen3.5/3.6 MoE: hybrid linear attention, shared expert, multimodal checkpoint naming.""" + + def test_multimodal_checkpoint_streams_text_model(self): + torch.manual_seed(17) + text_config = { + "vocab_size": 64, + "hidden_size": 16, + "num_hidden_layers": 4, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + "linear_key_head_dim": 8, + "linear_value_head_dim": 8, + "linear_num_key_heads": 2, + "linear_num_value_heads": 2, + "moe_intermediate_size": 8, + "shared_expert_intermediate_size": 8, + "num_experts": 4, + "num_experts_per_tok": 2, + } + config = Qwen3_5MoeConfig(text_config=text_config, vision_config={"depth": 1}) + reference_config = config.text_config + reference_config._experts_implementation = "eager" + reference_config._attn_implementation = "eager" + reference = Qwen3_5MoeForCausalLM(reference_config).eval() + self.assertIn("linear_attention", reference_config.layer_types) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) + config.save_pretrained(path) + # Same naming as the published Qwen3.6 checkpoints, including tensors the text model must skip. + state = { + (f"model.language_model.{name[6:]}" if name.startswith("model.") else name): value + for name, value in reference.state_dict().items() + } + state["model.visual.patch_embed.proj.weight"] = torch.zeros(2, 2) + state["mtp.fc.weight"] = torch.zeros(2, 2) + save_file(state, path / "model.safetensors") + # One expert is 3 * 8 * 16 * 4 = 1536 bytes in float32: force evictions. + streamed = load_moe_from_ssd(path, cache_bytes=1536) + self.assertFalse(any(".experts." in name for name, _ in streamed.named_parameters())) + ids = torch.tensor([[1, 2, 3, 4, 5]]) + with torch.inference_mode(): + expected = reference.generate(ids, max_new_tokens=4, do_sample=False, pad_token_id=0) + actual = streamed.generate(ids, max_new_tokens=4, do_sample=False, pad_token_id=0) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(streamed(ids).logits, reference(ids).logits, atol=1e-6, rtol=1e-5) + self.assertGreater(streamed.ssd_store.evictions, 0) + + +def quantize_q8_0(weight): + """A (rows, columns) float tensor -> (raw Q8_0 bytes, the weight those bytes dequantize to).""" + blocks = weight.float().reshape(weight.shape[0], -1, 32) + scales = (blocks.abs().amax(-1, keepdim=True) / 127).to(torch.float16) + values = torch.where(scales == 0, 0, torch.round(blocks / scales.float())).clamp(-127, 127).to(torch.int8) + raw = torch.cat((scales.view(torch.uint8), values.view(torch.uint8)), dim=-1) + return raw.reshape(-1), (values.float() * scales.float()).reshape(weight.shape) + + +def write_gguf(path, tensors): + """Minimal GGUF v3 with no metadata: tensors maps a name to (ggml dims, ggml type, raw bytes).""" + import struct + + header = bytearray(b"GGUF" + struct.pack("