Skip to main content

Module spill

Module spill 

Source
Expand description

SPILLING-PLAN: full tiered spilling (VRAM ↔ pinned-host ↔ mmap-disk).

Today memra has the VRAM↔pinned-host leg (the MoeSlotCache GPU slot cache + the pinned HostExps host store). This module adds the third tier: a HostBuf::Mmap arm (model.rs) so cold experts are demand-faulted from the GGUF file on disk instead of held in RAM, plus the runtime memory probe (MemBudget) that decides — per expert, at load — which tier each block lives in. Never hardcode: VRAM is queried via cuMemGetInfo, host RAM via /proc/meminfo.

THE GATE (SPILLING-PLAN §8): spilling is a memory-PLACEMENT change, never a numerics change. A Mmap expert and a Pinned expert feed qmatvec_view byte-for-byte identical GGUF bytes — the Pinned/Paged stores copied FROM exactly those on-disk bytes — so argmax is unchanged.

The disk tier is gated behind MEMRA_SPILL_DISK. Unset (default) = the current all-host behavior, byte-identical: HostExps::tiers stays None and every expert slices the single pinned/paged backing store. The daily models (9B/27B) fit 24 GB and NEVER trigger spill.

Structs§

MemBudget
Runtime free-memory budget (SPILLING-PLAN §2). Both numbers are QUERIED at load, never hardcoded — free host RAM “varies with other LLM servers”, so the split between pinned (Tier 1) and disk (Tier 2) must be decided against the live machine state.
SpillBlock
SPILLING-PLAN §3/§5: a single spillable weight block over the same {Pinned, Mmap} substrate. Lifted from the HostExps fields so dense weights (dense-70B case) can reuse the disk tier without the 256-expert stacking. Carried for the requested generalization; the MoE path uses HostExps directly (which now embeds the same tier machinery via HostBuf).
SpillCtx
Shared load-time spill context (SPILLING-PLAN §2 step 4). Built ONCE per model load when the disk tier is on, then handed by &mut to each HostExps::load so all layers/projections share ONE file mmap PER SHARD and draw down a single running pinned-RAM budget. Greedy in load order: pin until pinned_remaining is exhausted, then spill every later expert to Mmap.
Tiered
SPILLING-PLAN §3: the requested Tiered generalization. Structurally it is the existing HostExps (Tier 1/2 host backing, per-block) composed with the existing MoeSlotCache (Tier 0 GPU residency). Both seams are already present and unchanged; this names the composition. The MoE hot loop drives the two seams directly (expert_bytes() + with_moe_cache), so this is a documentation/structural alias, not a new hot path.

Functions§

disk_tier_enabled
Is the disk tier (Tier 2) enabled? Gated behind MEMRA_SPILL_DISK. Default (unset) = off => the unchanged all-host path (HostExps::tiers stays None). Set to anything to force-on.
place_expert
Build one expert’s HostBuf, choosing its tier under the running budget (SPILLING-PLAN §1.1): pin (Tier 1) while pinned_remaining covers the block, else Mmap it (Tier 2). file_off is this expert’s byte offset within ITS OWN SHARD’s file (= shards[t.shard].data_start + tensor.offset + e*stride), and shard selects the matching mmap. Returns the chosen HostBuf; the bytes are bit-identical whichever tier is picked.