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.