refactor(load): split loader.py into wire, mtp_target, and stream helpers

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

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:

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}

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} "

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)