[presentyourlove]

Description

This PR implements Frontier Dimension 3.1: Hierarchical Expert Offloading & Early-Layer Speculative Prefetching, enabling ultra-low-memory edge devices (e.g. 4GB RAM) to virtualize and serve giant MoE models (such as DeepSeek-V3 or Mixtral 8x22B) with negligible I/O stall overhead.

Background & Motivation

Serving cutting-edge Mixture-of-Experts (MoE) models on consumer edge devices is bottlenecked by physical memory capacity. Even when sharded, retaining active expert weights exceeds local RAM limits.
By partitioning MoE experts into tiered storage (resident VRAM/RAM for hot experts, NVMe SSD for cold experts) and employing an early-layer predictor (prefetching layer $L+4$ cold experts asynchronously during layer $L$ forward compute), transfer latencies from disk/host memory are completely hidden behind computation, enabling a 4GB device to "pseudo-host" a 32GB+ expert pool.

Key Changes

1. Tiered MoE Expert Models & Telemetry (src/exo/shared/types/moe_offload.py):

   • OffloadedExpertWeight: Immutable container capturing expert weights across tiered storage (vram, host_ram, nvme_ssd), access timestamps, and hot-pinning flags.
   • ExpertCacheSlot: Defines physical slot capacity and assigned keys.
   • MoEOffloadTelemetry: Tracks hot cache hits, prefetch hits, cold SSD misses, total swap-in volume, and I/O latency masking ratios.

2. Early-Layer Predictor & Hierarchical Engine (src/exo/worker/engines/moe/hierarchical_expert_manager.py):

   • EarlyLayerPrefetchPredictor: Extracts Top-K expert activations and projects candidate indices across lookahead layers ($L + 4$).
   • HierarchicalExpertEngine:
     • Asynchronous non-blocking background loading pipeline via thread pool.
     • LRU unpinned expert eviction from fast pool back to cold SSD storage upon capacity saturation.
     • Pinned hot experts permanently retained in memory.
     • Overlapped forward layer compute with background DMA prefetching.

3. Module Exports (src/exo/worker/engines/moe/__init__.py):

   • Exported HierarchicalExpertEngine and EarlyLayerPrefetchPredictor.

4. Unit Tests (src/exo/shared/tests/test_moe_hierarchical_offload.py):

   • Verified expert metadata and memory size accounting.
   • Verified lookahead prediction across list and dictionary routing probabilities.
   • Verified hot pool hits vs cold storage fallback and swap-in accounting.
   • Verified that LRU eviction safely offloads unpinned experts while protecting pinned hot experts.
   • Verified end-to-end early prefetch execution, compute overlap, and telemetry accuracy.

Verification

• Ran unit tests:

  .venv/Scripts/pytest.exe src/exo/shared/tests/test_moe_hierarchical_offload.py
  # 6 passed in 0.10s

• Code styling and linting:

  .venv/Scripts/ruff.exe check src/exo/shared/types/moe_offload.py src/exo/worker/engines/moe/hierarchical_expert_manager.py src/exo/worker/engines/moe/__init__.py src/exo/shared/tests/test_moe_hierarchical_offload.py
  # All checks passed!

Related Issues

• Resolves: Frontier Dimension 3.1 in 人人都能跑llm.md (Hierarchical Expert Offloading & Speculative Prefetching)