benchmark GPTQ implementations

gistfile1.txt

"""
Benchmark comparing eager vs compiled vs Triton GPTQ quantize_weight.

All three implementations are fully inlined:
  - quantize_weight_eager: the original loop from main
  - quantize_weight_compiled: the torch.compiled block path from this branch
  - quantize_weight_triton: hand-written Triton kernel for the block quantization

Generates random weights and a synthetic positive-definite Hessian,
then times all three.

Usage:
    python benchmarks/bench_gptq_quantize.py [--device cuda] [--rows 4096] [--cols 4096]
"""

import argparse
import math
import time
from copy import copy

import torch
import torch._dynamo.config
import triton
import triton.language as tl
from compressed_tensors.quantization import (
    QuantizationArgs,
    QuantizationScheme,
    QuantizationStrategy,
    fake_quantize,
)
from compressed_tensors.utils import patch_attr

from llmcompressor.modifiers.quantization.calibration import (
    initialize_observer,
    observe,
)

GPTQ_PRECISION = torch.float32
BLOCKSIZE = 128
PERCDAMP = 0.01
GROUP_SIZE = 128
NUM_BITS = 8
WARMUP = 2
ITERS = 5


# ── Shared helpers ───────────────────────────────────────────────────────────


def _prepare_weight_and_hessian(module, quant_args, hessian, blocksize, percdamp):
    """Common setup: clone weight, compute Hinv, build g_idx / scale / zp."""
    strategy = quant_args.strategy
    W = module.weight.clone().to(dtype=GPTQ_PRECISION)
    H = hessian.clone()
    num_rows = W.shape[0]
    num_columns = W.shape[1]

    observer = module.weight_observer

    g_idx = None
    if strategy in (
        QuantizationStrategy.GROUP,
        QuantizationStrategy.TENSOR_GROUP,
        QuantizationStrategy.BLOCK,
    ):
        divisor = (
            quant_args.group_size
            if strategy != QuantizationStrategy.BLOCK
            else quant_args.block_structure[1]
        )
        g_idx = (
            torch.arange(num_columns, device=W.device, dtype=torch.int) // divisor
        )

    qparams = observer.get_qparams()
    scale, zero_point, global_scale = (
        qparams["scale"],
        qparams["zero_point"],
        qparams["global_scale"],
    )

    losses = torch.zeros(num_rows, device=W.device)

    dead = torch.diag(H) == 0
    H[dead, dead] = 1
    W[:, dead] = 0

    damp = percdamp * torch.mean(torch.diag(H))
    diag = torch.arange(H.shape[0], device=H.device)
    H[diag, diag] += damp
    H = torch.linalg.cholesky(H)
    H = torch.cholesky_inverse(H)
    H = torch.linalg.cholesky(H, upper=True)
    Hinv = H

    return W, Hinv, losses, scale, zero_point, global_scale, g_idx


# ── Eager implementation (from main) ─────────────────────────────────────────


def quantize_weight_eager(
    module,
    quant_args,
    hessian,
    blocksize=128,
    percdamp=0.01,
):
    strategy = quant_args.strategy
    final_shape = module.weight.shape
    final_dtype = module.weight.dtype

    W, Hinv, losses, scale, zero_point, global_scale, g_idx = (
        _prepare_weight_and_hessian(module, quant_args, hessian, blocksize, percdamp)
    )
    num_columns = W.shape[1]

    for i1 in range(0, num_columns, blocksize):
        i2 = min(i1 + blocksize, num_columns)
        count = i2 - i1

        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]

        for i in range(count):
            w = W1[:, i]
            d = Hinv1[i, i]
            q = w.clone()

            if strategy == QuantizationStrategy.TENSOR:
                q = fake_quantize(
                    q, scale, zero_point, quant_args, global_scale=global_scale
                )
            elif strategy == QuantizationStrategy.CHANNEL:
                q = fake_quantize(
                    q,
                    scale[:, 0],
                    zero_point[:, 0],
                    quant_args,
                    global_scale=global_scale,
                )
            elif strategy in (
                QuantizationStrategy.GROUP,
                QuantizationStrategy.TENSOR_GROUP,
            ):
                column_idx = i1 + i
                group_index = g_idx[column_idx]
                altered_qargs = copy(quant_args)
                altered_qargs.strategy = QuantizationStrategy.CHANNEL
                q = fake_quantize(
                    q,
                    scale[:, group_index],
                    zero_point[:, group_index],
                    altered_qargs,
                    global_scale=global_scale,
                )
            elif strategy == QuantizationStrategy.BLOCK:
                column_idx = i1 + i
                block_column_idx = g_idx[column_idx]
                q = fake_quantize(
                    q.unsqueeze(1),
                    scale[:, block_column_idx : block_column_idx + 1],
                    zero_point[:, block_column_idx : block_column_idx + 1],
                    quant_args,
                    global_scale=global_scale,
                ).squeeze(1)
            else:
                raise ValueError(
                    f"Quantization strategy is not supported for GPTQ: {strategy}"
                )

            Q1[:, i] = q
            losses1[:, i] = (w - q) ** 2 / d**2

            err1 = (w - q) / d
            w1_err = err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
            W1[:, i:] -= w1_err
            Err1[:, i] = err1

        W[:, i1:i2] = Q1
        losses += torch.sum(losses1, 1) / 2
        W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

    W = W.reshape(final_shape).to(final_dtype)
    loss = torch.sum(losses).item()
    return loss, W


# ── Compiled implementation (from this branch) ──────────────────────────────


def _quantize_block(
    W1,
    Hinv1,
    scale_cols,
    zero_point_cols,
    quant_args,
    global_scale,
    count,
):
    Q1 = torch.zeros_like(W1)
    Err1 = torch.zeros_like(W1)
    losses1 = torch.zeros_like(W1)

    for i in range(count):
        w = W1[:, i]
        d = Hinv1[i, i]
        q = fake_quantize(
            w,
            scale_cols[:, i],
            zero_point_cols[:, i],
            quant_args,
            global_scale=global_scale,
        )
        Q1[:, i] = q
        losses1[:, i] = (w - q) ** 2 / d**2

        err1 = (w - q) / d
        w1_err = err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
        W1[:, i:] -= w1_err
        Err1[:, i] = err1

    return Q1, Err1, losses1


_quantize_block_compiled = torch.compile(_quantize_block, dynamic=True)


def quantize_weight_compiled(
    module,
    quant_args,
    hessian,
    blocksize=128,
    percdamp=0.01,
):
    strategy = quant_args.strategy
    final_shape = module.weight.shape
    final_dtype = module.weight.dtype

    W, Hinv, losses, scale, zero_point, global_scale, g_idx = (
        _prepare_weight_and_hessian(module, quant_args, hessian, blocksize, percdamp)
    )
    num_rows = W.shape[0]
    num_columns = W.shape[1]

    for i1 in range(0, num_columns, blocksize):
        i2 = min(i1 + blocksize, num_columns)
        count = i2 - i1

        W1 = W[:, i1:i2].clone()
        Hinv1 = Hinv[i1:i2, i1:i2]

        if strategy == QuantizationStrategy.TENSOR:
            scale_cols = scale.expand(num_rows, count)
            zero_point_cols = zero_point.expand(num_rows, count)
        elif strategy == QuantizationStrategy.CHANNEL:
            scale_cols = scale[:, :1].expand(-1, count)
            zero_point_cols = zero_point[:, :1].expand(-1, count)
        elif strategy in (
            QuantizationStrategy.GROUP,
            QuantizationStrategy.TENSOR_GROUP,
        ):
            cols = g_idx[i1:i2]
            scale_cols = scale[:, cols]
            zero_point_cols = zero_point[:, cols]
        elif strategy == QuantizationStrategy.BLOCK:
            row_block_size = quant_args.block_structure[0]
            row_idx = (
                torch.arange(num_rows, device=W.device, dtype=torch.int)
                // row_block_size
            )
            cols = g_idx[i1:i2]
            scale_cols = scale[row_idx][:, cols]
            zero_point_cols = zero_point[row_idx][:, cols]
        else:
            raise ValueError(
                f"Quantization strategy is not supported for GPTQ: {strategy}"
            )

        with (
            patch_attr(quant_args, "strategy", QuantizationStrategy.CHANNEL),
            torch._dynamo.config.patch(capture_scalar_outputs=True),
        ):
            Q1, Err1, losses1 = _quantize_block_compiled(
                W1,
                Hinv1.contiguous(),
                scale_cols,
                zero_point_cols,
                quant_args,
                global_scale,
                count,
            )

        W[:, i1:i2] = Q1
        losses += torch.sum(losses1, 1) / 2
        W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

    W = W.reshape(final_shape).to(final_dtype)
    loss = torch.sum(losses).item()
    return loss, W


# ── Triton implementation ────────────────────────────────────────────────────


@triton.jit
def _quantize_block_triton_kernel(
    W1_ptr,
    Hinv1_ptr,
    scale_ptr,
    zp_ptr,
    Q1_ptr,
    Err1_ptr,
    losses1_ptr,
    num_rows,
    count,
    stride_w_row,
    stride_h_row,
    stride_s_row,
    q_min,
    q_max,
    BLOCK_N: tl.constexpr,
):
    # one program per row — keeps entire W1 row in registers
    row = tl.program_id(0)
    if row >= num_rows:
        return

    cols = tl.arange(0, BLOCK_N)
    cmask = cols < count

    # load full row into registers (never spills back to DRAM)
    w = tl.load(W1_ptr + row * stride_w_row + cols, mask=cmask, other=0.0)
    s = tl.load(scale_ptr + row * stride_s_row + cols, mask=cmask, other=1.0)
    zp = tl.load(zp_ptr + row * stride_s_row + cols, mask=cmask, other=0.0)

    q_out = tl.zeros([BLOCK_N], dtype=tl.float32)
    err_out = tl.zeros([BLOCK_N], dtype=tl.float32)
    loss_out = tl.zeros([BLOCK_N], dtype=tl.float32)

    for i in range(count):
        imask = cols == i

        # extract scalar w[i], s[i], zp[i] via masked reduction
        wi = tl.sum(tl.where(imask, w, 0.0))
        si = tl.sum(tl.where(imask, s, 0.0))
        zpi = tl.sum(tl.where(imask, zp, 0.0))
        d = tl.load(Hinv1_ptr + i * stride_h_row + i)

        # inline fake_quantize
        q_int = tl.extra.cuda.libdevice.nearbyint(wi / si) + zpi
        q_int = tl.minimum(tl.maximum(q_int, q_min), q_max)
        qi = (q_int - zpi) * si

        diff = wi - qi
        err = diff / d

        q_out = tl.where(imask, qi, q_out)
        err_out = tl.where(imask, err, err_out)
        loss_out = tl.where(imask, diff * diff / (d * d), loss_out)

        # error propagation: w[i:] -= err * Hinv1[i, i:]
        # Hinv is tiny (128x128) and stays in L2 across all programs
        h_row = tl.load(Hinv1_ptr + i * stride_h_row + cols, mask=cmask, other=0.0)
        w = tl.where(cols >= i, w - err * h_row, w)

    tl.store(Q1_ptr + row * stride_w_row + cols, q_out, mask=cmask)
    tl.store(Err1_ptr + row * stride_w_row + cols, err_out, mask=cmask)
    tl.store(losses1_ptr + row * stride_w_row + cols, loss_out, mask=cmask)


def _quantize_block_triton(
    W1, Hinv1, scale_cols, zero_point_cols, quant_args, global_scale, count
):
    assert global_scale is None, "Triton kernel does not support global_scale"
    num_rows = W1.shape[0]
    Q1 = torch.zeros_like(W1)
    Err1 = torch.zeros_like(W1)
    losses1 = torch.zeros_like(W1)

    Hinv1 = Hinv1.contiguous()
    scale_cols = scale_cols.contiguous()
    zero_point_cols = zero_point_cols.to(dtype=W1.dtype).contiguous()

    q_min = -(2 ** (quant_args.num_bits - 1))
    q_max = 2 ** (quant_args.num_bits - 1) - 1

    BLOCK_N = triton.next_power_of_2(count)
    grid = (num_rows,)

    _quantize_block_triton_kernel[grid](
        W1,
        Hinv1,
        scale_cols,
        zero_point_cols,
        Q1,
        Err1,
        losses1,
        num_rows,
        count,
        W1.stride(0),
        Hinv1.stride(0),
        scale_cols.stride(0),
        float(q_min),
        float(q_max),
        BLOCK_N=BLOCK_N,
    )

    return Q1, Err1, losses1


def quantize_weight_triton(
    module,
    quant_args,
    hessian,
    blocksize=128,
    percdamp=0.01,
):
    strategy = quant_args.strategy
    final_shape = module.weight.shape
    final_dtype = module.weight.dtype

    W, Hinv, losses, scale, zero_point, global_scale, g_idx = (
        _prepare_weight_and_hessian(module, quant_args, hessian, blocksize, percdamp)
    )
    num_rows = W.shape[0]
    num_columns = W.shape[1]

    for i1 in range(0, num_columns, blocksize):
        i2 = min(i1 + blocksize, num_columns)
        count = i2 - i1

        W1 = W[:, i1:i2].clone()
        Hinv1 = Hinv[i1:i2, i1:i2]

        if strategy == QuantizationStrategy.TENSOR:
            scale_cols = scale.expand(num_rows, count)
            zero_point_cols = zero_point.expand(num_rows, count)
        elif strategy == QuantizationStrategy.CHANNEL:
            scale_cols = scale[:, :1].expand(-1, count)
            zero_point_cols = zero_point[:, :1].expand(-1, count)
        elif strategy in (
            QuantizationStrategy.GROUP,
            QuantizationStrategy.TENSOR_GROUP,
        ):
            cols = g_idx[i1:i2]
            scale_cols = scale[:, cols]
            zero_point_cols = zero_point[:, cols]
        elif strategy == QuantizationStrategy.BLOCK:
            row_block_size = quant_args.block_structure[0]
            row_idx = (
                torch.arange(num_rows, device=W.device, dtype=torch.int)
                // row_block_size
            )
            cols = g_idx[i1:i2]
            scale_cols = scale[row_idx][:, cols]
            zero_point_cols = zero_point[row_idx][:, cols]
        else:
            raise ValueError(
                f"Quantization strategy is not supported for GPTQ: {strategy}"
            )

        Q1, Err1, losses1 = _quantize_block_triton(
            W1,
            Hinv1.contiguous(),
            scale_cols,
            zero_point_cols,
            quant_args,
            global_scale,
            count,
        )

        W[:, i1:i2] = Q1
        losses += torch.sum(losses1, 1) / 2
        W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

    W = W.reshape(final_shape).to(final_dtype)
    loss = torch.sum(losses).item()
    return loss, W


# ── Triton codebook implementation (format-agnostic) ────────────────────────


def build_codebook(scale, zero_point, quant_args, global_scale=None):
    """Build a 3D tensor of representable dequantized values.

    Returns shape (r // r_b, c // c_b, 2^num_bits).  For integer formats the
    entries are evenly spaced; for float formats (FP4, FP8, …) the caller
    would populate the irregular grid instead.
    """
    num_bits = quant_args.num_bits
    if quant_args.symmetric:
        q_min = -(2 ** (num_bits - 1))
        q_max = 2 ** (num_bits - 1) - 1
    else:
        q_min = 0
        q_max = 2 ** num_bits - 1

    int_vals = torch.arange(
        q_min, q_max + 1, device=scale.device, dtype=torch.float32
    )

    s = scale.to(dtype=torch.float32)
    zp = zero_point.to(dtype=torch.float32)
    while s.dim() < 2:
        s = s.unsqueeze(0)
        zp = zp.unsqueeze(0)

    # (num_row_blocks, num_col_blocks, num_codes)
    codebook = (int_vals - zp.unsqueeze(-1)) * s.unsqueeze(-1)

    if global_scale is not None:
        codebook = codebook * global_scale

    return codebook.contiguous()


@triton.jit
def _quantize_block_triton_lut_kernel(
    W1_ptr,
    Hinv1_ptr,
    codebook_ptr,
    Q1_ptr,
    Err1_ptr,
    losses1_ptr,
    num_rows,
    count,
    col_offset,
    stride_w_row,
    stride_h_row,
    stride_cb_row,
    stride_cb_col,
    r_b,
    c_b,
    num_codes,
    BLOCK_N: tl.constexpr,
    BLOCK_C: tl.constexpr,
):
    row = tl.program_id(0)
    if row >= num_rows:
        return

    qrow = row // r_b

    cols = tl.arange(0, BLOCK_N)
    cmask = cols < count
    code_idx = tl.arange(0, BLOCK_C)
    code_mask = code_idx < num_codes

    w = tl.load(W1_ptr + row * stride_w_row + cols, mask=cmask, other=0.0)

    q_out = tl.zeros([BLOCK_N], dtype=tl.float32)
    err_out = tl.zeros([BLOCK_N], dtype=tl.float32)
    loss_out = tl.zeros([BLOCK_N], dtype=tl.float32)

    for i in range(count):
        imask = cols == i
        wi = tl.sum(tl.where(imask, w, 0.0))
        d = tl.load(Hinv1_ptr + i * stride_h_row + i)

        qcol = (col_offset + i) // c_b

        cb = tl.load(
            codebook_ptr + qrow * stride_cb_row + qcol * stride_cb_col + code_idx,
            mask=code_mask,
            other=float("inf"),
        )

        # nearest-neighbor lookup
        dist = tl.abs(wi - cb)
        min_dist = tl.min(dist, axis=0)
        is_min = dist == min_dist
        qi = tl.sum(tl.where(is_min, cb, 0.0)) / tl.sum(is_min.to(tl.float32))

        diff = wi - qi
        err = diff / d

        q_out = tl.where(imask, qi, q_out)
        err_out = tl.where(imask, err, err_out)
        loss_out = tl.where(imask, diff * diff / (d * d), loss_out)

        h_row = tl.load(Hinv1_ptr + i * stride_h_row + cols, mask=cmask, other=0.0)
        w = tl.where(cols >= i, w - err * h_row, w)

    tl.store(Q1_ptr + row * stride_w_row + cols, q_out, mask=cmask)
    tl.store(Err1_ptr + row * stride_w_row + cols, err_out, mask=cmask)
    tl.store(losses1_ptr + row * stride_w_row + cols, loss_out, mask=cmask)


def _quantize_block_triton_lut(W1, Hinv1, codebook, count, col_offset, r_b, c_b):
    num_rows = W1.shape[0]
    num_codes = codebook.shape[2]
    Q1 = torch.zeros_like(W1)
    Err1 = torch.zeros_like(W1)
    losses1 = torch.zeros_like(W1)

    BLOCK_N = triton.next_power_of_2(count)
    BLOCK_C = triton.next_power_of_2(num_codes)
    grid = (num_rows,)

    _quantize_block_triton_lut_kernel[grid](
        W1,
        Hinv1,
        codebook,
        Q1,
        Err1,
        losses1,
        num_rows,
        count,
        col_offset,
        W1.stride(0),
        Hinv1.stride(0),
        codebook.stride(0),
        codebook.stride(1),
        r_b,
        c_b,
        num_codes,
        BLOCK_N=BLOCK_N,
        BLOCK_C=BLOCK_C,
    )

    return Q1, Err1, losses1


def quantize_weight_triton_lut(
    module,
    quant_args,
    hessian,
    blocksize=128,
    percdamp=0.01,
):
    strategy = quant_args.strategy
    final_shape = module.weight.shape
    final_dtype = module.weight.dtype

    W, Hinv, losses, scale, zero_point, global_scale, g_idx = (
        _prepare_weight_and_hessian(module, quant_args, hessian, blocksize, percdamp)
    )
    num_rows = W.shape[0]
    num_columns = W.shape[1]

    if strategy == QuantizationStrategy.TENSOR:
        r_b, c_b = num_rows, num_columns
    elif strategy == QuantizationStrategy.CHANNEL:
        r_b, c_b = 1, num_columns
    elif strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP):
        r_b, c_b = 1, quant_args.group_size
    elif strategy == QuantizationStrategy.BLOCK:
        r_b, c_b = quant_args.block_structure
    else:
        raise ValueError(f"Unsupported strategy: {strategy}")

    codebook = build_codebook(scale, zero_point, quant_args, global_scale)

    for i1 in range(0, num_columns, blocksize):
        i2 = min(i1 + blocksize, num_columns)
        count = i2 - i1

        W1 = W[:, i1:i2].clone()
        Hinv1 = Hinv[i1:i2, i1:i2]

        Q1, Err1, losses1 = _quantize_block_triton_lut(
            W1, Hinv1.contiguous(), codebook, count, i1, r_b, c_b
        )

        W[:, i1:i2] = Q1
        losses += torch.sum(losses1, 1) / 2
        W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

    W = W.reshape(final_shape).to(final_dtype)
    loss = torch.sum(losses).item()
    return loss, W


# ── Triton cutoff implementation (format-agnostic, branchless bin) ───────────


def build_cutoffs_and_codes(scale, zero_point, quant_args, global_scale=None):
    """Build sorted codebook and midpoint cutoffs from quant params.

    Returns:
        codes:   (r // r_b, c // c_b, 2^num_bits)      — sorted representable values
        cutoffs: (r // r_b, c // c_b, 2^num_bits - 1)   — midpoints between adjacent codes
    """
    codes = build_codebook(scale, zero_point, quant_args, global_scale)
    cutoffs = (codes[..., :-1] + codes[..., 1:]) * 0.5
    return codes, cutoffs


@triton.jit
def _quantize_block_triton_cutoff_kernel(
    W1_ptr,
    Hinv1_ptr,
    codes_ptr,
    cutoffs_ptr,
    Q1_ptr,
    Err1_ptr,
    losses1_ptr,
    num_rows,
    count,
    col_offset,
    stride_w_row,
    stride_h_row,
    stride_codes_row,
    stride_codes_col,
    stride_cut_row,
    stride_cut_col,
    r_b,
    c_b,
    num_codes,
    BLOCK_N: tl.constexpr,
    BLOCK_C: tl.constexpr,
):
    row = tl.program_id(0)
    if row >= num_rows:
        return

    qrow = row // r_b

    cols = tl.arange(0, BLOCK_N)
    cmask = cols < count
    cut_idx = tl.arange(0, BLOCK_C)
    num_cutoffs = num_codes - 1
    cut_mask = cut_idx < num_cutoffs

    w = tl.load(W1_ptr + row * stride_w_row + cols, mask=cmask, other=0.0)

    q_out = tl.zeros([BLOCK_N], dtype=tl.float32)
    err_out = tl.zeros([BLOCK_N], dtype=tl.float32)
    loss_out = tl.zeros([BLOCK_N], dtype=tl.float32)

    for i in range(count):
        imask = cols == i
        wi = tl.sum(tl.where(imask, w, 0.0))
        d = tl.load(Hinv1_ptr + i * stride_h_row + i)

        qcol = (col_offset + i) // c_b

        # load cutoffs for this quantization block
        cuts = tl.load(
            cutoffs_ptr + qrow * stride_cut_row + qcol * stride_cut_col + cut_idx,
            mask=cut_mask,
            other=float("inf"),
        )

        # bin index = count of cutoffs that wi exceeds
        bin_idx = tl.sum((wi >= cuts).to(tl.int32), axis=0)

        # load the corresponding codebook value
        qi = tl.load(
            codes_ptr + qrow * stride_codes_row + qcol * stride_codes_col + bin_idx
        )

        diff = wi - qi
        err = diff / d

        q_out = tl.where(imask, qi, q_out)
        err_out = tl.where(imask, err, err_out)
        loss_out = tl.where(imask, diff * diff / (d * d), loss_out)

        h_row = tl.load(Hinv1_ptr + i * stride_h_row + cols, mask=cmask, other=0.0)
        w = tl.where(cols >= i, w - err * h_row, w)

    tl.store(Q1_ptr + row * stride_w_row + cols, q_out, mask=cmask)
    tl.store(Err1_ptr + row * stride_w_row + cols, err_out, mask=cmask)
    tl.store(losses1_ptr + row * stride_w_row + cols, loss_out, mask=cmask)


def _quantize_block_triton_cutoff(
    W1, Hinv1, codes, cutoffs, count, col_offset, r_b, c_b
):
    num_rows = W1.shape[0]
    num_codes = codes.shape[2]
    Q1 = torch.zeros_like(W1)
    Err1 = torch.zeros_like(W1)
    losses1 = torch.zeros_like(W1)

    BLOCK_N = triton.next_power_of_2(count)
    BLOCK_C = triton.next_power_of_2(num_codes - 1) if num_codes > 1 else 1
    grid = (num_rows,)

    _quantize_block_triton_cutoff_kernel[grid](
        W1,
        Hinv1,
        codes,
        cutoffs,
        Q1,
        Err1,
        losses1,
        num_rows,
        count,
        col_offset,
        W1.stride(0),
        Hinv1.stride(0),
        codes.stride(0),
        codes.stride(1),
        cutoffs.stride(0),
        cutoffs.stride(1),
        r_b,
        c_b,
        num_codes,
        BLOCK_N=BLOCK_N,
        BLOCK_C=BLOCK_C,
    )

    return Q1, Err1, losses1


def quantize_weight_triton_cutoff(
    module,
    quant_args,
    hessian,
    blocksize=128,
    percdamp=0.01,
):
    strategy = quant_args.strategy
    final_shape = module.weight.shape
    final_dtype = module.weight.dtype

    W, Hinv, losses, scale, zero_point, global_scale, g_idx = (
        _prepare_weight_and_hessian(module, quant_args, hessian, blocksize, percdamp)
    )
    num_rows = W.shape[0]
    num_columns = W.shape[1]

    if strategy == QuantizationStrategy.TENSOR:
        r_b, c_b = num_rows, num_columns
    elif strategy == QuantizationStrategy.CHANNEL:
        r_b, c_b = 1, num_columns
    elif strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP):
        r_b, c_b = 1, quant_args.group_size
    elif strategy == QuantizationStrategy.BLOCK:
        r_b, c_b = quant_args.block_structure
    else:
        raise ValueError(f"Unsupported strategy: {strategy}")

    codes, cutoffs = build_cutoffs_and_codes(
        scale, zero_point, quant_args, global_scale
    )

    for i1 in range(0, num_columns, blocksize):
        i2 = min(i1 + blocksize, num_columns)
        count = i2 - i1

        W1 = W[:, i1:i2].clone()
        Hinv1 = Hinv[i1:i2, i1:i2]

        Q1, Err1, losses1 = _quantize_block_triton_cutoff(
            W1, Hinv1.contiguous(), codes, cutoffs, count, i1, r_b, c_b
        )

        W[:, i1:i2] = Q1
        losses += torch.sum(losses1, 1) / 2
        W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

    W = W.reshape(final_shape).to(final_dtype)
    loss = torch.sum(losses).item()
    return loss, W


# ── Triton fused full-loop implementation ────────────────────────────────────


@triton.jit
def _quantize_weight_fused_kernel(
    W_ptr,
    Hinv_ptr,
    codes_ptr,
    cutoffs_ptr,
    Q_ptr,
    losses_ptr,
    num_rows,
    num_columns,
    blocksize,
    stride_w_row,
    stride_h_row,
    stride_codes_row,
    stride_codes_col,
    stride_cut_row,
    stride_cut_col,
    r_b,
    c_b,
    num_codes,
    BLOCK_BS: tl.constexpr,
    BLOCK_C: tl.constexpr,
    BLOCK_T: tl.constexpr,
):
    row = tl.program_id(0)
    if row >= num_rows:
        return

    qrow = row // r_b
    bs_range = tl.arange(0, BLOCK_BS)
    cut_idx = tl.arange(0, BLOCK_C)
    num_cutoffs = num_codes - 1
    cut_mask = cut_idx < num_cutoffs
    tile_range = tl.arange(0, BLOCK_T)

    loss_acc = 0.0

    for i1 in range(0, num_columns, blocksize):
        i2 = i1 + blocksize
        if i2 > num_columns:
            i2 = num_columns
        count = i2 - i1

        cmask = bs_range < count

        # load current GPTQ block of W into registers
        w = tl.load(
            W_ptr + row * stride_w_row + i1 + bs_range, mask=cmask, other=0.0
        )
        err_vec = tl.zeros([BLOCK_BS], dtype=tl.float32)

        # ── inner column loop: quantize + within-block error propagation ──
        for i in range(BLOCK_BS):
            if i < count:
                imask = bs_range == i
                wi = tl.sum(tl.where(imask, w, 0.0))
                d = tl.load(Hinv_ptr + (i1 + i) * stride_h_row + (i1 + i))

                qcol = (i1 + i) // c_b
                cuts = tl.load(
                    cutoffs_ptr
                    + qrow * stride_cut_row
                    + qcol * stride_cut_col
                    + cut_idx,
                    mask=cut_mask,
                    other=float("inf"),
                )
                bin_idx = tl.sum((wi >= cuts).to(tl.int32), axis=0)
                qi = tl.load(
                    codes_ptr
                    + qrow * stride_codes_row
                    + qcol * stride_codes_col
                    + bin_idx
                )

                diff = wi - qi
                err = diff / d

                tl.store(Q_ptr + row * stride_w_row + i1 + i, qi)
                loss_acc += diff * diff / (d * d)

                err_vec = tl.where(imask, err, err_vec)

                h_row = tl.load(
                    Hinv_ptr + (i1 + i) * stride_h_row + i1 + bs_range,
                    mask=cmask,
                    other=0.0,
                )
                w = tl.where(bs_range >= i, w - err * h_row, w)

        # ── inter-block error propagation ──
        # W[row, i2:] -= err_vec @ Hinv[i1:i2, i2:]
        # vectorized: load Hinv as 2D tile, broadcast multiply + reduce
        for j in range(i2, num_columns, BLOCK_T):
            tmask = tile_range < (num_columns - j)

            w_tail = tl.load(
                W_ptr + row * stride_w_row + j + tile_range,
                mask=tmask,
                other=0.0,
            )

            hinv_block = tl.load(
                Hinv_ptr
                + (i1 + bs_range[:, None]) * stride_h_row
                + (j + tile_range[None, :]),
                mask=(bs_range[:, None] < count) & tmask[None, :],
                other=0.0,
            )
            w_tail -= tl.sum(err_vec[:, None] * hinv_block, axis=0)

            tl.store(
                W_ptr + row * stride_w_row + j + tile_range,
                w_tail,
                mask=tmask,
            )

    tl.store(losses_ptr + row, loss_acc * 0.5)


def quantize_weight_triton_fused(
    module,
    quant_args,
    hessian,
    blocksize=128,
    percdamp=0.01,
):
    strategy = quant_args.strategy
    final_shape = module.weight.shape
    final_dtype = module.weight.dtype

    W, Hinv, losses, scale, zero_point, global_scale, g_idx = (
        _prepare_weight_and_hessian(module, quant_args, hessian, blocksize, percdamp)
    )
    num_rows = W.shape[0]
    num_columns = W.shape[1]

    if strategy == QuantizationStrategy.TENSOR:
        r_b, c_b = num_rows, num_columns
    elif strategy == QuantizationStrategy.CHANNEL:
        r_b, c_b = 1, num_columns
    elif strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP):
        r_b, c_b = 1, quant_args.group_size
    elif strategy == QuantizationStrategy.BLOCK:
        r_b, c_b = quant_args.block_structure
    else:
        raise ValueError(f"Unsupported strategy: {strategy}")

    codes, cutoffs = build_cutoffs_and_codes(
        scale, zero_point, quant_args, global_scale
    )

    Q = torch.zeros_like(W)
    grid = (num_rows,)

    BLOCK_BS = triton.next_power_of_2(blocksize)
    BLOCK_C = triton.next_power_of_2(codes.shape[2] - 1) if codes.shape[2] > 1 else 1
    BLOCK_T = 128

    _quantize_weight_fused_kernel[grid](
        W,
        Hinv,
        codes,
        cutoffs,
        Q,
        losses,
        num_rows,
        num_columns,
        blocksize,
        W.stride(0),
        Hinv.stride(0),
        codes.stride(0),
        codes.stride(1),
        cutoffs.stride(0),
        cutoffs.stride(1),
        r_b,
        c_b,
        codes.shape[2],
        BLOCK_BS=BLOCK_BS,
        BLOCK_C=BLOCK_C,
        BLOCK_T=BLOCK_T,
    )

    Q = Q.reshape(final_shape).to(final_dtype)
    loss = torch.sum(losses).item()
    return loss, Q


# ── Benchmark harness ────────────────────────────────────────────────────────


def make_inputs(rows, cols, device):
    print("creating inputs")
    torch.manual_seed(42)
    module = torch.nn.Linear(cols, rows, bias=False, device=device)
    quant_args = QuantizationArgs(
        num_bits=NUM_BITS,
        symmetric=True,
        strategy=QuantizationStrategy.GROUP,
        group_size=GROUP_SIZE,
    )
    module.quantization_scheme = QuantizationScheme(
        targets=["Linear"], weights=quant_args
    )
    initialize_observer(module, "weight")
    observe(module, "weight")

    hessian = torch.zeros(cols, cols, device=device, dtype=GPTQ_PRECISION)
    A = torch.randn(cols, cols, device=device, dtype=GPTQ_PRECISION)
    hessian = A @ A.t()
    hessian += torch.eye(cols, device=device, dtype=GPTQ_PRECISION)

    return module, quant_args, hessian


def time_fn(fn, module, quant_args, hessian, warmup, iters):
    is_cuda = module.weight.is_cuda

    for _ in range(warmup):
        fn(module, quant_args, hessian.clone(), BLOCKSIZE, PERCDAMP)
        if is_cuda:
            torch.cuda.synchronize()

    if is_cuda:
        torch.cuda.synchronize()
        torch.cuda.reset_peak_memory_stats()
        mem_before = torch.cuda.max_memory_allocated()

    times = []
    for _ in range(iters):
        if is_cuda:
            torch.cuda.synchronize()
        t0 = time.perf_counter()
        fn(module, quant_args, hessian.clone(), BLOCKSIZE, PERCDAMP)
        if is_cuda:
            torch.cuda.synchronize()
        times.append(time.perf_counter() - t0)

    peak_delta_mb = 0.0
    if is_cuda:
        peak_delta_mb = (
            torch.cuda.max_memory_allocated() - mem_before
        ) / (1024 * 1024)

    return times, peak_delta_mb


def main():
    parser = argparse.ArgumentParser(description="Benchmark GPTQ quantize_weight")
    parser.add_argument(
        "--device", default="cuda" if torch.cuda.is_available() else "cpu"
    )
    parser.add_argument("--rows", type=int, default=4096*2)
    parser.add_argument("--cols", type=int, default=4096*2)
    # parser.add_argument("--rows", type=int, default=1024)
    # parser.add_argument("--cols", type=int, default=1024)
    parser.add_argument("--warmup", type=int, default=WARMUP)
    parser.add_argument("--iters", type=int, default=ITERS)
    args = parser.parse_args()

    print(f"Device:    {args.device}")
    print(f"Weight:    ({args.rows}, {args.cols})")
    print(f"Blocksize: {BLOCKSIZE}")
    print(f"Percdamp:  {PERCDAMP}")
    print(f"Group:     {GROUP_SIZE}")
    print(f"Bits:      {NUM_BITS}")
    print(f"Warmup:    {args.warmup}  Iters: {args.iters}")
    print()

    module, quant_args, hessian = make_inputs(args.rows, args.cols, args.device)

    variants = [
        ("eager", quantize_weight_eager),
        ("compiled", quantize_weight_compiled),
        ("triton", quantize_weight_triton),
        ("triton_lut", quantize_weight_triton_lut),
        ("triton_cutoff", quantize_weight_triton_cutoff),
        ("triton_fused", quantize_weight_triton_fused),
    ]

    results = {}
    for name, fn in variants:
        print(f"Running {name} ...")
        times, peak_mb = time_fn(fn, module, quant_args, hessian, args.warmup, args.iters)
        results[name] = {"times": times, "peak_mb": peak_mb}

    medians = {n: sorted(r["times"])[len(r["times"]) // 2] for n, r in results.items()}

    print()
    print(f"{'':>16} {'median':>10} {'min':>10} {'max':>10} {'peak_mem':>10}")
    for name in results:
        t = results[name]["times"]
        med = medians[name]
        peak = results[name]["peak_mb"]
        mem_str = f"{peak:.1f}MB" if peak > 0 else "n/a"
        print(
            f"{name:>16} {med:>10.3f}s "
            f"{min(t):>10.3f}s {max(t):>10.3f}s "
            f"{mem_str:>10}"
        )

    print()
    base = medians["eager"]
    for name in list(results)[1:]:
        if medians[name] > 0:
            print(f"Speedup {name} vs eager: {base / medians[name]:.2f}x")


if __name__ == "__main__":
    with torch.no_grad():
        main()

                     median        min        max   peak_mem
           eager      2.191s      2.190s      2.196s   1282.6MB
        compiled      0.294s      0.294s      0.295s   1282.6MB
          triton      0.135s      0.135s      0.135s   1282.6MB (wont work for float quant)
      triton_lut      0.151s      0.151s      0.151s   1796.6MB
   triton_cutoff      0.129s      0.129s      0.129s   2302.6MB.  <<<
    triton_fused      0.562s      0.562s      0.562s   2302.6MB