feat(optim): per-head Muon for attention q/k/v (heads_per_block, default off)

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -439,6 +439,9 @@ class MuonConfig(BaseOptimizerConfig):
     betas2: float = Field(0.95, ge=0)
     """β2 for the AdamW/Lion sub-optimizer used on non-Muon params."""
 
+    heads_per_block: int | None = Field(None, ge=1)
+    """Orthogonalize the attention query, key and value projections in blocks of this many heads, each block on its own (per-head Muon, Kimi K3 report, section 2.5; ``1`` = per head). GQA key/value projections are split by their own heads; ``o_proj`` stays one matrix. If None, each projection is one matrix."""
+
 
 class SignSGDConfig(BaseOptimizerConfig):
     type: Literal["sign_sgd"] = "sign_sgd"

src/prime_rl/trainer/models/afmoe/modeling_afmoe.py

@@ -15,9 +15,11 @@
 
 from prime_rl.trainer.models.base import PreTrainedModelPrimeRL
 from prime_rl.trainer.models.layers.attn import (
+    attention_head_partitions,
     flash_attn_2_varlen_op,
     flash_attn_3_varlen_op,
     flash_attn_4_varlen_op,
+    head_blocks,
 )
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput
 from prime_rl.trainer.models.layers.mlp import FeedForward
@@ -82,6 +84,10 @@ def __init__(self, config: AfmoeAttentionConfig):
         self.q_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps))
         self.k_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps))
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        projections = (self.q_proj, self.k_proj, self.v_proj, self.gate_proj)
+        return {proj.weight: head_blocks(proj.out_features, self.head_dim, heads_per_block) for proj in projections}
+
 
 class AfmoeFlashAttention(AfmoeAttentionBase):
     """AFMoE attention using Flash Attention varlen functions."""
@@ -476,6 +482,9 @@ def init_buffers_post_meta(self):
             )
             rotary_emb.inv_freq.copy_(inv_freq)
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        return attention_head_partitions((layer.self_attn for layer in self.model.layers), heads_per_block)
+
 
 __all__ = [
     "AfmoeForCausalLM",

src/prime_rl/trainer/models/base.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import Literal
 
-from torch import Tensor
+from torch import Tensor, nn
 from transformers.modeling_utils import PreTrainedModel
 
 from prime_rl.utils.cp import CPContext
@@ -155,5 +155,10 @@ def init_buffers_post_meta(self) -> None:
         """
         raise NotImplementedError(f"init_buffers_post_meta is not implemented for {self.__class__.__name__}")
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        """Muon row partitions that split every attention head projection into blocks of ``heads_per_block`` heads
+        (``optim.heads_per_block``). Each model covers all of its attention projections or raises."""
+        raise NotImplementedError(f"optim.heads_per_block is not supported for {self.__class__.__name__}")
+
 
 __all__ = ["ALL_CP_STYLES", "CPStyle", "CPSupport", "PreTrainedModelPrimeRL"]

src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py

@@ -27,7 +27,7 @@
 from prime_rl.trainer.models.base import PreTrainedModelPrimeRL
 from prime_rl.trainer.models.glm4_moe.configuration_glm4_moe import Glm4MoeConfig
 from prime_rl.trainer.models.glm4_moe.converting_glm4_moe import conversion_chain
-from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig
+from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig, attention_head_partitions
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput
 from prime_rl.trainer.models.layers.mlp import FeedForward
 from prime_rl.trainer.models.layers.moe import MoE, MoEArgs
@@ -342,5 +342,8 @@ def init_buffers_post_meta(self):
         # TODO: Init TT MoE buffers
         # I think .to_empty() on gpu by default fills 0 so we are ok but this might not be guaranteed behavior
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        return attention_head_partitions((layer.self_attn for layer in self.model.layers), heads_per_block)
+
 
 __all__ = ["Glm4MoeConfig", "Glm4MoePreTrainedModel", "Glm4MoeModel", "Glm4MoeForCausalLM"]

src/prime_rl/trainer/models/gpt_oss/attention.py

@@ -10,7 +10,7 @@
 
 from prime_rl.trainer.models.fusions import fuse_qkv_projections
 from prime_rl.trainer.models.gpt_oss.configuration_gpt_oss import GptOssConfig
-from prime_rl.trainer.models.layers.attn import flash_attn_4_varlen_func
+from prime_rl.trainer.models.layers.attn import flash_attn_4_varlen_func, qkv_head_partitions
 from prime_rl.trainer.models.layers.rotary_emb import apply_rotary_pos_emb
 
 
@@ -51,6 +51,9 @@ def __init__(self, config: GptOssConfig, layer_idx: int) -> None:
         nn.init.normal_(self.sinks, mean=0.0, std=config.initializer_range)
         self.flash_attn = flash_attn_4_varlen_func
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        return qkv_head_partitions(self, heads_per_block, self.head_dim)
+
     def compute_attention(
         self,
         query: torch.Tensor,

src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py

@@ -10,6 +10,7 @@
     is_hf_state_dict,
     is_prime_state_dict,
 )
+from prime_rl.trainer.models.layers.attn import attention_head_partitions
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, VanillaOutputLinear
 from prime_rl.trainer.models.layers.moe import GroupedExperts, MoE, TokenChoiceTopKRouter
 from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig
@@ -217,6 +218,9 @@ def init_buffers_post_meta(self) -> None:
                 module.tokens_per_expert.zero_()
                 module.routing_confidence_sum.zero_()
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        return attention_head_partitions((layer.self_attn for layer in self.model.layers), heads_per_block)
+
 
 __all__ = [
     "GptOssDecoderLayer",

src/prime_rl/trainer/models/laguna/modeling_laguna.py

@@ -16,7 +16,7 @@
 from prime_rl.trainer.models.base import PreTrainedModelPrimeRL
 from prime_rl.trainer.models.laguna.configuration_laguna import LagunaConfig
 from prime_rl.trainer.models.laguna.converting_laguna import conversion_chain
-from prime_rl.trainer.models.layers.attn import AttentionConfig, FlashAttention
+from prime_rl.trainer.models.layers.attn import AttentionConfig, FlashAttention, attention_head_partitions
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput
 from prime_rl.trainer.models.layers.mlp import FeedForward
 from prime_rl.trainer.models.layers.moe import MoE, MoEArgs
@@ -119,6 +119,11 @@ def __init__(self, config: LagunaConfig, layer_idx: int, num_heads: int, flash_a
             self.g_proj = nn.Linear(config.hidden_size, gate_size, bias=False)
         self.o_proj = nn.Linear(num_heads * self.head_dim, config.hidden_size, bias=config.attention_bias)
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        if self.gating:
+            raise NotImplementedError("optim.heads_per_block does not cover Laguna's attention output gate (g_proj)")
+        return super().muon_head_partitions(heads_per_block)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -427,6 +432,9 @@ def init_buffers_post_meta(self) -> None:
                 if module.router.selection_bias is not None:
                     module.router.selection_bias.zero_()
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        return attention_head_partitions((layer.self_attn for layer in self.model.layers), heads_per_block)
+
 
 __all__ = [
     "LagunaForCausalLM",

src/prime_rl/trainer/models/layers/attn.py

@@ -1,4 +1,5 @@
 import functools
+from collections.abc import Iterable
 from dataclasses import dataclass
 from typing import Literal
 
@@ -418,6 +419,37 @@ class AttentionConfig:
     output_bias: bool = False
 
 
+def head_blocks(rows: int, head_rows: int, heads_per_block: int) -> tuple[int, ...]:
+    """Split a projection's ``rows`` (``head_rows`` per head) into blocks of ``heads_per_block`` heads."""
+    block_rows = heads_per_block * head_rows
+    if rows % block_rows:
+        raise ValueError(f"optim.heads_per_block = {heads_per_block} does not divide {rows // head_rows} heads")
+    return (block_rows,) * (rows // block_rows)
+
+
+def qkv_head_partitions(attn: nn.Module, heads_per_block: int, q_head_rows: int) -> dict[nn.Parameter, tuple[int, ...]]:
+    """Muon head blocks of an attention module's query, key and value projections, packed in ``qkv_proj`` or
+    separate. ``q_head_rows`` are the query projection's rows per head."""
+    head_rows = (q_head_rows, attn.head_dim, attn.head_dim)
+    blocks = [
+        head_blocks(rows, rows_per_head, heads_per_block) for rows, rows_per_head in zip(attn.qkv_sizes, head_rows)
+    ]
+    if attn.qkv_proj is not None:
+        return {attn.qkv_proj.weight: sum(blocks, ())}
+    return {proj.weight: proj_blocks for proj, proj_blocks in zip((attn.q_proj, attn.k_proj, attn.v_proj), blocks)}
+
+
+def attention_head_partitions(
+    attention_modules: Iterable[nn.Module], heads_per_block: int
+) -> dict[nn.Parameter, tuple[int, ...]]:
+    """The Muon head blocks of all ``attention_modules``."""
+    return {
+        param: blocks
+        for attn in attention_modules
+        for param, blocks in attn.muon_head_partitions(heads_per_block).items()
+    }
+
+
 # TODO: Does torch compile support config._attn_implementation forking?
 # If so, we can combine FlashAttention variants into one class
 # Otherwise, do ABC or something to make the signatures match
@@ -474,6 +506,9 @@ def project_qkv(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.
             return self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
         return self.qkv_proj(hidden_states).split(self.qkv_sizes, dim=-1)
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        return qkv_head_partitions(self, heads_per_block, self.head_dim)
+
     def _compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens, max_seqlen):
         """Run the flash attention kernel. q/k/v are [total_tokens, heads, dim]."""
         kwargs: dict = {"causal": True}

src/prime_rl/trainer/models/llama/modeling_llama.py

@@ -34,7 +34,7 @@
 from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
 
 from prime_rl.trainer.models.base import PreTrainedModelPrimeRL
-from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig
+from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig, attention_head_partitions
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput
 from prime_rl.trainer.models.layers.mlp import FeedForward
 from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig
@@ -319,3 +319,6 @@ def init_buffers_post_meta(self):
                 rotary_emb.config, rotary_emb.inv_freq.device
             )
             rotary_emb.inv_freq.copy_(inv_freq)
+
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        return attention_head_partitions((layer.self_attn for layer in self.model.layers), heads_per_block)

src/prime_rl/trainer/models/minimax_m2/modeling_minimax_m2.py

@@ -10,7 +10,7 @@
 from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
 
 from prime_rl.trainer.models.base import PreTrainedModelPrimeRL
-from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig
+from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig, attention_head_partitions
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput
 from prime_rl.trainer.models.layers.moe import MoE, MoEArgs
 from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig
@@ -277,6 +277,9 @@ def init_buffers_post_meta(self):
             )
             rotary_emb.inv_freq.copy_(inv_freq)
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        return attention_head_partitions((layer.self_attn for layer in self.model.layers), heads_per_block)
+
 
 __all__ = [
     "MiniMaxM2ForCausalLM",

src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py

@@ -3,7 +3,7 @@
 from transformers.modeling_outputs import BaseModelOutput
 
 from prime_rl.trainer.models.base import CPSupport, PreTrainedModelPrimeRL
-from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig
+from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig, attention_head_partitions
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, VanillaOutputLinear
 from prime_rl.trainer.models.layers.mlp import FeedForward
 from prime_rl.trainer.models.layers.moe import GroupedExperts, MoE, TokenChoiceTopKRouter
@@ -239,6 +239,10 @@ def init_buffers_post_meta(self) -> None:
                 module.tokens_per_expert.zero_()
                 module.routing_confidence_sum.zero_()
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        attention_layers = (layer for layer in self.model.layers if layer.layer_type == "attention")
+        return attention_head_partitions((layer.self_attn for layer in attention_layers), heads_per_block)
+
 
 __all__ = [
     "NemotronHDecoderLayer",

src/prime_rl/trainer/models/qwen3/modeling_qwen3.py

@@ -11,7 +11,7 @@
 from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
 
 from prime_rl.trainer.models.base import PreTrainedModelPrimeRL
-from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig
+from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig, attention_head_partitions
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput
 from prime_rl.trainer.models.layers.mlp import FeedForward
 from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig
@@ -275,5 +275,8 @@ def init_buffers_post_meta(self):
             )
             rotary_emb.inv_freq.copy_(inv_freq)
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        return attention_head_partitions((layer.self_attn for layer in self.model.layers), heads_per_block)
+
 
 __all__ = ["Qwen3ForCausalLM", "Qwen3Model", "Qwen3PreTrainedModel"]

src/prime_rl/trainer/models/qwen3_5/attention.py

@@ -1,7 +1,7 @@
 import torch
 from torch import nn
 
-from prime_rl.trainer.models.layers.attn import AttentionConfig, FlashAttention
+from prime_rl.trainer.models.layers.attn import AttentionConfig, FlashAttention, qkv_head_partitions
 from prime_rl.trainer.models.layers.rotary_emb import apply_rotary_pos_emb
 from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
 from prime_rl.trainer.models.qwen3_5.norm import Qwen3_5RMSNorm
@@ -37,6 +37,10 @@ def __init__(self, config: Qwen3_5TextConfig, attention_implementation: str) ->
         self.q_norm = Qwen3_5RMSNorm(self.head_dim, config.rms_norm_eps)
         self.k_norm = Qwen3_5RMSNorm(self.head_dim, config.rms_norm_eps)
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        # With the output gate, q_proj holds [query; gate] for each head.
+        return qkv_head_partitions(self, heads_per_block, self.qkv_sizes[0] // self.num_heads)
+
     def forward(
         self,
         hidden_states: torch.Tensor,

src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py

@@ -3,6 +3,7 @@
 from transformers.modeling_outputs import BaseModelOutput
 
 from prime_rl.trainer.models.base import ALL_CP_STYLES, CPSupport, PreTrainedModelPrimeRL
+from prime_rl.trainer.models.layers.attn import attention_head_partitions
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, VanillaOutputLinear
 from prime_rl.trainer.models.layers.mlp import FeedForward
 from prime_rl.trainer.models.layers.moe import GroupedExperts, MoE, TokenChoiceTopKRouter
@@ -384,6 +385,12 @@ def init_buffers_post_meta(self) -> None:
                 if module.router.selection_bias is not None:
                     module.router.selection_bias.zero_()
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        if self.is_vlm:
+            raise NotImplementedError("optim.heads_per_block does not cover the Qwen3.5 vision encoder")
+        attention_layers = (layer for layer in self.model.layers if layer.layer_type == "full_attention")
+        return attention_head_partitions((layer.self_attn for layer in attention_layers), heads_per_block)
+
 
 __all__ = [
     "Qwen3_5Attention",

src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py

@@ -1,6 +1,7 @@
 import torch
 from torch import nn
 
+from prime_rl.trainer.models.layers.attn import head_blocks
 from prime_rl.trainer.models.qwen3_8_flash_next.indexer import SparseAttentionIndexer
 from prime_rl.trainer.models.qwen3_8_flash_next.norm import RMSNorm
 from prime_rl.trainer.models.qwen3_8_flash_next.rotary_embedding import apply_rotary_embedding
@@ -48,6 +49,14 @@ def __init__(
         # Resolve once before torch.compile traces the layer; the loader caches the module.
         self.indexed_attention = prime_kernels.load("indexed_attention").indexed_attention
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        # q_proj holds [query; gate] for each head.
+        return {
+            self.q_proj.weight: head_blocks(self.q_proj.out_features, 2 * self.head_dim, heads_per_block),
+            self.k_proj.weight: head_blocks(self.k_proj.out_features, self.head_dim, heads_per_block),
+            self.v_proj.weight: head_blocks(self.v_proj.out_features, self.head_dim, heads_per_block),
+        }
+
     def forward(
         self,
         hidden_states: torch.Tensor,

src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py

@@ -3,6 +3,7 @@
 from transformers.modeling_outputs import BaseModelOutput
 
 from prime_rl.trainer.models.base import CPSupport, PreTrainedModelPrimeRL
+from prime_rl.trainer.models.layers.attn import attention_head_partitions
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, VanillaOutputLinear
 from prime_rl.trainer.models.layers.moe import MoE
 from prime_rl.trainer.models.qwen3_8_flash_next.attention import IndexedGatedAttention
@@ -306,6 +307,11 @@ def init_buffers_post_meta(self) -> None:
                 if module.router.selection_bias is not None:
                     module.router.selection_bias.zero_()
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        text_model = getattr(self.model, "language_model", self.model)
+        attention_layers = (layer for layer in text_model.layers if layer.layer_type != "linear_attention")
+        return attention_head_partitions((layer.self_attn for layer in attention_layers), heads_per_block)
+
 
 __all__ = [
     "Qwen3_8FlashNextDecoderLayer",

src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py

@@ -27,7 +27,7 @@
 from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
 
 from prime_rl.trainer.models.base import PreTrainedModelPrimeRL
-from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig
+from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig, attention_head_partitions
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput
 from prime_rl.trainer.models.layers.mlp import FeedForward
 from prime_rl.trainer.models.layers.moe import MoE, MoEArgs
@@ -349,6 +349,9 @@ def init_buffers_post_meta(self):
         # TODO: Init TT MoE buffers
         # I think .to_empty() on gpu by default fills 0 so we are ok but this might not be guaranteed behavior
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        return attention_head_partitions((layer.self_attn for layer in self.model.layers), heads_per_block)
+
 
 __all__ = [
     "Qwen3MoeForCausalLM",