use std::cell::Cell;
use std::collections::{HashMap, HashSet, VecDeque};
use std::fs::File;
#[cfg(any(target_os = "android", target_os = "linux"))]
use std::fs::OpenOptions;
use std::io;
use std::ops::Range;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::slice;
#[cfg(any(target_os = "android", target_os = "linux"))]
use std::sync::OnceLock;
#[cfg(feature = "bench")]
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Condvar, Mutex, MutexGuard};
use std::thread::{self, JoinHandle};
use anyhow::{Context as _, Result, anyhow, ensure};
use memmap2::{MmapMut, MmapOptions};
use super::gguf::ByteRange;
const IO_WORKERS: usize = 2;
const EXPERT_RANGE_COUNT: usize = 6;
const SLOT_CHARGE_ALIGNMENT: usize = 64 * 1024;
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub(super) struct ExpertKey {
layer: usize,
expert: usize,
}
impl ExpertKey {
pub(super) const fn new(layer: usize, expert: usize) -> Self {
Self { layer, expert }
}
}
struct ExpertData {
bytes: MmapMut,
parts: [Range<usize>; 6],
}
struct ExpertView {
ptr: *const u8,
len: usize,
parts: [Range<usize>; 6],
}
unsafe impl Send for ExpertView {}
unsafe impl Sync for ExpertView {}
impl ExpertData {
fn view(&self) -> ExpertView {
ExpertView {
ptr: self.bytes.as_ptr(),
len: self.bytes.len(),
parts: self.parts.clone(),
}
}
}
impl ExpertView {
fn part(&self, index: usize) -> Result<&[u8]> {
let range = self
.parts
.get(index)
.ok_or_else(|| anyhow!("expert part {index} is out of range"))?;
ensure!(
range.start <= range.end && range.end <= self.len,
"expert part range {}..{} exceeds mapping length {}",
range.start,
range.end,
self.len
);
Ok(unsafe { slice::from_raw_parts(self.ptr.add(range.start), range.len()) })
}
}
pub(super) struct ExpertLease {
key: ExpertKey,
generation: u64,
data: Option<ExpertView>,
shared: Arc<Shared>,
}
impl ExpertLease {
pub(super) fn part(&self, index: usize) -> Result<&[u8]> {
self.data
.as_ref()
.ok_or_else(|| anyhow!("expert lease has been released"))?
.part(index)
}
}
impl Drop for ExpertLease {
fn drop(&mut self) {
let mut state = lock(&self.shared.state);
self.data.take();
release_pin(&mut state, self.key, self.generation);
self.shared.changed.notify_all();
}
}
pub(super) struct ExpertRequest {
key: ExpertKey,
generation: u64,
active: bool,
shared: Arc<Shared>,
}
impl ExpertRequest {
pub(super) fn wait(mut self) -> Result<ExpertLease> {
let data = loop {
let mut state = lock(&self.shared.state);
match state.entries.get(&self.key) {
Some(CacheEntry::Ready {
data, generation, ..
}) if *generation == self.generation => break data.view(),
Some(CacheEntry::Loading { generation, .. })
if *generation == self.generation && !state.shutdown =>
{
drop(wait(&self.shared.changed, state));
}
Some(CacheEntry::Loading { generation, .. }) if *generation == self.generation => {
release_pin(&mut state, self.key, self.generation);
self.active = false;
self.shared.changed.notify_all();
return Err(anyhow!("expert cache is shut down"));
}
Some(CacheEntry::Failed {
generation, error, ..
}) if *generation == self.generation => {
let error = Arc::clone(error);
release_pin(&mut state, self.key, self.generation);
self.active = false;
self.shared.changed.notify_all();
return Err(anyhow!(error.to_string()));
}
Some(
CacheEntry::Ready { .. }
| CacheEntry::Loading { .. }
| CacheEntry::Failed { .. },
) => return Err(anyhow!("expert cache entry changed while requested")),
None => return Err(anyhow!("loaded expert is missing from the cache")),
}
};
let lease = ExpertLease {
key: self.key,
generation: self.generation,
data: Some(data),
shared: Arc::clone(&self.shared),
};
self.active = false;
Ok(lease)
}
}
impl Drop for ExpertRequest {
fn drop(&mut self) {
if !self.active {
return;
}
let mut state = lock(&self.shared.state);
release_pin(&mut state, self.key, self.generation);
self.shared.changed.notify_all();
}
}
enum CacheEntry {
Loading {
generation: u64,
pins: Cell<usize>,
},
Ready {
generation: u64,
data: ExpertData,
last_used: Cell<u64>,
pins: Cell<usize>,
},
Failed {
generation: u64,
error: Arc<str>,
pins: Cell<usize>,
},
}
struct LoadJob {
key: ExpertKey,
generation: u64,
ranges: [ByteRange; 6],
buffer: MmapMut,
}
struct StoreState {
entries: HashMap<ExpertKey, CacheEntry>,
queue: VecDeque<LoadJob>,
free_buffers: Vec<MmapMut>,
allocated_slots: usize,
max_slots: usize,
slot_bytes: usize,
clock: u64,
next_generation: u64,
shutdown: bool,
}
struct Shared {
state: Mutex<StoreState>,
work_available: Condvar,
changed: Condvar,
#[cfg(feature = "bench")]
queued_loads: AtomicUsize,
#[cfg(feature = "bench")]
slot_waits: AtomicUsize,
}
struct DirectFile {
file: File,
alignment: usize,
file_len: u64,
}
enum DirectLoad {
Complete([Range<usize>; 6]),
Buffered,
Unavailable,
}
struct ExpertSource {
buffered: File,
direct: Option<DirectFile>,
direct_enabled: AtomicBool,
}
impl ExpertSource {
fn new(buffered: File) -> Self {
let direct = open_direct_file(&buffered);
let direct_enabled = AtomicBool::new(direct.is_some());
Self {
buffered,
direct,
direct_enabled,
}
}
fn slot_capacity(&self, data_bytes: usize) -> Result<usize> {
let Some(direct) = &self.direct else {
return Ok(data_bytes);
};
let padding = direct
.alignment
.checked_mul(2 * EXPERT_RANGE_COUNT)
.context("expert direct-I/O padding overflow")?;
data_bytes
.checked_add(padding)
.context("expert cache slot size overflow")
}
}
pub(super) struct ExpertStore {
shared: Arc<Shared>,
request_gate: Mutex<()>,
workers: Vec<JoinHandle<()>>,
}
impl ExpertStore {
pub(super) fn new(
file: File,
slot_bytes: usize,
byte_budget: usize,
minimum_slots: usize,
maximum_slots: usize,
) -> Result<Self> {
ensure!(slot_bytes != 0, "expert cache slot must not be empty");
ensure!(
cfg!(any(unix, windows)),
"expert cache requires positional file reads"
);
ensure!(
minimum_slots != 0,
"expert cache requires at least one slot"
);
ensure!(
minimum_slots <= maximum_slots,
"expert cache slot bounds are invalid"
);
let source = Arc::new(ExpertSource::new(file));
let slot_bytes = source.slot_capacity(slot_bytes)?;
let slot_charge = align_up(slot_bytes, SLOT_CHARGE_ALIGNMENT)?;
let max_slots = (byte_budget / slot_charge).min(maximum_slots);
ensure!(
max_slots >= minimum_slots,
"expert cache budget provides {max_slots} slots, but inference requires at least {minimum_slots}"
);
let shared = Arc::new(Shared {
state: Mutex::new(StoreState {
entries: HashMap::new(),
queue: VecDeque::new(),
free_buffers: Vec::new(),
allocated_slots: 0,
max_slots,
slot_bytes,
clock: 0,
next_generation: 1,
shutdown: false,
}),
work_available: Condvar::new(),
changed: Condvar::new(),
#[cfg(feature = "bench")]
queued_loads: AtomicUsize::new(0),
#[cfg(feature = "bench")]
slot_waits: AtomicUsize::new(0),
});
let mut store = Self {
shared,
request_gate: Mutex::new(()),
workers: Vec::new(),
};
for index in 0..IO_WORKERS.min(max_slots) {
let source = Arc::clone(&source);
let shared = Arc::clone(&store.shared);
let worker = thread::Builder::new()
.name(format!("expert-io-{index}"))
.spawn(move || worker_loop(&shared, &source))
.context("start expert I/O worker")?;
store.workers.push(worker);
}
Ok(store)
}
#[cfg(feature = "bench")]
pub(super) fn queued_load_count(&self) -> usize {
self.shared.queued_loads.load(Ordering::Relaxed)
}
#[cfg(feature = "bench")]
pub(super) fn slot_wait_count(&self) -> usize {
self.shared.slot_waits.load(Ordering::Relaxed)
}
pub(super) fn request_many(
&self,
requests: impl IntoIterator<Item = (ExpertKey, [ByteRange; 6])>,
) -> Result<Vec<ExpertRequest>> {
let requests = requests.into_iter().collect::<Vec<_>>();
let _request_guard = lock(&self.request_gate);
let unique_keys = requests.iter().map(|(key, _)| *key).collect::<HashSet<_>>();
let state = lock(&self.shared.state);
ensure!(!state.shutdown, "expert cache is shut down");
ensure!(
unique_keys.len() <= state.max_slots,
"expert request needs {} slots, but the cache has {}",
unique_keys.len(),
state.max_slots
);
drop(state);
requests
.into_iter()
.map(|(key, ranges)| self.request(key, ranges))
.collect()
}
fn request(&self, key: ExpertKey, ranges: [ByteRange; 6]) -> Result<ExpertRequest> {
loop {
let mut state = lock(&self.shared.state);
ensure!(!state.shutdown, "expert cache is shut down");
state.clock = state.clock.saturating_add(1);
let last_used = state.clock;
if let Some(entry) = state.entries.get(&key) {
let generation = match entry {
CacheEntry::Loading { generation, pins }
| CacheEntry::Failed {
generation, pins, ..
} => {
pins.set(
pins.get()
.checked_add(1)
.context("expert pin count overflow")?,
);
*generation
}
CacheEntry::Ready {
generation,
last_used: entry_last_used,
pins,
..
} => {
pins.set(
pins.get()
.checked_add(1)
.context("expert pin count overflow")?,
);
entry_last_used.set(last_used);
*generation
}
};
return Ok(ExpertRequest {
key,
generation,
active: true,
shared: Arc::clone(&self.shared),
});
}
let next_generation = state
.next_generation
.checked_add(1)
.context("expert cache generation overflow")?;
if let Some(buffer) = take_buffer(&mut state)? {
let generation = state.next_generation;
state.next_generation = next_generation;
state.entries.insert(
key,
CacheEntry::Loading {
generation,
pins: Cell::new(1),
},
);
state.queue.push_back(LoadJob {
key,
generation,
ranges,
buffer,
});
#[cfg(feature = "bench")]
self.shared.queued_loads.fetch_add(1, Ordering::Relaxed);
self.shared.work_available.notify_one();
return Ok(ExpertRequest {
key,
generation,
active: true,
shared: Arc::clone(&self.shared),
});
}
#[cfg(feature = "bench")]
self.shared.slot_waits.fetch_add(1, Ordering::Relaxed);
drop(wait(&self.shared.changed, state));
}
}
}
impl Drop for ExpertStore {
fn drop(&mut self) {
{
let mut state = lock(&self.shared.state);
state.shutdown = true;
self.shared.work_available.notify_all();
self.shared.changed.notify_all();
}
for worker in self.workers.drain(..) {
let _result = worker.join();
}
fail_incomplete_loads(&self.shared);
}
}
fn release_pin(state: &mut StoreState, key: ExpertKey, generation: u64) {
let Some(entry) = state.entries.get(&key) else {
return;
};
let remove = match entry {
CacheEntry::Loading {
generation: entry_generation,
pins,
}
| CacheEntry::Ready {
generation: entry_generation,
pins,
..
} => {
if *entry_generation == generation {
pins.set(pins.get().saturating_sub(1));
}
false
}
CacheEntry::Failed {
generation: entry_generation,
pins,
..
} => {
if *entry_generation == generation {
pins.set(pins.get().saturating_sub(1));
pins.get() == 0
} else {
false
}
}
};
if remove {
state.entries.remove(&key);
}
}
fn take_buffer(state: &mut StoreState) -> Result<Option<MmapMut>> {
if let Some(buffer) = state.free_buffers.pop() {
return Ok(Some(buffer));
}
if state.allocated_slots < state.max_slots {
let buffer = MmapOptions::new()
.len(state.slot_bytes)
.map_anon()
.context("allocate expert cache slot")?;
state.allocated_slots += 1;
return Ok(Some(buffer));
}
let victim = state
.entries
.iter()
.filter_map(|(key, entry)| match entry {
CacheEntry::Ready {
last_used, pins, ..
} if pins.get() == 0 => Some((*key, last_used.get())),
CacheEntry::Loading { .. } | CacheEntry::Ready { .. } | CacheEntry::Failed { .. } => {
None
}
})
.min_by_key(|(_, last_used)| *last_used)
.map(|(key, _)| key);
let Some(victim) = victim else {
return Ok(None);
};
let Some(CacheEntry::Ready { data, .. }) = state.entries.remove(&victim) else {
return Ok(None);
};
Ok(Some(data.bytes))
}
fn worker_loop(shared: &Shared, source: &ExpertSource) {
loop {
let job = {
let mut state = lock(&shared.state);
while state.queue.is_empty() && !state.shutdown {
state = wait(&shared.work_available, state);
}
match state.queue.pop_front() {
Some(job) => Some(job),
None if state.shutdown => None,
None => continue,
}
};
let Some(job) = job else {
return;
};
complete_job(shared, source, job);
}
}
fn complete_job(shared: &Shared, source: &ExpertSource, mut job: LoadJob) {
let loaded = catch_unwind(AssertUnwindSafe(|| {
load_ranges(source, &job.ranges, &mut job.buffer)
}));
let result = match loaded {
Ok(result) => result
.with_context(|| format!("load layer {} expert {}", job.key.layer, job.key.expert)),
Err(_) => Err(anyhow!(
"expert I/O worker panicked while loading layer {} expert {}",
job.key.layer,
job.key.expert
)),
};
match result {
Ok(parts) => finish_load(shared, job, parts),
Err(error) => fail_load(shared, job, Arc::from(error.to_string())),
}
}
fn finish_load(shared: &Shared, job: LoadJob, parts: [Range<usize>; 6]) {
let mut state = lock(&shared.state);
let pins = match state.entries.get(&job.key) {
Some(CacheEntry::Loading { generation, pins }) if *generation == job.generation => {
pins.get()
}
Some(CacheEntry::Loading { .. } | CacheEntry::Ready { .. } | CacheEntry::Failed { .. })
| None => {
state.free_buffers.push(job.buffer);
shared.changed.notify_all();
return;
}
};
state.entries.remove(&job.key);
let data = ExpertData {
bytes: job.buffer,
parts,
};
state.clock = state.clock.saturating_add(1);
let last_used = state.clock;
state.entries.insert(
job.key,
CacheEntry::Ready {
generation: job.generation,
data,
last_used: Cell::new(last_used),
pins: Cell::new(pins),
},
);
shared.changed.notify_all();
}
fn fail_load(shared: &Shared, job: LoadJob, error: Arc<str>) {
let mut state = lock(&shared.state);
let pins = match state.entries.get(&job.key) {
Some(CacheEntry::Loading { generation, pins }) if *generation == job.generation => {
Some(pins.get())
}
Some(CacheEntry::Loading { .. } | CacheEntry::Ready { .. } | CacheEntry::Failed { .. })
| None => None,
};
state.free_buffers.push(job.buffer);
if let Some(pins) = pins {
state.entries.remove(&job.key);
if pins != 0 {
state.entries.insert(
job.key,
CacheEntry::Failed {
generation: job.generation,
error,
pins: Cell::new(pins),
},
);
}
}
shared.changed.notify_all();
}
fn fail_incomplete_loads(shared: &Shared) {
let mut state = lock(&shared.state);
while let Some(job) = state.queue.pop_front() {
state.free_buffers.push(job.buffer);
}
let error = Arc::<str>::from("expert cache shut down");
let loading = state
.entries
.iter()
.filter_map(|(key, entry)| match entry {
CacheEntry::Loading { .. } => Some(*key),
CacheEntry::Ready { .. } | CacheEntry::Failed { .. } => None,
})
.collect::<Vec<_>>();
for key in loading {
let Some(CacheEntry::Loading { generation, pins }) = state.entries.remove(&key) else {
continue;
};
state.entries.insert(
key,
CacheEntry::Failed {
generation,
error: Arc::clone(&error),
pins,
},
);
}
shared.changed.notify_all();
}
fn load_ranges(
source: &ExpertSource,
ranges: &[ByteRange; 6],
buffer: &mut MmapMut,
) -> Result<[Range<usize>; 6]> {
if source.direct_enabled.load(Ordering::Relaxed)
&& let Some(direct) = &source.direct
{
match load_ranges_direct(direct, ranges, buffer)? {
DirectLoad::Complete(parts) => return Ok(parts),
DirectLoad::Buffered => {}
DirectLoad::Unavailable => {
source.direct_enabled.store(false, Ordering::Relaxed);
}
}
}
load_ranges_buffered(&source.buffered, ranges, buffer)
}
fn load_ranges_direct(
direct: &DirectFile,
ranges: &[ByteRange; 6],
buffer: &mut MmapMut,
) -> Result<DirectLoad> {
if !(buffer.as_ptr() as usize).is_multiple_of(direct.alignment) {
return Ok(DirectLoad::Unavailable);
}
let mut parts = std::array::from_fn(|_| 0..0);
let mut offset = 0usize;
for (index, range) in ranges.iter().copied().enumerate() {
let file_start = range.start() / direct.alignment * direct.alignment;
let prefix = range.start() - file_start;
let required = prefix
.checked_add(range.len())
.context("expert direct-I/O range overflow")?;
let read_len = align_up(required, direct.alignment)?;
let end = offset
.checked_add(read_len)
.context("expert direct-I/O slot offset overflow")?;
ensure!(end <= buffer.len(), "expert data exceeds its cache slot");
let file_start = u64::try_from(file_start).context("expert file offset exceeds u64")?;
let file_end = file_start
.checked_add(u64::try_from(read_len).context("expert read length exceeds u64")?)
.context("expert direct-I/O file range overflow")?;
if file_end > direct.file_len {
return Ok(DirectLoad::Buffered);
}
if let Err(error) = read_exact_at(&direct.file, &mut buffer[offset..end], file_start) {
if direct_io_unavailable(&error) {
return Ok(DirectLoad::Unavailable);
}
return Err(error).context("read expert range with O_DIRECT");
}
let part_start = offset
.checked_add(prefix)
.context("expert cache part offset overflow")?;
let part_end = part_start
.checked_add(range.len())
.context("expert cache part offset overflow")?;
ensure!(
part_start.is_multiple_of(std::mem::align_of::<f32>()),
"expert cache part {index} is not F32-aligned"
);
parts[index] = part_start..part_end;
offset = end;
}
Ok(DirectLoad::Complete(parts))
}
fn load_ranges_buffered(
file: &File,
ranges: &[ByteRange; 6],
buffer: &mut MmapMut,
) -> Result<[Range<usize>; 6]> {
let mut parts = std::array::from_fn(|_| 0..0);
let mut offset = 0usize;
for (index, range) in ranges.iter().copied().enumerate() {
ensure!(
offset.is_multiple_of(std::mem::align_of::<f32>()),
"expert cache part {index} is not F32-aligned"
);
let end = offset
.checked_add(range.len())
.context("expert cache part offset overflow")?;
ensure!(end <= buffer.len(), "expert data exceeds its cache slot");
read_exact_at(
file,
&mut buffer[offset..end],
u64::try_from(range.start()).context("expert file offset exceeds u64")?,
)?;
discard_file_pages(file, range);
parts[index] = offset..end;
offset = end;
}
Ok(parts)
}
#[cfg(any(target_os = "android", target_os = "linux"))]
fn direct_io_unavailable(error: &io::Error) -> bool {
matches!(
error.raw_os_error(),
Some(libc::EINVAL | libc::EOPNOTSUPP | libc::ENOSYS)
) || error.kind() == io::ErrorKind::InvalidInput
}
#[cfg(not(any(target_os = "android", target_os = "linux")))]
fn direct_io_unavailable(_error: &io::Error) -> bool {
false
}
#[cfg(any(target_os = "android", target_os = "linux"))]
fn open_direct_file(source: &File) -> Option<DirectFile> {
use std::os::fd::AsRawFd as _;
use std::os::unix::fs::OpenOptionsExt as _;
let alignment = system_page_size()?;
let path = format!("/proc/self/fd/{}", source.as_raw_fd());
let file = OpenOptions::new()
.read(true)
.custom_flags(libc::O_DIRECT)
.open(path)
.ok()?;
let file_len = file.metadata().ok()?.len();
Some(DirectFile {
file,
alignment,
file_len,
})
}
#[cfg(not(any(target_os = "android", target_os = "linux")))]
fn open_direct_file(_source: &File) -> Option<DirectFile> {
None
}
#[cfg(unix)]
fn read_exact_at(file: &File, buffer: &mut [u8], offset: u64) -> io::Result<()> {
use std::os::unix::fs::FileExt as _;
file.read_exact_at(buffer, offset)
}
#[cfg(windows)]
fn read_exact_at(file: &File, mut buffer: &mut [u8], mut offset: u64) -> io::Result<()> {
use std::os::windows::fs::FileExt as _;
while !buffer.is_empty() {
let read = file.seek_read(buffer, offset)?;
if read == 0 {
return Err(io::Error::from(io::ErrorKind::UnexpectedEof));
}
offset = offset
.checked_add(read as u64)
.ok_or_else(|| io::Error::from(io::ErrorKind::InvalidInput))?;
buffer = &mut buffer[read..];
}
Ok(())
}
#[cfg(not(any(unix, windows)))]
fn read_exact_at(_file: &File, _buffer: &mut [u8], _offset: u64) -> io::Result<()> {
Err(io::Error::from(io::ErrorKind::Unsupported))
}
#[cfg(any(target_os = "android", target_os = "linux"))]
fn discard_file_pages(file: &File, range: ByteRange) {
use std::os::fd::AsRawFd as _;
let Some(page_size) = system_page_size() else {
return;
};
let Some(start) = range
.start()
.checked_add(page_size - 1)
.map(|start| start / page_size * page_size)
else {
return;
};
let end = range.end() / page_size * page_size;
if start >= end {
return;
}
let (Ok(offset), Ok(length)) = (
libc::off_t::try_from(start),
libc::off_t::try_from(end - start),
) else {
return;
};
let _result =
unsafe { libc::posix_fadvise(file.as_raw_fd(), offset, length, libc::POSIX_FADV_DONTNEED) };
}
#[cfg(any(target_os = "android", target_os = "linux"))]
fn system_page_size() -> Option<usize> {
static PAGE_SIZE: OnceLock<Option<usize>> = OnceLock::new();
*PAGE_SIZE.get_or_init(|| {
let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
usize::try_from(page_size)
.ok()
.filter(|size| size.is_power_of_two())
})
}
#[cfg(not(any(target_os = "android", target_os = "linux")))]
fn discard_file_pages(_file: &File, _range: ByteRange) {}
fn align_up(value: usize, alignment: usize) -> Result<usize> {
let remainder = value % alignment;
if remainder == 0 {
return Ok(value);
}
value
.checked_add(alignment - remainder)
.context("expert cache slot size overflow")
}
fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
match mutex.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
}
}
fn wait<'a, T>(condition: &Condvar, guard: MutexGuard<'a, T>) -> MutexGuard<'a, T> {
match condition.wait(guard) {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
}
}