Lazy-load GGUF MoE experts from DISK on demand
Motivation
MoE models only route each token to a small subset of experts, but tinygrad.llm currently brings the entire GGUF payload — including every expert — onto the default compute device before inference.
This makes models such as qwen3:30b-a3b, qwen3.6:35b-a3b, OLMoE, Moonlight, and GLM-4.7-Flash require memory for all packed experts even though only top-k experts are active for a token. It also increases startup cost and prevents running larger MoE GGUFs with a bounded expert cache.
Current behavior
The current path is eager at the GGUF-file level:
tinygrad/llm/gguf.py::_gguf_parsecallstensor.to(None).realize()before parsing metadata, so the whole GGUF leavesDISKand is copied toDevice.DEFAULT.- The parser immediately builds a state-dict entry for every tensor and has no tensor/expert selection API.
ExpertWeightsstores each routed projection as one packed(num_experts, out_features, in_features)tensor.- Routing happens later in
FFNBlock._feed_forward;self.weight[sel]selects top-k experts only after all expert weights have already been represented by the loaded model state. REALIZE=1additionally realizes every parameter, whileREALIZE=0still starts from a fully device-resident GGUF byte tensor.
I could not find an existing issue or PR implementing selected-expert GGUF loading/offload.
Related groundwork: #17616 makes ordinary DISK tensor assignment schedulable/lazy, but GGUF parsing and dynamic MoE dispatch still need an expert-aware loading layer.
Proposed feature
Add optional on-demand expert loading for MoE models loaded through tinygrad.llm from GGUF:
- Keep GGUF storage DISK-backed. Parse metadata and tensor descriptors without realizing the complete file on
Device.DEFAULT. Preserve(split file, byte offset, shape, GGML type)descriptors. - Make expert slices independently addressable. For packed expert tensors, map
(layer, projection, expert_id)to the exact quantized byte range/block layout needed to decode that expert. This must work for supported quantized GGUF types and split GGUF files. - Route before materializing experts. Compute top-k routing first, take the union of experts selected by the current batch/tokens, and materialize only cache misses before the expert matmuls.
- Use a bounded cache. Cache decoded experts by layer/projection/expert/dtype/device with a configurable byte or expert-count budget and deterministic eviction (for example LRU). Router, dense layers, norms, embeddings, output weights, and shared experts should remain resident.
- Support warm-path reuse and optional prefetch. Repeated expert selections should not reread or requantize weights. A later optimization could prefetch based on recent routing while preserving correctness.
- Keep an eager fallback. Models/devices where dynamic expert staging is unsupported should retain the current behavior.
The implementation will need to account for dynamic device-side sel, JIT/capture stability, batched tokens selecting different experts, and the current packed 3-D expert representation. Fixed staging buffers may be preferable to changing kernel buffer identities on every token.
Acceptance criteria
- Works through
python -m tinygrad.llm --model <MoE GGUF>and supports at least one existing preset such asqwen3:30b-a3b. - Supports the relevant quantized GGUF expert layouts and multi-part GGUF files.
- Produces deterministic output/logit parity with eager loading for the same model, prompt, and sampling configuration.
- Provides an API or environment option for eager versus lazy experts and a configurable cache/VRAM budget.
- Peak device memory is bounded by resident non-expert weights plus active/cached experts rather than all experts.
- Instrumentation demonstrates that cold inference reads only selected cache misses and warm inference does not reread cached experts.
- Handles the union of selected experts for batched/prefill tokens and preserves shared-expert behavior.
- Includes unit tests using a small synthetic MoE/GGUF fixture for cache hits, misses, eviction, byte ranges, and eager/lazy parity.
Initial test status
The existing MoE routing suite passes unchanged:
DEV=CPU python -m pytest -n12 -q test/unit/test_llm_moe.py
5 passedSource: tinygrad/tinygrad