E2: MLA attention + MLA KV cache

[tlkahn]

MLA attention + MLA KV cache

Multi-head Latent Attention base path for MODEL_DEEPSEEK_V3 (shared primitive for V3.2/V4). New attention path - not a variant of fwd_attention.

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

Part of: #105 (E2) · Epic: #56
Depends on: MultiLinear sibling issue

Tasks - attention (fwd_mla_attention)

• Q LoRA compression: q_a_proj -> q_a_layernorm -> q_b_proj; first layer may use direct q_proj (no LoRA)
• Compressed KV: kv_a_proj_with_mqa -> split kv_latent (normed) + k_pe
• Split Q into q_nope + q_pe; RoPE only on pe parts
• Per-head MultiLinear: embed_q and unembed_out
• SDPA with pe_scores as mask (q_pe @ k_pe^T)
• Decode-path optimization: embed_q on q_nope (not kv_latent), unembed_out post-SDPA
• Attention scale: q_head_dim^(-0.5), optional mscale_all_dim from rope_scaling

Tasks - KV cache

• Store compressed kv_latent ([B, 1, T, kv_lora_rank]) + k_pe ([B, 1, T, qk_rope_head_dim]), not full K/V
• Prefill append + decode step update paths
• Wire into engine cache lifetime (create/reset/free on engine thread only)

Notes

• Cross-check against mlx-lm mla.py / deepseek_v3.py, not mlx-serve ds4
• V3.2 dual-cache (CacheList) is a separate issue on top of this base cache