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:
@@ -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
@@ -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",
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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<()> {
|
||||
|
||||
@@ -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 => {
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
Reference in New Issue
Block a user