Struct SpeculativePipeline

pub struct SpeculativePipeline { /* private fields */ }

Speculative decoding pipeline: https://arxiv.org/pdf/2211.17192

Algorithm

Given draft model q and target model p with probability distributions
q_i(x) and p_i(x) for each token

• Keep the sample for token i if q_i(x) <= p_i(x)
  • This means the target model agrees
• Else (q_i(x) > p_i(x)) accept that token with prob p_i(x)/q_i(x)
  • If rejected, sample token from from p’_i(x) = norm(max(0, p(x) − q(x))) and do not take any more’

Implementations

impl SpeculativePipeline

pub fn new( target: Arc<Mutex<dyn Pipeline>>, draft: Arc<Mutex<dyn Pipeline>>, config: SpeculativeConfig, ) -> Result<SpeculativePipeline, Error>

Trait Implementations

impl Pipeline for SpeculativePipeline

fn forward_inputs( &mut self, _inputs: Box<dyn Any>, _return_raw_logits: bool, ) -> Result<ForwardInputsResult, Error>

fn sample_causal_gen<'life0, 'life1, 'life2, 'life3, 'async_trait>( &'life0 self, _seqs: &'life1 mut [&'life2 mut Sequence], _logits: Vec<Tensor>, _prefix_cacher: &'life3 mut PrefixCacheManagerV2, _disable_eos_stop: bool, _rng: Arc<Mutex<Isaac64Rng>>, ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + 'async_trait>>

where 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait, SpeculativePipeline: 'async_trait,

fn step<'life0, 'life1, 'life2, 'life3, 'life4, 'async_trait>( &'life0 mut self, input_seqs: &'life1 mut [&'life2 mut Sequence], is_prompt: bool, _return_raw_logits: bool, prefix_cacher: &'life3 mut PrefixCacheManagerV2, disable_eos_stop: bool, rng: Arc<Mutex<Isaac64Rng>>, backend_metadata: CacheBackendMetadata<'life4>, ) -> Pin<Box<dyn Future<Output = Result<Duration, Error>> + Send + 'async_trait>>

where 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait, 'life4: 'async_trait, SpeculativePipeline: 'async_trait,

Returns the total of model execution time.

fn category(&self) -> ModelCategory