feat: Phase 41 — Multi-Head Latent Attention (MLA)

[aarambh-darshan]

Summary

Implements Phase 41 — Multi-Head Latent Attention (MLA) from
ROADMAP_V4.md §41 (ARCHITECTURE_V4.md §55). MLA is the third and
final attention kind in the v4 hybrid schedule, completing the family v3
began: Gated DeltaNet (linear), DSA (sparse), and now MLA
(latent KV compression). MLA layers cache a single low-rank latent
vector (c_kv) plus a small dedicated rotary key slice per token, and
reconstruct per-head keys and values at attention time through trained
up-projection weights that are never cached — cutting KV-cache
memory per token substantially at long context while preserving
per-head expressiveness.

A schedule with zero mla_layers reproduces v3.0.0 exactly, holding
the backward-compatibility discipline every attention change since v1
has kept.

Bumps the workspace version 3.0.0 → 4.0.0-alpha.1 (the Phase 41
milestone tag).

---

What changed

aarambh-studio-core

• AttentionKind::LatentMLA variant.
• QatTarget::Mla variant (added to QatConfig::default().targets).
• New MlaConfig { latent_dim, nope_head_dim, rope_head_dim, n_heads, value_head_dim } with resolve() / validate() / cache_width().
• HybridAttentionSchedule extended with mla_layers: Vec<usize> and
  mla: Option<MlaConfig>. kind_for_layer gives MLA precedence
  over both the full_attention_every_n rule and the DSA override.
  New validate() + resolved_mla() helpers.

aarambh-studio-nn (new mla.rs)

• MlaCache — compressed c_kv latent + k_rope slice, dynamic and
  preallocated paths (mirror of KVCache).
• MlaAttention with decoupled RoPE (dedicated rope_head_dim
  RoPE cache; the host transformer's head-dim RoPE is not reused because
  a compressed latent cannot carry an already-rotated key). Projections:
  q_proj, kv_a_proj (down → c_kv), kv_a_norm (RMSNorm over the
  latent), up_k / up_v (per-head up-projection weights, trained but
  never cached), k_rope_proj (rotary key slice, shared across heads,
  cached), o_proj.
  Methods: forward, forward_train, forward_decode_batch,
  forward_with_capture, named_tensors, get_weight.
• TokenMixer::Mla + HybridKvCache::Mla wired through block.rs,
  kvcache.rs, lib.rs (as_mla / as_mla_mut accessors). Reuses the
  candle fallback attention kernel, which tolerates value_head_dim != nope_head_dim + rope_head_dim.

aarambh-studio-model

• build_mla() helper + LatentMLA build arm.
• empty_kv_cache_with_capacity() allocates MlaCache for MLA layers.
• insert_mixer_tensors() exports blocks.{i}.mla.* checkpoint names.
• get_weight() resolves the mla. prefix.
• New kv_cache_report() + KvCacheLayerReport powering
  --kv-cache-report.

aarambh-studio-weights

• RetrofitLoadReport.initialized_mla_tensors.
• .mla. tensor init path in load_retrofit_into_varmap_with_moe
  (shared tensors still load bit-exactly; MLA tensors are freshly
  initialised, same pattern as .deltanet. / .dsa.).

aarambh-studio (CLI)

• eval --kv-cache-report flag + run_kv_cache_report() — prints
  per-layer bytes/token by attention kind. No checkpoint required,
  only the config.

Configs / scripts / docs

• configs/mla_smoke.toml, configs/medium_hybrid_mla.toml,
  configs/large_hybrid_mla.toml.
• scripts/phase41_prepare_mla_retrofit.sh, scripts/phase41_smoke.sh.
• docs/phase41_mla.md (new).
• CHANGELOG.md §4.0.0-alpha.1; README (MLA capability row + version +
  docs links); ROADMAP_V4.md §41 status note; ARCHITECTURE_V4.md §55
  status note; SELF_LEARNING_V4.md §42 status note.

---

Backward compatibility

A schedule with an empty mla_layers and no mla block reproduces
v3.0.0 kind_for_layer exactly, and resolved_mla() returns None —
verified by schedule_with_zero_mla_layers_matches_v3_exactly. Five
existing HybridAttentionSchedule construction sites in tests were
updated with the new (defaulted) fields; behaviour unchanged.

---

Tests

All roadmap-required Phase 41 proof obligations are implemented and pass:

Test	Crate	Proves
schedule_with_zero_mla_layers_matches_v3_exactly	core	empty mla_layers == v3.0.0
mla_layers_take_precedence_over_every_n_and_dsa_override	core	MLA wins over every-n + DSA
mla_config_resolves_derived_dimensions / _rejects_odd_rope_head_dim	core	MlaConfig derivation + validation
mla_reconstructed_kv_matches_reference_full_attention_within_tolerance	nn	latent round-trip bounded; train == inference
decoupled_rope_nope_split_preserves_relative_position_encoding	nn	nope half position-invariant, rope half changes with offset
mla_kv_cache_bytes_per_token_is_smaller_than_full_or_gqa_baseline	nn	(latent_dim + rope_head_dim) < 2 * n_kv_heads * head_dim
partial_checkpoint_load_preserves_non_mla_layer_weights_exactly	weights	retrofit loads shared tensors bit-exactly, initialises 7 MLA tensors
mla_model_forwards_and_cached_forward_matches_full_forward	model	cached decode matches full forward; MLA cache grows per token
mla_training_backward_reaches_mla_parameters	model	gradients reach kv_a_proj / kv_a_norm / up_v / o_proj (SELF_LEARNING_V4 §42 reachability)
mla_kv_cache_report_shows_compressed_footprint	model	--kv-cache-report reports the compressed MLA footprint

---

CI gates (all green locally)

• cargo fmt --all --check
• cargo check --workspace --all-targets --locked
• cargo clippy --workspace --all-targets --locked -- -D warnings -D clippy::undocumented_unsafe_blocks
• cargo test --workspace --no-fail-fast --locked → 444 passed, 0 failed
• RUSTDOCFLAGS="-D warnings -D missing_docs" cargo doc --workspace --no-deps --locked
• scripts/phase28_release_audit.sh → passed for aarambh-studio 4.0.0-alpha.1
• find scripts -name '*.sh' | bash -n
• cargo build --release -p aarambh-studio --locked
• CLI smoke: --version + all 10 subcommand --help + eval --kv-cache-report

Note: the release LTO link (lto="thin", codegen-units=1) was verified with
LTO disabled to fit this sandbox's per-command time cap; the source is
identical and the full LTO build is what runs on GitHub Actions.

---

MLA smoke (verified end-to-end)

$ aarambh-studio eval --config configs/mla_smoke.toml --kv-cache-report
layer  kind             bytes/tok  note
0      latent_mla           320    latent_dim + rope_head_dim per token (compressed)
1      gated_deltanet         0    fixed recurrent state (not per-token)
total bytes/token: 320 | all-full baseline: 1024 | ratio: 0.312 (3.1x smaller)

• 2-step CPU training (mla_smoke.toml): loss 4.91 → 4.73, ppl 136 → 113.
• Trained checkpoint contains all 7 blocks.0.mla.* tensors +
  blocks.1.deltanet.* (verified via safetensors header parse).
• For the Medium hybrid config: MLA per-token cache = 528 elements vs the
  1024-element GQA baseline — a ~1.94× reduction on retrofitted layers.

---

How to review

1. git fetch && git checkout feat/phase-41-mla
2. cargo test --workspace --no-fail-fast --locked (444 pass)
3. cargo run -p aarambh-studio -- eval --config configs/mla_smoke.toml --kv-cache-report
4. Read docs/phase41_mla.md for the full mechanism, config, and retrofit recipe.

---

Scope / non-goals

• CUDA flash / fused MLA kernels are future work; MLA reuses the candle
  fallback attention path (mechanism + memory win are in place).
• YaRN/NTK long-context scaling applies unchanged to the host transformer's
  full-attention layers; MLA's dedicated rope_head_dim slice uses base RoPE.
• This PR is Phase 41 only. Phases 42–55 (audio, sparse MoE dispatch,
  multi-node, test-time scaling, RLAIF, tool execution, multi-agent, RAG,
  model merging, public server, chat templating, red-team, model card, final
  release) remain on the v4 roadmap.

Refs: ROADMAP_V4.md §41 · ARCHITECTURE_V4.md §55 · SELF_LEARNING_V4.md §42