diff --git a/families/personaplex/decoder_tp_builder.py b/families/personaplex/decoder_tp_builder.py index 7adf423e75..8ba1fdf943 100644 --- a/families/personaplex/decoder_tp_builder.py +++ b/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") diff --git a/families/personaplex/default_decoder.py b/families/personaplex/default_decoder.py index 272a57a595..4ef8f4061b 100644 --- a/families/personaplex/default_decoder.py +++ b/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( diff --git a/families/personaplex/default_dual_profile_decoder.py b/families/personaplex/default_dual_profile_decoder.py index 19fca8d359..18c77ca393 100644 --- a/families/personaplex/default_dual_profile_decoder.py +++ b/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). diff --git a/families/personaplex/graph_ops.py b/families/personaplex/graph_ops.py index d2af56e069..8396fa8eff 100644 --- a/families/personaplex/graph_ops.py +++ b/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( diff --git a/families/personaplex/tests/test_rotary_precision.py b/families/personaplex/tests/test_rotary_precision.py new file mode 100644 index 0000000000..46ef64b37c --- /dev/null +++ b/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)