[muon] Per-head Newton-Schulz for fused QKV projections

[0z5a]

Implements the fused-QKV part of #8367, on top of #8384.

What this PR does

A fused QKV weight holds Q, K and V in one matrix. #8384 left it on the full-matrix path because the three sections do not share a head count under GQA. They do not have to.

Every head is head_dim contiguous rows of dim 0 in both layouts that ship:

• sectioned (cat([q_proj, k_proj, v_proj]), what most training stacks store): n_q query heads, then n_kv key heads, then n_kv value heads, each head_dim rows;
• interleaved per KV group (Falcon's view(..., num_kv_heads, num_heads // num_kv_heads + 2, head_dim), GPT-NeoX's per-head q, k, v): each group is n_q/n_kv query heads, one key head, one value head, each head_dim rows.

Both layouts present the same row partition in a different order, and per-head Newton-Schulz is block-wise with the same max(1, head_dim / in_features)**0.5 scale on every block - so the update does not depend on which layout the weight uses, and the Q/K/V labelling never enters it. GQA changes how many heads a section holds, not where the blocks are.

So a fused weight needs no section bookkeeping: it is tagged with rows // head_dim head blocks, which for 8 query heads over 2 KV heads is 12 rather than 8 or 3 * 8. The candidate requires the exact fused total (num_attention_heads + 2 * num_kv_heads) * head_dim, which is also what declines a transposed Conv1D weight (GPT-2's c_attn, [hidden, 3*hidden], whose dim 0 is the input axis).

The tag rides on muon_num_heads, so muon_update, the six ZeRO/DDP call sites and the config surface need no new API.

Not in this PR

• Falcon's query_key_value stays on the full-matrix path, and a test pins why: its config reports num_attention_heads key/value heads while multi_query=True makes the module build a single one, so the config's fused total (192 rows) does not describe the weight (80). That is a pre-existing AutoTPMeta reading rather than something the per-head split introduces; the exact-total check sends it to the full-matrix path instead of splitting the wrong blocks.
• AutoTP. It repacks a fused weight into this rank's piece of each section, or interleaves it, so a rank's layout is not asserted by this change. The data-parallel path is what is trained and measured below.

Tests

$ pytest tests/unit/runtime/zero/test_per_head_muon.py                     # CPU
64 passed

$ pytest tests/unit/v1/ops/muon/test_per_head_muon_accelerator.py          # 2 x L20, ZeRO 1/2/3
26 passed

$ pre-commit run --files <changed files>
all hooks passed

CPU coverage added: fused tagging under GQA (8 query heads over 2 KV heads is 12 blocks, not 8), the exact-total guard, the Conv1D decline, the real architectures, and sectioned/interleaved equivalence on a real gradient - the test builds one fused weight in each layout and requires the updates to be the same permutation of each other, which is the property the change rests on. Accelerated coverage adds a fused training model across ZeRO 1/2/3 at world_size 2.

End-to-end training

2 x NVIDIA L20 46 GB, ZeRO-2, bf16, torch 2.10.0+cu128, dp=2. A 17,177,088-parameter LLaMA-shaped model whose attention holds Q|K|V in one qkv_proj (6 layers, hidden 512, 8 query heads over 2 KV heads, head_dim 64), 120 optimizer steps on 983,040 tokens of real text, global batch 16, both arms from an identical initialisation.

arm	mean loss (first 10)	mean loss (last 10)	tok/s	peak GiB
per_head_muon: false	4.0266	2.6789	85,411	1.68
per_head_muon: true	4.0227	2.6781	82,368	1.77

• With the flag on, blocks.*.qkv_proj.weight carries muon_num_heads = 12 and is the only tagged parameter; with the flag off every tag is None.
• One optimizer step from the shared initialisation moves 34/34 trainable weights differently between the arms (fused-weight L1 difference 0.235 against a weight L1 of 8679.4), so the per-head branch executes in the real training path instead of being tagged and skipped.
• On a real backward-pass gradient for the fused weight, the per-head update norms land within 1.30x of each other where the full-matrix ones span 2.78x; the two updates differ by 2.46 in relative norm.

The two loss curves are indistinguishable, and that is expected rather than a result: both arms start from an untrained initialisation whose heads are interchangeable, so the split has nothing to separate - the same outcome #8384 reported for its models. The mechanism measurement is what shows the feature does something; a convergence claim needs trained checkpoints and is not made here.

[alanhuangyoo]

Owner geometry is lost for a fused leaf once this rebases on #8436 (blocking)

_owner_candidate only answers for QUERY and KV, so a module in _LINEAR_ATTENTION_OWNERS with a fused projection gets [] and falls back to the config candidates — and in_proj_qkv is in _FUSED_QKV_LEAVES. Those leaves were declined outright before this PR, so this is the coincidence the whitelist was added to prevent. Either extend the owner path to the fused kind, or decline fused leaves under a listed owner. The rebase itself is clean — I ran it locally, 106 tests pass on 2 GPUs.

Fused and separate layouts now agree (non-blocking, worth putting in the description)

The full-matrix path scales a fused qkv by max(1, rows/cols)**0.5, up to sqrt(3) for [3h, h], while separately stored q/k/v each get 1.0. Today the step size depends on how the checkpoint happens to store the projection; per-head makes both 1.0. That's a correctness argument, stronger than the coverage one.

Also checked CodeGen — its 4-block reshape still lands on head boundaries when 4 divides num_heads, so it's covered rather than mis-split.