Merge pull request #32 from MotifTechnologies/feat/qk-clip-k-ratio

feat: add k_ratio for Q/K split control in MLA QK clipping [skip-build]

torch-ext/optimizer/muon.py

@@ -362,7 +362,7 @@ class Muon(torch.optim.Optimizer):
             qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits)
 
             scales_full = compute_scales(
-                p, qk_clip_state) if qk_clip_state is not None else None
+                p, qk_clip_state, n) if qk_clip_state is not None else None
             if scales_full is not None:
                 qk_clip(p, scales_full, qk_clip_state)
 
@@ -409,7 +409,7 @@ class Muon(torch.optim.Optimizer):
                 qk_clip_state = get_qk_clip_info(self.clip_config, n,
                                                  qk_logits)
                 scales_full = compute_scales(
-                    p, qk_clip_state) if qk_clip_state is not None else None
+                    p, qk_clip_state, n) if qk_clip_state is not None else None
                 if scales_full is not None:
                     qk_clip(p, scales_full, qk_clip_state)
 
@@ -476,7 +476,7 @@ class Muon(torch.optim.Optimizer):
                         qk_clip_state = get_qk_clip_info(
                             self.clip_config, n, qk_logits)
                         scales_full = compute_scales(
-                            p, qk_clip_state
+                            p, qk_clip_state, n
                         ) if qk_clip_state is not None else None
                         if scales_full is not None:
                             ratio = p.shape[0] // scales_full.shape[0]
@@ -555,7 +555,7 @@ class Muon(torch.optim.Optimizer):
                         qk_clip_state = get_qk_clip_info(
                             self.clip_config, n, qk_logits)
                         scales_full = compute_scales(
-                            p, qk_clip_state
+                            p, qk_clip_state, n
                         ) if qk_clip_state is not None else None
                         if scales_full is not None:
                             ratio = p.shape[0] // scales_full.shape[0]

torch-ext/optimizer/pipeline.py

@@ -322,7 +322,7 @@ def _update_params(
         state = param_to_state[id(p)]
         if state.qk_clip_state is None:
             continue
-        scales_full = compute_scales(p, state.qk_clip_state)
+        scales_full = compute_scales(p, state.qk_clip_state, state.name)
         if scales_full is not None:
             ratio = p.shape[0] // scales_full.shape[0]
             idx0 = state.rank_indices[rank][0]

torch-ext/optimizer/qk_clip.py

@@ -63,6 +63,9 @@ class QKClipInfo:
     qk_nope_head_dim: int = 0
     qk_rope_head_dim: int = 0
     v_head_dim: int = 0
+    # Fraction of log-reduction assigned to K (wkv_b).  Q gets (1-k_ratio).
+    # 0.0 = all Q (K not clipped), 0.5 = equal split (MHA-like), 1.0 = all K.
+    k_ratio: float = 0.0
 
 
 def get_qk_clip_info(clip_config, n, qk_logits):
@@ -85,6 +88,7 @@ def get_qk_clip_info(clip_config, n, qk_logits):
     threshold = clip_config.get('threshold')
     kind, layer_idx = parse_qk_layer(n)
     is_mla = clip_config.get('is_mla', False)
+    k_ratio = float(clip_config.get('k_ratio', 0.0))
 
     logit, indices = None, []
     if qk_logits is not None and kind is not None:
@@ -110,6 +114,7 @@ def get_qk_clip_info(clip_config, n, qk_logits):
             qk_nope_head_dim=clip_config['qk_nope_head_dim'],
             qk_rope_head_dim=clip_config['qk_rope_head_dim'],
             v_head_dim=clip_config['v_head_dim'],
+            k_ratio=k_ratio,
         )
     else:
         return QKClipInfo(
@@ -121,35 +126,62 @@ def get_qk_clip_info(clip_config, n, qk_logits):
         )
 
 
-def compute_scales(p, qk_clip_state):
+def compute_scales(p, qk_clip_state, name: str = ''):
     """Compute per-head scaling factors for QK clipping.
 
-    Returns scales tensor (√γ per head) if any head exceeds threshold, else None.
-    For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim.
+    Returns scales tensor per head if any head exceeds threshold, else None.
+
+    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 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.
     """
     kind = qk_clip_state.kind
+
+    # MLA wkv_b: skip entirely when k_ratio=0 (all reduction on Q).
+    if qk_clip_state.is_mla and kind == 'wkv_b' and qk_clip_state.k_ratio == 0.0:
+        return None
+
     indices = qk_clip_state.indices
     head_dim = qk_clip_state.head_dim
     threshold = qk_clip_state.threshold
     logit = qk_clip_state.logit
 
-    # Check if any head exceeds threshold before allocating.
+    # Exponent for (threshold/v_ele)^exponent:
+    #   MLA wq_b  → 1 - k_ratio   (Q carries this fraction of log-reduction)
+    #   MLA wkv_b → k_ratio        (K carries the rest; applied to k_nope rows only)
+    #   MHA/GQA   → 0.5            (equal split; equivalent to k_ratio=0.5)
+    #
+    # Principled k_ratio for Q:K = G:1 GQA: 1 / (1 + sqrt(G))
+    #   K's gradient is a sum of G Q-head terms that partially cancel, so its
+    #   effective recovery signal scales as 1/sqrt(G) (random-walk argument).
+    #   Setting clip/recovery equal for both sides:
+    #     (1-r)/1 = r/(1/sqrt(G))  =>  r = 1/(1+sqrt(G))
+    #   G=1 → 0.5 (recovers MHA), G=4 → ≈0.33, G=8 → ≈0.26.
+    if qk_clip_state.is_mla:
+        exponent = (1.0 - qk_clip_state.k_ratio) if kind == 'wq_b' else qk_clip_state.k_ratio
+    else:
+        exponent = 0.5
+
     head_scales = {}
     for logit_idx, head_idx in enumerate(indices):
         v_ele = float(logit[logit_idx])
         if v_ele > threshold:
-            new_scale = math.sqrt(threshold / v_ele)
+            new_scale = (threshold / v_ele) ** exponent
             if head_idx not in head_scales or new_scale < head_scales[head_idx]:
                 head_scales[head_idx] = new_scale
                 logger.info(
-                    f"[{kind}] Head {head_idx} exceeded threshold "
+                    f"[{kind}] {name} Head {head_idx} exceeded threshold "
                     f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}"
                 )
 
     if not head_scales:
         return None
 
-    # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows
+    # wkv_b row-stride per head = qk_nope + v_head_dim (not head_dim)
     if qk_clip_state.is_mla and kind == 'wkv_b':
         effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim
     else:
@@ -167,12 +199,13 @@ def qk_clip(p, scales, info):
 
     Args:
         p: Parameter (nn.Parameter or raw tensor).
-        scales: [n_heads] tensor, each element = √γ_h.
+        scales: [n_heads] tensor of per-head scale factors (γ^exponent).
+            Exponent semantics depend on kind — see compute_scales docstring.
         info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions.
 
-    MLA sub-region scaling per Algorithm 1 (MuonClip):
-        wq_b: q_nope rows → √γ,  q_pe rows → γ
-        wkv_b: k_nope rows → √γ, v rows → unchanged
+    MLA wq_b: γ^(1-k_ratio) applied uniformly to q_nope and q_pe rows.
+    MLA wkv_b: γ^k_ratio applied to k_nope rows only; v rows untouched.
+    MHA/GQA: √γ applied to all rows in each head.
     """
     W = p.data if isinstance(p, torch.nn.Parameter) else p
 
@@ -186,13 +219,12 @@ def qk_clip(p, scales, info):
         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.view(-1, 1, 1))  # q_nope → √γ
-        W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1,
-                                                         1))  # q_pe   → γ
+        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)
 
     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.view(-1, 1, 1))  # k_nope → √γ
+        W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1))  # k_nope → γ^k_ratio
         # v rows: not touched (k_R shared rotary unchanged)