Epic: Engine Inference Core Stage E - MoE + hybrid families

[tlkahn]

Epic: Stage E - MoE + hybrid families

The two non-dense forward paths from the reference: forwardMoeWith (expert dispatch via mlx_gather_qmm/mlx_gather_mm) and forwardHybridWith (conv/SSM layers with their own cache entries). Green + demoable: all remaining generate families pass parity gates.

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

Part of: #51
Depends on: Stage D (#55, done)

Decision 5 superseded (2026-07-20): mlx-lm implements DeepSeek V3/V3.2/V4 natively on MLX using safetensors with standard mlx ops - no GGUF or ds4 engine needed. Port from mlx-lm's deepseek_v3.py / deepseek_v32.py + mla.py, not from ds4.zig. Target checkpoint: mlx-community/DeepSeek-V4-Flash-4bit. The Stage A safetensors loader works. Also: qwen3_5_moe is the separate MODEL_QWEN3_5_MOE enum value from Stage A1. See spec - Resolved design decisions.

Decision 11 (2026-07-20): do not collapse DeepSeek variants into MODEL_DEEPSEEK_V4. mlx-lm keeps separate modules; V3.2 adds a DSA Indexer + dual cache (real forward/cache fork). Enum grain:

• deepseek_v3 -> MODEL_DEEPSEEK_V3 (MLA + MoE)
• deepseek_v3_2 / deepseek_v32 -> MODEL_DEEPSEEK_V32 (MLA + MoE + DSA)
• deepseek_v4 -> MODEL_DEEPSEEK_V4 (own value; characterize Flash before any alias)

Shared code is primitives (fwd_mla_*, MoE gate, sanitize), not one family tag. A1 interim debt still maps all four strings to MODEL_DEEPSEEK_V4 - correct before or at start of E2 (#109). Reject messages must name the actual family.

---

Phases

E1 - MoE infrastructure

• Engine Inference Core E1: MoE infrastructure #104 MoE infrastructure (shared by deepseek + qwen3_5_moe)
  • E1: MoE block (fwd_moe) - router, gather_qmm/mm, shared experts #107 MoE block (fwd_moe)
  • E1: DeepSeek-style MoE gate (sigmoid, group top-k, correction bias) #108 DeepSeek-style MoE gate

E2 - DeepSeek V3 / V3.2 / V4 (MLA + MoE; separate families)

• Engine Inference Core E2: DeepSeek V3 / V3.2 / V4 (MLA + MoE) #105 DeepSeek V3 / V3.2 / V4
  • E2: Family split correction (decision 11) + DeepSeek config fields #109 Family split correction (decision 11) + config fields
  • E2: MultiLinear (fwd_multi_linear) primitive for MLA #110 MultiLinear (fwd_multi_linear)
  • E2: MLA attention + MLA KV cache #111 MLA attention + MLA KV cache
  • E2: DeepSeek weight sanitization + weights_expected_names #112 Weight sanitization + weights_expected_names
  • E2: Dense/MoE layer mixing + DeepSeek V3 forward wiring #113 Dense/MoE layer mixing + V3 forward wiring
  • E2: V3.2 DSA indexer (MODEL_DEEPSEEK_V32) #114 V3.2 DSA indexer
  • E2: V4 characterization + forward (DeepSeek-V4-Flash-4bit) #115 V4 characterization + forward
  • E2: DeepSeek parity gates (mlx-lm oracle) #116 DeepSeek parity gates (mlx-lm oracle)

E3 - qwen3_5_moe, lfm2, nemotron_h

• Engine Inference Core E3: qwen3_5_moe, lfm2, nemotron_h #106 qwen3_5_moe, lfm2, nemotron_h
  • E3: Hybrid cache (SSM/conv) + unified hybrid support gate #117 Hybrid cache + unified hybrid support gate
  • E3: qwen3_5_moe (linear attention + MoE FFN) #118 qwen3_5_moe
  • E3: lfm2 (short-conv + attention hybrid) #119 lfm2
  • E3: nemotron_h (mamba-style SSM + attention/MLP) #120 nemotron_h
  • E3: Parity gates (qwen3_5_moe, lfm2, nemotron_h) #121 E3 parity gates

---

Dependency sketch

E1 (#104)
 ├─ #107 fwd_moe
 └─ #108 DeepSeek MoE gate
        │
        ▼
E2 (#105) ─────────────────────────────────────┐
 ├─ #109 family split + config (first)          │
 ├─ #110 MultiLinear ──► #111 MLA+cache         │
 ├─ #112 sanitize                               │
 ├─ #113 V3 forward ◄── E1, #111, #112          │
 ├─ #114 V3.2 DSA ◄── #111, #113                │
 ├─ #115 V4 ◄── #113 (+ characterization)       │
 └─ #116 parity ◄── as families land            │
                                                 │
E3 (#106)  (MoE part needs E1; hybrid is local) │
 ├─ #117 hybrid cache + gate unify              │
 ├─ #118 qwen3_5_moe ◄── E1, #117               │
 ├─ #119 lfm2 ◄── #117                          │
 ├─ #120 nemotron_h ◄── #117                    │
 └─ #121 parity ◄── as families land            │

Tracking notes (from prior #56 comments)

• Strict weight validation: when implementing each family, extend weights_expected_names() (PR feat(engine): stage A - loading foundations (mlxbridge safetensors surface, weight loader) #62 pattern). Owned by E2: DeepSeek weight sanitization + weights_expected_names #112 (DeepSeek) and per-family E3 issues (E3: qwen3_5_moe (linear attention + MoE FFN) #118-E3: nemotron_h (mamba-style SSM + attention/MLP) #120).
• Hybrid gate unification: collapse pure-dense triple-reject when linear attention lands (PR feat(engine): stage D6 - qwen3_5 dense full-attn deltas #99 finding 6). Owned by E3: Hybrid cache (SSM/conv) + unified hybrid support gate #117.

Out of scope (still under #51, not Stage E)

• model/engine: bert embeddings path (CMD_EMBED + /v1/embeddings) #58 bert embeddings
• engine: KV cache quantization (port kv_quant from mlx-serve) #59 KV cache quantization
• engine: prefix-cache reuse across turns (port prefix_cache from mlx-serve) #60 prefix-cache reuse
• Stage F lifecycle (Engine Inference Core Stage F: lifecycle + memory hardening #57)

[tlkahn]

Tracking note: extend strict weight validation

PR #62 (d85e5b2) added weights_expected_names() with strict expected-tensor-set validation for MODEL_QWEN3. When implementing this stage's families, extend weights_expected_names to emit the correct expected tensor set for each new family (the function currently returns 0 / no-strict-set for non-qwen3, falling back to generic checks).

Key points:

• Each family needs its own list of per-layer matmul suffixes, norm suffixes, and any family-specific tensors
• The orphan-triplet sweep and norm-dtype check already apply to all families
• See qwen3_layer_matmuls / qwen3_layer_norms arrays in weights.c for the pattern

[tlkahn]

Decision 5 superseded - DeepSeek via safetensors/mlx-c, not GGUF/ds4

Investigation found that mlx-lm (Apple's official Python reference) implements DeepSeek V3, V3.2, and V4 natively using standard MLX operations on safetensors checkpoints. The mlx-serve approach of wrapping a separate ds4 engine (GGUF, custom Metal kernels) is mlx-serve-specific, not an MLX platform limitation.

Key architectural components needed:

1. MLA (Multi-head Latent Attention) - compressed KV via LoRA-like bottleneck projections; KV cache stores compressed latent (kv_lora_rank dim, typically 512) + positional embedding, not full K/V. Per-head MultiLinear transforms replace the standard K/V decompression.

2. MoE with sigmoid-scored gating - group expert selection, shared experts always active, dense/MoE layer mixing.

3. Weight sanitization - fp8 dequant, kv_b_proj split into embed_q/unembed_out, expert stacking. This happens at load time in weights.c.

Target checkpoint: mlx-community/DeepSeek-V4-Flash-4bit (safetensors). Parity oracle: mlx-lm (not mlx-serve).

Spec updated: doc/plans/engine-inference-core.md - Stage E rewritten with E1/E2/E3 phasing. weights.c rejection message updated from "requires GGUF" to "MLA not yet implemented (Stage E)".

[tlkahn]

Decision 11 - DeepSeek family enum grain (do not collapse into MODEL_DEEPSEEK_V4)

Reflecting on mlx-lm rather than the interim A1 shortcut: V3 / V3.2 / V4 must not share one model_family_t.

Why the collapse is wrong as an end state

• mlx-lm keeps separate modules (deepseek_v3.py vs deepseek_v32.py); V3.2 is not "V3 with different dims"
• V3.2 adds a DSA Indexer, dual cache (CacheList), and sparse gather/mask - a real attention/cache fork
• Naming everything MODEL_DEEPSEEK_V4 makes reject/log strings lie when the user loaded V3
• There is no deepseek_v4.py in current mlx-lm; V4 must be characterized from Flash weights before any alias
• Same rule as decision 4 (qwen3_5 dense vs MoE): when paths fork, keep separate enum cases

Target mapping

model_type	model_family_t	Notes
deepseek_v3	MODEL_DEEPSEEK_V3	MLA + MoE
deepseek_v3_2, deepseek_v32	MODEL_DEEPSEEK_V32	MLA + MoE + DSA indexer (aliases map together)
deepseek_v4	MODEL_DEEPSEEK_V4	Own value; characterize before aliasing

Shared code = primitives (fwd_mla_*, MultiLinear, MoE gate, sanitize helpers), not one family tag.

A1 interim debt: code still maps all four strings to MODEL_DEEPSEEK_V4. Correct enum + model_family_from_type() + tests + reject messages before or at the start of E2.

Spec + CLAUDE.md updated accordingly (doc/plans/engine-inference-core.md decision 11 + Stage E2 rewrite).

[tlkahn]

Tracking note: unify hybrid/dense gate checks when linear attention lands

Carried forward from PR #99 second review (finding 6, deferred by design).

Today engine_model_check_supported for pure-dense qwen3_5 rejects hybrid via three overlapping checks:

1. linear_num_key_heads > 0 || linear_num_value_heads > 0 (field presence)
2. has_hybrid_layers (parser-detected hybrid; already folded into shared reject_dense_common by ceeb350)
3. Loop over layer_is_global when has_explicit_layer_types (reject any non-full-attention layer)

Check (3) is the most precise for well-formed configs; (1)/(2) are useful early-exits for malformed input but add maintenance weight.

When Stage E adds linear attention / hybrid layers: collapse these into a single unified hybrid gate (one shared helper, family-specific allow-lists for which hybrid shapes are accepted). Do not keep the pure-dense triple-reject as-is once hybrid is a first-class path - the allow path for Stage E families will need the inverse of today's blanket rejects.

Refs: PR #99 review comment 5018410477 finding 6; reply on issue-83 / 2abb872.