osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
60 lines
1.9 KiB
Rust
60 lines
1.9 KiB
Rust
//! Speculative Decoding for Inference Acceleration
|
|
//!
|
|
//! This module implements speculative decoding to improve inference throughput by using
|
|
//! a smaller, faster "draft" model to generate candidate tokens that are then verified
|
|
//! by the target model. This can provide significant speedups for autoregressive generation.
|
|
//!
|
|
//! Key features:
|
|
//! - Multiple draft model types (small models, pruned models, cached predictions)
|
|
//! - Configurable acceptance thresholds and draft lengths
|
|
//! - Performance tracking and adaptive optimization
|
|
//! - Safe error handling throughout the pipeline
|
|
//!
|
|
//! # Layering
|
|
//!
|
|
//! This module is the *orchestration* layer: strategy traits
|
|
//! ([`MedusaDraftModel`], [`SelfSpeculativeModel`]), configs, candidate
|
|
//! trees, and the [`AdvancedSpeculativeDecoder`] that combines them. The
|
|
//! concrete numeric implementations live at the crate root:
|
|
//! [`crate::medusa`] (real multi-head FFN prediction over hidden states,
|
|
//! arXiv:2401.10774) and [`crate::lookahead`] (n-gram lookahead decoding).
|
|
//! The configs here (`MedusaConfig`, `LookaheadConfig`) parameterize the
|
|
//! orchestration strategies; they are distinct from the implementation
|
|
//! configs (`crate::medusa::MedusaConfig` is re-exported from the crate
|
|
//! root as `MedusaHeadsConfig` to disambiguate).
|
|
|
|
mod advanced;
|
|
mod config;
|
|
mod decoder;
|
|
mod eagle;
|
|
pub mod eagle3;
|
|
mod error;
|
|
mod lookahead;
|
|
mod medusa;
|
|
mod metrics;
|
|
mod self_spec;
|
|
pub mod streaming;
|
|
mod traits;
|
|
mod tree;
|
|
mod types;
|
|
|
|
pub use advanced::*;
|
|
pub use config::*;
|
|
pub use decoder::*;
|
|
pub use eagle::*;
|
|
pub use eagle3::{DynamicDraftTree, Eagle3Config, Eagle3Decoder};
|
|
pub use error::*;
|
|
pub use lookahead::*;
|
|
pub use medusa::*;
|
|
pub use metrics::*;
|
|
pub use self_spec::*;
|
|
pub use streaming::{
|
|
SpeculativeStreamConfig, SpeculativeStreamer, StreamStats, StreamedToken, collect_stream,
|
|
};
|
|
pub use traits::*;
|
|
pub use tree::*;
|
|
pub use types::*;
|
|
|
|
#[cfg(test)]
|
|
mod tests;
|