llama/rope: gate fp64 hf_precompute_freqs_cis on cos/sin scaling (pytorch#19308)

### Summary
a79521b ("Add LongRoPE support and fp64 RoPE precompute for Phi-3 /
Phi-4 family") unconditionally moved hf_precompute_freqs_cis to fp64
cos/sin precompute with a final cast to fp32. That works for the Phi-4
device validation that motivated the commit, but it broke
test_static_attention.py::test_within_transformer on the Linux unittest
runners (pull, pull-editable, trunk-release have been 100% red since the
commit landed).

The test compares mha_transformer (built with use_hf_rope=False, taking
the pure-fp32 precompute_freqs_cis path) against static_transformer
(built with use_hf_rope=True, taking hf_precompute_freqs_cis) at
rtol=1e-3, with shared weights. Before a79521b, both paths produced
bit-identical fp32 cos/sin tables (verified empirically: 0/192 entries
differed). After the commit, HF cos/sin diverge from non-HF by ~1 ULP in
38/192 entries; that drift compounds across 4 transformer layers and
tips past rtol=1e-3 on the CI runners (Python 3.10, source-built torch).
Local Python 3.12 stayed just barely within tolerance, which is why
review missed it.

Gate the fp64 precompute on the property the original commit was
actually protecting: a non-trivial cos/sin scale being applied. That is
either LongRoPE active (Phi-3 / Phi-4 set short_factor and long_factor
via config) or an explicit attention_factor != 1.0 passed through. Both
cases preserve fp64; vanilla HF RoPE (Llama family, the test config)
goes back to fp32 throughout and re-establishes bit-identical agreement
with the non-HF path.

Authored with Claude Code.

### Test plan
CI

Author: rascani

examples/models/llama/rope.py

@@ -154,21 +154,32 @@ def hf_precompute_freqs_cis(
     # Partial rotary embeddings.
     dim = int(dim * partial_rotary_factor)
 
-    # Compute the RoPE table in fp64 to minimize ULP-level drift; cast to fp32
-    # once at the end. Phi-4 Mini's narrow decode-time logit margins make the
-    # exported model sensitive to 1-ULP differences in freqs_cos / freqs_sin
-    # under sampling, especially on the Vulkan delegate.
+    # fp64 precompute is required whenever cos/sin will be scaled by a
+    # non-trivial attention_factor (LongRoPE on Phi-3 / Phi-4 family). There,
+    # fp32 ULP-level rounding in the table is load-bearing on Vulkan under
+    # sampling -- a fp32-only regression manifests as decode-time n-gram
+    # looping, not a unit-test red. For vanilla HF RoPE, fp32 throughout
+    # produces cos/sin tables bit-identical to the non-HF precompute_freqs_cis
+    # path, which the static-attention vs MHA parity tests rely on.
+    #
+    # If you add a new model that needs cos/sin scaling but does not set
+    # short_factor / long_factor / attention_factor, extend the gate below.
+    longrope_active = (short_factor is not None) or (long_factor is not None)
+    needs_fp64 = longrope_active or (
+        attention_factor is not None and attention_factor != 1.0
+    )
+    compute_dtype = torch.float64 if needs_fp64 else torch.float32
+
     inv_freq = 1.0 / (
         theta
         ** (
-            torch.arange(0, dim, 2, device=device, dtype=torch.int64).to(torch.float64)
+            torch.arange(0, dim, 2, device=device, dtype=torch.int64).to(compute_dtype)
             / dim
         )
     )
 
     # LongRoPE: divide inv_freq element-wise by short_factor or long_factor.
     # Selection mirrors HF: long_factor when seq_len > original_max_position_embeddings.
-    longrope_active = (short_factor is not None) or (long_factor is not None)
     if longrope_active:
         chosen = (
             long_factor
@@ -178,7 +189,7 @@ def hf_precompute_freqs_cis(
         if chosen is None:
             # Fall back to whichever factor was provided.
             chosen = short_factor if long_factor is None else long_factor
-        ext_factors = torch.tensor(chosen, dtype=torch.float64, device=device)
+        ext_factors = torch.tensor(chosen, dtype=compute_dtype, device=device)
         assert ext_factors.numel() == inv_freq.numel(), (
             f"LongRoPE factor length {ext_factors.numel()} must equal dim/2 "
             f"({inv_freq.numel()})"
@@ -200,8 +211,8 @@ def hf_precompute_freqs_cis(
                 )
 
     # pyre-ignore Undefined attribute [16]: `float` has no attribute `device`.
-    t = torch.arange(end, device=inv_freq.device, dtype=torch.int64).to(torch.float64)
-    freqs = torch.outer(t, inv_freq).to(torch.float64)  # pyre-ignore
+    t = torch.arange(end, device=inv_freq.device, dtype=torch.int64).to(compute_dtype)
+    freqs = torch.outer(t, inv_freq).to(compute_dtype)  # pyre-ignore
     emb = torch.cat((freqs, freqs), dim=-1)
     cos_tab = torch.cos(emb)
     sin_tab = torch.sin(emb)