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)