From 8a011cf9b2bf3d7f9d53f8d8c4b20754f107a330 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Sun, 18 Oct 2026 09:17:35 +0000 Subject: [PATCH] fix(personaplex): round the SwiGLU gate product once The reference gating kernel computes SiLU(gate) * up in FP32 from the Half (or BF16) projections and rounds only the product. The family graph applied the sigmoid, the SiLU product and the up product in the work dtype, rounding SiLU(gate) to Half before the multiply. add_swiglu_activation now casts the projections to FP32 for the pointwise math and rounds the result once to their dtype; add_swiglu_mlp and the dual-profile _swiglu_mlp both use it. FP32 networks are unchanged. A GPU test drives the MLP's normal output with gate 7/64 and up 1.5, where the reference gives 0.08648681640625 and the old rounding 0.0865478515625. Fixes #1793 Signed-off-by: chaofengw --- .../default_dual_profile_decoder.py | 8 +- families/personaplex/graph_blocks.py | 29 ++++-- .../tests/test_swiglu_precision.py | 95 +++++++++++++++++++ 3 files changed, 120 insertions(+), 12 deletions(-) create mode 100644 families/personaplex/tests/test_swiglu_precision.py diff --git a/families/personaplex/default_dual_profile_decoder.py b/families/personaplex/default_dual_profile_decoder.py index 19fca8d359..0221ea3270 100644 --- a/families/personaplex/default_dual_profile_decoder.py +++ b/families/personaplex/default_dual_profile_decoder.py @@ -83,12 +83,8 @@ def _swiglu_mlp( ) -> trt.ITensor: gate = matmul(inp, hidden, mlp_size, weights[f"{prefix}.w_gate"], f"{prefix}.w_gate") up = matmul(inp, hidden, mlp_size, weights[f"{prefix}.w_up"], f"{prefix}.w_up") - sigmoid = network.add_activation(gate, trt.ActivationType.SIGMOID) - swish = network.add_elementwise(gate, sigmoid.get_output(0), trt.ElementWiseOperation.PROD) - gated = network.add_elementwise(swish.get_output(0), up, trt.ElementWiseOperation.PROD) - mlp_out = matmul( - gated.get_output(0), mlp_size, hidden, weights[f"{prefix}.w_down"], f"{prefix}.w_down" - ) + gated = graph_blocks.add_swiglu_activation(network, gate, up) + mlp_out = matmul(gated, mlp_size, hidden, weights[f"{prefix}.w_down"], f"{prefix}.w_down") return mlp_out diff --git a/families/personaplex/graph_blocks.py b/families/personaplex/graph_blocks.py index 2fc6187026..e81d482606 100644 --- a/families/personaplex/graph_blocks.py +++ b/families/personaplex/graph_blocks.py @@ -313,6 +313,27 @@ def add_attention_block( } +def add_swiglu_activation( + network: trt.INetworkDefinition, gate: trt.ITensor, up: trt.ITensor +) -> trt.ITensor: + """SiLU(gate) * up in FP32, rounded once to the projections' dtype. + + The reference gating kernel computes the activation and the product in FP32 + from Half/BF16 projections and rounds only the result; rounding SiLU(gate) + first changes the product. + """ + work_dtype = gate.dtype + if work_dtype != trt.float32: + gate = network.add_cast(gate, trt.float32).get_output(0) + up = network.add_cast(up, trt.float32).get_output(0) + sigmoid = network.add_activation(gate, trt.ActivationType.SIGMOID).get_output(0) + swish = network.add_elementwise(gate, sigmoid, trt.ElementWiseOperation.PROD).get_output(0) + gated = network.add_elementwise(swish, up, trt.ElementWiseOperation.PROD).get_output(0) + if work_dtype != trt.float32: + gated = network.add_cast(gated, work_dtype).get_output(0) + return gated + + def add_swiglu_mlp( network: trt.INetworkDefinition, inp: trt.ITensor, @@ -332,13 +353,9 @@ def add_swiglu_mlp( gate = matmul(inp, hidden_size, mlp_size, weights[f"{prefix}.w_gate"], f"{_lp}.w_gate") up = matmul(inp, hidden_size, mlp_size, weights[f"{prefix}.w_up"], f"{_lp}.w_up") - sigmoid = network.add_activation(gate, trt.ActivationType.SIGMOID) - swish = network.add_elementwise(gate, sigmoid.get_output(0), trt.ElementWiseOperation.PROD) - gated = network.add_elementwise(swish.get_output(0), up, trt.ElementWiseOperation.PROD) + gated = add_swiglu_activation(network, gate, up) - mlp_out = matmul( - gated.get_output(0), mlp_size, hidden_size, weights[f"{prefix}.w_down"], f"{_lp}.w_down" - ) + mlp_out = matmul(gated, mlp_size, hidden_size, weights[f"{prefix}.w_down"], f"{_lp}.w_down") return mlp_out diff --git a/families/personaplex/tests/test_swiglu_precision.py b/families/personaplex/tests/test_swiglu_precision.py new file mode 100644 index 0000000000..14e0faf8b0 --- /dev/null +++ b/families/personaplex/tests/test_swiglu_precision.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Exercise the SwiGLU rounding boundary through the MLP's normal output.""" + +import numpy as np +import pytest + +trt = pytest.importorskip("tensorrt") +torch = pytest.importorskip("torch") + +from families.personaplex import graph_blocks # noqa: E402 +from families.personaplex.default_dual_profile_decoder import _swiglu_mlp # noqa: E402 + +WIDTH = 128 +# Binary values: exact in Half and BF16. Rounding SiLU(gate) to Half before the +# product gives 0.0865478515625; the FP32 product rounded once gives +# 0.08648681640625, about 600 FP32 ulps from a Half rounding boundary. +GATE, UP = 7.0 / 64.0, 1.5 + + +def _weights(): + # One-hot input selects row 0; each projection has a single unit path, so + # every GEMM output is exact and only the activation boundary can differ. + w_gate = np.zeros((WIDTH, WIDTH), np.float32) + w_up = np.zeros((WIDTH, WIDTH), np.float32) + w_down = np.zeros((WIDTH, WIDTH), np.float32) + w_gate[0, 0], w_up[0, 0], w_down[0, 0] = GATE, UP, 1.0 + return {"mlp.w_gate": w_gate, "mlp.w_up": w_up, "mlp.w_down": w_down} + + +def _build(path): + logger = trt.Logger(trt.Logger.ERROR) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + x = network.add_input("x", trt.float16, (1, WIDTH)) + weights = _weights() + if path == "graph_blocks": + y = graph_blocks.add_swiglu_mlp( + network, + x, + weights=weights, + prefix="mlp", + hidden_size=WIDTH, + mlp_size=WIDTH, + dtype=np.float16, + ) + else: + matmul = graph_blocks.make_matmul_fn(network, np.float16, None) + y = _swiglu_mlp( + network, + x, + matmul=matmul, + weights=weights, + prefix="mlp", + hidden=WIDTH, + mlp_size=WIDTH, + ) + y.name = "y" + network.mark_output(y) + config = builder.create_builder_config() + plan = builder.build_serialized_network(network, config) + assert plan is not None + runtime = trt.Runtime(logger) + return runtime, runtime.deserialize_cuda_engine(plan) + + +@pytest.mark.gpu +@pytest.mark.trt +@pytest.mark.parametrize("rebuild", range(3)) +@pytest.mark.parametrize("path", ["graph_blocks", "dual_profile"]) +def test_swiglu_rounds_the_product_once(path, rebuild): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + del rebuild # Each parametrized invocation builds a separate engine. + runtime, engine = _build(path) + assert engine is not None + context = engine.create_execution_context() + assert context is not None + x = torch.zeros((1, WIDTH), dtype=torch.float16, device="cuda") + x[0, 0] = 1.0 + y = torch.empty((1, WIDTH), dtype=torch.float16, device="cuda") + assert context.set_tensor_address("x", x.data_ptr()) + assert context.set_tensor_address("y", y.data_ptr()) + assert context.execute_async_v3(torch.cuda.current_stream().cuda_stream) + torch.cuda.synchronize() + actual = y.cpu().reshape(-1) + del context, engine, runtime + # Pinned Moshi's compiled gating kernel: Half projections, SiLU and the + # product in FP32, then one conversion to Half. + gate = torch.tensor(GATE, dtype=torch.float32, device="cuda") + up = torch.tensor(UP, dtype=torch.float32, device="cuda") + expected = torch.zeros(WIDTH, dtype=torch.float16) + expected[0] = (torch.nn.functional.silu(gate) * up).half().cpu() + torch.testing.assert_close(actual, expected, rtol=0, atol=0)