fix(qk_clip): scale q_pe rows by full γ, not γ^(1-k_ratio) [skip-build]

MLA logit = q_nope·k_nope (content) + q_pe·k_rope (rotary). The decoupled
rotary key k_rope is shared across heads and left unchanged, so the rotary
term is reduced by the query side alone. Since PR #32 introduced k_ratio,
wq_b scaled q_pe by γ^(1-k_ratio) uniformly with q_nope, making the rotary
term shrink by γ^(1-k_ratio) instead of γ — an under-clip for any k_ratio > 0.

Give q_pe the full γ (exponent 1) regardless of k_ratio; the k_ratio split
only applies to the content (nope) rows, whose K side is scalable. k_ratio=0
is unaffected (Q already got full γ there).

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

torch-ext/optimizer/qk_clip.py

@@ -133,11 +133,15 @@ def compute_scales(p, qk_clip_state, name: str = ''):
 
     Scale convention (let γ = threshold/value):
       MHA/GQA wq/wk : scale = √γ on both Q and K; product = γ.
-      MLA wq_b      : scale = γ^(1 - k_ratio).  Q carries (1-k_ratio) of the
-                       log-reduction; k_ratio=0 → full γ on Q only (K unclipped).
+      MLA wq_b      : scale = γ^(1 - k_ratio) for the q_nope rows.  The q_pe
+                       (rotary) rows instead take the full γ in qk_clip, since
+                       the shared rotary key is left unchanged and the rotary
+                       query must carry the whole reduction.  k_ratio=0 → full
+                       γ on all Q rows (K unclipped).
       MLA wkv_b     : scale = γ^k_ratio applied to k_nope rows only (v rows
                        untouched). Returns None when k_ratio=0 (K unclipped).
-    Together wq_b × wkv_b deliver γ^(1-k_ratio) × γ^k_ratio = γ end-to-end.
+    Per logit component: content = γ^(1-k_ratio)·γ^k_ratio = γ;
+    rotary = γ·1 = γ.  Both are reduced by exactly γ end-to-end.
     """
     kind = qk_clip_state.kind
 
@@ -203,7 +207,8 @@ def qk_clip(p, scales, info):
             Exponent semantics depend on kind — see compute_scales docstring.
         info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions.
 
-    MLA wq_b: γ^(1-k_ratio) applied uniformly to q_nope and q_pe rows.
+    MLA wq_b: γ^(1-k_ratio) on q_nope rows; full γ on q_pe rows (the shared
+        rotary key is unchanged, so the rotary query carries the whole γ).
     MLA wkv_b: γ^k_ratio applied to k_nope rows only; v rows untouched.
     MHA/GQA: √γ applied to all rows in each head.
     """
@@ -220,7 +225,16 @@ def qk_clip(p, scales, info):
         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.view(-1, 1, 1))  # q_nope → γ^(1-k_ratio)
-        W_3d[:, qk_nope:, :].mul_(scales.view(-1, 1, 1))  # q_pe   → γ^(1-k_ratio)
+        # q_pe: the decoupled rotary key is shared across heads and left
+        # unchanged, so the rotary logit term q_pe·k_rope is reduced by the
+        # query side alone. The rotary query rows must therefore carry the full
+        # γ (exponent 1), independent of k_ratio; scaling them by γ^(1-k_ratio)
+        # like q_nope would under-clip the rotary logit whenever k_ratio > 0.
+        if info.k_ratio < 1.0:
+            q_pe_scales = scales**(1.0 / (1.0 - info.k_ratio))  # γ^(1-k) → γ
+        else:
+            q_pe_scales = scales  # degenerate: all burden on K, Q unclipped
+        W_3d[:, qk_nope:, :].mul_(q_pe_scales.view(-1, 1, 1))  # q_pe → γ
 
     elif info.kind == 'wkv_b':
         qk_nope = info.qk_nope_head_dim