llmcompressor.modifiers.gptq.gptq_quantize

def fused_gptq_block_update(
    work: torch.Tensor,
    hinv: torch.Tensor,
    scale: torch.Tensor,
    zero_point: torch.Tensor | None,
    quantized: torch.Tensor,
    errors: torch.Tensor,
    q_min: float,
    q_max: float,
    quant_type: int,
) -> None:
    """Run the fused Triton GPTQ block update on FP32 working tensors."""
    if not HAS_TRITON:
        raise RuntimeError("Triton is unavailable")
    if (
        work.dim() != 3
        or work.device.type != "cuda"
        or work.dtype != torch.float32
        or hinv.dtype != torch.float32
        or quantized.dtype != torch.float32
        or errors.dtype != torch.float32
        or scale.dtype not in (torch.float16, torch.bfloat16, torch.float32)
        or any(
            tensor.device != work.device for tensor in (hinv, scale, quantized, errors)
        )
        or (zero_point is not None and zero_point.device != work.device)
    ):
        raise ValueError("invalid tensors for fused GPTQ block update")

    batch, out_rows, width = work.shape
    if (
        width <= 0
        or width > 256
        or width & (width - 1)
        or hinv.shape != (batch, width, width)
        or scale.shape != work.shape
        or quantized.shape != work.shape
        or errors.shape != work.shape
    ):
        raise ValueError("invalid shapes for fused GPTQ block update")
    if zero_point is not None and zero_point.shape != work.shape:
        raise ValueError("zero_point must have the same shape as work")

    dequant_dtype = {
        torch.float32: 0,
        torch.bfloat16: 1,
        torch.float16: 2,
    }[scale.dtype]
    has_zp = zero_point is not None
    use_fp8_e4b15 = quant_type == 2 and not _is_fp8_supported(work.device)
    if has_zp:
        zero_point = zero_point.to(torch.float32)

    block_rows = 16
    _gptq_block_update_kernel[(batch, triton.cdiv(out_rows, block_rows))](
        work,
        hinv,
        scale,
        zero_point if has_zp else scale,
        quantized,
        errors,
        out_rows,
        *work.stride(),
        *hinv.stride(),
        *scale.stride(),
        *(zero_point.stride() if has_zp else (0, 0, 0)),
        *quantized.stride(),
        *errors.stride(),
        float(q_min),
        float(q_max),
        WIDTH=width,
        QUANT_TYPE=quant_type,
        DEQUANT_DTYPE=dequant_dtype,
        HAS_ZP=has_zp,
        USE_FP8_E4B15=use_fp8_e4b15,
        BLOCK_ROWS=block_rows,
        num_warps=4,
        num_stages=2,
    )

@ImplBackend.entrypoint("gptq_block_update")
def gptq_block_update(
    W1: torch.Tensor,
    Hinv1: torch.Tensor,
    Q1: torch.Tensor,
    Err1: torch.Tensor,
    losses1: torch.Tensor,
    *,
    scale: torch.Tensor,
    zero_point: torch.Tensor | None,
    global_scale: torch.Tensor | None,
    g_idx: torch.Tensor | None,
    quant_args: QuantizationArgs,
    i1: int,
) -> None:
    """Run one GPTQ block with the eager Torch implementation."""
    if W1.dim() != 3:
        raise ValueError("The eager GPTQ block backend requires a 3D weight block")

    block_width = W1.shape[-1]
    altered_qargs = copy(quant_args)
    altered_qargs.strategy = QuantizationStrategy.CHANNEL
    eff, zp = column_scale_window(
        scale,
        zero_point,
        global_scale,
        g_idx,
        quant_args,
        num_rows=W1.shape[-2],
        i1=i1,
        i2=i1 + block_width,
    )
    for i in range(block_width):
        w = W1[:, :, i]
        # Flatten batch into CT's column dimension, allowing one QDQ call while
        # retaining an independent scale for every batch item and output row.
        q = fake_quantize(
            w.transpose(0, 1),
            eff[:, :, i].transpose(0, 1),
            None if zp is None else zp[:, :, i].transpose(0, 1),
            altered_qargs,
        ).transpose(0, 1)

        diagonal = Hinv1[:, i, i]
        error = (w - q) / diagonal[:, None]
        Q1[:, :, i] = q
        Err1[:, :, i] = error
        losses1[:, :, i] = error.square()
        W1[:, :, i:] -= error.unsqueeze(-1) * Hinv1[:, i, i:].unsqueeze(1)

def quantize_weight(
    weights: torch.Tensor,
    hessians: torch.Tensor,
    scale: torch.Tensor,
    zero_point: torch.Tensor,
    global_scale: torch.Tensor | None,
    quant_args: QuantizationArgs,
    blocksize: int = 128,
    percdamp: float = 0.01,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Quantize a batch of weights according to the GPTQ algorithm.

    :param weights: weights with shape [batch, rows, columns]
    :param hessians: normalized Hessians with shape [batch, columns, columns]
    :param scale: stacked observer scales
    :param zero_point: stacked observer zero points
    :param global_scale: optional stacked observer global scales
    :param quant_args: quantization arguments used to find quantization parameters
    :param blocksize: chunk size of quantization updates
    :param percdamp: dampening factor on hessian diagonal
    :return: quantized weights, per-batch losses, and RTN fallback flags
    """
    if weights.ndim != 3 or hessians.ndim != 3:
        raise ValueError("weights and hessians must have shape [batch, ...]")
    if weights.shape[0] != hessians.shape[0]:
        raise ValueError("weights and hessians must have matching batch sizes")

    batch_size, num_rows, num_columns = weights.shape
    strategy = quant_args.strategy
    actorder = quant_args.actorder
    final_dtype = weights.dtype
    device = weights.device
    # The caller provides a disposable stacked weight tensor, so use it as the
    # working buffer when it is already FP32 instead of allocating another copy.
    W = weights.to(device=device, dtype=GPTQ_PRECISION)
    # The stacked Hessian is the disposable working buffer.
    H = hessians.to(device=device, dtype=GPTQ_PRECISION)
    del weights, hessians
    scale = scale.to(device=device)
    zero_point = zero_point.to(device=device)
    if global_scale is not None:
        global_scale = global_scale.to(device=device)

    W, H, perm = apply_activation_ordering(W, H, quant_args.actorder)
    # handle g_idx
    g_idx = None
    if strategy in (
        QuantizationStrategy.GROUP,
        QuantizationStrategy.TENSOR_GROUP,
        QuantizationStrategy.BLOCK,
    ):
        # mapping from column index to group index
        divisor = (
            quant_args.group_size
            if strategy != QuantizationStrategy.BLOCK
            else quant_args.block_structure[1]
        )
        g_idx = torch.arange(num_columns, device=device, dtype=torch.int) // divisor

        if actorder == ActivationOrdering.WEIGHT:
            g_idx = torch.gather(g_idx.unsqueeze(0).expand(batch_size, -1), 1, perm)

    losses = torch.zeros(batch_size, num_rows, device=device)
    used_rtn_fallback = torch.zeros(batch_size, dtype=torch.bool, device=device)
    Hinv = factorize_hessian(W, H, percdamp, used_rtn_fallback)
    # See section 3.4 of https://arxiv.org/abs/2203.07259
    for i1 in range(0, num_columns, blocksize):
        i2 = min(i1 + blocksize, num_columns)

        W1 = W[:, :, i1:i2].clone()
        Q1 = torch.zeros_like(W1)
        Err1 = torch.zeros_like(W1)
        losses1 = torch.zeros_like(W1)
        Hinv1 = Hinv[:, i1:i2, i1:i2]

        gptq_block_update(
            W1,
            Hinv1,
            Q1,
            Err1,
            losses1,
            scale=scale,
            zero_point=zero_point,
            global_scale=global_scale,
            g_idx=g_idx,
            quant_args=quant_args,
            i1=i1,
        )

        # propagate block error
        W[:, :, i1:i2] = Q1
        losses += torch.sum(losses1, 2) / 2

        w_err = torch.bmm(Err1, Hinv[:, i1:i2, i2:])
        W[:, :, i2:] -= w_err

    if perm is not None:
        # restore original permutation
        # Release block-local tensors before allocating the restored output.
        del W1, Q1, Err1, losses1, Hinv1, w_err
        invperm = torch.argsort(perm, dim=-1)
        W = torch.gather(W, -1, invperm.unsqueeze(-2).expand(-1, num_rows, -1))

    return W.to(final_dtype), losses.sum(dim=1), used_rtn_fallback