Module kimi_loader 

Loads ferrox-models::kimi_decoder weights from Kimi K3’s real safetensors checkpoint (via ferrox-safetensors::ShardedSafetensors), using the exact real tensor names/shapes/dtypes fetched live from huggingface.co/moonshotai/Kimi-K3 (a real shard header, not guessed) – confirmed to match this crate’s kda/mla/latent_moe/ kimi_decoder struct field names and shapes exactly.

One real, non-obvious fact confirmed by reading actual tensor shapes rather than assuming they match modeling_kimi_linear.py’s KimiDeltaAttention.__init__ literally: self_attn.A_log’s real on-disk shape is [128], not [num_heads] = [96] (confirmed independently by self_attn.b_proj.weight’s real shape [96, 7168], which is unambiguously [num_heads, hidden_dim]). The real fused_recurrent_kda kernel only ever indexes A_log[i_hv] for i_hv in 0..num_heads, so this is real, harmless padding (likely to a GPU-friendly round size) rather than a spec mismatch – this loader reads the real 128-element tensor but only uses its first num_heads elements, matching what the real kernel actually consumes.

Dequantizes small F32/BF16 tensors (per-head attention parameters, layernorms, dense-FFN and shared-expert projections) to owned f32 eagerly at load time – matching this project’s established BF16-handling convention, e.g. ferrox-models::loader’s GGUF path – since these are cheap regardless. Routed-expert MXFP4 weights are the one format this loader does not eagerly dequantize: load_mxfp4_weight_matrix builds a zero-copy WeightMatrix::Mxfp4 (mmap-backed packed/scale buffers, real Kimi K3 stores these as two separate tensors per projection – see ferrox_quant::dot_mxfp4_row_f32’s doc comment), matching the zero-copy-mmap-plus-fused-dot discipline every other quantized format in this codebase already uses. This is a real fix, not just a design preference: real-hardware testing (rented 62GB instance, see docs/MODELS.md) found that eagerly dequantizing all 896 of a real MoE layer’s routed experts to owned f32 needs roughly 117GB of RAM and reproducibly OOM-killed a 62GB rented instance – the zero-copy path here keeps a loaded layer’s resident memory close to its on-disk size (the real MXFP4 packing ratio: 2 values/byte plus one scale byte per 32 values) instead of expanding every value to 4 bytes whether it’s ever used by real top-k routing or not.

Structs

BlockResidualWeights

The four block-residual weight vectors real Kimi K3 attaches to every layer ({prefix}.self_attention_res_{norm,proj}.weight, {prefix}.mlp_res_{norm,proj}.weight). Real on-disk *_proj.weight shape is [1, hidden_dim] (a Linear(hidden_dim, 1)’s weight); the raw bytes are already exactly [hidden_dim] flattened.

KimiExpertSource

ExpertSource over a Kimi safetensors checkpoint: each expert’s six tensors (three matrices’ packed+scale buffers) are read positionally from the owning shard files and concatenated in KimiStoredExpertLayout’s fixed order.

KimiRealHparams

Kimi K3’s real per-layer hyperparameters needed to load any layer (not tied to ferrox_moe::MoeLayerConfig/ferrox_models::ModelConfig, neither of which model the “latent MoE” down-projected dimension or the dense leading layer’s own intermediate size – kept as a small, dedicated struct here rather than widening those shared types for one model’s real values).

KimiStoredExpertLayout

Loads one routed expert’s real MXFP4 weights ({prefix}.experts.{expert_idx}.{w1,w2,w3}.{weight_packed,weight_scale}). moe_hidden_dim is the real latent dimension (routed_expert_hidden_size=3584 for Kimi K3, not the outer hidden_dim=7168 – see ferrox-models::latent_moe’s module doc comment); moe_intermediate_dim is the per-expert FFN size (3072). Per-layer byte layout of one store-backed Kimi routed expert’s combined buffer: w1_packed, w1_scale, w2_packed, w2_scale, w3_packed, w3_scale concatenated in that fixed order. Every expert in a Kimi layer has identical dims, so one layout serves the layer.

Enums

KimiLoadError

Functions

load_block_residual

load_dense_mlp

Loads the dense leading layer’s feed-forward block (real tensor names under {prefix}.mlp.*) – Kimi K3’s layer 0 only (first_k_dense_replace=1).

load_f32_vec

Reads any real tensor as an owned f32 vector, dispatching on its real declared dtype (F32 direct, BF16 dequantized) – exposed pub since not every real weight (e.g. the per-layer input_layernorm.weight/post_attention_layernorm.weight, which aren’t nested inside KdaAttnWeights/MlaAttnWeights/ DenseMlpWeights/BlockResidualWeights) has a dedicated loader function above.

load_kda_attn

Loads one KDA-attention layer’s weights (real tensor names under {prefix}.self_attn.*). num_heads/head_dim/hidden_dim must match the real config (Kimi K3: 96/128/7168) – passed explicitly rather than hardcoded so this loader can also be exercised against small synthetic on-disk fixtures in tests.

load_kimi_checkpoint

load_kimi_checkpoint_with_expert_cache

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.

load_kimi_expert

load_kimi_layer

Loads any one layer (KDA or Gated-MLA attention, dense or MoE FFN), dispatching on kind/is_dense – pass ModelConfig::layer_attention_kind(layer_idx)/ ModelConfig::layer_is_dense(layer_idx) for Kimi K3’s real per-layer topology.

load_latent_moe

Loads one full MoE layer: the gate (with its real aux-loss-free e_score_correction_bias), the shared down/up latent projections + norm, every routed expert (n_experts, real MXFP4), and the shared expert (real BF16, on the full hidden_dim, not the latent space – see ferrox-models::latent_moe’s module doc comment).

load_mla_attn

Loads one Gated-MLA-attention layer’s weights (real tensor names under {prefix}.self_attn.*).