fix(personaplex): preserve residual input rounding

families/personaplex/decoder_tp_builder.py

@@ -418,9 +418,7 @@ def _add_tp_decoder_layer(
         else:
             norm2 = attn["normed"]
     else:
-        residual1 = network.add_elementwise(
-            hidden, attn_out, trt.ElementWiseOperation.SUM
-        ).get_output(0)
+        residual1 = graph_ops.add_residual_sum(network, hidden, attn_out)
         norm2 = _apply_norm(
             network,
             residual1,
@@ -447,17 +445,11 @@ def _add_tp_decoder_layer(
     mlp_out = add_all_reduce_sum(network, mlp_out, tp_size)
 
     if parallel_residual:
-        sum_attn = network.add_elementwise(
-            hidden, attn_out, trt.ElementWiseOperation.SUM
-        ).get_output(0)
-        residual2 = network.add_elementwise(
-            sum_attn, mlp_out, trt.ElementWiseOperation.SUM
-        ).get_output(0)
+        sum_attn = graph_ops.add_residual_sum(network, hidden, attn_out)
+        residual2 = graph_ops.add_residual_sum(network, sum_attn, mlp_out)
         post_attn_tensor = sum_attn
     else:
-        residual2 = network.add_elementwise(
-            residual1, mlp_out, trt.ElementWiseOperation.SUM
-        ).get_output(0)
+        residual2 = graph_ops.add_residual_sum(network, residual1, mlp_out)
         post_attn_tensor = residual1
 
     return {

families/personaplex/default_decoder.py

@@ -667,10 +667,10 @@ def _add_decoder_layer(
         else:
             norm2 = attn["normed"]
     else:
-        residual1 = network.add_elementwise(hidden, attn_out, trt.ElementWiseOperation.SUM)
+        residual1 = graph_ops.add_residual_sum(network, hidden, attn_out)
         norm2 = _apply_norm(
             network,
-            residual1.get_output(0),
+            residual1,
             hidden_size,
             weights[f"{prefix}.post_attn_norm"],
             weights.get(f"{prefix}.post_attn_norm_beta"),
@@ -709,19 +709,15 @@ def _add_decoder_layer(
 
     # Final residual connection
     if parallel_residual:
-        sum_attn = network.add_elementwise(hidden, attn_out, trt.ElementWiseOperation.SUM)
-        residual2 = network.add_elementwise(
-            sum_attn.get_output(0), mlp_out, trt.ElementWiseOperation.SUM
-        )
-        post_attn_tensor = sum_attn.get_output(0)
+        sum_attn = graph_ops.add_residual_sum(network, hidden, attn_out)
+        residual2 = graph_ops.add_residual_sum(network, sum_attn, mlp_out)
+        post_attn_tensor = sum_attn
     else:
-        residual2 = network.add_elementwise(
-            residual1.get_output(0), mlp_out, trt.ElementWiseOperation.SUM
-        )
-        post_attn_tensor = residual1.get_output(0)
+        residual2 = graph_ops.add_residual_sum(network, residual1, mlp_out)
+        post_attn_tensor = residual1
 
     return {
-        "hidden": residual2.get_output(0),
+        "hidden": residual2,
         "post_attn": post_attn_tensor,
         "present_k": present_k,
         "present_v": present_v,

families/personaplex/default_dual_profile_decoder.py

@@ -527,12 +527,10 @@ def _add_profile(opt_sq: int, max_sq: int, *, fixed: bool = False):
             else:
                 norm2 = normed
         else:
-            residual1 = network.add_elementwise(
-                hidden_state, attn_out, trt.ElementWiseOperation.SUM
-            )
+            residual1 = graph_ops.add_residual_sum(network, hidden_state, attn_out)
             norm2 = _norm_multi(
                 network,
-                residual1.get_output(0),
+                residual1,
                 hidden,
                 weights[f"{prefix}.post_attn_norm"],
                 weights.get(f"{prefix}.post_attn_norm_beta"),
@@ -567,15 +565,11 @@ def _add_profile(opt_sq: int, max_sq: int, *, fixed: bool = False):
 
         # Final residual.
         if parallel_residual:
-            sum_attn = network.add_elementwise(hidden_state, attn_out, trt.ElementWiseOperation.SUM)
-            residual2 = network.add_elementwise(
-                sum_attn.get_output(0), mlp_out, trt.ElementWiseOperation.SUM
-            )
+            sum_attn = graph_ops.add_residual_sum(network, hidden_state, attn_out)
+            residual2 = graph_ops.add_residual_sum(network, sum_attn, mlp_out)
         else:
-            residual2 = network.add_elementwise(
-                residual1.get_output(0), mlp_out, trt.ElementWiseOperation.SUM
-            )
-        hidden_state = residual2.get_output(0)
+            residual2 = graph_ops.add_residual_sum(network, residual1, mlp_out)
+        hidden_state = residual2
 
     # ---- Final norm + LM head -------------------------------------------
     final_norm = weights.get("final_norm")

families/personaplex/graph_ops.py

@@ -23,6 +23,21 @@ def _cast_back_to_trt_dtype(
     return network.add_cast(tensor, target_dtype).get_output(0)
 
 
+def add_residual_sum(
+    network: trt.INetworkDefinition, residual: trt.ITensor, update: trt.ITensor
+) -> trt.ITensor:
+    """Add already rounded branch outputs and restore the residual dtype.
+
+    Summing in FP32 keeps the producer's Half rounding boundary distinct from
+    the residual add. Otherwise GEMM fusion can add the unrounded accumulator.
+    """
+    output_dtype = residual.dtype
+    residual = _cast_back_to_trt_dtype(network, residual, trt.float32)
+    update = _cast_back_to_trt_dtype(network, update, trt.float32)
+    total = network.add_elementwise(residual, update, trt.ElementWiseOperation.SUM).get_output(0)
+    return _cast_back_to_trt_dtype(network, total, output_dtype)
+
+
 def _add_matrix_multiply_with_fp32_accumulation(
     network: trt.INetworkDefinition,
     lhs: trt.ITensor,

families/personaplex/tests/test_residual_precision.py

@@ -0,0 +1,136 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Keep the linear-output rounding boundary inside the actual decoder layer."""
+
+import numpy as np
+import pytest
+
+trt = pytest.importorskip("tensorrt")
+torch = pytest.importorskip("torch")
+
+from families.personaplex import graph_ops  # noqa: E402
+from families.personaplex.default_decoder import _add_decoder_layer  # noqa: E402
+
+
+@pytest.mark.gpu
+@pytest.mark.trt
+@pytest.mark.parametrize("rebuild", range(3))
+@pytest.mark.parametrize("precision", ["fp16", "fp32"])
+def test_decoder_residual_rounds_the_projection_first(precision, rebuild):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    del rebuild
+    width, mlp_size, heads, cache_length = 1024, 2816, 16, 18
+    np_dtype = np.float16 if precision == "fp16" else np.float32
+    trt_dtype = trt.float16 if precision == "fp16" else trt.float32
+    torch_dtype = torch.float16 if precision == "fp16" else torch.float32
+    rng = np.random.default_rng(42)
+
+    def dense():
+        return (
+            rng.integers(-128, 129, size=(width, width), dtype=np.int16).astype(np.float32) / 1024
+        )
+
+    w_v, w_o = dense(), dense()
+    w_v[:, :2] = 0
+    w_v[0, 0], w_v[1, 1] = 1.0, 1.0
+    w_o[:, 0] = 0
+    w_o[0, 0], w_o[1, 0] = 1.0, 2**-12
+    alpha = np.ones(width, np.float32)
+    alpha[0] = -1.0
+    weights = {
+        "L.input_norm": alpha,
+        "L.post_attn_norm": np.ones(width, np.float32),
+        "L.w_q": np.zeros((width, width), np.float32),
+        "L.w_k": np.zeros((width, width), np.float32),
+        "L.w_v": w_v,
+        "L.w_o": w_o,
+        "L.w_gate": np.zeros((width, mlp_size), np.float32),
+        "L.w_up": np.zeros((width, mlp_size), np.float32),
+        "L.w_down": np.zeros((mlp_size, width), np.float32),
+    }
+    logger = trt.Logger(trt.Logger.ERROR)
+    builder = trt.Builder(logger)
+    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+    hidden = network.add_input("hidden", trt_dtype, (1, width))
+    cache_k = network.add_input("cache_k", trt_dtype, (cache_length, width))
+    cache_v = network.add_input("cache_v", trt_dtype, (cache_length, width))
+    mask = network.add_input("mask", trt.float32, (1, cache_length + 1))
+    if mask.dtype != trt_dtype:
+        mask = network.add_cast(mask, trt_dtype).get_output(0)
+    position = graph_ops.add_constant(network, (1,), np.zeros(1, np.int32), dtype=np.int32)
+    epsilon = graph_ops.add_constant(network, (1, 1), np.array([[1e-8]], np.float32))
+    layer = _add_decoder_layer(
+        network=network,
+        hidden=hidden,
+        cache_k=cache_k,
+        cache_v=cache_v,
+        attention_mask=mask,
+        position_id=position,
+        attention_scale=1 / 8,
+        eps_tensor=epsilon,
+        eps=1e-8,
+        weights=weights,
+        prefix="L",
+        hidden_size=width,
+        attention_size=width,
+        kv_attention_size=width,
+        mlp_size=mlp_size,
+        num_heads=heads,
+        num_kv_heads=heads,
+        head_dim=64,
+        max_cache_length=cache_length,
+        position_type="learned",
+        dtype=np_dtype,
+    )
+    for name in ("hidden", "present_k", "present_v"):
+        tensor = layer[name]
+        tensor.name = "output_" + name
+        network.mark_output(tensor)
+    config = builder.create_builder_config()
+    config.clear_flag(trt.BuilderFlag.TF32)
+    plan = builder.build_serialized_network(network, config)
+    assert plan is not None
+    runtime = trt.Runtime(logger)
+    engine = runtime.deserialize_cuda_engine(plan)
+    assert engine is not None
+    execution = engine.create_execution_context()
+    assert execution is not None
+    feeds = {
+        "hidden": torch.ones((1, width), dtype=torch_dtype, device="cuda"),
+        "cache_k": torch.zeros((cache_length, width), dtype=torch_dtype, device="cuda"),
+        "cache_v": torch.zeros((cache_length, width), dtype=torch_dtype, device="cuda"),
+        "mask": torch.full((1, cache_length + 1), -10000.0, device="cuda"),
+    }
+    feeds["hidden"][0, 0] = -1.0
+    feeds["mask"][0, -1] = 0.0
+    outputs = {}
+    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(engine.get_tensor_shape(name)), dtype=torch_dtype, device="cuda"
+            )
+            outputs[name] = tensor
+        assert execution.set_tensor_address(name, tensor.data_ptr())
+    assert execution.execute_async_v3(torch.cuda.current_stream().cuda_stream)
+    torch.cuda.synchronize()
+    # RMSNorm gives an all-one vector; one valid key makes context exactly V.
+    # Native rounds linear(V) to the model dtype before adding the residual.
+    normalized = torch.ones((1, width), dtype=torch_dtype, device="cuda")
+    native_v = torch.nn.functional.linear(
+        normalized, torch.from_numpy(w_v.T.copy()).to("cuda", torch_dtype)
+    )
+    native_projection = torch.nn.functional.linear(
+        native_v, torch.from_numpy(w_o.T.copy()).to("cuda", torch_dtype)
+    )
+    expected = (native_projection + feeds["hidden"])[0, 0].cpu()
+    actual = outputs["output_hidden"][0, 0].cpu()
+    actual_dtype = outputs["output_hidden"].dtype
+    del execution, engine, runtime
+    assert float(expected) == (0.0 if precision == "fp16" else 2**-12)
+    assert actual_dtype == torch_dtype
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)