Source code for embedl_deploy._internal.tensorrt.modules.linear

# Copyright (C) 2026 Embedl AB

"""Fused ``nn.Module`` replacements for Linear and LayerNorm patterns."""

from typing import TypeAlias

import torch
import torch.nn.functional as F
from torch import nn

from embedl_deploy._internal.core.modules import (
    ActivationLike,
    FusedModule,
    PrecisionSynced,
)
from embedl_deploy._internal.core.quantize.config import Precision
from embedl_deploy._internal.core.quantize.stubs import (
    SmoothQuantObserver,
    WeightFakeQuantize,
)




class FusedLinear(FusedModule):
    """Fused wrapper for a standalone ``Linear`` layer.

    :param linear:
        The ``nn.Linear`` from the matched chain.
    """

    inputs_to_quantize: set[int] = {0}

    def __init__(self, linear: nn.Linear) -> None:
        super().__init__()
        self.linear = linear
        self.weight_fake_quant = WeightFakeQuantize({self})
        self.precision_deferred.upstream = True

    @property
    def quantized_weight(self) -> torch.Tensor | None:
        return self.linear.weight



    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply ``linear``, fake-quantizing the weight."""
        weight = self.weight_fake_quant(self.linear.weight)
        return F.linear(x, weight, self.linear.bias)


    def __repr__(self) -> str:  # pragma: no cover
        return (
            f"FusedLinear("
            f"{self.linear.in_features}→{self.linear.out_features})"
        )





class FusedLinearAct(FusedModule):
    """Fused ``Linear → Activation``.

    :param linear:
        The ``nn.Linear`` from the matched chain.
    :param act:
        The activation module from the matched chain.
    """

    inputs_to_quantize: set[int] = {0}

    def __init__(self, linear: nn.Linear, act: ActivationLike) -> None:
        super().__init__()
        self.linear = linear
        self.act = act
        self.weight_fake_quant = WeightFakeQuantize({self})
        self.precision_deferred.upstream = True
        if isinstance(act, nn.SiLU):
            self.precision_synced = PrecisionSynced(output=False)

    @property
    def quantized_weight(self) -> torch.Tensor | None:
        return self.linear.weight



    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply ``linear → activation``, fake-quantizing the weight."""
        weight = self.weight_fake_quant(self.linear.weight)
        x = F.linear(x, weight, self.linear.bias)
        return self.act(x)


    def __repr__(self) -> str:  # pragma: no cover
        act_name = type(self.act).__name__
        return (
            f"FusedLinearAct("
            f"{self.linear.in_features}→{self.linear.out_features}, "
            f"act={act_name})"
        )



FusedLinearLike: TypeAlias = FusedLinear | FusedLinearAct




class FusedLayerNorm(FusedModule):
    """Fused wrapper for a standalone ``LayerNorm``.

    Weight quantization is disabled by default. LayerNorm's learnable
    ``weight`` is an element-wise affine scale, so quantizing it yields
    negligible savings while hurting accuracy.

    :param layer_norm:
        The ``nn.LayerNorm`` from the matched chain.
    """

    inputs_to_quantize: set[int] = set()

    def __init__(self, layer_norm: nn.LayerNorm) -> None:
        super().__init__()
        self.layer_norm = layer_norm
        self.precision_synced = PrecisionSynced(output=False)
        self.output_precision = Precision.INT8
        self.smooth_quant_observer = SmoothQuantObserver(
            consumers={self},
            layer_norm=layer_norm,
        )



    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply ``layer_norm``."""
        return self.layer_norm(x)


    def __repr__(self) -> str:  # pragma: no cover
        return (
            f"FusedLayerNorm("
            f"normalized_shape={self.layer_norm.normalized_shape}, "
            f"eps={self.layer_norm.eps})"
        )