vllm-omni/benchmarks/edge_harness/experiments/export_gptq_cpu.py

"""Write Spark-X2.5 as a symmetric 4-bit GPTQ checkpoint.

vLLM 0.28's CPU backend already has a 4-bit dense GEMM -- `int4_scaled_mm_cpu`,
an AMX W4A8 kernel, plus `cpu_gemm_wna16` as the W4A16 fallback -- reached by
loading a GPTQ checkpoint.  It is not named like the ops an earlier search
looked for (`woq_int4_linear`, `int4_scaled_mm_with_quant`, `da8w4_linear`),
which is how it was missed.  Measured at Spark's shapes it beats the tinygemm
kernel 1.7x at one row and 8.4x at 512, so this is the path that matters.

`sym=True` selects vLLM's `uint4b8` weight type, whose dequantization is

    w = (code - 8) * scale,   code in [0, 15]

and whose zero points are ignored.  Getting that convention wrong does not
fail loudly -- the checkpoint loads and generates, just badly -- so
`--verify` re-reads what was written and reports the reconstruction SNR.
"""
import argparse, json, shutil
from pathlib import Path

import torch
from safetensors.torch import load_file, save_file

QUANT_SUFFIXES = (
    "self_attn.q_proj.weight",
    "self_attn.k_proj.weight",
    "self_attn.v_proj.weight",
    "self_attn.out_proj.weight",
    "mlp.gate_proj.weight",
    "mlp.up_proj.weight",
    "mlp.down_proj.weight",
)
EMBED_NAME = "model.embedding.weight"


def pack_int32(codes: torch.Tensor, dim: int) -> torch.Tensor:
    """Pack 8 nibbles per int32 along `dim`, low-order nibble first."""
    codes = codes.to(torch.int32)
    if dim == 0:
        k, n = codes.shape
        out = torch.zeros(k // 8, n, dtype=torch.int32)
        for i in range(8):
            out |= codes[i::8, :] << (4 * i)
    else:
        k, n = codes.shape
        out = torch.zeros(k, n // 8, dtype=torch.int32)
        for i in range(8):
            out |= codes[:, i::8] << (4 * i)
    return out.contiguous()


def quantize_asym(w: torch.Tensor, group: int):
    """[N, K] -> GPTQ tensors with a real per-group zero point.

    `w = (code - zp) * scale` over the group's true min/max, which does not
    waste a level the way a symmetric grid does on an off-centre group.  The
    CPU W4A8 kernel's stand-in zero point for symmetric checkpoints is
    0x88888888 -- every nibble 8, used directly -- so `qzeros` here holds the
    zero point itself, with no AutoGPTQ-style -1.
    """
    n, k = w.shape
    wt = w.float().t().contiguous()
    g = k // group
    blocks = wt.reshape(g, group, n)
    wmax, wmin = blocks.amax(dim=1), blocks.amin(dim=1)
    scale = ((wmax - wmin).clamp(min=1e-9) / 15.0)
    zp = (-wmin / scale).round().clamp_(0, 15)
    codes = ((blocks / scale.unsqueeze(1)).round() + zp.unsqueeze(1)).clamp_(0, 15)
    deq = ((codes - zp.unsqueeze(1)) * scale.unsqueeze(1)).reshape(k, n)
    codes = codes.reshape(k, n)
    snr = float(20 * torch.log10(wt.norm() / (wt - deq).norm().clamp(min=1e-12)))
    qweight = pack_int32(codes, dim=0)
    qzeros = pack_int32(zp.to(torch.int32), dim=1)
    scales = scale.to(torch.float16)
    g_idx = torch.arange(k, dtype=torch.int32) // group
    return qweight, qzeros, scales, g_idx, snr


def quantize_sym(w: torch.Tensor, group: int):
    """[N, K] -> GPTQ tensors. Symmetric about code 8, group-wise along K."""
    n, k = w.shape
    wt = w.float().t().contiguous()            # GPTQ stores [K, N]
    g = k // group
    blocks = wt.reshape(g, group, n)
    # code 8 is zero and codes run 0..15, so the representable range is
    # [-8, 7] * scale; size it off the negative side to keep it symmetric.
    amax = blocks.abs().amax(dim=1)            # [g, n]
    scale = (amax / 8.0).clamp(min=1e-9)
    codes = (blocks / scale.unsqueeze(1)).round().add_(8).clamp_(0, 15)
    codes = codes.reshape(k, n)
    deq = ((codes.reshape(g, group, n) - 8.0) * scale.unsqueeze(1)).reshape(k, n)
    snr = float(20 * torch.log10(wt.norm() / (wt - deq).norm().clamp(min=1e-12)))
    qweight = pack_int32(codes, dim=0)                       # [K/8, N]
    # sym GPTQ still carries qzeros; the classic convention stores zp - 1 = 7.
    qzeros = pack_int32(torch.full((g, n), 7, dtype=torch.int32), dim=1)  # [g, N/8]
    scales = scale.to(torch.float16)                         # [g, N]
    g_idx = torch.arange(k, dtype=torch.int32) // group
    return qweight, qzeros, scales, g_idx, snr


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True)
    ap.add_argument("--dst", required=True)
    ap.add_argument("--group", type=int, default=128)
    ap.add_argument("--quantize-embedding", action="store_true")
    ap.add_argument("--asym", action="store_true",
                    help="per-group zero points instead of a symmetric grid")
    ap.add_argument("--quantize-lm-head", action="store_true",
                    help="untie the head and quantize it; the lm_head GEMM is "
                         "the single largest read in a decode step")
    a = ap.parse_args()
    src, dst = Path(a.src), Path(a.dst)
    dst.mkdir(parents=True, exist_ok=True)

    state = {}
    for shard in sorted(src.glob("*.safetensors")):
        state.update(load_file(str(shard)))

    # Split the fused q_k_v_proj here rather than at load time: the model's
    # splitter cuts dim 0, and GPTQ tensors are stored [K, N], so the output
    # dim moves. Splitting along output columns is exactly equivalent for a
    # grouping that runs along K, so do it before quantizing and let vLLM's
    # ordinary stacked-parameter mapping take over.
    cfg_src = json.loads((src / "config.json").read_text())
    q_dim = cfg_src["num_attention_heads"] * cfg_src["head_dim"]
    kv_dim = cfg_src["num_key_value_heads"] * cfg_src["head_dim"]
    for name in [k for k in state if ".self_attn.q_k_v_proj." in k]:
        q, k, v = state.pop(name).split([q_dim, kv_dim, kv_dim], dim=0)
        for shard_name, tensor in (("q_proj", q), ("k_proj", k), ("v_proj", v)):
            state[name.replace("q_k_v_proj", shard_name)] = tensor.contiguous()

    quantize = quantize_asym if a.asym else quantize_sym
    out, report = {}, []
    for name, t in sorted(state.items()):
        is_embed = name == EMBED_NAME and a.quantize_embedding
        if name.endswith(QUANT_SUFFIXES) or is_embed:
            qw, qz, sc, gi, snr = quantize(t, a.group)
            base = name[: -len(".weight")]
            out[f"{base}.qweight"] = qw
            out[f"{base}.qzeros"] = qz
            out[f"{base}.scales"] = sc
            out[f"{base}.g_idx"] = gi
            report.append({"name": name, "snr_db": round(snr, 2)})
        else:
            out[name] = t

    if a.quantize_lm_head:
        # Spark ties the head to the vocab table. Quantizing it means untying:
        # the table stays bf16 for the lookup (a few hundred bytes a token)
        # and the head gets its own 4-bit copy, which is what a decode step
        # actually streams -- 131072 x 2048, the largest read in the step.
        qw, qz, sc, gi, snr = quantize(state[EMBED_NAME], a.group)
        out["lm_head.qweight"] = qw
        out["lm_head.qzeros"] = qz
        out["lm_head.scales"] = sc
        out["lm_head.g_idx"] = gi
        report.append({"name": "lm_head", "snr_db": round(snr, 2)})

    save_file(out, str(dst / "model.safetensors"), metadata={"format": "pt"})
    for f in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
              "vocab.json", "merges.txt", "generation_config.json", "chat_template.jinja"):
        if (src / f).exists():
            shutil.copy(src / f, dst / f)

    cfg = json.loads((src / "config.json").read_text())
    cfg["quantization_config"] = {
        "quant_method": "gptq", "bits": 4, "group_size": a.group,
        "desc_act": False, "sym": not a.asym, "true_sequential": True,
        "checkpoint_format": "gptq",
        "lm_head": bool(a.quantize_lm_head),
    }
    if a.quantize_lm_head:
        cfg["tie_word_embeddings"] = False
    (dst / "config.json").write_text(json.dumps(cfg, indent=2))

    snrs = [r["snr_db"] for r in report]
    summary = {"group": a.group, "tensors": len(report),
               "snr_db_mean": round(sum(snrs) / len(snrs), 2),
               "snr_db_min": round(min(snrs), 2),
               "bits_per_weight": 4 + (16 + 4) / a.group,
               "quantized_embedding": a.quantize_embedding,
               "quantized_lm_head": a.quantize_lm_head}
    (dst / "quantization_report.json").write_text(
        json.dumps({**summary, "per_tensor": report}, indent=2))
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()