Method __init__

tensorrt_llm/quantization/layers.py:899–931

(
            self,
            num_embeddings: int,
            embedding_dim: int,
            dtype: Optional[str] = None,
            tp_size: int = 1,
            tp_group: Optional[list] = None,
            sharding_dim: int = 0,
            tp_rank: Optional[int] = None,
            quant_mode=QuantMode.use_weight_only(),
    )

class WeightOnlyQuantEmbedding(Embedding):

    def __init__(
            self,
            num_embeddings: int,
            embedding_dim: int,
            dtype: Optional[str] = None,
            tp_size: int = 1,
            tp_group: Optional[list] = None,
            sharding_dim: int = 0,
            tp_rank: Optional[int] = None,
            quant_mode=QuantMode.use_weight_only(),
    ):
        super().__init__(
            num_embeddings,
            embedding_dim,
            dtype,  # dtype,
            tp_size,
            tp_group,
            sharding_dim,
            tp_rank)
        # only support int8 wo now
        # TODO support int4 wo
        self.quant_mode = quant_mode
        self.per_token_scale = Parameter(shape=(self.num_embeddings, ),
                                         dtype=dtype)

        if sharding_dim == 1:
            self.weight = Parameter(shape=(self.num_embeddings,
                                           self.embedding_dim // self.tp_size),
                                    dtype="int8")
        elif sharding_dim == 0:
            self.weight = Parameter(shape=(math.ceil(
                self.num_embeddings / self.tp_size), self.embedding_dim),
                                    dtype="int8")

    def forward(self, x):
        result = embedding(x,