Class Linear

tensorrt_llm/layers/linear.py:313–436

class Linear(LinearBase):

    def __init__(
        self,
        in_features,
        out_features,
        bias=True,
        dtype=None,
        tp_group=None,
        tp_size=1,
        gather_output=True,
        share_weight=None,
        strict_dtype=False,
        pad_lda=0,
        pad_ldc=0,
        prefer_managed_weight=True,
        is_qkv=False,
    ):
        super().__init__(
            local_in_features=in_features,
            local_out_features=out_features // tp_size,
            bias=bias,
            dtype=dtype,
            tp_group=tp_group,
            tp_size=tp_size,
            share_weight=share_weight,
            strict_dtype=strict_dtype,
            pad_lda=pad_lda,
            pad_ldc=pad_ldc,
            prefer_managed_weight=prefer_managed_weight,
        )
        self.gather_output = gather_output
        self.is_qkv = is_qkv
        self.tp_dim = 0
        if bias:
            set_obj_attrs(
                self.bias,
                {
                    "weight_loader": self.weight_loader,
                },
            )

    @classmethod
    def tp_split_dim(cls) -> int:
        return 0

    def collect_and_bias(self, x, **kwargs):

        if self.bias is not None:
            bias = cast(self.bias.value, x.dtype)
            x = x + bias

        if self.gather_output and self.tp_size > 1 and self.tp_group is not None:
            # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size]
            x = allgather(x, self.tp_group, gather_dim=-1)

        return x