From 9fbe70195f45a3ef71d5586d31238143867a5f1a Mon Sep 17 00:00:00 2001 From: Disty0 Date: Mon, 9 Mar 2026 23:05:08 +0300 Subject: [PATCH] SDNQ add 14, 12 and 10 bit support --- modules/sdnq/common.py | 56 ++++++ modules/sdnq/dequantizer.py | 10 +- modules/sdnq/layers/conv/conv_int8.py | 4 +- modules/sdnq/layers/linear/linear_int8.py | 4 +- modules/sdnq/packed_float.py | 13 +- modules/sdnq/packed_int.py | 204 ++++++++++++++++++---- modules/sdnq/quantizer.py | 9 +- 7 files changed, 244 insertions(+), 56 deletions(-) diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py index 8edc3bcf4..6e8951505 100644 --- a/modules/sdnq/common.py +++ b/modules/sdnq/common.py @@ -13,6 +13,9 @@ dtype_dict = { "int16": {"min": -32768, "max": 32767, "num_bits": 16, "sign": 1, "exponent": 0, "mantissa": 15, "target_dtype": torch.int16, "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": True, "is_packed": False}, "int8": {"min": -128, "max": 127, "num_bits": 8, "sign": 1, "exponent": 0, "mantissa": 7, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.int8, "is_unsigned": False, "is_integer": True, "is_packed": False}, ### Custom Integers + "int14": {"min": -8192, "max": 8191, "num_bits": 14, "sign": 1, "exponent": 0, "mantissa": 13, "target_dtype": "int14", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": True, "is_packed": True}, + "int12": {"min": -2048, "max": 2047, "num_bits": 12, "sign": 1, "exponent": 0, "mantissa": 11, "target_dtype": "int12", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": True, "is_packed": True}, + "int10": {"min": -512, "max": 511, "num_bits": 10, "sign": 1, "exponent": 0, "mantissa": 9, "target_dtype": "int10", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": True, "is_packed": True}, "int7": {"min": -64, "max": 63, "num_bits": 7, "sign": 1, "exponent": 0, "mantissa": 6, "target_dtype": "int7", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True}, "int6": {"min": -32, "max": 31, "num_bits": 6, "sign": 1, "exponent": 0, "mantissa": 5, "target_dtype": "int6", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True}, "int5": {"min": -16, "max": 15, "num_bits": 5, "sign": 1, "exponent": 0, "mantissa": 4, "target_dtype": "int5", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True}, @@ -24,6 +27,9 @@ dtype_dict = { "uint16": {"min": 0, "max": 65535, "num_bits": 16, "sign": 0, "exponent": 0, "mantissa": 16, "target_dtype": torch.uint16, "torch_dtype": torch.uint16, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": True, "is_packed": False}, "uint8": {"min": 0, "max": 255, "num_bits": 8, "sign": 0, "exponent": 0, "mantissa": 8, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": False}, ### Custom Unsigned Integers + "uint14": {"min": 0, "max": 16384, "num_bits": 14, "sign": 0, "exponent": 0, "mantissa": 14, "target_dtype": "uint14", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": True, "is_packed": True}, + "uint12": {"min": 0, "max": 4096, "num_bits": 12, "sign": 0, "exponent": 0, "mantissa": 12, "target_dtype": "uint12", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": True, "is_packed": True}, + "uint10": {"min": 0, "max": 1024, "num_bits": 10, "sign": 0, "exponent": 0, "mantissa": 10, "target_dtype": "uint10", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": True, "is_packed": True}, "uint7": {"min": 0, "max": 127, "num_bits": 7, "sign": 0, "exponent": 0, "mantissa": 7, "target_dtype": "uint7", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True}, "uint6": {"min": 0, "max": 63, "num_bits": 6, "sign": 0, "exponent": 0, "mantissa": 6, "target_dtype": "uint6", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True}, "uint5": {"min": 0, "max": 31, "num_bits": 5, "sign": 0, "exponent": 0, "mantissa": 5, "target_dtype": "uint5", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True}, @@ -43,11 +49,31 @@ dtype_dict = { "float16_e3m12fn": {"min": -31.99609375, "max": 31.99609375, "num_bits": 16, "sign": 1, "exponent": 3, "mantissa": 12, "min_normal": 0.125030517578125, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True}, "float16_e4m11fn": {"min": -511.875, "max": 511.875, "num_bits": 16, "sign": 1, "exponent": 4, "mantissa": 11, "min_normal": 0.007816314697265625, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True}, "float16_e5m10fn": {"min": -131008.0, "max": 131008.0, "num_bits": 16, "sign": 1, "exponent": 5, "mantissa": 10, "min_normal": 3.0547380447387695e-05, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + # + "float14_e1m12fn": {"min": -3.99951171875, "max": 3.99951171875, "num_bits": 14, "sign": 1, "exponent": 1, "mantissa": 12, "min_normal": 1.000244140625, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + "float14_e2m11fn": {"min": -7.998046875, "max": 7.998046875, "num_bits": 14, "sign": 1, "exponent": 2, "mantissa": 11, "min_normal": 0.500244140625, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + "float14_e3m10fn": {"min": -31.984375, "max": 31.984375, "num_bits": 14, "sign": 1, "exponent": 3, "mantissa": 10, "min_normal": 0.1251220703125, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + "float14_e4m9fn": {"min": -511.5, "max": 511.5, "num_bits": 14, "sign": 1, "exponent": 4, "mantissa": 9, "min_normal": 0.0078277587890625, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + "float14_e5m8fn": {"min": -130816.0, "max": 130816.0, "num_bits": 14, "sign": 1, "exponent": 5, "mantissa": 8, "min_normal": 3.063678741455078e-05, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + # + "float12_e1m10fn": {"min": -3.998046875, "max": 3.998046875, "num_bits": 12, "sign": 1, "exponent": 1, "mantissa": 10, "min_normal": 1.0009765625, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + "float12_e2m9fn": {"min": -7.9921875, "max": 7.9921875, "num_bits": 12, "sign": 1, "exponent": 2, "mantissa": 9, "min_normal": 0.5009765625, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + "float12_e3m8fn": {"min": -31.9375, "max": 31.9375, "num_bits": 12, "sign": 1, "exponent": 3, "mantissa": 8, "min_normal": 0.12548828125, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + "float12_e4m7fn": {"min": -510.0, "max": 510.0, "num_bits": 12, "sign": 1, "exponent": 4, "mantissa": 7, "min_normal": 0.00787353515625, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + "float12_e5m6fn": {"min": -130048.0, "max": 130048.0, "num_bits": 12, "sign": 1, "exponent": 5, "mantissa": 6, "min_normal": 3.0994415283203125e-05, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + # + "float10_e1m8fn": {"min": -3.9921875, "max": 3.9921875, "num_bits": 10, "sign": 1, "exponent": 1, "mantissa": 8, "min_normal": 1.00390625, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + "float10_e2m7fn": {"min": -7.96875, "max": 7.96875, "num_bits": 10, "sign": 1, "exponent": 2, "mantissa": 7, "min_normal": 0.50390625, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + "float10_e3m6fn": {"min": -31.75, "max": 31.75, "num_bits": 10, "sign": 1, "exponent": 3, "mantissa": 6, "min_normal": 0.126953125, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + "float10_e4m5fn": {"min": -504.0, "max": 504.0, "num_bits": 10, "sign": 1, "exponent": 4, "mantissa": 5, "min_normal": 0.008056640625, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + "float10_e5m4fn": {"min": -126976.0, "max": 126976.0, "num_bits": 10, "sign": 1, "exponent": 5, "mantissa": 4, "min_normal": 3.24249267578125e-05, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True}, + # "float8_e1m6fn": {"min": -3.96875, "max": 3.96875, "num_bits": 8, "sign": 1, "exponent": 1, "mantissa": 6, "min_normal": 1.015625, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True}, "float8_e2m5fn": {"min": -7.875, "max": 7.875, "num_bits": 8, "sign": 1, "exponent": 2, "mantissa": 5, "min_normal": 0.515625, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True}, "float8_e3m4fn": {"min": -31.0, "max": 31.0, "num_bits": 8, "sign": 1, "exponent": 3, "mantissa": 4, "min_normal": 0.1328125, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True}, "float8_e4m3fn_sdnq": {"min": -480.0, "max": 480.0, "num_bits": 8, "sign": 1, "exponent": 4, "mantissa": 3, "min_normal": 0.0087890625, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True}, "float8_e5m2fn": {"min": -114688.0, "max": 114688.0, "num_bits": 8, "sign": 1, "exponent": 5, "mantissa": 2, "min_normal": 3.814697265625e-05, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True}, + # "float7_e1m5fn": {"min": -3.9375, "max": 3.9375, "num_bits": 7, "sign": 1, "exponent": 1, "mantissa": 5, "min_normal": 1.03125, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True}, "float7_e2m4fn": {"min": -7.75, "max": 7.75, "num_bits": 7, "sign": 1, "exponent": 2, "mantissa": 4, "min_normal": 0.53125, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True}, "float7_e3m3fn": {"min": -30.0, "max": 30.0, "num_bits": 7, "sign": 1, "exponent": 3, "mantissa": 3, "min_normal": 0.140625, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True}, @@ -80,6 +106,24 @@ dtype_dict = { "float16_e4m12fnu": {"min": 0, "max": 511.9375, "num_bits": 16, "sign": 0, "exponent": 4, "mantissa": 12, "min_normal": 0.007814407348632812, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True}, "float16_e5m11fnu": {"min": 0, "max": 131040.0, "num_bits": 16, "sign": 0, "exponent": 5, "mantissa": 11, "min_normal": 3.053247928619385e-05, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True}, # + "float14_e1m13fnu": {"min": 0, "max": 3.999755859375, "num_bits": 14, "sign": 0, "exponent": 1, "mantissa": 13, "min_normal": 1.0001220703125, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + "float14_e2m12fnu": {"min": 0, "max": 7.9990234375, "num_bits": 14, "sign": 0, "exponent": 2, "mantissa": 12, "min_normal": 0.5001220703125, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + "float14_e3m11fnu": {"min": 0, "max": 31.9921875, "num_bits": 14, "sign": 0, "exponent": 3, "mantissa": 11, "min_normal": 0.12506103515625, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + "float14_e4m10fnu": {"min": 0, "max": 511.75, "num_bits": 14, "sign": 0, "exponent": 4, "mantissa": 10, "min_normal": 0.00782012939453125, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + "float14_e5m9fnu": {"min": 0, "max": 130944.0, "num_bits": 14, "sign": 0, "exponent": 5, "mantissa": 9, "min_normal": 3.057718276977539e-05, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + # + "float12_e1m11fnu": {"min": 0, "max": 3.9990234375, "num_bits": 12, "sign": 0, "exponent": 1, "mantissa": 11, "min_normal": 1.00048828125, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + "float12_e2m10fnu": {"min": 0, "max": 7.99609375, "num_bits": 12, "sign": 0, "exponent": 2, "mantissa": 10, "min_normal": 0.50048828125, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + "float12_e3m9fnu": {"min": 0, "max": 31.96875, "num_bits": 12, "sign": 0, "exponent": 3, "mantissa": 9, "min_normal": 0.125244140625, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + "float12_e4m8fnu": {"min": 0, "max": 511.0, "num_bits": 12, "sign": 0, "exponent": 4, "mantissa": 8, "min_normal": 0.007843017578125, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + "float12_e5m7fnu": {"min": 0, "max": 130560.0, "num_bits": 12, "sign": 0, "exponent": 5, "mantissa": 7, "min_normal": 3.075599670410156e-05, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + # + "float10_e1m9fnu": {"min": 0, "max": 3.99609375, "num_bits": 10, "sign": 0, "exponent": 1, "mantissa": 9, "min_normal": 1.001953125, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + "float10_e2m8fnu": {"min": 0, "max": 7.984375, "num_bits": 10, "sign": 0, "exponent": 2, "mantissa": 8, "min_normal": 0.501953125, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + "float10_e3m7fnu": {"min": 0, "max": 31.875, "num_bits": 10, "sign": 0, "exponent": 3, "mantissa": 7, "min_normal": 0.1259765625, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + "float10_e4m6fnu": {"min": 0, "max": 508.0, "num_bits": 10, "sign": 0, "exponent": 4, "mantissa": 6, "min_normal": 0.0079345703125, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + "float10_e5m5fnu": {"min": 0, "max": 129024.0, "num_bits": 10, "sign": 0, "exponent": 5, "mantissa": 5, "min_normal": 3.147125244140625e-05, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True}, + # "float8_e1m7fnu": {"min": 0, "max": 3.984375, "num_bits": 8, "sign": 0, "exponent": 1, "mantissa": 7, "min_normal": 1.0078125, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True}, "float8_e2m6fnu": {"min": 0, "max": 7.9375, "num_bits": 8, "sign": 0, "exponent": 2, "mantissa": 6, "min_normal": 0.5078125, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True}, "float8_e3m5fnu": {"min": 0, "max": 31.5, "num_bits": 8, "sign": 0, "exponent": 3, "mantissa": 5, "min_normal": 0.12890625, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True}, @@ -122,6 +166,9 @@ dtype_dict = { dtype_dict["fp32"] = dtype_dict["float32"] dtype_dict["bf16"] = dtype_dict["bfloat16"] dtype_dict["fp16"] = dtype_dict["float16"] +dtype_dict["fp14"] = dtype_dict["float14_e5m8fn"] +dtype_dict["fp12"] = dtype_dict["float12_e5m6fn"] +dtype_dict["fp10"] = dtype_dict["float10_e5m4fn"] dtype_dict["fp8"] = dtype_dict["float8_e4m3fn"] dtype_dict["fp7"] = dtype_dict["float7_e3m3fn"] dtype_dict["fp6"] = dtype_dict["float6_e3m2fn"] @@ -147,6 +194,9 @@ torch_dtype_dict = { torch.float8_e5m2: "float8_e5m2", } +if hasattr(torch, "float8_e8m0fnu"): + dtype_dict["float8_e8m0fnu"] = {"min": -1.70141e+38, "max": 1.70141e+38, "num_bits": 8, "sign": 1, "exponent": 8, "mantissa": 0, "target_dtype": "fp8", "torch_dtype": torch.float8_e8m0fnu, "storage_dtype": torch.float8_e8m0fnu, "is_unsigned": False, "is_integer": False, "is_packed": False} + torch_dtype_dict[torch.float8_e8m0fnu] = "float8_e8m0fnu" if hasattr(torch, "float8_e4m3fnuz"): dtype_dict["float8_e4m3fnuz"] = {"min": -240.0, "max": 240.0, "num_bits": 8, "sign": 1, "exponent": 4, "mantissa": 3, "target_dtype": "fp8", "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False, "is_packed": False} torch_dtype_dict[torch.float8_e4m3fnuz] = "float8_e4m3fnuz" @@ -178,8 +228,14 @@ weights_dtype_order = [ "uint7", "float7_e1m6fnu", "float7_e2m5fnu", "float7_e3m4fnu", "float7_e4m3fnu", "float7_e5m2fnu", "int8", "float8_e4m3fn", "float8_e5m2", "float8_e1m6fn", "float8_e2m5fn", "float8_e3m4fn", "float8_e4m3fn_sdnq", "float8_e5m2fn", "uint8", "float8_e1m7fnu", "float8_e2m6fnu", "float8_e3m5fnu", "float8_e4m4fnu", "float8_e5m3fnu", + "int10", "float10_e1m8fn", "float10_e2m7fn", "float10_e3m6fn", "float10_e4m5fn_sdnq", "float10_e5m4fn", + "uint10", "float10_e1m9fnu", "float10_e2m8fnu", "float10_e3m7fnu", "float10_e4m6fnu", "float10_e5m5fnu", + "int12", "float12_e1m10fn", "float12_e2m9fn", "float12_e3m8fn", "float12_e4m7fn_sdnq", "float12_e5m6fn", + "uint12", "float12_e1m11fnu", "float12_e2m10fnu", "float12_e3m9fnu", "float12_e4m8fnu", "float12_e5m7fnu", ] weights_dtype_order_fp32 = weights_dtype_order + [ + "int14", "float14_e1m12fn", "float14_e2m11fn", "float14_e3m10fn", "float14_e4m9fn_sdnq", "float14_e5m8fn", + "uint14", "float14_e1m13fnu", "float14_e2m12fnu", "float14_e3m11fnu", "float14_e4m10fnu", "float14_e5m9fnu", "int16", "float16", "float16_e1m14fn", "float16_e2m13fn", "float16_e3m12fn", "float16_e4m11fn", "float16_e5m10fn", "uint16", "float16_e1m15fnu", "float16_e2m14fnu", "float16_e3m13fnu", "float16_e4m12fnu", "float16_e5m11fnu", ] diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py index f9a51ce6f..dae9bdfea 100644 --- a/modules/sdnq/dequantizer.py +++ b/modules/sdnq/dequantizer.py @@ -6,7 +6,7 @@ import torch from modules import devices from .common import dtype_dict, compile_func, use_contiguous_mm, use_tensorwise_fp8_matmul -from .packed_int import unpack_int_symetric, unpack_int_asymetric +from .packed_int import unpack_int from .packed_float import unpack_float from .layers import SDNQLayer @@ -67,12 +67,12 @@ def dequantize_symmetric_with_bias(weight: torch.CharTensor, scale: torch.FloatT @devices.inference_context() def dequantize_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, weights_dtype: str, svd_up: torch.FloatTensor = None, svd_down: torch.FloatTensor = None, dtype: torch.dtype = None, result_shape: torch.Size = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor: - return dequantize_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, svd_up=svd_up, svd_down=svd_down, dtype=dtype, result_shape=result_shape, skip_quantized_matmul=skip_quantized_matmul) + return dequantize_asymmetric(unpack_int(weight, weights_dtype, shape), scale, zero_point, svd_up=svd_up, svd_down=svd_down, dtype=dtype, result_shape=result_shape, skip_quantized_matmul=skip_quantized_matmul) @devices.inference_context() def dequantize_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, weights_dtype: str, svd_up: torch.FloatTensor = None, svd_down: torch.FloatTensor = None, dtype: torch.dtype = None, result_shape: torch.Size = None, skip_quantized_matmul: bool = False, re_quantize_for_matmul: bool = False) -> torch.FloatTensor: - return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, svd_up=svd_up, svd_down=svd_down, dtype=dtype, result_shape=result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) + return dequantize_symmetric(unpack_int(weight, weights_dtype, shape, dtype=scale.dtype), scale, svd_up=svd_up, svd_down=svd_down, dtype=dtype, result_shape=result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) @devices.inference_context() @@ -159,12 +159,12 @@ def re_quantize_matmul_symmetric(weight: torch.CharTensor, scale: torch.FloatTen @devices.inference_context() def re_quantize_matmul_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, weights_dtype: str, matmul_dtype: str, result_shape: torch.Size, svd_up: torch.FloatTensor = None, svd_down: torch.FloatTensor = None) -> tuple[torch.Tensor, torch.FloatTensor]: - return re_quantize_matmul_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, matmul_dtype, svd_up=svd_up, svd_down=svd_down, result_shape=result_shape) + return re_quantize_matmul_asymmetric(unpack_int(weight, weights_dtype, shape), scale, zero_point, matmul_dtype, svd_up=svd_up, svd_down=svd_down, result_shape=result_shape) @devices.inference_context() def re_quantize_matmul_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, weights_dtype: str, matmul_dtype: str, result_shape: torch.Size = None, svd_up: torch.FloatTensor = None, svd_down: torch.FloatTensor = None) -> tuple[torch.Tensor, torch.FloatTensor]: - return re_quantize_matmul_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, matmul_dtype, svd_up=svd_up, svd_down=svd_down, result_shape=result_shape) + return re_quantize_matmul_symmetric(unpack_int(weight, weights_dtype, shape, dtype=scale.dtype), scale, matmul_dtype, svd_up=svd_up, svd_down=svd_down, result_shape=result_shape) @devices.inference_context() diff --git a/modules/sdnq/layers/conv/conv_int8.py b/modules/sdnq/layers/conv/conv_int8.py index 9e59ba542..df54830ef 100644 --- a/modules/sdnq/layers/conv/conv_int8.py +++ b/modules/sdnq/layers/conv/conv_int8.py @@ -3,8 +3,8 @@ import torch from ...common import compile_func, int_mm_func # noqa: TID252 -from ...packed_int import unpack_int_symetric # noqa: TID252 from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 +from ...packed_int import unpack_int # noqa: TID252 from .forward import get_conv_args, process_conv_input from ..linear.linear_int8 import quantize_int_mm_input # noqa: TID252 @@ -36,7 +36,7 @@ def conv_int8_matmul( bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up) if quantized_weight_shape is not None: - weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8).t_() + weight = unpack_int(weight, weights_dtype, quantized_weight_shape, dtype=torch.int8).t_() scale = scale.t() input, scale = quantize_int_mm_input(input, scale) input, weight = check_mats(input, weight) diff --git a/modules/sdnq/layers/linear/linear_int8.py b/modules/sdnq/layers/linear/linear_int8.py index bedfc95c3..21eed8e10 100644 --- a/modules/sdnq/layers/linear/linear_int8.py +++ b/modules/sdnq/layers/linear/linear_int8.py @@ -3,8 +3,8 @@ import torch from ...common import compile_func, int_mm_func # noqa: TID252 -from ...packed_int import unpack_int_symetric # noqa: TID252 from ...dequantizer import quantize_int_mm, dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 +from ...packed_int import unpack_int # noqa: TID252 from .forward import check_mats @@ -29,7 +29,7 @@ def int8_matmul( weights_dtype: str = None, ) -> torch.FloatTensor: if quantized_weight_shape is not None: - weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8).t_() + weight = unpack_int(weight, weights_dtype, quantized_weight_shape, dtype=torch.int8).t_() scale = scale.t() return_dtype = input.dtype output_shape = (*input.shape[:-1], weight.shape[-1]) diff --git a/modules/sdnq/packed_float.py b/modules/sdnq/packed_float.py index 665c21187..66b03c873 100644 --- a/modules/sdnq/packed_float.py +++ b/modules/sdnq/packed_float.py @@ -1,7 +1,7 @@ import torch from .common import dtype_dict -from .packed_int import pack_int_asymetric, unpack_int_asymetric +from .packed_int import pack_int, unpack_int float_bits_to_uint_dict = { @@ -12,6 +12,9 @@ float_bits_to_uint_dict = { 5: "uint5", 6: "uint6", 7: "uint7", + 10: "uint10", + 12: "uint12", + 14: "uint14", } @@ -51,8 +54,8 @@ def pack_float(x: torch.FloatTensor, weights_dtype: str) -> torch.Tensor: ~(-(1 << total_bits)), ).view(torch.uint32) - if total_bits < 8: - x = pack_int_asymetric(x, float_bits_to_uint_dict[total_bits]) + if total_bits not in {8, 16}: + x = pack_int(x, float_bits_to_uint_dict[total_bits]) else: x = x.to(dtype=dtype_dict[weights_dtype]["storage_dtype"]) @@ -72,8 +75,8 @@ def unpack_float(x: torch.Tensor, shape: torch.Size, weights_dtype: str) -> torc mantissa_difference = 23 - mantissa_bits exponent_difference = 8 - exponent_bits - if total_bits < 8: - x = unpack_int_asymetric(x, shape, float_bits_to_uint_dict[total_bits]) + if total_bits not in {8, 16}: + x = unpack_int(x, float_bits_to_uint_dict[total_bits], shape) x = x.to(dtype=torch.uint32).view(torch.int32) x = torch.bitwise_left_shift( diff --git a/modules/sdnq/packed_int.py b/modules/sdnq/packed_int.py index 2322a4e2f..f5e92b53f 100644 --- a/modules/sdnq/packed_int.py +++ b/modules/sdnq/packed_int.py @@ -5,22 +5,84 @@ import torch from .common import dtype_dict -def pack_int_symetric(tensor: torch.CharTensor, weights_dtype: str) -> torch.ByteTensor: - return packed_int_function_dict[weights_dtype]["pack"](tensor.sub_(dtype_dict[weights_dtype]["min"]).to(dtype=dtype_dict[weights_dtype]["storage_dtype"])) - - -def pack_int_asymetric(tensor: torch.CharTensor, weights_dtype: str) -> torch.ByteTensor: +def pack_int(tensor: torch.Tensor, weights_dtype: str) -> torch.Tensor: + if not dtype_dict[weights_dtype]["is_unsigned"]: + tensor = tensor.sub(dtype_dict[weights_dtype]["min"]) return packed_int_function_dict[weights_dtype]["pack"](tensor.to(dtype=dtype_dict[weights_dtype]["storage_dtype"])) -def unpack_int_symetric(packed_tensor: torch.ByteTensor, shape: torch.Size, weights_dtype: str, dtype: torch.dtype = None) -> torch.CharTensor: - if dtype is None: - dtype = dtype_dict[weights_dtype]["torch_dtype"] - return packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape).to(dtype=dtype).add_(dtype_dict[weights_dtype]["min"]) +def unpack_int(packed_tensor: torch.Tensor, weights_dtype: str, shape: torch.Size, dtype: torch.dtype = None) -> torch.Tensor: + packed_tensor = packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape) + if not dtype_dict[weights_dtype]["is_unsigned"]: + packed_tensor = packed_tensor.to(dtype=dtype_dict[weights_dtype]["torch_dtype"] if dtype is None else dtype).add_(dtype_dict[weights_dtype]["min"]) + return packed_tensor -def unpack_int_asymetric(packed_tensor: torch.ByteTensor, shape: torch.Size, weights_dtype: str) -> torch.CharTensor: - return packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape) +def pack_uint14(tensor: torch.Tensor) -> torch.Tensor: + packed_tensor = tensor.contiguous().view(-1, 8) + packed_tensor = torch.bitwise_or( + packed_tensor[:, :7], + torch.bitwise_and( + torch.stack( + ( + torch.bitwise_left_shift(packed_tensor[:, 7], 2), + torch.bitwise_left_shift(packed_tensor[:, 7], 4), + torch.bitwise_left_shift(packed_tensor[:, 7], 6), + torch.bitwise_left_shift(packed_tensor[:, 7], 8), + torch.bitwise_left_shift(packed_tensor[:, 7], 10), + torch.bitwise_left_shift(packed_tensor[:, 7], 12), + torch.bitwise_left_shift(packed_tensor[:, 7], 14), + ), + dim=-1 + ), + 49152 + ), + ) + return packed_tensor + + +def pack_uint12(tensor: torch.Tensor) -> torch.Tensor: + packed_tensor = tensor.contiguous().view(-1, 4) + packed_tensor = torch.bitwise_or( + packed_tensor[:, :3], + torch.bitwise_and( + torch.stack( + ( + torch.bitwise_left_shift(packed_tensor[:, 3], 4), + torch.bitwise_left_shift(packed_tensor[:, 3], 8), + torch.bitwise_left_shift(packed_tensor[:, 3], 12), + ), + dim=-1 + ), + 61440 + ) + ) + return packed_tensor + + +def pack_uint10(tensor: torch.ByteTensor) -> torch.ByteTensor: + packed_tensor = tensor.contiguous().view(-1, 8) + packed_tensor = torch.cat( + ( + torch.bitwise_or(packed_tensor[:, :3], torch.bitwise_left_shift(packed_tensor[:, 5:8], 10)), + torch.bitwise_or( + packed_tensor[:, 3], + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 5], 4), 15360), + torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 6), 49152), + ), + ).unsqueeze(-1), + torch.bitwise_or( + packed_tensor[:, 4], + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 4), 15360), + torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 8), 49152), + ), + ).unsqueeze(-1), + ), + dim=-1 + ) + return packed_tensor def pack_uint7(tensor: torch.ByteTensor) -> torch.ByteTensor: @@ -48,24 +110,19 @@ def pack_uint7(tensor: torch.ByteTensor) -> torch.ByteTensor: def pack_uint6(tensor: torch.ByteTensor) -> torch.ByteTensor: packed_tensor = tensor.contiguous().view(-1, 4) - packed_tensor = torch.cat( - ( - torch.bitwise_or( - packed_tensor[:, :2], - torch.bitwise_and( - torch.stack( - ( - torch.bitwise_left_shift(packed_tensor[:, 3], 2), - torch.bitwise_left_shift(packed_tensor[:, 3], 4), - ), - dim=-1 - ), - 192 - ) + packed_tensor = torch.bitwise_or( + packed_tensor[:, :3], + torch.bitwise_and( + torch.stack( + ( + torch.bitwise_left_shift(packed_tensor[:, 3], 2), + torch.bitwise_left_shift(packed_tensor[:, 3], 4), + torch.bitwise_left_shift(packed_tensor[:, 3], 6), + ), + dim=-1 ), - torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 6)).unsqueeze(-1), - ), - dim=-1 + 192 + ) ) return packed_tensor @@ -143,6 +200,74 @@ def pack_uint1(tensor: torch.Tensor) -> torch.Tensor: return packed_tensor +def unpack_uint14(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: + result = torch.cat( + ( + torch.bitwise_and(packed_tensor[:, :7], 16383), + torch.bitwise_or( + torch.bitwise_or( + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 2), 12288), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 4), 3072), + ), + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 6), 768), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 8), 192), + ), + ), + torch.bitwise_or( + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 10), 48), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 5], 12), 12), + ), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 6], 14), 3), + ), + ).unsqueeze(-1) + ), + dim=-1 + ).view(shape) + return result + + +def unpack_uint12(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: + result = torch.cat( + ( + torch.bitwise_and(packed_tensor[:, :3], 4095), + torch.bitwise_or( + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 4), 3840), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 8), 240), + ), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 12), 15) + ).unsqueeze(-1) + ), + dim=-1 + ).view(shape) + return result + + +def unpack_uint10(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: + result_bitwise_right_shift = torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, :3], 10), 63) + result = torch.cat( + ( + torch.bitwise_and(packed_tensor[:, :5], 1023), + torch.bitwise_or( + result_bitwise_right_shift[:, :2], + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3:5], 4), 960), + ), + torch.bitwise_or( + result_bitwise_right_shift[:, 2], + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 6), 768), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 8), 192), + ), + ).unsqueeze(-1), + ), + dim=-1 + ).view(shape) + return result + + def unpack_uint7(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor: result = torch.cat( ( @@ -175,7 +300,7 @@ def unpack_uint7(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By def unpack_uint6(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor: result = torch.cat( ( - torch.bitwise_and(packed_tensor[:, 0:3], 63), + torch.bitwise_and(packed_tensor[:, :3], 63), torch.bitwise_or( torch.bitwise_or( torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 2), 48), @@ -280,12 +405,9 @@ def unpack_uint1(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor packed_int_function_dict = { - "int7": {"pack": pack_uint7, "unpack": unpack_uint7}, - "int6": {"pack": pack_uint6, "unpack": unpack_uint6}, - "int5": {"pack": pack_uint5, "unpack": unpack_uint5}, - "int4": {"pack": pack_uint4, "unpack": unpack_uint4}, - "int3": {"pack": pack_uint3, "unpack": unpack_uint3}, - "int2": {"pack": pack_uint2, "unpack": unpack_uint2}, + "uint14": {"pack": pack_uint14, "unpack": unpack_uint14}, + "uint12": {"pack": pack_uint12, "unpack": unpack_uint12}, + "uint10": {"pack": pack_uint10, "unpack": unpack_uint10}, "uint7": {"pack": pack_uint7, "unpack": unpack_uint7}, "uint6": {"pack": pack_uint6, "unpack": unpack_uint6}, "uint5": {"pack": pack_uint5, "unpack": unpack_uint5}, @@ -293,5 +415,15 @@ packed_int_function_dict = { "uint3": {"pack": pack_uint3, "unpack": unpack_uint3}, "uint2": {"pack": pack_uint2, "unpack": unpack_uint2}, "uint1": {"pack": pack_uint1, "unpack": unpack_uint1}, - "bool": {"pack": pack_uint1, "unpack": unpack_uint1}, } + +packed_int_function_dict["int14"] = packed_int_function_dict["uint14"] +packed_int_function_dict["int12"] = packed_int_function_dict["uint12"] +packed_int_function_dict["int10"] = packed_int_function_dict["uint10"] +packed_int_function_dict["int7"] = packed_int_function_dict["uint7"] +packed_int_function_dict["int6"] = packed_int_function_dict["uint6"] +packed_int_function_dict["int5"] = packed_int_function_dict["uint5"] +packed_int_function_dict["int4"] = packed_int_function_dict["uint4"] +packed_int_function_dict["int3"] = packed_int_function_dict["uint3"] +packed_int_function_dict["int2"] = packed_int_function_dict["uint2"] +packed_int_function_dict["bool"] = packed_int_function_dict["uint1"] diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py index 3dd2b9976..61f17f84f 100644 --- a/modules/sdnq/quantizer.py +++ b/modules/sdnq/quantizer.py @@ -17,7 +17,7 @@ from accelerate import init_empty_weights from modules import devices, shared from .common import sdnq_version, dtype_dict, common_skip_keys, module_skip_keys_dict, accepted_weight_dtypes, accepted_matmul_dtypes, weights_dtype_order, weights_dtype_order_fp32, allowed_types, linear_types, conv_types, conv_transpose_types, compile_func, use_tensorwise_fp8_matmul, use_contiguous_mm, check_torch_compile from .dequantizer import SDNQDequantizer, dequantize_sdnq_model -from .packed_int import pack_int_symetric, pack_int_asymetric +from .packed_int import pack_int from .packed_float import pack_float from .forward import get_forward_func from .layers import get_sdnq_wrapper_class @@ -298,7 +298,7 @@ def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int if ( not dequantize_fp32 - and dtype_dict[weights_dtype]["num_bits"] <= 8 + and dtype_dict[weights_dtype]["max"] <= 16384 # 1/fp16_min_normal and not ( use_quantized_matmul and not dtype_dict[quantized_matmul_dtype]["is_integer"] @@ -403,10 +403,7 @@ def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int if dtype_dict[weights_dtype]["is_packed"]: if dtype_dict[weights_dtype]["is_integer"]: - if dtype_dict[weights_dtype]["is_unsigned"]: - weight = pack_int_asymetric(weight, weights_dtype) - else: - weight = pack_int_symetric(weight, weights_dtype) + weight = pack_int(weight, weights_dtype) else: weight = pack_float(weight, weights_dtype) else: