Class TestQuantize4BitFunctional

tests/test_functional.py:575–1034

class TestQuantize4BitFunctional:
    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype)
    @pytest.mark.parametrize("quant_type", ["fp4", "nf4"])
    @pytest.mark.parametrize(
        "blocksize",
        [32, 64, 128, 256, 512, 1024, 2048, 4096],
    )
    def test_4bit_quant(self, device, dtype, quant_type, blocksize):
        if device == "hpu" and not is_supported_on_hpu(quant_type, dtype):
            pytest.skip("This configuration is not supported on HPU.")

        A1 = torch.randn(1024, 1024, device=device, dtype=dtype)
        qa, SA = F.quantize_4bit(A1, blocksize=blocksize, quant_type=quant_type)
        d = SA.as_dict()
        SA = F.QuantState.from_dict(d, device=torch.device(device))
        A2 = F.dequantize_4bit(qa, SA, blocksize=blocksize, quant_type=quant_type)
        del qa, SA

        assert A2.dtype == dtype

        err = (A1 - A2).abs().float()
        del A2

        relerr = (err / (A1.abs().float() + 1e-8)).mean()
        err = err.mean()

        # Expected (mean, std) per configuration, from 200 samples on RTX 4090.
        # Thresholds are set at mean + N_SIGMA * std to avoid flaky failures
        # while still catching real regressions. Worst-case std across dtypes is used.
        N_SIGMA = 7
        error_stats = {
            "fp4": {
                "err": {
                    32: (0.088925, 0.000091),
                    64: (0.096543, 0.000111),
                    128: (0.102969, 0.000134),
                    256: (0.108684, 0.000182),
                    512: (0.114115, 0.000234),
                    1024: (0.119333, 0.000320),
                    2048: (0.124556, 0.000455),
                    4096: (0.129536, 0.000612),
                },
                "rel_err": {
                    32: (0.242443, 0.000330),
                    64: (0.260125, 0.000379),
                    128: (0.275817, 0.000433),
                    256: (0.289831, 0.000497),
                    512: (0.302881, 0.000583),
                    1024: (0.315000, 0.000757),
                    2048: (0.326607, 0.000955),
                    4096: (0.337169, 0.001239),
                },
            },
            "nf4": {
                "err": {
                    32: (0.067746, 0.000069),
                    64: (0.072798, 0.000074),