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

ARCHITECTURE_V4.md

@@ -149,6 +149,11 @@ another.
 
 **Crate:** `aarambh-studio-nn` (`mla.rs`) | **Depends on:** v1 §6.3 (GQA/RoPE), v2 §21 (YaRN/NTK), v3 §29 (`HybridAttentionSchedule`)
 
+> **Status: Implemented in v4.0.0-alpha.1 (Phase 41).** `MlaAttention` and
+> `MlaCache` ship in `crates/aarambh-studio-nn/src/mla.rs`; `AttentionKind::LatentMLA`
+> and `MlaConfig` extend the schedule; the partial-checkpoint retrofit and
+> `--kv-cache-report` are wired through. See `docs/phase41_mla.md` for usage.
+
 ### The Problem
 
 v3 gave the model two ways to reduce the cost of a growing KV cache:

CHANGELOG.md

@@ -2,6 +2,42 @@
 
 > From first principles. From zero. From Rust.
 
+## [4.0.0-alpha.1] - 2026-07-31
+
+### Added
+
+- **Phase 41 — Multi-Head Latent Attention (MLA):** A third attention kind
+  (`AttentionKind::LatentMLA`) joins Full and Gated DeltaNet in the
+  `HybridAttentionSchedule`, completing the attention family v3 began (linear,
+  sparse, latent-compressed). MLA layers cache a single low-rank latent vector
+  (`c_kv`, width `latent_dim`) plus a small dedicated rotary key slice
+  (`rope_head_dim`) per token, reconstructing per-head keys and values at
+  attention time through trained up-projection weights that are never cached.
+  - New `aarambh-studio-nn::mla` module (`MlaAttention`, `MlaCache`) with
+    decoupled RoPE (nope half from the latent, rope half separately cached),
+    inference/training/batched-decode/capture paths, and QAT-wrapped
+    projections (`QatTarget::Mla`).
+  - `HybridAttentionSchedule` extended with `mla_layers: Vec<usize>` and
+    `mla: Option<MlaConfig>`; `mla_layers` takes precedence over the
+    `full_attention_every_n` rule and the DSA override. A schedule with an
+    empty `mla_layers` reproduces v3.0.0 exactly.
+  - New `MlaConfig` (`latent_dim`, `nope_head_dim`, `rope_head_dim`, `n_heads`,
+    `value_head_dim`) with dimension derivation and validation.
+  - Model integration: per-layer MLA build, `HybridKvCache::Mla` allocation,
+    named-tensor export (`blocks.{i}.mla.*`), and weight lookup.
+  - Partial-checkpoint retrofit extended: `.mla.` tensors are freshly
+    initialised alongside `.deltanet.` and `.dsa.` while every shared tensor
+    loads bit-exactly (`RetrofitLoadReport.initialized_mla_tensors`).
+  - `aarambh-studio eval --kv-cache-report` prints per-layer bytes/token by
+    attention kind (no checkpoint required).
+  - New configs: `configs/mla_smoke.toml`, `configs/medium_hybrid_mla.toml`,
+    `configs/large_hybrid_mla.toml`; new scripts
+    `scripts/phase41_prepare_mla_retrofit.sh`, `scripts/phase41_smoke.sh`;
+    new doc `docs/phase41_mla.md`.
+  - For the Medium hybrid MLA config, MLA per-token cache = 528 elements vs the
+    1024-element GQA baseline — a ~1.94× reduction on retrofitted layers at
+    long context.
+
 ## [3.0.0] - 2026-07-25
 
 ### Added

Cargo.toml

@@ -23,7 +23,7 @@ members = [
 resolver = "2"
 
 [workspace.package]
-version = "3.0.0"
+version = "4.0.0-alpha.1"
 edition = "2024"
 rust-version = "1.89"
 description = "From first principles. From zero. From Rust."

README.md

@@ -18,7 +18,9 @@ with hybrid Gated DeltaNet, DeepSeek Sparse Attention,
 fine-grained MoE with shared experts, Multi-Token Prediction (MTP), on-policy
 distillation, native quantization-aware training, native video/document input,
 bounded long-horizon tool-use chains, persistent forgetting diagnostics, and
-Max thinking mode (16,384-token budget).
+Max thinking mode (16,384-token budget). **v4.0.0-alpha.1** begins the v4 arc
+with Multi-Head Latent Attention (Phase 41) — a third attention kind that
+compresses the KV cache into a single low-rank latent per token.
 
 > [!IMPORTANT]
 > This is a source and engineering project. It does not publish crates to
@@ -30,7 +32,7 @@ Max thinking mode (16,384-token budget).
 | Area | Capabilities |
 |---|---|
 | Model | RMSNorm, RoPE, GQA, SwiGLU, KV cache, tied embeddings, Tiny to Large configs |
-| Efficient architecture | YaRN/NTK/linear RoPE scaling, Gated DeltaNet, learned block-sparse DSA, fine-grained MoE, MTP |
+| Efficient architecture | YaRN/NTK/linear RoPE scaling, Gated DeltaNet, learned block-sparse DSA, Multi-Head Latent Attention (MLA), fine-grained MoE, MTP |
 | Training | BPE data pipeline, AdamW, cosine schedule, gradient accumulation/clipping, checkpoint resume, BF16 CUDA, single-node multi-GPU, on-policy distillation, native INT4/INT8 QAT |
 | Fine-tuning | SFT, LoRA, QLoRA, DoRA, QDoRA, VLM adapters, GRPO, DPO, QDPO, tool-call tuning |
 | Inference | Greedy/sampled decoding, streaming, thinking budgets, external or one-checkpoint MTP speculation, tool grammar, caller-executed chains |
@@ -119,7 +121,9 @@ See the phase-specific docs for full walkthroughs with smoke fixtures:
 | Document understanding | [docs/phase36_document.md](docs/phase36_document.md) |
 | OpenAI-compatible server | [docs/inference-server.md](docs/inference-server.md) |
 | Evaluation & forgetting diagnostics | [docs/phase38_forgetting.md](docs/phase38_forgetting.md) |
+| Multi-Head Latent Attention (MLA) | [docs/phase41_mla.md](docs/phase41_mla.md) |
 | Quantization (GPTQ, QAT, GGUF) | [docs/phase34_qat.md](docs/phase34_qat.md) |
+| MLA hybrid attention & KV-cache report | `aarambh-studio eval --kv-cache-report` + [docs/phase41_mla.md](docs/phase41_mla.md) |
 | Fine-tuning (SFT, adapters, GRPO, DPO) | `aarambh-studio finetune --help` |
 | Self-learning | [SELF_LEARNING_V3.md](SELF_LEARNING_V3.md) |
 

ROADMAP_V4.md

@@ -186,6 +186,14 @@ beyond what each phase's Dependency Policy note allows.
 
 **Duration:** 10–14 days | **Hardware:** Kaggle (free quota)
 
+> **Status: Implemented in v4.0.0-alpha.1.** The `mla.rs` module, `MlaConfig`,
+> the three-way `HybridAttentionSchedule`, the partial-checkpoint retrofit path,
+> `--kv-cache-report`, the smoke/retrofit scripts, and the full test suite
+> (reconstruction tolerance, decoupled-RoPE split, cache-size, partial-load,
+> backward-reachability) are all in place. See `docs/phase41_mla.md` and
+> `CHANGELOG.md` §4.0.0-alpha.1. The checkbox list below is the original plan,
+> preserved for traceability.
+
 ### Goal
 A third attention kind — latent KV compression — addable to the
 `HybridAttentionSchedule` v3 §29 introduced, so a model can now mix

SELF_LEARNING_V4.md

@@ -106,6 +106,16 @@ their own pass:
 
 The short version: **nothing changes, and that is the point.**
 
+> **Status: Verified for v4.0.0-alpha.1 (Phase 41).** The
+> `mla_training_backward_reaches_mla_parameters` test confirms gradients reach
+> the MLA down-projection (`kv_a_proj`), latent norm, value up-projection
+> (`up_v`), and output projection (`o_proj`) — the §42 reachability argument.
+> MLA's Q/K gradient behaviour on the CPU candle-fallback attention path
+> matches the existing GQA path (full Q/K gradients flow under the CUDA/flash
+> path used in real training); MLA is wired into the identical attention path,
+> so self-learning's gradient orthogonalisation reaches MLA weights
+> consistently with every other attention kind.
+
 Online GRPO's math (`SELF_LEARNING.md` §5) operates on log-probabilities
 of generated tokens under the current policy. It has no dependency on
 *how* those log-probabilities were computed internally — whether a