Skip to content

streaming.md

Latest commit

 

History

History
87 lines (58 loc) · 7.44 KB

streaming.md

File metadata and controls

87 lines (58 loc) · 7.44 KB
Raw

SSD Streaming Expert Layers (streaming)

The Problem

Qwen3.6-35B-A3B has 40 layers × 256 experts, with each layer's MoE weights around 310MB after 4-bit quantization — keeping them all resident far exceeds the Apple Silicon unified-memory budget. edge0's approach: keep the weights on SSD, mmap them into the GPU on demand, and use LRU + prefetch + fixed slots to turn "tens of MB moved per step" into "almost nothing moved per step".

Core components (edge0/streaming/):

Module Responsibility
mmap.py SafetensorsMmap: single-file byte-range mmap, lazily reads raw bytes by tensor name
cache.py SharedExpertCache: cross-layer shared LRU with capacity cache_slots
layer.py StreamingSwitchGLU: one instance per layer, holding all execution paths and state
options.py LayerOptions: typed configuration (a replacement for deployment-time env knobs)

Math Contract (must be bit-identical to the vendored model)

The MoE block math is down(silu(gate(x)) * up(x)), i.e. _swiglu(up, gate) = nn.silu(gate) * up. The bundle order maps directly onto the order of the math parameters, so the weight stack needs no reordering:

  • separate (unfused): the bundle order is ("up_proj", "gate_proj", "down_proj") — up first! After stacking, the wargs order matches the _swiglu(up, gate) signature position by position: (w_u, s_u, b_u, w_g, s_g, b_g, w_d, s_d, b_d).
  • fused gate+up: the bundle order is ("gate_up_proj", "down_proj"); the row layout of gate_up is gate rows on top, up rows below (concatenated along the output feature axis), and mx.split(x_gu, 2, -1) yields (gate, up).

Any reordering (e.g. the historical gate-first bundle order) is exposed by the reference comparison in tests/test_streaming_math.py (relative L2 jumps from ≈0.2% to ≈100%).

Quantized Layout

  • switch_mlp.<proj>.weight: packed u32 [E, out, in/8] (4-bit affine, group size 64);
  • scales / biases: bf16 dtype, per group [E, out, in/64];
  • Kernel: backends.quant.gather_qmm (MLX quantized gather matmul), with bf16 internal precision after dequantization (relative L2 ≈0.24%; the test tolerance is calibrated against this).

Execution Paths (from fast to full)

1. exact (on-demand bundle)

Deduplicate the routing indices → _get_bundles(unique) builds per expert (read from the LRU or mmap) → stack → gather_qmm. The slowest but most general path; it is the correctness baseline for the other paths.

2. staged (fixed-slot double buffering, the decode workhorse)

  • staged_n fixed slots plus an overflow zero slot;
  • Routing indices → mx.take on the slot table, indices never leave the GPU (zero host synchronization per layer per step);
  • staged_sync: fill synchronized at step boundaries; asm_cache: caches the slot table and the stacked graph nodes per expert set, so repeated sets are not rebuilt; incr_stack: replaces the 9 mx.stack nodes with incremental-stack put_along_axis row writes (deduplicated against the LRU while staged is active, and returned by incr_writeback);
  • Missing experts map to the overflow zero slot and their contribution is dropped; if a slot is not ready, the path falls back to exact.

3. hot (LRU-resident top-N)

  • Each layer picks its top-N (hot_per_layer) from the decayed counts _hot_counts; experts that hit in the LRU live as a resident stack (load_hot_layer / materialize_hot, with the sliding window hot_window controlling how many layers stay resident at once);
  • Hits take the stacked gather; misses take exact with a scatter-add.

4. full-layer (whole-layer loading for E3b, the prefill workhorse)

  • load_full_layer(): loads the layer's 9 tensors directly (the checkpoint already stores each layer as a single stacked tensor; the mmap dtype conversion costs ≈9ms per layer, and the CPU load is hidden under the previous layer's GPU execution);
  • _gather_sort folds the batch into the token dimension → sorted gather → _scatter_unsort un-sorts and restores the batch dimension;
  • full_layer_prefill turns whole-layer prefill on per tier, and prefill_full_layers limits it to the leading N layers (0 = every layer, the edge0-8b setting); layers outside that range go through hot/exact. With full_layer_prefill=False the whole-layer path is off entirely — no layer is loaded whole, whatever prefill_full_layers says — and prefill runs the prefill_hot window if one is set, else the plain per-expert on-demand path;
  • After use, clear_full_layer() frees the GPU copy, and the page cache carries the hot data.

Sorting and Compilation

With use_compile, the staged/exact paths are wrapped in mx.compile:

  • set sizes ≥64 use the sorted variant _moe_math_sorted (_gather_sort preprocessing + sorted gather + _scatter_unsort restore);
  • otherwise the direct gather variant _moe_math is used;
  • the two are mathematically equivalent, and the tests compare every path against exact.

Key Options (LayerOptions)

Field Meaning Default
staged / staged_n / staged_trigger Fixed-slot decode toggle / slot count / trigger top-k False / 8 / 8
staged_replace The staged set replaces routing (paired with the prerouter; zero drops) False
staged_sync / asm_cache / incr_stack / incr_writeback Fill synchronization / graph-node cache / incremental stack / writeback to the LRU True / True / False / False
hot_per_layer / hot_update_interval / hot_decay / pin_bonus Hot-expert residency count and refresh 0 / 4 / 0.75 / 2.0
cache_slots / prefetch_cap LRU capacity / prefetch buffer 64 / 48
load_threads / prefetch_threads Build / prefetch thread counts 8 / 4
full_layer_prefill / prefill_full_layers Whole-layer prefill loading / number of leading layers False / 0
prefill_hot hot stack size during prefill 0
warm_willneed Kernel bulk readahead (madvise WILLNEED) over the expert ranges a prefetch/stage is about to touch False
use_compile / top_k compile wrapping / routing top-k override True / None

Presets: staged_k4() (edge0-35b: staged decode with 4 slots, prefill hot stack 32, on-demand prefill), prod_k8() (edge0-8b: the reference deployment profile — staged decode off, E3b whole-layer prefill), and staged_k8() (the plain K=8 staged variant). Both tiers share cache_slots=64.

The whole-layer prefill is the fastest path when the checkpoint stays in the page cache (warm 27-token prefill: 0.24 s vs 0.37 s on-demand on an M4 Pro), and the slowest one when it does not (cold: 5.3 s / 4.06 GiB read vs 0.6-1.1 s / 0.4-0.8 GiB; the on-demand figure varies with how many distinct experts the prompt routes to). edge0 demo|chat|serve --prefill-ondemand selects the on-demand path for machines in the second group.

Why Whole-Layer Loading Is Also Fast

The MoE weights in the checkpoint are already single per-layer stacked tensors (such as [256, 512, 256]), so a whole-layer load is 9 direct reads + a dtype view, not 256×9 per-expert builds. CPU loading and GPU execution are overlapped through the before_layer_cb + async_eval_per_layer pipeline.

Note that a fused gate_up whole-layer load must interleave the gate/up rows per expert (consistent with the per-expert concatenation in _build); concatenating the two [E, ...] tensors wholesale and then reshaping produces misaligned rows (a real bug in an early version, now pinned down by test_full_layer_prefill_matches_exact[fused]).