Trait ModelWeightLoader 

pub trait ModelWeightLoader {
    // 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 defer_predicate(&self, _config: &ModelConfig) -> Option<DeferHook> { ... }
    fn load_vision_encoder(
        &self,
        _store: &WeightStore,
        _config: &ModelConfig,
        _gpu: &dyn GpuBackend,
    ) -> Result<Option<VisionTower>> { ... }
}

Loads weights from a WeightStore into typed layer objects.

Required Methods

fn supports_tp(&self) -> bool

Whether this loader’s weight slicing is TP-aware. No default — every loader MUST declare this explicitly so adding a new model architecture cannot accidentally inherit a false and silently regress users who pass --tp-size > 1.

Loaders that honour config.tp_world_size / config.tp_rank when loading attention Q/K/V/O, MoE gate/up/down, head-parallel SSM components, and lm_head return true. Loaders that always load full replicated weights return false.

The startup path in spark-server/src/main.rs consults this method to fail-fast at load time when --tp-size > 1 is requested against a TP-unaware loader. Extending TP to a new architecture requires:

1. Wire slice_for_rank (in crate::tp_shard) per Q/K/V/O, gate/up/down, and any head-parallel SSM tensors.
2. Divide num_attention_heads / num_key_value_heads per the same axis when constructing layer state.
3. Return true from this method.

See weight_loader/minimax.rs for the reference implementation.

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

Load all transformer layers from the weight store.

layer_kv_dtypes is indexed by attention layer index (0-based sequential counter over full-attention layers only). Each attention layer receives its own KV cache dtype, enabling mixed-precision KV caching where boundary layers use higher precision.

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>

Load the final RMSNorm weight used before the LM head.

gpu is passed so model-specific loaders can do on-device weight transforms at load time (e.g. Gemma-4 shifts the learned absolute- scale weight by -1 into the offset-from-1 convention expected by Atlas’s rms_norm kernel). Loaders that don’t need it should ignore the argument.

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>>

Load MTP head weights (returns None if no MTP weights in store).

Provided Methods

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

Drop store tensors this loader has finished with, after every load_* reader has run and before the buffer arena / KV cache are sized.

Default: keep everything. That is correct for the loaders that bind zero-copy from the store’s device pointers — the store IS the model’s weights, and TransformerModel releases it at teardown.

Override only when the loader uploads its own copies (a TP shard, a host round-trip, a dtype conversion), because then the store’s originals are dead the moment the binder returns. On unified-memory GB10 that duplicate comes straight out of the KV budget.

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

Per-(layer, role) weight precision schedule (C.3, 2026-04-25). Default impl returns the empty schedule (every lookup yields Dtype::Inherit), preserving the existing per-checkpoint dtype logic byte-for-byte. Loader-specific implementations can override to honour MODEL.toml’s [precision] block.

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

Build the n-gram embedding, when this architecture fuses hashed n-gram lookups into the input embedding (LongCat / Qwen3.8-Flash-Next).

Separate from load_embedding because the result is NOT a weight: it is a small engine that needs the sequence’s CONTEXT token ids at forward time, not just the id being embedded. Returning None — the default — leaves the plain embed_tokens gather in place.

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

Load MTP weights for multi-module MTP (DeepSeek-V3 / MiniMax-M2 style: N independent transformer modules, each with its own attention + MoE + KV cache). Returns an empty Vec when the checkpoint has no MTP modules, a 1-element Vec for single-module MTP (Qwen3.5 family), or N elements for multi-module.

Default impl adapts load_mtp_weights so existing single-module loaders don’t need to change. MiniMax overrides this directly.

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

Per-layer (num_kv_heads, head_dim) overrides for heterogeneous attention models (e.g. Gemma-4 with sliding 16×256 and full 4×512). Default empty — homogeneous models skip per-layer dims and the KV cache allocator uses the global (num_kv_heads, head_dim). Populated by loaders whose models have different attention geometries per layer. Indexed by attention layer index (same as layer_kv_dtypes).

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

Load DFlash drafter weights from a separate WeightStore pointing at the drafter checkpoint (z-lab/Qwen3.6-{27B,35B-A3B}-DFlash). Default impl returns None so loaders that don’t yet support DFlash silently fall through to the existing MTP path. Override in loaders whose target models pair with a DFlash drafter (Qwen3.5/3.6 family). The same drafter format works across both 27B-dense and 35B-A3B-MoE targets — only the target_hidden_size validated against the drafter’s fc input dimension differs.

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

Load one or more startup-static PEFT LoRA adapters from their own WeightStores (the adapter_model.safetensors tensors, already on-device BF16) into the fixed-address rank-padded pool (one slot each).

Unlike load_dflash_weights’ vestigial Ok(None) default, the default here is a WORKING model-agnostic implementation (the remap needs only ModelConfig::layer_type + projection dims); families needing a bespoke key remap override it. Called from factory::build_model BEFORE the buffer arena + KV sizing so the pool bytes are budgeted against the KV cache. A single-element slice is byte-identical to the pre-multi-adapter single-adapter path.

fn binds_vision_encoder(&self) -> bool

Will this loader ever bind a vision encoder for a multimodal checkpoint?

Default true — “load everything” is the safe answer, so a loader that forgets to override this can never lose weights it needs. A loader whose port is deliberately text-only overrides it to false, and the weight loader then skips the tower’s tensors instead of reading a gigabyte of unified memory that nothing will bind. build_model still frees an unbound tower afterwards (keyed off the bind result, not off this), so this is a peak-memory optimisation, not the correctness gate.

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

Which checkpoint tensors will this loader read from the HOST at bind time, so the weight loader must record their location instead of uploading them?

Default None — “upload everything” is the safe answer, exactly as with Self::binds_vision_encoder, so a loader that never overrides this can never lose a tensor it needs. Plumbed the same way: serve_phases::weights asks once, before the load, and hands the answer to whichever checkpoint loader runs.

This is NOT “skip”: a deferred tensor’s (path, offset, shape, dtype) is kept in WeightStore::deferred, and the loader that asked for it is the only thing that will ever read it. Override only when the bytes that belong on the device are DERIVED from the on-disk ones — a full-width export of a weight whose only forward is quantised, say — because then the upload is pure transient and on a unified-memory box the transient is the whole problem. See spark_runtime::weights::DeferHook for the full contract, including the half that says the OOM pre-flight must agree.

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

Load vision encoder weights (returns None for text-only models).

Implementors

impl ModelWeightLoader for MistralWeightLoader

impl ModelWeightLoader for Glm5NextWeightLoader

impl ModelWeightLoader for KimiK3WeightLoader

impl ModelWeightLoader for DeepSeekV4WeightLoader

impl ModelWeightLoader for Gemma4WeightLoader

impl ModelWeightLoader for LagunaWeightLoader

impl ModelWeightLoader for LongcatWeightLoader

impl ModelWeightLoader for MinimaxM2WeightLoader

impl ModelWeightLoader for NemotronHWeightLoader

impl ModelWeightLoader for NllbWeightLoader

impl ModelWeightLoader for Qwen3VLWeightLoader

impl ModelWeightLoader for Qwen3WeightLoader

impl ModelWeightLoader for Qwen4ExpWeightLoader

impl ModelWeightLoader for Qwen35DenseWeightLoader

impl ModelWeightLoader for Qwen35WeightLoader

impl ModelWeightLoader for Step3p7WeightLoader