Struct GenericMdlm 

pub struct GenericMdlm { /* private fields */ }

MDLM masked-diffusion backend.

Loads kuleshov-group/mdlm-owt-format weights (prefix backbone.) and runs the bidirectional DiT forward pass with full hook support, exactly matching the fp32 reference.

Implementations

impl GenericMdlm

pub fn load( config: MdlmConfig, device: &Device, dtype: DType, vb: VarBuilder<'_>, ) -> Result<Self>

Load an MDLM model from a VarBuilder.

The caller constructs the VarBuilder (buffered or mmap) and provides the parsed MdlmConfig. The constant conditioning vector and the RoPE cache are pre-computed here.

Shapes

• returns: a model whose forward maps [batch, seq] token ids to [batch, seq, vocab_size] logits.

Memory

Loads the single mdlm-owt safetensors (~648 MB) through the caller’s VarBuilder: near-zero extra copy with the mmap feature, or ~648 MB CPU when buffered. The conditioning vector and RoPE cache are negligible.

Errors

Returns MIError::Config if the checkpoint sets time_conditioning = true (unsupported), or MIError::Model if weight loading fails or a dimension is inconsistent with the checkpoint.

pub const fn config(&self) -> &MdlmConfig

Access the model configuration.

Trait Implementations

impl MIBackend for GenericMdlm

fn num_layers(&self) -> usize

Number of layers (transformer blocks or RWKV blocks).

fn hidden_size(&self) -> usize

Hidden dimension (d_model).

fn vocab_size(&self) -> usize

Vocabulary size.

fn num_heads(&self) -> usize

Number of attention heads (or RWKV heads).

fn forward(&self, input_ids: &Tensor, hooks: &HookSpec) -> Result<HookCache>

Unified forward pass with optional hook capture and interventions.

fn project_to_vocab(&self, hidden: &Tensor) -> Result<Tensor>

Project a hidden-state tensor to vocabulary logits.

fn embedding_vector(&self, token_id: u32) -> Result<Tensor>

Return the raw embedding vector for a single token.

fn chat_template( &self, _prompt: &str, _system_prompt: Option<&str>, ) -> Option<String>

Format a prompt with the model’s chat template, if any.