qk_clip: value-only fix — keep original structure, drop per-row rewrite [skip-build]

Revert the per-row expand_row_scales restructure. qk_clip keeps its exact
original reshape + named-slice shape; only the multiplied values change:
compute_scales now returns raw per-head γ, and each sub-region is raised to
its exponent inline (q_nope γ^(1-r), q_pe γ, k_nope γ^r, MHA √γ). The q_pe
line is byte-identical to before — it is now correct because scales is raw γ.

muon.py sharded path stays per-head (no per-row vectors); it just applies the
content exponent to match the raw-γ contract, preserving its prior behavior.
The sharded path's pre-existing MLA sub-region limitation (q_pe/v scaled by the
content exponent) is left as-is and documented, not addressed here.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

torch-ext/optimizer/muon.py

@@ -20,8 +20,7 @@ from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO,
                             zeropower_via_newtonschulz5,
                             zeropower_via_newtonschulz5_batched)
 from .pipeline import muon_chunk_pipeline, prelaunch_first_gather
-from .qk_clip import (compute_scales, expand_row_scales, get_qk_clip_info,
-                      qk_clip)
+from .qk_clip import compute_scales, get_qk_clip_info, qk_clip
 
 logger = logging.getLogger(__name__)
 
@@ -480,20 +479,26 @@ class Muon(torch.optim.Optimizer):
                             p, qk_clip_state, n
                         ) if qk_clip_state is not None else None
                         if scales_full is not None:
-                            # Per-row γ**exponent over the whole param, indexed
-                            # by this rank's global row slice. Unlike the old
-                            # per-head form, this scales the MLA sub-regions
-                            # correctly (q_pe → full γ, v rows → untouched).
-                            full_row_scales = expand_row_scales(
-                                scales_full, qk_clip_state)
+                            # scales_full is raw per-head γ; apply the content
+                            # exponent per head (MHA √γ, wq_b γ^(1-r), wkv_b
+                            # γ^r). NOTE: this per-head form cannot resolve the
+                            # MLA rotary/value sub-regions (q_pe → full γ, v
+                            # untouched) — see qk_clip; the sharded path scales
+                            # q_pe/v by the content exponent instead.
+                            if not qk_clip_state.is_mla:
+                                e = 0.5
+                            elif qk_clip_state.kind == 'wq_b':
+                                e = 1.0 - qk_clip_state.k_ratio
+                            else:
+                                e = qk_clip_state.k_ratio
+                            ratio = p.shape[0] // scales_full.shape[0]
                             idx0 = local_indices[0]
                             if isinstance(idx0, slice):
-                                idx0 = torch.arange(
-                                    idx0.start or 0,
-                                    idx0.stop,
-                                    device=full_row_scales.device)
-                            row_scales = full_row_scales[idx0.to(
-                                full_row_scales.device)]
+                                start = idx0.start or 0
+                                idx0 = torch.arange(start,
+                                                    idx0.stop,
+                                                    device=scales_full.device)
+                            row_scales = (scales_full**e)[idx0 // ratio]
                             p._local_tensor.mul_(row_scales.view(-1, 1))
 
                 if not batch_params:
@@ -565,20 +570,26 @@ class Muon(torch.optim.Optimizer):
                             p, qk_clip_state, n
                         ) if qk_clip_state is not None else None
                         if scales_full is not None:
-                            # Per-row γ**exponent over the whole param, indexed
-                            # by this rank's global row slice. Unlike the old
-                            # per-head form, this scales the MLA sub-regions
-                            # correctly (q_pe → full γ, v rows → untouched).
-                            full_row_scales = expand_row_scales(
-                                scales_full, qk_clip_state)
+                            # scales_full is raw per-head γ; apply the content
+                            # exponent per head (MHA √γ, wq_b γ^(1-r), wkv_b
+                            # γ^r). NOTE: this per-head form cannot resolve the
+                            # MLA rotary/value sub-regions (q_pe → full γ, v
+                            # untouched) — see qk_clip; the sharded path scales
+                            # q_pe/v by the content exponent instead.
+                            if not qk_clip_state.is_mla:
+                                e = 0.5
+                            elif qk_clip_state.kind == 'wq_b':
+                                e = 1.0 - qk_clip_state.k_ratio
+                            else:
+                                e = qk_clip_state.k_ratio
+                            ratio = p.shape[0] // scales_full.shape[0]
                             idx0 = local_indices[0]
                             if isinstance(idx0, slice):
-                                idx0 = torch.arange(
-                                    idx0.start or 0,
-                                    idx0.stop,
-                                    device=full_row_scales.device)
-                            row_scales = full_row_scales[idx0.to(
-                                full_row_scales.device)]
+                                start = idx0.start or 0
+                                idx0 = torch.arange(start,
+                                                    idx0.stop,
+                                                    device=scales_full.device)
+                            row_scales = (scales_full**e)[idx0 // ratio]
                             p._local_tensor.mul_(row_scales.view(-1, 1))
 
     def _setup_parallel(self, names, params, group, qk_logits):

torch-ext/optimizer/qk_clip.py

@@ -130,9 +130,9 @@ def compute_scales(p, qk_clip_state, name: str = ''):
 
     Returns a per-head tensor of γ (1.0 for heads under the threshold) if any
     head requires clipping, else None. The per-region exponent on γ is NOT
-    applied here — it is applied at multiply time by ``expand_row_scales`` so
-    that a single per-head γ can drive different sub-regions (q_nope, q_pe,
-    k_nope, v) with different exponents. See ``expand_row_scales`` / ``qk_clip``.
+    applied here — it is applied at multiply time (``qk_clip`` for the whole
+    tensor, or the sharded path in muon.py) so that one per-head γ can drive
+    different sub-regions (q_nope, q_pe, k_nope, v) with different exponents.
 
     wkv_b with k_ratio=0 returns None: the whole reduction is placed on Q and
     the key projection is left unclipped.
@@ -178,37 +178,9 @@ def compute_scales(p, qk_clip_state, name: str = ''):
     return scales_full
 
 
-def expand_row_scales(gamma, info):
-    """Per-head raw γ ([H]) → per-row multiplier ([H*stride]) = γ ** exponent,
-    where the exponent depends on the row's sub-region within its head.
-
-    Each logit component ends up reduced by exactly γ:
-      MHA        : √γ on every row                       (√γ · √γ = γ).
-      MLA q_nope : γ^(1-k_ratio) ┐ content: γ^(1-k_ratio)·γ^k_ratio = γ.
-          k_nope : γ^k_ratio     ┘
-      MLA q_pe   : γ  (exponent 1) — the decoupled rotary key is shared and
-                   left unchanged, so the rotary query carries the whole γ for
-                   every k_ratio; γ^(1-k_ratio) here would under-clip rotary.
-      MLA v rows : γ^0 = 1, i.e. untouched.
-
-    Shared by both apply paths: qk_clip (whole tensor) and the sharded DTensor
-    path in muon.py (which indexes the returned vector by global row).
-    """
-    r, n = info.k_ratio, info.qk_nope_head_dim
-    if not info.is_mla:
-        exps = torch.full((info.head_dim, ), 0.5)
-    elif info.kind == 'wq_b':
-        exps = torch.cat(
-            [torch.full((n, ), 1.0 - r),
-             torch.ones(info.qk_rope_head_dim)])
-    else:  # wkv_b: k_nope → γ^r, v rows → γ^0 = 1
-        exps = torch.cat([torch.full((n, ), r), torch.zeros(info.v_head_dim)])
-    exps = exps.to(gamma.device)  # [stride]
-    return (gamma[:, None]**exps[None, :]).reshape(-1)  # [H*stride], head-major
-
-
 def qk_clip(p, scales, info):
-    """Apply QK-clip to a Q/K projection weight, scaling each row by γ**exponent.
+    """Apply QK-clip to a Q/K projection weight; each sub-region is raised to
+    its own exponent so every logit component is reduced by exactly γ.
 
     Args:
         p: Parameter (nn.Parameter or raw tensor).
@@ -216,10 +188,31 @@ def qk_clip(p, scales, info):
             unclipped heads, as returned by ``compute_scales``.
         info: QKClipInfo with kind, head_dim, k_ratio and MLA sub-head dims.
 
-    The per-region exponent (√γ for MHA; γ^(1-k_ratio)/γ^k_ratio content split;
-    full γ on the rotary query; value rows untouched) is applied by
-    ``expand_row_scales``.
+    Exponents:
+        MHA        : √γ on all rows.
+        MLA q_nope : γ^(1-k_ratio),  k_nope : γ^k_ratio   (content product = γ).
+        MLA q_pe   : γ (full) — the decoupled rotary key is shared and left
+                     unchanged, so the rotary query carries the whole γ for
+                     every k_ratio; γ^(1-k_ratio) here would under-clip rotary.
+        MLA v rows : untouched.
     """
     W = p.data if isinstance(p, torch.nn.Parameter) else p
-    row_scales = expand_row_scales(scales, info)  # [num_rows], fp32
-    W.mul_(row_scales.view(-1, 1))  # in-place downcast to W.dtype
+    r = info.k_ratio
+
+    if not info.is_mla:
+        W.view(-1, info.head_dim, W.shape[1]).mul_((scales**0.5).view(-1, 1, 1))
+        return
+
+    if info.kind == 'wq_b':
+        qk_nope = info.qk_nope_head_dim
+        qk_head_dim = qk_nope + info.qk_rope_head_dim
+        W_3d = W.view(-1, qk_head_dim, W.shape[1])  # [H, qk_head_dim, in_dim]
+        W_3d[:, :qk_nope, :].mul_((scales**(1.0 - r)).view(-1, 1, 1))  # q_nope → γ^(1-r)
+        W_3d[:, qk_nope:, :].mul_(scales.view(-1, 1, 1))  # q_pe → γ (full)
+
+    elif info.kind == 'wkv_b':
+        qk_nope = info.qk_nope_head_dim
+        kv_stride = qk_nope + info.v_head_dim
+        W_3d = W.view(-1, kv_stride, W.shape[1])  # [H, kv_stride, in_dim]
+        W_3d[:, :qk_nope, :].mul_((scales**r).view(-1, 1, 1))  # k_nope → γ^r
+        # v rows: untouched (shared rotary key k_rope is unchanged)