[gpt_oss] Move YaRN attention mscale from rope cos/sin to softmax scale

CosSinRoPE is now a pure rotation (no mscale baked into cos/sin).
gpt-oss, which was the only CosSinRoPE+YaRN user, instead folds
mscale^2 into its softmax_scale.

Mathematically equivalent: applying mscale to cos/sin scaled q and k
each by m, so the attention logit was scaled by m^2. Applying m^2 to
the softmax scale produces the same logit. Verified logit-neutral to
~1e-6 on gpt-oss params.

Author: sdmyzlp

torchtitan/models/common/rope.py

@@ -106,7 +106,6 @@ class Config(Module.Config):
         beta_fast: float = 32.0
         beta_slow: float = 1.0
         original_seq_len: int = 4096
-        mscale: float = 0.0
         truncate: bool = True
 
     def __init__(self, config: Config):
@@ -282,13 +281,10 @@ def _precompute_cache(self) -> torch.Tensor:
         max_seq_len = cfg.max_seq_len
         base = cfg.theta
 
-        mscale = 1.0
-
         if cfg.scaling == "llama":
             raise NotImplementedError("Cos/sin RoPE does not support Llama scaling.")
 
         if cfg.scaling == "yarn" and cfg.rope_factor > 1.0:
-            mscale = 0.1 * math.log(cfg.rope_factor) + 1.0
             inv_freq = _yarn_inv_freq(
                 dim,
                 base,
@@ -307,8 +303,8 @@ def _precompute_cache(self) -> torch.Tensor:
         freqs = torch.outer(t, inv_freq).float()
         theta = torch.cat([freqs, freqs], dim=-1)
 
-        cos = theta.cos() * mscale
-        sin = theta.sin() * mscale
+        cos = theta.cos()
+        sin = theta.sin()
         return torch.cat([cos, sin], dim=-1)
 
     def _reshape_cache(

torchtitan/models/gpt_oss/model.py

@@ -104,6 +104,11 @@ def __init__(self, config: Config):
 
         # Standard attention softmax scale (1/sqrt(head_dim))
         self.softmax_scale = 1.0 / math.sqrt(self.head_dim)
+        if config.rope.scaling == "yarn" and config.rope.rope_factor > 1.0:
+            mscale = 0.1 * math.log(config.rope.rope_factor) + 1.0
+            # Merge YaRN attention mscale into softmax_scale, with
+            # m**2 being equivalent to scaling q / k each by mscale.
+            self.softmax_scale *= mscale * mscale
 
         self.qkv_linear = config.qkv_linear.build()
         self.wo = config.wo.build()