gistfile1.txt

"""Minimal repro: packed integer zero-point failures for uint7/uint8.

The humming kernel handles integer ZPs correctly when tensors are
unpacked. But the packed path (packed=True in prepare_humming_weight
and prepare_humming_zero_point) — which is what vLLM uses when loading
pack-quantized checkpoints — produces wrong results for num_bits > 6.

This causes catastrophic accuracy loss for W7A16asym and W7A8asym
models (20M word perplexity vs ~10 for symmetric).
"""

import pytest
import torch
from humming import dtypes, ops
from humming.kernel.humming import HummingKernel
from humming.schema.humming import HummingWeightSchema
from humming.utils.weight import (
    prepare_humming_weight,
    prepare_humming_weight_scale,
    prepare_humming_zero_point,
    dequantize_weight,
)

N, K, M = 1024, 1024, 1


def _make_packed_data(b_dtype_str):
    """Quantize random weights into packed format with integer ZPs,
    matching what a pack-quantized checkpoint provides."""
    b_dtype = dtypes.DataType.from_str(b_dtype_str)
    schema = HummingWeightSchema(
        b_dtype=b_dtype,
        bs_dtype=dtypes.bfloat16,
        weight_scale_group_size=0,
        has_zero_point=True,
        is_fp_zero_point=False,
    )

    torch.manual_seed(42)
    fp_weight = torch.randn((N, K), dtype=torch.bfloat16, device="cuda:0")
    tensors = HummingWeightSchema.quant_tensor(fp_weight, schema, torch.bfloat16)

    return tensors["weight"], tensors["weight_scale"], tensors["zero_point"]


def run_packed_int_zp_gemm(b_dtype_str):
    """Run GEMM using the packed path (packed=True), which is the
    code path vLLM takes when loading pack-quantized checkpoints."""
    b_dtype = dtypes.DataType.from_str(b_dtype_str)
    a_dtype = dtypes.DataType.from_str("bfloat16")

    weight_packed, weight_scale, zp_packed = _make_packed_data(b_dtype_str)

    weight_ref = dequantize_weight(
        weight_packed, weight_scale=weight_scale, zero_point=zp_packed,
        global_scale=None, dtype=b_dtype, packed=True,
    )

    weight = prepare_humming_weight(
        weight_packed, b_dtype, a_dtype, zp_packed, packed=True,
    )
    ws = prepare_humming_weight_scale(weight_scale, to_apply_on_c=True)
    zp = prepare_humming_zero_point(zp_packed, b_dtype, packed=True)

    inputs = torch.ones((M, K), dtype=torch.bfloat16, device="cuda:0")

    kernel = HummingKernel(
        shape_n=N, shape_k=K,
        block_shape=(16, a_dtype.num_bits * 16, 512 // a_dtype.num_bits),
        warp_shape=(16, a_dtype.num_bits * 4, 512 // a_dtype.num_bits),
        a_dtype=a_dtype, b_dtype=b_dtype, c_dtype=a_dtype, bs_dtype=a_dtype,
        num_stages=3, use_warp_spec=False, use_tma=False, use_cp_async=False,
        has_zero_point=True, is_fp_zero_point=False,
        mma_type="mma", use_stream_k=False,
    )

    outputs = torch.zeros((M, N), dtype=torch.bfloat16, device=inputs.device)
    outputs = ops.launch_kernel(
        configs=[kernel.kernel_id], inputs=inputs, weight=weight,
        outputs=outputs, zero_point=zp, input_scale=None, weight_scale=ws,
    )

    outputs_ref = inputs.float().matmul(weight_ref.T).to(torch.bfloat16)
    return outputs, outputs_ref


@pytest.mark.parametrize("b_dtype_str", ["uint3", "uint4", "uint5", "uint7", "uint8"])
def test_packed_int_zero_point(b_dtype_str):
    """Packed integer ZP path — broken for uint7/uint8 (num_bits > 6)."""
    outputs, outputs_ref = run_packed_int_zp_gemm(b_dtype_str)
    torch.testing.assert_close(outputs, outputs_ref, rtol=0.03, atol=1.0)


if __name__ == "__main__":
    print("=" * 80)
    print("Packed integer zero-point GEMM repro (packed=True)")
    print("=" * 80)

    for b_dtype_str in ["uint3", "uint4", "uint5", "uint7", "uint8"]:
        try:
            outputs, outputs_ref = run_packed_int_zp_gemm(b_dtype_str)
            diff = (outputs.float() - outputs_ref.float()).abs()
            max_abs = diff.max().item()
            mean_abs = diff.mean().item()
            try:
                torch.testing.assert_close(outputs, outputs_ref, rtol=0.03, atol=1.0)
                status = "PASS"
            except AssertionError:
                status = "FAIL"
            print(f"  {b_dtype_str:5s}  {status}  max_abs={max_abs:10.2f}  mean_abs={mean_abs:.4f}")
        except Exception as e:
            print(f"  {b_dtype_str:5s}  ERROR  {e}")

    print("=" * 80)

================================================================================
Packed integer zero-point GEMM repro (packed=True)
================================================================================
  uint3  PASS  max_abs=      0.00  mean_abs=0.0000
  uint4  PASS  max_abs=      0.00  mean_abs=0.0000
  uint5  PASS  max_abs=      0.50  mean_abs=0.0088
  uint7  FAIL  max_abs=   2413.25  mean_abs=171.8042
  uint8  PASS  max_abs=      0.50  mean_abs=0.0351
================================================================================