Commit History
-
feat: add Per-Head Muon (Kimi K3 §2.5) via head_keys [skip-build]
Partition attention Q/K/V momentum matrices along the head dimension (dim 0) and orthogonalize each head's block separately, instead of running Newton-Schulz on the full projection matrix. Equalizes the update scale across heads (K3 tech report §2.5). - head_keys={"q_proj": 128, ...}: component-level name matching (same rule as expert_keys) mapping to per-head dim-0 row stride; MLA strides supported (e.g. wq_b/wkv_b) - NS granularity only: momentum stays elementwise; matching params are viewed [H*d, N] -> [H, d, N] and run through the existing batched bmm Newton-Schulz; adjusted LR uses the head-block shape (consistent with the per-expert path) - Sharding-agnostic: all three paths (base / distributed_muon / parallel) apply the head partition to the full gathered matrix, so results are identical under any DP/TP/HSDP placement with no communication changes; _muon_state.head_dim switches the compute call - conftest: make transformers import lazy so synthetic tests run in…ebeb3a1
