diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py
index f8f5a813b38..ca98033b0ca 100644
--- a/megatron/core/optimizer/__init__.py
+++ b/megatron/core/optimizer/__init__.py
@@ -66,6 +66,8 @@
     HAVE_EMERGING_OPTIMIZERS,
     _create_emerging_optimizer,
     _get_qkv_split_shapes,
+    _localize_qkv_split_shapes,
+    _qkv_split_groups_are_complete,
 )
 from .fully_sharded_optimizer import FullyShardedOptimizer
 from .grad_scaler import ConstantGradScaler, DynamicGradScaler
@@ -771,6 +773,8 @@ def _get_megatron_emerging_optimizer(
         raise ValueError(f"Unsupported emerging optimizer: {eopt_name}")
     if config.fp16:
         raise ValueError('emerging optimizer with fp16 is not supported.')
+    if config.muon_split_qkv_per_head and not config.muon_split_qkv:
+        raise ValueError("muon_split_qkv_per_head requires muon_split_qkv=True")
 
     if pg_collection is None:
         pg_collection = ProcessGroupCollection.use_mpu_process_groups()
@@ -779,25 +783,108 @@ def _get_megatron_emerging_optimizer(
 
     # Tag parameters with optimizer-specific attributes (expert_tp, is_qkv).
     for model_chunk in model_chunks:
-        qkv_split_shapes = None
         for name, param in model_chunk.named_parameters():
             if not param.requires_grad:
                 continue
             if 'experts' in name and 'shared' not in name:
                 param.expert_tp = True
-            # TODO(deyuf): support MLA
-            if 'linear_qkv.weight' in name and len(param.shape) == 2:
-                if qkv_split_shapes is None:
-                    qkv_split_shapes = _get_qkv_split_shapes(model_chunk.config)
-                if param.shape[0] % sum(qkv_split_shapes) == 0:
-                    param.is_qkv = True
-                    param.qkv_split_shapes = qkv_split_shapes
+            qkv_layout = getattr(param, 'qkv_layout', None)
+            if (qkv_layout is not None or 'linear_qkv.weight' in name) and len(param.shape) == 2:
+                if qkv_layout is not None:
+                    qkv_split_shapes = _get_qkv_split_shapes(
+                        qkv_layout, split_qkv_per_head=config.muon_split_qkv_per_head
+                    )
+                    logical_split_shapes = (
+                        qkv_split_shapes
+                        if config.muon_split_qkv_per_head
+                        else qkv_split_shapes * qkv_layout.num_groups
+                    )
                 else:
+                    # Backward compatibility for custom QKV modules that do not annotate
+                    # their weight with the owning attention layer's logical layout.
+                    qkv_split_shapes = _get_qkv_split_shapes(
+                        model_chunk.config, split_qkv_per_head=config.muon_split_qkv_per_head
+                    )
+                    logical_split_shapes = (
+                        qkv_split_shapes
+                        if config.muon_split_qkv_per_head
+                        else qkv_split_shapes * model_chunk.config.num_query_groups
+                    )
+
+                tp_group = (
+                    pg_collection.expt_tp
+                    if getattr(param, 'expert_tp', False)
+                    else pg_collection.tp
+                )
+                tp_size = get_pg_size(tp_group)
+                tp_rank = get_pg_rank(tp_group)
+                gtp_remat_group = (
+                    (
+                        pg_collection.expt_gtp_remat
+                        if getattr(param, 'expert_tp', False)
+                        else pg_collection.gtp_remat
+                    )
+                    if getattr(param, 'is_gtp_weight_remat', False)
+                    else None
+                )
+                gtp_size = get_pg_size(gtp_remat_group)
+                gtp_rank = get_pg_rank(gtp_remat_group)
+
+                qkv_gtp_pad_length = (
+                    int(getattr(param, 'pad_length', 0))
+                    if getattr(param, 'is_gtp_weight_remat', False)
+                    else 0
+                )
+                physical_tp_local_rows = param.shape[0] * gtp_size
+                if not 0 <= qkv_gtp_pad_length < physical_tp_local_rows:
+                    raise RuntimeError(
+                        f"Invalid Muon QKV GTP padding for {name}: "
+                        f"pad_length={qkv_gtp_pad_length}, "
+                        f"physical_tp_local_rows={physical_tp_local_rows}"
+                    )
+                logical_tp_local_rows = physical_tp_local_rows - qkv_gtp_pad_length
+                expected_logical_rows = logical_tp_local_rows * tp_size
+                if expected_logical_rows != sum(logical_split_shapes):
                     log_single_rank(
                         logger,
                         logging.DEBUG,
                         f"Emerging optimizer QKV split skipped for {name}: "
-                        f"shape={tuple(param.shape)}, split_shapes={qkv_split_shapes}",
+                        f"logical_rows={sum(logical_split_shapes)}, "
+                        f"local_rows={param.shape[0]}, tp_size={tp_size}, "
+                        f"gtp_remat_size={gtp_size}, "
+                        f"gtp_pad_length={qkv_gtp_pad_length}",
+                    )
+                    param.is_qkv = False
+                    param.qkv_split_shapes = None
+                    param.qkv_split_shapes_global = None
+                    param.qkv_gtp_pad_length = 0
+                    param.qkv_split_groups_are_complete = False
+                    param.qkv_split_heads_are_complete = False
+                    continue
+
+                param.is_qkv = True
+                param.qkv_split_shapes_global = logical_split_shapes
+                param.qkv_gtp_pad_length = qkv_gtp_pad_length
+                local_start = tp_rank * logical_tp_local_rows + gtp_rank * param.shape[0]
+                if config.muon_split_qkv_per_head:
+                    if qkv_gtp_pad_length > 0:
+                        # A padded GTP shard does not have a purely logical local layout.
+                        # Force the per-head path to use qkv_split_shapes_global instead.
+                        param.qkv_split_shapes = None
+                        param.qkv_split_heads_are_complete = False
+                    else:
+                        param.qkv_split_shapes, param.qkv_split_heads_are_complete = (
+                            _localize_qkv_split_shapes(
+                                qkv_split_shapes, local_start=local_start, local_rows=param.shape[0]
+                            )
+                        )
+                else:
+                    param.qkv_split_shapes = qkv_split_shapes
+                    param.qkv_split_groups_are_complete = (
+                        qkv_gtp_pad_length == 0
+                        and _qkv_split_groups_are_complete(
+                            qkv_split_shapes, local_start=local_start, local_rows=param.shape[0]
+                        )
                     )
 
     # Apply optimizer-specific default param overrides (e.g. muon: non-linear -> adam).
diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py
index 53ac956b35c..8bce97892b4 100644
--- a/megatron/core/optimizer/emerging_optimizers.py
+++ b/megatron/core/optimizer/emerging_optimizers.py
@@ -11,9 +11,12 @@
 import inspect
 import logging
 from dataclasses import dataclass, field
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as package_version
 from typing import Any, Callable, Dict, Literal, Optional, get_args
 
 import torch
+from packaging.version import Version
 from torch.optim.optimizer import ParamsT
 
 from megatron.core.process_groups_config import ProcessGroupCollection
@@ -43,6 +46,18 @@
 
 logger = logging.getLogger(__name__)
 
+try:
+    EMERGING_OPTIMIZERS_VERSION = Version(package_version("emerging-optimizers"))
+except PackageNotFoundError:
+    EMERGING_OPTIMIZERS_VERSION = Version("0")
+
+_BATCHED_NEWTON_SCHULZ_MIN_VERSION = Version("0.3.0")
+
+
+def _supports_batched_newton_schulz() -> bool:
+    """Return whether Emerging-Optimizers supports batched Newton-Schulz."""
+    return EMERGING_OPTIMIZERS_VERSION >= _BATCHED_NEWTON_SCHULZ_MIN_VERSION
+
 
 def get_supported_coefficient_types() -> tuple[str, ...]:
     """Return the coefficient types supported by the installed emerging_optimizers.
@@ -130,11 +145,29 @@ def _is_nonlinear_or_embedding(param):
     return getattr(param, 'is_embedding_or_output_parameter', False) or len(param.shape) != 2
 
 
-def _get_qkv_split_shapes(model_cfg) -> list[int]:
-    """Compute QKV split shapes from model config."""
+def _get_qkv_split_shapes(model_cfg, split_qkv_per_head: bool = False) -> list[int]:
+    """Compute fused QKV split shapes from logical attention layout metadata.
+
+    Args:
+        model_cfg: Transformer config or an owning attention layer's ``QKVLayout`` metadata.
+        split_qkv_per_head: Return one split size per physical attention head. When false,
+            return the per-query-group Q, gate (if present), K, and V projection widths.
+    """
+    if hasattr(model_cfg, 'projection_split_shapes'):
+        if split_qkv_per_head:
+            return list(model_cfg.per_head_split_shapes) * model_cfg.num_groups
+        return list(model_cfg.projection_split_shapes)
+
     query_projection_size = (
         model_cfg.num_attention_heads // model_cfg.num_query_groups * model_cfg.kv_channels
     )
+    if split_qkv_per_head:
+        num_query_heads_per_group = model_cfg.num_attention_heads // model_cfg.num_query_groups
+        per_group_shapes = [model_cfg.kv_channels] * num_query_heads_per_group
+        if getattr(model_cfg, 'attention_output_gate', False):
+            per_group_shapes += [model_cfg.kv_channels] * num_query_heads_per_group
+        per_group_shapes += [model_cfg.kv_channels, model_cfg.kv_channels]
+        return per_group_shapes * model_cfg.num_query_groups
     if getattr(model_cfg, 'attention_output_gate', False):
         return [
             query_projection_size,
@@ -145,6 +178,46 @@ def _get_qkv_split_shapes(model_cfg) -> list[int]:
     return [query_projection_size, model_cfg.kv_channels, model_cfg.kv_channels]
 
 
+def _localize_qkv_split_shapes(
+    global_split_shapes: list[int], local_start: int, local_rows: int
+) -> tuple[list[int], bool]:
+    """Intersect global per-head split sizes with a rank-local contiguous row range.
+
+    Returns:
+        The physical rank-local split sizes and whether every intersected head is complete.
+    """
+    local_stop = local_start + local_rows
+    local_split_shapes = []
+    all_heads_complete = True
+    head_start = 0
+    for head_rows in global_split_shapes:
+        head_stop = head_start + head_rows
+        overlap_start = max(head_start, local_start)
+        overlap_stop = min(head_stop, local_stop)
+        if overlap_start < overlap_stop:
+            overlap_rows = overlap_stop - overlap_start
+            local_split_shapes.append(overlap_rows)
+            all_heads_complete &= overlap_rows == head_rows
+        head_start = head_stop
+
+    if sum(local_split_shapes) != local_rows:
+        raise RuntimeError(
+            f"Muon per-head QKV local range [{local_start}, {local_stop}) is outside "
+            f"the global split shape with {sum(global_split_shapes)} rows"
+        )
+    return local_split_shapes, all_heads_complete
+
+
+def _qkv_split_groups_are_complete(
+    split_shapes: list[int], local_start: int, local_rows: int
+) -> bool:
+    """Return whether a local row range contains only complete fused QKV groups."""
+    split_width = sum(split_shapes)
+    if split_width <= 0:
+        raise ValueError(f"Muon QKV split shapes must sum to a positive size: {split_shapes}")
+    return local_start % split_width == 0 and local_rows % split_width == 0
+
+
 # ===========================================================================
 # Registry – populated below only when emerging_optimizers is installed.
 # ===========================================================================
@@ -169,6 +242,7 @@ def __init__(
         weight_decay: float = 0.01,
         use_decoupled_weight_decay: bool = True,
         split_qkv: bool = False,
+        split_qkv_per_head: bool = False,
         is_qkv_fn: Callable[[torch.Tensor], bool] | None = None,
         qkv_split_shapes: list[int] | None = None,
         fp32_matmul_prec: str = "medium",
@@ -181,6 +255,8 @@ def __init__(
     ) -> None:
         if num_ns_steps < 1:
             raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}")
+        if split_qkv_per_head and not split_qkv:
+            raise ValueError("split_qkv_per_head requires split_qkv=True")
 
         def scaled_orthogonalize_fn(
             grad: torch.Tensor,
@@ -211,6 +287,7 @@ def scaled_orthogonalize_fn(
         self.pg_collection = pg_collection
         self.tp_mode = tp_mode
         self.split_qkv = split_qkv
+        self.split_qkv_per_head = split_qkv_per_head
         self.is_qkv_fn = is_qkv_fn
         self.qkv_split_shapes = qkv_split_shapes
 
@@ -230,6 +307,130 @@ def scaled_orthogonalize_fn(
             scaled_orthogonalize_fn=scaled_orthogonalize_fn,
         )
 
+    def _get_gtp_remat_group(self, p):
+        """Return the GTP-remat process group for a parameter, if configured."""
+        is_expert = getattr(p, 'expert_tp', False)
+        return (
+            (self.pg_collection.expt_gtp_remat if is_expert else self.pg_collection.gtp_remat)
+            if self.pg_collection
+            else None
+        )
+
+    def _gather_qkv_grad(self, p, grad, tp_group, expected_rows, gather_gtp=True):
+        """Reconstruct a fused QKV gradient and record how to restore its local shard."""
+        gathered_grad = grad
+        gtp_slice = None
+        gtp_remat_group = self._get_gtp_remat_group(p)
+        gtp_pad_length = int(getattr(p, "qkv_gtp_pad_length", 0))
+        if gtp_pad_length < 0:
+            raise RuntimeError(f"Muon QKV GTP padding must be non-negative: {gtp_pad_length}")
+        if (
+            gather_gtp
+            and gtp_remat_group is not None
+            and get_pg_size(gtp_remat_group) > 1
+            and getattr(p, 'is_gtp_weight_remat', False)
+        ):
+            gtp_size = get_pg_size(gtp_remat_group)
+            gtp_rank = get_pg_rank(gtp_remat_group)
+            gtp_local_rows = gathered_grad.shape[0]
+            shards = [torch.empty_like(gathered_grad) for _ in range(gtp_size)]
+            torch.distributed.all_gather(shards, gathered_grad, gtp_remat_group)
+            gathered_grad = torch.cat(shards, dim=0)
+            if gtp_pad_length >= gathered_grad.shape[0]:
+                raise RuntimeError(
+                    "Invalid Muon QKV GTP padding after gathering: "
+                    f"pad_length={gtp_pad_length}, gathered_rows={gathered_grad.shape[0]}"
+                )
+            if gtp_pad_length > 0:
+                gathered_grad = gathered_grad[:-gtp_pad_length]
+            gtp_slice = (gtp_rank, gtp_local_rows, gtp_pad_length)
+        elif gtp_pad_length > 0:
+            raise RuntimeError(
+                "Muon QKV has GTP padding but its GTP-remat shards were not gathered"
+            )
+
+        tp_slice = None
+        if gathered_grad.shape[0] != expected_rows:
+            partition_dim = getattr(p, "partition_dim", None)
+            if partition_dim != 0 or tp_group is None:
+                raise RuntimeError(
+                    f"Muon QKV split shape mismatch: grad_shape={tuple(gathered_grad.shape)}, "
+                    f"expected_rows={expected_rows}, partition_dim={partition_dim}"
+                )
+            tp_size = get_pg_size(tp_group)
+            if gathered_grad.shape[0] * tp_size != expected_rows:
+                raise RuntimeError(
+                    "Muon QKV split cannot reconstruct the global tensor: "
+                    f"local_grad_shape={tuple(gathered_grad.shape)}, tp_size={tp_size}, "
+                    f"expected_rows={expected_rows}"
+                )
+            tp_rank = get_pg_rank(tp_group)
+            tp_local_rows = gathered_grad.shape[0]
+            shards = [torch.empty_like(gathered_grad) for _ in range(tp_size)]
+            torch.distributed.all_gather(shards, gathered_grad, tp_group)
+            gathered_grad = torch.cat(shards, dim=0)
+            tp_slice = (tp_rank, tp_local_rows)
+
+        if gathered_grad.shape[0] != expected_rows:
+            raise RuntimeError(
+                "Muon QKV split shape mismatch after gathering: "
+                f"grad_shape={tuple(gathered_grad.shape)}, expected_rows={expected_rows}"
+            )
+        return gathered_grad, tp_slice, gtp_slice
+
+    @staticmethod
+    def _restore_local_qkv_grad(gathered_grad, tp_slice, gtp_slice):
+        """Restore the TP and GTP-remat shards recorded by ``_gather_qkv_grad``."""
+        if tp_slice is not None:
+            tp_rank, tp_local_rows = tp_slice
+            gathered_grad = gathered_grad[tp_rank * tp_local_rows : (tp_rank + 1) * tp_local_rows]
+        if gtp_slice is not None:
+            gtp_rank, gtp_local_rows, gtp_pad_length = gtp_slice
+            if gtp_pad_length > 0:
+                gathered_grad = torch.nn.functional.pad(gathered_grad, (0, 0, 0, gtp_pad_length))
+            gathered_grad = gathered_grad[
+                gtp_rank * gtp_local_rows : (gtp_rank + 1) * gtp_local_rows
+            ]
+        return gathered_grad.contiguous()
+
+    def _orthogonalize_split_qkv(self, grad, split_shapes, orthogonalize_fn):
+        """Split and reconstruct Megatron's interleaved fused QKV update."""
+        if grad.ndim != 2:
+            raise RuntimeError(f"Muon QKV gradient must be 2D, got {grad.ndim}D")
+        if not split_shapes or any(size <= 0 for size in split_shapes):
+            raise RuntimeError(f"Muon QKV split shapes must be positive: {split_shapes}")
+        split_width = sum(split_shapes)
+        if self.split_qkv_per_head:
+            if grad.shape[0] != split_width:
+                raise RuntimeError(
+                    f"Muon per-head QKV split shape mismatch: grad_shape={tuple(grad.shape)}, "
+                    f"split_shapes={split_shapes}"
+                )
+            if len(set(split_shapes)) == 1 and _supports_batched_newton_schulz():
+                # A 3D input selects Emerging-Optimizers' batched Newton-Schulz path.
+                head_rows = split_shapes[0]
+                return orthogonalize_fn(grad.view(len(split_shapes), head_rows, -1)).view_as(grad)
+            return torch.cat(
+                [orthogonalize_fn(head) for head in torch.split(grad, split_shapes, dim=0)], dim=0
+            )
+
+        if grad.shape[0] % split_width != 0:
+            raise RuntimeError(
+                f"Muon QKV split shape mismatch: grad_shape={tuple(grad.shape)}, "
+                f"split_shapes={split_shapes}"
+            )
+        num_query_groups = grad.shape[0] // split_width
+        grouped_grad = grad.view(num_query_groups, split_width, -1)
+        projection_grads = torch.split(grouped_grad, split_shapes, dim=1)
+        projection_grads = [
+            projection.reshape(-1, grad.shape[-1]) for projection in projection_grads
+        ]
+        projection_grads = [
+            orthogonalize_fn(projection).view(num_query_groups, -1, grad.shape[-1])
+            for projection in projection_grads
+        ]
+        return torch.cat(projection_grads, dim=1).view_as(grad)
+
     def scaled_orthogonalize_fn_with_gtp_remat(self, p, grad, tp_group, partition_dim):
         """All-gather grad along GTP_remat/EGTP_remat dim 0, orthogonalize, then slice back.
 
@@ -239,12 +440,7 @@ def scaled_orthogonalize_fn_with_gtp_remat(self, p, grad, tp_group, partition_di
         When GTP_remat is inactive this is a plain passthrough to scaled_orthogonalize_fn.
         """
         # TODO: Clean up code that determines if parameter is a MoE layer and which TP group to use
-        is_expert = getattr(p, 'expert_tp', False)
-        gtp_remat_group = (
-            (self.pg_collection.expt_gtp_remat if is_expert else self.pg_collection.gtp_remat)
-            if self.pg_collection
-            else None
-        )
+        gtp_remat_group = self._get_gtp_remat_group(p)
 
         if gtp_remat_group is None or get_pg_size(gtp_remat_group) <= 1:
             return self.scaled_orthogonalize_fn(grad, tp_group, partition_dim)
@@ -266,6 +462,74 @@ def scaled_orthogonalize_fn_with_gtp_remat(self, p, grad, tp_group, partition_di
         shard_size = gathered_grad.shape[0] // gtp_remat_size
         return gathered_grad[gtp_rank * shard_size : (gtp_rank + 1) * shard_size].contiguous()
 
+    def _orthogonalize_qkv_per_head(self, p, grad, tp_group):
+        """Orthogonalize every Q, gate, K, and V head independently.
+
+        Split sizes may describe complete heads in the local tensor or the global fused
+        QKV tensor. For a global layout, reconstruct GTP-remat and TP dimension 0 before
+        splitting so heads crossing rank boundaries remain complete.
+        """
+        local_split_shapes = getattr(p, "qkv_split_shapes", None)
+        heads_are_complete = getattr(p, "qkv_split_heads_are_complete", None)
+        has_gtp_padding = int(getattr(p, "qkv_gtp_pad_length", 0)) > 0
+        use_local_layout = not has_gtp_padding and (
+            heads_are_complete is True
+            or (
+                heads_are_complete is None
+                and local_split_shapes is not None
+                and sum(local_split_shapes) == grad.shape[0]
+            )
+        )
+        if use_local_layout:
+            qkv_split_shapes = local_split_shapes
+        else:
+            qkv_split_shapes = getattr(p, "qkv_split_shapes_global", None)
+            if qkv_split_shapes is None:
+                qkv_split_shapes = self.qkv_split_shapes
+        if qkv_split_shapes is None:
+            raise RuntimeError("Muon per-head QKV split requested but qkv_split_shapes is not set")
+        if not qkv_split_shapes or any(size <= 0 for size in qkv_split_shapes):
+            raise RuntimeError(
+                f"Muon per-head QKV split shapes must be positive: {qkv_split_shapes}"
+            )
+
+        expected_rows = sum(qkv_split_shapes)
+        gathered_grad, tp_slice, gtp_slice = self._gather_qkv_grad(
+            p, grad, tp_group, expected_rows, gather_gtp=not use_local_layout
+        )
+
+        gathered_grad = self._orthogonalize_split_qkv(
+            gathered_grad,
+            qkv_split_shapes,
+            lambda head_grad: self.scaled_orthogonalize_fn(
+                head_grad, tp_group=None, partition_dim=None
+            ),
+        )
+
+        return self._restore_local_qkv_grad(gathered_grad, tp_slice, gtp_slice)
+
+    def _orthogonalize_fragmented_qkv(self, p, grad, tp_group, split_shapes):
+        """Orthogonalize projections after reconstructing fragmented query-group blocks."""
+        global_split_shapes = getattr(p, "qkv_split_shapes_global", None)
+        if global_split_shapes is None:
+            raise RuntimeError("Muon fragmented QKV split requires global split shapes")
+        expected_rows = sum(global_split_shapes)
+        if expected_rows % sum(split_shapes) != 0:
+            raise RuntimeError(
+                f"Muon global QKV layout does not contain complete query groups: "
+                f"global_split_shapes={global_split_shapes}, split_shapes={split_shapes}"
+            )
+
+        gathered_grad, tp_slice, gtp_slice = self._gather_qkv_grad(p, grad, tp_group, expected_rows)
+        gathered_grad = self._orthogonalize_split_qkv(
+            gathered_grad,
+            split_shapes,
+            lambda projection_grad: self.scaled_orthogonalize_fn(
+                projection_grad, tp_group=None, partition_dim=None
+            ),
+        )
+        return self._restore_local_qkv_grad(gathered_grad, tp_slice, gtp_slice)
+
     def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> torch.Tensor:
         """Orthogonalize the momentum.
 
@@ -292,36 +556,31 @@ def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> t
             partition_dim = None
 
         if self.split_qkv and self.is_qkv_fn(p):  # type: ignore[misc]
-            grad_shape = grad.shape
+            if self.split_qkv_per_head:
+                return self._orthogonalize_qkv_per_head(p, grad, tp_group)
+
             qkv_split_shapes = getattr(p, "qkv_split_shapes", None)
             if qkv_split_shapes is None:
                 qkv_split_shapes = self.qkv_split_shapes
             if qkv_split_shapes is None:
                 raise RuntimeError("Muon QKV split requested but qkv_split_shapes is not set")
-            qkv_split_dim = sum(qkv_split_shapes)
-            if grad_shape[0] % qkv_split_dim != 0:
-                raise RuntimeError(
-                    f"Muon QKV split shape mismatch: grad_shape={tuple(grad_shape)}, "
-                    f"split_shapes={qkv_split_shapes}"
-                )
+            if (
+                getattr(p, "qkv_split_groups_are_complete", None) is False
+                or int(getattr(p, "qkv_gtp_pad_length", 0)) > 0
+            ) and getattr(p, "qkv_split_shapes_global", None) is not None:
+                return self._orthogonalize_fragmented_qkv(p, grad, tp_group, qkv_split_shapes)
             log_single_rank(
                 logger,
                 logging.DEBUG,
-                f'qkv split grad shape {grad_shape}, split shapes {qkv_split_shapes}',
+                f'qkv split grad shape {grad.shape}, split shapes {qkv_split_shapes}',
             )
-            num_query_groups = grad_shape[0] // qkv_split_dim
-            qkv_grads = torch.split(
-                grad.view(num_query_groups, qkv_split_dim, -1), qkv_split_shapes, dim=1
+            grad = self._orthogonalize_split_qkv(
+                grad,
+                qkv_split_shapes,
+                lambda projection_grad: self.scaled_orthogonalize_fn_with_gtp_remat(
+                    p, projection_grad, tp_group, partition_dim
+                ),
             )
-            qkv_grads = [g.reshape(-1, grad_shape[-1]) for g in qkv_grads]
-
-            qkv_grads = [
-                self.scaled_orthogonalize_fn_with_gtp_remat(p, g, tp_group, partition_dim).view(
-                    num_query_groups, -1, grad_shape[-1]
-                )
-                for g in qkv_grads
-            ]
-            grad = torch.cat(qkv_grads, dim=1).view(grad_shape)
         else:
             grad = self.scaled_orthogonalize_fn_with_gtp_remat(p, grad, tp_group, partition_dim)
         return grad
@@ -344,6 +603,7 @@ class TensorParallelAdaptiveMuon(TensorParallelMuon, AdaptiveMuon):
         weight_decay: Weight decay coefficient.
         use_decoupled_weight_decay: Whether to use decoupled weight decay.
         split_qkv: Whether to split QKV weights for orthogonalization.
+        split_qkv_per_head: Whether to orthogonalize individual Q, gate, K, and V heads.
         is_qkv_fn: Function to determine if a tensor is a QKV weight.
         qkv_split_shapes: Shapes for splitting QKV weights.
         fp32_matmul_prec: Precision for FP32 matrix multiplication.
@@ -367,6 +627,7 @@ def __init__(
         weight_decay: float = 0.01,
         use_decoupled_weight_decay: bool = True,
         split_qkv: bool = False,
+        split_qkv_per_head: bool = False,
         is_qkv_fn: Callable[[torch.Tensor], bool] | None = None,
         qkv_split_shapes: list[int] | None = None,
         fp32_matmul_prec: str = "medium",
@@ -389,6 +650,7 @@ def __init__(
             weight_decay=weight_decay,
             use_decoupled_weight_decay=use_decoupled_weight_decay,
             split_qkv=split_qkv,
+            split_qkv_per_head=split_qkv_per_head,
             is_qkv_fn=is_qkv_fn,
             qkv_split_shapes=qkv_split_shapes,
             fp32_matmul_prec=fp32_matmul_prec,
@@ -436,7 +698,9 @@ def _muon_config_to_kwargs(config, model_chunks, pg_collection) -> Dict[str, Any
     """Convert OptimizerConfig to TensorParallelMuon constructor kwargs."""
     kwargs = _kwargs_from_config(TensorParallelMuon, "muon", config)
     kwargs["is_qkv_fn"] = lambda p: getattr(p, "is_qkv", False)
-    kwargs["qkv_split_shapes"] = _get_qkv_split_shapes(model_chunks[0].config)
+    kwargs["qkv_split_shapes"] = _get_qkv_split_shapes(
+        model_chunks[0].config, split_qkv_per_head=kwargs.get("split_qkv_per_head", False)
+    )
     kwargs["pg_collection"] = pg_collection
     return kwargs
 
diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py
index 24f9a032c47..44bcdd74e85 100644
--- a/megatron/core/optimizer/optimizer_config.py
+++ b/megatron/core/optimizer/optimizer_config.py
@@ -263,6 +263,12 @@ class OptimizerConfig:
     muon_split_qkv: bool = True
     """Whether to split QKV parameters for Muon optimizer."""
 
+    muon_split_qkv_per_head: bool = False
+    """Whether to orthogonalize each Q, gate, K, and V head independently. Requires
+    ``muon_split_qkv``. Batched execution requires Emerging-Optimizers 0.3.0 or newer;
+    older versions process heads individually. By default, Muon orthogonalizes the Q,
+    gate, K, and V projection matrices separately."""
+
     muon_nesterov: bool = False
     """Whether to use Nesterov-style momentum in the internal SGD."""
 
diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py
index 0f248bcf399..8e5e76cbc11 100644
--- a/megatron/core/tensor_parallel/layers.py
+++ b/megatron/core/tensor_parallel/layers.py
@@ -63,7 +63,12 @@
 _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS = {
     "expert_tp": False,
     "is_qkv": False,
+    "qkv_layout": None,
     "qkv_split_shapes": None,
+    "qkv_split_shapes_global": None,
+    "qkv_gtp_pad_length": 0,
+    "qkv_split_groups_are_complete": False,
+    "qkv_split_heads_are_complete": False,
     "tensor_model_parallel": False,
     "partition_dim": -1,
     "partition_stride": 1,
diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py
index 682b75fb701..750d20c4e4c 100644
--- a/megatron/core/transformer/attention.py
+++ b/megatron/core/transformer/attention.py
@@ -136,6 +136,49 @@
     HAVE_FUSED_QKV_ROPE = False
 
 
+@dataclass(frozen=True)
+class QKVLayout:
+    """Logical row layout for a packed attention projection weight.
+
+    ``projection_split_shapes`` describes the projection slices repeated in every group.
+    ``per_head_split_shapes`` describes the independently orthogonalizable head slices in the
+    same group. Standard fused QKV has one group per query group, while MLA up-projections have
+    one group per attention head.
+    """
+
+    num_groups: int
+    projection_split_shapes: tuple[int, ...]
+    per_head_split_shapes: tuple[int, ...]
+
+    @classmethod
+    def from_transformer_config(cls, config: TransformerConfig) -> 'QKVLayout':
+        """Build the fused QKV row layout described by a transformer config."""
+        assert config.num_query_groups is not None
+        assert config.kv_channels is not None
+        num_query_heads_per_group = config.num_attention_heads // config.num_query_groups
+        projection_split_shapes = [num_query_heads_per_group * config.kv_channels]
+        per_head_split_shapes = [config.kv_channels] * num_query_heads_per_group
+        if config.attention_output_gate:
+            projection_split_shapes.append(num_query_heads_per_group * config.kv_channels)
+            per_head_split_shapes += [config.kv_channels] * num_query_heads_per_group
+        projection_split_shapes += [config.kv_channels, config.kv_channels]
+        per_head_split_shapes += [config.kv_channels, config.kv_channels]
+        return cls(
+            num_groups=config.num_query_groups,
+            projection_split_shapes=tuple(projection_split_shapes),
+            per_head_split_shapes=tuple(per_head_split_shapes),
+        )
+
+    @classmethod
+    def from_splits(cls, num_groups: int, split_shapes: tuple[int, ...]) -> 'QKVLayout':
+        """Build a layout whose projection slices are repeated once per attention head."""
+        return cls(
+            num_groups=num_groups,
+            projection_split_shapes=split_shapes,
+            per_head_split_shapes=split_shapes,
+        )
+
+
 class LinearQkvInterface(Protocol):
     """Interface for linear_qkv modules."""
 
@@ -1693,6 +1736,7 @@ def __init__(
             pg_collection=self.pg_collection,
             name=(name + ".linear_qkv") if name is not None else None,
         )
+        self.linear_qkv.weight.qkv_layout = QKVLayout.from_transformer_config(self.config)
 
         # Resolve which norm class to use for Q and K.
         # Config selects the default norm class; spec overrides if set.
diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py
index e0b6af7aa7f..f776c102b97 100644
--- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py
+++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py
@@ -33,7 +33,7 @@
     gather_from_tensor_model_parallel_region,
     scatter_to_sequence_parallel_region,
 )
-from megatron.core.transformer.attention import Attention
+from megatron.core.transformer.attention import Attention, QKVLayout
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.core.transformer.mla_qk_norm_config import QKNormConfigResolver
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
@@ -372,6 +372,15 @@ def __init__(
             name=(name + ".linear_kv_up_proj") if name is not None else None,
         )
 
+        q_up_proj = self.linear_q_proj if self.config.q_lora_rank is None else self.linear_q_up_proj
+        q_up_proj.weight.qkv_layout = QKVLayout.from_splits(
+            self.config.num_attention_heads,
+            (self.config.qk_head_dim, self.config.qk_pos_emb_head_dim),
+        )
+        self.linear_kv_up_proj.weight.qkv_layout = QKVLayout.from_splits(
+            self.config.num_attention_heads, (self.config.qk_head_dim, self.config.v_head_dim)
+        )
+
         if self.config.q_lora_rank is not None:
             self.q_layernorm = build_module(
                 layer_classes["q_layernorm"],
diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py
index aa21e78ce86..aff59ff7c74 100644
--- a/megatron/core/transformer/multi_latent_attention.py
+++ b/megatron/core/transformer/multi_latent_attention.py
@@ -34,7 +34,7 @@
     gather_from_tensor_model_parallel_region,
     scatter_to_sequence_parallel_region,
 )
-from megatron.core.transformer.attention import Attention, LinearProjBuilder
+from megatron.core.transformer.attention import Attention, LinearProjBuilder, QKVLayout
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.core.transformer.mla_qk_norm_config import QKNormConfigResolver
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
@@ -616,6 +616,15 @@ def __init__(
             name=(name + ".linear_kv_up_proj") if name is not None else None,
         )
 
+        q_up_proj = self.linear_q_proj if self.config.q_lora_rank is None else self.linear_q_up_proj
+        q_up_proj.weight.qkv_layout = QKVLayout.from_splits(
+            self.config.num_attention_heads,
+            (self.config.qk_head_dim, self.config.qk_pos_emb_head_dim),
+        )
+        self.linear_kv_up_proj.weight.qkv_layout = QKVLayout.from_splits(
+            self.config.num_attention_heads, (self.config.qk_head_dim, self.config.v_head_dim)
+        )
+
         if self.config.q_lora_rank is not None:
             self.q_layernorm = layer_classes["q_layernorm"](
                 hidden_size=self.config.q_lora_rank,
@@ -1329,6 +1338,14 @@ def __init__(
             name=(name + ".linear_kv_up_proj") if name is not None else None,
         )
 
+        self.linear_q_up_proj.weight.qkv_layout = QKVLayout.from_splits(
+            self.config.num_attention_heads,
+            (self.config.qk_head_dim, self.config.qk_pos_emb_head_dim),
+        )
+        self.linear_kv_up_proj.weight.qkv_layout = QKVLayout.from_splits(
+            self.config.num_attention_heads, (self.config.qk_head_dim, self.config.v_head_dim)
+        )
+
         self.q_layernorm = layer_classes["q_layernorm"](
             hidden_size=self.config.q_lora_rank,
             config=self.config,
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index d4f9eb9c0de..31d361cc022 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -2517,6 +2517,11 @@ def _add_regularization_args(parser):
     group.add_argument('--muon-no-split-qkv', action='store_false', default=True,
                        dest='muon_split_qkv',
                        help='Whether to split QKV parameters for Muon optimizer')
+    group.add_argument('--muon-split-qkv-per-head', action='store_true',
+                       help='Orthogonalize each Q, gate, K, and V head independently. '
+                       'Batched execution requires emerging-optimizers>=0.3.0; older versions '
+                       'process heads individually. By default, Q, gate, K, and V projections '
+                       'are orthogonalized separately')
     group.add_argument('--muon-nesterov', action='store_true',
                        help='Whether to use Nesterov-style momentum in the internal SGD')
     group.add_argument('--muon-scale-mode', type=str, default='spectral',
diff --git a/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py b/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py
index 961ab070556..e8858bfbcb2 100644
--- a/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py
+++ b/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py
@@ -15,6 +15,8 @@
 2.  TestTPGTPColumnParallelLinear - column-parallel Linear: fwd/bwd correctness (weight shape verified inline)
 3.  TestTPGTPRowParallelLinear    - row-parallel Linear: fwd/bwd smoke test + numerical correctness
 4.  TestTPGTPLayerNormLinear      - LayerNormLinear column-parallel smoke test
+5.  TestTPGTPPaddingAlignment     - alignment padding is applied independently per TP slice
+6.  TestTPGTPMuonQKVPadding       - Muon excludes GTP padding from QKV orthogonalization
 
 Tests use (tp_size, gtp_remat_size) = (2, 2) → world_size = 4 (runs on 4-GPU machines).
 
@@ -27,7 +29,10 @@
 import pytest
 import torch
 import torch.distributed as dist
+import torch.nn.functional as F
 
+from megatron.core.optimizer.emerging_optimizers import HAVE_EMERGING_OPTIMIZERS, TensorParallelMuon
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel.gtp_api import HAVE_GTP
 
 if not HAVE_GTP:
@@ -395,3 +400,104 @@ def test_pre_init_pads_per_tp_slice(self, tp_size, gtp_remat_size):
         world_size = tp_size * gtp_remat_size
         _requires_multi_gpu(world_size)
         _run_distributed(_worker_pre_init_tp_padding, world_size, tp_size, gtp_remat_size)
+
+
+def _worker_muon_qkv_padding(
+    rank, world_size, port, tp_size, gtp_remat_size, split_per_head, split_shapes
+):
+    """Muon must exclude per-TP GTP padding from QKV orthogonalization."""
+    del port
+    tp_group, gtp_remat_group, tp_rank, gtp_rank = _build_groups(
+        rank, world_size, tp_size, gtp_remat_size
+    )
+
+    logical_rows = sum(split_shapes)
+    logical_tp_rows = logical_rows // tp_size
+    pad_length = 3
+    physical_tp_rows = logical_tp_rows + pad_length
+    gtp_local_rows = physical_tp_rows // gtp_remat_size
+    hidden_size = 4
+
+    # Full logical gradient before TP/GTP sharding; physical GTP padding is not included.
+    logical_grad = torch.arange(
+        logical_rows * hidden_size, dtype=torch.float32, device="cuda"
+    ).view(logical_rows, hidden_size)
+    tp_grad = logical_grad[tp_rank * logical_tp_rows : (tp_rank + 1) * logical_tp_rows]
+    # Use a sentinel so the test detects padding entering the orthogonalization.
+    physical_tp_grad = F.pad(tp_grad, (0, 0, 0, pad_length), value=10_000.0)
+    local_grad = physical_tp_grad[
+        gtp_rank * gtp_local_rows : (gtp_rank + 1) * gtp_local_rows
+    ].clone()
+
+    param = torch.nn.Parameter(torch.zeros_like(local_grad))
+    param.partition_dim = 0
+    param.is_qkv = True
+    param.is_gtp_weight_remat = True
+    param.qkv_gtp_pad_length = pad_length
+    # Padded per-head shards intentionally have no rank-local logical split metadata.
+    param.qkv_split_shapes = None if split_per_head else split_shapes
+    param.qkv_split_shapes_global = split_shapes
+    param.qkv_split_heads_are_complete = False
+    param.qkv_split_groups_are_complete = False
+
+    optimizer = TensorParallelMuon(
+        params=[param],
+        split_qkv=True,
+        split_qkv_per_head=split_per_head,
+        is_qkv_fn=lambda p: getattr(p, 'is_qkv', False),
+        qkv_split_shapes=split_shapes,
+        pg_collection=ProcessGroupCollection(tp=tp_group, gtp_remat=gtp_remat_group),
+        tp_mode="blockwise",
+    )
+
+    def center_rows(x, tp_group=None, partition_dim=None):
+        del tp_group, partition_dim
+        return x - x.mean(dim=-2, keepdim=True)
+
+    optimizer.scaled_orthogonalize_fn = center_rows
+    actual = optimizer.orthogonalize(param, local_grad)
+
+    if split_per_head:
+        expected_logical_grad = torch.cat(
+            [center_rows(head) for head in torch.split(logical_grad, split_shapes, dim=0)], dim=0
+        )
+    else:
+        grouped_grad = logical_grad.view(1, logical_rows, hidden_size)
+        projections = torch.split(grouped_grad, split_shapes, dim=1)
+        expected_logical_grad = torch.cat(
+            [
+                center_rows(projection.reshape(-1, hidden_size)).view_as(projection)
+                for projection in projections
+            ],
+            dim=1,
+        ).view_as(logical_grad)
+
+    expected_tp = expected_logical_grad[tp_rank * logical_tp_rows : (tp_rank + 1) * logical_tp_rows]
+    expected_tp = F.pad(expected_tp, (0, 0, 0, pad_length))
+    expected = expected_tp[gtp_rank * gtp_local_rows : (gtp_rank + 1) * gtp_local_rows]
+
+    torch.testing.assert_close(actual, expected)
+    assert actual.shape == local_grad.shape
+    if gtp_rank == gtp_remat_size - 1:
+        assert torch.count_nonzero(actual[-pad_length:]) == 0
+
+
+class TestTPGTPMuonQKVPadding:
+    @pytest.mark.parametrize(
+        "split_per_head,split_shapes", [(True, [2, 2, 2, 2, 2]), (False, [6, 2, 2])]
+    )
+    def test_padding_is_excluded_from_qkv_orthogonalization(self, split_per_head, split_shapes):
+        if not HAVE_EMERGING_OPTIMIZERS:
+            pytest.skip("emerging_optimizers package is not installed")
+        tp_size = 2
+        gtp_remat_size = 2
+        world_size = tp_size * gtp_remat_size
+        _requires_multi_gpu(world_size)
+        _run_distributed(
+            _worker_muon_qkv_padding,
+            world_size,
+            tp_size,
+            gtp_remat_size,
+            split_per_head,
+            split_shapes,
+        )
diff --git a/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py b/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py
index a76746d7674..09a592110d4 100644
--- a/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py
+++ b/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py
@@ -10,6 +10,7 @@
     copy_tensor_model_parallel_attributes,
     param_is_not_tensor_parallel_duplicate,
 )
+from megatron.core.transformer.attention import QKVLayout
 from megatron.core.transformer.transformer_config import TransformerConfig
 from tests.unit_tests.test_utilities import Utils
 
@@ -91,16 +92,30 @@ def test_row_parallel_linear_tp_attrs_no_init(self, use_cpu_init):
         assert hasattr(w, "partition_stride") and w.partition_stride == 1
 
 
-def test_copy_tensor_model_parallel_attributes_preserves_qkv_split_shapes():
+def test_copy_tensor_model_parallel_attributes_preserves_qkv_metadata():
     source = torch.empty(4, 4)
     destination = torch.empty_like(source)
     source.is_qkv = True
-    source.qkv_split_shapes = [256, 64, 64]
+    source.qkv_layout = QKVLayout(
+        num_groups=2,
+        projection_split_shapes=(256, 64, 64),
+        per_head_split_shapes=(64, 64, 64, 64, 64, 64),
+    )
+    source.qkv_split_shapes = [2, 2]
+    source.qkv_split_shapes_global = [2] * 4
+    source.qkv_gtp_pad_length = 3
+    source.qkv_split_groups_are_complete = True
+    source.qkv_split_heads_are_complete = True
 
     copy_tensor_model_parallel_attributes(destination, source)
 
     assert destination.is_qkv is True
+    assert destination.qkv_layout == source.qkv_layout
     assert destination.qkv_split_shapes == source.qkv_split_shapes
+    assert destination.qkv_split_shapes_global == source.qkv_split_shapes_global
+    assert destination.qkv_gtp_pad_length == 3
+    assert destination.qkv_split_groups_are_complete is True
+    assert destination.qkv_split_heads_are_complete is True
 
 
 def test_non_allreduce_param_uses_expert_tp_group_for_duplicate_filter():
diff --git a/tests/unit_tests/test_emerging_optimizers.py b/tests/unit_tests/test_emerging_optimizers.py
index e3b9f666fb2..e9688528f89 100644
--- a/tests/unit_tests/test_emerging_optimizers.py
+++ b/tests/unit_tests/test_emerging_optimizers.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import json
 import os
 
 import pytest
@@ -10,18 +11,30 @@
 
 from megatron.core import parallel_state
 from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
+from megatron.core.models.gpt.gpt_model import GPTModel
+from megatron.core.models.gpt.heterogeneous.heterogeneous_layer_specs import (
+    get_gpt_heterogeneous_layer_spec,
+)
 from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer
 from megatron.core.optimizer.emerging_optimizers import (
     HAVE_EMERGING_OPTIMIZERS,
     TensorParallelAdaptiveMuon,
     TensorParallelMuon,
     _get_qkv_split_shapes,
+    _localize_qkv_split_shapes,
+    _qkv_split_groups_are_complete,
     get_supported_coefficient_types,
     validate_coefficient_type,
 )
 from megatron.core.optimizer.muon import get_megatron_muon_optimizer
 from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.transformer import TransformerConfig
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.core.transformer import MLATransformerConfig, TransformerConfig
+from megatron.core.transformer.attention import QKVLayout
+from megatron.core.transformer.heterogeneous.heterogeneous_config import (
+    HeterogeneousTransformerConfig,
+)
 from tests.unit_tests.test_utilities import Utils
 
 if HAVE_EMERGING_OPTIMIZERS:
@@ -80,6 +93,44 @@ def test_muon_qkv_split_shapes():
 
     assert _get_qkv_split_shapes(config) == [128, 64, 64]
     assert _get_qkv_split_shapes(gated_config) == [128, 128, 64, 64]
+    assert _get_qkv_split_shapes(config, split_qkv_per_head=True) == [64] * 32
+    assert _get_qkv_split_shapes(gated_config, split_qkv_per_head=True) == [64] * 48
+
+    mla_layout = QKVLayout.from_splits(4, (128, 64))
+    assert _get_qkv_split_shapes(mla_layout) == [128, 64]
+    assert _get_qkv_split_shapes(mla_layout, split_qkv_per_head=True) == [128, 64] * 4
+
+
+def test_muon_local_qkv_head_split_shapes_can_differ_by_tp_rank():
+    """Rank-local per-head layouts report complete and fragmented heads."""
+    global_split_shapes = [64] * 20
+
+    rank_0_shapes, rank_0_complete = _localize_qkv_split_shapes(
+        global_split_shapes, local_start=0, local_rows=160
+    )
+    rank_1_shapes, rank_1_complete = _localize_qkv_split_shapes(
+        global_split_shapes, local_start=160, local_rows=160
+    )
+    aligned_shapes, aligned_complete = _localize_qkv_split_shapes(
+        global_split_shapes, local_start=0, local_rows=640
+    )
+
+    assert rank_0_shapes == [64, 64, 32]
+    assert rank_1_shapes == [32, 64, 64]
+    assert not rank_0_complete
+    assert not rank_1_complete
+    assert aligned_shapes == [64] * 10
+    assert aligned_complete
+
+
+def test_muon_qkv_query_group_layout_localization():
+    """Projection splitting detects query groups fragmented by TP row ranges."""
+    split_shapes = [256, 64, 64]
+
+    assert _qkv_split_groups_are_complete(split_shapes, local_start=0, local_rows=384)
+    assert _qkv_split_groups_are_complete(split_shapes, local_start=384, local_rows=768)
+    assert not _qkv_split_groups_are_complete(split_shapes, local_start=0, local_rows=192)
+    assert not _qkv_split_groups_are_complete(split_shapes, local_start=192, local_rows=192)
 
 
 def test_muon_optimizer_smoke():
@@ -475,6 +526,243 @@ def create_tp_model_and_optimizer(self, mode):
 
         return model, optimizer
 
+    def test_optimizer_factory_tags_qkv_when_tp_exceeds_query_groups(self):
+        """The public optimizer factory tags query groups fragmented across TP ranks."""
+        pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+        tp_size = pg_collection.tp.size()
+        assert tp_size == 2
+        model_parallel_cuda_manual_seed(123)
+        transformer_config = TransformerConfig(
+            num_layers=1,
+            hidden_size=8,
+            num_attention_heads=2,
+            num_query_groups=1,
+            kv_channels=4,
+            tensor_model_parallel_size=tp_size,
+            use_cpu_initialization=False,
+            add_bias_linear=False,
+        )
+        model = GPTModel(
+            config=transformer_config,
+            transformer_layer_spec=get_gpt_layer_local_spec(),
+            vocab_size=32,
+            max_sequence_length=8,
+            pre_process=False,
+            post_process=False,
+            pg_collection=pg_collection,
+        )
+        optimizer_config = OptimizerConfig(
+            optimizer='muon',
+            lr=0.01,
+            use_distributed_optimizer=False,
+            muon_split_qkv=True,
+            muon_tp_mode="blockwise",
+        )
+
+        optimizer = get_megatron_optimizer(
+            config=optimizer_config,
+            model_chunks=[model],
+            use_gloo_process_groups=False,
+            pg_collection=pg_collection,
+        )
+
+        qkv_weight = model.decoder.layers[0].self_attention.linear_qkv.weight
+        assert optimizer is not None
+        assert qkv_weight.shape[0] == 8
+        assert qkv_weight.is_qkv
+        assert qkv_weight.qkv_split_shapes == [8, 4, 4]
+        assert qkv_weight.qkv_split_shapes_global == [8, 4, 4]
+        assert not qkv_weight.qkv_split_groups_are_complete
+
+    @pytest.mark.parametrize("split_per_head", [False, True])
+    def test_optimizer_factory_uses_heterogeneous_layer_qkv_layout(self, split_per_head):
+        """Each heterogeneous attention layer supplies its own logical QKV layout."""
+        pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+        tp_size = pg_collection.tp.size()
+        assert tp_size == 2
+        model_parallel_cuda_manual_seed(123)
+        block_configs = {
+            "block_configs": [
+                {
+                    "attention": {
+                        "no_op": False,
+                        "replace_with_linear": False,
+                        "num_query_groups": 2,
+                    },
+                    "mlp": {"no_op": False, "replace_with_linear": False, "ffn_hidden_size": 16},
+                },
+                {
+                    "attention": {
+                        "no_op": False,
+                        "replace_with_linear": False,
+                        "num_query_groups": 1,
+                    },
+                    "mlp": {"no_op": False, "replace_with_linear": False, "ffn_hidden_size": 16},
+                },
+            ]
+        }
+        transformer_config = HeterogeneousTransformerConfig(
+            num_layers=2,
+            hidden_size=8,
+            num_attention_heads=2,
+            kv_channels=4,
+            tensor_model_parallel_size=tp_size,
+            use_cpu_initialization=False,
+            add_bias_linear=False,
+            heterogeneous_layers_config_encoded_json=json.dumps(block_configs),
+        )
+        model = GPTModel(
+            config=transformer_config,
+            transformer_layer_spec=get_gpt_heterogeneous_layer_spec(transformer_config),
+            vocab_size=32,
+            max_sequence_length=8,
+            pre_process=False,
+            post_process=False,
+            pg_collection=pg_collection,
+        )
+        optimizer_config = OptimizerConfig(
+            optimizer='muon',
+            lr=0.01,
+            use_distributed_optimizer=False,
+            muon_split_qkv=True,
+            muon_split_qkv_per_head=split_per_head,
+            muon_tp_mode="blockwise",
+        )
+
+        optimizer = get_megatron_optimizer(
+            config=optimizer_config,
+            model_chunks=[model],
+            use_gloo_process_groups=False,
+            pg_collection=pg_collection,
+        )
+
+        first_qkv = model.decoder.layers[0].self_attention.linear_qkv.weight
+        second_qkv = model.decoder.layers[1].self_attention.linear_qkv.weight
+        assert optimizer is not None
+        assert transformer_config.num_query_groups == 2
+        assert first_qkv.qkv_layout.num_groups == 2
+        assert second_qkv.qkv_layout.num_groups == 1
+        assert first_qkv.shape[0] == 12
+        assert second_qkv.shape[0] == 8
+        assert first_qkv.is_qkv
+        assert second_qkv.is_qkv
+        if split_per_head:
+            assert first_qkv.qkv_split_shapes_global == [4] * 6
+            assert second_qkv.qkv_split_shapes_global == [4] * 4
+            assert first_qkv.qkv_split_heads_are_complete
+            assert second_qkv.qkv_split_heads_are_complete
+        else:
+            assert first_qkv.qkv_split_shapes_global == [4, 4, 4] * 2
+            assert second_qkv.qkv_split_shapes_global == [8, 4, 4]
+            assert first_qkv.qkv_split_groups_are_complete
+            assert not second_qkv.qkv_split_groups_are_complete
+
+    @pytest.mark.parametrize("split_per_head", [False, True])
+    @pytest.mark.parametrize("q_lora_rank", [None, 4])
+    def test_optimizer_factory_uses_mla_up_projection_layouts(self, split_per_head, q_lora_rank):
+        """MLA up-projections use module-owned layouts with TP-aware Muon splitting."""
+        pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+        tp_size = pg_collection.tp.size()
+        assert tp_size == 2
+        model_parallel_cuda_manual_seed(123)
+        transformer_config = MLATransformerConfig(
+            num_layers=1,
+            hidden_size=8,
+            num_attention_heads=2,
+            q_lora_rank=q_lora_rank,
+            kv_lora_rank=4,
+            qk_head_dim=4,
+            qk_pos_emb_head_dim=2,
+            v_head_dim=3,
+            tensor_model_parallel_size=tp_size,
+            use_cpu_initialization=False,
+            add_bias_linear=False,
+            multi_latent_attention=True,
+            rope_type="rope",
+            rotary_base=10000,
+            original_max_position_embeddings=8,
+        )
+        model = GPTModel(
+            config=transformer_config,
+            transformer_layer_spec=get_gpt_layer_local_spec(multi_latent_attention=True),
+            vocab_size=32,
+            max_sequence_length=8,
+            pre_process=False,
+            post_process=False,
+            pg_collection=pg_collection,
+        )
+        optimizer_config = OptimizerConfig(
+            optimizer='muon',
+            lr=0.01,
+            use_distributed_optimizer=False,
+            muon_split_qkv=True,
+            muon_split_qkv_per_head=split_per_head,
+            muon_tp_mode="blockwise",
+        )
+
+        optimizer = get_megatron_optimizer(
+            config=optimizer_config,
+            model_chunks=[model],
+            use_gloo_process_groups=False,
+            pg_collection=pg_collection,
+        )
+
+        attention = model.decoder.layers[0].self_attention
+        q_up_weight = (
+            attention.linear_q_proj.weight
+            if q_lora_rank is None
+            else attention.linear_q_up_proj.weight
+        )
+        kv_up_weight = attention.linear_kv_up_proj.weight
+        assert optimizer is not None
+        assert q_up_weight.is_qkv
+        assert kv_up_weight.is_qkv
+        assert q_up_weight.qkv_split_shapes_global == [4, 2] * 2
+        assert kv_up_weight.qkv_split_shapes_global == [4, 3] * 2
+        if split_per_head:
+            assert q_up_weight.qkv_split_heads_are_complete
+            assert kv_up_weight.qkv_split_heads_are_complete
+        else:
+            assert q_up_weight.qkv_split_groups_are_complete
+            assert kv_up_weight.qkv_split_groups_are_complete
+        if q_lora_rank is not None:
+            assert not getattr(attention.linear_q_down_proj.weight, 'is_qkv', False)
+        assert not getattr(attention.linear_kv_down_proj.weight, 'is_qkv', False)
+
+    def test_optimizer_factory_skips_mismatched_qkv_layout(self):
+        """A QKV layout mismatch falls back to whole-matrix Muon orthogonalization."""
+        pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+        tp_size = pg_collection.tp.size()
+        transformer_config = TransformerConfig(
+            num_layers=1,
+            hidden_size=8,
+            num_attention_heads=2,
+            num_query_groups=1,
+            kv_channels=4,
+            tensor_model_parallel_size=tp_size,
+        )
+        model = torch.nn.Module()
+        model.config = transformer_config
+        model.linear_qkv = torch.nn.Linear(8, 7, bias=False, dtype=torch.float32, device='cuda')
+        model.linear_qkv.weight.tensor_model_parallel = True
+        model.linear_qkv.weight.partition_dim = 0
+        optimizer_config = OptimizerConfig(
+            optimizer='muon', lr=0.01, use_distributed_optimizer=False, muon_split_qkv=True
+        )
+
+        optimizer = get_megatron_optimizer(
+            config=optimizer_config,
+            model_chunks=[model],
+            use_gloo_process_groups=False,
+            pg_collection=pg_collection,
+        )
+
+        qkv_weight = model.linear_qkv.weight
+        assert optimizer is not None
+        assert not qkv_weight.is_qkv
+        assert qkv_weight.qkv_split_shapes is None
+        assert qkv_weight.qkv_split_shapes_global is None
+
     @pytest.mark.parametrize("mode", ["duplicated", "distributed"])
     def test_muon_optimizer_modes_multirank_same_result(self, mode):
         """Test that duplicated and distributed modes produce same results with TP > 1."""
@@ -516,6 +804,99 @@ def test_muon_optimizer_blockwise_mode_different_result(self):
             model.weight.data, original_weight
         ), "Weight should be updated with mode=blockwise"
 
+    def test_muon_optimizer_per_head_split_gathers_fragmented_heads(self, monkeypatch):
+        """Per-head splitting reconstructs heads that cross TP rank boundaries."""
+        monkeypatch.setattr(
+            "megatron.core.optimizer.emerging_optimizers.EMERGING_OPTIMIZERS_VERSION",
+            Version("0.3.0"),
+        )
+        pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+        tp_group = pg_collection.tp
+        tp_rank = tp_group.rank()
+        local_grad = torch.arange(12, dtype=torch.float32, device='cuda').view(3, 4)
+        local_grad = local_grad + tp_rank * local_grad.numel()
+        param = torch.nn.Parameter(torch.zeros_like(local_grad))
+        param.partition_dim = 0
+        param.is_qkv = True
+        param.qkv_split_shapes = [2, 2]
+        param.qkv_split_shapes_global = [2, 2, 2]
+        param.qkv_split_heads_are_complete = False
+
+        optimizer = TensorParallelMuon(
+            params=[param],
+            split_qkv=True,
+            split_qkv_per_head=True,
+            is_qkv_fn=lambda p: getattr(p, 'is_qkv', False),
+            qkv_split_shapes=[2, 2, 2],
+            pg_collection=pg_collection,
+            tp_mode="blockwise",
+        )
+
+        def center_rows(x, tp_group=None, partition_dim=None):
+            del tp_group, partition_dim
+            return x - x.mean(dim=-2, keepdim=True)
+
+        optimizer.scaled_orthogonalize_fn = center_rows
+        actual = optimizer.orthogonalize(param, local_grad)
+
+        shards = [torch.empty_like(local_grad) for _ in range(tp_group.size())]
+        torch.distributed.all_gather(shards, local_grad, tp_group)
+        global_grad = torch.cat(shards, dim=0)
+        expected_global = torch.cat(
+            [center_rows(head) for head in torch.split(global_grad, [2, 2, 2], dim=0)], dim=0
+        )
+        expected = expected_global[tp_rank * 3 : (tp_rank + 1) * 3]
+        torch.testing.assert_close(actual, expected)
+
+    @pytest.mark.parametrize("split_shapes", ([4, 2, 2], [4, 4, 2, 2], [3, 5]))
+    def test_muon_optimizer_projection_split_gathers_fragmented_query_group(self, split_shapes):
+        """Projection splitting reconstructs a query group split over TP ranks."""
+        pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+        tp_group = pg_collection.tp
+        tp_rank = tp_group.rank()
+        global_rows = sum(split_shapes)
+        assert global_rows % tp_group.size() == 0
+        local_rows = global_rows // tp_group.size()
+        local_grad = torch.arange(local_rows * 4, dtype=torch.float32, device='cuda').view(
+            local_rows, 4
+        )
+        local_grad = local_grad + tp_rank * local_grad.numel()
+        param = torch.nn.Parameter(torch.zeros_like(local_grad))
+        param.partition_dim = 0
+        param.is_qkv = True
+        param.qkv_split_shapes = split_shapes
+        param.qkv_split_shapes_global = split_shapes
+        param.qkv_split_groups_are_complete = False
+
+        optimizer = TensorParallelMuon(
+            params=[param],
+            split_qkv=True,
+            is_qkv_fn=lambda p: getattr(p, 'is_qkv', False),
+            qkv_split_shapes=split_shapes,
+            pg_collection=pg_collection,
+            tp_mode="blockwise",
+        )
+
+        def center_rows(x, tp_group=None, partition_dim=None):
+            del tp_group, partition_dim
+            return x - x.mean(dim=-2, keepdim=True)
+
+        optimizer.scaled_orthogonalize_fn = center_rows
+        actual = optimizer.orthogonalize(param, local_grad)
+
+        shards = [torch.empty_like(local_grad) for _ in range(tp_group.size())]
+        torch.distributed.all_gather(shards, local_grad, tp_group)
+        global_grad = torch.cat(shards, dim=0)
+        expected_global = torch.cat(
+            [
+                center_rows(projection)
+                for projection in torch.split(global_grad, split_shapes, dim=0)
+            ],
+            dim=0,
+        )
+        expected = expected_global[tp_rank * local_rows : (tp_rank + 1) * local_rows]
+        torch.testing.assert_close(actual, expected)
+
 
 # All non-custom coefficient types supported by emerging_optimizers.
 _TESTABLE_COEFFICIENT_TYPES = (
@@ -720,6 +1101,125 @@ def test_muon_optimizer_qkv_split():
     ), "Weights should be different between split_qkv=True and split_qkv=False"
 
 
+def test_muon_optimizer_qkv_split_per_head_is_opt_in(monkeypatch):
+    """Per-head splitting is guarded and differs from projection splitting."""
+    monkeypatch.setattr(
+        "megatron.core.optimizer.emerging_optimizers.EMERGING_OPTIMIZERS_VERSION", Version("0.3.0")
+    )
+    grad = torch.arange(48, dtype=torch.float32, device='cuda').view(16, 3)
+    projection_param = torch.nn.Parameter(torch.zeros_like(grad))
+    projection_param.is_qkv = True
+    projection_param.qkv_split_shapes = [4, 2, 2]
+    head_param = torch.nn.Parameter(torch.zeros_like(grad))
+    head_param.is_qkv = True
+    head_param.qkv_split_shapes = [2] * 8
+    orthogonalize_call_shapes = []
+
+    def center_rows(x, tp_group=None, partition_dim=None):
+        del tp_group, partition_dim
+        orthogonalize_call_shapes.append(tuple(x.shape))
+        return x - x.mean(dim=-2, keepdim=True)
+
+    projection_optimizer = TensorParallelMuon(
+        params=[projection_param],
+        split_qkv=True,
+        is_qkv_fn=lambda p: getattr(p, 'is_qkv', False),
+        qkv_split_shapes=[4, 2, 2],
+        pg_collection=None,
+    )
+    projection_optimizer.scaled_orthogonalize_fn = center_rows
+    projection_out = projection_optimizer.orthogonalize(projection_param, grad)
+    orthogonalize_call_shapes.clear()
+
+    head_optimizer = TensorParallelMuon(
+        params=[head_param],
+        split_qkv=True,
+        split_qkv_per_head=True,
+        is_qkv_fn=lambda p: getattr(p, 'is_qkv', False),
+        qkv_split_shapes=[2] * 8,
+        pg_collection=None,
+    )
+    head_optimizer.scaled_orthogonalize_fn = center_rows
+    head_out = head_optimizer.orthogonalize(head_param, grad)
+    assert orthogonalize_call_shapes == [(8, 2, 3)]
+
+    expected_head_out = torch.cat(
+        [center_rows(head) for head in torch.split(grad, [2] * 8, dim=0)], dim=0
+    )
+    torch.testing.assert_close(head_out, expected_head_out)
+    assert not torch.equal(projection_out, head_out)
+
+
+def test_muon_optimizer_qkv_split_per_head_requires_split_qkv():
+    """The per-head switch cannot enable QKV splitting by itself."""
+    param = torch.nn.Parameter(torch.zeros(4, 4, dtype=torch.float32, device='cuda'))
+    with pytest.raises(ValueError, match="split_qkv_per_head requires split_qkv=True"):
+        TensorParallelMuon(
+            params=[param], split_qkv=False, split_qkv_per_head=True, pg_collection=None
+        )
+
+
+def test_muon_optimizer_uniform_per_head_splits_fall_back_without_batched_ns(monkeypatch):
+    """Uniform per-head splits use individual 2D calls before Emerging-Optimizers 0.3."""
+    monkeypatch.setattr(
+        "megatron.core.optimizer.emerging_optimizers.EMERGING_OPTIMIZERS_VERSION", Version("0.2.0")
+    )
+    grad = torch.arange(16, dtype=torch.float32, device='cuda').view(4, 4)
+    param = torch.nn.Parameter(torch.zeros_like(grad))
+    param.is_qkv = True
+    param.qkv_split_shapes = [2, 2]
+    optimizer = TensorParallelMuon(
+        params=[param],
+        split_qkv=True,
+        split_qkv_per_head=True,
+        is_qkv_fn=lambda p: getattr(p, 'is_qkv', False),
+        qkv_split_shapes=[2, 2],
+        pg_collection=None,
+    )
+    call_shapes = []
+
+    def center_rows(x, tp_group=None, partition_dim=None):
+        del tp_group, partition_dim
+        call_shapes.append(tuple(x.shape))
+        return x - x.mean(dim=-2, keepdim=True)
+
+    optimizer.scaled_orthogonalize_fn = center_rows
+    actual = optimizer.orthogonalize(param, grad)
+    assert call_shapes == [(2, 4), (2, 4)]
+    expected = torch.cat(
+        [head - head.mean(dim=-2, keepdim=True) for head in torch.split(grad, [2, 2])]
+    )
+    torch.testing.assert_close(actual, expected)
+
+
+def test_muon_optimizer_nonuniform_per_head_splits_do_not_require_batched_ns(monkeypatch):
+    """Nonuniform per-head splits keep using the unbatched compatibility path."""
+    monkeypatch.setattr(
+        "megatron.core.optimizer.emerging_optimizers.EMERGING_OPTIMIZERS_VERSION", Version("0.2.0")
+    )
+    grad = torch.arange(12, dtype=torch.float32, device='cuda').view(3, 4)
+    param = torch.nn.Parameter(torch.zeros_like(grad))
+    param.is_qkv = True
+    param.qkv_split_shapes = [2, 1]
+    optimizer = TensorParallelMuon(
+        params=[param],
+        split_qkv=True,
+        split_qkv_per_head=True,
+        is_qkv_fn=lambda p: getattr(p, 'is_qkv', False),
+        qkv_split_shapes=[2, 1],
+        pg_collection=None,
+    )
+
+    def center_rows(x, tp_group=None, partition_dim=None):
+        del tp_group, partition_dim
+        return x - x.mean(dim=-2, keepdim=True)
+
+    optimizer.scaled_orthogonalize_fn = center_rows
+    actual = optimizer.orthogonalize(param, grad)
+    expected = torch.cat([center_rows(head) for head in torch.split(grad, [2, 1])])
+    torch.testing.assert_close(actual, expected)
+
+
 def test_muon_optimizer_extra_scale_factor():
     """Test TensorParallelMuon optimizer with different extra_scale_factor values."""
     model = torch.nn.Linear(80, 40, bias=False, dtype=torch.float32, device='cuda')
diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py b/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py
index fc1778f649f..7f2c2afa16f 100644
--- a/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py
+++ b/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py
@@ -430,6 +430,11 @@ def test_functionality(tp_cp: List[int], qkv_format: str, down_proj_use_column_p
         cp_comm_type="all_gather" if cp_size > 1 else None,
         pg_collection=None,
     ).cuda()
+    assert absorbed_mla.linear_q_up_proj.weight.qkv_layout.num_groups == config.num_attention_heads
+    assert absorbed_mla.linear_kv_up_proj.weight.qkv_layout.projection_split_shapes == (
+        config.qk_head_dim,
+        config.v_head_dim,
+    )
 
     state_dict = standard_mla.state_dict()
     absorbed_mla.load_state_dict(state_dict)
diff --git a/tests/unit_tests/transformer/test_multi_latent_attention.py b/tests/unit_tests/transformer/test_multi_latent_attention.py
index 95517f7176d..5cc21bbfd26 100644
--- a/tests/unit_tests/transformer/test_multi_latent_attention.py
+++ b/tests/unit_tests/transformer/test_multi_latent_attention.py
@@ -219,6 +219,16 @@ def test_dynamic_inference_forwards_decode_only_to_flash_attention(self, is_deco
     def test_constructor(self):
         assert isinstance(self.parallel_attention, MLASelfAttention)
         assert self.parallel_attention.layer_number == 1
+        assert self.parallel_attention.linear_q_up_proj.weight.qkv_layout.num_groups == 4
+        assert (
+            self.parallel_attention.linear_q_up_proj.weight.qkv_layout.projection_split_shapes
+            == (128, 64)
+        )
+        assert self.parallel_attention.linear_kv_up_proj.weight.qkv_layout.num_groups == 4
+        assert (
+            self.parallel_attention.linear_kv_up_proj.weight.qkv_layout.projection_split_shapes
+            == (128, 128)
+        )
 
         num_weights = sum([p.numel() for p in self.parallel_attention.parameters()])
         assert num_weights == 65036
@@ -1700,6 +1710,9 @@ def test_constructor(self):
         assert isinstance(self.fused_attention, MLASelfAttention)
         assert self.fused_attention.layer_number == 1
         assert hasattr(self.fused_attention, 'linear_qkv_down_proj')
+        assert self.fused_attention.linear_q_up_proj.weight.qkv_layout.num_groups == 4
+        assert self.fused_attention.linear_kv_up_proj.weight.qkv_layout.num_groups == 4
+        assert getattr(self.fused_attention.linear_qkv_down_proj.weight, 'qkv_layout', None) is None
 
     def test_fused_weight_shape(self):
         config = self.transformer_config