#17617·tinygrad

Lazy-load GGUF MoE experts from DISK on demand

Author: EthanBnnttCreated Aug 20, 2026Updated Sep 2, 2026

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_parse calls tensor.to(None).realize() before parsing metadata, so the whole GGUF leaves DISK and is copied to Device.DEFAULT.
  • The parser immediately builds a state-dict entry for every tensor and has no tensor/expert selection API.
  • ExpertWeights stores 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=1 additionally realizes every parameter, while REALIZE=0 still 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:

  1. 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.
  2. 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.
  3. 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.
  4. 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.
  5. 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.
  6. 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 as qwen3: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 passed