From 1f71d8f302b39f20b32367fdefdf0a09ec5d6111 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Sun, 18 Oct 2026 08:35:02 +0000 Subject: [PATCH 1/2] fix(personaplex): preserve FP32 rotary arithmetic Keep cosine and sine constants in FP32 in the single-step, dual-profile and tensor-parallel decoder builders. Rotate in FP32 and round queries and keys back before caching, matching the official model arithmetic boundaries. Add an independent Torch regression through the normal KV-cache output, with repeated fresh tiny engines. The original Native FP16 benchmark and accuracy gates remain unchanged. Fixes #1790 Signed-off-by: chaofengw --- families/personaplex/decoder_tp_builder.py | 8 +- families/personaplex/default_decoder.py | 6 +- .../default_dual_profile_decoder.py | 8 +- families/personaplex/graph_ops.py | 9 +- .../tests/test_rotary_precision.py | 95 +++++++++++++++++++ 5 files changed, 112 insertions(+), 14 deletions(-) create mode 100644 families/personaplex/tests/test_rotary_precision.py 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..d6edb40b2f 100644 --- a/families/personaplex/graph_ops.py +++ b/families/personaplex/graph_ops.py @@ -979,7 +979,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 +1000,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) From 97197612aa3cb0eb762f7a26bba6147bdb4aa924 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Sun, 18 Oct 2026 08:50:31 +0000 Subject: [PATCH 2/2] fix(personaplex): generate FP32 rotary coefficients Match the reference frequency, angle and trigonometric evaluation precision. Evaluating in FP64 and only converting afterwards changes cached keys even for a decoder configured as FP32. The same family helper also supplies Mimi encoder tables. The original and corrected decoders were each independently built three times per precision and graph path. All 12 original cases fail numerical assertions; all 12 corrected cases pass the identical CUDA Torch oracle without changing existing accuracy gates. Refs: #1790 Signed-off-by: chaofengw --- families/personaplex/graph_ops.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/families/personaplex/graph_ops.py b/families/personaplex/graph_ops.py index d6edb40b2f..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(