pub enum ServingState<Entry, Cache> {
Prefill {
positions: Vec<Entry>,
cache: Cache,
},
Decode {
last: Entry,
cache: Cache,
},
Verify {
draft: Vec<Entry>,
snapshot: Cache,
cache: Cache,
},
Accept {
n: usize,
next: Entry,
cache: Cache,
},
Rollback {
snapshot: Cache,
to: Entry,
},
Done {
cache: Cache,
},
}Expand description
Which evaluation shape is legal right now. An enum whose transitions
consume self states that directly; no existing proxima primitive (pipe,
cache or config) expresses it more cheaply.
Variants§
Prefill
positions: the prompt token ids to place in one M-row program
evaluation. cache starts empty – prefill is the only state that
does not inherit recurrent state from a predecessor.
Decode
last: the most recently accepted token id, fed back as the next
single-row evaluation’s input.
Verify
draft: candidate token ids proposed by a speculator, scored in one
K-row evaluation with logits returned at every row. snapshot is
the pre-draft cache, restored verbatim on rejection.
Accept
n: how many leading draft tokens the verifier accepted; next is
the token id decoding resumes from (the confirmed final draft token
when n == draft.len(), otherwise the target’s own row-n
prediction that replaced the rejected draft token).
Rollback
to: the token id decoding resumes from after snapshot is
restored in place of a rejected draft’s advanced cache.
Done
Terminal: no further evaluation is legal.
Fields
cache: CacheImplementations§
Source§impl<Entry, Cache> ServingState<Entry, Cache>
impl<Entry, Cache> ServingState<Entry, Cache>
Sourcepub fn start(positions: Vec<Entry>, cache: Cache) -> Self
pub fn start(positions: Vec<Entry>, cache: Cache) -> Self
Enter the state machine at Prefill with an empty cache.
Sourcepub fn advance_prefill(
self,
next_token: Entry,
cache: Cache,
) -> Result<Self, ServingFsmError>
pub fn advance_prefill( self, next_token: Entry, cache: Cache, ) -> Result<Self, ServingFsmError>
Prefill { positions } -> Decode { last }: the one M-row program
evaluation prefill performs produced next_token and left cache
holding the resulting recurrent/KV state.
Sourcepub fn advance_decode(
self,
next_token: Entry,
cache: Cache,
) -> Result<Self, ServingFsmError>
pub fn advance_decode( self, next_token: Entry, cache: Cache, ) -> Result<Self, ServingFsmError>
Decode { last } -> Decode { last }: one single-row evaluation.
Sourcepub fn enter_verify(self, draft: Vec<Entry>) -> Result<Self, ServingFsmError>where
Cache: Clone,
pub fn enter_verify(self, draft: Vec<Entry>) -> Result<Self, ServingFsmError>where
Cache: Clone,
Decode { last } -> Verify { draft, snapshot }: the current cache
becomes snapshot, restorable verbatim if the draft is rejected.
Sourcepub fn resume(self) -> Result<Self, ServingFsmError>
pub fn resume(self) -> Result<Self, ServingFsmError>
Accept { n, next, cache } -> Decode { last }: no evaluation, just
unwrapping the placement accept already chose.
Sourcepub fn rollback(self) -> Result<Self, ServingFsmError>
pub fn rollback(self) -> Result<Self, ServingFsmError>
Rollback { snapshot, to } -> Decode { last }: placement copy only,
no program evaluation – snapshot is whichever row’s Cache
accept selected as the correct resume point.
Source§impl<Entry: PartialEq + Clone, Cache> ServingState<Entry, Cache>
impl<Entry: PartialEq + Clone, Cache> ServingState<Entry, Cache>
Sourcepub fn accept(
self,
row_tokens: &[Entry],
row_caches: Vec<Cache>,
) -> Result<Self, ServingFsmError>
pub fn accept( self, row_tokens: &[Entry], row_caches: Vec<Cache>, ) -> Result<Self, ServingFsmError>
Verify { draft, snapshot } -> Accept { n, next } | Rollback { snapshot, to }:
scores the K-row verify evaluation’s per-row greedy argmax
(row_tokens[i] predicts the token that follows having consumed
input row i, where row 0 follows last and row i>0 follows
draft[i - 1]) against draft itself. row_caches[i] is the exact
Cache the same evaluation produced at row i – accepting n
draft tokens means resuming from row_caches[n] verbatim (a
placement copy the caller already computed), never replaying
anything. n == draft.len(): every drafted token matched, Accept
resumes from the confirmed final draft token. n < draft.len():
draft[n] was wrong, so this returns Rollback carrying the
correct placement and the target’s own token in its place.
Sourcepub fn accept_rows(
self,
accepted: usize,
row_tokens: &[Entry],
row_caches: Vec<Cache>,
) -> Result<Self, ServingFsmError>
pub fn accept_rows( self, accepted: usize, row_tokens: &[Entry], row_caches: Vec<Cache>, ) -> Result<Self, ServingFsmError>
Self::accept with the accepted count supplied by the caller (see
crate::accept_rule::AcceptRule::accepted_rows). A verifier that stops
sampling at the first row that differs passes only the tokens it chose:
row_tokens needs to reach row accepted (and nothing past it) unless
every draft row was accepted, while row_caches always holds one entry
per draft row.