memra_engine/

spill.rs

//! 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.

use crate::model::HostBuf;
use crate::Engine;
use memmap2::Mmap;
use std::sync::Arc;

/// 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.
#[derive(Clone, Copy, Debug)]
pub struct MemBudget {
    /// Free VRAM in bytes, from `cuMemGetInfo` (authoritative; accounts for other GPU processes).
    pub free_vram: usize,
    /// Bytes of host RAM safe to pin: `/proc/meminfo MemAvailable` × `pinned_frac` (default 0.60).
    /// Capped so `cudaHostAlloc` can neither OOM nor evict the page cache the Tier-2 mmap depends on.
    pub free_pinnable_ram: usize,
}

impl MemBudget {
    pub fn probe(e: &Engine) -> Result<Self, Box<dyn std::error::Error>> {
        let (free_vram, _total) = e.ctx().mem_get_info()?; // same call moe_cache.rs:77 uses
        let avail = read_meminfo_kb("MemAvailable")? * 1024; // MemAvailable (NOT MemFree)
        let frac = std::env::var("MEMRA_SPILL_PINNED_FRAC")
            .ok()
            .and_then(|s| s.parse::<f64>().ok())
            .unwrap_or(0.60);
        Ok(MemBudget {
            free_vram,
            free_pinnable_ram: (avail as f64 * frac) as usize,
        })
    }
}

/// Parse one `/proc/meminfo` field (a value in kB) by key, e.g. "MemAvailable".
fn read_meminfo_kb(key: &str) -> Result<usize, Box<dyn std::error::Error>> {
    let s = std::fs::read_to_string("/proc/meminfo")?;
    for line in s.lines() {
        // line form: "MemAvailable:   12345678 kB"
        if let Some(rest) = line.strip_prefix(key) {
            let rest = rest.trim_start_matches(':').trim();
            let kb: usize = rest
                .split_whitespace()
                .next()
                .ok_or("malformed /proc/meminfo line")?
                .parse()?;
            return Ok(kb);
        }
    }
    Err(format!("/proc/meminfo: key {key} not found").into())
}

/// 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.
#[inline]
pub fn disk_tier_enabled() -> bool {
    std::env::var("MEMRA_SPILL_DISK").is_ok()
}

/// 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`.
pub struct SpillCtx {
    /// One `MAP_SHARED` mmap per physical GGUF shard, shared (`Arc`) across every spilled expert
    /// block that lives in that shard. Index = `TensorInfo::shard`. Single-file models have len 1.
    /// PER-SHARD, not one map: a split model's `tensor_file_range` offsets are relative to the
    /// OWNING shard's file, so pairing them with shard 0's mmap would read the wrong bytes (and
    /// would index out of bounds for any shard larger than shard 0).
    pub file_maps: Vec<Arc<Mmap>>,
    /// The opened inodes backing `file_maps`, same indexing, retained for positioned expert reads.
    pub files: Vec<Arc<std::fs::File>>,
    /// Pinned-RAM budget still available (bytes); decremented as experts are pinned.
    pub pinned_remaining: usize,
    /// Diagnostics: how many experts landed pinned vs. mmap'd, and total disk-tier bytes.
    pub n_pinned: usize,
    pub n_mmap: usize,
    pub mmap_bytes: usize,
}

impl SpillCtx {
    /// Clone each parsed shard's opened inode, create a `MAP_SHARED` mmap per shard, and seed the
    /// pinned budget from a live `MemBudget` probe.
    /// The whole-map expert advice defaults to random (the historical behavior); setting
    /// `MEMRA_MOE_MMAP_ADVICE=normal` restores ordinary Linux readahead. SPILLING-PLAN §1.
    pub fn open(
        g: &memra_gguf::GgufFile,
        budget: &MemBudget,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let mut files = Vec::with_capacity(g.n_shards());
        let mut file_maps = Vec::with_capacity(g.n_shards());
        for i in 0..g.n_shards() {
            let file = g.shard_file(i).clone();
            // MAP_SHARED, no MAP_POPULATE (memmap2's default Mmap::map): zero upfront copy,
            // demand-fault.
            let map = unsafe { Mmap::map(file.as_ref())? };
            let _ = memra_gguf::source::apply_expert_mmap_advice(&map);
            files.push(file);
            file_maps.push(Arc::new(map));
        }
        Ok(SpillCtx {
            file_maps,
            files,
            pinned_remaining: budget.free_pinnable_ram,
            n_pinned: 0,
            n_mmap: 0,
            mmap_bytes: 0,
        })
    }
}

/// 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.
pub fn place_expert(
    ctx: &mut SpillCtx,
    e: &Engine,
    raw: &[u8],
    file_off: usize,
    shard: usize,
) -> Result<HostBuf, Box<dyn std::error::Error>> {
    let len = raw.len();
    if ctx.pinned_remaining >= len {
        // Tier 1: pinned host memory — true async DMA at full PCIe (matches the no-spill path).
        ctx.pinned_remaining -= len;
        ctx.n_pinned += 1;
        let mut p = unsafe { e.ctx().alloc_pinned::<u8>(len)? };
        {
            let dst = p.as_mut_slice()?;
            dst.copy_from_slice(raw);
        }
        let base = p.as_ptr()? as *const u8;
        Ok(HostBuf::Pinned {
            slice: std::sync::Arc::new(p),
            base,
            len,
        })
    } else {
        // Tier 2: mmap the GGUF region — demand-faulted from NVMe on first H2D. Zero RAM cost.
        ctx.n_mmap += 1;
        ctx.mmap_bytes += len;
        Ok(HostBuf::Mmap {
            map: ctx.file_maps[shard].clone(),
            file: ctx.files[shard].clone(),
            off: file_off,
            len,
        })
    }
}

/// 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`).
pub struct SpillBlock {
    pub host: HostBuf,
    pub qtype: i32,
    pub in_f: usize,
    pub out_f: usize,
    pub row_bytes: usize,
}

impl SpillBlock {
    /// The H2D DMA source for this block — resolves the tier (`Pinned` fast / `Mmap` demand-fault).
    #[inline]
    pub fn bytes(&self) -> &[u8] {
        self.host.as_bytes()
    }
}

/// 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.
pub struct Tiered {
    pub host: crate::model::HostExps, // Tier 1/2 (Pinned hot / Mmap cold), per-expert
    pub slots: crate::moe_cache::MoeSlotCache, // Tier 0 GPU residency (existing slot cache)
}

#[cfg(all(test, unix))]
mod tests {
    use super::{MemBudget, SpillCtx};
    use memra_gguf::{GgufFile, GGUF_MAGIC};

    #[test]
    fn spill_ctx_keeps_parsed_gguf_inode_after_path_replacement() {
        let path =
            std::env::temp_dir().join(format!("memra-spill-inode-{}.gguf", std::process::id()));
        let mut original = Vec::new();
        original.extend_from_slice(&GGUF_MAGIC.to_le_bytes());
        original.extend_from_slice(&3u32.to_le_bytes());
        original.extend_from_slice(&0i64.to_le_bytes());
        original.extend_from_slice(&0i64.to_le_bytes());
        original.resize(32, 0);
        std::fs::write(&path, &original).unwrap();

        let gguf = GgufFile::open(&path).unwrap();
        std::fs::remove_file(&path).unwrap();
        std::fs::write(&path, vec![0xA5u8; original.len()]).unwrap();

        let budget = MemBudget {
            free_vram: 0,
            free_pinnable_ram: 0,
        };
        let spill = SpillCtx::open(&gguf, &budget).unwrap();
        assert_eq!(spill.files.len(), 1, "single-file GGUF must yield exactly one shard map");
        assert!(std::sync::Arc::ptr_eq(&spill.files[0], gguf.opened_file()));
        assert_eq!(&spill.file_maps[0][..], original.as_slice());
        assert_eq!(std::fs::read(&path).unwrap(), vec![0xA5u8; original.len()]);

        std::fs::remove_file(path).ok();
    }
}