fix(personaplex): preserve reference RoPE precision

families/personaplex/decoder_tp_builder.py

@@ -189,11 +189,11 @@ def _cast_work_dtype(tensor: trt.ITensor) -> trt.ITensor:
             partial_rotary_factor,
             interleaved=interleaved_rope,
         )
-        cos_half_tensor = _cast_work_dtype(
-            graph_ops.add_constant(network, cos_half_np.shape, cos_half_np, dtype=work_np_dtype)
+        cos_half_tensor = graph_ops.add_constant(
+            network, cos_half_np.shape, cos_half_np, dtype=np.float32
         )
-        sin_half_tensor = _cast_work_dtype(
-            graph_ops.add_constant(network, sin_half_np.shape, sin_half_np, dtype=work_np_dtype)
+        sin_half_tensor = graph_ops.add_constant(
+            network, sin_half_np.shape, sin_half_np, dtype=np.float32
         )
     else:
         raise NotImplementedError("PersonaPlex tensor-parallel builds require RoPE")

families/personaplex/default_decoder.py

@@ -279,13 +279,11 @@ def _cast_work_dtype(tensor: trt.ITensor) -> trt.ITensor:
             interleaved=interleaved_rope,
         )
         cos_half_tensor = graph_ops.add_constant(
-            network, cos_half_np.shape, cos_half_np, dtype=work_np_dtype
+            network, cos_half_np.shape, cos_half_np, dtype=np.float32
         )
-        cos_half_tensor = _cast_work_dtype(cos_half_tensor)
         sin_half_tensor = graph_ops.add_constant(
-            network, sin_half_np.shape, sin_half_np, dtype=work_np_dtype
+            network, sin_half_np.shape, sin_half_np, dtype=np.float32
         )
-        sin_half_tensor = _cast_work_dtype(sin_half_tensor)
     elif position_type == "learned":
         pos_embed_np = weights["position_embedding"]
         position_embed_table = const_in_work_dtype(

families/personaplex/default_dual_profile_decoder.py

@@ -290,11 +290,11 @@ def _add_profile(opt_sq: int, max_sq: int, *, fixed: bool = False):
             partial_rotary_factor,
             interleaved=interleaved_rope,
         )
-        cos_half_table = _const_in_work_dtype(
-            network, cos_half_np.shape, cos_half_np, work_np_dtype, work_trt_dtype
+        cos_half_table = graph_ops.add_constant(
+            network, cos_half_np.shape, cos_half_np, dtype=np.float32
         )
-        sin_half_table = _const_in_work_dtype(
-            network, sin_half_np.shape, sin_half_np, work_np_dtype, work_trt_dtype
+        sin_half_table = graph_ops.add_constant(
+            network, sin_half_np.shape, sin_half_np, dtype=np.float32
         )
 
     # Learned position embedding (GPT-2 / OPT / GPT-Neo / XGLM).

families/personaplex/graph_ops.py

@@ -8,6 +8,8 @@
 
 from __future__ import annotations
 
+import math
+
 import numpy as np
 import tensorrt as trt
 
@@ -789,17 +791,13 @@ def make_rope_table_half_dim(
     default = 1.0 if cosine else 0.0
     if max_cache_length <= 0 or rope_theta <= 0.0:
         return np.full((max(max_cache_length, 1), max(half, 1)), default, dtype=np.float32)
-    table = np.full((max_cache_length, half), default, dtype=np.float32)
-    for pos in range(max_cache_length):
-        for d in range(half):
-            # For both interleaved and rotate-half the frequency index is d
-            # (the distinction only affects which input pair is rotated; the
-            # freq assignment per half-dim is the same).
-            exponent = (2.0 * d) / rotary_ndims
-            inv_freq = rope_theta ** (-exponent)
-            angle = pos * inv_freq
-            table[pos, d] = np.cos(angle) if cosine else np.sin(angle)
-    return table
+    # Moshi forms frequencies, angles and trigonometric values in FP32.
+    # Computing them in float64 and converting afterwards changes coefficients.
+    frequency_scale = np.float32(-math.log(rope_theta) * 2.0 / rotary_ndims)
+    frequencies = np.exp(np.arange(half, dtype=np.float32) * frequency_scale)
+    positions = np.arange(max_cache_length, dtype=np.float32)
+    angles = positions[:, None] * frequencies[None, :]
+    return np.cos(angles) if cosine else np.sin(angles)
 
 
 def reshape_rows_to_heads_4d(
@@ -979,7 +977,11 @@ def add_apply_rope_native(
     """
     rotary_embedding_dim = validate_native_rope_dim(rotary_embedding_dim)
     attention_size = num_heads * head_dim
-
+    output_dtype = inp.dtype
+    # Moshi rotates FP32 real/imaginary components, then rounds to the input dtype.
+    inp = _cast_back_to_trt_dtype(network, inp, trt.float32)
+    cos_cache_2d = _cast_back_to_trt_dtype(network, cos_cache_2d, trt.float32)
+    sin_cache_2d = _cast_back_to_trt_dtype(network, sin_cache_2d, trt.float32)
     inp_4d = reshape_rows_to_heads_4d(network, inp, num_heads, head_dim, sequence_length)
 
     # Reshape position_id [Sq] -> [1, Sq] (batch=1).
@@ -996,7 +998,8 @@ def add_apply_rope_native(
     )
     rope.set_input(3, pos_2d.get_output(0))
 
-    return reshape_heads_4d_to_rows(network, rope.get_output(0), attention_size, sequence_length)
+    rotated = reshape_heads_4d_to_rows(network, rope.get_output(0), attention_size, sequence_length)
+    return _cast_back_to_trt_dtype(network, rotated, output_dtype)
 
 
 def add_apply_rope_native_sequence(

families/personaplex/tests/test_rotary_precision.py

@@ -0,0 +1,95 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Exercise RoPE through the decoder's normal KV-cache output."""
+
+import numpy as np
+import pytest
+
+trt = pytest.importorskip("tensorrt")
+torch = pytest.importorskip("torch")
+
+from families.personaplex.checkpoint_mapper import WeightDict  # noqa: E402
+from families.personaplex.config import ModelConfig  # noqa: E402
+from families.personaplex.default_decoder import build_standard_decoder_engine  # noqa: E402
+
+
+@pytest.mark.gpu
+@pytest.mark.trt
+@pytest.mark.parametrize("rebuild", range(3))
+@pytest.mark.parametrize("precision", ["fp16", "fp32"])
+@pytest.mark.parametrize("role", ["decode", "dual_profile"])
+def test_decoder_rotates_before_rounding_kv_cache(precision, role, rebuild):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    del rebuild  # Each parametrized invocation builds a separate engine.
+    width, cache_length = 128, 2
+    config = ModelConfig(
+        vocab_size=2,
+        hidden_size=width,
+        intermediate_size=width,
+        num_hidden_layers=1,
+        num_attention_heads=1,
+        num_key_value_heads=1,
+        raw={"_decoder_engine_role": role},
+    )
+    weights = WeightDict({
+        "embedding": np.ones((2, width), np.float32),
+        "final_norm": np.ones(width, np.float32),
+        "w_out": np.zeros((width, 2), np.float32),
+        "layer.0.input_norm": np.ones(width, np.float32),
+        "layer.0.post_attn_norm": np.ones(width, np.float32),
+    })
+    for name in ("w_q", "w_k", "w_v", "w_o", "w_gate", "w_up", "w_down"):
+        weights[f"layer.0.{name}"] = np.zeros((width, width), np.float32)
+    # Constant K=(1,1,0,...), independent of attention and GEMM accumulation.
+    weights["layer.0.k_bias"] = np.zeros(width, np.float32)
+    weights["layer.0.k_bias"][:2] = 1.0
+    plan = build_standard_decoder_engine(
+        config, weights, cache_length, precision=precision, interleaved_rope=True
+    )
+    logger = trt.Logger(trt.Logger.ERROR)
+    runtime = trt.Runtime(logger)
+    engine = runtime.deserialize_cuda_engine(plan)
+    assert engine is not None
+    context = engine.create_execution_context()
+    assert context is not None
+    dtype = torch.float16 if precision == "fp16" else torch.float32
+    feeds = {
+        "token_id": torch.zeros(1, dtype=torch.int32, device="cuda"),
+        "position_id": torch.ones(1, dtype=torch.int32, device="cuda"),
+        "attention_mask": torch.zeros((1, cache_length + 1), device="cuda"),
+        "cache_k_0": torch.zeros((cache_length, width), dtype=dtype, device="cuda"),
+        "cache_v_0": torch.zeros((cache_length, width), dtype=dtype, device="cuda"),
+    }
+    for index in range(engine.num_io_tensors):
+        name = engine.get_tensor_name(index)
+        if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
+            if -1 in tuple(engine.get_tensor_shape(name)):
+                assert context.set_input_shape(name, tuple(feeds[name].shape))
+    outputs = {}
+    torch_dtypes = {trt.float16: torch.float16, trt.float32: torch.float32}
+    for index in range(engine.num_io_tensors):
+        name = engine.get_tensor_name(index)
+        if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
+            tensor = feeds[name]
+        else:
+            tensor = torch.empty(
+                tuple(context.get_tensor_shape(name)),
+                dtype=torch_dtypes[engine.get_tensor_dtype(name)],
+                device="cuda",
+            )
+            outputs[name] = tensor
+        assert context.set_tensor_address(name, tensor.data_ptr())
+    assert context.execute_async_v3(torch.cuda.current_stream().cuda_stream)
+    torch.cuda.synchronize()
+    actual = outputs["present_k_0"].cpu().reshape(-1)
+    del context, engine, runtime
+    # Pinned Moshi applies FP32 real/imaginary products and only then converts
+    # to the input dtype. Frequency zero at position one has angle exactly one.
+    angle = torch.tensor(1.0, dtype=torch.float32, device="cuda")
+    expected = torch.zeros(width, dtype=dtype, device="cuda")
+    expected[0] = (angle.cos() - angle.sin()).to(dtype)
+    expected[1] = (angle.sin() + angle.cos()).to(dtype)
+    assert actual.dtype == dtype  # The KV cache retains the model's precision.
+    torch.testing.assert_close(actual, expected.cpu(), rtol=0, atol=0)