diff --git a/bench/serve-bench.py b/bench/serve-bench.py index 17964de7..3828dec1 100755 --- a/bench/serve-bench.py +++ b/bench/serve-bench.py @@ -285,7 +285,7 @@ def _build_tokenizer(spec): from transformers import AutoTokenizer return AutoTokenizer.from_pretrained(t) src = spec.gguf if (not t or t.lower() == "gguf") else _expand(t) - from gmlx.load.loader import load_gguf_wire_bytes + from gmlx.load.wire import load_gguf_wire_bytes from gmlx.load.tokenizer import load_tokenizer_from_gguf arrays, kquant_meta, arch_meta, meta, _shapes = load_gguf_wire_bytes(src, zero_copy=True) del arrays, kquant_meta # release mmap views; only meta+arch needed diff --git a/gmlx/commands/cli.py b/gmlx/commands/cli.py index 69afe8e6..41c35357 100644 --- a/gmlx/commands/cli.py +++ b/gmlx/commands/cli.py @@ -1155,12 +1155,8 @@ def print_family_note(args) -> None: def _report_only(args) -> int: """Load wire bytes + remap, print the inventory and the rendered prompt.""" from gmlx.load.gguf_meta import first_nonzero_int, read_int - from gmlx.load.loader import ( - _resolve_chat_template, - load_gguf_wire_bytes, - print_inventory, - remap_arrays, - ) + from gmlx.load.loader import _resolve_chat_template, print_inventory + from gmlx.load.wire import load_gguf_wire_bytes, remap_arrays # Codec preflight so an IQ / unsupported-codec GGUF refuses cleanly here # instead of crashing kq.load_gguf. The arch gate is *skipped* - report-only @@ -1470,7 +1466,7 @@ def _apply_placement(args, model) -> None: feeder_decode=getattr(args, "decode_feeder", None), ) if stream_cpu: - from gmlx.load.loader import configure_stream_cpu + from gmlx.stream.expert_streaming import configure_stream_cpu n, _ = configure_stream_cpu(model, gguf_path=gguf_path, **feeders) if n == 0: @@ -1479,7 +1475,7 @@ def _apply_placement(args, model) -> None: "(dense?) model on the CPU device" ) else: - from gmlx.load.loader import install_expert_streaming + from gmlx.stream.expert_streaming import install_expert_streaming n, _ = install_expert_streaming( model, gguf_path=gguf_path, @@ -1492,7 +1488,7 @@ def _apply_placement(args, model) -> None: return if getattr(args, "moe_experts", None) is not None: - from gmlx.load.loader import install_moe_experts_override + from gmlx.stream.expert_streaming import install_moe_experts_override install_moe_experts_override(model, args.moe_experts) if getattr(args, "moe_expert_mass", None) is not None: diff --git a/gmlx/gen/benchmarks.py b/gmlx/gen/benchmarks.py index b5ad7519..bdb22139 100644 --- a/gmlx/gen/benchmarks.py +++ b/gmlx/gen/benchmarks.py @@ -21,7 +21,7 @@ generate_speculative, generate_speculative_owned, ) -import gmlx.load.loader as loader +import gmlx.stream.expert_streaming as expert_streaming def _synth_prompt_ids(tokenizer, n: int) -> list[int]: @@ -262,7 +262,7 @@ def bench( """ import mlx_lm - step, defaulted = loader._resolve_prefill_step(model, prefill_step_size) + step, defaulted = expert_streaming._resolve_prefill_step(model, prefill_step_size) if defaulted: print(f"[bench] streaming model: prefill chunk size defaults to {step}") pf_kwargs = {} if step is None else {"prefill_step_size": step} @@ -422,7 +422,7 @@ def _seed_len(D: int) -> int: # Same prefill-width policy as deployed generation (explicit > streaming # 8192 > stock). The mlx-lm path takes it as a stream_generate kwarg; the # drafter A/B baseline chunks through _bench_ar_tps(prefill_chunk=...). - step, defaulted = loader._resolve_prefill_step(model, prefill_step_size) + step, defaulted = expert_streaming._resolve_prefill_step(model, prefill_step_size) if defaulted: print(f"[bench] streaming model: prefill chunk size defaults to {step}") pf_kwargs = {} if step is None else {"prefill_step_size": step} diff --git a/gmlx/gen/generation.py b/gmlx/gen/generation.py index 0354a2c7..9c627b2b 100644 --- a/gmlx/gen/generation.py +++ b/gmlx/gen/generation.py @@ -15,7 +15,7 @@ import mlx.core as mx -import gmlx.load.loader as loader +import gmlx.stream.expert_streaming as expert_streaming # Tokens per target prefill forward on the speculative path only. mlx-vlm forces @@ -514,8 +514,8 @@ def generate( if prompt_cache is not None: gen_kwargs["prompt_cache"] = prompt_cache # Module-attribute lookup so the monkeypatch seam - # gmlx.load.loader._resolve_prefill_step stays live for this path. - step, defaulted = loader._resolve_prefill_step(model, prefill_step_size) + # gmlx.stream.expert_streaming._resolve_prefill_step stays live for this path. + step, defaulted = expert_streaming._resolve_prefill_step(model, prefill_step_size) if defaulted and verbose: print( f"[prefill] streaming model: chunk size defaults to {step} " diff --git a/gmlx/load/adapter.py b/gmlx/load/adapter.py index 78930ea0..019c7a95 100644 --- a/gmlx/load/adapter.py +++ b/gmlx/load/adapter.py @@ -144,7 +144,7 @@ def load_lora_adapter(adapter_path: str, """Read a GGUF LoRA adapter from disk and build its apply plan. The adapter's a/b tensors are full-precision (F32), so the wire-byte reader returns them as plain arrays (no kquant codec).""" - from .loader import load_gguf_wire_bytes + from .wire import load_gguf_wire_bytes arrays, _kquant_meta, _arch, meta, _shapes = load_gguf_wire_bytes( adapter_path, expect_quant=False) diff --git a/gmlx/load/loader.py b/gmlx/load/loader.py index 1b117c58..cd6feed3 100644 --- a/gmlx/load/loader.py +++ b/gmlx/load/loader.py @@ -15,8 +15,6 @@ from __future__ import annotations import os -import random -import re import time import mlx.core as mx @@ -30,11 +28,7 @@ from .dtypes import activation_dtype, activation_dtype_name from gmlx.envflags import env_bool, env_choice, env_int from gmlx.upstream.attn_hd512 import install_hd512_sdpa -from gmlx.gen.prefill_decay import ( - deduct_untracked_weights, - install_prefill_decay, - note_untracked_weights, -) +from gmlx.gen.prefill_decay import install_prefill_decay, note_untracked_weights import gmlx.upstream.gpt_oss_prefill as gpt_oss_prefill # noqa: F401 (registers gpt_oss score profile) from .modules import install_fused_moe_glu, install_hyv3_shexp_fold from gmlx.upstream.occupancy_fuse import install_occupancy_fuse @@ -60,7 +54,7 @@ start_populate, wait_for as wait_for_populate, ) -from .preflight import find_split_shards, preflight +from .preflight import preflight from gmlx.upstream.dsv32_patches import ( _patch_dsv32_dense_default, _patch_dsv32_indexer_fp32, @@ -77,886 +71,11 @@ _tiled_v_patch_applied, ) from .gguf_meta import first_nonzero_int, read_int -from .native_fp import _strip_weight -from .remap import RemapDecision, parse_gguf_name -from .transforms import ( - coalesce_split_experts, - fuse_shexp_gate_up, - qk_permute_wire, - retarget, - split_fused_gate_up_kquant, -) - - -# GGUF wire-byte loading - - -def load_gguf_wire_bytes( - gguf_path: str, - zero_copy: bool = True, - shards: list[str] | None = None, - expect_quant: bool = True, -) -> tuple[dict[str, mx.array], dict[str, str], str | None, dict, dict]: - """Load GGUF tensors as raw kquant wire bytes via the C++ ``kq.load_gguf``. - - ``kq.load_gguf`` reads every supported quant codec (K-quant, legacy, IQ) - as uint8 wire - bytes with a vestigial ``.scales`` placeholder, and F32/F16/BF16/ - I8/I16/I32 tensors with their native dtype. By default (``zero_copy=True``) - each tensor is a no-copy view over gguflib's mmap; ``zero_copy=False`` - memcpy's every tensor out of the mmap in C++. It also decodes all GGUF KV - metadata, so no gguf-py GGUFReader is opened in the load path. - - Returns ``(arrays, kquant_meta, arch, meta, tensor_shapes)``: - - ``arch`` is ``general.architecture`` from the first shard's metadata, or - None if absent (caller may override). - - ``meta`` is the decoded GGUF KV dict (key -> int/float/bool/str/list). - - ``tensor_shapes`` is tensor name -> logical shape (GGUF native order). - - Handles split GGUFs by loading all shards and merging; metadata + - tensor_shapes come from the first shard. ``shards`` may be passed (e.g. from - a prior preflight pass) to skip re-discovery. - """ - if shards is None: - shards = find_split_shards(gguf_path) - arrays: dict[str, mx.array] = {} - kquant_meta: dict[str, str] = {} - meta: dict = {} - tensor_shapes: dict = {} - for i, shard in enumerate(shards): - s_arrays, s_codecs, s_meta, s_shapes = kq.load_gguf(shard, zero_copy) - arrays.update(s_arrays) - kquant_meta.update(s_codecs) - tensor_shapes.update(s_shapes) - if i == 0: - meta = s_meta - if len(shards) > 1: - loadlog.verbose_print( - f"[gguf] loaded {len(shards)} shards, {len(arrays)} total tensors" - ) - - if expect_quant and not kquant_meta: - loadlog.warn( - "WARNING: no quantized tensors found - is this actually a K-quant GGUF?" - ) - - arch = meta.get("general.architecture") - return arrays, kquant_meta, arch, meta, tensor_shapes - - -# Tensor-name remap + layout transforms - - -class _RemapDict(dict): - """Weight sink that refuses silent clobbers: two GGUF tensors remapping to - the same target name is a table bug, never a legitimate overwrite.""" - - def __setitem__(self, key, value): - if key in self: - raise ValueError( - f"tensor remap collision: two source tensors map to {key!r}") - dict.__setitem__(self, key, value) - - -def _own(arr: mx.array) -> mx.array: - """Return an owned copy of ``arr`` decoupled from the source GGUF mapping. - - With zero-copy loading, native (non-quantized) tensors are views over a - file-backed shared mapping. An in-place elementwise transform on such a view - can be fused by the array library's buffer-donation optimization into a - write *through* the mapping, mutating the file on disk. Copying the data out - to host first breaks that aliasing, so the transform result is computed in a - private buffer and the source file is never touched. Used only by the small - arithmetic transforms (RMSNorm-unbake, SSM ``A``), where the cost is - negligible; bulk quantized tensors stay zero-copy. - """ - if arr.dtype == mx.bfloat16: - # numpy has no bf16 buffer format; both transform call sites compute - # in f32 anyway. astype allocates a fresh buffer, never the mapping. - arr = arr.astype(mx.float32) - return mx.array(np.array(arr)) - - -def remap_arrays( - arrays: dict[str, mx.array], - kquant_meta: dict[str, str], - arch: str, - *, - no_remap: bool = False, - target_prefix: str = "", - fail_on_unknown: bool = False, - n_head: int | None = None, - n_head_kv: int | None = None, - owned_names: set[str] | None = None, -) -> tuple[dict[str, mx.array], dict[str, str], dict[str, int]]: - """Apply name remap + layout transforms to GGUF arrays. - - Returns ``(hf_weights, hf_kquant_meta, stats)`` where ``hf_kquant_meta`` - maps the post-remap tensor name to its codec string. - - ``n_head`` / ``n_head_kv`` are required when any tensor needs the LLAMA Q/K - permute applied. When omitted, the qk_permute transform falls back to a - pass-through with a warning (the resulting model mis-attends). - - ``owned_names``, when given, collects the post-remap names of arithmetic - transform results (qk_permute, SSM A, gemma norm-unbake): arrays that must - own their buffers, never alias the source mapping (donation tripwire; see - ``_verify_zero_copy_views``). Shape-op transforms legitimately alias and - are not collected. - """ - hf_weights: dict[str, mx.array] = _RemapDict() - hf_kquant_meta: dict[str, str] = {} - stats = { - "mapped": 0, - "skipped": 0, - "split": 0, - "failed": 0, - "passthrough": 0, - "qk_permute_applied": 0, - "qk_permute_skipped": 0, - "conv1d_unsqueeze": 0, - "kda_conv_weight": 0, - "gemma_norm_minus_one": 0, - } - - # We process weight tensors; .scales sibling placeholders produced by the - # wire-byte loader get re-emitted alongside their weight under the HF name. - for name, arr in arrays.items(): - if name.endswith(".scales") or name.endswith(".biases"): - continue - codec = kquant_meta.get(name) - - if no_remap: - hf_name = name - transform = "passthrough" - else: - dec = parse_gguf_name(arch, name) - if dec.kind == RemapDecision.KIND_SKIP: - stats["skipped"] += 1 - continue - if dec.kind == RemapDecision.KIND_FAIL: - if fail_on_unknown: - raise RuntimeError(f"unmapped tensor {name!r}: {dec.reason}") - loadlog.warn( - f"WARNING: skipping unmapped tensor {name!r}: {dec.reason}" - ) - stats["failed"] += 1 - continue - hf_name = retarget(dec.hf_name, target_prefix) - transform = dec.transform - - if transform == "passthrough": - hf_weights[hf_name] = arr - if codec is not None: - hf_weights[_strip_weight(hf_name) + ".scales"] = arrays[ - _strip_weight(name) + ".scales" - ] - hf_kquant_meta[hf_name] = codec - stats["passthrough"] += 1 - stats["mapped"] += 1 - - elif transform == "moe_split_gate_up": - base = hf_name[: -len("gate_up_proj.weight")].rstrip(".") - gate_name = f"{base}.gate_proj.weight" - up_name = f"{base}.up_proj.weight" - gate, up = split_fused_gate_up_kquant(arr) - hf_weights[gate_name] = gate - hf_weights[up_name] = up - if codec is not None: - # Both halves get a vestigial scales entry under their own name. - hf_weights[_strip_weight(gate_name) + ".scales"] = mx.zeros( - (1,), dtype=mx.uint8 - ) - hf_weights[_strip_weight(up_name) + ".scales"] = mx.zeros( - (1,), dtype=mx.uint8 - ) - hf_kquant_meta[gate_name] = codec - hf_kquant_meta[up_name] = codec - stats["split"] += 1 - stats["mapped"] += 2 - - elif transform == "altup_split": - # gemma-3n stores the AltUp (un)projections as one stacked 3-D - # tensor; the MLX-native layout (GGUF dims reversed) is - # (altup_num_inputs-1, out, in). mlx_lm wants a list of separate - # Linears, so emit `{base}.{i}.weight` per stack slice. These are - # plain F16 tensors (not kquant), so a pure array slice suffices. - base = _strip_weight(hf_name) - for i in range(arr.shape[0]): - hf_weights[f"{base}.{i}.weight"] = arr[i] - stats["mapped"] += 1 - stats["split"] += 1 - - elif transform == "qk_permute": - # llama.cpp's convert_hf_to_gguf::LlamaModel.permute reorders Q/K - # rows so ggml's interleaved-pairs RoPE matches HF's concat-half - # RoPE. mlx-lm's llama/mistral3 attention uses the HF layout, so we - # undo the permute when loading from GGUF directly. - is_k = hf_name.endswith("k_proj.weight") - n_heads_for = n_head_kv if (is_k and n_head_kv is not None) else n_head - if n_heads_for is None: - loadlog.warn( - f"WARNING: qk_permute requested for {hf_name!r} but " - f"n_head/n_head_kv not provided; loading without " - f"permute (attention will be wrong)." - ) - hf_weights[hf_name] = arr - stats["qk_permute_skipped"] += 1 - else: - hf_weights[hf_name] = qk_permute_wire(arr, n_heads_for) - stats["qk_permute_applied"] += 1 - if owned_names is not None: - owned_names.add(hf_name) - if codec is not None: - hf_weights[_strip_weight(hf_name) + ".scales"] = arrays[ - _strip_weight(name) + ".scales" - ] - hf_kquant_meta[hf_name] = codec - stats["mapped"] += 1 - - elif transform == "conv1d_unsqueeze": - # Pure shape op; works on any dtype. Mamba conv weights are - # F32/BF16 (not kquant), so codec is None here. - hf_weights[hf_name] = arr[..., None] - stats["conv1d_unsqueeze"] += 1 - stats["mapped"] += 1 - - elif transform == "kda_conv_weight": - # Kimi-K3 KDA depthwise short conv: GGUF ships (1, d_inner, 1, - # d_conv) numpy order, or (d_inner, 1, d_conv) when quantization - # drops the trailing 1. conv_step varies fastest in both, so a - # pure reshape to (d_inner, d_conv) is exact; mlx Conv1d wants - # (out_channels=d_inner, kernel=d_conv, in/groups=1). - hf_weights[hf_name] = arr.reshape(-1, arr.shape[-1])[..., None] - stats["kda_conv_weight"] += 1 - stats["mapped"] += 1 - - elif transform == "ssm_a_to_a_log": - # GGUF stores SSM_A as -exp(A_log); invert to recover A_log. - # Squeeze extra leading dim (nemotron_h stores as [1, N]). - # _own() first: the negate/log would otherwise be donated into the - # source mapping (see _own docstring). - out = mx.log(-_own(arr).astype(mx.float32)) - hf_weights[hf_name] = out.reshape(-1) if out.ndim > 1 else out - stats["mapped"] += 1 - if owned_names is not None: - owned_names.add(hf_name) - - elif transform == "flatten": - # Reshape multi-dim tensor to 1D (e.g. nemotron_h ssm_norm stored as - # [n_groups, group_size], ssm_d stored as [1, N]). - hf_weights[hf_name] = arr.reshape(-1) - stats["mapped"] += 1 - - elif transform == "gate_1d_unsqueeze": - # Shared expert gate: GGUF stores 1D [hidden_size], but - # nn.Linear(hidden_size, 1, bias=False) has weight [1, hidden_size]. - hf_weights[hf_name] = arr.reshape(1, -1) if arr.ndim == 1 else arr - stats["mapped"] += 1 - - elif transform == "gemma_norm_minus_one": - # llama.cpp bakes +1 into gemma RMSNorm weights at conversion (the - # GGUF stores hf_weight + 1, used directly by ggml). mlx_lm's - # gemma/gemma2/gemma3 RMSNorm computes rms_norm(x, 1.0 + weight), - # i.e. it expects the *raw* HF weight - so undo the bake here. - # (gemma4_text uses its norm weight directly and is not tagged.) - # _own() first so the subtract isn't donated back into the source - # mapping (see _own docstring). - hf_weights[hf_name] = _own(arr).astype(mx.float32) - 1.0 - stats["gemma_norm_minus_one"] += 1 - stats["mapped"] += 1 - if owned_names is not None: - owned_names.add(hf_name) - - else: - raise RuntimeError(f"unknown transform {transform!r} for {name!r}") - - # Hand back a plain dict: the anti-clobber guard applies to remap - # population only. Later stages (native-fp repack, transforms) - # legitimately replace entries in place. - return dict(hf_weights), hf_kquant_meta, stats - - -def strip_nextn_trunk_overflow( - hf_weights: dict, hf_kquant_meta: dict, meta, arch: str -) -> int: - """Drop remapped weights of trailing NextN/MTP block(s) from the trunk tree. - - nemotron_h_moe GGUFs carry the MTP layer as ``blk.{block_count - 1}`` with - the same tensor names as trunk blocks, so the trunk remap emits - ``backbone.layers.{N}.*`` entries for a layer index the trunk model does not - have (llama.cpp likewise excludes nextn layers from the trunk graph; the - stock nemotron_h ``sanitize`` only strips HF-named ``mtp.*`` keys). The MTP - drafter loads that block separately. Returns the number of entries dropped. - """ - if arch != "nemotron_h_moe": - return 0 - nextn = read_int(meta, f"{arch}.nextn_predict_layers") or 0 - block_count = read_int(meta, f"{arch}.block_count") or 0 - if nextn <= 0 or block_count <= nextn: - return 0 - trunk = block_count - nextn - # backbone.*: the NEMOTRON_H_MOE override table; model.*: MTP-block - # tensors the override table does not claim (post_attention_norm) fall - # through to the canonical map's model.layers.{N}.* naming. - pat = re.compile(r"^(?:backbone|model)\.layers\.(\d+)\.") - dropped = 0 - for name in list(hf_weights): - m = pat.match(name) - if m and int(m.group(1)) >= trunk: - del hf_weights[name] - hf_kquant_meta.pop(name, None) - dropped += 1 - return dropped - - -# MTP / "nextn" drafter remap (native-head: the drafter weights live in the -# GGUF's own MTP block, i.e. block index >= num_hidden_layers) - -# The four ``nextn.*`` extras -> the mlx-vlm ``Qwen3_5MTPDraftModel`` param tree. -# The MTP block's *standard* decoder tensors (attn_*, ffn_*, the two block norms) -# reuse the canonical text remap (``parse_gguf_name``) with ``model.layers.{N}.`` -# rewritten to the drafter's ``layers.{i}.``. The embed table + LM head are not -# here - the drafter binds the target's at runtime (qwen3.5/3.6 GGUFs carry no -# ``nextn.embed_tokens`` / ``nextn.shared_head_head``). -_MTP_NEXTN_MAP = { - "eh_proj": "fc.weight", - "enorm": "pre_fc_norm_embedding.weight", - "hnorm": "pre_fc_norm_hidden.weight", - "shared_head_norm": "norm.weight", -} - - -def remap_mtp_arrays( - arrays: dict[str, mx.array], - kquant_meta: dict[str, str], - arch: str, - *, - first_mtp_block: int, - num_mtp_layers: int = 1, - n_head: int | None = None, - n_head_kv: int | None = None, -) -> tuple[dict[str, mx.array], dict[str, str], dict[str, int]]: - """Remap a GGUF's native MTP block(s) onto the drafter's ``mtp.*`` tree. - - ``first_mtp_block`` is the GGUF block index of the first MTP block (equals - the target's ``num_hidden_layers``); block ``first_mtp_block + i`` maps to the - drafter's ``layers.{i}``. Returns drafter-relative names (no ``model.`` - prefix); the caller builds the drafter and ``load_weights`` these onto it. - - Self-contained (does not touch the text-path ``remap_arrays``): it reuses - ``parse_gguf_name`` for the standard decoder tensors' name+transform decision - and the shared standalone transforms for emit. - """ - hf_weights: dict[str, mx.array] = _RemapDict() - hf_kquant_meta: dict[str, str] = {} - stats = { - "mapped": 0, - "skipped": 0, - "split": 0, - "passthrough": 0, - "qk_permute_applied": 0, - "qk_permute_skipped": 0, - "conv1d_unsqueeze": 0, - } - - def _emit(hf_name: str, transform: str, arr, codec, src_name: str) -> None: - if transform == "passthrough": - hf_weights[hf_name] = arr - if codec is not None: - hf_weights[_strip_weight(hf_name) + ".scales"] = arrays[ - _strip_weight(src_name) + ".scales" - ] - hf_kquant_meta[hf_name] = codec - stats["passthrough"] += 1 - stats["mapped"] += 1 - elif transform == "moe_split_gate_up": - base = hf_name[: -len("gate_up_proj.weight")].rstrip(".") - gate_name = f"{base}.gate_proj.weight" - up_name = f"{base}.up_proj.weight" - gate, up = split_fused_gate_up_kquant(arr) - hf_weights[gate_name] = gate - hf_weights[up_name] = up - if codec is not None: - hf_weights[_strip_weight(gate_name) + ".scales"] = mx.zeros( - (1,), dtype=mx.uint8 - ) - hf_weights[_strip_weight(up_name) + ".scales"] = mx.zeros( - (1,), dtype=mx.uint8 - ) - hf_kquant_meta[gate_name] = codec - hf_kquant_meta[up_name] = codec - stats["split"] += 1 - stats["mapped"] += 2 - elif transform == "gate_1d_unsqueeze": - hf_weights[hf_name] = arr.reshape(1, -1) if arr.ndim == 1 else arr - stats["mapped"] += 1 - elif transform == "flatten": - hf_weights[hf_name] = arr.reshape(-1) - stats["mapped"] += 1 - elif transform == "qk_permute": - is_k = hf_name.endswith("k_proj.weight") - nh = n_head_kv if (is_k and n_head_kv is not None) else n_head - if nh is None: - loadlog.warn( - f"WARNING: qk_permute requested for {hf_name!r} but " - f"n_head/n_head_kv not provided; loading without " - f"permute (attention will be wrong)." - ) - hf_weights[hf_name] = arr - stats["qk_permute_skipped"] += 1 - else: - hf_weights[hf_name] = qk_permute_wire(arr, nh) - stats["qk_permute_applied"] += 1 - if codec is not None: - hf_weights[_strip_weight(hf_name) + ".scales"] = arrays[ - _strip_weight(src_name) + ".scales" - ] - hf_kquant_meta[hf_name] = codec - stats["mapped"] += 1 - elif transform == "conv1d_unsqueeze": - hf_weights[hf_name] = arr[..., None] - stats["conv1d_unsqueeze"] += 1 - stats["mapped"] += 1 - else: - raise RuntimeError( - f"MTP remap: unsupported transform {transform!r} for {src_name!r}" - ) - - mtp_blocks = {first_mtp_block + i: i for i in range(num_mtp_layers)} - for name, arr in arrays.items(): - if name.endswith(".scales") or name.endswith(".biases"): - continue - m = re.match(r"^blk\.(\d+)\.(.+)$", name) - if not m: - continue - blk = int(m.group(1)) - if blk not in mtp_blocks: - continue - layer_i = mtp_blocks[blk] - rest = m.group(2) - codec = kquant_meta.get(name) - if rest.startswith("nextn."): - key = rest[len("nextn.") :] - base = key[: -len(".weight")] if key.endswith(".weight") else key - target = _MTP_NEXTN_MAP.get(base) - if target is None: - # e.g. nextn.embed_tokens / shared_head_head - shared from target. - stats["skipped"] += 1 - continue - _emit(target, "passthrough", arr, codec, name) - else: - dec = parse_gguf_name(arch, name) - if dec.kind != RemapDecision.KIND_MAP: - stats["skipped"] += 1 - continue - marker = f"model.layers.{blk}." - if marker not in dec.hf_name: - stats["skipped"] += 1 - continue - inner = dec.hf_name.split(marker, 1)[1] - _emit(f"layers.{layer_i}.{inner}", dec.transform, arr, codec, name) - # Hand back a plain dict: the anti-clobber guard applies to remap - # population only. Later stages (native-fp repack, transforms) - # legitimately replace entries in place. - return dict(hf_weights), hf_kquant_meta, stats - - -def remap_gemma4_assistant_arrays(arrays: dict, kquant_meta: dict): - """Remap a gemma4 assistant-drafter GGUF onto the mlx-vlm - ``Gemma4AssistantDraftModel`` param tree. - - The standard decoder / embed / norm tensors reuse the canonical gemma4 remap - (``parse_gguf_name`` already emits the exact ``model.*`` names the drafter - uses, including ``layer_output_scale -> layers.N.layer_scalar`` and - ``output_norm -> model.norm``); only the two bridge projections need renaming - and ``rope_freqs`` is dropped (it's computed, not a param). Every gemma4 - tensor maps as a passthrough (no qk-permute), so the emit is direct. - """ - hf_weights: dict[str, mx.array] = _RemapDict() - hf_kquant_meta: dict[str, str] = {} - stats = {"mapped": 0, "skipped": 0} - for name, arr in arrays.items(): - if name.endswith(".scales") or name.endswith(".biases"): - continue - base = name[: -len(".weight")] if name.endswith(".weight") else name - if base.endswith("pre_proj") or base.endswith("pre_projection"): - hf = "pre_projection.weight" - elif base.endswith("post_proj") or base.endswith("post_projection"): - hf = "post_projection.weight" - elif base.endswith("centroids"): - # ordered-embeddings sparse head (E2B/E4B); Q8_0, swapped by kquant. - hf = "masked_embedding.centroids.weight" - elif base.endswith("token_ordering"): - # I32 index vector, no .weight suffix on the param, never quantized. - hf_weights["masked_embedding.token_ordering"] = arr.astype(mx.int32) - stats["mapped"] += 1 - continue - elif base == "rope_freqs": - stats["skipped"] += 1 - continue - else: - dec = parse_gguf_name("gemma4", name) - if dec.kind != RemapDecision.KIND_MAP: - stats["skipped"] += 1 - continue - if dec.transform != "passthrough": - raise RuntimeError( - f"gemma4 assistant remap: unexpected transform " - f"{dec.transform!r} for {name!r}" - ) - hf = dec.hf_name - codec = kquant_meta.get(name) - hf_weights[hf] = arr - if codec is not None: - hf_weights[_strip_weight(hf) + ".scales"] = arrays[ - _strip_weight(name) + ".scales" - ] - hf_kquant_meta[hf] = codec - stats["mapped"] += 1 - # Hand back a plain dict: the anti-clobber guard applies to remap - # population only. Later stages (native-fp repack, transforms) - # legitimately replace entries in place. - return dict(hf_weights), hf_kquant_meta, stats - - -# MTP target wrapper + capability resolver - -# The hooks the mlx-vlm MTP engine probes on the *target*'s ``language_model``. -# Only ``rollback_speculative_cache`` is hard-required by the engine; the rest are -# pinned as version tripwires (a mlx-vlm bump that renames/drops one fails the -# hook-contract smoke loudly instead of corrupting decode). qwen3.5/3.6 expose -# the full verify_* set; gemma4 has a leaner set (it drafts via an assistant -# model + ``speculative_draft_hidden`` rather than the verify_* hooks). -_MTP_TARGET_HOOKS = ( - "rollback_speculative_cache", - "speculative_logits_from_hidden", - "speculative_argmax_from_hidden", - "speculative_verify_logits", - "speculative_verify_hidden", -) -_MTP_TARGET_HOOKS_BY_TYPE = { - "gemma4_text": ( - "rollback_speculative_cache", - "speculative_logits_from_hidden", - "speculative_draft_hidden", - ), - # DeepseekV4SpecLM (vendored mlx-lm class, not mlx-vlm): no - # speculative_verify_logits -- verify goes through verify_hidden and the - # walk computes logits/argmax from the raw 4D hidden. - "deepseek_v4": ( - "rollback_speculative_cache", - "speculative_logits_from_hidden", - "speculative_argmax_from_hidden", - "speculative_verify_hidden", - ), - # HyV3SpecLM (vendored mlx-lm class): same lean set as deepseek_v4. - "hy_v3": ( - "rollback_speculative_cache", - "speculative_logits_from_hidden", - "speculative_argmax_from_hidden", - "speculative_verify_hidden", - ), - # MuseGlimmerSpecLM (vendored mlx-lm class): same lean set as deepseek_v4. - "muse_glimmer": ( - "rollback_speculative_cache", - "speculative_logits_from_hidden", - "speculative_argmax_from_hidden", - "speculative_verify_hidden", - ), - # Qwen4ExpSpecLM (vendored mlx-lm class): same lean set as deepseek_v4. - "qwen4_exp": ( - "rollback_speculative_cache", - "speculative_logits_from_hidden", - "speculative_argmax_from_hidden", - "speculative_verify_hidden", - ), - # Glm5NextSpecLM (vendored mlx-lm class): same lean set as deepseek_v4. - "glm5_next": ( - "rollback_speculative_cache", - "speculative_logits_from_hidden", - "speculative_argmax_from_hidden", - "speculative_verify_hidden", - ), - # NemotronHSpecLM (stock mlx-lm class + hooks): same lean set. - "nemotron_h": ( - "rollback_speculative_cache", - "speculative_logits_from_hidden", - "speculative_argmax_from_hidden", - "speculative_verify_hidden", - ), -} - - -class MTPTextTarget(nn.Module): - """Expose an mlx-vlm text ``LanguageModel`` as ``.language_model``. +from .mtp_target import _build_mtp_target +from .transforms import coalesce_split_experts, fuse_shexp_gate_up +from .wire import load_gguf_wire_bytes, remap_arrays, strip_nextn_trunk_overflow - The MTP engine reaches the target through ``model.language_model`` (for the - ``speculative_*`` hooks + ``hidden_states``), and the drafter's ``bind`` - walks ``.language_model.model.embed_tokens``. This is deliberately not the - serving ``TextOnlyModel`` wrapper, whose ``.language_model`` is a logits-only - adapter with none of those hooks. - """ - - def __init__(self, language_model, config: dict): - super().__init__() - self.language_model = language_model - self.config = config - - def make_cache(self): - return self.language_model.make_cache() - - def get_input_embeddings(self, input_ids=None, pixel_values=None, **kwargs): - """Text-only embedding lookup the MTP engine calls on the top-level - model (``mlx_vlm.generate.ar.generate_step``). Mirrors the qwen3.5 VLM - ``Model``'s text-only branch - a GGUF text target has no vision tower - - returning an ``InputEmbeddingsFeatures`` whose ``inputs_embeds`` is the - token embedding. Clears ``_position_ids`` so mrope falls back to the - plain text positions.""" - from mlx_vlm.models.base import InputEmbeddingsFeatures - - self.language_model._position_ids = None - embeds = self.language_model.model.embed_tokens(input_ids) - # gemma4 scales token embeddings by sqrt(hidden) in its input_ids path - # (Gemma4Model.__call__), but the inputs_embeds path does not - so a - # target fed embeds must get them pre-scaled. qwen has no such scale. - if self.config.get("model_type") == "gemma4_text": - embeds = embeds * self.language_model.model.embed_scale - return InputEmbeddingsFeatures(inputs_embeds=embeds) - - def __call__(self, *args, **kwargs): - out = self.language_model(*args, **kwargs) - if isinstance(out, mx.array): - # mlx-lm-style SpecLMs (deepseek_v4) return raw logits on the - # plain path; mlx-vlm's AR engine expects ``outputs.logits``. - from mlx_vlm.models.base import LanguageModelOutput - - return LanguageModelOutput(logits=out) - return out - - -def _mtp_target_classes(model_type: str): - """Capability-resolver row ``(arch, speculative) -> (LanguageModel, build)``. - - ``build(config_dict) -> language_model`` encapsulates the per-arch - constructor signature: qwen3.5/3.6 take ``(TextConfig, ModelConfig)`` (the - mrope ``get_rope_index`` reads ``vision_config`` even for text, so we pass a - default one - the text forward never touches the tower); gemma4 takes a bare - ``TextConfig``. The stock mlx-lm classes gmlx builds for the plain text - capability ship none of the ``speculative_*`` hooks, which is why MTP - escalates to mlx-vlm here. Extend with new rows (deepseek_v4, ...) as drafters - land, together with ``_vlm_spec_language_model`` when the arch also has a - VLM shape (mmproj) so the two paths stay in lockstep. - """ - import importlib - - if model_type in ("qwen3_5", "qwen3_5_moe"): - sub = model_type - lang = importlib.import_module(f"mlx_vlm.models.{sub}.language") - cfg = importlib.import_module(f"mlx_vlm.models.{sub}.config") - - # Owned model-level forward by default: the model-level control - # flow (masks, left-padding walk, batched-padded prefill, capture) - # runs from gmlx code on a subclass; layers stay stock so the - # fused-kernel seams keep engaging. GMLX_QWEN_OWNED=0 reverts to - # the stock class wholesale. - language_model_cls = lang.LanguageModel - if env_bool("GMLX_QWEN_OWNED", True): - import gmlx.models.qwen35.owned as qwen35_owned - - language_model_cls = qwen35_owned.language_model_class(sub) - - def build(config): - text_config = cfg.TextConfig.from_dict(config) - model_config = cfg.ModelConfig.from_dict( - { - "model_type": sub, - "text_config": dict(config), - "vision_config": {}, - "vocab_size": config.get("vocab_size"), - } - ) - return language_model_cls(text_config, model_config) - - return language_model_cls, build - if model_type == "gemma4_text": - lang = importlib.import_module("mlx_vlm.models.gemma4.language") - cfg = importlib.import_module("mlx_vlm.models.gemma4.config") - - # Owned mask builder + attention by default: the nosync offset - # semantics and the hd512 row-route dispatch run from gmlx - # subclasses; layers stay stock so the fused-MoE swap keeps - # engaging. GMLX_GEMMA_OWNED=0 reverts to the stock class (the - # patch regime still installs and covers it). - language_model_cls = lang.LanguageModel - if env_bool("GMLX_GEMMA_OWNED", True): - import gmlx.models.gemma4.owned as gemma4_owned - - language_model_cls = gemma4_owned.OwnedGemma4LanguageModel - - def build(config): - return language_model_cls(cfg.TextConfig.from_dict(config)) - - return language_model_cls, build - if model_type == "deepseek_v4": - # Vendored mlx-lm-class target (no mlx-vlm counterpart): the SpecLM - # subclass carries the speculative_* hooks + rotating-undo arming. - import gmlx.models.deepseek_v4.mtp as deepseek_v4_mtp - from gmlx.models.deepseek_v4.model import ModelArgs, ensure_registered - - ensure_registered() - - def build(config): - return deepseek_v4_mtp.DeepseekV4SpecLM(ModelArgs.from_dict(config)) - - return deepseek_v4_mtp.DeepseekV4SpecLM, build - if model_type == "hy_v3": - import gmlx.models.hy_v3.mtp as hy_v3_mtp - import gmlx.models.hy_v3.tools as hy_v3_tools - from gmlx.models.hy_v3.model import ModelArgs, ensure_registered - - ensure_registered() - hy_v3_tools.ensure_registered() - - def build(config): - return hy_v3_mtp.HyV3SpecLM(ModelArgs.from_dict(config)) - - return hy_v3_mtp.HyV3SpecLM, build - if model_type == "qwen4_exp": - import gmlx.models.qwen4_exp.mtp as qwen4_exp_mtp - from gmlx.models.qwen4_exp.model import ModelArgs, ensure_registered - - ensure_registered() - - def build(config): - return qwen4_exp_mtp.Qwen4ExpSpecLM(ModelArgs.from_dict(config)) - - return qwen4_exp_mtp.Qwen4ExpSpecLM, build - if model_type == "muse_glimmer": - import gmlx.models.muse_glimmer.mtp as muse_glimmer_mtp - import gmlx.models.muse_glimmer.tools as muse_glimmer_tools - from gmlx.models.muse_glimmer.model import ModelArgs, ensure_registered - - ensure_registered() - muse_glimmer_tools.ensure_registered() - def build(config): - return muse_glimmer_mtp.MuseGlimmerSpecLM(ModelArgs.from_dict(config)) - - return muse_glimmer_mtp.MuseGlimmerSpecLM, build - if model_type == "glm5_next": - import gmlx.models.glm5_next.mtp as glm5_next_mtp - from gmlx.models.glm5_next.model import ModelArgs, ensure_registered - - ensure_registered() - - def build(config): - return glm5_next_mtp.Glm5NextSpecLM(ModelArgs.from_dict(config)) - - return glm5_next_mtp.Glm5NextSpecLM, build - if model_type == "nemotron_h": - import gmlx.models.nemotron_h.mtp as nemotron_h_mtp - from mlx_lm.models.nemotron_h import ModelArgs - - def build(config): - return nemotron_h_mtp.NemotronHSpecLM(ModelArgs.from_dict(config)) - - return nemotron_h_mtp.NemotronHSpecLM, build - from .arch_table import MTP_WIRED_MODEL_TYPES - - raise NotImplementedError( - f"MTP target class for model_type {model_type!r} not wired " - f"(supported: {' / '.join(sorted(MTP_WIRED_MODEL_TYPES))})" - ) - - -# VLM model_types whose spec-capability row lives under a different text -# model_type in the tables above (hook sets + target classes). -_VLM_SPEC_MODEL_TYPE_ALIASES = { - "gemma4": "gemma4_text", - "gemma4_unified": "gemma4_text", - # The Vision-Exp container's language model carries the deepseek_v4 - # text hooks (DeepseekV4SpecHooks). - "deepseek_v4_vl": "deepseek_v4", -} - - -def _spec_hook_key(model_type: str) -> str: - """The `_MTP_TARGET_HOOKS_BY_TYPE` / target-class key for a VLM - model_type.""" - return _VLM_SPEC_MODEL_TYPE_ALIASES.get(model_type, model_type) - - -def _vlm_spec_language_model(model_type: str): - """Spec-capable ``language_model`` row for a VLM model_type, or None. - - ``build(model_config) -> language_model`` constructs from the VLM's REAL - mlx-vlm ModelConfig (real vision_config, so mrope ``get_rope_index`` on - image turns stays correct - unlike ``_mtp_target_classes``'s text build, - which passes an empty one). None means the built ``.language_model`` - already carries the hooks (muse_glimmer / glm5_next / qwen4_exp wire - their mixins in their own vlm_model), or the arch has no spec support - (the per-arch hook check downstream fails loud). The env gates - (GMLX_QWEN_OWNED / GMLX_GEMMA_OWNED) are consulted at call time via - ``_mtp_target_classes``, so =0 resolves to the stock class and the swap - no-ops against the already-stock tree.""" - if model_type in ("qwen3_5", "qwen3_5_moe"): - cls, _ = _mtp_target_classes(model_type) - - def build(model_config): - return cls(model_config.text_config, model_config) - - return cls, build - if model_type in ("gemma4", "gemma4_unified"): - cls, _ = _mtp_target_classes("gemma4_text") - - def build(model_config): - return cls(model_config.text_config) - - return cls, build - return None - - -def _ensure_argmax_hook(language_model) -> None: - """Batched greedy verify walk: under greedy the engine takes the - per-position deferred walk (one CPU<->GPU sync per draft position) unless - the target exposes speculative_argmax_from_hidden, which lets it argmax - all block+1 verify positions in a single op (zero per-position syncs -> - _speculative_walk). gemma4's LanguageModel ships only - speculative_logits_from_hidden, so synthesize the argmax wrapper from it - - lossless (same tokens), just fewer syncs (~+8% decode on a small target - whose round is sync-bound). Only ever fires for the gemma4 row; every - other type's hook table already lists speculative_argmax_from_hidden.""" - if not hasattr(language_model, "speculative_argmax_from_hidden") and hasattr( - language_model, "speculative_logits_from_hidden" - ): - _lm = language_model - _lm.speculative_argmax_from_hidden = lambda hidden: mx.argmax( - _lm.speculative_logits_from_hidden(hidden), axis=-1 - ) - - -def _build_mtp_target(config_dict: dict): - """Build the MTP target as an mlx-vlm text ``LanguageModel`` (seam 1).""" - config = dict(config_dict) - config.pop("quantization", None) - config.pop("quantization_config", None) - model_type = config.get("model_type", "") - LanguageModel, build = _mtp_target_classes(model_type) - hooks = _MTP_TARGET_HOOKS_BY_TYPE.get(model_type, _MTP_TARGET_HOOKS) - missing = [h for h in hooks if not hasattr(LanguageModel, h)] - if missing: - raise RuntimeError( - f"mlx-vlm {model_type} LanguageModel missing MTP hooks {missing} " - f"- version drift; pin mlx-vlm or update the hook set" - ) - language_model = build(config) - _ensure_argmax_hook(language_model) - wrapper = MTPTextTarget(language_model, config) - loadlog.verbose_print( - f"[build] {model_type} -> {type(language_model).__name__} (MTP target wrapper)" - ) - return wrapper, config # Model construction (bypassing nn.quantize) @@ -1159,52 +278,6 @@ def __call__(self, x): loadlog.verbose_print(f"[patch] hunyuan: norm_topk_prob router rescale on {n} MoE layers") -# MoE expert CPU offload (hybrid GPU+CPU inference) -# -# On unified memory the GPU constraint is the wired limit, not a separate -# VRAM pool: Metal-resident buffers must be wired, while CPU-consumed mmap -# pages ride the page cache (evictable, can exceed RAM). For fine-grained MoE -# the routed expert stacks are ~90-95% of the bytes but each expert is read -# with probability top_k/n_experts per token, while the every-token layers -# (attention, norms, routers, shared experts, embeddings, KV cache) are read -# every token. -# Running the SwitchGLU expert containers on the CPU stream therefore keeps -# those hot layers + KV on GPU while the expert wire bytes stay file-backed in -# the page cache when the GPU is idle, and the kquant gather op executes on -# its threaded CPU path. MLX's cross-stream dependency tracking handles the -# GPU->CPU->GPU handoff inside each MoE layer (zero-copy - same pages). -# -# Residency (measured): Metal wires only what GPU work references or what -# sits in MLX's residency set - unreferenced file-backed buffers stay -# evictable page cache even under full memory pressure. The one hazard is -# mlx-lm's generation-time wired-limit bump (see -# _neutralize_wired_limit_sweep): MLX services a raised wired limit by -# sweeping every live buffer into the residency set, offloaded experts -# included. Models larger than the wired budget therefore run in streaming -# mode: the sweep is neutralized and GPU prefill routing is forced off, so -# the GPU never references (and never wires) expert bytes, and the page -# cache streams them from disk. -# -# Prefill staging: at decode each expert sees ~top_k/n_experts of one token, -# but a prefill chunk makes every expert hot with tens of rows each - a GEMM -# workload where the CPU (~1.5 TFLOP/s) is the wrong device. Calls with at -# least GMLX_STREAM_GPU_TOKENS tokens therefore run on the default (GPU) -# stream against the same zero-copy buffers - no copies, no staging; the -# driver wires the touched expert bytes for the duration of the work and -# releases them when the GPU goes idle. Threshold 0 disables GPU routing -# (pure CPU experts, the conservative choice when the model is far larger -# than RAM and prefill-wiring every expert is undesirable). -# -# Cost model (measured): offloaded decode pays a per-layer surcharge of -# genuine CPU dot compute plus per-layer stream fences and CPU-pool -# wake-from-idle (3 wakes per layer, one per gather; the wake cost grows -# when the pool sits idle between layers while the GPU runs the -# every-token layers). -# Routing every call to the GPU stream instead (GMLX_STREAM_GPU_TOKENS=1, -# no CPU hop) runs ~3.7x faster on a fits-in-RAM MoE, so in-RAM the CPU -# offload is for the over-budget regime, not the fast path. - -_CPU_OFFLOAD_CLASS_CACHE: dict = {} # Tokens-per-call at or above which an offloaded expert forward runs on the # GPU stream (prefill regime). Decode calls (1-few tokens) stay on CPU. @@ -1236,360 +309,12 @@ def _arena_split_max_tokens() -> int: return env_int("GMLX_ARENA_SPLIT_MAX_TOKENS", 256) -def _available_ram_bytes(include_inactive: bool = True) -> int | None: - """RAM this process can take without swapping anyone's anonymous memory: - free + purgeable + the file-backed page cache (macOS ``vm_stat``). A - load-time snapshot of the machine's offer - a machine already - half-occupied by other workloads offers the arena half a machine, - whatever the hardware total says. File-backed pages drop without IO - whatever queue they sit on; counting only the *inactive* queue (the old - formula) missed the tens of GB of recently-read GGUF cache still on the - active queue and made a mostly-cache machine look nearly full. A - ``vm_stat`` without the ``File-backed pages`` line falls back to - inactive + speculative. - - ``include_inactive=False`` is the stricter set (free + purgeable + - speculative only) for a caller that must not take the page cache.""" - from gmlx.serve import kernel_vm - - s = kernel_vm.snapshot() - if s is not None: - return s["free"] + s["purgeable"] + ( - s["filebacked"] if include_inactive else s["speculative"]) - import subprocess - - try: - # posix_spawn (absolute path, close_fds=False): a fork beside a - # Metal-mapped buffer copies the buffer first. - out = subprocess.run( - ["/usr/bin/vm_stat"], capture_output=True, text=True, timeout=5, - close_fds=False, - ).stdout - except Exception: - return None - m = re.search(r"page size of (\d+)", out) - if not m: - return None - keys = ["free", "purgeable"] - pages = 0 - found = False - if include_inactive: - mm = re.search(r"File-backed pages:\s+(\d+)\.", out) - if mm: - pages += int(mm.group(1)) - found = True - else: - keys += ["speculative", "inactive"] - else: - keys.append("speculative") - for key in keys: - mm = re.search(rf"Pages {key}:\s+(\d+)\.", out) - if mm: - pages += int(mm.group(1)) - found = True - return pages * int(m.group(1)) if found else None - - -def _ram_floor_bytes(ram: int | None) -> int: - """``gmlx.stream.budget.host_floor_bytes``.""" - from gmlx.stream.budget import host_floor_bytes - - return host_floor_bytes(ram) - - -def _decode_arena_bytes( - total_bytes: int, offsets, budget: int | None, room_bytes: int | None = None, - pinned_bytes: int = 0, streamable_bytes: int = 0, - cast_dead_bytes: int = 0, ring_bytes: int = 0, -) -> int: - """Arena budget for the decode feeder: what the memory ceiling leaves - after the non-expert weights, the KV room, the prefill ring and the - host floor, clamped to the RAM reclaimable right now, and capped at - the expert bytes themselves (a model whose experts fit goes fully - resident). - - The ceiling is the serve governor's (``gmlx.stream.budget``), so the - arena, the prefill ring and the KV cache share one budget: at decode - the governor's headroom is the room minus live KV, whatever the box's - working-set ratio or the quant's ring size. ``ring_bytes`` keeps the - ring's room out of the arena for good: a ring rebuilt on top of a - full wired arena, or lent out of it with a copy of every layer, is a - transient the kernel has to swap for on a box the arena has already - filled. ``room_bytes`` is the priced KV room - (``budget.kv_room_bytes``); None keeps the flat legacy reserve. - ``GMLX_DECODE_ARENA_RAM_FRAC`` caps the ceiling at a fraction of - physical RAM when set. ``GMLX_DECODE_ARENA_GB`` overrides the - ceilings but is still clamped to what is reclaimable minus the floor - - an arena wired past that starves the page cache every buffered - read path depends on (``GMLX_DECODE_ARENA_FORCE=1`` restores the - unclamped behavior). - - A second live streaming install needs no term here: mlock moves a page - out of the file-backed count and an arena is anonymous, so the - reclaimable snapshot already excludes both.""" - from gmlx.stream.budget import ceiling_bytes, legacy_room_bytes - - env = os.environ.get("GMLX_DECODE_ARENA_GB") - if env: - want = int(float(env) * (1 << 30)) - if env_bool("GMLX_DECODE_ARENA_FORCE", False): - return want - avail = _available_ram_bytes() - if avail is None: - return want - try: - ram = int(mx.device_info()["memory_size"]) - except Exception: - ram = avail - cap = max(0, avail - _ram_floor_bytes(ram)) - if want > cap: - print( - f"[stream] GMLX_DECODE_ARENA_GB={env} exceeds reclaimable" - f" RAM minus the floor; clamping the arena to" - f" {cap / (1 << 30):.1f}GB (GMLX_DECODE_ARENA_FORCE=1" - f" overrides)" - ) - return cap - return want - if budget is None: - return 0 - ceiling = int(ceiling_bytes() or budget) - ram = None - try: - ram = int(mx.device_info()["memory_size"]) - except Exception: - pass - frac = os.environ.get("GMLX_DECODE_ARENA_RAM_FRAC", "") - if frac and ram: - try: - ceiling = min(ceiling, int(float(frac) * ram)) - except (ValueError, OverflowError): - pass - expert_bytes = sum(r[2] for ranges in offsets.values() for r in ranges) - # Streamable components are page-cache citizens like the experts; - # charging them as non-expert would zero the arena. Cast tensors cost - # what their converted copy weighs, not what the wire does: the wire - # range is unpinned and never read again (gmlx.stream.pin_weights - # .cast_copies), so charging it would cancel the pin it just freed. - non_expert_bytes = max( - 0, total_bytes - expert_bytes - streamable_bytes - cast_dead_bytes) - room = int(room_bytes) if room_bytes is not None else legacy_room_bytes() - # The floor on both measures: the ceiling is a share of the Metal - # working set, and the OS side (the page cache, other processes) is - # not in it. - arena = (ceiling - non_expert_bytes - room - int(ring_bytes) - - _ram_floor_bytes(ram)) - # Second ceiling: what is reclaimable right now. The governor ceiling - # assumes an otherwise idle machine; co-resident workloads shrink the - # offer, and a wired arena sized past it would evict them to swap. The - # floor keeps a breathing margin for the system. This is a live - # post-pin snapshot: already-wired weights are out of it, so only the - # still-unwired share of the non-expert set is charged (charging all - # of it double-counted the pin and zeroed the arena on exactly the - # models that need it). - avail = _available_ram_bytes() - if avail is not None: - unpinned = max(0, non_expert_bytes - pinned_bytes) - arena = min( - arena, - avail - _ram_floor_bytes(ram or avail) - room - unpinned - - int(ring_bytes), - ) - return min(max(0, arena), expert_bytes) - - -def _prefill_ring_reason(offsets, left: int | None) -> str | None: - """Why the prefill ring must not be built, or None. The ring (two - slots of the largest layer's expert stacks, sized by the model) takes - its room under the memory ceiling before the decode arena; ``left`` - is what the ceiling leaves after the every-token weights and the KV - room. A ring larger than that would sit on top of them and take the - ceiling with it at the first prefill. An explicit GMLX_DECODE_ARENA_GB - is the user's budget and the ring is not judged against it.""" - if left is None or os.environ.get("GMLX_DECODE_ARENA_GB"): - return None - from gmlx.stream.prefill_feeder import ring_bytes - - ring = ring_bytes(offsets) - if ring <= left: - return None - return (f"ring 2 x {ring / 2e9:.1f} GB exceeds the {left / 1e9:.1f} GB " - "left under the memory ceiling after the every-token weights " - "and the KV room") - - -def _neutralize_wired_limit_sweep(): - """Pin the MLX wired limit at its default for the rest of the process. - - mlx-lm wraps generation in a context manager that raises the wired limit - to the device's max recommended working set. MLX services that by adding - every live buffer - file-backed zero-copy weight views included - to its - Metal residency set, which wires them all. For a model larger than the - wired budget that sweep exhausts wired memory within seconds of the - first GPU command (hard-panic territory). There is no per-buffer - opt-out, so streaming mode no-ops ``mx.set_wired_limit`` instead: Metal - then wires only what GPU work actually references (the every-token - layers + KV), and - expert pages stay plain evictable page cache. Covers every caller - (generate, server batch path, trainer) since all resolve the function - through ``mx.`` at call time. Idempotent. - - mlx-lm's ``wired_limit()`` context manager also prints a per-generation - large-model warning sized against the limit this function just pinned - - meaningless in streaming mode, and noisy (once per chat turn). Swap it - for a quiet context that keeps the exit synchronize (the original syncs - the generation stream before restoring the limit; callers may rely on - that barrier at generator teardown). NB: patched via importlib - - ``import mlx_lm.generate`` binds the function mlx_lm re-exports in - ``__init__``, not the submodule. - """ - if getattr(mx.set_wired_limit, "_kq_no_sweep", False): - return - # A generator that started before this call left the limit raised, and - # its exit restore is a no-op from here on. Lower it now: with it up, - # the next streaming load's walk wires the file's resident pages as it - # creates the views, and every command buffer in the process fails - # with a Metal out-of-memory error. - try: - prev = mx.set_wired_limit(0) - except Exception: - prev = 0 - if prev: - print(f"[stream] wired limit lowered from {prev / 1e9:.1f} GB to 0: " - "a raised limit wires every live buffer, zero-copy views included") - def _no_sweep(*_a, **_k): - return 0 - _no_sweep._kq_no_sweep = True - mx.set_wired_limit = _no_sweep - import contextlib - import importlib - @contextlib.contextmanager - def _quiet_wired_limit(model, streams=None): - try: - yield - finally: - if streams is not None: - for s in streams: - mx.synchronize(s) - else: - mx.synchronize() - - _quiet_wired_limit._kq_no_sweep = True - for mod_name in ("mlx_lm.generate", "mlx_lm.utils"): - try: - mod = importlib.import_module(mod_name) - except ImportError: - continue - if hasattr(mod, "wired_limit"): - mod.wired_limit = _quiet_wired_limit - - -def _install_wired_limit_warn_once(): - """Cap mlx-lm's large-model warning at one print per process. - - ``stream_generate`` enters mlx-lm's ``wired_limit()`` context on every - call - at least once per chat turn - and on entry the context prints its - near-the-wired-budget warning unconditionally, so a resident model just - over the 0.9x threshold re-warns every turn. There is no seam around the - print, so swap in a re-implementation with identical wiring behavior - (raise the limit, synchronize on exit, restore) that warns only the - first time. - - Installed at the end of every ``load_model`` (the resident path); the - streaming / CPU replacements above are stricter (they drop the sweep - entirely), so this never overwrites them - and they overwrite this when - they engage, which is always after load. Idempotent. NB: patched via - importlib - ``import mlx_lm.generate`` binds the function mlx_lm - re-exports in ``__init__``, not the submodule. - """ - import contextlib - import importlib - - from mlx.utils import tree_reduce - - state = {"warned": False} - - @contextlib.contextmanager - def _warn_once_wired_limit(model, streams=None): - if not mx.metal.is_available(): - yield - return - model_bytes = tree_reduce( - lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, - model, 0) - max_rec_size = mx.device_info()["max_recommended_working_set_size"] - if model_bytes > 0.9 * max_rec_size and not state["warned"]: - state["warned"] = True - model_mb = model_bytes // 2**20 - max_rec_mb = max_rec_size // 2**20 - print( - f"[WARNING] Generating with a model that requires {model_mb} " - f"MB which is close to the maximum recommended size of " - f"{max_rec_mb} MB. This can be slow. See the documentation " - "for possible work-arounds: " - "https://github.com/ml-explore/mlx-lm/tree/main#large-models" - ) - old_limit = mx.set_wired_limit(max_rec_size) - try: - yield - finally: - if streams is not None: - for s in streams: - mx.synchronize(s) - else: - mx.synchronize() - mx.set_wired_limit(old_limit) - - _warn_once_wired_limit._kq_warn_once = True - for mod_name in ("mlx_lm.generate", "mlx_lm.utils"): - try: - mod = importlib.import_module(mod_name) - except ImportError: - continue - fn = getattr(mod, "wired_limit", None) - if fn is None or getattr(fn, "_kq_no_sweep", False) or \ - getattr(fn, "_kq_warn_once", False): - continue - mod.wired_limit = _warn_once_wired_limit -def configure_cpu_device(): - """Run everything on the CPU device (``--stream-cpu``): mmap-streamed weights. - - Besides setting the default device this (a) keeps the graph - single-device - the fused-GDN runtime patch dispatches Metal kernels - regardless of the default device - and (b) no-ops mlx-lm's - ``wired_limit`` context: it reads - ``device_info()["max_recommended_working_set_size"]``, absent on the - CPU device, and wiring is meaningless on CPU. NB: patched via importlib - - ``import mlx_lm.generate`` binds the function mlx_lm re-exports in - ``__init__``, not the submodule. - """ - import contextlib - import importlib - - mx.set_default_device(mx.cpu) - os.environ.setdefault("GMLX_FUSED_GDN", "0") - - @contextlib.contextmanager - def _wired_noop(model, streams=None): - yield - - # No sweep at all on CPU; the marker keeps a later load_model's - # warn-once variant (_install_wired_limit_warn_once) from clobbering it. - _wired_noop._kq_no_sweep = True - for mod_name in ("mlx_lm.generate", "mlx_lm.utils"): - try: - mod = importlib.import_module(mod_name) - except ImportError: - continue - if hasattr(mod, "wired_limit"): - mod.wired_limit = _wired_noop - print("[device] cpu (mmap-streamed weights; fused-GDN Metal patch off)") def _resolve_feeder_defaults( @@ -1611,40 +336,6 @@ def _resolve_feeder_defaults( return feeder_prefill, feeder_decode -def configure_stream_cpu( - model, - gguf_path: str | None = None, - feeder_prefill: bool | None = None, - feeder_decode: bool | None = None, -): - """Whole-model CPU streaming (``--stream-cpu``): run the model on the CPU - device with the streaming-expert machinery always engaged. - - ``--stream-cpu`` is an explicit opt-in into the CPU/over-RAM path, so it - forces streaming (``force_stream=True``) regardless of model size - experts - run on the CPU stream whether or not the model fits the wired budget (a - fits-in-RAM model is then served from the page cache rather than faulting - from disk; for the faster all-GPU path on a model that fits, omit - ``--stream-cpu``). The GPU - working-set budget is still captured before switching the default device to - CPU so the over-/under-budget log line stays accurate (the CPU device would - otherwise report a budget that hides the condition). Returns - ``(n_wrapped, offloaded_bytes)``. - """ - try: - gpu_info = dict(mx.device_info()) - except Exception: - gpu_info = None - configure_cpu_device() - if gpu_info and "max_recommended_working_set_size" in gpu_info: - mx.device_info = lambda: gpu_info - return install_expert_streaming( - model, - gguf_path=gguf_path, - force_stream=True, - feeder_prefill=feeder_prefill, - feeder_decode=feeder_decode, - ) def _kq_expert_gpu_ok(module) -> bool: @@ -1748,1022 +439,8 @@ def _lookahead_default(model) -> bool: model, "model_type", None) not in _LA_DEFAULT_OFF_FAMILIES -def _install_gpu_residency(model, moe_modules, *, - skip_ids=frozenset(), - include_expert_stacks: bool = False) -> None: - """Wire every non-expert weight buffer into the Metal residency set, - so command buffers stop re-wiring the every-token weights' pages on - every use (the - per-use wiring is what an unswept streaming install pays instead of - the neutralized wire-everything sweep). - - ``skip_ids``: arrays that must NOT be inserted - streamed lookup - tables (a residency insert wires the buffer as surely as a GPU op). - ``include_expert_stacks``: table-only streaming keeps the experts - resident, so the GB-scale-stack belt is lifted and they are wired - with everything else.""" - import mlx_kquant as kq - - if not getattr(kq, "residency_insert", None): - print("[stream] gpu-resident weights unavailable " - "(mlx-kquant lacks residency ops)") - return - skip = set(skip_ids) - for mods in moe_modules.values(): - for m in mods: - for attr in ("gate_proj", "up_proj", "down_proj"): - w = getattr(getattr(m, attr, None), "weight", None) - if w is not None: - skip.add(id(w)) - inserted = [] - nbytes = 0 - for _, a in tree_flatten(model.parameters()): - if id(a) in skip: - continue - if (not include_expert_stacks - and a.ndim == 3 and a.nbytes > (1 << 30)): - continue # belt: any GB-scale stack is an expert container - if kq.residency_insert(a): - inserted.append(a) - nbytes += a.nbytes - kq.residency_commit() - model._kq_resident_arrays = inserted - n = len(inserted) - print(f"[stream] gpu-resident weights: {n} buffers " - f"({nbytes / 1e9:.1f} GB) in the Metal residency set " - "(GMLX_GPU_RESIDENT=0 disables)") - - -def install_expert_streaming( - model, - n_layers: int | None = None, - gguf_path: str | None = None, - force_stream: bool = False, - feeder_prefill: bool | None = None, - feeder_decode: bool | None = None, - stats_verbose: bool | None = None, -): - """Run routed-expert stacks (SwitchGLU) on the CPU stream. - - Wraps each ``SwitchGLU`` in the first ``n_layers`` decoder layers (all - layers when None) so its forward - the expert gather matmuls - executes - under ``mx.stream(mx.cpu)`` at decode shapes, and on the default (GPU) - stream for prefill-sized calls (see the staging note above). Per-instance - ``__class__`` swap; routers, shared experts, attention, and the KV cache - stay on the default (GPU) stream. Returns ``(n_wrapped, offloaded_bytes)``. - - ``gguf_path`` (the loaded checkpoint) enables sequential expert prefetch - for streaming-mode models - see ``gmlx.stream.prefetch``. Without it, - over-budget prefill demand-faults expert bytes at random-read bandwidth. - """ - from .modules import switch_layer_types - - _, glu_types = switch_layer_types() - - layers = getattr(model, "layers", None) - if layers is None: - layers = model.model.layers - - # Streaming mode: neutralize the generation-time residency sweep (which - # would otherwise wire the whole model - see _neutralize_wired_limit_sweep) - # and run every expert call on the CPU stream. Engaged when the model is - # over the wired budget (it must stream) or when force_stream is set: - # --stream-cpu (configure_stream_cpu) passes force_stream so the flag does - # what it says - experts on CPU regardless of model size; on a fits-in-RAM - # model the page cache then serves those bytes from RAM rather than faulting - # from disk. --stream-experts keeps the budget-keyed decision, so below the - # budget it still routes prefill-sized calls to the GPU stream - # (GMLX_STREAM_GPU_TOKENS) - the fast path in-RAM. - params = getattr(model, "parameters", None) - total_bytes = sum(a.nbytes for _, a in tree_flatten(params())) if params else 0 - try: - budget = int(0.9 * mx.device_info()["max_recommended_working_set_size"]) - except Exception: - budget = None - over_budget = budget is not None and total_bytes > budget - - # Selection ladder step 1 (docs/streaming.md): archs with a - # declared streamable lookup table (e.g. qwen4exp's 26.8 GiB PLE - # n-gram table) stream it instead of the experts when it alone brings - # the resident set under budget - table gathers touch ~1.4 KB/token - # against the experts' every-MoE-layer surcharge. The table wrap runs - # its row gather on a dedicated CPU stream so the buffer is never a - # GPU-stream input (a single GPU reference would wire all of it). - # When the post-table estimate is still over budget, v1 falls back to - # expert streaming with the table resident: streaming both at once - # (compose) needs the hot-row arena and is not shipped. - table_offloaded = 0 - if not force_stream: - from gmlx.stream.table_stream import ( - install_table_streaming, - table_bytes, - table_stream_selected, - ) - - compose = False - if table_stream_selected(model, total_bytes, budget): - post = total_bytes - table_bytes(model) - # Step 2 default: over budget even post-table streams both - # (compose). GMLX_STREAM_PLE_COMPOSE=0 keeps the table - # resident; the selection test already honors it in auto - # mode, so this only fires under GMLX_STREAM_PLE=1. - if budget is not None and post > budget: - if env_bool("GMLX_STREAM_PLE_COMPOSE", True): - compose = True - table_offloaded, table_names = ( - install_table_streaming(model)) - else: - print( - "[stream] table stays resident " - "(GMLX_STREAM_PLE_COMPOSE=0); experts stream" - ) - else: - table_offloaded, table_names = install_table_streaming(model) - if table_offloaded and compose: - loadlog.info( - f"[stream] compose: streamable table " - f"{'+'.join(table_names)} " - f"({table_offloaded / 2**30:.1f} GiB) on the CPU stream " - "AND experts streamed" - ) - key = getattr(model, "_kq_weights_key", None) - from gmlx.gen.prefill_decay import ( - note_streamed_tracked_bytes, - untracked_weight_bytes_for, - ) - tracked = max( - 0.0, total_bytes - untracked_weight_bytes_for(key)) - credit = min(float(table_offloaded), tracked) - if credit > 0: - note_streamed_tracked_bytes( - credit, key, source="table", cap=tracked) - deduct_untracked_weights(table_offloaded, key) - elif table_offloaded: - # The selection test admits the table only when the remainder - # clears the budget (or streaming is forced on a fits model), - # so experts are resident from here on. - over_budget = False - base = ("" if budget is None - else f" of {budget / 2**30:.1f} GiB budget") - loadlog.info( - f"[stream] streamable table {'+'.join(table_names)} " - f"({table_offloaded / 2**30:.1f} GiB) stays file-backed on " - "the CPU stream; experts resident (post-deduction " - f"{(total_bytes - table_offloaded) / 2**30:.1f} GiB{base})" - ) - key = getattr(model, "_kq_weights_key", None) - from gmlx.gen.prefill_decay import ( - note_streamed_tracked_bytes, - untracked_weight_bytes_for, - ) - tracked = max( - 0.0, total_bytes - untracked_weight_bytes_for(key)) - credit = min(float(table_offloaded), tracked) - if credit > 0: - note_streamed_tracked_bytes( - credit, key, source="table", cap=tracked) - deduct_untracked_weights(table_offloaded, key) - - streaming = force_stream or over_budget - prefetcher = None - cast_dead_bytes = 0 - held_wired = 0 - if streaming: - _neutralize_wired_limit_sweep() - # Reclaim the wired bytes of released streaming models (feeder and - # MoE modules reference each other, so unwiring waits for a - # collection), then charge what is still held against this weight - # pin. Wired pages are invisible to jetsam, so two pins that each - # size against the whole machine wire it solid. The arena needs no - # charge; see _decode_arena_bytes. - from gmlx.stream import installs as _installs - - freed = _installs.reclaim_dead() - held_wired = _installs.live_wired_bytes() - if freed: - loadlog.info( - f"[stream] reclaimed {freed / 1e9:.1f} GB wired from a " - "released streaming model") - if held_wired: - print( - f"[stream] another live streaming install holds " - f"{held_wired / 1e9:.1f} GB wired; this model sizes against " - "what is left (release the other model first for the full " - "budget)") - from gmlx.stream.prefetch import maybe_make_prefetcher - - prefetcher = maybe_make_prefetcher(gguf_path) - if prefetcher is not None: - object.__setattr__(model, "_kq_prefetcher", prefetcher) - # Wire the every-token weights before the decode feeder sizes its - # arena: pinned every-token pages come out of the same wired budget. - from gmlx.stream.pin_weights import cast_copies, maybe_pin_weights - from gmlx.stream.table_stream import streamable_tables_for - - # Declared streamable components never enter the pin set, streamed - # or resident: mlocking them starves the expert page cache. Cast - # tensors go too - their wire bytes have no view left to keep. - casts = cast_copies(model, gguf_path) - cast_dead_bytes = casts.dead_bytes - pin_exclude = frozenset( - t.gguf_name for t, _ in streamable_tables_for(model) - ) | casts.names - weights_pin = maybe_pin_weights( - gguf_path, exclude_names=pin_exclude, reserved_bytes=held_wired) - if weights_pin is not None: - object.__setattr__(model, "_kq_weights_pin", weights_pin) - _installs.record(model, weights_pin.pinned_bytes) - - def _wrapped_class(cls): - sub = _CPU_OFFLOAD_CLASS_CACHE.get(cls) - if sub is None: - # A fused base consumes routing scores itself (mix seam); a - # stock base (unrecognized activation, e.g. minimax-m3's - # SwiGLUOAI) takes (x, indices) only. The wrapper still - # advertises _kq_scores_sink so blocks hand scores over for - # miss-shed; it strips them before forwarding and applies - # the shed mix python-side. - _fwd_scores = bool(getattr(cls, "_kq_mix_scores", False)) - - class _CPUOffload(cls): - _kq_scores_sink = True - - def __call__(self, x, indices, *args, **kwargs): - # Extra args pass through untouched (e.g. deepseek-v4 - # hands the fused SwitchGLU its routing scores). A base - # without the mix seam takes (x, indices) only: keep the - # scores for the miss-shed hook and strip them from what - # gets forwarded. - scores_arg = args[0] if args else None - if args and not _fwd_scores: - args = args[1:] - # Threshold read per call (cheap; once per MoE layer per - # forward) so env changes A/B without a reload. Streaming - # mode pins everything to CPU: a GPU expert call would - # wire the buffers it references, which an over-budget - # model cannot afford. - cpu_only = getattr(self, "_kq_cpu_only", False) - gpu_tokens = _stream_gpu_tokens( - getattr( - self, "_kq_gpu_tokens_default", _STREAM_GPU_TOKENS_DEFAULT - ) - ) - n_tokens = indices.size // indices.shape[-1] - pf = getattr(self, "_kq_prefetcher", None) - fdr = getattr(self, "_kq_feeder", None) - dfr = getattr(self, "_kq_decode_feeder", None) - small = n_tokens <= _arena_stage_max_tokens() - la = getattr(self, "_kq_lookahead", None) - la_pred = None - ph = _PHASE - if ph is not None: - _phase_token( - ph, getattr(self, "_kq_li", None), n_tokens) - if n_tokens != 1: - ph = None - lsp = getattr(self, "_kq_layer_shed", None) - if lsp is not None and cpu_only and n_tokens == 1: - rng = getattr(self, "_kq_shed_rng", None) - if rng is None: - # per-layer seed: reproducible shed pattern - rng = random.Random( - 0x5EED ^ (getattr(self, "_kq_li", 0) or 0)) - object.__setattr__(self, "_kq_shed_rng", rng) - if rng.random() < lsp: - # Skip the routed path entirely (gather, stage - # and this layer's eval fence). The unmixed - # zeros return makes the block mix nothing and - # still add its shared expert. - if dfr is not None: - dfr._layer_shed_n += 1 - return mx.zeros( - (*x.shape[:-1], indices.shape[-1], - x.shape[-1]), dtype=x.dtype) - gt = getattr(self, "_kq_gpu_token", None) - gt_live = ( - gt is not None - and gt._route_shed is not None - and cpu_only - and n_tokens == 1 - and dfr is not None - and dfr.covers(self._kq_li) - ) - if gt_live: - dfr.ensure_wired() - # Token tick for EVERY covered decode layer, stage - # path included: boundary detection and the - # adaptive hot-set refresh live here. - gt.on_layer_entry( - self._kq_li, - None if getattr(self, "_kq_in_split", False) - else getattr(self, "_kq_miss_shed", None)) - if ( - gt_live - and scores_arg is not None - and not dfr.wedged_at(self._kq_li) - and gt.layer_autonomous(self._kq_li) - ): - # GPU-autonomous layer (gpu-dispatch Tier 2): no - # per-layer eval. route_shed remaps ids to arena - # slots and sheds non-resident experts on the GPU; - # the graph flushes at the next stage-path layer's - # eval or the logits, and the host consumes the - # recorded misses at the token boundary - # (popularity + prestage + fresh slot tables) - see - # gpu_token.py for the fence argument. In adaptive - # mode only layers with a measured hit rate above - # GMLX_AUTO_HOT_HIT run here, so the shed cost per - # layer is near zero. - tbl = gt.table(self._kq_li) - self._kq_cpu_only = False - try: - with dfr.swapped(self._kq_li): - with mx.stream(mx.gpu): - sc_f32 = scores_arg.astype(mx.float32) - slots, mix, m_ids, m_sc = ( - gt._route_shed( - indices.astype(mx.uint32), - sc_f32, tbl)) - mix_c = mix.astype(x.dtype) - if _fwd_scores: - y = super().__call__( - x, slots, mix_c, - *args[1:], **kwargs) - else: - y = super().__call__( - x, slots, *args, **kwargs) - if y.ndim == x.ndim + 1: - y = (y * mix_c[..., None]).sum( - axis=-2) - gt.record( - self._kq_li, indices, sc_f32, - m_ids, m_sc, y) - return y - finally: - self._kq_cpu_only = True - if la is not None and cpu_only and n_tokens == 1: - # Decode only: prefill prestage would fault the - # cold arena while the ring holds the wired budget. - # Lookahead: run the NEXT MoE layer's router on this - # layer's input and evaluate it together with the - # router read below (one sync either way). The - # prediction feeds nothing downstream - it only - # records recall (probe) or drives prestage reads. - # Latent-MoE blocks (kimi-k3) hand the full-width - # router input over out of band; x here is the - # expert container's latent-width input. - x_la = getattr(self, "_kq_la_input", None) - if x_la is None: - x_la = x - if ph is not None: - t_la = time.perf_counter() - la_pred = la.on_call(x_la, indices) - ph["la"] += time.perf_counter() - t_la - else: - la_pred = la.on_call(x_la, indices) - if ( - dfr is not None - and cpu_only - and small - and dfr.covers(self._kq_li) - ): - # Decode feeder: the routed experts are served from - # this layer's wired GPU arena; misses are pread from - # the GGUF into evicted slots first. Small prefill - # chunks take this path too when their routed set - # fits - the arena persists across requests, which is - # what makes repeat short-prompt TTFT cheap. The eval - # is both the router read and the arena-overwrite - # safety fence (see decode_feeder.py). ``stage`` - # returns None when the call routes to more distinct - # experts than the arena has slots - fall through. - t0 = time.perf_counter() if ph is not None else 0.0 - if n_tokens == 1: - dfr.ensure_wired() - # Miss-shed is decode-only: a single-token leaf of an - # arena token split is prefill work, and a shedding - # leaf would return a mixed rank-3 output next to a - # clean leaf's per-expert rank-4 - the reassembly - # concatenate cannot take both. - ms = (None if getattr(self, "_kq_in_split", False) - else getattr(self, "_kq_miss_shed", None)) - sc_f32 = None - if (ms is not None and scores_arg is not None - and n_tokens == 1): - # Shed reads the scores host-side; fold them into - # the router eval so the hook adds a small D2H - # copy, not a second per-layer graph flush. - sc_f32 = scores_arg.astype(mx.float32) - mx.eval(indices, sc_f32) - else: - mx.eval(indices) - if ph is not None: - t1 = time.perf_counter() - ph["ev"] += t1 - t0 - wait0 = getattr(dfr, "_t_demand", 0.0) - ids = np.array(indices) - shed_args = None - shed_mix = None - if sc_f32 is not None: - sc = np.asarray(sc_f32).reshape(-1) - keep = dfr.shed_misses( - self._kq_li, ids.reshape(-1), sc, ms) - if keep is not None: - # Arena-path only: the overflow fallback - # below keeps the original routed set. - kept = ids.reshape(-1)[keep] - shp = ids.shape[:-1] + (kept.size,) - ids = np.ascontiguousarray(kept.reshape(shp)) - scn = sc[keep] - # survivors keep the token's full mass - scn = scn * (sc.sum() / max(scn.sum(), 1e-20)) - sc_mx = mx.array(scn.reshape(shp)).astype( - scores_arg.dtype) - if _fwd_scores: - shed_args = (sc_mx,) + args[1:] - else: - # Stock base returns per-expert outputs; - # the block's weights still cover the - # full routed set, so mix the shed - # survivors here instead. - shed_mix = sc_mx - slots = dfr.stage(self._kq_li, ids) - if ph is not None: - t2 = time.perf_counter() - w = getattr(dfr, "_t_demand", 0.0) - wait0 - ph["stage_wait"] += w - ph["stage_book"] += (t2 - t1) - w - if la_pred: - # This layer's demand misses have joined - # (stage returned); the predicted layers' - # misses now read in the background while this - # layer's gather and the next layers' every-token - # work compute - speculation never competes with - # demand traffic for the SSD. - la_keep = ( - ms if getattr( - self, "_kq_prestage_keepers", False) - else None) - for _dst, (_ids, _sc) in la_pred.items(): - if la_keep is not None: - dfr.prestage( - _dst, _ids, keep_mass=la_keep, - pred_scores=_sc) - else: - dfr.prestage(_dst, _ids) - if ph is not None: - ph["prestage"] += time.perf_counter() - t2 - if slots is not None: - # arena call: weights are wired GPU views for - # this scope, so lift the streaming CPU pin - # and let the fused kq kernels run - if shed_args is not None: - args = shed_args - self._kq_cpu_only = False - try: - t3 = (time.perf_counter() - if ph is not None else 0.0) - with dfr.swapped(self._kq_li): - with mx.stream(mx.gpu): - y = super().__call__( - x, mx.array(slots), - *args, **kwargs) - if (shed_mix is not None - and y.ndim == x.ndim + 1): - y = (y * shed_mix[..., None]).sum( - axis=-2) - if ph is not None: - ph["build"] += time.perf_counter() - t3 - return y - finally: - self._kq_cpu_only = True - if ( - dfr is not None - and cpu_only - and 1 < n_tokens <= _arena_split_max_tokens() - and not kwargs - and dfr.covers(self._kq_li) - and dfr.can_stage_smaller(self._kq_li) - ): - # The chunk routes more distinct experts than the - # arena has slots (stage refused above, or the chunk - # is over the stage-size gate and was never tried). - # Halve along the token axis and recurse: pieces - # whose routed union fits are served from the wired - # arena's read pool, so a turn-transition prefill or - # a wide verify batch never drops to the CPU - # page-cache gather. Bottoms out at n_tokens == 1, - # which always takes a non-split path. - ax = x.ndim - 2 - orig = ((scores_arg,) + args - if scores_arg is not None and not _fwd_scores - else args) - sliceable = ( - x.shape[ax] == n_tokens - and indices.ndim == x.ndim - and all( - isinstance(a, mx.array) - and a.ndim == x.ndim - and a.shape[ax] == n_tokens - for a in orig) - ) - if sliceable: - half = n_tokens // 2 - parts = [] - prev_split = getattr( - self, "_kq_in_split", False) - object.__setattr__( - self, "_kq_in_split", True) - try: - for sl in (slice(0, half), - slice(half, n_tokens)): - t = tuple( - [slice(None)] * ax + [sl]) - parts.append(self.__call__( - x[t], indices[t], - *[a[t] for a in orig])) - # The pieces share one precomputed - # routing. Thus the stage-time eval of a - # later piece's indices does not wait for - # an earlier piece's gather. Staging - # could overwrite (or resize away) arena - # slots that the unexecuted gather - # references. Execute each piece before - # the next piece stages. - mx.eval(parts[-1]) - finally: - object.__setattr__( - self, "_kq_in_split", prev_split) - return mx.concatenate(parts, axis=ax) - wedged = dfr is not None and dfr.wedged_at(self._kq_li) - if wedged and dfr.has_dead(self._kq_li): - # A wedged read poisoned part of this layer's file - # range: no fallback below (mmap gather, advisory - # prefetch, prefill staging) may touch a dead - # expert's bytes - rewrite the routing ids first. - mx.eval(indices) - indices = mx.array(dfr.redirect_dead( - self._kq_li, np.array(indices))) - if ( - fdr is not None - and not wedged - and small - and n_tokens >= _STREAM_PREFETCH_MIN_TOKENS - and fdr.covers(self._kq_li) - ): - # Router-aware partial staging: a short chunk routes - # to a fraction of the experts, so stage only those - # slices into the ring slot instead of the whole - # layer (see feeder.prefill_partial_call). - mx.eval(indices) - ids = np.unique(np.array(indices)).tolist() - with fdr.prefill_partial_call(self, self._kq_li, ids): - with mx.stream(mx.gpu): - return super().__call__( - x, indices, *args, **kwargs) - if ( - fdr is not None - and not wedged - and n_tokens >= _STREAM_PREFETCH_MIN_TOKENS - and fdr.covers(self._kq_li) - ): - # Feeder prefill: this layer's expert stacks are - # staged straight from the GGUF into GPU-visible - # ring slots and the GEMM runs on the GPU stream - # from the slot - the page cache never sees the - # bytes. The eval is the ring protocol's slot-free - # proof (previous layer's compute has finished); - # see feeder.py. Wedged layers skip this (and the - # whole-layer advisory below): both sweep the full - # expert range, poisoned bytes included. - mx.eval(x) - with fdr.prefill_call(self, self._kq_li): - with mx.stream(mx.gpu): - return super().__call__( - x, indices, *args, **kwargs) - if ( - pf is not None - and not wedged - and pf.enabled - and n_tokens >= _STREAM_PREFETCH_MIN_TOKENS - ): - # Streaming prefill: materialize the lazy graph up to - # this layer so the advisory window advances at - # execution pace. Build-time would fire every layer's - # advisory at once, and an over-RAM advisory storm - # evicts its own earlier reads. - mx.eval(x) - pf.on_layer(self._kq_li) - elif ( - pf is not None - and pf.enabled - and cpu_only - and env_bool("GMLX_DECODE_PREFETCH", True) - ): - # Streaming decode: the router's top-k is tiny and - # the gather needs it anyway - evaluate it now and - # pull the selected experts' slices into the page - # cache at queue depth (on_decode) instead of - # demand-faulting 16 KB clusters from inside the - # gemv. GMLX_DECODE_PREFETCH=0 disables. - mx.eval(indices) - pf.on_decode( - self._kq_li, - np.unique(np.array(indices)).tolist(), - ) - if gpu_tokens > 0 and n_tokens >= gpu_tokens and not cpu_only: - # Prefill regime: GEMM on the GPU stream, same - # zero-copy buffers. - return super().__call__(x, indices, *args, **kwargs) - with mx.stream(mx.cpu): - return super().__call__(x, indices, *args, **kwargs) - - _CPUOffload.__name__ = cls.__name__ + "_CPUOffload" - _CPU_OFFLOAD_CLASS_CACHE[cls] = sub = _CPUOffload - return sub - - n_wrapped = 0 - offloaded = 0 - n_cpu_only_codec = 0 - moe_modules: dict[int, list] = {} - for li, layer in enumerate(layers): - if n_layers is not None and li >= n_layers: - break - for m in layer.modules(): - if not isinstance(m, glu_types): - continue - if m.__class__ in _CPU_OFFLOAD_CLASS_CACHE.values(): - continue # already wrapped (idempotent) - gpu_ok = _kq_expert_gpu_ok(m) - if not gpu_ok: - n_cpu_only_codec += 1 - m.__class__ = _wrapped_class(m.__class__) - if streaming: - m._kq_cpu_only = True - object.__setattr__(m, "_kq_li", li) - if gpu_ok: - moe_modules.setdefault(li, []).append(m) - if prefetcher is not None: - object.__setattr__(m, "_kq_prefetcher", prefetcher) - elif gpu_ok: - # All-GPU auto-policy: in-RAM, the residency sweep wires the - # whole model regardless of where expert calls run, so the - # CPU hop has no memory benefit and a large decode cost - # (measured ~4-5x). Route every call to the GPU stream; an - # explicit GMLX_STREAM_GPU_TOKENS (e.g. 0) overrides. - object.__setattr__(m, "_kq_gpu_tokens_default", 1) - else: - m._kq_cpu_only = True - offloaded += sum(a.nbytes for _, a in tree_flatten(m.parameters())) - n_wrapped += 1 - if over_budget and offloaded: - # Streamed expert bytes are page cache, never wired, and must not - # tax headroom_bytes() or the admission gate and request preflight - # starve every request. Which side of the accounting they sit on - # depends on how the load materialized them: registered untracked - # (zero-copy walk, small tracked delta) they need deducting; but a - # load whose arrays landed allocator-tracked (untracked registered - # ~0) has them inside mx.get_active_memory instead, and headroom - # needs the add-back credit. tracked = total - untracked splits - # the two regimes; the credit is clamped to the expert share. - # Same 0.9 x working-set budget test as _warm_touch_pass. - key = getattr(model, "_kq_weights_key", None) - from gmlx.gen.prefill_decay import ( - note_streamed_tracked_bytes, - untracked_weight_bytes_for, - ) - tracked = max(0.0, total_bytes - untracked_weight_bytes_for(key)) - credit = min(float(offloaded), tracked) - if credit > 0: - note_streamed_tracked_bytes( - credit, key, source="experts", cap=tracked) - print( - f"[stream] headroom credits {credit / 1e9:.1f} GB of " - "allocator-tracked expert bytes as reclaimable page cache" - ) - deduct_untracked_weights(offloaded, key) - if n_cpu_only_codec: - print( - f"[stream] {n_cpu_only_codec} expert stacks use a CPU-only codec " - "(no Metal matmul kernels yet): feeder/arena staging and GPU " - "prefill routing off - every expert call runs on the CPU stream" - ) - # Non-expert weights + KV run on the default device: CPU for --stream-cpu - # (configure_stream_cpu sets the default to CPU before this call), GPU for - # --stream-experts. - base_dev = "CPU" if "cpu" in str(mx.default_device()).lower() else "GPU" - if streaming: - head = ( - f"model {total_bytes / 1e9:.0f} GB > ~{budget / 1e9:.0f} GB " - "wired budget" - if over_budget - else f"model {total_bytes / 1e9:.0f} GB, streaming forced" - ) - loadlog.info( - f"[stream] streaming: {head} - {n_wrapped} MoE layers' experts " - f"({offloaded / 1e9:.1f} GB) stay file-backed; rest of the model " - f"+ KV on {base_dev}" - ) - feeder_prefill, feeder_decode = _resolve_feeder_defaults( - feeder_prefill, feeder_decode - ) - feeder = None - dfeeder = None - room = arena = None - if streaming and prefetcher is not None and moe_modules: - from gmlx.stream.budget import kv_room_bytes - from gmlx.stream.table_stream import streamed_table_bytes - - pin = getattr(model, "_kq_weights_pin", None) - room = kv_room_bytes(gguf_path) - arena_kw = dict( - room_bytes=room.bytes, - pinned_bytes=getattr(pin, "pinned_bytes", 0), - streamable_bytes=streamed_table_bytes(model), - cast_dead_bytes=cast_dead_bytes) - arena = _decode_arena_bytes( - total_bytes, prefetcher.offsets, budget, **arena_kw) - ring = 0 - if ( - streaming - and prefetcher is not None - and moe_modules - and feeder_prefill - ): - from gmlx.stream.prefill_feeder import ( - maybe_make_prefill_feeder, - ring_bytes, - ) - # No working-set budget (the CPU device): the ring is not judged. - reason = _prefill_ring_reason( - prefetcher.offsets, arena if budget is not None else None) - if reason: - print(f"[stream] feeder prefill unavailable ({reason}); " - "falling back to page-cache prefetch") - else: - feeder = maybe_make_prefill_feeder( - prefetcher.offsets, moe_modules) - if feeder is not None and arena is not None: - # The ring keeps its room for the process lifetime: a later - # prefill rebuilds it there, with no lend out of the arena. - ring = ring_bytes(prefetcher.offsets) - arena = _decode_arena_bytes( - total_bytes, prefetcher.offsets, budget, ring_bytes=ring, - **arena_kw) - if feeder is not None: - n_cov = sum(feeder.covers(li) for li in moe_modules) - for li, mods in moe_modules.items(): - if feeder.covers(li): - for m in mods: - object.__setattr__(m, "_kq_feeder", feeder) - object.__setattr__(model, "_kq_feeder", feeder) - cov = ( - "" if n_cov == len(moe_modules) - else f" on {n_cov}/{len(moe_modules)} layers" - ) - loadlog.info( - "[stream] feeder prefill: expert stacks staged straight " - f"from GGUF through 2 x {feeder.slot_bytes / 1e9:.1f} GB " - f"GPU-visible ring slots{cov} (--no-prefill-feeder disables)" - ) - if ( - streaming - and prefetcher is not None - and moe_modules - and feeder_decode - ): - from gmlx.stream.decode_feeder import maybe_make_decode_feeder - - from gmlx.stream.budget import ceiling_bytes - - # The ring's room is out of the arena's budget already, so the - # two never sum past the ceiling; the lend (DecodeFeeder - # .lend_for_ring) stays as the fallback for a box whose free RAM - # is gone when a later prefill rebuilds the ring. - dfeeder = maybe_make_decode_feeder( - prefetcher.offsets, moe_modules, arena, stats_verbose) - if dfeeder is not None: - dfeeder._room_bytes = room.bytes - n_cov = sum(dfeeder.covers(li) for li in moe_modules) - for li, mods in moe_modules.items(): - if dfeeder.covers(li): - for m in mods: - object.__setattr__(m, "_kq_decode_feeder", dfeeder) - object.__setattr__(model, "_kq_decode_feeder", dfeeder) - # Committed from here, not from the first decode: the arena - # wires itself the moment this model decodes, and a second - # install that sized against the unwired window would find the - # memory gone before it ever ran. - _installs.record(model, dfeeder.nominal_bytes) - _installs.record_arena(dfeeder) - if feeder is not None: - # The first decode call frees the ring (DecodeFeeder - # .ensure_wired). A later prefill pass rebuilds it in its - # own room; the rebuild asks the arena to lend only when - # the box has lost that room. - dfeeder._release_ring = feeder.release_slots - feeder._lend_hook = dfeeder.lend_for_ring - wired = ( - "fully wired at first decode" - if dfeeder._mlock_deferred - else f"{dfeeder.locked_bytes / 1e9:.1f} GB wired" - ) - cov = ( - "" if n_cov == len(moe_modules) - else f" on {n_cov}/{len(moe_modules)} layers" - ) - loadlog.info( - f"[stream] decode feeder: {dfeeder.nominal_bytes / 1e9:.1f} GB " - f"popularity-managed expert arena ({wired}){cov} " - "(--no-decode-feeder disables, GMLX_DECODE_ARENA_GB sizes)" - ) - ceiling = ceiling_bytes() or budget - expert_bytes = sum( - r[2] for rs in prefetcher.offsets.values() for r in rs) - room_how = ( - f"{room.depth} tokens x {room.width}: kv " - f"{room.kv_bytes / 1e9:.1f} + prefill " - f"{room.transient_bytes / 1e9:.1f} + reserve " - f"{room.reserve_bytes / 1e9:.1f}" - if room.priced else "flat GMLX_DECODE_KV_RESERVE_GB") - # Always visible, like the pin line: the one line a memory - # report needs. - try: - floor = _ram_floor_bytes(int(mx.device_info()["memory_size"])) - except Exception: - floor = _ram_floor_bytes(None) - print( - f"[stream] memory budget: ceiling {ceiling / 1e9:.1f} GB = " - f"every-token {(total_bytes - expert_bytes) / 1e9:.1f} + " - f"arena {dfeeder.nominal_bytes / 1e9:.1f} + ring " - f"{ring / 1e9:.1f} + kv room {room.bytes / 1e9:.1f} " - f"({room_how}) + floor {floor / 1e9:.1f}; " - "GMLX_STREAM_KV_CTX sizes the room" - ) - rate = getattr(dfeeder, "_probe_bps", 0.0) - measured = ( - f"drive reads {rate / 1e9:.1f} GB/s" - if rate else "fast-disk recipe forced") - if dfeeder._fast_disk: - loadlog.info( - f"[stream] decode feeder: {measured} - prefetch takes the" - " bandwidth (predictions evict by popularity, the barrier" - " joins only what the call routes to, prestage reads at" - " normal disk priority; --stream-fast-disk off restores" - " the conservative recipe)" - ) - elif rate: - loadlog.info( - f"[stream] decode feeder: {measured} - demand misses have" - " the bandwidth, prefetch stays out of their way" - " (--stream-fast-disk on overrides," - " GMLX_DECODE_FAST_DISK_GBPS sets the bar)" - ) - if (streaming or table_offloaded) and env_bool("GMLX_GPU_RESIDENT", True): - tskip = frozenset() - if table_offloaded: - from gmlx.stream.table_stream import streamed_table_array_ids - - tskip = streamed_table_array_ids(model) - _install_gpu_residency( - model, moe_modules, skip_ids=tskip, - include_expert_stacks=bool(table_offloaded) and not streaming) - if streaming and dfeeder is not None: - import gmlx.stream.gpu_token as gpu_token - - if gpu_token.autonomous_enabled(): - if gpu_token.route_shed_op() is None: - print( - "[stream] gpu-autonomous: requested but the installed " - "mlx_kquant has no route_shed op; falling back to " - "per-layer staging" - ) - else: - gt = gpu_token.GpuTokenState(dfeeder) - gpu_token.register_exit_stats(gt) - for li, mods in moe_modules.items(): - if dfeeder.covers(li): - for m in mods: - object.__setattr__(m, "_kq_gpu_token", gt) - object.__setattr__(model, "_kq_gpu_token", gt) - mode_note = ( - "all covered layers syncless (shed-heavy diagnostic)" - if gpu_token.autonomous_mode() == "all" - else "adaptive: layers above GMLX_AUTO_HOT_HIT go " - "syncless, the rest keep per-layer staging" - ) - loadlog.info( - "[stream] gpu-autonomous token: route_shed remaps + " - f"sheds on GPU; {mode_note}; misses prestage at " - "token boundaries (GMLX_GPU_AUTONOMOUS=1|all)" - ) - if streaming and dfeeder is not None and env_bool( - "GMLX_GPU_KEEPWARM", True): - import gmlx.stream.keepwarm as keepwarm - - keepwarm.start() - loadlog.info( - "[stream] gpu keep-warm: background heartbeat holds GPU " - "clocks between per-layer decode bursts, parked while no " - "decode is running (lossless, costs power only during " - "decode; GMLX_GPU_KEEPWARM=0 disables)" - ) - la_probe = env_bool("GMLX_DECODE_LOOKAHEAD_PROBE", False) - # Lookahead's replica router folds into the per-layer sync; whether its - # stall savings cover that tax is a per-family measurement. On - # glm_moe_dsa (GLM-5.2, 75 layers, top-8) it measured net negative - # (~40ms/tok sync for ~18ms of stalls), so those families default off. - # An explicit GMLX_DECODE_LOOKAHEAD always wins. - la_default = _lookahead_default(model) - la_prefetch = ( - env_bool("GMLX_DECODE_LOOKAHEAD", la_default) and dfeeder is not None) - if (streaming and dfeeder is not None and not la_default - and "GMLX_DECODE_LOOKAHEAD" not in os.environ): - loadlog.info( - "[stream] lookahead prestage: off by family default (replica-" - "router sync tax measured above its stall savings; " - "GMLX_DECODE_LOOKAHEAD=1 enables)" - ) - if streaming and (la_probe or la_prefetch): - from gmlx.stream.lookahead import install_lookahead - n_la = install_lookahead( - model, layers, probe=la_probe, prefetch=la_prefetch, - stats_verbose=stats_verbose, - ) - la_depth = max(1, min(3, env_int("GMLX_DECODE_LOOKAHEAD_DEPTH", 1))) - la_what = ( - "next-layer router predictions" - if la_depth == 1 - else f"router predictions {la_depth} layers deep" - ) - if n_la and la_prefetch: - loadlog.info( - f"[stream] lookahead prestage: {la_what} pre-read arena " - f"misses on {n_la} MoE layer pairs (lossless; " - "GMLX_DECODE_LOOKAHEAD=0 disables)" - ) - if n_la and la_probe: - loadlog.info( - f"[stream] lookahead probe: recording {la_what} recall " - f"on {n_la} MoE layer pairs (lossless; table at exit)" - ) - if streaming: - # The context line printed above and the feeder lines cover the - # normal story; what remains is the fallback mechanics for - # whatever the feeders don't handle. - fallback = [] - if dfeeder is None: - fallback.append( - "decode streams expert bytes from disk through the page " - "cache (disk-bound)" - if over_budget - else "decode reads experts through the page cache on the " - "CPU stream" - ) - if feeder is None: - fallback.append( - "prefill uses sequential page-cache prefetch" - if prefetcher is not None - else "prefill demand-faults (no gguf_path)" - ) - if fallback: - loadlog.info(f"[stream] {'; '.join(fallback)}") - if not over_budget: - b = f"~{budget / 1e9:.0f} GB" if budget else "unknown" - print( - f"[stream] --stream-cpu streams experts even though the " - f"{total_bytes / 1e9:.0f} GB model fits the wired budget " - f"({b}) - omit --stream-cpu for the faster all-GPU path on " - "a model that fits" - ) - else: - gpu_tokens = _stream_gpu_tokens(1) - if gpu_tokens == 1: - staging = ( - "model fits the wired budget - decode auto-routed to the " - "GPU stream (GMLX_STREAM_GPU_TOKENS=0 forces CPU decode)" - ) - elif gpu_tokens > 0: - staging = f"prefill calls >={gpu_tokens} tokens routed to GPU" - else: - staging = "GPU prefill routing disabled" - if table_offloaded: - # Table-only mode: the experts are GPU-resident and wired (the - # streamed table made room); "file-backed" would be wrong. - loadlog.info( - f"[stream] routed experts resident on GPU across " - f"{n_wrapped} layers ({offloaded / 1e9:.1f} GB wired; " - f"{staging})" - ) - else: - loadlog.info( - f"[stream] routed experts -> CPU stream on {n_wrapped} " - f"layers ({offloaded / 1e9:.1f} GB stays file-backed; rest " - f"of the model + KV on {base_dev}; {staging})" - ) - return n_wrapped, offloaded # Default prefill chunk width on streaming-mode (over-wired-budget) models. @@ -2801,17 +478,6 @@ def moe_streaming_active(model) -> bool: ) -def _resolve_prefill_step(model, requested: int | None) -> tuple[int | None, bool]: - """Pick the prefill chunk width: an explicit request always wins; a - streaming-mode model defaults to ``_STREAMING_PREFILL_STEP``, or to its - model_type's narrower entry; everything else keeps mlx-lm's own - default. Returns ``(step_or_none, defaulted)``.""" - if requested is not None or not moe_streaming_active(model): - return requested, False - mt = getattr(model, "model_type", None) or getattr( - getattr(model, "args", None), "model_type", None) - return _STREAMING_PREFILL_STEP_BY_MODEL_TYPE.get( - mt, _STREAMING_PREFILL_STEP), True def _switch_num_experts(glu) -> int: @@ -2870,75 +536,6 @@ def model_is_moe(model) -> bool: return False -def install_moe_experts_override(model, k: int) -> int: - """Experiment, lossy: route every token to ``k`` experts instead of the - trained top-k, on MoE blocks whose experts ``install_expert_streaming`` - wrapped (and only those - the knob exists to probe how router fan-out - shapes offloaded prefill/decode traffic, not as a general sampler). - - The override rewrites the router's own top-k attribute (``top_k`` / - ``num_experts_per_tok`` on the block, and on a DeepSeek-style gate - submodule when present - named ``gate``, or ``router`` on hy_v3), so - expert selection and the arch's weight renormalization run unchanged - at the new k. Outputs differ from the trained model by design; parity - gates will fail. Returns the number of MoE blocks overridden; raises - on k < 1 or k > the expert count. - """ - if k < 1: - raise ValueError(f"MoE top-k override must be >= 1, got {k}") - layers = getattr(model, "layers", None) - if layers is None: - layers = model.model.layers - overridden = 0 - trained_k = None - for layer in layers: - for owner in layer.modules(): - glu = None - for child in owner.children().values(): - candidates = child if isinstance(child, (list, tuple)) else [child] - for c in candidates: - if type(c).__name__.endswith("_CPUOffload"): - glu = c - break - if glu is not None: - break - if glu is None: - continue - n_experts = _switch_num_experts(glu) - if n_experts and k > n_experts: - raise ValueError( - f"MoE top-k override {k} exceeds the {n_experts}-expert " - "stack on an offloaded layer" - ) - hit = False - for target in ( - owner, - getattr(owner, "gate", None), - getattr(owner, "router", None), - ): - if target is None: - continue - for attr in ("top_k", "num_experts_per_tok"): - current = getattr(target, attr, None) - if isinstance(current, int): - if trained_k is None: - trained_k = current - setattr(target, attr, k) - hit = True - if hit: - overridden += 1 - if overridden: - print( - f"[stream] MoE top-k override: {trained_k}->{k} experts/token " - f"on {overridden} offloaded MoE layers (lossy - outputs differ " - "from the trained router)" - ) - else: - print( - "[stream] MoE top-k override found no offloaded MoE block " - "with a router top-k attribute - no effect" - ) - return overridden class _FactoredRoPE(nn.Module): @@ -4078,6 +1675,8 @@ def load_model( # Resident generation re-enters mlx-lm's wired_limit() every turn, and # its near-budget warning prints on every entry; cap it at one. + from gmlx.stream.wired_limit import _install_wired_limit_warn_once + _install_wired_limit_warn_once() return model, config, tokenizer diff --git a/gmlx/load/mtp_target.py b/gmlx/load/mtp_target.py new file mode 100644 index 00000000..12dfa3ce --- /dev/null +++ b/gmlx/load/mtp_target.py @@ -0,0 +1,347 @@ +"""MTP target wrapper and capability resolver. + +The hook tables the mlx-vlm MTP engine probes on a target's +``language_model``, the per-arch target classes +(``_mtp_target_classes`` / ``_vlm_spec_language_model``), and +``MTPTextTarget``, the wrapper that exposes an mlx-vlm text +``LanguageModel`` to the MTP engine and the drafter bind walk. +""" +from __future__ import annotations + +import mlx.core as mx +import mlx.nn as nn + +from gmlx.envflags import env_bool + +from . import loadlog + + +_MTP_TARGET_HOOKS = ( + "rollback_speculative_cache", + "speculative_logits_from_hidden", + "speculative_argmax_from_hidden", + "speculative_verify_logits", + "speculative_verify_hidden", +) +_MTP_TARGET_HOOKS_BY_TYPE = { + "gemma4_text": ( + "rollback_speculative_cache", + "speculative_logits_from_hidden", + "speculative_draft_hidden", + ), + # DeepseekV4SpecLM (vendored mlx-lm class, not mlx-vlm): no + # speculative_verify_logits -- verify goes through verify_hidden and the + # walk computes logits/argmax from the raw 4D hidden. + "deepseek_v4": ( + "rollback_speculative_cache", + "speculative_logits_from_hidden", + "speculative_argmax_from_hidden", + "speculative_verify_hidden", + ), + # HyV3SpecLM (vendored mlx-lm class): same lean set as deepseek_v4. + "hy_v3": ( + "rollback_speculative_cache", + "speculative_logits_from_hidden", + "speculative_argmax_from_hidden", + "speculative_verify_hidden", + ), + # MuseGlimmerSpecLM (vendored mlx-lm class): same lean set as deepseek_v4. + "muse_glimmer": ( + "rollback_speculative_cache", + "speculative_logits_from_hidden", + "speculative_argmax_from_hidden", + "speculative_verify_hidden", + ), + # Qwen4ExpSpecLM (vendored mlx-lm class): same lean set as deepseek_v4. + "qwen4_exp": ( + "rollback_speculative_cache", + "speculative_logits_from_hidden", + "speculative_argmax_from_hidden", + "speculative_verify_hidden", + ), + # Glm5NextSpecLM (vendored mlx-lm class): same lean set as deepseek_v4. + "glm5_next": ( + "rollback_speculative_cache", + "speculative_logits_from_hidden", + "speculative_argmax_from_hidden", + "speculative_verify_hidden", + ), + # NemotronHSpecLM (stock mlx-lm class + hooks): same lean set. + "nemotron_h": ( + "rollback_speculative_cache", + "speculative_logits_from_hidden", + "speculative_argmax_from_hidden", + "speculative_verify_hidden", + ), +} + + +class MTPTextTarget(nn.Module): + """Expose an mlx-vlm text ``LanguageModel`` as ``.language_model``. + + The MTP engine reaches the target through ``model.language_model`` (for the + ``speculative_*`` hooks + ``hidden_states``), and the drafter's ``bind`` + walks ``.language_model.model.embed_tokens``. This is deliberately not the + serving ``TextOnlyModel`` wrapper, whose ``.language_model`` is a logits-only + adapter with none of those hooks. + """ + + def __init__(self, language_model, config: dict): + super().__init__() + self.language_model = language_model + self.config = config + + def make_cache(self): + return self.language_model.make_cache() + + def get_input_embeddings(self, input_ids=None, pixel_values=None, **kwargs): + """Text-only embedding lookup the MTP engine calls on the top-level + model (``mlx_vlm.generate.ar.generate_step``). Mirrors the qwen3.5 VLM + ``Model``'s text-only branch - a GGUF text target has no vision tower - + returning an ``InputEmbeddingsFeatures`` whose ``inputs_embeds`` is the + token embedding. Clears ``_position_ids`` so mrope falls back to the + plain text positions.""" + from mlx_vlm.models.base import InputEmbeddingsFeatures + + self.language_model._position_ids = None + embeds = self.language_model.model.embed_tokens(input_ids) + # gemma4 scales token embeddings by sqrt(hidden) in its input_ids path + # (Gemma4Model.__call__), but the inputs_embeds path does not - so a + # target fed embeds must get them pre-scaled. qwen has no such scale. + if self.config.get("model_type") == "gemma4_text": + embeds = embeds * self.language_model.model.embed_scale + return InputEmbeddingsFeatures(inputs_embeds=embeds) + + def __call__(self, *args, **kwargs): + out = self.language_model(*args, **kwargs) + if isinstance(out, mx.array): + # mlx-lm-style SpecLMs (deepseek_v4) return raw logits on the + # plain path; mlx-vlm's AR engine expects ``outputs.logits``. + from mlx_vlm.models.base import LanguageModelOutput + + return LanguageModelOutput(logits=out) + return out + + +def _mtp_target_classes(model_type: str): + """Capability-resolver row ``(arch, speculative) -> (LanguageModel, build)``. + + ``build(config_dict) -> language_model`` encapsulates the per-arch + constructor signature: qwen3.5/3.6 take ``(TextConfig, ModelConfig)`` (the + mrope ``get_rope_index`` reads ``vision_config`` even for text, so we pass a + default one - the text forward never touches the tower); gemma4 takes a bare + ``TextConfig``. The stock mlx-lm classes gmlx builds for the plain text + capability ship none of the ``speculative_*`` hooks, which is why MTP + escalates to mlx-vlm here. Extend with new rows (deepseek_v4, ...) as drafters + land, together with ``_vlm_spec_language_model`` when the arch also has a + VLM shape (mmproj) so the two paths stay in lockstep. + """ + import importlib + + if model_type in ("qwen3_5", "qwen3_5_moe"): + sub = model_type + lang = importlib.import_module(f"mlx_vlm.models.{sub}.language") + cfg = importlib.import_module(f"mlx_vlm.models.{sub}.config") + + # Owned model-level forward by default: the model-level control + # flow (masks, left-padding walk, batched-padded prefill, capture) + # runs from gmlx code on a subclass; layers stay stock so the + # fused-kernel seams keep engaging. GMLX_QWEN_OWNED=0 reverts to + # the stock class wholesale. + language_model_cls = lang.LanguageModel + if env_bool("GMLX_QWEN_OWNED", True): + import gmlx.models.qwen35.owned as qwen35_owned + + language_model_cls = qwen35_owned.language_model_class(sub) + + def build(config): + text_config = cfg.TextConfig.from_dict(config) + model_config = cfg.ModelConfig.from_dict( + { + "model_type": sub, + "text_config": dict(config), + "vision_config": {}, + "vocab_size": config.get("vocab_size"), + } + ) + return language_model_cls(text_config, model_config) + + return language_model_cls, build + if model_type == "gemma4_text": + lang = importlib.import_module("mlx_vlm.models.gemma4.language") + cfg = importlib.import_module("mlx_vlm.models.gemma4.config") + + # Owned mask builder + attention by default: the nosync offset + # semantics and the hd512 row-route dispatch run from gmlx + # subclasses; layers stay stock so the fused-MoE swap keeps + # engaging. GMLX_GEMMA_OWNED=0 reverts to the stock class (the + # patch regime still installs and covers it). + language_model_cls = lang.LanguageModel + if env_bool("GMLX_GEMMA_OWNED", True): + import gmlx.models.gemma4.owned as gemma4_owned + + language_model_cls = gemma4_owned.OwnedGemma4LanguageModel + + def build(config): + return language_model_cls(cfg.TextConfig.from_dict(config)) + + return language_model_cls, build + if model_type == "deepseek_v4": + # Vendored mlx-lm-class target (no mlx-vlm counterpart): the SpecLM + # subclass carries the speculative_* hooks + rotating-undo arming. + import gmlx.models.deepseek_v4.mtp as deepseek_v4_mtp + from gmlx.models.deepseek_v4.model import ModelArgs, ensure_registered + + ensure_registered() + + def build(config): + return deepseek_v4_mtp.DeepseekV4SpecLM(ModelArgs.from_dict(config)) + + return deepseek_v4_mtp.DeepseekV4SpecLM, build + if model_type == "hy_v3": + import gmlx.models.hy_v3.mtp as hy_v3_mtp + import gmlx.models.hy_v3.tools as hy_v3_tools + from gmlx.models.hy_v3.model import ModelArgs, ensure_registered + + ensure_registered() + hy_v3_tools.ensure_registered() + + def build(config): + return hy_v3_mtp.HyV3SpecLM(ModelArgs.from_dict(config)) + + return hy_v3_mtp.HyV3SpecLM, build + if model_type == "qwen4_exp": + import gmlx.models.qwen4_exp.mtp as qwen4_exp_mtp + from gmlx.models.qwen4_exp.model import ModelArgs, ensure_registered + + ensure_registered() + + def build(config): + return qwen4_exp_mtp.Qwen4ExpSpecLM(ModelArgs.from_dict(config)) + + return qwen4_exp_mtp.Qwen4ExpSpecLM, build + if model_type == "muse_glimmer": + import gmlx.models.muse_glimmer.mtp as muse_glimmer_mtp + import gmlx.models.muse_glimmer.tools as muse_glimmer_tools + from gmlx.models.muse_glimmer.model import ModelArgs, ensure_registered + + ensure_registered() + muse_glimmer_tools.ensure_registered() + + def build(config): + return muse_glimmer_mtp.MuseGlimmerSpecLM(ModelArgs.from_dict(config)) + + return muse_glimmer_mtp.MuseGlimmerSpecLM, build + if model_type == "glm5_next": + import gmlx.models.glm5_next.mtp as glm5_next_mtp + from gmlx.models.glm5_next.model import ModelArgs, ensure_registered + + ensure_registered() + + def build(config): + return glm5_next_mtp.Glm5NextSpecLM(ModelArgs.from_dict(config)) + + return glm5_next_mtp.Glm5NextSpecLM, build + if model_type == "nemotron_h": + import gmlx.models.nemotron_h.mtp as nemotron_h_mtp + from mlx_lm.models.nemotron_h import ModelArgs + + def build(config): + return nemotron_h_mtp.NemotronHSpecLM(ModelArgs.from_dict(config)) + + return nemotron_h_mtp.NemotronHSpecLM, build + from .arch_table import MTP_WIRED_MODEL_TYPES + + raise NotImplementedError( + f"MTP target class for model_type {model_type!r} not wired " + f"(supported: {' / '.join(sorted(MTP_WIRED_MODEL_TYPES))})" + ) + + +# VLM model_types whose spec-capability row lives under a different text +# model_type in the tables above (hook sets + target classes). +_VLM_SPEC_MODEL_TYPE_ALIASES = { + "gemma4": "gemma4_text", + "gemma4_unified": "gemma4_text", + # The Vision-Exp container's language model carries the deepseek_v4 + # text hooks (DeepseekV4SpecHooks). + "deepseek_v4_vl": "deepseek_v4", +} + + +def _spec_hook_key(model_type: str) -> str: + """The `_MTP_TARGET_HOOKS_BY_TYPE` / target-class key for a VLM + model_type.""" + return _VLM_SPEC_MODEL_TYPE_ALIASES.get(model_type, model_type) + + +def _vlm_spec_language_model(model_type: str): + """Spec-capable ``language_model`` row for a VLM model_type, or None. + + ``build(model_config) -> language_model`` constructs from the VLM's REAL + mlx-vlm ModelConfig (real vision_config, so mrope ``get_rope_index`` on + image turns stays correct - unlike ``_mtp_target_classes``'s text build, + which passes an empty one). None means the built ``.language_model`` + already carries the hooks (muse_glimmer / glm5_next / qwen4_exp wire + their mixins in their own vlm_model), or the arch has no spec support + (the per-arch hook check downstream fails loud). The env gates + (GMLX_QWEN_OWNED / GMLX_GEMMA_OWNED) are consulted at call time via + ``_mtp_target_classes``, so =0 resolves to the stock class and the swap + no-ops against the already-stock tree.""" + if model_type in ("qwen3_5", "qwen3_5_moe"): + cls, _ = _mtp_target_classes(model_type) + + def build(model_config): + return cls(model_config.text_config, model_config) + + return cls, build + if model_type in ("gemma4", "gemma4_unified"): + cls, _ = _mtp_target_classes("gemma4_text") + + def build(model_config): + return cls(model_config.text_config) + + return cls, build + return None + + +def _ensure_argmax_hook(language_model) -> None: + """Batched greedy verify walk: under greedy the engine takes the + per-position deferred walk (one CPU<->GPU sync per draft position) unless + the target exposes speculative_argmax_from_hidden, which lets it argmax + all block+1 verify positions in a single op (zero per-position syncs -> + _speculative_walk). gemma4's LanguageModel ships only + speculative_logits_from_hidden, so synthesize the argmax wrapper from it - + lossless (same tokens), just fewer syncs (~+8% decode on a small target + whose round is sync-bound). Only ever fires for the gemma4 row; every + other type's hook table already lists speculative_argmax_from_hidden.""" + if not hasattr(language_model, "speculative_argmax_from_hidden") and hasattr( + language_model, "speculative_logits_from_hidden" + ): + _lm = language_model + _lm.speculative_argmax_from_hidden = lambda hidden: mx.argmax( + _lm.speculative_logits_from_hidden(hidden), axis=-1 + ) + + +def _build_mtp_target(config_dict: dict): + """Build the MTP target as an mlx-vlm text ``LanguageModel`` (seam 1).""" + config = dict(config_dict) + config.pop("quantization", None) + config.pop("quantization_config", None) + model_type = config.get("model_type", "") + LanguageModel, build = _mtp_target_classes(model_type) + hooks = _MTP_TARGET_HOOKS_BY_TYPE.get(model_type, _MTP_TARGET_HOOKS) + missing = [h for h in hooks if not hasattr(LanguageModel, h)] + if missing: + raise RuntimeError( + f"mlx-vlm {model_type} LanguageModel missing MTP hooks {missing} " + f"- version drift; pin mlx-vlm or update the hook set" + ) + language_model = build(config) + _ensure_argmax_hook(language_model) + wrapper = MTPTextTarget(language_model, config) + loadlog.verbose_print( + f"[build] {model_type} -> {type(language_model).__name__} (MTP target wrapper)" + ) + return wrapper, config diff --git a/gmlx/load/vlm.py b/gmlx/load/vlm.py index 2d220429..ef05ad9b 100644 --- a/gmlx/load/vlm.py +++ b/gmlx/load/vlm.py @@ -39,12 +39,11 @@ _FP32_KEEP_BY_MODEL_TYPE, _active_now, _install_and_load, - _vlm_spec_language_model, - load_gguf_wire_bytes, materialize_module_arrays, - remap_arrays, weights_source_key, ) +from .mtp_target import _vlm_spec_language_model +from .wire import load_gguf_wire_bytes, remap_arrays from .preflight import preflight from .transforms import coalesce_split_experts diff --git a/gmlx/load/wire.py b/gmlx/load/wire.py new file mode 100644 index 00000000..e9dfb841 --- /dev/null +++ b/gmlx/load/wire.py @@ -0,0 +1,551 @@ +"""GGUF wire-byte loading and array remapping. + +Reads a GGUF's raw kquant wire bytes (``load_gguf_wire_bytes``) and remaps +GGUF tensor names onto the mlx-lm/mlx-vlm parameter tree: the text path +(``remap_arrays``), MTP extras (``remap_mtp_arrays``), and the gemma4 +assistant head (``remap_gemma4_assistant_arrays``). +""" +from __future__ import annotations + +import re + +import mlx.core as mx +import numpy as np + +import mlx_kquant as kq + +from . import loadlog +from .gguf_meta import read_int +from .native_fp import _strip_weight +from .preflight import find_split_shards +from .remap import RemapDecision, parse_gguf_name +from .transforms import qk_permute_wire, retarget, split_fused_gate_up_kquant + + +def load_gguf_wire_bytes( + gguf_path: str, + zero_copy: bool = True, + shards: list[str] | None = None, + expect_quant: bool = True, +) -> tuple[dict[str, mx.array], dict[str, str], str | None, dict, dict]: + """Load GGUF tensors as raw kquant wire bytes via the C++ ``kq.load_gguf``. + + ``kq.load_gguf`` reads every supported quant codec (K-quant, legacy, IQ) + as uint8 wire + bytes with a vestigial ``.scales`` placeholder, and F32/F16/BF16/ + I8/I16/I32 tensors with their native dtype. By default (``zero_copy=True``) + each tensor is a no-copy view over gguflib's mmap; ``zero_copy=False`` + memcpy's every tensor out of the mmap in C++. It also decodes all GGUF KV + metadata, so no gguf-py GGUFReader is opened in the load path. + + Returns ``(arrays, kquant_meta, arch, meta, tensor_shapes)``: + - ``arch`` is ``general.architecture`` from the first shard's metadata, or + None if absent (caller may override). + - ``meta`` is the decoded GGUF KV dict (key -> int/float/bool/str/list). + - ``tensor_shapes`` is tensor name -> logical shape (GGUF native order). + + Handles split GGUFs by loading all shards and merging; metadata + + tensor_shapes come from the first shard. ``shards`` may be passed (e.g. from + a prior preflight pass) to skip re-discovery. + """ + if shards is None: + shards = find_split_shards(gguf_path) + arrays: dict[str, mx.array] = {} + kquant_meta: dict[str, str] = {} + meta: dict = {} + tensor_shapes: dict = {} + for i, shard in enumerate(shards): + s_arrays, s_codecs, s_meta, s_shapes = kq.load_gguf(shard, zero_copy) + arrays.update(s_arrays) + kquant_meta.update(s_codecs) + tensor_shapes.update(s_shapes) + if i == 0: + meta = s_meta + if len(shards) > 1: + loadlog.verbose_print( + f"[gguf] loaded {len(shards)} shards, {len(arrays)} total tensors" + ) + + if expect_quant and not kquant_meta: + loadlog.warn( + "WARNING: no quantized tensors found - is this actually a K-quant GGUF?" + ) + + arch = meta.get("general.architecture") + return arrays, kquant_meta, arch, meta, tensor_shapes + + +# Tensor-name remap + layout transforms + + +class _RemapDict(dict): + """Weight sink that refuses silent clobbers: two GGUF tensors remapping to + the same target name is a table bug, never a legitimate overwrite.""" + + def __setitem__(self, key, value): + if key in self: + raise ValueError( + f"tensor remap collision: two source tensors map to {key!r}") + dict.__setitem__(self, key, value) + + +def _own(arr: mx.array) -> mx.array: + """Return an owned copy of ``arr`` decoupled from the source GGUF mapping. + + With zero-copy loading, native (non-quantized) tensors are views over a + file-backed shared mapping. An in-place elementwise transform on such a view + can be fused by the array library's buffer-donation optimization into a + write *through* the mapping, mutating the file on disk. Copying the data out + to host first breaks that aliasing, so the transform result is computed in a + private buffer and the source file is never touched. Used only by the small + arithmetic transforms (RMSNorm-unbake, SSM ``A``), where the cost is + negligible; bulk quantized tensors stay zero-copy. + """ + if arr.dtype == mx.bfloat16: + # numpy has no bf16 buffer format; both transform call sites compute + # in f32 anyway. astype allocates a fresh buffer, never the mapping. + arr = arr.astype(mx.float32) + return mx.array(np.array(arr)) + + +def remap_arrays( + arrays: dict[str, mx.array], + kquant_meta: dict[str, str], + arch: str, + *, + no_remap: bool = False, + target_prefix: str = "", + fail_on_unknown: bool = False, + n_head: int | None = None, + n_head_kv: int | None = None, + owned_names: set[str] | None = None, +) -> tuple[dict[str, mx.array], dict[str, str], dict[str, int]]: + """Apply name remap + layout transforms to GGUF arrays. + + Returns ``(hf_weights, hf_kquant_meta, stats)`` where ``hf_kquant_meta`` + maps the post-remap tensor name to its codec string. + + ``n_head`` / ``n_head_kv`` are required when any tensor needs the LLAMA Q/K + permute applied. When omitted, the qk_permute transform falls back to a + pass-through with a warning (the resulting model mis-attends). + + ``owned_names``, when given, collects the post-remap names of arithmetic + transform results (qk_permute, SSM A, gemma norm-unbake): arrays that must + own their buffers, never alias the source mapping (donation tripwire; see + ``_verify_zero_copy_views``). Shape-op transforms legitimately alias and + are not collected. + """ + hf_weights: dict[str, mx.array] = _RemapDict() + hf_kquant_meta: dict[str, str] = {} + stats = { + "mapped": 0, + "skipped": 0, + "split": 0, + "failed": 0, + "passthrough": 0, + "qk_permute_applied": 0, + "qk_permute_skipped": 0, + "conv1d_unsqueeze": 0, + "kda_conv_weight": 0, + "gemma_norm_minus_one": 0, + } + + # We process weight tensors; .scales sibling placeholders produced by the + # wire-byte loader get re-emitted alongside their weight under the HF name. + for name, arr in arrays.items(): + if name.endswith(".scales") or name.endswith(".biases"): + continue + codec = kquant_meta.get(name) + + if no_remap: + hf_name = name + transform = "passthrough" + else: + dec = parse_gguf_name(arch, name) + if dec.kind == RemapDecision.KIND_SKIP: + stats["skipped"] += 1 + continue + if dec.kind == RemapDecision.KIND_FAIL: + if fail_on_unknown: + raise RuntimeError(f"unmapped tensor {name!r}: {dec.reason}") + loadlog.warn( + f"WARNING: skipping unmapped tensor {name!r}: {dec.reason}" + ) + stats["failed"] += 1 + continue + hf_name = retarget(dec.hf_name, target_prefix) + transform = dec.transform + + if transform == "passthrough": + hf_weights[hf_name] = arr + if codec is not None: + hf_weights[_strip_weight(hf_name) + ".scales"] = arrays[ + _strip_weight(name) + ".scales" + ] + hf_kquant_meta[hf_name] = codec + stats["passthrough"] += 1 + stats["mapped"] += 1 + + elif transform == "moe_split_gate_up": + base = hf_name[: -len("gate_up_proj.weight")].rstrip(".") + gate_name = f"{base}.gate_proj.weight" + up_name = f"{base}.up_proj.weight" + gate, up = split_fused_gate_up_kquant(arr) + hf_weights[gate_name] = gate + hf_weights[up_name] = up + if codec is not None: + # Both halves get a vestigial scales entry under their own name. + hf_weights[_strip_weight(gate_name) + ".scales"] = mx.zeros( + (1,), dtype=mx.uint8 + ) + hf_weights[_strip_weight(up_name) + ".scales"] = mx.zeros( + (1,), dtype=mx.uint8 + ) + hf_kquant_meta[gate_name] = codec + hf_kquant_meta[up_name] = codec + stats["split"] += 1 + stats["mapped"] += 2 + + elif transform == "altup_split": + # gemma-3n stores the AltUp (un)projections as one stacked 3-D + # tensor; the MLX-native layout (GGUF dims reversed) is + # (altup_num_inputs-1, out, in). mlx_lm wants a list of separate + # Linears, so emit `{base}.{i}.weight` per stack slice. These are + # plain F16 tensors (not kquant), so a pure array slice suffices. + base = _strip_weight(hf_name) + for i in range(arr.shape[0]): + hf_weights[f"{base}.{i}.weight"] = arr[i] + stats["mapped"] += 1 + stats["split"] += 1 + + elif transform == "qk_permute": + # llama.cpp's convert_hf_to_gguf::LlamaModel.permute reorders Q/K + # rows so ggml's interleaved-pairs RoPE matches HF's concat-half + # RoPE. mlx-lm's llama/mistral3 attention uses the HF layout, so we + # undo the permute when loading from GGUF directly. + is_k = hf_name.endswith("k_proj.weight") + n_heads_for = n_head_kv if (is_k and n_head_kv is not None) else n_head + if n_heads_for is None: + loadlog.warn( + f"WARNING: qk_permute requested for {hf_name!r} but " + f"n_head/n_head_kv not provided; loading without " + f"permute (attention will be wrong)." + ) + hf_weights[hf_name] = arr + stats["qk_permute_skipped"] += 1 + else: + hf_weights[hf_name] = qk_permute_wire(arr, n_heads_for) + stats["qk_permute_applied"] += 1 + if owned_names is not None: + owned_names.add(hf_name) + if codec is not None: + hf_weights[_strip_weight(hf_name) + ".scales"] = arrays[ + _strip_weight(name) + ".scales" + ] + hf_kquant_meta[hf_name] = codec + stats["mapped"] += 1 + + elif transform == "conv1d_unsqueeze": + # Pure shape op; works on any dtype. Mamba conv weights are + # F32/BF16 (not kquant), so codec is None here. + hf_weights[hf_name] = arr[..., None] + stats["conv1d_unsqueeze"] += 1 + stats["mapped"] += 1 + + elif transform == "kda_conv_weight": + # Kimi-K3 KDA depthwise short conv: GGUF ships (1, d_inner, 1, + # d_conv) numpy order, or (d_inner, 1, d_conv) when quantization + # drops the trailing 1. conv_step varies fastest in both, so a + # pure reshape to (d_inner, d_conv) is exact; mlx Conv1d wants + # (out_channels=d_inner, kernel=d_conv, in/groups=1). + hf_weights[hf_name] = arr.reshape(-1, arr.shape[-1])[..., None] + stats["kda_conv_weight"] += 1 + stats["mapped"] += 1 + + elif transform == "ssm_a_to_a_log": + # GGUF stores SSM_A as -exp(A_log); invert to recover A_log. + # Squeeze extra leading dim (nemotron_h stores as [1, N]). + # _own() first: the negate/log would otherwise be donated into the + # source mapping (see _own docstring). + out = mx.log(-_own(arr).astype(mx.float32)) + hf_weights[hf_name] = out.reshape(-1) if out.ndim > 1 else out + stats["mapped"] += 1 + if owned_names is not None: + owned_names.add(hf_name) + + elif transform == "flatten": + # Reshape multi-dim tensor to 1D (e.g. nemotron_h ssm_norm stored as + # [n_groups, group_size], ssm_d stored as [1, N]). + hf_weights[hf_name] = arr.reshape(-1) + stats["mapped"] += 1 + + elif transform == "gate_1d_unsqueeze": + # Shared expert gate: GGUF stores 1D [hidden_size], but + # nn.Linear(hidden_size, 1, bias=False) has weight [1, hidden_size]. + hf_weights[hf_name] = arr.reshape(1, -1) if arr.ndim == 1 else arr + stats["mapped"] += 1 + + elif transform == "gemma_norm_minus_one": + # llama.cpp bakes +1 into gemma RMSNorm weights at conversion (the + # GGUF stores hf_weight + 1, used directly by ggml). mlx_lm's + # gemma/gemma2/gemma3 RMSNorm computes rms_norm(x, 1.0 + weight), + # i.e. it expects the *raw* HF weight - so undo the bake here. + # (gemma4_text uses its norm weight directly and is not tagged.) + # _own() first so the subtract isn't donated back into the source + # mapping (see _own docstring). + hf_weights[hf_name] = _own(arr).astype(mx.float32) - 1.0 + stats["gemma_norm_minus_one"] += 1 + stats["mapped"] += 1 + if owned_names is not None: + owned_names.add(hf_name) + + else: + raise RuntimeError(f"unknown transform {transform!r} for {name!r}") + + # Hand back a plain dict: the anti-clobber guard applies to remap + # population only. Later stages (native-fp repack, transforms) + # legitimately replace entries in place. + return dict(hf_weights), hf_kquant_meta, stats + + +def strip_nextn_trunk_overflow( + hf_weights: dict, hf_kquant_meta: dict, meta, arch: str +) -> int: + """Drop remapped weights of trailing NextN/MTP block(s) from the trunk tree. + + nemotron_h_moe GGUFs carry the MTP layer as ``blk.{block_count - 1}`` with + the same tensor names as trunk blocks, so the trunk remap emits + ``backbone.layers.{N}.*`` entries for a layer index the trunk model does not + have (llama.cpp likewise excludes nextn layers from the trunk graph; the + stock nemotron_h ``sanitize`` only strips HF-named ``mtp.*`` keys). The MTP + drafter loads that block separately. Returns the number of entries dropped. + """ + if arch != "nemotron_h_moe": + return 0 + nextn = read_int(meta, f"{arch}.nextn_predict_layers") or 0 + block_count = read_int(meta, f"{arch}.block_count") or 0 + if nextn <= 0 or block_count <= nextn: + return 0 + trunk = block_count - nextn + # backbone.*: the NEMOTRON_H_MOE override table; model.*: MTP-block + # tensors the override table does not claim (post_attention_norm) fall + # through to the canonical map's model.layers.{N}.* naming. + pat = re.compile(r"^(?:backbone|model)\.layers\.(\d+)\.") + dropped = 0 + for name in list(hf_weights): + m = pat.match(name) + if m and int(m.group(1)) >= trunk: + del hf_weights[name] + hf_kquant_meta.pop(name, None) + dropped += 1 + return dropped + + +# MTP / "nextn" drafter remap (native-head: the drafter weights live in the +# GGUF's own MTP block, i.e. block index >= num_hidden_layers) + +# The four ``nextn.*`` extras -> the mlx-vlm ``Qwen3_5MTPDraftModel`` param tree. +# The MTP block's *standard* decoder tensors (attn_*, ffn_*, the two block norms) +# reuse the canonical text remap (``parse_gguf_name``) with ``model.layers.{N}.`` +# rewritten to the drafter's ``layers.{i}.``. The embed table + LM head are not +# here - the drafter binds the target's at runtime (qwen3.5/3.6 GGUFs carry no +# ``nextn.embed_tokens`` / ``nextn.shared_head_head``). +_MTP_NEXTN_MAP = { + "eh_proj": "fc.weight", + "enorm": "pre_fc_norm_embedding.weight", + "hnorm": "pre_fc_norm_hidden.weight", + "shared_head_norm": "norm.weight", +} + + +def remap_mtp_arrays( + arrays: dict[str, mx.array], + kquant_meta: dict[str, str], + arch: str, + *, + first_mtp_block: int, + num_mtp_layers: int = 1, + n_head: int | None = None, + n_head_kv: int | None = None, +) -> tuple[dict[str, mx.array], dict[str, str], dict[str, int]]: + """Remap a GGUF's native MTP block(s) onto the drafter's ``mtp.*`` tree. + + ``first_mtp_block`` is the GGUF block index of the first MTP block (equals + the target's ``num_hidden_layers``); block ``first_mtp_block + i`` maps to the + drafter's ``layers.{i}``. Returns drafter-relative names (no ``model.`` + prefix); the caller builds the drafter and ``load_weights`` these onto it. + + Self-contained (does not touch the text-path ``remap_arrays``): it reuses + ``parse_gguf_name`` for the standard decoder tensors' name+transform decision + and the shared standalone transforms for emit. + """ + hf_weights: dict[str, mx.array] = _RemapDict() + hf_kquant_meta: dict[str, str] = {} + stats = { + "mapped": 0, + "skipped": 0, + "split": 0, + "passthrough": 0, + "qk_permute_applied": 0, + "qk_permute_skipped": 0, + "conv1d_unsqueeze": 0, + } + + def _emit(hf_name: str, transform: str, arr, codec, src_name: str) -> None: + if transform == "passthrough": + hf_weights[hf_name] = arr + if codec is not None: + hf_weights[_strip_weight(hf_name) + ".scales"] = arrays[ + _strip_weight(src_name) + ".scales" + ] + hf_kquant_meta[hf_name] = codec + stats["passthrough"] += 1 + stats["mapped"] += 1 + elif transform == "moe_split_gate_up": + base = hf_name[: -len("gate_up_proj.weight")].rstrip(".") + gate_name = f"{base}.gate_proj.weight" + up_name = f"{base}.up_proj.weight" + gate, up = split_fused_gate_up_kquant(arr) + hf_weights[gate_name] = gate + hf_weights[up_name] = up + if codec is not None: + hf_weights[_strip_weight(gate_name) + ".scales"] = mx.zeros( + (1,), dtype=mx.uint8 + ) + hf_weights[_strip_weight(up_name) + ".scales"] = mx.zeros( + (1,), dtype=mx.uint8 + ) + hf_kquant_meta[gate_name] = codec + hf_kquant_meta[up_name] = codec + stats["split"] += 1 + stats["mapped"] += 2 + elif transform == "gate_1d_unsqueeze": + hf_weights[hf_name] = arr.reshape(1, -1) if arr.ndim == 1 else arr + stats["mapped"] += 1 + elif transform == "flatten": + hf_weights[hf_name] = arr.reshape(-1) + stats["mapped"] += 1 + elif transform == "qk_permute": + is_k = hf_name.endswith("k_proj.weight") + nh = n_head_kv if (is_k and n_head_kv is not None) else n_head + if nh is None: + loadlog.warn( + f"WARNING: qk_permute requested for {hf_name!r} but " + f"n_head/n_head_kv not provided; loading without " + f"permute (attention will be wrong)." + ) + hf_weights[hf_name] = arr + stats["qk_permute_skipped"] += 1 + else: + hf_weights[hf_name] = qk_permute_wire(arr, nh) + stats["qk_permute_applied"] += 1 + if codec is not None: + hf_weights[_strip_weight(hf_name) + ".scales"] = arrays[ + _strip_weight(src_name) + ".scales" + ] + hf_kquant_meta[hf_name] = codec + stats["mapped"] += 1 + elif transform == "conv1d_unsqueeze": + hf_weights[hf_name] = arr[..., None] + stats["conv1d_unsqueeze"] += 1 + stats["mapped"] += 1 + else: + raise RuntimeError( + f"MTP remap: unsupported transform {transform!r} for {src_name!r}" + ) + + mtp_blocks = {first_mtp_block + i: i for i in range(num_mtp_layers)} + for name, arr in arrays.items(): + if name.endswith(".scales") or name.endswith(".biases"): + continue + m = re.match(r"^blk\.(\d+)\.(.+)$", name) + if not m: + continue + blk = int(m.group(1)) + if blk not in mtp_blocks: + continue + layer_i = mtp_blocks[blk] + rest = m.group(2) + codec = kquant_meta.get(name) + if rest.startswith("nextn."): + key = rest[len("nextn.") :] + base = key[: -len(".weight")] if key.endswith(".weight") else key + target = _MTP_NEXTN_MAP.get(base) + if target is None: + # e.g. nextn.embed_tokens / shared_head_head - shared from target. + stats["skipped"] += 1 + continue + _emit(target, "passthrough", arr, codec, name) + else: + dec = parse_gguf_name(arch, name) + if dec.kind != RemapDecision.KIND_MAP: + stats["skipped"] += 1 + continue + marker = f"model.layers.{blk}." + if marker not in dec.hf_name: + stats["skipped"] += 1 + continue + inner = dec.hf_name.split(marker, 1)[1] + _emit(f"layers.{layer_i}.{inner}", dec.transform, arr, codec, name) + # Hand back a plain dict: the anti-clobber guard applies to remap + # population only. Later stages (native-fp repack, transforms) + # legitimately replace entries in place. + return dict(hf_weights), hf_kquant_meta, stats + + +def remap_gemma4_assistant_arrays(arrays: dict, kquant_meta: dict): + """Remap a gemma4 assistant-drafter GGUF onto the mlx-vlm + ``Gemma4AssistantDraftModel`` param tree. + + The standard decoder / embed / norm tensors reuse the canonical gemma4 remap + (``parse_gguf_name`` already emits the exact ``model.*`` names the drafter + uses, including ``layer_output_scale -> layers.N.layer_scalar`` and + ``output_norm -> model.norm``); only the two bridge projections need renaming + and ``rope_freqs`` is dropped (it's computed, not a param). Every gemma4 + tensor maps as a passthrough (no qk-permute), so the emit is direct. + """ + hf_weights: dict[str, mx.array] = _RemapDict() + hf_kquant_meta: dict[str, str] = {} + stats = {"mapped": 0, "skipped": 0} + for name, arr in arrays.items(): + if name.endswith(".scales") or name.endswith(".biases"): + continue + base = name[: -len(".weight")] if name.endswith(".weight") else name + if base.endswith("pre_proj") or base.endswith("pre_projection"): + hf = "pre_projection.weight" + elif base.endswith("post_proj") or base.endswith("post_projection"): + hf = "post_projection.weight" + elif base.endswith("centroids"): + # ordered-embeddings sparse head (E2B/E4B); Q8_0, swapped by kquant. + hf = "masked_embedding.centroids.weight" + elif base.endswith("token_ordering"): + # I32 index vector, no .weight suffix on the param, never quantized. + hf_weights["masked_embedding.token_ordering"] = arr.astype(mx.int32) + stats["mapped"] += 1 + continue + elif base == "rope_freqs": + stats["skipped"] += 1 + continue + else: + dec = parse_gguf_name("gemma4", name) + if dec.kind != RemapDecision.KIND_MAP: + stats["skipped"] += 1 + continue + if dec.transform != "passthrough": + raise RuntimeError( + f"gemma4 assistant remap: unexpected transform " + f"{dec.transform!r} for {name!r}" + ) + hf = dec.hf_name + codec = kquant_meta.get(name) + hf_weights[hf] = arr + if codec is not None: + hf_weights[_strip_weight(hf) + ".scales"] = arrays[ + _strip_weight(name) + ".scales" + ] + hf_kquant_meta[hf] = codec + stats["mapped"] += 1 + # Hand back a plain dict: the anti-clobber guard applies to remap + # population only. Later stages (native-fp repack, transforms) + # legitimately replace entries in place. + return dict(hf_weights), hf_kquant_meta, stats diff --git a/gmlx/serve/bridge_vlm.py b/gmlx/serve/bridge_vlm.py index 429ccb69..db5aa4e3 100644 --- a/gmlx/serve/bridge_vlm.py +++ b/gmlx/serve/bridge_vlm.py @@ -541,17 +541,17 @@ def _install_stream_placement( # The whole model runs on the CPU device, and this device change # applies to the process. Use it for one over-RAM model. Do not mix # it with GPU-resident models in one config-mode server. - from gmlx.load.loader import configure_stream_cpu + from gmlx.stream.expert_streaming import configure_stream_cpu configure_stream_cpu( text_model, gguf_path=gguf_path, feeder_prefill=feeder_prefill, feeder_decode=feeder_decode) elif stream: # "experts": routed experts stream; rest of model + KV on GPU - from gmlx.load.loader import install_expert_streaming + from gmlx.stream.expert_streaming import install_expert_streaming install_expert_streaming( text_model, gguf_path=gguf_path, feeder_prefill=feeder_prefill, feeder_decode=feeder_decode) if moe_experts is not None: - from gmlx.load.loader import install_moe_experts_override + from gmlx.stream.expert_streaming import install_moe_experts_override install_moe_experts_override(text_model, moe_experts) if moe_expert_mass is not None: from gmlx.stream.moe_experts import install_moe_expert_mass diff --git a/gmlx/serve/residency.py b/gmlx/serve/residency.py index 6ba2178e..8566633a 100644 --- a/gmlx/serve/residency.py +++ b/gmlx/serve/residency.py @@ -922,7 +922,7 @@ def _build(self, cache_key, model_path, adapter_path, model_kind, footprint, # Before the walk, not at the install after it: a resident # model's generator left the wired limit raised, and the walk # under it wires the file's pages (see the loader). - from gmlx.load.loader import _neutralize_wired_limit_sweep + from gmlx.stream.wired_limit import _neutralize_wired_limit_sweep _neutralize_wired_limit_sweep() # The walk wraps the routed experts as tracked MLX views. They # are page cache, but the install credits them only at the diff --git a/gmlx/spec/mtp_load.py b/gmlx/spec/mtp_load.py index e199780e..a60978bf 100644 --- a/gmlx/spec/mtp_load.py +++ b/gmlx/spec/mtp_load.py @@ -29,22 +29,26 @@ from gmlx.load.gguf_meta import first_nonzero_int, read_int, read_string from gmlx.load.loader import ( _FP32_KEEP_BY_MODEL_TYPE, - _MTP_TARGET_HOOKS, - _MTP_TARGET_HOOKS_BY_TYPE, _active_now, - _ensure_argmax_hook, _install_and_load, _resolve_chat_template, - _spec_hook_key, build_model, - load_gguf_wire_bytes, materialize_module_arrays, model_is_moe, print_inventory, + weights_source_key, +) +from gmlx.load.mtp_target import ( + _MTP_TARGET_HOOKS, + _MTP_TARGET_HOOKS_BY_TYPE, + _ensure_argmax_hook, + _spec_hook_key, +) +from gmlx.load.wire import ( + load_gguf_wire_bytes, remap_arrays, remap_gemma4_assistant_arrays, remap_mtp_arrays, - weights_source_key, ) from gmlx.load.native_fp import _strip_weight from gmlx.load.populate import maybe_populate_for_load @@ -1757,7 +1761,7 @@ def load_mtp_model( n_head_kv=n_head_kv, owned_names=owned_names, ) - from gmlx.load.loader import strip_nextn_trunk_overflow + from gmlx.load.wire import strip_nextn_trunk_overflow n_nextn_dropped = strip_nextn_trunk_overflow(hf_weights, hf_kquant_meta, meta, arch) if n_nextn_dropped: diff --git a/gmlx/stream/budget.py b/gmlx/stream/budget.py index 1cc270c6..05d23db0 100644 --- a/gmlx/stream/budget.py +++ b/gmlx/stream/budget.py @@ -18,9 +18,12 @@ from __future__ import annotations import os +import re from dataclasses import dataclass -from gmlx.envflags import env_int +import mlx.core as mx + +from gmlx.envflags import env_bool, env_int _DEFAULT_CTX = 32768 _LEGACY_RESERVE_GB = 8.0 @@ -59,8 +62,6 @@ def reclaimable_ram_bytes() -> int | None: v = None if v is not None: return int(v) - from gmlx.load.loader import _available_ram_bytes - return _available_ram_bytes() @@ -153,3 +154,181 @@ def kv_room_bytes(gguf_path: str | None, env: dict | None = None) -> KvRoom: _cap_bytes()) except Exception: return price_room(None, None, 0, 0) + + +def _available_ram_bytes(include_inactive: bool = True) -> int | None: + """RAM this process can take without swapping anyone's anonymous memory: + free + purgeable + the file-backed page cache (macOS ``vm_stat``). A + load-time snapshot of the machine's offer - a machine already + half-occupied by other workloads offers the arena half a machine, + whatever the hardware total says. File-backed pages drop without IO + whatever queue they sit on; counting only the *inactive* queue (the old + formula) missed the tens of GB of recently-read GGUF cache still on the + active queue and made a mostly-cache machine look nearly full. A + ``vm_stat`` without the ``File-backed pages`` line falls back to + inactive + speculative. + + ``include_inactive=False`` is the stricter set (free + purgeable + + speculative only) for a caller that must not take the page cache.""" + from gmlx.serve import kernel_vm + + s = kernel_vm.snapshot() + if s is not None: + return s["free"] + s["purgeable"] + ( + s["filebacked"] if include_inactive else s["speculative"]) + import subprocess + + try: + # posix_spawn (absolute path, close_fds=False): a fork beside a + # Metal-mapped buffer copies the buffer first. + out = subprocess.run( + ["/usr/bin/vm_stat"], capture_output=True, text=True, timeout=5, + close_fds=False, + ).stdout + except Exception: + return None + m = re.search(r"page size of (\d+)", out) + if not m: + return None + keys = ["free", "purgeable"] + pages = 0 + found = False + if include_inactive: + mm = re.search(r"File-backed pages:\s+(\d+)\.", out) + if mm: + pages += int(mm.group(1)) + found = True + else: + keys += ["speculative", "inactive"] + else: + keys.append("speculative") + for key in keys: + mm = re.search(rf"Pages {key}:\s+(\d+)\.", out) + if mm: + pages += int(mm.group(1)) + found = True + return pages * int(m.group(1)) if found else None + + +def _ram_floor_bytes(ram: int | None) -> int: + """``gmlx.stream.budget.host_floor_bytes``.""" + return host_floor_bytes(ram) + + +def _decode_arena_bytes( + total_bytes: int, offsets, budget: int | None, room_bytes: int | None = None, + pinned_bytes: int = 0, streamable_bytes: int = 0, + cast_dead_bytes: int = 0, ring_bytes: int = 0, +) -> int: + """Arena budget for the decode feeder: what the memory ceiling leaves + after the non-expert weights, the KV room, the prefill ring and the + host floor, clamped to the RAM reclaimable right now, and capped at + the expert bytes themselves (a model whose experts fit goes fully + resident). + + The ceiling is the serve governor's (``gmlx.stream.budget``), so the + arena, the prefill ring and the KV cache share one budget: at decode + the governor's headroom is the room minus live KV, whatever the box's + working-set ratio or the quant's ring size. ``ring_bytes`` keeps the + ring's room out of the arena for good: a ring rebuilt on top of a + full wired arena, or lent out of it with a copy of every layer, is a + transient the kernel has to swap for on a box the arena has already + filled. ``room_bytes`` is the priced KV room + (``budget.kv_room_bytes``); None keeps the flat legacy reserve. + ``GMLX_DECODE_ARENA_RAM_FRAC`` caps the ceiling at a fraction of + physical RAM when set. ``GMLX_DECODE_ARENA_GB`` overrides the + ceilings but is still clamped to what is reclaimable minus the floor + - an arena wired past that starves the page cache every buffered + read path depends on (``GMLX_DECODE_ARENA_FORCE=1`` restores the + unclamped behavior). + + A second live streaming install needs no term here: mlock moves a page + out of the file-backed count and an arena is anonymous, so the + reclaimable snapshot already excludes both.""" + env = os.environ.get("GMLX_DECODE_ARENA_GB") + if env: + want = int(float(env) * (1 << 30)) + if env_bool("GMLX_DECODE_ARENA_FORCE", False): + return want + avail = _available_ram_bytes() + if avail is None: + return want + try: + ram = int(mx.device_info()["memory_size"]) + except Exception: + ram = avail + cap = max(0, avail - _ram_floor_bytes(ram)) + if want > cap: + print( + f"[stream] GMLX_DECODE_ARENA_GB={env} exceeds reclaimable" + f" RAM minus the floor; clamping the arena to" + f" {cap / (1 << 30):.1f}GB (GMLX_DECODE_ARENA_FORCE=1" + f" overrides)" + ) + return cap + return want + if budget is None: + return 0 + ceiling = int(ceiling_bytes() or budget) + ram = None + try: + ram = int(mx.device_info()["memory_size"]) + except Exception: + pass + frac = os.environ.get("GMLX_DECODE_ARENA_RAM_FRAC", "") + if frac and ram: + try: + ceiling = min(ceiling, int(float(frac) * ram)) + except (ValueError, OverflowError): + pass + expert_bytes = sum(r[2] for ranges in offsets.values() for r in ranges) + # Streamable components are page-cache citizens like the experts; + # charging them as non-expert would zero the arena. Cast tensors cost + # what their converted copy weighs, not what the wire does: the wire + # range is unpinned and never read again (gmlx.stream.pin_weights + # .cast_copies), so charging it would cancel the pin it just freed. + non_expert_bytes = max( + 0, total_bytes - expert_bytes - streamable_bytes - cast_dead_bytes) + room = int(room_bytes) if room_bytes is not None else legacy_room_bytes() + # The floor on both measures: the ceiling is a share of the Metal + # working set, and the OS side (the page cache, other processes) is + # not in it. + arena = (ceiling - non_expert_bytes - room - int(ring_bytes) + - _ram_floor_bytes(ram)) + # Second ceiling: what is reclaimable right now. The governor ceiling + # assumes an otherwise idle machine; co-resident workloads shrink the + # offer, and a wired arena sized past it would evict them to swap. The + # floor keeps a breathing margin for the system. This is a live + # post-pin snapshot: already-wired weights are out of it, so only the + # still-unwired share of the non-expert set is charged (charging all + # of it double-counted the pin and zeroed the arena on exactly the + # models that need it). + avail = _available_ram_bytes() + if avail is not None: + unpinned = max(0, non_expert_bytes - pinned_bytes) + arena = min( + arena, + avail - _ram_floor_bytes(ram or avail) - room - unpinned + - int(ring_bytes), + ) + return min(max(0, arena), expert_bytes) + + +def _prefill_ring_reason(offsets, left: int | None) -> str | None: + """Why the prefill ring must not be built, or None. The ring (two + slots of the largest layer's expert stacks, sized by the model) takes + its room under the memory ceiling before the decode arena; ``left`` + is what the ceiling leaves after the every-token weights and the KV + room. A ring larger than that would sit on top of them and take the + ceiling with it at the first prefill. An explicit GMLX_DECODE_ARENA_GB + is the user's budget and the ring is not judged against it.""" + if left is None or os.environ.get("GMLX_DECODE_ARENA_GB"): + return None + from gmlx.stream.prefill_feeder import ring_bytes + + ring = ring_bytes(offsets) + if ring <= left: + return None + return (f"ring 2 x {ring / 2e9:.1f} GB exceeds the {left / 1e9:.1f} GB " + "left under the memory ceiling after the every-token weights " + "and the KV room") diff --git a/gmlx/stream/decode_feeder.py b/gmlx/stream/decode_feeder.py index 50eb41c3..7f6a42c3 100644 --- a/gmlx/stream/decode_feeder.py +++ b/gmlx/stream/decode_feeder.py @@ -1494,7 +1494,7 @@ def _regrow_headroom_ok(self) -> bool: it, so a regrow cannot trip the floor that shrank the arena.""" need = self._arena_bytes_at(self._pressure_steps - 1) - self.arena_bytes try: - from gmlx.load.loader import _ram_floor_bytes + from gmlx.stream.budget import _ram_floor_bytes from gmlx.stream.budget import kernel_floor_bytes, reclaimable_ram_bytes avail = reclaimable_ram_bytes() diff --git a/gmlx/stream/expert_streaming.py b/gmlx/stream/expert_streaming.py new file mode 100644 index 00000000..3396d4f1 --- /dev/null +++ b/gmlx/stream/expert_streaming.py @@ -0,0 +1,1219 @@ +"""MoE expert streaming: CPU offload install, GPU residency, prefill step.""" +from __future__ import annotations + +import os +import random +import time + +import mlx.core as mx +import numpy as np +from mlx.utils import tree_flatten + +from gmlx.envflags import env_bool, env_int +from gmlx.gen.prefill_decay import deduct_untracked_weights +from gmlx.load import loadlog +from gmlx.load.loader import ( + _PHASE, + _STREAM_GPU_TOKENS_DEFAULT, + _STREAM_PREFETCH_MIN_TOKENS, + _STREAMING_PREFILL_STEP, + _STREAMING_PREFILL_STEP_BY_MODEL_TYPE, + _arena_split_max_tokens, + _arena_stage_max_tokens, + _kq_expert_gpu_ok, + _lookahead_default, + _phase_token, + _resolve_feeder_defaults, + _stream_gpu_tokens, + _switch_num_experts, + moe_streaming_active, +) + +from .budget import _decode_arena_bytes, _prefill_ring_reason, _ram_floor_bytes +from .wired_limit import _neutralize_wired_limit_sweep, configure_cpu_device + + +# MoE expert CPU offload (hybrid GPU+CPU inference) +# +# On unified memory the GPU constraint is the wired limit, not a separate +# VRAM pool: Metal-resident buffers must be wired, while CPU-consumed mmap +# pages ride the page cache (evictable, can exceed RAM). For fine-grained MoE +# the routed expert stacks are ~90-95% of the bytes but each expert is read +# with probability top_k/n_experts per token, while the every-token layers +# (attention, norms, routers, shared experts, embeddings, KV cache) are read +# every token. +# Running the SwitchGLU expert containers on the CPU stream therefore keeps +# those hot layers + KV on GPU while the expert wire bytes stay file-backed in +# the page cache when the GPU is idle, and the kquant gather op executes on +# its threaded CPU path. MLX's cross-stream dependency tracking handles the +# GPU->CPU->GPU handoff inside each MoE layer (zero-copy - same pages). +# +# Residency (measured): Metal wires only what GPU work references or what +# sits in MLX's residency set - unreferenced file-backed buffers stay +# evictable page cache even under full memory pressure. The one hazard is +# mlx-lm's generation-time wired-limit bump (see +# _neutralize_wired_limit_sweep): MLX services a raised wired limit by +# sweeping every live buffer into the residency set, offloaded experts +# included. Models larger than the wired budget therefore run in streaming +# mode: the sweep is neutralized and GPU prefill routing is forced off, so +# the GPU never references (and never wires) expert bytes, and the page +# cache streams them from disk. +# +# Prefill staging: at decode each expert sees ~top_k/n_experts of one token, +# but a prefill chunk makes every expert hot with tens of rows each - a GEMM +# workload where the CPU (~1.5 TFLOP/s) is the wrong device. Calls with at +# least GMLX_STREAM_GPU_TOKENS tokens therefore run on the default (GPU) +# stream against the same zero-copy buffers - no copies, no staging; the +# driver wires the touched expert bytes for the duration of the work and +# releases them when the GPU goes idle. Threshold 0 disables GPU routing +# (pure CPU experts, the conservative choice when the model is far larger +# than RAM and prefill-wiring every expert is undesirable). +# +# Cost model (measured): offloaded decode pays a per-layer surcharge of +# genuine CPU dot compute plus per-layer stream fences and CPU-pool +# wake-from-idle (3 wakes per layer, one per gather; the wake cost grows +# when the pool sits idle between layers while the GPU runs the +# every-token layers). +# Routing every call to the GPU stream instead (GMLX_STREAM_GPU_TOKENS=1, +# no CPU hop) runs ~3.7x faster on a fits-in-RAM MoE, so in-RAM the CPU +# offload is for the over-budget regime, not the fast path. + +_CPU_OFFLOAD_CLASS_CACHE: dict = {} + + +def configure_stream_cpu( + model, + gguf_path: str | None = None, + feeder_prefill: bool | None = None, + feeder_decode: bool | None = None, +): + """Whole-model CPU streaming (``--stream-cpu``): run the model on the CPU + device with the streaming-expert machinery always engaged. + + ``--stream-cpu`` is an explicit opt-in into the CPU/over-RAM path, so it + forces streaming (``force_stream=True``) regardless of model size - experts + run on the CPU stream whether or not the model fits the wired budget (a + fits-in-RAM model is then served from the page cache rather than faulting + from disk; for the faster all-GPU path on a model that fits, omit + ``--stream-cpu``). The GPU + working-set budget is still captured before switching the default device to + CPU so the over-/under-budget log line stays accurate (the CPU device would + otherwise report a budget that hides the condition). Returns + ``(n_wrapped, offloaded_bytes)``. + """ + try: + gpu_info = dict(mx.device_info()) + except Exception: + gpu_info = None + configure_cpu_device() + if gpu_info and "max_recommended_working_set_size" in gpu_info: + mx.device_info = lambda: gpu_info + return install_expert_streaming( + model, + gguf_path=gguf_path, + force_stream=True, + feeder_prefill=feeder_prefill, + feeder_decode=feeder_decode, + ) + + +def _install_gpu_residency(model, moe_modules, *, + skip_ids=frozenset(), + include_expert_stacks: bool = False) -> None: + """Wire every non-expert weight buffer into the Metal residency set, + so command buffers stop re-wiring the every-token weights' pages on + every use (the + per-use wiring is what an unswept streaming install pays instead of + the neutralized wire-everything sweep). + + ``skip_ids``: arrays that must NOT be inserted - streamed lookup + tables (a residency insert wires the buffer as surely as a GPU op). + ``include_expert_stacks``: table-only streaming keeps the experts + resident, so the GB-scale-stack belt is lifted and they are wired + with everything else.""" + import mlx_kquant as kq + + if not getattr(kq, "residency_insert", None): + print("[stream] gpu-resident weights unavailable " + "(mlx-kquant lacks residency ops)") + return + skip = set(skip_ids) + for mods in moe_modules.values(): + for m in mods: + for attr in ("gate_proj", "up_proj", "down_proj"): + w = getattr(getattr(m, attr, None), "weight", None) + if w is not None: + skip.add(id(w)) + inserted = [] + nbytes = 0 + for _, a in tree_flatten(model.parameters()): + if id(a) in skip: + continue + if (not include_expert_stacks + and a.ndim == 3 and a.nbytes > (1 << 30)): + continue # belt: any GB-scale stack is an expert container + if kq.residency_insert(a): + inserted.append(a) + nbytes += a.nbytes + kq.residency_commit() + model._kq_resident_arrays = inserted + n = len(inserted) + print(f"[stream] gpu-resident weights: {n} buffers " + f"({nbytes / 1e9:.1f} GB) in the Metal residency set " + "(GMLX_GPU_RESIDENT=0 disables)") + + +def install_expert_streaming( + model, + n_layers: int | None = None, + gguf_path: str | None = None, + force_stream: bool = False, + feeder_prefill: bool | None = None, + feeder_decode: bool | None = None, + stats_verbose: bool | None = None, +): + """Run routed-expert stacks (SwitchGLU) on the CPU stream. + + Wraps each ``SwitchGLU`` in the first ``n_layers`` decoder layers (all + layers when None) so its forward - the expert gather matmuls - executes + under ``mx.stream(mx.cpu)`` at decode shapes, and on the default (GPU) + stream for prefill-sized calls (see the staging note above). Per-instance + ``__class__`` swap; routers, shared experts, attention, and the KV cache + stay on the default (GPU) stream. Returns ``(n_wrapped, offloaded_bytes)``. + + ``gguf_path`` (the loaded checkpoint) enables sequential expert prefetch + for streaming-mode models - see ``gmlx.stream.prefetch``. Without it, + over-budget prefill demand-faults expert bytes at random-read bandwidth. + """ + from gmlx.load.modules import switch_layer_types + + _, glu_types = switch_layer_types() + + layers = getattr(model, "layers", None) + if layers is None: + layers = model.model.layers + + # Streaming mode: neutralize the generation-time residency sweep (which + # would otherwise wire the whole model - see _neutralize_wired_limit_sweep) + # and run every expert call on the CPU stream. Engaged when the model is + # over the wired budget (it must stream) or when force_stream is set: + # --stream-cpu (configure_stream_cpu) passes force_stream so the flag does + # what it says - experts on CPU regardless of model size; on a fits-in-RAM + # model the page cache then serves those bytes from RAM rather than faulting + # from disk. --stream-experts keeps the budget-keyed decision, so below the + # budget it still routes prefill-sized calls to the GPU stream + # (GMLX_STREAM_GPU_TOKENS) - the fast path in-RAM. + params = getattr(model, "parameters", None) + total_bytes = sum(a.nbytes for _, a in tree_flatten(params())) if params else 0 + try: + budget = int(0.9 * mx.device_info()["max_recommended_working_set_size"]) + except Exception: + budget = None + over_budget = budget is not None and total_bytes > budget + + # Selection ladder step 1 (docs/streaming.md): archs with a + # declared streamable lookup table (e.g. qwen4exp's 26.8 GiB PLE + # n-gram table) stream it instead of the experts when it alone brings + # the resident set under budget - table gathers touch ~1.4 KB/token + # against the experts' every-MoE-layer surcharge. The table wrap runs + # its row gather on a dedicated CPU stream so the buffer is never a + # GPU-stream input (a single GPU reference would wire all of it). + # When the post-table estimate is still over budget, v1 falls back to + # expert streaming with the table resident: streaming both at once + # (compose) needs the hot-row arena and is not shipped. + table_offloaded = 0 + if not force_stream: + from gmlx.stream.table_stream import ( + install_table_streaming, + table_bytes, + table_stream_selected, + ) + + compose = False + if table_stream_selected(model, total_bytes, budget): + post = total_bytes - table_bytes(model) + # Step 2 default: over budget even post-table streams both + # (compose). GMLX_STREAM_PLE_COMPOSE=0 keeps the table + # resident; the selection test already honors it in auto + # mode, so this only fires under GMLX_STREAM_PLE=1. + if budget is not None and post > budget: + if env_bool("GMLX_STREAM_PLE_COMPOSE", True): + compose = True + table_offloaded, table_names = ( + install_table_streaming(model)) + else: + print( + "[stream] table stays resident " + "(GMLX_STREAM_PLE_COMPOSE=0); experts stream" + ) + else: + table_offloaded, table_names = install_table_streaming(model) + if table_offloaded and compose: + loadlog.info( + f"[stream] compose: streamable table " + f"{'+'.join(table_names)} " + f"({table_offloaded / 2**30:.1f} GiB) on the CPU stream " + "AND experts streamed" + ) + key = getattr(model, "_kq_weights_key", None) + from gmlx.gen.prefill_decay import ( + note_streamed_tracked_bytes, + untracked_weight_bytes_for, + ) + tracked = max( + 0.0, total_bytes - untracked_weight_bytes_for(key)) + credit = min(float(table_offloaded), tracked) + if credit > 0: + note_streamed_tracked_bytes( + credit, key, source="table", cap=tracked) + deduct_untracked_weights(table_offloaded, key) + elif table_offloaded: + # The selection test admits the table only when the remainder + # clears the budget (or streaming is forced on a fits model), + # so experts are resident from here on. + over_budget = False + base = ("" if budget is None + else f" of {budget / 2**30:.1f} GiB budget") + loadlog.info( + f"[stream] streamable table {'+'.join(table_names)} " + f"({table_offloaded / 2**30:.1f} GiB) stays file-backed on " + "the CPU stream; experts resident (post-deduction " + f"{(total_bytes - table_offloaded) / 2**30:.1f} GiB{base})" + ) + key = getattr(model, "_kq_weights_key", None) + from gmlx.gen.prefill_decay import ( + note_streamed_tracked_bytes, + untracked_weight_bytes_for, + ) + tracked = max( + 0.0, total_bytes - untracked_weight_bytes_for(key)) + credit = min(float(table_offloaded), tracked) + if credit > 0: + note_streamed_tracked_bytes( + credit, key, source="table", cap=tracked) + deduct_untracked_weights(table_offloaded, key) + + streaming = force_stream or over_budget + prefetcher = None + cast_dead_bytes = 0 + held_wired = 0 + if streaming: + _neutralize_wired_limit_sweep() + # Reclaim the wired bytes of released streaming models (feeder and + # MoE modules reference each other, so unwiring waits for a + # collection), then charge what is still held against this weight + # pin. Wired pages are invisible to jetsam, so two pins that each + # size against the whole machine wire it solid. The arena needs no + # charge; see _decode_arena_bytes. + from gmlx.stream import installs as _installs + + freed = _installs.reclaim_dead() + held_wired = _installs.live_wired_bytes() + if freed: + loadlog.info( + f"[stream] reclaimed {freed / 1e9:.1f} GB wired from a " + "released streaming model") + if held_wired: + print( + f"[stream] another live streaming install holds " + f"{held_wired / 1e9:.1f} GB wired; this model sizes against " + "what is left (release the other model first for the full " + "budget)") + from gmlx.stream.prefetch import maybe_make_prefetcher + + prefetcher = maybe_make_prefetcher(gguf_path) + if prefetcher is not None: + object.__setattr__(model, "_kq_prefetcher", prefetcher) + # Wire the every-token weights before the decode feeder sizes its + # arena: pinned every-token pages come out of the same wired budget. + from gmlx.stream.pin_weights import cast_copies, maybe_pin_weights + from gmlx.stream.table_stream import streamable_tables_for + + # Declared streamable components never enter the pin set, streamed + # or resident: mlocking them starves the expert page cache. Cast + # tensors go too - their wire bytes have no view left to keep. + casts = cast_copies(model, gguf_path) + cast_dead_bytes = casts.dead_bytes + pin_exclude = frozenset( + t.gguf_name for t, _ in streamable_tables_for(model) + ) | casts.names + weights_pin = maybe_pin_weights( + gguf_path, exclude_names=pin_exclude, reserved_bytes=held_wired) + if weights_pin is not None: + object.__setattr__(model, "_kq_weights_pin", weights_pin) + _installs.record(model, weights_pin.pinned_bytes) + + def _wrapped_class(cls): + sub = _CPU_OFFLOAD_CLASS_CACHE.get(cls) + if sub is None: + # A fused base consumes routing scores itself (mix seam); a + # stock base (unrecognized activation, e.g. minimax-m3's + # SwiGLUOAI) takes (x, indices) only. The wrapper still + # advertises _kq_scores_sink so blocks hand scores over for + # miss-shed; it strips them before forwarding and applies + # the shed mix python-side. + _fwd_scores = bool(getattr(cls, "_kq_mix_scores", False)) + + class _CPUOffload(cls): + _kq_scores_sink = True + + def __call__(self, x, indices, *args, **kwargs): + # Extra args pass through untouched (e.g. deepseek-v4 + # hands the fused SwitchGLU its routing scores). A base + # without the mix seam takes (x, indices) only: keep the + # scores for the miss-shed hook and strip them from what + # gets forwarded. + scores_arg = args[0] if args else None + if args and not _fwd_scores: + args = args[1:] + # Threshold read per call (cheap; once per MoE layer per + # forward) so env changes A/B without a reload. Streaming + # mode pins everything to CPU: a GPU expert call would + # wire the buffers it references, which an over-budget + # model cannot afford. + cpu_only = getattr(self, "_kq_cpu_only", False) + gpu_tokens = _stream_gpu_tokens( + getattr( + self, "_kq_gpu_tokens_default", _STREAM_GPU_TOKENS_DEFAULT + ) + ) + n_tokens = indices.size // indices.shape[-1] + pf = getattr(self, "_kq_prefetcher", None) + fdr = getattr(self, "_kq_feeder", None) + dfr = getattr(self, "_kq_decode_feeder", None) + small = n_tokens <= _arena_stage_max_tokens() + la = getattr(self, "_kq_lookahead", None) + la_pred = None + ph = _PHASE + if ph is not None: + _phase_token( + ph, getattr(self, "_kq_li", None), n_tokens) + if n_tokens != 1: + ph = None + lsp = getattr(self, "_kq_layer_shed", None) + if lsp is not None and cpu_only and n_tokens == 1: + rng = getattr(self, "_kq_shed_rng", None) + if rng is None: + # per-layer seed: reproducible shed pattern + rng = random.Random( + 0x5EED ^ (getattr(self, "_kq_li", 0) or 0)) + object.__setattr__(self, "_kq_shed_rng", rng) + if rng.random() < lsp: + # Skip the routed path entirely (gather, stage + # and this layer's eval fence). The unmixed + # zeros return makes the block mix nothing and + # still add its shared expert. + if dfr is not None: + dfr._layer_shed_n += 1 + return mx.zeros( + (*x.shape[:-1], indices.shape[-1], + x.shape[-1]), dtype=x.dtype) + gt = getattr(self, "_kq_gpu_token", None) + gt_live = ( + gt is not None + and gt._route_shed is not None + and cpu_only + and n_tokens == 1 + and dfr is not None + and dfr.covers(self._kq_li) + ) + if gt_live: + dfr.ensure_wired() + # Token tick for EVERY covered decode layer, stage + # path included: boundary detection and the + # adaptive hot-set refresh live here. + gt.on_layer_entry( + self._kq_li, + None if getattr(self, "_kq_in_split", False) + else getattr(self, "_kq_miss_shed", None)) + if ( + gt_live + and scores_arg is not None + and not dfr.wedged_at(self._kq_li) + and gt.layer_autonomous(self._kq_li) + ): + # GPU-autonomous layer (gpu-dispatch Tier 2): no + # per-layer eval. route_shed remaps ids to arena + # slots and sheds non-resident experts on the GPU; + # the graph flushes at the next stage-path layer's + # eval or the logits, and the host consumes the + # recorded misses at the token boundary + # (popularity + prestage + fresh slot tables) - see + # gpu_token.py for the fence argument. In adaptive + # mode only layers with a measured hit rate above + # GMLX_AUTO_HOT_HIT run here, so the shed cost per + # layer is near zero. + tbl = gt.table(self._kq_li) + self._kq_cpu_only = False + try: + with dfr.swapped(self._kq_li): + with mx.stream(mx.gpu): + sc_f32 = scores_arg.astype(mx.float32) + slots, mix, m_ids, m_sc = ( + gt._route_shed( + indices.astype(mx.uint32), + sc_f32, tbl)) + mix_c = mix.astype(x.dtype) + if _fwd_scores: + y = super().__call__( + x, slots, mix_c, + *args[1:], **kwargs) + else: + y = super().__call__( + x, slots, *args, **kwargs) + if y.ndim == x.ndim + 1: + y = (y * mix_c[..., None]).sum( + axis=-2) + gt.record( + self._kq_li, indices, sc_f32, + m_ids, m_sc, y) + return y + finally: + self._kq_cpu_only = True + if la is not None and cpu_only and n_tokens == 1: + # Decode only: prefill prestage would fault the + # cold arena while the ring holds the wired budget. + # Lookahead: run the NEXT MoE layer's router on this + # layer's input and evaluate it together with the + # router read below (one sync either way). The + # prediction feeds nothing downstream - it only + # records recall (probe) or drives prestage reads. + # Latent-MoE blocks (kimi-k3) hand the full-width + # router input over out of band; x here is the + # expert container's latent-width input. + x_la = getattr(self, "_kq_la_input", None) + if x_la is None: + x_la = x + if ph is not None: + t_la = time.perf_counter() + la_pred = la.on_call(x_la, indices) + ph["la"] += time.perf_counter() - t_la + else: + la_pred = la.on_call(x_la, indices) + if ( + dfr is not None + and cpu_only + and small + and dfr.covers(self._kq_li) + ): + # Decode feeder: the routed experts are served from + # this layer's wired GPU arena; misses are pread from + # the GGUF into evicted slots first. Small prefill + # chunks take this path too when their routed set + # fits - the arena persists across requests, which is + # what makes repeat short-prompt TTFT cheap. The eval + # is both the router read and the arena-overwrite + # safety fence (see decode_feeder.py). ``stage`` + # returns None when the call routes to more distinct + # experts than the arena has slots - fall through. + t0 = time.perf_counter() if ph is not None else 0.0 + if n_tokens == 1: + dfr.ensure_wired() + # Miss-shed is decode-only: a single-token leaf of an + # arena token split is prefill work, and a shedding + # leaf would return a mixed rank-3 output next to a + # clean leaf's per-expert rank-4 - the reassembly + # concatenate cannot take both. + ms = (None if getattr(self, "_kq_in_split", False) + else getattr(self, "_kq_miss_shed", None)) + sc_f32 = None + if (ms is not None and scores_arg is not None + and n_tokens == 1): + # Shed reads the scores host-side; fold them into + # the router eval so the hook adds a small D2H + # copy, not a second per-layer graph flush. + sc_f32 = scores_arg.astype(mx.float32) + mx.eval(indices, sc_f32) + else: + mx.eval(indices) + if ph is not None: + t1 = time.perf_counter() + ph["ev"] += t1 - t0 + wait0 = getattr(dfr, "_t_demand", 0.0) + ids = np.array(indices) + shed_args = None + shed_mix = None + if sc_f32 is not None: + sc = np.asarray(sc_f32).reshape(-1) + keep = dfr.shed_misses( + self._kq_li, ids.reshape(-1), sc, ms) + if keep is not None: + # Arena-path only: the overflow fallback + # below keeps the original routed set. + kept = ids.reshape(-1)[keep] + shp = ids.shape[:-1] + (kept.size,) + ids = np.ascontiguousarray(kept.reshape(shp)) + scn = sc[keep] + # survivors keep the token's full mass + scn = scn * (sc.sum() / max(scn.sum(), 1e-20)) + sc_mx = mx.array(scn.reshape(shp)).astype( + scores_arg.dtype) + if _fwd_scores: + shed_args = (sc_mx,) + args[1:] + else: + # Stock base returns per-expert outputs; + # the block's weights still cover the + # full routed set, so mix the shed + # survivors here instead. + shed_mix = sc_mx + slots = dfr.stage(self._kq_li, ids) + if ph is not None: + t2 = time.perf_counter() + w = getattr(dfr, "_t_demand", 0.0) - wait0 + ph["stage_wait"] += w + ph["stage_book"] += (t2 - t1) - w + if la_pred: + # This layer's demand misses have joined + # (stage returned); the predicted layers' + # misses now read in the background while this + # layer's gather and the next layers' every-token + # work compute - speculation never competes with + # demand traffic for the SSD. + la_keep = ( + ms if getattr( + self, "_kq_prestage_keepers", False) + else None) + for _dst, (_ids, _sc) in la_pred.items(): + if la_keep is not None: + dfr.prestage( + _dst, _ids, keep_mass=la_keep, + pred_scores=_sc) + else: + dfr.prestage(_dst, _ids) + if ph is not None: + ph["prestage"] += time.perf_counter() - t2 + if slots is not None: + # arena call: weights are wired GPU views for + # this scope, so lift the streaming CPU pin + # and let the fused kq kernels run + if shed_args is not None: + args = shed_args + self._kq_cpu_only = False + try: + t3 = (time.perf_counter() + if ph is not None else 0.0) + with dfr.swapped(self._kq_li): + with mx.stream(mx.gpu): + y = super().__call__( + x, mx.array(slots), + *args, **kwargs) + if (shed_mix is not None + and y.ndim == x.ndim + 1): + y = (y * shed_mix[..., None]).sum( + axis=-2) + if ph is not None: + ph["build"] += time.perf_counter() - t3 + return y + finally: + self._kq_cpu_only = True + if ( + dfr is not None + and cpu_only + and 1 < n_tokens <= _arena_split_max_tokens() + and not kwargs + and dfr.covers(self._kq_li) + and dfr.can_stage_smaller(self._kq_li) + ): + # The chunk routes more distinct experts than the + # arena has slots (stage refused above, or the chunk + # is over the stage-size gate and was never tried). + # Halve along the token axis and recurse: pieces + # whose routed union fits are served from the wired + # arena's read pool, so a turn-transition prefill or + # a wide verify batch never drops to the CPU + # page-cache gather. Bottoms out at n_tokens == 1, + # which always takes a non-split path. + ax = x.ndim - 2 + orig = ((scores_arg,) + args + if scores_arg is not None and not _fwd_scores + else args) + sliceable = ( + x.shape[ax] == n_tokens + and indices.ndim == x.ndim + and all( + isinstance(a, mx.array) + and a.ndim == x.ndim + and a.shape[ax] == n_tokens + for a in orig) + ) + if sliceable: + half = n_tokens // 2 + parts = [] + prev_split = getattr( + self, "_kq_in_split", False) + object.__setattr__( + self, "_kq_in_split", True) + try: + for sl in (slice(0, half), + slice(half, n_tokens)): + t = tuple( + [slice(None)] * ax + [sl]) + parts.append(self.__call__( + x[t], indices[t], + *[a[t] for a in orig])) + # The pieces share one precomputed + # routing. Thus the stage-time eval of a + # later piece's indices does not wait for + # an earlier piece's gather. Staging + # could overwrite (or resize away) arena + # slots that the unexecuted gather + # references. Execute each piece before + # the next piece stages. + mx.eval(parts[-1]) + finally: + object.__setattr__( + self, "_kq_in_split", prev_split) + return mx.concatenate(parts, axis=ax) + wedged = dfr is not None and dfr.wedged_at(self._kq_li) + if wedged and dfr.has_dead(self._kq_li): + # A wedged read poisoned part of this layer's file + # range: no fallback below (mmap gather, advisory + # prefetch, prefill staging) may touch a dead + # expert's bytes - rewrite the routing ids first. + mx.eval(indices) + indices = mx.array(dfr.redirect_dead( + self._kq_li, np.array(indices))) + if ( + fdr is not None + and not wedged + and small + and n_tokens >= _STREAM_PREFETCH_MIN_TOKENS + and fdr.covers(self._kq_li) + ): + # Router-aware partial staging: a short chunk routes + # to a fraction of the experts, so stage only those + # slices into the ring slot instead of the whole + # layer (see feeder.prefill_partial_call). + mx.eval(indices) + ids = np.unique(np.array(indices)).tolist() + with fdr.prefill_partial_call(self, self._kq_li, ids): + with mx.stream(mx.gpu): + return super().__call__( + x, indices, *args, **kwargs) + if ( + fdr is not None + and not wedged + and n_tokens >= _STREAM_PREFETCH_MIN_TOKENS + and fdr.covers(self._kq_li) + ): + # Feeder prefill: this layer's expert stacks are + # staged straight from the GGUF into GPU-visible + # ring slots and the GEMM runs on the GPU stream + # from the slot - the page cache never sees the + # bytes. The eval is the ring protocol's slot-free + # proof (previous layer's compute has finished); + # see feeder.py. Wedged layers skip this (and the + # whole-layer advisory below): both sweep the full + # expert range, poisoned bytes included. + mx.eval(x) + with fdr.prefill_call(self, self._kq_li): + with mx.stream(mx.gpu): + return super().__call__( + x, indices, *args, **kwargs) + if ( + pf is not None + and not wedged + and pf.enabled + and n_tokens >= _STREAM_PREFETCH_MIN_TOKENS + ): + # Streaming prefill: materialize the lazy graph up to + # this layer so the advisory window advances at + # execution pace. Build-time would fire every layer's + # advisory at once, and an over-RAM advisory storm + # evicts its own earlier reads. + mx.eval(x) + pf.on_layer(self._kq_li) + elif ( + pf is not None + and pf.enabled + and cpu_only + and env_bool("GMLX_DECODE_PREFETCH", True) + ): + # Streaming decode: the router's top-k is tiny and + # the gather needs it anyway - evaluate it now and + # pull the selected experts' slices into the page + # cache at queue depth (on_decode) instead of + # demand-faulting 16 KB clusters from inside the + # gemv. GMLX_DECODE_PREFETCH=0 disables. + mx.eval(indices) + pf.on_decode( + self._kq_li, + np.unique(np.array(indices)).tolist(), + ) + if gpu_tokens > 0 and n_tokens >= gpu_tokens and not cpu_only: + # Prefill regime: GEMM on the GPU stream, same + # zero-copy buffers. + return super().__call__(x, indices, *args, **kwargs) + with mx.stream(mx.cpu): + return super().__call__(x, indices, *args, **kwargs) + + _CPUOffload.__name__ = cls.__name__ + "_CPUOffload" + _CPU_OFFLOAD_CLASS_CACHE[cls] = sub = _CPUOffload + return sub + + n_wrapped = 0 + offloaded = 0 + n_cpu_only_codec = 0 + moe_modules: dict[int, list] = {} + for li, layer in enumerate(layers): + if n_layers is not None and li >= n_layers: + break + for m in layer.modules(): + if not isinstance(m, glu_types): + continue + if m.__class__ in _CPU_OFFLOAD_CLASS_CACHE.values(): + continue # already wrapped (idempotent) + gpu_ok = _kq_expert_gpu_ok(m) + if not gpu_ok: + n_cpu_only_codec += 1 + m.__class__ = _wrapped_class(m.__class__) + if streaming: + m._kq_cpu_only = True + object.__setattr__(m, "_kq_li", li) + if gpu_ok: + moe_modules.setdefault(li, []).append(m) + if prefetcher is not None: + object.__setattr__(m, "_kq_prefetcher", prefetcher) + elif gpu_ok: + # All-GPU auto-policy: in-RAM, the residency sweep wires the + # whole model regardless of where expert calls run, so the + # CPU hop has no memory benefit and a large decode cost + # (measured ~4-5x). Route every call to the GPU stream; an + # explicit GMLX_STREAM_GPU_TOKENS (e.g. 0) overrides. + object.__setattr__(m, "_kq_gpu_tokens_default", 1) + else: + m._kq_cpu_only = True + offloaded += sum(a.nbytes for _, a in tree_flatten(m.parameters())) + n_wrapped += 1 + if over_budget and offloaded: + # Streamed expert bytes are page cache, never wired, and must not + # tax headroom_bytes() or the admission gate and request preflight + # starve every request. Which side of the accounting they sit on + # depends on how the load materialized them: registered untracked + # (zero-copy walk, small tracked delta) they need deducting; but a + # load whose arrays landed allocator-tracked (untracked registered + # ~0) has them inside mx.get_active_memory instead, and headroom + # needs the add-back credit. tracked = total - untracked splits + # the two regimes; the credit is clamped to the expert share. + # Same 0.9 x working-set budget test as _warm_touch_pass. + key = getattr(model, "_kq_weights_key", None) + from gmlx.gen.prefill_decay import ( + note_streamed_tracked_bytes, + untracked_weight_bytes_for, + ) + tracked = max(0.0, total_bytes - untracked_weight_bytes_for(key)) + credit = min(float(offloaded), tracked) + if credit > 0: + note_streamed_tracked_bytes( + credit, key, source="experts", cap=tracked) + print( + f"[stream] headroom credits {credit / 1e9:.1f} GB of " + "allocator-tracked expert bytes as reclaimable page cache" + ) + deduct_untracked_weights(offloaded, key) + if n_cpu_only_codec: + print( + f"[stream] {n_cpu_only_codec} expert stacks use a CPU-only codec " + "(no Metal matmul kernels yet): feeder/arena staging and GPU " + "prefill routing off - every expert call runs on the CPU stream" + ) + # Non-expert weights + KV run on the default device: CPU for --stream-cpu + # (configure_stream_cpu sets the default to CPU before this call), GPU for + # --stream-experts. + base_dev = "CPU" if "cpu" in str(mx.default_device()).lower() else "GPU" + if streaming: + head = ( + f"model {total_bytes / 1e9:.0f} GB > ~{budget / 1e9:.0f} GB " + "wired budget" + if over_budget + else f"model {total_bytes / 1e9:.0f} GB, streaming forced" + ) + loadlog.info( + f"[stream] streaming: {head} - {n_wrapped} MoE layers' experts " + f"({offloaded / 1e9:.1f} GB) stay file-backed; rest of the model " + f"+ KV on {base_dev}" + ) + feeder_prefill, feeder_decode = _resolve_feeder_defaults( + feeder_prefill, feeder_decode + ) + feeder = None + dfeeder = None + room = arena = None + if streaming and prefetcher is not None and moe_modules: + from gmlx.stream.budget import kv_room_bytes + from gmlx.stream.table_stream import streamed_table_bytes + + pin = getattr(model, "_kq_weights_pin", None) + room = kv_room_bytes(gguf_path) + arena_kw = dict( + room_bytes=room.bytes, + pinned_bytes=getattr(pin, "pinned_bytes", 0), + streamable_bytes=streamed_table_bytes(model), + cast_dead_bytes=cast_dead_bytes) + arena = _decode_arena_bytes( + total_bytes, prefetcher.offsets, budget, **arena_kw) + ring = 0 + if ( + streaming + and prefetcher is not None + and moe_modules + and feeder_prefill + ): + from gmlx.stream.prefill_feeder import ( + maybe_make_prefill_feeder, + ring_bytes, + ) + + # No working-set budget (the CPU device): the ring is not judged. + reason = _prefill_ring_reason( + prefetcher.offsets, arena if budget is not None else None) + if reason: + print(f"[stream] feeder prefill unavailable ({reason}); " + "falling back to page-cache prefetch") + else: + feeder = maybe_make_prefill_feeder( + prefetcher.offsets, moe_modules) + if feeder is not None and arena is not None: + # The ring keeps its room for the process lifetime: a later + # prefill rebuilds it there, with no lend out of the arena. + ring = ring_bytes(prefetcher.offsets) + arena = _decode_arena_bytes( + total_bytes, prefetcher.offsets, budget, ring_bytes=ring, + **arena_kw) + if feeder is not None: + n_cov = sum(feeder.covers(li) for li in moe_modules) + for li, mods in moe_modules.items(): + if feeder.covers(li): + for m in mods: + object.__setattr__(m, "_kq_feeder", feeder) + object.__setattr__(model, "_kq_feeder", feeder) + cov = ( + "" if n_cov == len(moe_modules) + else f" on {n_cov}/{len(moe_modules)} layers" + ) + loadlog.info( + "[stream] feeder prefill: expert stacks staged straight " + f"from GGUF through 2 x {feeder.slot_bytes / 1e9:.1f} GB " + f"GPU-visible ring slots{cov} (--no-prefill-feeder disables)" + ) + if ( + streaming + and prefetcher is not None + and moe_modules + and feeder_decode + ): + from gmlx.stream.decode_feeder import maybe_make_decode_feeder + + from gmlx.stream.budget import ceiling_bytes + + # The ring's room is out of the arena's budget already, so the + # two never sum past the ceiling; the lend (DecodeFeeder + # .lend_for_ring) stays as the fallback for a box whose free RAM + # is gone when a later prefill rebuilds the ring. + dfeeder = maybe_make_decode_feeder( + prefetcher.offsets, moe_modules, arena, stats_verbose) + if dfeeder is not None: + dfeeder._room_bytes = room.bytes + n_cov = sum(dfeeder.covers(li) for li in moe_modules) + for li, mods in moe_modules.items(): + if dfeeder.covers(li): + for m in mods: + object.__setattr__(m, "_kq_decode_feeder", dfeeder) + object.__setattr__(model, "_kq_decode_feeder", dfeeder) + # Committed from here, not from the first decode: the arena + # wires itself the moment this model decodes, and a second + # install that sized against the unwired window would find the + # memory gone before it ever ran. + _installs.record(model, dfeeder.nominal_bytes) + _installs.record_arena(dfeeder) + if feeder is not None: + # The first decode call frees the ring (DecodeFeeder + # .ensure_wired). A later prefill pass rebuilds it in its + # own room; the rebuild asks the arena to lend only when + # the box has lost that room. + dfeeder._release_ring = feeder.release_slots + feeder._lend_hook = dfeeder.lend_for_ring + wired = ( + "fully wired at first decode" + if dfeeder._mlock_deferred + else f"{dfeeder.locked_bytes / 1e9:.1f} GB wired" + ) + cov = ( + "" if n_cov == len(moe_modules) + else f" on {n_cov}/{len(moe_modules)} layers" + ) + loadlog.info( + f"[stream] decode feeder: {dfeeder.nominal_bytes / 1e9:.1f} GB " + f"popularity-managed expert arena ({wired}){cov} " + "(--no-decode-feeder disables, GMLX_DECODE_ARENA_GB sizes)" + ) + ceiling = ceiling_bytes() or budget + expert_bytes = sum( + r[2] for rs in prefetcher.offsets.values() for r in rs) + room_how = ( + f"{room.depth} tokens x {room.width}: kv " + f"{room.kv_bytes / 1e9:.1f} + prefill " + f"{room.transient_bytes / 1e9:.1f} + reserve " + f"{room.reserve_bytes / 1e9:.1f}" + if room.priced else "flat GMLX_DECODE_KV_RESERVE_GB") + # Always visible, like the pin line: the one line a memory + # report needs. + try: + floor = _ram_floor_bytes(int(mx.device_info()["memory_size"])) + except Exception: + floor = _ram_floor_bytes(None) + print( + f"[stream] memory budget: ceiling {ceiling / 1e9:.1f} GB = " + f"every-token {(total_bytes - expert_bytes) / 1e9:.1f} + " + f"arena {dfeeder.nominal_bytes / 1e9:.1f} + ring " + f"{ring / 1e9:.1f} + kv room {room.bytes / 1e9:.1f} " + f"({room_how}) + floor {floor / 1e9:.1f}; " + "GMLX_STREAM_KV_CTX sizes the room" + ) + rate = getattr(dfeeder, "_probe_bps", 0.0) + measured = ( + f"drive reads {rate / 1e9:.1f} GB/s" + if rate else "fast-disk recipe forced") + if dfeeder._fast_disk: + loadlog.info( + f"[stream] decode feeder: {measured} - prefetch takes the" + " bandwidth (predictions evict by popularity, the barrier" + " joins only what the call routes to, prestage reads at" + " normal disk priority; --stream-fast-disk off restores" + " the conservative recipe)" + ) + elif rate: + loadlog.info( + f"[stream] decode feeder: {measured} - demand misses have" + " the bandwidth, prefetch stays out of their way" + " (--stream-fast-disk on overrides," + " GMLX_DECODE_FAST_DISK_GBPS sets the bar)" + ) + if (streaming or table_offloaded) and env_bool("GMLX_GPU_RESIDENT", True): + tskip = frozenset() + if table_offloaded: + from gmlx.stream.table_stream import streamed_table_array_ids + + tskip = streamed_table_array_ids(model) + _install_gpu_residency( + model, moe_modules, skip_ids=tskip, + include_expert_stacks=bool(table_offloaded) and not streaming) + if streaming and dfeeder is not None: + import gmlx.stream.gpu_token as gpu_token + + if gpu_token.autonomous_enabled(): + if gpu_token.route_shed_op() is None: + print( + "[stream] gpu-autonomous: requested but the installed " + "mlx_kquant has no route_shed op; falling back to " + "per-layer staging" + ) + else: + gt = gpu_token.GpuTokenState(dfeeder) + gpu_token.register_exit_stats(gt) + for li, mods in moe_modules.items(): + if dfeeder.covers(li): + for m in mods: + object.__setattr__(m, "_kq_gpu_token", gt) + object.__setattr__(model, "_kq_gpu_token", gt) + mode_note = ( + "all covered layers syncless (shed-heavy diagnostic)" + if gpu_token.autonomous_mode() == "all" + else "adaptive: layers above GMLX_AUTO_HOT_HIT go " + "syncless, the rest keep per-layer staging" + ) + loadlog.info( + "[stream] gpu-autonomous token: route_shed remaps + " + f"sheds on GPU; {mode_note}; misses prestage at " + "token boundaries (GMLX_GPU_AUTONOMOUS=1|all)" + ) + if streaming and dfeeder is not None and env_bool( + "GMLX_GPU_KEEPWARM", True): + import gmlx.stream.keepwarm as keepwarm + + keepwarm.start() + loadlog.info( + "[stream] gpu keep-warm: background heartbeat holds GPU " + "clocks between per-layer decode bursts, parked while no " + "decode is running (lossless, costs power only during " + "decode; GMLX_GPU_KEEPWARM=0 disables)" + ) + la_probe = env_bool("GMLX_DECODE_LOOKAHEAD_PROBE", False) + # Lookahead's replica router folds into the per-layer sync; whether its + # stall savings cover that tax is a per-family measurement. On + # glm_moe_dsa (GLM-5.2, 75 layers, top-8) it measured net negative + # (~40ms/tok sync for ~18ms of stalls), so those families default off. + # An explicit GMLX_DECODE_LOOKAHEAD always wins. + la_default = _lookahead_default(model) + la_prefetch = ( + env_bool("GMLX_DECODE_LOOKAHEAD", la_default) and dfeeder is not None) + if (streaming and dfeeder is not None and not la_default + and "GMLX_DECODE_LOOKAHEAD" not in os.environ): + loadlog.info( + "[stream] lookahead prestage: off by family default (replica-" + "router sync tax measured above its stall savings; " + "GMLX_DECODE_LOOKAHEAD=1 enables)" + ) + if streaming and (la_probe or la_prefetch): + from gmlx.stream.lookahead import install_lookahead + + n_la = install_lookahead( + model, layers, probe=la_probe, prefetch=la_prefetch, + stats_verbose=stats_verbose, + ) + la_depth = max(1, min(3, env_int("GMLX_DECODE_LOOKAHEAD_DEPTH", 1))) + la_what = ( + "next-layer router predictions" + if la_depth == 1 + else f"router predictions {la_depth} layers deep" + ) + if n_la and la_prefetch: + loadlog.info( + f"[stream] lookahead prestage: {la_what} pre-read arena " + f"misses on {n_la} MoE layer pairs (lossless; " + "GMLX_DECODE_LOOKAHEAD=0 disables)" + ) + if n_la and la_probe: + loadlog.info( + f"[stream] lookahead probe: recording {la_what} recall " + f"on {n_la} MoE layer pairs (lossless; table at exit)" + ) + if streaming: + # The context line printed above and the feeder lines cover the + # normal story; what remains is the fallback mechanics for + # whatever the feeders don't handle. + fallback = [] + if dfeeder is None: + fallback.append( + "decode streams expert bytes from disk through the page " + "cache (disk-bound)" + if over_budget + else "decode reads experts through the page cache on the " + "CPU stream" + ) + if feeder is None: + fallback.append( + "prefill uses sequential page-cache prefetch" + if prefetcher is not None + else "prefill demand-faults (no gguf_path)" + ) + if fallback: + loadlog.info(f"[stream] {'; '.join(fallback)}") + if not over_budget: + b = f"~{budget / 1e9:.0f} GB" if budget else "unknown" + print( + f"[stream] --stream-cpu streams experts even though the " + f"{total_bytes / 1e9:.0f} GB model fits the wired budget " + f"({b}) - omit --stream-cpu for the faster all-GPU path on " + "a model that fits" + ) + else: + gpu_tokens = _stream_gpu_tokens(1) + if gpu_tokens == 1: + staging = ( + "model fits the wired budget - decode auto-routed to the " + "GPU stream (GMLX_STREAM_GPU_TOKENS=0 forces CPU decode)" + ) + elif gpu_tokens > 0: + staging = f"prefill calls >={gpu_tokens} tokens routed to GPU" + else: + staging = "GPU prefill routing disabled" + if table_offloaded: + # Table-only mode: the experts are GPU-resident and wired (the + # streamed table made room); "file-backed" would be wrong. + loadlog.info( + f"[stream] routed experts resident on GPU across " + f"{n_wrapped} layers ({offloaded / 1e9:.1f} GB wired; " + f"{staging})" + ) + else: + loadlog.info( + f"[stream] routed experts -> CPU stream on {n_wrapped} " + f"layers ({offloaded / 1e9:.1f} GB stays file-backed; rest " + f"of the model + KV on {base_dev}; {staging})" + ) + return n_wrapped, offloaded + + +def _resolve_prefill_step(model, requested: int | None) -> tuple[int | None, bool]: + """Pick the prefill chunk width: an explicit request always wins; a + streaming-mode model defaults to ``_STREAMING_PREFILL_STEP``, or to its + model_type's narrower entry; everything else keeps mlx-lm's own + default. Returns ``(step_or_none, defaulted)``.""" + if requested is not None or not moe_streaming_active(model): + return requested, False + mt = getattr(model, "model_type", None) or getattr( + getattr(model, "args", None), "model_type", None) + return _STREAMING_PREFILL_STEP_BY_MODEL_TYPE.get( + mt, _STREAMING_PREFILL_STEP), True + + +def install_moe_experts_override(model, k: int) -> int: + """Experiment, lossy: route every token to ``k`` experts instead of the + trained top-k, on MoE blocks whose experts ``install_expert_streaming`` + wrapped (and only those - the knob exists to probe how router fan-out + shapes offloaded prefill/decode traffic, not as a general sampler). + + The override rewrites the router's own top-k attribute (``top_k`` / + ``num_experts_per_tok`` on the block, and on a DeepSeek-style gate + submodule when present - named ``gate``, or ``router`` on hy_v3), so + expert selection and the arch's weight renormalization run unchanged + at the new k. Outputs differ from the trained model by design; parity + gates will fail. Returns the number of MoE blocks overridden; raises + on k < 1 or k > the expert count. + """ + if k < 1: + raise ValueError(f"MoE top-k override must be >= 1, got {k}") + layers = getattr(model, "layers", None) + if layers is None: + layers = model.model.layers + overridden = 0 + trained_k = None + for layer in layers: + for owner in layer.modules(): + glu = None + for child in owner.children().values(): + candidates = child if isinstance(child, (list, tuple)) else [child] + for c in candidates: + if type(c).__name__.endswith("_CPUOffload"): + glu = c + break + if glu is not None: + break + if glu is None: + continue + n_experts = _switch_num_experts(glu) + if n_experts and k > n_experts: + raise ValueError( + f"MoE top-k override {k} exceeds the {n_experts}-expert " + "stack on an offloaded layer" + ) + hit = False + for target in ( + owner, + getattr(owner, "gate", None), + getattr(owner, "router", None), + ): + if target is None: + continue + for attr in ("top_k", "num_experts_per_tok"): + current = getattr(target, attr, None) + if isinstance(current, int): + if trained_k is None: + trained_k = current + setattr(target, attr, k) + hit = True + if hit: + overridden += 1 + if overridden: + print( + f"[stream] MoE top-k override: {trained_k}->{k} experts/token " + f"on {overridden} offloaded MoE layers (lossy - outputs differ " + "from the trained router)" + ) + else: + print( + "[stream] MoE top-k override found no offloaded MoE block " + "with a router top-k attribute - no effect" + ) + return overridden diff --git a/gmlx/stream/table_stream.py b/gmlx/stream/table_stream.py index 5a753882..b615f458 100644 --- a/gmlx/stream/table_stream.py +++ b/gmlx/stream/table_stream.py @@ -194,7 +194,7 @@ def install_table_streaming(model) -> tuple[int, list[str]]: neutralize the wired-limit sweep. Returns ``(offloaded_bytes, gguf_names)``; the caller (the loader's selection ladder) owns the residency deduction and the decision log line. Idempotent.""" - from gmlx.load.loader import _neutralize_wired_limit_sweep + from gmlx.stream.wired_limit import _neutralize_wired_limit_sweep offloaded = 0 names: list[str] = [] diff --git a/gmlx/stream/wired_limit.py b/gmlx/stream/wired_limit.py new file mode 100644 index 00000000..dfddee68 --- /dev/null +++ b/gmlx/stream/wired_limit.py @@ -0,0 +1,181 @@ +"""Wired-limit sweep control and CPU-stream device setup for streaming +(over-wired-budget) models.""" +from __future__ import annotations + +import os + +import mlx.core as mx + + +def _neutralize_wired_limit_sweep(): + """Pin the MLX wired limit at its default for the rest of the process. + + mlx-lm wraps generation in a context manager that raises the wired limit + to the device's max recommended working set. MLX services that by adding + every live buffer - file-backed zero-copy weight views included - to its + Metal residency set, which wires them all. For a model larger than the + wired budget that sweep exhausts wired memory within seconds of the + first GPU command (hard-panic territory). There is no per-buffer + opt-out, so streaming mode no-ops ``mx.set_wired_limit`` instead: Metal + then wires only what GPU work actually references (the every-token + layers + KV), and + expert pages stay plain evictable page cache. Covers every caller + (generate, server batch path, trainer) since all resolve the function + through ``mx.`` at call time. Idempotent. + + mlx-lm's ``wired_limit()`` context manager also prints a per-generation + large-model warning sized against the limit this function just pinned - + meaningless in streaming mode, and noisy (once per chat turn). Swap it + for a quiet context that keeps the exit synchronize (the original syncs + the generation stream before restoring the limit; callers may rely on + that barrier at generator teardown). NB: patched via importlib - + ``import mlx_lm.generate`` binds the function mlx_lm re-exports in + ``__init__``, not the submodule. + """ + if getattr(mx.set_wired_limit, "_kq_no_sweep", False): + return + # A generator that started before this call left the limit raised, and + # its exit restore is a no-op from here on. Lower it now: with it up, + # the next streaming load's walk wires the file's resident pages as it + # creates the views, and every command buffer in the process fails + # with a Metal out-of-memory error. + try: + prev = mx.set_wired_limit(0) + except Exception: + prev = 0 + if prev: + print(f"[stream] wired limit lowered from {prev / 1e9:.1f} GB to 0: " + "a raised limit wires every live buffer, zero-copy views included") + + def _no_sweep(*_a, **_k): + return 0 + + _no_sweep._kq_no_sweep = True + mx.set_wired_limit = _no_sweep + + import contextlib + import importlib + + @contextlib.contextmanager + def _quiet_wired_limit(model, streams=None): + try: + yield + finally: + if streams is not None: + for s in streams: + mx.synchronize(s) + else: + mx.synchronize() + + _quiet_wired_limit._kq_no_sweep = True + for mod_name in ("mlx_lm.generate", "mlx_lm.utils"): + try: + mod = importlib.import_module(mod_name) + except ImportError: + continue + if hasattr(mod, "wired_limit"): + mod.wired_limit = _quiet_wired_limit + + +def _install_wired_limit_warn_once(): + """Cap mlx-lm's large-model warning at one print per process. + + ``stream_generate`` enters mlx-lm's ``wired_limit()`` context on every + call - at least once per chat turn - and on entry the context prints its + near-the-wired-budget warning unconditionally, so a resident model just + over the 0.9x threshold re-warns every turn. There is no seam around the + print, so swap in a re-implementation with identical wiring behavior + (raise the limit, synchronize on exit, restore) that warns only the + first time. + + Installed at the end of every ``load_model`` (the resident path); the + streaming / CPU replacements above are stricter (they drop the sweep + entirely), so this never overwrites them - and they overwrite this when + they engage, which is always after load. Idempotent. NB: patched via + importlib - ``import mlx_lm.generate`` binds the function mlx_lm + re-exports in ``__init__``, not the submodule. + """ + import contextlib + import importlib + + from mlx.utils import tree_reduce + + state = {"warned": False} + + @contextlib.contextmanager + def _warn_once_wired_limit(model, streams=None): + if not mx.metal.is_available(): + yield + return + model_bytes = tree_reduce( + lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, + model, 0) + max_rec_size = mx.device_info()["max_recommended_working_set_size"] + if model_bytes > 0.9 * max_rec_size and not state["warned"]: + state["warned"] = True + model_mb = model_bytes // 2**20 + max_rec_mb = max_rec_size // 2**20 + print( + f"[WARNING] Generating with a model that requires {model_mb} " + f"MB which is close to the maximum recommended size of " + f"{max_rec_mb} MB. This can be slow. See the documentation " + "for possible work-arounds: " + "https://github.com/ml-explore/mlx-lm/tree/main#large-models" + ) + old_limit = mx.set_wired_limit(max_rec_size) + try: + yield + finally: + if streams is not None: + for s in streams: + mx.synchronize(s) + else: + mx.synchronize() + mx.set_wired_limit(old_limit) + + _warn_once_wired_limit._kq_warn_once = True + for mod_name in ("mlx_lm.generate", "mlx_lm.utils"): + try: + mod = importlib.import_module(mod_name) + except ImportError: + continue + fn = getattr(mod, "wired_limit", None) + if fn is None or getattr(fn, "_kq_no_sweep", False) or \ + getattr(fn, "_kq_warn_once", False): + continue + mod.wired_limit = _warn_once_wired_limit + + +def configure_cpu_device(): + """Run everything on the CPU device (``--stream-cpu``): mmap-streamed weights. + + Besides setting the default device this (a) keeps the graph + single-device - the fused-GDN runtime patch dispatches Metal kernels + regardless of the default device - and (b) no-ops mlx-lm's + ``wired_limit`` context: it reads + ``device_info()["max_recommended_working_set_size"]``, absent on the + CPU device, and wiring is meaningless on CPU. NB: patched via importlib + - ``import mlx_lm.generate`` binds the function mlx_lm re-exports in + ``__init__``, not the submodule. + """ + import contextlib + import importlib + + mx.set_default_device(mx.cpu) + os.environ.setdefault("GMLX_FUSED_GDN", "0") + + @contextlib.contextmanager + def _wired_noop(model, streams=None): + yield + + # No sweep at all on CPU; the marker keeps a later load_model's + # warn-once variant (_install_wired_limit_warn_once) from clobbering it. + _wired_noop._kq_no_sweep = True + for mod_name in ("mlx_lm.generate", "mlx_lm.utils"): + try: + mod = importlib.import_module(mod_name) + except ImportError: + continue + if hasattr(mod, "wired_limit"): + mod.wired_limit = _wired_noop + print("[device] cpu (mmap-streamed weights; fused-GDN Metal patch off)") diff --git a/gmlx/tui/chat.py b/gmlx/tui/chat.py index ece538d9..0299a5f3 100644 --- a/gmlx/tui/chat.py +++ b/gmlx/tui/chat.py @@ -2981,7 +2981,7 @@ def _post_text_load(): _apply_cli_adapter(args, b.model, b.config) _apply_placement(args, b.model) - from gmlx.load.loader import _resolve_prefill_step + from gmlx.stream.expert_streaming import _resolve_prefill_step step, defaulted = _resolve_prefill_step( b.model, kv_kwargs.get("prefill_step_size") diff --git a/tests/commands/test_cli_run_flags.py b/tests/commands/test_cli_run_flags.py index 100bef2d..44b9b754 100644 --- a/tests/commands/test_cli_run_flags.py +++ b/tests/commands/test_cli_run_flags.py @@ -137,16 +137,17 @@ def report_stubs(monkeypatch): import gmlx.load.tokenizer as tok_mod import gmlx.gen.generation as generation import gmlx.load.loader as loader + import gmlx.load.wire as wire monkeypatch.setattr(preflight_mod, "preflight", lambda path, arch=None: None) monkeypatch.setattr( - loader, "load_gguf_wire_bytes", + wire, "load_gguf_wire_bytes", lambda path, zero_copy=True: ({"t": 0}, {"t": {}}, "gemma4", {}, {})) import gmlx.load.gguf_meta as gguf_meta monkeypatch.setattr(gguf_meta, "read_int", lambda meta, key: None) monkeypatch.setattr(gguf_meta, "first_nonzero_int", lambda meta, key: None) monkeypatch.setattr( - loader, "remap_arrays", + wire, "remap_arrays", lambda arrays, kq, arch, no_remap=False, n_head=None, n_head_kv=None: (arrays, {}, {})) monkeypatch.setattr(loader, "print_inventory", diff --git a/tests/e2e/run_stream_e2e.py b/tests/e2e/run_stream_e2e.py index 04733563..c393d3cc 100644 --- a/tests/e2e/run_stream_e2e.py +++ b/tests/e2e/run_stream_e2e.py @@ -517,7 +517,7 @@ def phase_warmth(a) -> None: recl1 = kernel_reclaimable() arena1 = (s2.get("memory") or {}).get("arena_bytes") or 0 need = 0.25 * nominal - from gmlx.load.loader import _ram_floor_bytes + from gmlx.stream.budget import _ram_floor_bytes gate = need + _ram_floor_bytes(ram_bytes()) + (floor or 0) expected = recl1 is not None and recl1 >= gate print(f"warmth: released at +{time.monotonic() - t_rel:.0f}s; arena {gb(arena1)}, " diff --git a/tests/gen/test_long_context.py b/tests/gen/test_long_context.py index 31bee564..9bf84d10 100644 --- a/tests/gen/test_long_context.py +++ b/tests/gen/test_long_context.py @@ -147,7 +147,7 @@ def _load(path): model, config, tok = load_model(path, verbose=False) if _over_wired_budget(model): - from gmlx.load.loader import install_expert_streaming + from gmlx.stream.expert_streaming import install_expert_streaming install_expert_streaming(model, gguf_path=path) return model, config, tok diff --git a/tests/load/test_loader_guards.py b/tests/load/test_loader_guards.py index 248045b7..a24ad56e 100644 --- a/tests/load/test_loader_guards.py +++ b/tests/load/test_loader_guards.py @@ -8,7 +8,7 @@ import mlx.core as mx import pytest -from gmlx.load.loader import _own, _RemapDict +from gmlx.load.wire import _own, _RemapDict def test_own_bf16_copies_via_f32(): @@ -39,7 +39,7 @@ def test_remap_result_allows_native_fp_repack(): # a phantom "remap collision". import numpy as np - from gmlx.load.loader import remap_arrays + from gmlx.load.wire import remap_arrays from gmlx.load.native_fp import repack_native_fp_weights raw = mx.array(np.zeros((4, 2 * 17), dtype=np.uint8)) # 2 mxfp4 blocks/row diff --git a/tests/load/test_nemotron_lightning.py b/tests/load/test_nemotron_lightning.py index 9c7b5c9a..1bb47097 100644 --- a/tests/load/test_nemotron_lightning.py +++ b/tests/load/test_nemotron_lightning.py @@ -63,7 +63,7 @@ def test_pattern_no_nextn_unchanged(): def test_strip_nextn_trunk_overflow(): - from gmlx.load.loader import strip_nextn_trunk_overflow + from gmlx.load.wire import strip_nextn_trunk_overflow meta = {f"{ARCH}.nextn_predict_layers": 1, f"{ARCH}.block_count": 5} w = { @@ -80,7 +80,7 @@ def test_strip_nextn_trunk_overflow(): def test_strip_noop_other_arch(): - from gmlx.load.loader import strip_nextn_trunk_overflow + from gmlx.load.wire import strip_nextn_trunk_overflow w = {"backbone.layers.4.mixer.in_proj.weight": 1} assert strip_nextn_trunk_overflow(w, {}, {}, "qwen3") == 0 diff --git a/tests/load/test_offload.py b/tests/load/test_offload.py index 347a0159..ab9f6f3a 100644 --- a/tests/load/test_offload.py +++ b/tests/load/test_offload.py @@ -13,10 +13,10 @@ from mlx_lm.utils import _get_classes # noqa: E402 from gmlx.load.config_synth import synthesize_config # noqa: E402 -from gmlx.load.loader import ( # noqa: E402 +from gmlx.load.loader import moe_streaming_active # noqa: E402 +from gmlx.stream.expert_streaming import ( # noqa: E402 _resolve_prefill_step, install_expert_streaming, - moe_streaming_active, ) from test_config_synth import _qwen3next_meta # noqa: E402 @@ -364,7 +364,7 @@ def test_moe_experts_override_targets_offloaded_router(monkeypatch): raise; a model with no offloaded experts is a no-op.""" import mlx.nn as nn - from gmlx.load.loader import install_moe_experts_override + from gmlx.stream.expert_streaming import install_moe_experts_override class _Gate(nn.Module): # DeepSeek-style: top_k lives on the gate def __init__(self): diff --git a/tests/load/test_remap.py b/tests/load/test_remap.py index 176cb1a5..2067c4e7 100644 --- a/tests/load/test_remap.py +++ b/tests/load/test_remap.py @@ -1263,7 +1263,7 @@ def test_kimi_k3_kda_conv_weight_array_transform(): # Both wire layouts reshape to mlx Conv1d (d_inner, d_conv, 1) exactly. import mlx.core as mx import numpy as np - from gmlx.load.loader import remap_arrays + from gmlx.load.wire import remap_arrays d_inner, d_conv = 6, 4 base = np.arange(d_inner * d_conv, dtype=np.float32) arrays = { diff --git a/tests/load/test_vlm_deepseek4v.py b/tests/load/test_vlm_deepseek4v.py index afc67594..e371511b 100644 --- a/tests/load/test_vlm_deepseek4v.py +++ b/tests/load/test_vlm_deepseek4v.py @@ -264,7 +264,7 @@ def test_image_processor_geometry_matches_reference_rules(): @pytest.mark.skipif(not os.path.exists(MMPROJ), reason="local mmproj absent") def test_real_mmproj_remaps_onto_the_tower(): - from gmlx.load.loader import load_gguf_wire_bytes + from gmlx.load.wire import load_gguf_wire_bytes arrays, codecs, _arch, mm_meta, _shapes = load_gguf_wire_bytes( MMPROJ, zero_copy=True, expect_quant=False) out, skipped, kq = remap_vision_arrays(arrays, "deepseek_v4_vl") diff --git a/tests/load/test_vlm_spec_target_swap.py b/tests/load/test_vlm_spec_target_swap.py index e7888156..9eeed143 100644 --- a/tests/load/test_vlm_spec_target_swap.py +++ b/tests/load/test_vlm_spec_target_swap.py @@ -22,7 +22,7 @@ import gmlx.models.gemma4.owned as gemma4_owned import gmlx.models.qwen35.owned as qwen35_owned -from gmlx.load.loader import _spec_hook_key, _vlm_spec_language_model +from gmlx.load.mtp_target import _spec_hook_key, _vlm_spec_language_model from gmlx.load.vlm import _swap_spec_language_model diff --git a/tests/load/test_wired_limit_patch.py b/tests/load/test_wired_limit_patch.py index 19ab2e93..802ec195 100644 --- a/tests/load/test_wired_limit_patch.py +++ b/tests/load/test_wired_limit_patch.py @@ -13,7 +13,7 @@ import mlx.core as mx import pytest -from gmlx.load.loader import _install_wired_limit_warn_once +from gmlx.stream.wired_limit import _install_wired_limit_warn_once # `import mlx_lm.generate` binds the function mlx_lm re-exports in # __init__, not the submodule - same trap the loader patches around. @@ -117,7 +117,7 @@ def test_neutralize_lowers_a_raised_wired_limit(monkeypatch, capsys): # lowers the real limit itself, or the next walk wires the file's pages. import mlx.core as mx - from gmlx.load.loader import _neutralize_wired_limit_sweep + from gmlx.stream.wired_limit import _neutralize_wired_limit_sweep calls = [] diff --git a/tests/models/test_deepseek_v4_mtp.py b/tests/models/test_deepseek_v4_mtp.py index f5b90a68..2e4e98c3 100644 --- a/tests/models/test_deepseek_v4_mtp.py +++ b/tests/models/test_deepseek_v4_mtp.py @@ -59,7 +59,7 @@ def _randomize_zero_params(mod) -> None: def _build_target(cfg): - from gmlx.load.loader import MTPTextTarget + from gmlx.load.mtp_target import MTPTextTarget lm = DeepseekV4SpecLM(v4.ModelArgs.from_dict(cfg)) mx.eval(lm.parameters()) @@ -477,7 +477,7 @@ def test_default_drafter_block_size_is_4(): def test_speclm_hooks_match_loader_contract(): - from gmlx.load.loader import _MTP_TARGET_HOOKS_BY_TYPE + from gmlx.load.mtp_target import _MTP_TARGET_HOOKS_BY_TYPE hooks = _MTP_TARGET_HOOKS_BY_TYPE["deepseek_v4"] for hook in hooks: diff --git a/tests/models/test_gemma4_owned.py b/tests/models/test_gemma4_owned.py index 984e5727..cd59c556 100644 --- a/tests/models/test_gemma4_owned.py +++ b/tests/models/test_gemma4_owned.py @@ -78,18 +78,18 @@ def _ids(*toks): def test_loader_selects_owned_by_default(monkeypatch): - import gmlx.load.loader as loader + import gmlx.load.mtp_target as mtp_target monkeypatch.delenv("GMLX_GEMMA_OWNED", raising=False) - cls, build = loader._mtp_target_classes("gemma4_text") + cls, build = mtp_target._mtp_target_classes("gemma4_text") assert cls is OwnedGemma4LanguageModel def test_loader_env_reverts_to_stock(monkeypatch): - import gmlx.load.loader as loader + import gmlx.load.mtp_target as mtp_target monkeypatch.setenv("GMLX_GEMMA_OWNED", "0") - cls, build = loader._mtp_target_classes("gemma4_text") + cls, build = mtp_target._mtp_target_classes("gemma4_text") assert cls is _G.LanguageModel diff --git a/tests/models/test_glm5_next_mtp.py b/tests/models/test_glm5_next_mtp.py index 74d2f211..b4d9a88b 100644 --- a/tests/models/test_glm5_next_mtp.py +++ b/tests/models/test_glm5_next_mtp.py @@ -174,7 +174,7 @@ def test_logits_from_hidden_matches_forward(): def _wrap_target(lm): - from gmlx.load.loader import MTPTextTarget + from gmlx.load.mtp_target import MTPTextTarget return MTPTextTarget(lm, {"model_type": "glm5_next"}) @@ -295,7 +295,7 @@ def test_mtp_greedy_identity_accept_path(): def test_speclm_hooks_match_loader_contract(): - from gmlx.load.loader import _MTP_TARGET_HOOKS_BY_TYPE, _mtp_target_classes + from gmlx.load.mtp_target import _MTP_TARGET_HOOKS_BY_TYPE, _mtp_target_classes for hook in _MTP_TARGET_HOOKS_BY_TYPE["glm5_next"]: assert callable(getattr(Glm5NextSpecLM, hook, None)), hook @@ -372,7 +372,7 @@ def test_mtp_remap_covers_closed_tensor_set(): params, covering the full drafter tree (both directions closed).""" from mlx.utils import tree_flatten - from gmlx.load.loader import remap_mtp_arrays + from gmlx.load.wire import remap_mtp_arrays args = _tiny_args() drafter = _build_drafter(args) diff --git a/tests/models/test_hy_v3_mtp.py b/tests/models/test_hy_v3_mtp.py index a1686593..12f24b92 100644 --- a/tests/models/test_hy_v3_mtp.py +++ b/tests/models/test_hy_v3_mtp.py @@ -46,7 +46,7 @@ def _tiny_config() -> dict: def _build_target(cfg): - from gmlx.load.loader import MTPTextTarget + from gmlx.load.mtp_target import MTPTextTarget lm = HyV3SpecLM(ModelArgs.from_dict(cfg)) mx.eval(lm.parameters()) @@ -186,7 +186,7 @@ def test_mtp_greedy_identity_accept_path(block): def test_speclm_hooks_match_loader_contract(): - from gmlx.load.loader import _MTP_TARGET_HOOKS_BY_TYPE + from gmlx.load.mtp_target import _MTP_TARGET_HOOKS_BY_TYPE for hook in _MTP_TARGET_HOOKS_BY_TYPE["hy_v3"]: assert callable(getattr(HyV3SpecLM, hook, None)), hook @@ -236,7 +236,7 @@ def test_mtp_remap_covers_closed_tensor_set(): hy3-1M-MTP-IQ2_M.gguf, with tiny shapes.""" from mlx.utils import tree_flatten - from gmlx.load.loader import remap_mtp_arrays + from gmlx.load.wire import remap_mtp_arrays cfg = _tiny_config() drafter = _build_drafter(cfg) diff --git a/tests/models/test_hy_v4_model.py b/tests/models/test_hy_v4_model.py index b153a237..7623223f 100644 --- a/tests/models/test_hy_v4_model.py +++ b/tests/models/test_hy_v4_model.py @@ -292,14 +292,15 @@ def test_streaming_prefill_step_is_narrowed_for_hy_v4(): def test_prefill_step_resolves_from_the_model_type(monkeypatch): from gmlx.load import loader + from gmlx.stream import expert_streaming - monkeypatch.setattr(loader, "moe_streaming_active", lambda _m: True) + monkeypatch.setattr(expert_streaming, "moe_streaming_active", lambda _m: True) model = _model() - step, defaulted = loader._resolve_prefill_step(model, None) + step, defaulted = expert_streaming._resolve_prefill_step(model, None) assert defaulted is True assert step == loader._STREAMING_PREFILL_STEP_BY_MODEL_TYPE["hy_v4"] # An explicit request always wins. - assert loader._resolve_prefill_step(model, 1024) == (1024, False) + assert expert_streaming._resolve_prefill_step(model, 1024) == (1024, False) # --- forward integrity ------------------------------------------------------- diff --git a/tests/models/test_qwen35_owned.py b/tests/models/test_qwen35_owned.py index 639aa068..2802e922 100644 --- a/tests/models/test_qwen35_owned.py +++ b/tests/models/test_qwen35_owned.py @@ -28,7 +28,7 @@ _patch_gated_delta_tiled_v, _patch_mlxvlm_gated_delta_tiled_v, ) -from gmlx.load.loader import _mtp_target_classes +from gmlx.load.mtp_target import _mtp_target_classes ATOL = 2e-3 # differing-route bound (shortcut removed / kernel path) TIGHT_ATOL = 1e-5 # same-ops bound diff --git a/tests/models/test_qwen4exp_mtp.py b/tests/models/test_qwen4exp_mtp.py index 1357752d..e0365001 100644 --- a/tests/models/test_qwen4exp_mtp.py +++ b/tests/models/test_qwen4exp_mtp.py @@ -176,7 +176,7 @@ def test_remap_strips_prefix_and_threads_codecs(): def test_arch_table_and_loader_rows(): import gmlx.load.arch_table as arch_table - from gmlx.load.loader import _MTP_TARGET_HOOKS_BY_TYPE, _mtp_target_classes + from gmlx.load.mtp_target import _MTP_TARGET_HOOKS_BY_TYPE, _mtp_target_classes from gmlx.spec.mtp_load import _assistant_kind assert arch_table.drafter_arches("qwen4_exp") == (MTP_ARCH,) diff --git a/tests/serve/test_residency.py b/tests/serve/test_residency.py index a5e38f16..d2761799 100644 --- a/tests/serve/test_residency.py +++ b/tests/serve/test_residency.py @@ -316,7 +316,7 @@ def test_streaming_build_lowers_the_wired_limit_and_credits_the_walk(monkeypatch import pytest import gmlx.gen.prefill_decay as pd - import gmlx.load.loader as loader + import gmlx.stream.wired_limit as wired_limit import gmlx.serve.capacity as cap monkeypatch.setattr(pd, "_STREAMED_TRACKED", {}) @@ -332,7 +332,7 @@ def test_streaming_build_lowers_the_wired_limit_and_credits_the_walk(monkeypatch monkeypatch.setattr(cap, "preload_gate", lambda *a, **k: None) monkeypatch.setenv("GMLX_DECODE_ARENA_GB", "60") seen = [] - monkeypatch.setattr(loader, "_neutralize_wired_limit_sweep", + monkeypatch.setattr(wired_limit, "_neutralize_wired_limit_sweep", lambda: seen.append("neutralized")) proxy = _RuntimeProxy(_FakeOriginal()) diff --git a/tests/serve/test_serve_vlm.py b/tests/serve/test_serve_vlm.py index 26fc1cfe..ea8a8271 100644 --- a/tests/serve/test_serve_vlm.py +++ b/tests/serve/test_serve_vlm.py @@ -120,12 +120,12 @@ def spy_placement(monkeypatch): """Record the placement/lever installers instead of running them, so the branch is testable without a streamed model. Each entry is ``(target, args, kwargs)`` keyed by installer name.""" - import gmlx.load.loader as loader + import gmlx.stream.expert_streaming as expert_streaming import gmlx.stream.moe_experts as moe_experts seen: dict[str, tuple] = {} - def _spy(name, module=loader): + def _spy(name, module=expert_streaming): def _fake(target, *args, **kwargs): seen[name] = (target, args, kwargs) monkeypatch.setattr(module, name, _fake, raising=True) diff --git a/tests/spec/test_mtp.py b/tests/spec/test_mtp.py index f5e56c0f..8f1eed06 100644 --- a/tests/spec/test_mtp.py +++ b/tests/spec/test_mtp.py @@ -21,6 +21,8 @@ import gmlx.upstream.gdn_patches as gdn_patches # noqa: E402 import gmlx.load.loader as loader # noqa: E402 +import gmlx.load.mtp_target as mtp_target # noqa: E402 +import gmlx.load.wire as wire # noqa: E402 @pytest.fixture @@ -35,16 +37,16 @@ def cpu_device(): # target hook contract (version tripwire) @pytest.mark.parametrize("model_type", ["qwen3_5", "qwen3_5_moe"]) def test_mtp_target_exposes_speculative_hooks(model_type): - LanguageModel, _build = loader._mtp_target_classes(model_type) - missing = [h for h in loader._MTP_TARGET_HOOKS if not hasattr(LanguageModel, h)] + LanguageModel, _build = mtp_target._mtp_target_classes(model_type) + missing = [h for h in mtp_target._MTP_TARGET_HOOKS if not hasattr(LanguageModel, h)] assert not missing, ( f"mlx-vlm {model_type} LanguageModel missing hooks {missing}; " - f"the MTP engine needs all of {loader._MTP_TARGET_HOOKS}") + f"the MTP engine needs all of {mtp_target._MTP_TARGET_HOOKS}") def test_mtp_target_resolver_rejects_unknown_arch(): with pytest.raises(NotImplementedError): - loader._mtp_target_classes("llama") + mtp_target._mtp_target_classes("llama") # seam 3: mlx-vlm gated_delta tiled-V patch @@ -116,7 +118,7 @@ def _find_mtp_gguf(gguf_dir): except Exception: continue try: - _a, _k, _am, meta, shapes = loader.load_gguf_wire_bytes( + _a, _k, _am, meta, shapes = wire.load_gguf_wire_bytes( str(path), shards=pf.shards) cfg = synthesize_config(meta, shapes) except Exception: @@ -136,7 +138,7 @@ def test_mtp_drafter_remap_full_coverage(gguf_dir, cpu_device): pytest.skip("no native-head MTP GGUF found in KQUANT_TEST_GGUF_DIR") pf = preflight(path) - arrays, kqm, _am, meta, shapes = loader.load_gguf_wire_bytes( + arrays, kqm, _am, meta, shapes = wire.load_gguf_wire_bytes( path, shards=pf.shards) config = synthesize_config(meta, shapes) @@ -152,7 +154,7 @@ def test_mtp_drafter_remap_full_coverage(gguf_dir, cpu_device): n_head = gguf_meta.read_int(meta, f"{arch}.attention.head_count") n_head_kv = gguf_meta.first_nonzero_int( meta, f"{arch}.attention.head_count_kv") - d_w, d_m, _stats = loader.remap_mtp_arrays( + d_w, d_m, _stats = wire.remap_mtp_arrays( arrays, kqm, arch, first_mtp_block=int(config["num_hidden_layers"]), num_mtp_layers=int(config.get("mtp_num_hidden_layers", 1)), diff --git a/tests/spec/test_mtp_dispatch.py b/tests/spec/test_mtp_dispatch.py index 941491cb..2ce0e76d 100644 --- a/tests/spec/test_mtp_dispatch.py +++ b/tests/spec/test_mtp_dispatch.py @@ -15,7 +15,7 @@ import gmlx.load.arch_table as arch_table import gmlx.models.qwen35.owned as qwen35_owned import gmlx.spec.mtp_load as mtp_load -from gmlx.load.loader import ( +from gmlx.load.mtp_target import ( _MTP_TARGET_HOOKS, _MTP_TARGET_HOOKS_BY_TYPE, _spec_hook_key, diff --git a/tests/stream/test_budget.py b/tests/stream/test_budget.py index 4d449925..d8fad4e5 100644 --- a/tests/stream/test_budget.py +++ b/tests/stream/test_budget.py @@ -128,10 +128,9 @@ def test_transient_bytes_follows_the_decay_cap(monkeypatch): def test_reclaimable_ram_bytes_prefers_the_kernel_counters(monkeypatch): import gmlx.serve.kernel_vm as kv - import gmlx.load.loader as loader monkeypatch.setattr(kv, "reclaimable_bytes", lambda: 12.5e9) - monkeypatch.setattr(loader, "_available_ram_bytes", lambda include_inactive=True: 3) + monkeypatch.setattr(budget, "_available_ram_bytes", lambda include_inactive=True: 3) assert budget.reclaimable_ram_bytes() == int(12.5e9) monkeypatch.setattr(kv, "reclaimable_bytes", lambda: None) assert budget.reclaimable_ram_bytes() == 3 @@ -148,7 +147,7 @@ def test_available_ram_reads_the_mach_counters_without_a_spawn(monkeypatch): arena copies the arena, so the serve process must not fork here.""" import subprocess - from gmlx.load.loader import _available_ram_bytes + from gmlx.stream.budget import _available_ram_bytes from gmlx.serve import kernel_vm if kernel_vm.snapshot() is None: diff --git a/tests/stream/test_decode_feeder.py b/tests/stream/test_decode_feeder.py index bed45769..98a1b2a4 100644 --- a/tests/stream/test_decode_feeder.py +++ b/tests/stream/test_decode_feeder.py @@ -14,13 +14,11 @@ import mlx.core as mx from mlx_lm.models.switch_layers import SwitchGLU -import gmlx.load.loader +import gmlx.stream.budget import gmlx.serve.kernel_vm -from gmlx.load.loader import ( - _decode_arena_bytes, - _resolve_feeder_defaults, - install_expert_streaming, -) +from gmlx.load.loader import _resolve_feeder_defaults +from gmlx.stream.budget import _decode_arena_bytes +from gmlx.stream.expert_streaming import install_expert_streaming _KINDS = ("gate", "up", "down") _E = 4 # experts per layer @@ -235,7 +233,7 @@ def test_regrow_leaves_both_floors_behind(monkeypatch, tmp_path): _pressure_setup(monkeypatch, level, regrow_polls=1) import gmlx.stream.budget as budget - monkeypatch.setattr(gmlx.load.loader, "_ram_floor_bytes", lambda ram: 10 << 30) + monkeypatch.setattr(gmlx.stream.budget, "_ram_floor_bytes", lambda ram: 10 << 30) monkeypatch.setattr(budget, "kernel_floor_bytes", lambda: float(4 << 30)) avail = {"v": 0} monkeypatch.setattr(budget, "reclaimable_ram_bytes", lambda: avail["v"]) @@ -262,7 +260,7 @@ def test_pressure_regrow_after_sustained_normal(monkeypatch, tmp_path): # The regrow reads the governor's reclaimable measure (file-backed # pages included), never the free-pages-only set. - monkeypatch.setattr(gmlx.load.loader, "_available_ram_bytes", + monkeypatch.setattr(gmlx.stream.budget, "_available_ram_bytes", lambda include_inactive=True: 0) monkeypatch.setattr(budget, "reclaimable_ram_bytes", lambda: avail["v"]) monkeypatch.setattr(budget, "kernel_floor_bytes", lambda: 4e9) @@ -315,7 +313,7 @@ def test_arena_budget_math(monkeypatch): monkeypatch.delenv("GMLX_DECODE_PAGECACHE_GB", raising=False) monkeypatch.delenv("GMLX_DECODE_ARENA_FORCE", raising=False) monkeypatch.setattr( - gmlx.load.loader, "_available_ram_bytes", lambda: None + gmlx.stream.budget, "_available_ram_bytes", lambda: None ) # available-RAM ceiling out of the way for the deterministic cases monkeypatch.setattr( mx, "device_info", lambda: {"memory_size": 1000 << 30} @@ -356,7 +354,7 @@ def test_arena_budget_math(monkeypatch): # Available-RAM ceiling binds when the machine is busy: 40 GB # reclaimable minus the floor beats the fraction of a 100 GB machine. monkeypatch.setattr( - gmlx.load.loader, "_available_ram_bytes", lambda: 40 << 30 + gmlx.stream.budget, "_available_ram_bytes", lambda: 40 << 30 ) monkeypatch.setenv("GMLX_DECODE_RAM_FLOOR_GB", "5") monkeypatch.delenv("GMLX_DECODE_PAGECACHE_GB", raising=False) @@ -386,7 +384,7 @@ def test_arena_budget_math(monkeypatch): monkeypatch.setenv("GMLX_DECODE_ARENA_FORCE", "1") assert _decode_arena_bytes(60 << 30, offsets, budget=None) == 200 << 30 monkeypatch.delenv("GMLX_DECODE_ARENA_FORCE", raising=False) - monkeypatch.setattr(gmlx.load.loader, "_available_ram_bytes", lambda: None) + monkeypatch.setattr(gmlx.stream.budget, "_available_ram_bytes", lambda: None) assert _decode_arena_bytes(60 << 30, offsets, budget=None) == 200 << 30 @@ -1576,9 +1574,9 @@ class _R: monkeypatch.setattr(subprocess, "run", lambda *a, **k: _R()) monkeypatch.setattr(gmlx.serve.kernel_vm, "snapshot", lambda: None) # vm_stat fallback page = 16384 - assert gmlx.load.loader._available_ram_bytes() == (1000 + 500 + 700000) * page + assert gmlx.stream.budget._available_ram_bytes() == (1000 + 500 + 700000) * page # Strict no-victims set: free + purgeable + speculative only. - assert gmlx.load.loader._available_ram_bytes(include_inactive=False) == \ + assert gmlx.stream.budget._available_ram_bytes(include_inactive=False) == \ (1000 + 500 + 5000) * page @@ -1597,7 +1595,7 @@ class _R: monkeypatch.setattr(subprocess, "run", lambda *a, **k: _R()) monkeypatch.setattr(gmlx.serve.kernel_vm, "snapshot", lambda: None) # vm_stat fallback - assert gmlx.load.loader._available_ram_bytes() == \ + assert gmlx.stream.budget._available_ram_bytes() == \ (1000 + 500 + 5000 + 200000) * 16384 @@ -1798,7 +1796,7 @@ def test_governor_shrink_regrows_with_pressure_polling_off(monkeypatch, tmp_path def test_prefill_ring_reason(monkeypatch, tmp_path): - from gmlx.load.loader import _prefill_ring_reason + from gmlx.stream.budget import _prefill_ring_reason from gmlx.stream.prefill_feeder import ring_bytes monkeypatch.delenv("GMLX_DECODE_ARENA_GB", raising=False) diff --git a/tests/stream/test_gpu_token.py b/tests/stream/test_gpu_token.py index c1cd8431..e4df2fbc 100644 --- a/tests/stream/test_gpu_token.py +++ b/tests/stream/test_gpu_token.py @@ -17,7 +17,7 @@ from mlx_lm.models.switch_layers import SwitchGLU from gmlx.stream.gpu_token import GpuTokenState, route_shed_op -from gmlx.load.loader import install_expert_streaming +from gmlx.stream.expert_streaming import install_expert_streaming from test_decode_feeder import _make_feeder diff --git a/tests/stream/test_lookahead.py b/tests/stream/test_lookahead.py index 7b85ee26..e7b9de4e 100644 --- a/tests/stream/test_lookahead.py +++ b/tests/stream/test_lookahead.py @@ -6,7 +6,7 @@ import numpy as np import pytest -from gmlx.load.loader import install_expert_streaming +from gmlx.stream.expert_streaming import install_expert_streaming from gmlx.stream.lookahead import ( LookaheadProbe, _gate_module_select, diff --git a/tests/stream/test_moe_experts.py b/tests/stream/test_moe_experts.py index 16fa9973..eb44f38f 100644 --- a/tests/stream/test_moe_experts.py +++ b/tests/stream/test_moe_experts.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from gmlx.load.loader import install_expert_streaming, install_moe_experts_override +from gmlx.stream.expert_streaming import install_expert_streaming, install_moe_experts_override from gmlx.stream.moe_experts import ( _INLINE_SWAPS, ExpertProbe, diff --git a/tests/stream/test_pin_cast_copies.py b/tests/stream/test_pin_cast_copies.py index 834aea94..a88b166f 100644 --- a/tests/stream/test_pin_cast_copies.py +++ b/tests/stream/test_pin_cast_copies.py @@ -22,7 +22,7 @@ import mlx.core as mx -from gmlx.load.loader import _decode_arena_bytes +from gmlx.stream.budget import _decode_arena_bytes from gmlx.stream.pin_weights import cast_copies VOCAB, HIDDEN = 120832, 6144 @@ -110,10 +110,10 @@ def test_dead_bytes_are_the_wire_minus_the_copy(): def test_the_arena_gains_what_the_dead_wire_gave_up(monkeypatch): """Without this the sizer re-charges every byte the pin released.""" - import gmlx.load.loader as loader + import gmlx.stream.budget as budget monkeypatch.delenv("GMLX_DECODE_ARENA_GB", raising=False) - monkeypatch.setattr(loader, "_available_ram_bytes", lambda *a, **k: None) + monkeypatch.setattr(budget, "_available_ram_bytes", lambda *a, **k: None) monkeypatch.setattr(mx, "device_info", lambda: {"memory_size": 137 * 10**9}) offsets = {0: [(0, 0, 200 << 30)]} diff --git a/tests/stream/test_prefill_feeder.py b/tests/stream/test_prefill_feeder.py index cb204272..37eb50f0 100644 --- a/tests/stream/test_prefill_feeder.py +++ b/tests/stream/test_prefill_feeder.py @@ -14,7 +14,7 @@ import mlx.core as mx from mlx_lm.models.switch_layers import SwitchGLU -from gmlx.load.loader import install_expert_streaming +from gmlx.stream.expert_streaming import install_expert_streaming from test_decode_feeder import ( _KINDS, diff --git a/tests/stream/test_table_stream.py b/tests/stream/test_table_stream.py index 1bf21c22..7c8c360e 100644 --- a/tests/stream/test_table_stream.py +++ b/tests/stream/test_table_stream.py @@ -12,7 +12,7 @@ import mlx.core as mx import gmlx.stream.table_stream as ts -from gmlx.load.loader import install_expert_streaming +from gmlx.stream.expert_streaming import install_expert_streaming def _kquant_table(rows=8, dims=32, seed=7): @@ -251,7 +251,7 @@ def test_ladder_streams_table_and_keeps_experts_resident(monkeypatch): deducted = [] monkeypatch.setattr( - "gmlx.load.loader.deduct_untracked_weights", + "gmlx.stream.expert_streaming.deduct_untracked_weights", lambda n, key=None: deducted.append(n)) n, offloaded = install_expert_streaming(model) @@ -476,11 +476,11 @@ def test_streamed_table_bytes_keys_on_streamed_state(monkeypatch): def test_arena_sizing_excludes_streamable_bytes(monkeypatch): - from gmlx.load.loader import _decode_arena_bytes - import gmlx.load.loader as loader + from gmlx.stream.budget import _decode_arena_bytes + import gmlx.stream.budget as budget monkeypatch.delenv("GMLX_DECODE_ARENA_GB", raising=False) - monkeypatch.setattr(loader, "_available_ram_bytes", lambda *a, **k: None) + monkeypatch.setattr(budget, "_available_ram_bytes", lambda *a, **k: None) monkeypatch.setattr( mx, "device_info", lambda: {"memory_size": 137 * 10**9, diff --git a/tests/tui/test_chat_e2e.py b/tests/tui/test_chat_e2e.py index 4ef084e5..d45204f6 100644 --- a/tests/tui/test_chat_e2e.py +++ b/tests/tui/test_chat_e2e.py @@ -148,7 +148,7 @@ def fake_mpc(model, max_kv_size=None): monkeypatch.setattr("gmlx.commands.cli.maybe_load_from_config", lambda *a, **k: None) monkeypatch.setattr("gmlx.load.loader.load_model", load_model or (lambda *a, **k: (object(), {}, _FakeTok()))) - monkeypatch.setattr("gmlx.load.loader._resolve_prefill_step", + monkeypatch.setattr("gmlx.stream.expert_streaming._resolve_prefill_step", lambda model, step: (None, False)) monkeypatch.setattr("gmlx.commands.cli._apply_placement", lambda args, model: None) monkeypatch.setattr("mlx_lm.models.cache.make_prompt_cache", fake_mpc) @@ -217,7 +217,7 @@ def fake_mpc(model, max_kv_size=None): monkeypatch.setattr("gmlx.commands.cli.maybe_load_from_config", lambda *a, **k: None) monkeypatch.setattr("gmlx.load.loader.load_model", lambda *a, **k: (object(), {}, _FakeTok())) - monkeypatch.setattr("gmlx.load.loader._resolve_prefill_step", + monkeypatch.setattr("gmlx.stream.expert_streaming._resolve_prefill_step", lambda model, step: (None, False)) monkeypatch.setattr("gmlx.commands.cli._apply_placement", lambda args, model: None) monkeypatch.setattr("mlx_lm.models.cache.make_prompt_cache", fake_mpc)