ModelWeightLoader

metrale_model_arch::weight_loader

Trait ModelWeightLoader 

Source 
pub trait ModelWeightLoader {

Show 17 methods    // Required methods
    fn supports_tp(&self) -> bool;
    fn load_layers(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
        layer_kv_dtypes: &[KvCacheDtype],
    ) -> Result<Vec<Box<dyn TransformerLayer>>>;
    fn load_embedding(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<DenseWeight>;
    fn load_final_norm(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<DenseWeight>;
    fn load_lm_head(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<DenseWeight>;
    fn load_mtp_weights(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<Option<MtpWeights>>;

    // Provided methods
    fn prune_after_load(
        &self,
        _store: &mut WeightStore,
        _config: &ModelConfig,
        _gpu: &dyn GpuBackend,
    ) -> Result<()> { ... }
    fn precision_schedule(&self, _config: &ModelConfig) -> PrecisionSchedule { ... }
    fn load_ngram_embedding(
        &self,
        _store: &WeightStore,
        _config: &ModelConfig,
        _gpu: &dyn GpuBackend,
        _max_tokens: usize,
    ) -> Result<Option<NgramEmbedding>> { ... }
    fn load_mtp_weights_multi(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<Vec<MtpWeights>> { ... }
    fn kv_layer_dims(&self, _config: &ModelConfig) -> Vec<(usize, usize)> { ... }
    fn load_dflash_weights(
        &self,
        _drafter_store: &WeightStore,
        _config: &ModelConfig,
        _gpu: &dyn GpuBackend,
        _tp_size: usize,
    ) -> Result<Option<DflashWeights>> { ... }
    fn load_lora_adapters(
        &self,
        adapters: &[LoraAdapterInput<'_>],
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
        max_loras: usize,
        max_lora_rank: usize,
    ) -> Result<Option<LoraWeights>> { ... }
    fn binds_vision_encoder(&self) -> bool { ... }
    fn reads_expert_table_plan(&self) -> bool { ... }
    fn defer_predicate(&self, _config: &ModelConfig) -> Option<DeferHook> { ... }
    fn load_vision_encoder(
        &self,
        _store: &WeightStore,
        _config: &ModelConfig,
        _gpu: &dyn GpuBackend,
    ) -> Result<Option<VisionTower>> { ... }

}
Expand description

2026-09-25: Loads weights from a WeightStore into typed layer objects.

Required Methods§

Source

fn supports_tp(&self) -> bool

2026-09-25: Whether this loader’s weight slicing honours config.tp_world_size / config.tp_rank. There is no default, so every loader declares it.

With --tp-size > 1, startup (serve_phases/topology.rs) refuses a loader that returns false. weight_loader/minimax.rs is a loader that returns true.

Source

fn load_layers( &self, store: &WeightStore, config: &ModelConfig, gpu: &dyn GpuBackend, layer_kv_dtypes: &[KvCacheDtype], ) -> Result<Vec<Box<dyn TransformerLayer>>>

2026-09-25: Load all transformer layers from the weight store.

layer_kv_dtypes holds one KV cache dtype per attention layer, indexed by the layer’s 0-based ordinal among the attention layers.

Source

fn load_embedding( &self, store: &WeightStore, config: &ModelConfig, gpu: &dyn GpuBackend, ) -> Result<DenseWeight>

Source

fn load_final_norm( &self, store: &WeightStore, config: &ModelConfig, gpu: &dyn GpuBackend, ) -> Result<DenseWeight>

2026-09-25: Load the final RMSNorm weight used before the LM head.

Source

fn load_lm_head( &self, store: &WeightStore, config: &ModelConfig, gpu: &dyn GpuBackend, ) -> Result<DenseWeight>

Source

fn load_mtp_weights( &self, store: &WeightStore, config: &ModelConfig, gpu: &dyn GpuBackend, ) -> Result<Option<MtpWeights>>

2026-09-25: Load single-module MTP head weights, or None.

Provided Methods§

Source

fn prune_after_load( &self, _store: &mut WeightStore, _config: &ModelConfig, _gpu: &dyn GpuBackend, ) -> Result<()>

2026-09-25: Drop store tensors this loader has finished with. factory::build calls it after every load_* reader and before the KV budget is computed, so the freed memory counts.

The default keeps everything, which is right for a loader that binds the store’s device pointers zero-copy. A loader that uploads its own copies (a TP shard, a host round-trip, a dtype conversion) overrides it to free the originals.

Source

fn precision_schedule(&self, _config: &ModelConfig) -> PrecisionSchedule

2026-09-25: Per-(layer, role) weight precision schedule. The default is PrecisionSchedule::default(), in which every role is Dtype::Inherit.

Source

fn load_ngram_embedding( &self, _store: &WeightStore, _config: &ModelConfig, _gpu: &dyn GpuBackend, _max_tokens: usize, ) -> Result<Option<NgramEmbedding>>

2026-09-25: Build the n-gram embedding, for an architecture that fuses hashed n-gram lookups into the input embedding (the LongCat loader). The default None leaves the plain embed_tokens gather in place.

Source

fn load_mtp_weights_multi( &self, store: &WeightStore, config: &ModelConfig, gpu: &dyn GpuBackend, ) -> Result<Vec<MtpWeights>>

2026-09-25: Load every MTP module; factory::build calls this one. The default wraps load_mtp_weights (zero or one module). The MiniMax and Step3p7 loaders override it.

Source

fn kv_layer_dims(&self, _config: &ModelConfig) -> Vec<(usize, usize)>

2026-09-25: Per-attention-layer (num_kv_heads, head_dim), indexed like layer_kv_dtypes, for models whose attention geometry differs per layer (the Gemma-4 loader). The default is empty: every layer uses the global dims.

Source

fn load_dflash_weights( &self, _drafter_store: &WeightStore, _config: &ModelConfig, _gpu: &dyn GpuBackend, _tp_size: usize, ) -> Result<Option<DflashWeights>>

2026-09-25: Load DFlash drafter weights from the drafter’s own WeightStore. The default returns None, and no loader overrides it: factory::build loads the drafter with the free function dflash_loader::load_dflash_weights.

Source

fn load_lora_adapters( &self, adapters: &[LoraAdapterInput<'_>], config: &ModelConfig, gpu: &dyn GpuBackend, max_loras: usize, max_lora_rank: usize, ) -> Result<Option<LoraWeights>>

2026-09-25: Load startup-static PEFT LoRA adapters into the LoRA pool, through metrale_model_layers::lora::load_lora_adapters_multi. factory::build calls it before the buffer arena and the KV budget, so the pool counts as used memory.

Source

fn binds_vision_encoder(&self) -> bool

2026-09-25: Will this loader bind a vision encoder for a multimodal checkpoint?

Default true, so a loader that does not override it never loses a weight. On false the checkpoint loader skips the tower’s tensors (serve_phases/weights.rs). factory::build frees an unbound tower afterwards either way, keyed off the bind result, so this only lowers peak memory.

Source

fn reads_expert_table_plan(&self) -> bool

2026-10-02: Does this loader read the serve’s MoE expert-table decision (metrale_model_layers::layers::MoeExpertTables)? The serve plans and publishes it before load only for a loader that does.

Source

fn defer_predicate(&self, _config: &ModelConfig) -> Option<DeferHook>

2026-09-25: Which checkpoint tensors will this loader read from the host at bind time, so the checkpoint loader records their location instead of uploading them?

Default None: upload everything. serve_phases/weights.rs asks once, before the load. A deferred tensor is not skipped: it is recorded in WeightStore::deferred, and the loader that asked reads it. See metrale_model_weights::weights::DeferHook for the contract.

Source

fn load_vision_encoder( &self, _store: &WeightStore, _config: &ModelConfig, _gpu: &dyn GpuBackend, ) -> Result<Option<VisionTower>>

2026-09-25: Load vision encoder weights; the default returns None.

Implementors§