Skip to main content

load_kimi_checkpoint_with_expert_cache

Function load_kimi_checkpoint_with_expert_cache 

Source
pub fn load_kimi_checkpoint_with_expert_cache(
    shard: &ShardedSafetensors,
    model_cfg: &ModelConfig,
    hp: &KimiRealHparams,
    expert_cache_bytes: Option<u64>,
) -> Result<KimiDecoderWeights, KimiLoadError>
Expand description

Loads a complete KimiDecoderWeights – every one of model_cfg’s real layers (dispatched per-layer via model_cfg.layer_attention_kind/ layer_is_dense, driven by hp’s per-layer dimensions), plus the real top-level tensors (real names confirmed against a real shard header: language_model.model.embed_tokens.weight, language_model.lm_head.weight, language_model.model.norm.weight, language_model.model.output_attn_res_{norm,proj}.weight). This is the assembly step load_kimi_layer itself doesn’t do – calling it once per real layer and building the surrounding KimiDecoderWeights – analogous to ferrox-models::loader::Decoder::from_gguf, but for Kimi K3’s real safetensors format. Not blocked on anything (the zero-copy MXFP4 fix removes the memory obstacle a full loader would otherwise hit for every non-dense layer’s routed experts); simply not runnable against the real 2.8T-parameter checkpoint in this environment (96 shards, 1.56TB) – tested here against small synthetic on-disk fixtures instead, real safetensors bytes and real tensor names throughout. Like load_kimi_checkpoint, but with expert_cache_bytes: Some(budget) every MoE layer’s routed experts are converted to store-backed lazy materialization after loading: one bounded, lease-protected ExpertStore shared by the whole model reads each expert’s six tensors positionally from the owning shard files on miss, instead of holding 896 expert objects per layer resident. Attention, dense layers, shared experts, router/projections, embeddings, and the output head are untouched. Bit-identical to the eager path (same bytes, same kernels) – pinned by the equivalence test against the synthetic multi-layer checkpoint.