Struct LLaDA2Runner 

pub struct LLaDA2Runner {
    pub cfg: LLaDA2MoeConfig,
    pub weights: LLaDA2Weights,
    /* private fields */
}

Fields

cfg: LLaDA2MoeConfig

weights: LLaDA2Weights

Implementations

impl LLaDA2Runner

pub fn builder() -> LLaDA2RunnerBuilder

pub fn config(&self) -> &LLaDA2MoeConfig

pub fn device(&self) -> Device

pub fn max_seq(&self) -> usize

pub fn predictive_offload_enabled(&self) -> bool

pub fn jump_steps(&self) -> usize

pub fn predictive_offload_info(&self) -> Option<PredictiveOffloadInfo>

pub fn moe_offload(&self) -> Option<&MoeOffloadState>

pub fn moe_store(&self) -> Option<&MoeExpertStore>

pub fn sync_moe_residency(&self, compiled: &mut CompiledGraph)

pub fn forward_logits( &mut self, input_ids: &[f32], position_ids: &[f32], attn_mask: &[f32], ) -> Result<Vec<f32>, Error>

pub fn block_denoise_loop( &mut self, cfg: BlockDenoiseConfig, ) -> BlockDenoiseLoop<LLaDA2RunnerForward<'_>>

pub fn get_offload_stats(&mut self) -> TideOffloadStats

pub fn reset_offload_step_stats(&mut self)

pub fn generate( &mut self, gen_cfg: &GenerateConfig, prompt_ids: &[u32], ) -> Result<(Vec<u32>, Vec<BlockDenoiseStepStats>), Error>

pub fn offload_stats( &self, residency: Option<&MoeResidencyStats>, ) -> Option<TideOffloadStats>

Trait Implementations

impl BlockDenoiseSampler for LLaDA2Runner

fn sample_block( &mut self, x: &[u32], window_end: usize, block_length: usize, refresh_experts: bool, gen_cfg: &GenerateConfig, model_cfg: &LLaDA2MoeConfig, step_ctx: DenoiseStepCtx, ) -> Result<(Vec<u32>, Vec<f32>), Error>