Function embedding_4bit

exir/passes/_quant_patterns_and_replacements.py:425–460

(
    weight: torch.Tensor,
    weight_scales: torch.Tensor,
    weight_zero_points: Optional[torch.Tensor],
    weight_quant_min: int,
    weight_quant_max: int,
    indices: torch.Tensor,
)

@impl(quantized_decomposed_lib, "embedding_4bit", "CompositeExplicitAutograd")
def embedding_4bit(
    weight: torch.Tensor,
    weight_scales: torch.Tensor,
    weight_zero_points: Optional[torch.Tensor],
    weight_quant_min: int,
    weight_quant_max: int,
    indices: torch.Tensor,
) -> torch.Tensor:
    assert (
        weight_quant_min == -8
    ), "embedding_4bit in ExecuTorch expects weight_quant_min == -8"
    assert (
        weight_quant_max == 7
    ), "embedding_4bit in ExecuTorch expects weight_quant_max == 7"

    embedding_weight_checks(weight, weight_scales, weight_zero_points)
    group_size = (2 * weight.size(1)) // (
        weight_scales.size(1) if weight_scales.dim() == 2 else 1
    )
    weight_even = weight.div(16, rounding_mode="trunc")
    weight_odd = weight.remainder(16)
    weight_unpacked = torch.stack((weight_even, weight_odd), dim=-1)
    weight = weight_unpacked.view(weight.shape[0], -1)
    weight = weight.view(torch.int8).add(-8)

    weight = torch.ops.quantized_decomposed.dequantize_per_channel_group.default(
        weight,
        weight_scales,
        weight_zero_points,
        weight_quant_min,
        weight_quant_max,
        weight.dtype,
        group_size,
        weight_scales.dtype,
    )
    return torch.ops.aten.embedding.default(weight, indices)


@register_fake("quantized_decomposed::embedding_4bit")