[Per-Head Muon] Head geometry that lives outside num_attention_heads/head_dim: Kimi-K3 KDA layers and GLM-5.2's DSA indexer

[alanhuangyoo]

Follow-up to #8367 / #8384, opened at @delock's suggestion so the layout question has its own place to be discussed.

What is not covered today

#8384 tags a projection for per-head Newton-Schulz by reading the head count from AutoTPMeta and the per-head width from the config. That covers standard attention and MLA. Two head-blocked matrices in the very models #8384 was tested against are left on the full-matrix path, because their head geometry is not in the fields the tagger reads.

Kimi-K3's linear-attention layers. inference-optimization/Kimi-K3-0.40B is a hybrid: linear_attn_config gives full_attn_layers: [4, 8] and kda_layers: [1, 2, 3, 5, 6, 7], so 2 of its 8 language layers are MLA and 6 are linear attention with convolution. Instantiated, 318 Muon parameters, 4 tagged — both MLA layers, correctly:

layer kind	parameter	shape	today
MLA (layers 4, 8)	q_b_proj	(768, 256)	per-head, 8 × 96
MLA (layers 4, 8)	kv_b_proj	(1024, 128)	per-head, 8 × 128
KDA (6 layers)	q_proj / k_proj / v_proj	(256, 1024)	full-matrix

The KDA projections are genuinely head-blocked. From the model's own code:

# modeling_kimi_k3_linear.py
self.head_dim = config.linear_attn_config["head_dim"]      # 32
self.num_heads = config.linear_attn_config["num_heads"]    # 8
projection_k_size = self.head_k_dim * self.num_k_heads
self.q_proj = nn.Linear(self.hidden_size, projection_k_size, bias=False)

8 × 32 = 256, which is the shape. They are rejected only because num_attention_heads (8) and head_dim (74) give 8 × 74 = 592, so the shape check declines rather than the tagger recognising them. The head count happens to be right and the width happens to be wrong.

GLM-5.2's sparse-attention indexer. inference-optimization/GLM-5.2-0.8B-A0.8B has indexer.wq_b at (512, 512), blocked on dim 0 by index_n_heads 8 × index_head_dim 64:

# modeling_glm_moe_dsa.py
self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False)

Its count is in index_n_heads, which AutoTPMeta does not read. The indexer's other two matrices are not candidates: wk is (64, 2048), a single shared K with no head axis, and weights_proj is (8, 2048), one row per head.

The questions this needs answered first

1. Does Muon Split apply to linear-attention heads in K3, or is the split specific to full attention?
2. Does it apply to the DSA indexer in GLM-5, or only to the attention proper?

Both are model questions rather than implementation ones, which is why #8384 declines rather than guessing. If either answer is yes, the code change is small; if the answer is no, the current behaviour is already correct and this issue closes.

If the answer is yes

The general fix is probably to read the head count off the owning module rather than adding config keys one architecture at a time. Every case above carries it as an attribute:

module	attributes
KimiLinearAttention	num_heads, head_dim (from linear_attn_config)
GlmMoeDsaIndexer	n_heads, head_dim (from index_*)
MLA / standard attention	num_heads, and a width derivable per kind

Walking named_modules() to map a parameter to its owner and asking the module is architecture-agnostic, and the shape confirmation in #8384 stays as the guard. It is a larger change than #8384 and would want its own review, which is why it is not folded in there.

Not in scope

Whether per-head helps on these layers is a separate question from whether the split is well-defined on them. The measurements in #8384 show the mechanism only has something to act on once heads have differentiated: on trained checkpoints the gradient carries a 6–8x median head-norm imbalance that the whole-matrix path reduces to ~3x and per-head takes to ~1.1x, while both mini checkpoints ship untrained weights whose heads are interchangeable.

I am happy to take this if it is wanted.

[delock]

For KDA, I think from the sprit of muon split, it should also appy to KDA as well. For indexer in GLM, there is no mention that indexer needs to be applied muon split, so they should not be included.

[alanhuangyoo]

Both answered, so this is implemented as you scoped it: PR #8436.

KDA is in — the head count and width come from the module that built the projection, so KimiDeltaAttention's q/k/v go from full-matrix to 8 heads of 32 and Kimi-K3 goes 4 tagged -> 22. The indexer is out, and out for the right reason: wq_b at (512, 512) does confirm on index_n_heads 8 x index_head_dim 64, so it is declined by the module not being attention rather than by the shape, and there is a test asserting the geometry confirms so that stays true if someone changes the guard.

GLM-5.2 is unchanged at 12 tagged. Both models run end to end on their real checkpoints with the loss falling.

Detail worth your view is in the PR: I made the module question opt-in (ask modules that identify as attention) rather than opt-out (ask everything, skip a named list). Opt-out would have tagged the indexer by default, which is the opposite of what you concluded here; the cost is that a linear-attention module not named for attention — Qwen3-Next's Qwen3NextGatedDeltaNet — keeps today's behaviour until a marker is added. There is a test pinning that so the cost is visible rather than silent.

[alanhuangyoo]

@delock — both halves of your ruling are implemented in #8436 (0a1fbe1, the same day you answered). I should have closed the loop here rather than only on the PR; doing that now.

KDA is in. Kimi-K3's linear-attention q_proj / k_proj / v_proj are tagged by reading the head count and width off the owning module rather than the top-level config, because linear_attn_config's num_heads: 8, head_dim: 32 is what builds them and num_attention_heads: 8, head_dim: 74 is what the config-derived candidates describe. The other matrices of the same layer (g_proj, f_b_proj, b_proj, o_proj) stay on the full-matrix path.

The indexer is out. GLM-5.2's indexer.wq_b is exactly index_n_heads 8 x index_head_dim 64, so an unconditional read of a module's own n_heads/head_dim would have tagged it. Only modules that say they are attention are asked, and the comment in deepspeed/__init__.py cites this issue for why:

#8420 settled that Muon Split covers attention, not the indexer that picks which keys attention will see.

Tests, 1×H20:

test_linear_attention_geometry_comes_from_the_owning_module[q_proj/k_proj/v_proj]   PASSED
test_the_other_matrices_of_a_linear_attention_layer_stay_whole[g_proj/f_b_proj/b_proj/o_proj]  PASSED
test_the_sparse_attention_indexer_is_not_tagged                                    PASSED
12 passed

test_the_sparse_attention_indexer_is_not_tagged builds the real wq_b at (512, 512) and asserts it is declined — so if someone later makes the module read unconditional, this issue's answer fails as a test rather than being forgotten.

#8436 stacks on #8384 and is the only one of my Muon PRs that has to wait for it. Happy to close this issue now, or leave it open until #8436 lands — whichever you prefer for tracking.