[muon] Per-head Newton-Schulz for attention projections - #8384

Open

alanhuangyoo wants to merge 12 commits into

deepspeedai:masterfrom

alanhuangyoo:feat/per-head-muon

Commits on Sep 3, 2026

[muon] Support MLA up-projections and read head counts from AutoTPMeta
alanhuangyoo
committed
870ee00

Commits on Sep 4, 2026

[muon] Cover MLA models that have no q_lora_rank
alanhuangyoo
committed
c6a7793
[muon] Keep muon_update compiled
alanhuangyoo
committed
c3e9334

Commits on Sep 6, 2026

[muon] Resolve head geometry by candidate, and report what the flag did
alanhuangyoo
committed
37740f6
[muon] Re-resolve the head count against the shard under tensor parallelism
alanhuangyoo
committed
943fd03
[muon] Read the layer's shape, not the ZeRO-3 partition's
alanhuangyoo
committed
5d16495

Commits on Sep 7, 2026

Merge remote-tracking branch 'upstream/master' into feat/per-head-muon
alanhuangyoo
committed
61b4253

Commits on Sep 8, 2026

Consolidate the per-head Muon tests into two files
alanhuangyoo
committed
789a1e7