pub struct ExpertPager { /* private fields */ }Expand description
A byte-budgeted, disk-backed cache of MoE experts.
Model-agnostic: it knows byte ranges, not architectures. A model registers where its banks live and asks for the experts a token fired.
Cheap to share — gather takes &self, so several decode threads can page
against one cache.
Implementations§
Source§impl ExpertPager
impl ExpertPager
Sourcepub fn new(budget_bytes: u64) -> Self
pub fn new(budget_bytes: u64) -> Self
A pager holding at most budget_bytes of expert data.
A budget of 0 disables caching: every read goes to disk. That is a legitimate configuration (it is what the planner’s cost model assumes, being the pessimistic case) rather than an error.
Sourcepub fn register(
&mut self,
layer: usize,
bank: impl Into<String>,
loc: BankLocation,
) -> Result<BankIx>
pub fn register( &mut self, layer: usize, bank: impl Into<String>, loc: BankLocation, ) -> Result<BankIx>
Tell the pager where one bank lives.
Rejects a location whose bytes do not divide evenly into its experts:
that means bytes_per_expert was computed from a wrong expert count or
a wrong tensor, and every read after the first would be misaligned —
which does not fail, it returns plausible-looking garbage.
Sourcepub fn bank_ix(&self, layer: usize, bank: &str) -> Option<BankIx>
pub fn bank_ix(&self, layer: usize, bank: &str) -> Option<BankIx>
Handle for a registered bank, or None if it was never registered.
Hoist this out of a per-token loop: the handle is Copy, and every
lookup taking one is allocation-free.
Sourcepub fn location(&self, ix: BankIx) -> &BankLocation
pub fn location(&self, ix: BankIx) -> &BankLocation
Where a bank lives.
Sourcepub fn bank_count(&self) -> usize
pub fn bank_count(&self) -> usize
Banks registered so far.
Sourcepub fn total_bank_bytes(&self) -> u64
pub fn total_bank_bytes(&self) -> u64
Bytes the whole registered checkpoint would occupy if held resident — what paging is avoiding.
pub fn stats(&self) -> PagerStats
Sourcepub fn resident_bytes(&self) -> u64
pub fn resident_bytes(&self) -> u64
Bytes currently cached.
Sourcepub fn expert(
&self,
layer: usize,
bank: &str,
id: usize,
) -> Result<Arc<Vec<u8>>>
pub fn expert( &self, layer: usize, bank: &str, id: usize, ) -> Result<Arc<Vec<u8>>>
One expert’s packed bytes, from cache or disk.
Resolves the bank by name each call; Self::expert_at skips that.
Sourcepub fn expert_at(&self, ix: BankIx, id: usize) -> Result<Arc<Vec<u8>>>
pub fn expert_at(&self, ix: BankIx, id: usize) -> Result<Arc<Vec<u8>>>
One expert’s packed bytes, by handle.
Sourcepub fn gather(&self, layer: usize, bank: &str, ids: &[usize]) -> Result<Vec<u8>>
pub fn gather(&self, layer: usize, bank: &str, ids: &[usize]) -> Result<Vec<u8>>
The fired experts’ bytes, concatenated in the order given.
Laid out exactly like the full bank, so the result is a GroupedMatMul
operand with num_experts = ids.len(). Order is the caller’s: it has to
match the routing weights the caller will apply, so this must not sort
or deduplicate. Repeating an id is legal and copies it twice.
Sourcepub fn gather_at(&self, ix: BankIx, ids: &[usize]) -> Result<Vec<u8>>
pub fn gather_at(&self, ix: BankIx, ids: &[usize]) -> Result<Vec<u8>>
Self::gather by handle — the name is resolved once, not per expert.
Sourcepub fn prefetch(&self, layer: usize, bank: &str, ids: &[usize])
pub fn prefetch(&self, layer: usize, bank: &str, ids: &[usize])
Pull experts into the cache ahead of needing them.
Useful when the router for step t+1 is known while step t is still
computing. Errors are swallowed: a prefetch that fails costs a later
cache miss, and turning a speculative read into a hard failure would be
worse than the miss.
Auto Trait Implementations§
impl !Freeze for ExpertPager
impl RefUnwindSafe for ExpertPager
impl Send for ExpertPager
impl Sync for ExpertPager
impl Unpin for ExpertPager
impl UnsafeUnpin for ExpertPager
impl UnwindSafe for ExpertPager
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more