Skip to content

E2: DeepSeek weight sanitization + weights_expected_names #112

Description

@tlkahn

DeepSeek weight sanitization + weights_expected_names

Load-time transforms so mlx-lm safetensors checkpoints become the tensors the MLA/MoE forward path expects.

Spec: doc/plans/engine-inference-core.md § E2

Part of: #105 (E2) · Epic: #56
Depends on: Family split + config fields (knows ranks/dims/expert counts)

Tasks

  • fp8 dequant (block-128 weight_scale_inv)
  • kv_b_proj -> embed_q + unembed_out split (reshape [n_heads, qk_nope_head_dim + v_head_dim, kv_lora_rank], split, requantize)
  • Expert stacking (experts.N.gate/up/down_proj -> stacked switch_mlp.gate/up/down_proj)
  • MTP layer removal
  • V3.2 keeps indexer tensors (wq_b, wk, weights_proj, ...)
  • Extend weights_expected_names() for MODEL_DEEPSEEK_V3 / V32 / V4 (per-layer matmul suffixes, norms, family-specific tensors)
  • CPU tests for split/stack/dequant on fixture or synthetic tensors

Tracking note

PR #62 added strict expected-tensor-set validation for qwen3. Non-qwen3 currently returns 0 (no strict set). Each DeepSeek family needs its own list; orphan-triplet sweep and norm-dtype check already apply globally.

Metadata

Metadata

Assignees

Labels

engine-inferenceEngine inference core (issue #51) feature areaenhancementNew feature or request

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions