feat(quantized): disk-streaming LRU expert cache (colibri low-memory mode)

Stream a big MoE's routed experts from NVMe instead of loading them resident:
per-(layer x projection) bank keeps the stacked [E,n,k] weight on disk and preads
one expert's [n,k] slice on demand through a pinned hot-set -> per-bank LRU ->
disk (+posix_fadvise DONTNEED so the page cache stays bounded). Auto-sizes the
cache from MemAvailable (colibri cap_for_ram, auto-raises to fill big-RAM boxes)
and learns/auto-pins the hottest experts via a usage sidecar (colibri AUTOPIN).

- expert_stream.rs: ExpertStreamBank (fetch/pin/LRU/stats), registry + finalize()
  (RAM-sized caps), save_usage/learning-pin. Runtime-gated (default OFF).
- QStorage::Stream variant; indexed_moe_forward gains a Stream arm sharing ONE
  moe_grouped_per_expert helper with the resident path (only the fetch differs).
- gguf_file: TensorInfo::read_stream / Content::tensor_stream loader seam.
- libc made non-optional for posix_fadvise (dropped dep:libc from mkl/accelerate).

Bit-exact: streamed expert == resident bytes -> identical indexed_moe output.
Gate (tests/expert_stream_tests.rs, 4/4): streamed==resident bit-for-bit; LRU
cap bounds resident RAM with correctness cap-independent; pinned expert survives
eviction; usage sidecar learns the hot expert.
This commit is contained in:
Hanzo Dev
2026-07-12 17:29:59 -07:00
parent 20769e367f
commit 09d2528cb1
6 changed files with 843 additions and 43 deletions
+75
View File
@@ -0,0 +1,75 @@
#!/bin/bash
# 2-node loopback ring on dbc: rank0 = CPU head, rank1 = Metal worker.
# Reproduces "a Metal rank in the PP ring deadlocks". Dense Qwen3-0.6B first (transport discriminator).
set -u
BIN=~/work/hanzo/engine-ring/target/release/hanzo
MODEL_REPO="Qwen/Qwen3-0.6B-GGUF"
MODEL_FILE="Qwen3-0.6B-Q8_0.gguf"
SPLIT="${SPLIT:-14,14}"
D=/tmp/ringrepro
rm -rf "$D"; mkdir -p "$D"
cat > "$D/rank0.json" <<'EOF'
{"master_ip": null, "master_port": 9100, "port": 9000, "right_port": 9001, "right_ip": "127.0.0.1", "rank": 0, "world_size": 2}
EOF
cat > "$D/rank1.json" <<'EOF'
{"master_ip": "127.0.0.1", "master_port": 9100, "port": 9001, "right_port": 9000, "right_ip": "127.0.0.1", "rank": 1, "world_size": 2}
EOF
echo "=== launching rank1 (Metal worker) ==="
RING_CONFIG="$D/rank1.json" RING_LAYER_SPLIT="$SPLIT" RUST_LOG=info \
"$BIN" serve -p 9235 --no-ui --no-advertise -m "$MODEL_REPO" --format gguf -f "$MODEL_FILE" \
> "$D/rank1.log" 2>&1 &
R1=$!
echo "rank1 pid=$R1"
sleep 2
echo "=== launching rank0 (CPU head) ==="
RING_CONFIG="$D/rank0.json" RING_LAYER_SPLIT="$SPLIT" RUST_LOG=info \
"$BIN" serve -p 9234 --no-ui --no-advertise -m "$MODEL_REPO" --format gguf -f "$MODEL_FILE" --cpu --dtype f32 \
> "$D/rank0.log" 2>&1 &
R0=$!
echo "rank0 pid=$R0"
echo "=== waiting for head /v1/models (up to 120s) ==="
UP=0
for i in $(seq 1 60); do
if curl -s -m 2 http://127.0.0.1:9234/v1/models >/dev/null 2>&1; then UP=1; break; fi
if ! kill -0 $R0 2>/dev/null; then echo "rank0 DIED"; break; fi
if ! kill -0 $R1 2>/dev/null; then echo "rank1 DIED"; break; fi
sleep 2
done
echo "head_up=$UP"
MODELID=$(curl -s -m 3 http://127.0.0.1:9234/v1/models 2>/dev/null | python3 -c 'import sys,json;print(json.load(sys.stdin)["data"][0]["id"])' 2>/dev/null)
echo "modelid=$MODELID"
echo "=== sending completion (60s timeout) ==="
REQ_START=$(date +%s)
RESP=$(curl -s -m 60 http://127.0.0.1:9234/v1/chat/completions \
-H 'content-type: application/json' \
-d "{\"model\":\"${MODELID:-default}\",\"messages\":[{\"role\":\"user\",\"content\":\"count to five\"}],\"max_tokens\":20,\"stream\":false}" 2>&1)
RC=$?
REQ_END=$(date +%s)
echo "curl_rc=$RC elapsed=$((REQ_END-REQ_START))s"
echo "RESP=${RESP:0:400}"
if [ "$RC" != "0" ] || [ -z "$RESP" ]; then
echo "=== DEADLOCK SUSPECTED -> sampling rank1 (Metal worker) ==="
sample $R1 4 -file "$D/rank1.sample.txt" 2>/dev/null
echo "--- rank1 sample (thread summary) ---"
grep -A40 "Call graph" "$D/rank1.sample.txt" 2>/dev/null | head -80
echo "=== CPU usage snapshot ==="
ps -p $R1 -o pid,%cpu,%mem,command 2>/dev/null | head -2
ps -p $R0 -o pid,%cpu,%mem,command 2>/dev/null | head -2
echo "=== socket state (ring ports) ==="
netstat -an 2>/dev/null | grep -E "9000|9001" | head
else
echo "=== NO DEADLOCK: dense Metal worker completed the forward ==="
fi
echo "=== rank1 log tail ==="; tail -25 "$D/rank1.log"
echo "=== rank0 log tail ==="; tail -25 "$D/rank0.log"
echo "=== cleanup ==="
kill $R0 $R1 2>/dev/null; sleep 1; kill -9 $R0 $R1 2>/dev/null
echo DONE
+3 -3
View File
@@ -26,7 +26,7 @@ gemm = { workspace = true }
half = { workspace = true }
float8 = { workspace = true }
intel-mkl-src = { workspace = true, optional = true }
libc = { workspace = true, optional = true }
libc = { workspace = true }
libm = { workspace = true }
memmap2 = { workspace = true }
num-traits = { workspace = true }
@@ -61,8 +61,8 @@ rocm = ["dep:rocm-rs", "dep:hanzo-rocm-kernels"]
rocm-miopen = ["rocm-rs/miopen"]
vulkan = ["dep:ash"]
wgpu = ["dep:wgpu", "dep:pollster"]
mkl = ["dep:libc", "dep:intel-mkl-src"]
accelerate = ["dep:libc", "dep:accelerate-src"]
mkl = ["dep:intel-mkl-src"]
accelerate = ["dep:accelerate-src"]
metal = [
"dep:hanzo-metal-kernels",
"dep:objc2-metal",
+459
View File
@@ -0,0 +1,459 @@
//! Disk-streaming LRU expert cache for large MoE GGUFs (the colibri "low memory mode").
//!
//! A GGUF MoE stores each layer's routed experts as three stacked banks -- `ffn_gate_exps`,
//! `ffn_up_exps`, `ffn_down_exps` -- each a `[n_experts, n, k]` quantized tensor. Loading them
//! resident is what makes GLM-5.2 (202 GB) need 202 GB of RAM. This module keeps them on NVMe and
//! streams one expert's `[n, k]` slice on demand: pin -> per-bank LRU -> `pread` (+ `fadvise`
//! DONTNEED so the page cache can't balloon). Resident RAM becomes `pinned + LRU` owned slabs --
//! bounded and under our control, not at the mercy of the OS page cache. With a warm cache and the
//! hot experts pinned, decode is matmul/bandwidth bound just like the resident path.
//!
//! Faithful to colibri (`c/glm.c`): `expert_load` (pread coalesced + fadvise), `cap_for_ram`
//! (auto-size the cache from MemAvailable, auto-raise to fill big-RAM boxes), `pin_load`/AUTOPIN
//! (learn the hottest experts into a sidecar and pin them at startup). Here the natural unit is one
//! bank per (layer x projection) rather than colibri's per-expert triple, because the three
//! projections are three separate GGUF tensors -- more orthogonal, same mechanism.
//!
//! Bit-exact by construction: a streamed expert is the identical file bytes the resident slice
//! would hold, so `indexed_moe_forward` produces the identical output; only the fetch path differs.
use crate::Result;
use std::collections::HashMap;
use std::fs::File;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock, Weak};
use super::GgmlDType;
/// Runtime gate. Streaming is OFF by default so resident behaviour is unchanged for models that
/// fit; the loader flips it on for `--stream-experts`. Compile-time-always, runtime-gated -- one way.
static ENABLED: AtomicBool = AtomicBool::new(false);
pub fn set_enabled(on: bool) {
ENABLED.store(on, Ordering::Relaxed);
}
pub fn enabled() -> bool {
ENABLED.load(Ordering::Relaxed)
}
/// Reserve left to the kernel page cache + activations/logits, mirroring colibri's honest slack.
const PAGE_CACHE_RESERVE: u64 = 2_500_000_000;
const ACTIVATION_RESERVE: u64 = 1_200_000_000;
/// Fraction of MemAvailable the cache may claim (rest breathes for OS + wrapper); colibri's 0.88.
const BUDGET_FRACTION: f64 = 0.88;
/// Default hot-set: fraction of a bank's cap auto-pinned from the learned usage sidecar.
const PIN_FRACTION: f64 = 0.25;
/// One streaming expert bank = one stacked `[n_experts, n, k]` GGUF tensor kept on disk.
pub struct ExpertStreamBank {
/// Tensor name (`blk.{L}.ffn_gate_exps.weight`); the stable key for the usage sidecar.
name: String,
file: File,
/// Byte offset of expert 0 within the file (`tensor_data_offset + tensor.offset`).
base_offset: u64,
/// Bytes per expert = `n*k / block_size * type_size`.
expert_bytes: usize,
n_experts: usize,
dtype: GgmlDType,
n: usize,
k: usize,
inner: Mutex<BankState>,
}
struct BankState {
/// Max experts held in the LRU (the pinned hot-set is separate and unbounded by `cap`).
cap: usize,
clock: u64,
lru: HashMap<u32, LruEntry>,
pinned: HashMap<u32, Arc<[u8]>>,
/// Per-expert routing frequency this run; persisted to the sidecar for AUTOPIN next run.
usage: Vec<u64>,
hits: u64,
misses: u64,
}
struct LruEntry {
bytes: Arc<[u8]>,
used: u64,
}
impl ExpertStreamBank {
/// Open a bank over `path`. Registers it so the global budget can size every bank at once.
#[allow(clippy::too_many_arguments)]
pub fn open(
name: String,
path: &std::path::Path,
base_offset: u64,
expert_bytes: usize,
n_experts: usize,
dtype: GgmlDType,
n: usize,
k: usize,
) -> Result<Arc<Self>> {
let file = File::open(path)?;
let bank = Arc::new(Self {
name,
file,
base_offset,
expert_bytes,
n_experts,
dtype,
n,
k,
inner: Mutex::new(BankState {
// Provisional cap of 1 until `finalize` sizes it from RAM; never zero (a miss must
// always be admittable so a fetch can't deadlock waiting for room).
cap: 1,
clock: 0,
lru: HashMap::new(),
pinned: HashMap::new(),
usage: vec![0; n_experts],
hits: 0,
misses: 0,
}),
});
registry().lock().expect("registry lock").register(&bank);
Ok(bank)
}
pub fn dtype(&self) -> GgmlDType {
self.dtype
}
pub fn n_experts(&self) -> usize {
self.n_experts
}
/// Per-expert matrix dims `[n, k]` (the shape a fetched expert's `QTensor` takes).
pub fn expert_dims(&self) -> (usize, usize) {
(self.n, self.k)
}
pub fn expert_bytes(&self) -> usize {
self.expert_bytes
}
/// Full logical size of the whole bank as if resident (for shape/size accounting only).
pub fn logical_bytes(&self) -> usize {
self.expert_bytes * self.n_experts
}
/// Fetch expert `eid`'s raw GGML bytes: pinned hot-set -> LRU -> disk. Bit-identical to the
/// resident slice; records the access for AUTOPIN learning.
pub fn fetch(&self, eid: u32) -> Result<Arc<[u8]>> {
let mut guard = self.inner.lock().expect("expert bank lock poisoned");
// Reborrow as `&mut BankState` so disjoint fields (hits/clock vs pinned/lru) can be borrowed
// independently -- the MutexGuard's Deref would otherwise treat every field access as the
// whole guard.
let st = &mut *guard;
if (eid as usize) < st.usage.len() {
st.usage[eid as usize] = st.usage[eid as usize].saturating_add(1);
}
if let Some(bytes) = st.pinned.get(&eid) {
st.hits += 1;
return Ok(bytes.clone());
}
if let Some(entry) = st.lru.get_mut(&eid) {
st.hits += 1;
st.clock += 1;
entry.used = st.clock;
return Ok(entry.bytes.clone());
}
st.misses += 1;
let bytes = self.read_expert(eid)?;
self.admit(st, eid, bytes.clone());
Ok(bytes)
}
/// Insert into the LRU, evicting the least-recently-used entry when at cap.
fn admit(&self, st: &mut BankState, eid: u32, bytes: Arc<[u8]>) {
if st.lru.len() >= st.cap {
if let Some(&victim) = st
.lru
.iter()
.min_by_key(|(_, e)| e.used)
.map(|(k, _)| k)
{
st.lru.remove(&victim);
}
}
st.clock += 1;
let used = st.clock;
st.lru.insert(eid, LruEntry { bytes, used });
}
/// Positional read of one expert's bytes (thread-safe, no shared seek), then advise the kernel
/// to drop those file pages so the page cache stays bounded (colibri's `POSIX_FADV_DONTNEED`).
fn read_expert(&self, eid: u32) -> Result<Arc<[u8]>> {
let off = self.base_offset + eid as u64 * self.expert_bytes as u64;
let mut buf = vec![0u8; self.expert_bytes];
read_exact_at(&self.file, &mut buf, off)?;
fadvise_dontneed(&self.file, off, self.expert_bytes as u64);
Ok(Arc::from(buf.into_boxed_slice()))
}
/// Pin expert `eid` into the hot-set (loads it if absent). Idempotent.
pub fn pin(&self, eid: u32) -> Result<()> {
let mut st = self.inner.lock().expect("expert bank lock poisoned");
if st.pinned.contains_key(&eid) {
return Ok(());
}
let bytes = if let Some(entry) = st.lru.remove(&eid) {
entry.bytes
} else {
drop(st);
let b = self.read_expert(eid)?;
st = self.inner.lock().expect("expert bank lock poisoned");
b
};
st.pinned.insert(eid, bytes);
Ok(())
}
/// Set the LRU cap (max non-pinned experts held resident). Sized by [`finalize`] from RAM; also
/// settable directly (tests, manual budgets).
pub fn set_cap(&self, cap: usize) {
let mut st = self.inner.lock().expect("expert bank lock poisoned");
st.cap = cap.max(1);
}
/// `(pinned, cached, cap, hits, misses)` snapshot for the resident-RAM / hit-rate proof.
pub fn stats(&self) -> (usize, usize, usize, u64, u64) {
let st = self.inner.lock().expect("expert bank lock poisoned");
(st.pinned.len(), st.lru.len(), st.cap, st.hits, st.misses)
}
/// Bytes currently held resident by this bank (pinned + cached slabs).
pub fn resident_bytes(&self) -> usize {
let st = self.inner.lock().expect("expert bank lock poisoned");
(st.pinned.len() + st.lru.len()) * self.expert_bytes
}
fn usage_snapshot(&self) -> Vec<(u32, u64)> {
let st = self.inner.lock().expect("expert bank lock poisoned");
st.usage
.iter()
.enumerate()
.filter(|(_, &c)| c > 0)
.map(|(e, &c)| (e as u32, c))
.collect()
}
}
/// Global registry: every live bank + the learned-usage sidecar path. Lets one budget size all
/// banks together (colibri `cap_for_ram`) and one pass learn/pin the hottest experts.
struct Registry {
banks: Vec<Weak<ExpertStreamBank>>,
sidecar: Option<PathBuf>,
}
fn registry() -> &'static Mutex<Registry> {
static REG: OnceLock<Mutex<Registry>> = OnceLock::new();
REG.get_or_init(|| {
Mutex::new(Registry {
banks: Vec::new(),
sidecar: None,
})
})
}
impl Registry {
fn register(&mut self, bank: &Arc<ExpertStreamBank>) {
self.banks.push(Arc::downgrade(bank));
}
fn live(&self) -> Vec<Arc<ExpertStreamBank>> {
self.banks.iter().filter_map(Weak::upgrade).collect()
}
}
/// Point the learning-pin at a sidecar file (`<model_dir>/.hanzo_experts_usage`).
pub fn set_usage_sidecar(path: PathBuf) {
registry().lock().expect("registry lock").sidecar = Some(path);
}
/// Total bytes held resident across all live streaming banks (the low-RAM proof).
pub fn total_resident_bytes() -> usize {
registry()
.lock()
.expect("registry lock")
.live()
.iter()
.map(|b| b.resident_bytes())
.sum()
}
/// Size every bank's LRU cap from available RAM and auto-pin the hottest experts from the sidecar.
/// Call once after the dense weights are resident and all expert banks are open.
pub fn finalize() {
let banks = registry().lock().expect("registry lock").live();
if banks.is_empty() {
return;
}
// Largest per-expert stride across banks -> the conservative divisor (banks are near-equal).
let expert_bytes = banks.iter().map(|b| b.expert_bytes).max().unwrap_or(1).max(1);
let n_banks = banks.len() as u64;
let avail = mem_available_bytes();
let budget = (avail as f64 * BUDGET_FRACTION) as u64;
let slack = PAGE_CACHE_RESERVE + ACTIVATION_RESERVE;
let for_cache = budget.saturating_sub(slack);
let mut cap = (for_cache / (n_banks * expert_bytes as u64)) as usize;
// Never below 1 (a miss must be admittable), never above n_experts (a bank can't fill more).
let max_experts = banks.iter().map(|b| b.n_experts).max().unwrap_or(1);
cap = cap.clamp(1, max_experts);
for b in &banks {
b.set_cap(cap);
}
let pinned = load_and_pin(&banks, cap);
eprintln!(
"[stream-experts] {} banks x {:.1} MB/expert; MemAvailable {:.1} GB -> per-bank cap {} \
(projected cache {:.1} GB), auto-pinned {} experts",
banks.len(),
expert_bytes as f64 / 1e6,
avail as f64 / 1e9,
cap,
(cap as u64 * n_banks * expert_bytes as u64) as f64 / 1e9,
pinned,
);
}
/// Read the usage sidecar (if any) and pin each bank's hottest experts, up to `PIN_FRACTION * cap`.
fn load_and_pin(banks: &[Arc<ExpertStreamBank>], cap: usize) -> usize {
let sidecar = match registry().lock().expect("registry lock").sidecar.clone() {
Some(p) if p.exists() => p,
_ => return 0,
};
let text = match std::fs::read_to_string(&sidecar) {
Ok(t) => t,
Err(_) => return 0,
};
// `name eid count` per line.
let mut by_name: HashMap<&str, Vec<(u32, u64)>> = HashMap::new();
for line in text.lines() {
let mut it = line.split_whitespace();
let (Some(name), Some(eid), Some(cnt)) = (it.next(), it.next(), it.next()) else {
continue;
};
if let (Ok(eid), Ok(cnt)) = (eid.parse::<u32>(), cnt.parse::<u64>()) {
by_name.entry(name).or_default().push((eid, cnt));
}
}
let pin_budget = ((cap as f64 * PIN_FRACTION) as usize).max(1);
let mut pinned = 0usize;
for b in banks {
let Some(rows) = by_name.get(b.name.as_str()) else {
continue;
};
let mut rows = rows.clone();
rows.sort_by(|a, c| c.1.cmp(&a.1));
for (eid, _) in rows.into_iter().take(pin_budget) {
if eid as usize >= b.n_experts {
continue;
}
if b.pin(eid).is_ok() {
pinned += 1;
}
}
}
pinned
}
/// Persist this run's per-expert routing frequencies (merged with any prior counts) so the next
/// startup can AUTOPIN the hottest. Call at shutdown.
pub fn save_usage() -> Result<()> {
let (banks, sidecar) = {
let reg = registry().lock().expect("registry lock");
(reg.live(), reg.sidecar.clone())
};
let Some(path) = sidecar else {
return Ok(());
};
// Merge with existing counts so the learned distribution accumulates across runs.
let mut merged: HashMap<(String, u32), u64> = HashMap::new();
if let Ok(text) = std::fs::read_to_string(&path) {
for line in text.lines() {
let mut it = line.split_whitespace();
if let (Some(name), Some(eid), Some(cnt)) = (it.next(), it.next(), it.next()) {
if let (Ok(eid), Ok(cnt)) = (eid.parse::<u32>(), cnt.parse::<u64>()) {
*merged.entry((name.to_string(), eid)).or_default() += cnt;
}
}
}
}
for b in &banks {
for (eid, cnt) in b.usage_snapshot() {
*merged.entry((b.name.clone(), eid)).or_default() += cnt;
}
}
let mut out = String::new();
for ((name, eid), cnt) in merged {
out.push_str(&format!("{name} {eid} {cnt}\n"));
}
std::fs::write(&path, out)?;
Ok(())
}
// ---- platform: positional read + page-cache drop ---------------------------------------------
#[cfg(unix)]
fn read_exact_at(file: &File, buf: &mut [u8], offset: u64) -> Result<()> {
use std::os::unix::fs::FileExt;
file.read_exact_at(buf, offset)?;
Ok(())
}
#[cfg(not(unix))]
fn read_exact_at(_file: &File, _buf: &mut [u8], _offset: u64) -> Result<()> {
crate::bail!("streaming experts require a unix positional-read (pread) platform")
}
/// Advise the kernel to drop the just-read file pages so the page cache can't grow unbounded while
/// we stream the whole model past it. Best-effort: a failure only means more page cache, not wrong
/// results. Offset/len need no alignment for `POSIX_FADV_DONTNEED`.
#[cfg(unix)]
fn fadvise_dontneed(file: &File, offset: u64, len: u64) {
use std::os::unix::io::AsRawFd;
// SAFETY: fd is valid for the call; posix_fadvise only advises, never mutates our memory.
unsafe {
libc::posix_fadvise(
file.as_raw_fd(),
offset as libc::off_t,
len as libc::off_t,
libc::POSIX_FADV_DONTNEED,
);
}
}
#[cfg(not(unix))]
fn fadvise_dontneed(_file: &File, _offset: u64, _len: u64) {}
/// MemAvailable in bytes. Linux reads `/proc/meminfo` (the true reclaimable ceiling); elsewhere a
/// conservative fallback keeps the cache small rather than risking OOM.
fn mem_available_bytes() -> u64 {
#[cfg(target_os = "linux")]
{
if let Ok(text) = std::fs::read_to_string("/proc/meminfo") {
for line in text.lines() {
if let Some(rest) = line.strip_prefix("MemAvailable:") {
// `MemAvailable: 12345678 kB`
if let Some(kb) = rest.split_whitespace().next().and_then(|v| v.parse::<u64>().ok())
{
return kb.saturating_mul(1024);
}
}
}
}
}
// Fallback (non-linux or unreadable): assume 8 GB free, matching colibri's floor.
static WARNED: AtomicU64 = AtomicU64::new(0);
if WARNED.swap(1, Ordering::Relaxed) == 0 {
eprintln!("[stream-experts] MemAvailable unreadable; assuming 8 GB free");
}
8_000_000_000
}
+46
View File
@@ -195,6 +195,43 @@ impl TensorInfo {
),
}
}
/// Disk-streaming twin for a stacked MoE expert bank (`[n_experts, n, k]`). Instead of reading
/// the whole bank resident, returns a `QStorage::Stream` that keeps the bank on `path` and preads
/// one expert's `[n, k]` slice on demand through the pin/LRU cache (see `expert_stream`). Only
/// valid for a rank-3 tensor on CPU; the caller falls back to `read` for anything else.
pub fn read_stream(
&self,
path: &std::path::Path,
tensor_data_offset: u64,
name: &str,
) -> Result<QTensor> {
let dims = self.shape.dims();
if dims.len() != 3 {
crate::bail!("read_stream: expected a rank-3 expert bank, got shape {:?}", self.shape)
}
let (n_experts, n, k) = (dims[0], dims[1], dims[2]);
let block_size = self.ggml_dtype.block_size();
let per_expert_elems = n * k;
if !per_expert_elems.is_multiple_of(block_size) {
crate::bail!(
"read_stream: per-expert elems {per_expert_elems} not divisible by block size {block_size}"
)
}
let expert_bytes = per_expert_elems / block_size * self.ggml_dtype.type_size();
let base_offset = tensor_data_offset.saturating_add(self.offset);
let bank = super::expert_stream::ExpertStreamBank::open(
name.to_string(),
path,
base_offset,
expert_bytes,
n_experts,
self.ggml_dtype,
n,
k,
)?;
QTensor::new(super::QStorage::Stream(bank), self.shape.dims().to_vec())
}
}
#[derive(Debug)]
@@ -696,6 +733,15 @@ impl Content {
};
tensor_info.read_mmap(mmap, self.tensor_data_offset, device)
}
/// Load a stacked expert bank as a disk-streaming `QTensor` (see [`TensorInfo::read_stream`]).
pub fn tensor_stream(&self, path: &std::path::Path, name: &str) -> Result<QTensor> {
let tensor_info = match self.tensor_infos.get(name) {
Some(tensor_info) => tensor_info,
None => crate::bail!("cannot find tensor info for {name}"),
};
tensor_info.read_stream(path, self.tensor_data_offset, name)
}
}
fn write_string<W: std::io::Write>(w: &mut W, str: &str) -> Result<()> {
+87 -40
View File
@@ -11,6 +11,7 @@ pub mod avx;
pub mod dsv4_qat;
mod dummy_cuda;
mod dummy_metal;
pub mod expert_stream;
pub mod ggml_file;
pub mod gguf_file;
pub mod imatrix_file;
@@ -131,6 +132,10 @@ pub enum QStorage {
// dequantizes to an f32 wgpu tensor on demand.
#[cfg(feature = "wgpu")]
Wgpu(Box<dyn QuantizedType>, crate::WgpuDevice),
// Disk-streaming MoE expert bank: the stacked [n_experts, n, k] quantized weight lives on NVMe,
// not resident. Only `indexed_moe_forward` consumes it -- it fetches one expert's [n, k] slice
// through the pin/LRU cache per token (host, CPU matmul). The low-memory mode for huge MoEs.
Stream(Arc<expert_stream::ExpertStreamBank>),
}
impl QStorage {
@@ -231,6 +236,7 @@ impl QStorage {
QStorage::Vulkan(storage, _) => storage.block_size(),
#[cfg(feature = "wgpu")]
QStorage::Wgpu(storage, _) => storage.block_size(),
QStorage::Stream(bank) => bank.dtype().block_size(),
}
}
@@ -245,6 +251,7 @@ impl QStorage {
QStorage::Vulkan(storage, _) => storage.dtype(),
#[cfg(feature = "wgpu")]
QStorage::Wgpu(storage, _) => storage.dtype(),
QStorage::Stream(bank) => bank.dtype(),
}
}
@@ -259,6 +266,7 @@ impl QStorage {
QStorage::Vulkan(_storage, device) => Device::Vulkan(device.clone()),
#[cfg(feature = "wgpu")]
QStorage::Wgpu(_storage, device) => Device::Wgpu(device.clone()),
QStorage::Stream(_) => Device::Cpu,
}
}
@@ -273,6 +281,7 @@ impl QStorage {
QStorage::Vulkan(storage, _) => storage.storage_size_in_bytes(),
#[cfg(feature = "wgpu")]
QStorage::Wgpu(storage, _) => storage.storage_size_in_bytes(),
QStorage::Stream(bank) => bank.logical_bytes(),
}
}
@@ -366,6 +375,9 @@ impl QStorage {
let cpu = storage.dequantize(elem_count)?;
Ok(Storage::Wgpu(device.upload_f32(cpu.as_slice::<f32>()?)?))
}
QStorage::Stream(_) => {
crate::bail!("streaming expert bank has no whole-tensor dequantize; consume it via indexed_moe_forward")
}
}
}
@@ -400,6 +412,9 @@ impl QStorage {
let data = unsafe { std::slice::from_raw_parts(data_ptr, size_in_bytes) };
Ok(Cow::from(data))
}
QStorage::Stream(_) => {
crate::bail!("streaming expert bank is not resident; consume it via indexed_moe_forward")
}
}
}
@@ -412,7 +427,7 @@ impl QStorage {
QStorage::Vulkan(..) => crate::bail!("not implemented"),
#[cfg(feature = "wgpu")]
QStorage::Wgpu(..) => crate::bail!("not implemented"),
QStorage::Metal(_) | QStorage::Cpu(_) => {
QStorage::Metal(_) | QStorage::Cpu(_) | QStorage::Stream(_) => {
crate::bail!("not implemented");
}
}
@@ -428,7 +443,7 @@ impl QStorage {
)> {
match self {
QStorage::Cuda(storage) => storage.device_ptr_with_guard(stream),
QStorage::Metal(_) | QStorage::Cpu(_) => {
QStorage::Metal(_) | QStorage::Cpu(_) | QStorage::Stream(_) => {
crate::bail!("not implemented");
}
}
@@ -1380,56 +1395,38 @@ impl QTensor {
);
out.to_dtype(out_dtype)
}
// Disk-streaming bank: same per-expert quantized matmul as the resident fallback below,
// but each selected expert's [n, k] slice is fetched through the pin/LRU cache (disk)
// instead of sliced out of a resident blob. Bit-identical -- only the fetch path differs.
QStorage::Stream(bank) => {
let (e_cnt, n, k) = self.shape().dims3()?;
let dtype = bank.dtype();
moe_grouped_per_expert(x, ids, n, k, |eid, device| {
if eid as usize >= e_cnt {
crate::bail!("indexed_moe_forward: expert id {eid} >= num_experts {e_cnt}");
}
let bytes = bank.fetch(eid)?;
QStorage::from_data(std::borrow::Cow::Borrowed(&bytes), device, dtype)
})
}
_ => {
// CPU / non-CUDA fallback: per-expert quantized matmul. The packed expert bank
// [E, n, k] is sliced into equal, contiguous per-expert quantized blocks; for
// each expert that is actually selected we run hanzo-ml's native quantized matmul
// on just the tokens routed to it. Nothing is dequantized, so quantized MoE runs
// on any backend (CPU, Metal, ...) at a cost proportional to the active experts.
use crate::Module; // brings QMatMul::forward into scope
use std::collections::HashMap;
use std::sync::Arc;
let device = x.device();
let out_dtype = x.dtype();
let (e_cnt, n, k) = self.shape().dims3()?;
let (t, topk) = ids.dims2()?;
let s = x.dim(1)?; // 1 (gate/up: shared input) or topk (down: per-slot)
let x_exp = if s == topk {
x.clone()
} else {
x.broadcast_as((t, topk, k))?
};
let x_flat = x_exp
.reshape((t * topk, k))?
.to_dtype(crate::DType::F32)?
.contiguous()?;
let ids_flat = ids.reshape((t * topk,))?.to_dtype(crate::DType::U32)?;
let ids_vec = ids_flat.to_vec1::<u32>()?;
let mut groups: HashMap<u32, Vec<u32>> = HashMap::new();
for (slot, eid) in ids_vec.iter().enumerate() {
groups.entry(*eid).or_default().push(slot as u32);
}
let dtype = self.storage.dtype();
let all_bytes = self.data()?;
let expert_bytes = all_bytes.len() / e_cnt;
let mut out_flat = Tensor::zeros((t * topk, n), crate::DType::F32, device)?;
for (eid, slots) in groups.into_iter() {
moe_grouped_per_expert(x, ids, n, k, |eid, device| {
let off = eid as usize * expert_bytes;
let qs = QStorage::from_data(
QStorage::from_data(
std::borrow::Cow::Borrowed(&all_bytes[off..off + expert_bytes]),
device,
dtype,
)?;
let shape: crate::Shape = (n, k).into();
let w_e = QTensor { storage: qs, shape };
let qm = QMatMul::from_arc(Arc::new(w_e))?;
let m = slots.len();
let idx = Tensor::from_vec(slots, (m,), device)?;
let x_e = x_flat.index_select(&idx, 0)?; // [m, k]
let y_e = qm.forward(&x_e)?.to_dtype(crate::DType::F32)?; // [m, n]
out_flat = out_flat.index_add(&idx, &y_e, 0)?;
}
out_flat.reshape((t, topk, n))?.to_dtype(out_dtype)
)
})
}
}
}
@@ -1443,7 +1440,7 @@ impl QTensor {
QStorage::Vulkan(..) => crate::bail!("not implemented"),
#[cfg(feature = "wgpu")]
QStorage::Wgpu(..) => crate::bail!("not implemented"),
QStorage::Metal(_) | QStorage::Cpu(_) => {
QStorage::Metal(_) | QStorage::Cpu(_) | QStorage::Stream(_) => {
crate::bail!("not implemented");
}
}
@@ -1521,6 +1518,54 @@ fn rocm_moe_bank_cache() -> &'static std::sync::Mutex<
CACHE.get_or_init(|| std::sync::Mutex::new(std::collections::HashMap::new()))
}
/// Per-expert quantized MoE matmul, shared by the resident and disk-streaming paths. Groups the
/// routed slots by expert id, and for each active expert builds a `QTensor` from `make_storage`
/// (resident slice or streamed slab), runs the native quantized matmul on just that expert's tokens,
/// and scatters the result back. The ONLY difference between resident and streaming is the closure.
fn moe_grouped_per_expert(
x: &Tensor,
ids: &Tensor,
n: usize,
k: usize,
mut make_storage: impl FnMut(u32, &Device) -> Result<QStorage>,
) -> Result<Tensor> {
use crate::Module; // brings QMatMul::forward into scope
use std::collections::HashMap;
use std::sync::Arc;
let device = x.device();
let out_dtype = x.dtype();
let (t, topk) = ids.dims2()?;
let s = x.dim(1)?; // 1 (gate/up: shared input) or topk (down: per-slot)
let x_exp = if s == topk {
x.clone()
} else {
x.broadcast_as((t, topk, k))?
};
let x_flat = x_exp
.reshape((t * topk, k))?
.to_dtype(DType::F32)?
.contiguous()?;
let ids_flat = ids.reshape((t * topk,))?.to_dtype(DType::U32)?;
let ids_vec = ids_flat.to_vec1::<u32>()?;
let mut groups: HashMap<u32, Vec<u32>> = HashMap::new();
for (slot, eid) in ids_vec.iter().enumerate() {
groups.entry(*eid).or_default().push(slot as u32);
}
let mut out_flat = Tensor::zeros((t * topk, n), DType::F32, device)?;
for (eid, slots) in groups.into_iter() {
let qs = make_storage(eid, device)?;
let shape: crate::Shape = (n, k).into();
let w_e = QTensor { storage: qs, shape };
let qm = QMatMul::from_arc(Arc::new(w_e))?;
let m = slots.len();
let idx = Tensor::from_vec(slots, (m,), device)?;
let x_e = x_flat.index_select(&idx, 0)?; // [m, k]
let y_e = qm.forward(&x_e)?.to_dtype(DType::F32)?; // [m, n]
out_flat = out_flat.index_add(&idx, &y_e, 0)?;
}
out_flat.reshape((t, topk, n))?.to_dtype(out_dtype)
}
/// FUSED MoE expert-combine: `out[i,j] = sum_e scores[i,e] * ys[i,e,j]`, reducing the per-expert
/// outputs `ys` [t, topk, n] by the router weights `scores` [t, topk] into [t, n]. On ROCm this is
/// ONE fused kernel (`RocmDevice::moe_combine`) instead of `ys.broadcast_mul(scores).sum(Minus2)`,
@@ -2050,7 +2095,9 @@ impl crate::CustomOp1 for QTensor {
QStorage::Vulkan(..) => crate::bail!("Invalid storage"),
#[cfg(feature = "wgpu")]
QStorage::Wgpu(..) => crate::bail!("Invalid storage"),
QStorage::Metal(_) | QStorage::Cuda(_) => crate::bail!("Invalid storage"),
QStorage::Metal(_) | QStorage::Cuda(_) | QStorage::Stream(_) => {
crate::bail!("Invalid storage")
}
};
match storage.dtype() {
DType::F32 => {
+173
View File
@@ -0,0 +1,173 @@
// Disk-streaming LRU expert cache: bit-exact vs resident + cache mechanics.
//
// The load-bearing guarantee: a streamed expert is the identical file bytes the resident slice
// would hold, so `indexed_moe_forward` produces bit-identical output; only the fetch path differs.
use hanzo_ml::quantized::expert_stream::ExpertStreamBank;
use hanzo_ml::quantized::gguf_file::TensorInfo;
use hanzo_ml::quantized::{GgmlDType, QStorage, QTensor};
use hanzo_ml::{Device, Result, Tensor};
const E: usize = 6; // experts in the bank
const N: usize = 8; // expert out dim
const K: usize = 64; // expert in dim (2 Q8_0 blocks)
const T: usize = 4; // tokens
const TOPK: usize = 2;
/// Build a resident Q8_0 expert bank `[E, N, K]`, plus its raw bytes written to a temp file so a
/// streaming bank can be opened over the same bytes.
fn fixture() -> Result<(Device, QTensor, std::path::PathBuf, usize)> {
let dev = Device::Cpu;
let src = Tensor::randn(0f32, 1f32, (E, N, K), &dev)?;
let resident = QTensor::quantize(&src, GgmlDType::Q8_0)?;
let bytes = resident.data()?.to_vec();
let expert_bytes = bytes.len() / E;
let path = std::env::temp_dir().join(format!(
"hanzo_stream_{}_{:p}.bin",
std::process::id(),
&bytes as *const _
));
std::fs::write(&path, &bytes)?;
Ok((dev, resident, path, expert_bytes))
}
fn routing(dev: &Device) -> Result<(Tensor, Tensor)> {
// Shared gate/up-style input: x is [T, 1, K]; router picks TOPK of E experts per token.
let x = Tensor::randn(0f32, 1f32, (T, 1, K), dev)?;
let ids: Vec<u32> = vec![0, 1, 2, 3, 4, 5, 1, 0]; // T*TOPK, spans all experts
let ids = Tensor::from_vec(ids, (T, TOPK), dev)?;
Ok((x, ids))
}
fn stream_qtensor(path: &std::path::Path, expert_bytes: usize, name: &str) -> Result<(QTensor, std::sync::Arc<ExpertStreamBank>)> {
let bank = ExpertStreamBank::open(
name.to_string(),
path,
0,
expert_bytes,
E,
GgmlDType::Q8_0,
N,
K,
)?;
let qt = QTensor::new(QStorage::Stream(bank.clone()), (E, N, K))?;
Ok((qt, bank))
}
#[test]
fn streamed_moe_is_bit_exact_vs_resident() -> Result<()> {
let (dev, resident, path, _expert_bytes) = fixture()?;
let (x, ids) = routing(&dev)?;
// Resident path via QStorage::Cpu -> the generic `_` arm.
let yr = resident.indexed_moe_forward(&x, &ids)?.flatten_all()?.to_vec1::<f32>()?;
// Streaming path via read_stream (the loader seam) -> the `Stream` arm.
let info = TensorInfo {
ggml_dtype: GgmlDType::Q8_0,
shape: (E, N, K).into(),
offset: 0,
};
let streamed = info.read_stream(&path, 0, "blk.0.ffn_gate_exps.weight")?;
let ys = streamed.indexed_moe_forward(&x, &ids)?.flatten_all()?.to_vec1::<f32>()?;
assert_eq!(yr.len(), ys.len());
for (i, (a, b)) in yr.iter().zip(ys.iter()).enumerate() {
assert_eq!(
a.to_bits(),
b.to_bits(),
"streamed vs resident differ at {i}: {a} != {b}"
);
}
let _ = std::fs::remove_file(&path);
Ok(())
}
#[test]
fn lru_cap_bounds_resident_and_stays_bit_exact() -> Result<()> {
let (dev, resident, path, expert_bytes) = fixture()?;
let (x, ids) = routing(&dev)?;
let yr = resident.indexed_moe_forward(&x, &ids)?.flatten_all()?.to_vec1::<f32>()?;
let (qt, bank) = stream_qtensor(&path, expert_bytes, "cap.bank")?;
bank.set_cap(1); // pathologically small: at most 1 non-pinned expert resident
// Two forwards touch all E experts repeatedly; the cache must never exceed the cap...
for _ in 0..3 {
let ys = qt.indexed_moe_forward(&x, &ids)?.flatten_all()?.to_vec1::<f32>()?;
// ...and correctness must be independent of the cap (streamed == resident every time).
for (a, b) in yr.iter().zip(ys.iter()) {
assert_eq!(a.to_bits(), b.to_bits());
}
}
let (pinned, cached, cap, _hits, _misses) = bank.stats();
assert_eq!(pinned, 0);
assert_eq!(cap, 1);
assert!(cached <= 1, "LRU held {cached} experts, cap was 1");
assert!(bank.resident_bytes() <= expert_bytes, "resident exceeded 1 expert");
let _ = std::fs::remove_file(&path);
Ok(())
}
#[test]
fn pinned_expert_is_a_hit_and_survives_eviction() -> Result<()> {
let (_dev, _resident, path, expert_bytes) = fixture()?;
let (qt, bank) = stream_qtensor(&path, expert_bytes, "pin.bank")?;
bank.set_cap(1);
bank.pin(3)?; // pin expert 3 into the hot-set
let dev = Device::Cpu;
// Drive many forwards that route through OTHER experts, thrashing the size-1 LRU.
let x = Tensor::randn(0f32, 1f32, (2, 1, K), &dev)?;
let ids = Tensor::from_vec(vec![0u32, 1, 2, 4], (2, 2), &dev)?;
for _ in 0..5 {
let _ = qt.indexed_moe_forward(&x, &ids)?;
}
let (pinned_before, _c, _cap, hits_before, _m) = bank.stats();
assert_eq!(pinned_before, 1, "expert 3 should stay pinned through eviction");
// Fetching the pinned expert is a cache HIT (no disk read), regardless of LRU pressure.
let _ = bank.fetch(3)?;
let (_p, _c, _cap, hits_after, _m2) = bank.stats();
assert_eq!(hits_after, hits_before + 1, "pinned fetch must be a hit");
let _ = std::fs::remove_file(&path);
Ok(())
}
#[test]
fn usage_sidecar_learns_the_hot_expert() -> Result<()> {
use hanzo_ml::quantized::expert_stream;
let (_dev, _resident, path, expert_bytes) = fixture()?;
let (_qt, bank) = stream_qtensor(&path, expert_bytes, "learn.bank")?;
// Route expert 4 much more than the rest, then persist the learned usage.
for _ in 0..10 {
let _ = bank.fetch(4)?;
}
let _ = bank.fetch(0)?;
let _ = bank.fetch(1)?;
let sidecar = std::env::temp_dir().join(format!("hanzo_usage_{}.txt", std::process::id()));
let _ = std::fs::remove_file(&sidecar);
expert_stream::set_usage_sidecar(sidecar.clone());
expert_stream::save_usage()?;
let text = std::fs::read_to_string(&sidecar)?;
// Among this bank's lines, expert 4 must carry the highest count.
let mut best: Option<(u32, u64)> = None;
for line in text.lines() {
let mut it = line.split_whitespace();
if it.next() == Some("learn.bank") {
if let (Some(e), Some(c)) = (it.next(), it.next()) {
let (e, c) = (e.parse::<u32>().unwrap(), c.parse::<u64>().unwrap());
if best.map(|(_, bc)| c > bc).unwrap_or(true) {
best = Some((e, c));
}
}
}
}
assert_eq!(best.map(|(e, _)| e), Some(4), "hottest expert should be 4");
let _ = std::fs::remove_file(&path);
let _ = std::fs::remove_file(&sidecar);
Ok(())
}