Skip to content

Commit 8d7d5f3

Browse files
MarioSiegMarioSieg
authored
Merge pull request #28 from PrimeIntellect-ai/int2
Int2
2 parents fc64491 + e6eed19 commit 8d7d5f3

33 files changed

Lines changed: 1858 additions & 1869 deletions

‎benchmark/bench.cpp‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,18 +16,18 @@ auto main() -> int {
1616
const std::size_t nt {std::max(1u, std::thread::hardware_concurrency())};
1717
volatile std::size_t numel {1024*1024*1024/4};
1818
std::vector<float> data_in {};
19-
std::vector<std::uint8_t> data_out {};
19+
std::vector<piquant::uint4_t> data_out {};
2020
data_in.resize(numel);
21-
data_out.resize(numel);
21+
data_out.resize((numel+1)/2);
2222
std::random_device rd {};
2323
std::mt19937 gen {rd()};
2424
std::uniform_real_distribution dist {-1.0f, 1.0f};
2525
std::ranges::generate(data_in, [&] { return dist(gen); });
2626
ankerl::nanobench::Bench bench {};
2727
piquant::context ctx {nt};
28-
ctx.quantize_generic<float, std::uint8_t>(data_in, data_out, 0.2f, 127, piquant::round_mode::nearest);
29-
bench.run("dequantize", [&] {
30-
ctx.dequantize_generic<std::uint8_t, float>(data_out, data_in, 0.2f, 127, piquant::reduce_op::add);
28+
bench.run("requantize", [&] {
29+
ctx.quantize_generic<float, piquant::uint4_t>(data_in, data_out, 0.2f, 127, piquant::round_mode::nearest);
30+
ctx.dequantize_generic<piquant::uint4_t, float>(data_out, data_in, 0.2f, 127, piquant::reduce_op::add);
3131
});
3232
return 0;
3333
}

‎include/piquant.h‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,11 @@ typedef enum piquant_reduce_op_t {
2828
PIQUANT_REDUCE_OP_ADD, /* output[i] += quantize(input[i]) */
2929
} piquant_reduce_op_t;
3030

31-
typedef enum piquant_dtype_t {
31+
typedef enum piquant_dtype_t { // Order must match dtype enum class in piquant.hpp
32+
PIQUANT_DTYPE_F32 = 0,
33+
PIQUANT_DTYPE_F64,
34+
PIQUANT_DTYPE_UINT2,
35+
PIQUANT_DTYPE_INT2,
3236
PIQUANT_DTYPE_UINT4,
3337
PIQUANT_DTYPE_INT4,
3438
PIQUANT_DTYPE_UINT8,
@@ -39,8 +43,6 @@ typedef enum piquant_dtype_t {
3943
PIQUANT_DTYPE_INT32,
4044
PIQUANT_DTYPE_UINT64,
4145
PIQUANT_DTYPE_INT64,
42-
PIQUANT_DTYPE_F32,
43-
PIQUANT_DTYPE_F64
4446
} piquant_dtype_t;
4547

4648
extern PIQUANT_EXPORT piquant_context_t* piquant_context_create(size_t num_threads);
@@ -54,7 +56,7 @@ extern PIQUANT_EXPORT void piquant_quantize(
5456
piquant_dtype_t dtype_out,
5557
size_t numel,
5658
float scale,
57-
int32_t zero_point,
59+
int64_t zero_point,
5860
piquant_round_mode_t mode
5961
);
6062

@@ -66,7 +68,7 @@ extern PIQUANT_EXPORT void piquant_dequantize(
6668
piquant_dtype_t dtype_out,
6769
size_t numel,
6870
float scale,
69-
int32_t zero_point,
71+
int64_t zero_point,
7072
piquant_reduce_op_t op
7173
);
7274

‎include/piquant.hpp‎

Lines changed: 90 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -17,23 +17,26 @@
1717
#endif
1818

1919
namespace piquant {
20-
// (All types except u/int4) Amount of standard deviations/sigmas (left and right of 0) to use for the quantization range
21-
static constexpr double stddev_scale {12.0};
22-
23-
// (u/int4 only) Amount of standard deviations/sigmas (left and right of 0) to use for the quantization range
24-
static constexpr double stddev_scale_int4 {2.7};
25-
2620
enum class round_mode {
2721
nearest,
28-
stochastic
22+
stochastic,
23+
24+
count_
2925
};
3026

3127
enum class reduce_op {
3228
set, // output[i] = dequantize(input[i])
3329
add, // output[i] += qdeuantize(input[i])
30+
31+
count_
3432
};
3533

34+
// All supported data types for quantization and dequantization. Order matters.
3635
enum class dtype {
36+
f32 = 0,
37+
f64,
38+
uint2,
39+
int2,
3740
uint4,
3841
int4,
3942
uint8,
@@ -44,46 +47,65 @@ namespace piquant {
4447
int32,
4548
uint64,
4649
int64,
47-
f32,
48-
f64,
4950

50-
num_
51+
count_
52+
};
53+
static_assert(static_cast<std::underlying_type_t<dtype>>(dtype::count_) <= 0xff);
54+
static_assert(static_cast<std::underlying_type_t<dtype>>(dtype::f32) == 0);
55+
static_assert(static_cast<std::underlying_type_t<dtype>>(dtype::f64) == 1);
56+
57+
constexpr std::size_t float_dtype_count {2}; // Exclude quantized types
58+
59+
struct uint2_t final {
60+
using packed_storage = std::uint8_t;
61+
packed_storage bits;
62+
63+
constexpr uint2_t() noexcept : bits{} {}
64+
constexpr uint2_t(int u8) noexcept : bits{static_cast<packed_storage>(u8)} {}
65+
constexpr auto operator == (uint2_t rhs) const noexcept -> bool { return bits == rhs.bits; }
66+
constexpr auto operator != (uint2_t rhs) const noexcept -> bool { return !(*this == rhs); }
67+
constexpr auto operator == (packed_storage rhs) const noexcept -> bool { return bits == rhs; }
68+
constexpr auto operator != (packed_storage rhs) const noexcept -> bool { return !(*this == rhs); }
69+
};
70+
71+
struct int2_t final {
72+
using packed_storage = std::int8_t;
73+
packed_storage bits;
74+
75+
constexpr int2_t() noexcept : bits{} {}
76+
constexpr int2_t(int u8) noexcept : bits{static_cast<packed_storage>(u8)} {}
77+
constexpr auto operator == (int2_t rhs) const noexcept -> bool { return bits == rhs.bits; }
78+
constexpr auto operator != (int2_t rhs) const noexcept -> bool { return !(*this == rhs); }
79+
constexpr auto operator == (packed_storage rhs) const noexcept -> bool { return bits == rhs; }
80+
constexpr auto operator != (packed_storage rhs) const noexcept -> bool { return !(*this == rhs); }
5181
};
52-
static_assert(static_cast<std::underlying_type_t<dtype>>(dtype::num_) <= 0xff);
5382

5483
struct uint4_t final {
55-
std::uint8_t u8;
56-
constexpr uint4_t() noexcept : u8 {} {}
57-
constexpr uint4_t(int u8) noexcept : u8 {static_cast<std::uint8_t>(u8)} {}
58-
constexpr auto operator == (std::uint8_t y) const noexcept -> bool { return this->u8 == y; }
59-
constexpr auto operator != (std::uint8_t y) const noexcept -> bool { return !(*this == y); }
60-
constexpr auto operator == (uint4_t y) const noexcept -> bool { return this->u8 == y.u8; }
61-
constexpr auto operator != (uint4_t y) const noexcept -> bool { return !(*this == y); }
62-
constexpr auto pack(std::uint8_t lo, std::uint8_t hi) noexcept -> void {
63-
u8 = static_cast<std::uint8_t>((lo & 15) | ((hi & 15) << 4));
64-
}
65-
[[nodiscard]] constexpr auto unpack() const noexcept -> std::array<std::uint8_t, 2> {
66-
return {static_cast<std::uint8_t>(u8 & 15), static_cast<std::uint8_t>(u8 >> 4)};
67-
}
84+
using packed_storage = std::uint8_t;
85+
packed_storage bits;
86+
87+
constexpr uint4_t() noexcept : bits {} {}
88+
constexpr uint4_t(int u8) noexcept : bits {static_cast<packed_storage>(u8)} {}
89+
constexpr auto operator == (uint4_t rhs) const noexcept -> bool { return bits == rhs.bits; }
90+
constexpr auto operator != (uint4_t rhs) const noexcept -> bool { return !(*this == rhs); }
91+
constexpr auto operator == (packed_storage rhs) const noexcept -> bool { return bits == rhs; }
92+
constexpr auto operator != (packed_storage rhs) const noexcept -> bool { return !(*this == rhs); }
6893
};
6994

7095
struct int4_t final {
71-
std::int8_t u8;
72-
constexpr int4_t() noexcept : u8 {} {}
73-
constexpr int4_t(int u8) noexcept : u8 {static_cast<std::int8_t>(u8)} {}
74-
constexpr auto operator == (std::int8_t u8) const noexcept -> bool { return this->u8 == u8; }
75-
constexpr auto operator != (std::int8_t y) const noexcept -> bool { return !(*this == y); }
76-
constexpr auto operator == (int4_t y) const noexcept -> bool { return this->u8 == y.u8; }
77-
constexpr auto operator != (int4_t y) const noexcept -> bool { return !(*this == y); }
78-
constexpr auto pack(std::int8_t lo, std::int8_t hi) noexcept -> void {
79-
u8 = static_cast<std::int8_t>((lo & 15) | ((hi & 15) << 4));
80-
}
81-
[[nodiscard]] constexpr auto unpack() const noexcept -> std::array<std::int8_t, 2> {
82-
constexpr auto snex4 {[](std::int8_t x) noexcept -> std::int8_t { return x & 0x8 ? static_cast<std::int8_t>(x|0xF0) : x; }};
83-
return {(snex4(u8 & 15)), (snex4(u8 >> 4))};
84-
}
96+
using packed_storage = std::int8_t;
97+
packed_storage bits;
98+
99+
constexpr int4_t() noexcept : bits {} {}
100+
constexpr int4_t(int u8) noexcept : bits {static_cast<packed_storage>(u8)} {}
101+
constexpr auto operator == (int4_t rhs) const noexcept -> bool { return bits == rhs.bits; }
102+
constexpr auto operator != (int4_t rhs) const noexcept -> bool { return !(*this == rhs); }
103+
constexpr auto operator == (packed_storage rhs) const noexcept -> bool { return bits == rhs; }
104+
constexpr auto operator != (packed_storage rhs) const noexcept -> bool { return !(*this == rhs); }
85105
};
86106

107+
static_assert(sizeof(uint2_t) == 1);
108+
static_assert(sizeof(int2_t) == 1);
87109
static_assert(sizeof(uint4_t) == 1);
88110
static_assert(sizeof(int4_t) == 1);
89111

@@ -106,6 +128,10 @@ namespace piquant {
106128
};
107129

108130
constexpr std::array dtype_infos {
131+
dtype_info{.name="f32", .stride=4, .bit_size=32, .flags=dtype_flags::is_float+dtype_flags::is_signed}, // f32
132+
dtype_info{.name="f64", .stride=8, .bit_size=64, .flags=dtype_flags::is_float+dtype_flags::is_signed}, // f64
133+
dtype_info{.name="uint2", .stride=1, .bit_size=2, .flags=dtype_flags::is_quant+dtype_flags::is_int+dtype_flags::is_packed}, // uint2
134+
dtype_info{.name="int2", .stride=1, .bit_size=2, .flags=dtype_flags::is_quant+dtype_flags::is_int+dtype_flags::is_packed+dtype_flags::is_signed}, // int2
109135
dtype_info{.name="uint4", .stride=1, .bit_size=4, .flags=dtype_flags::is_quant+dtype_flags::is_int+dtype_flags::is_packed}, // uint4
110136
dtype_info{.name="int4", .stride=1, .bit_size=4, .flags=dtype_flags::is_quant+dtype_flags::is_int+dtype_flags::is_packed+dtype_flags::is_signed}, // int4
111137
dtype_info{.name="uint8", .stride=1, .bit_size=8, .flags=dtype_flags::is_quant+dtype_flags::is_int}, // uint8
@@ -115,9 +141,7 @@ namespace piquant {
115141
dtype_info{.name="uint32", .stride=4, .bit_size=32, .flags=dtype_flags::is_quant+dtype_flags::is_int}, // uint32
116142
dtype_info{.name="int32", .stride=4, .bit_size=32, .flags=dtype_flags::is_quant+dtype_flags::is_int+dtype_flags::is_signed}, // int32
117143
dtype_info{.name="uint64", .stride=8, .bit_size=64, .flags=dtype_flags::is_quant+dtype_flags::is_int}, // uint64
118-
dtype_info{.name="int64", .stride=8, .bit_size=64, .flags=dtype_flags::is_quant+dtype_flags::is_int+dtype_flags::is_signed}, // int64
119-
dtype_info{.name="f32", .stride=4, .bit_size=32, .flags=dtype_flags::is_float+dtype_flags::is_signed}, // f32
120-
dtype_info{.name="f64", .stride=8, .bit_size=64, .flags=dtype_flags::is_float+dtype_flags::is_signed}, // f64
144+
dtype_info{.name="int64", .stride=8, .bit_size=64, .flags=dtype_flags::is_quant+dtype_flags::is_int+dtype_flags::is_signed} // int64
121145
};
122146
static_assert([]() -> bool {
123147
for (auto&& info : dtype_infos) {
@@ -132,19 +156,31 @@ namespace piquant {
132156
static constexpr T min{std::numeric_limits<T>::min()};
133157
static constexpr T max{std::numeric_limits<T>::max()};
134158
};
159+
template<> struct dtype_limits<uint2_t> final {
160+
static constexpr std::uint8_t min{0};
161+
static constexpr std::uint8_t max{3};
162+
};
163+
template<> struct dtype_limits<int2_t> final {
164+
static constexpr std::int8_t min{-2};
165+
static constexpr std::int8_t max{1};
166+
};
135167
template<> struct dtype_limits<uint4_t> final {
136168
static constexpr std::uint8_t min{0};
137-
static constexpr std::uint8_t max{0xf};
169+
static constexpr std::uint8_t max{15};
138170
};
139171
template<> struct dtype_limits<int4_t> final {
140-
static constexpr std::int8_t min{-0x8};
141-
static constexpr std::int8_t max{0x7};
172+
static constexpr std::int8_t min{-8};
173+
static constexpr std::int8_t max{7};
142174
};
143175

176+
template<typename T> concept is_int2 = std::is_same_v<T, uint2_t> || std::is_same_v<T, int2_t>;
144177
template<typename T> concept is_int4 = std::is_same_v<T, uint4_t> || std::is_same_v<T, int4_t>;
145-
template<typename T> concept is_dtype = std::is_arithmetic_v<T> || is_int4<T>;
178+
template<typename T> concept is_packed_int = is_int2<T> || is_int4<T>;
179+
template<typename T> concept is_dtype = std::is_arithmetic_v<T> || is_packed_int<T>;
146180
template<typename T> requires is_dtype<T> struct dtype_traits final {};
147181

182+
template<> struct dtype_traits<uint2_t> { static constexpr dtype type_code = dtype::uint2; };
183+
template<> struct dtype_traits<int2_t> { static constexpr dtype type_code = dtype::int2; };
148184
template<> struct dtype_traits<uint4_t> { static constexpr dtype type_code = dtype::uint4; };
149185
template<> struct dtype_traits<int4_t> { static constexpr dtype type_code = dtype::int4; };
150186
template<> struct dtype_traits<std::int8_t> { static constexpr dtype type_code = dtype::int8; };
@@ -283,16 +319,16 @@ namespace piquant {
283319
};
284320

285321
struct quant_descriptor final {
286-
command_type type{command_type::quant};
287-
const std::byte* in{};
288-
std::byte* out{};
289-
std::int64_t numel{};
322+
command_type type {command_type::quant};
323+
const std::byte* in {};
324+
std::byte* out {};
325+
std::int64_t numel {};
290326
float scale{};
291-
std::int64_t zero_point{};
292-
dtype dt_in{};
293-
dtype dt_out{};
294-
round_mode rnd_mode{};
295-
reduce_op reduce{};
327+
std::int64_t zero_point {};
328+
dtype dt_in {};
329+
dtype dt_out {};
330+
round_mode rounding {};
331+
reduce_op reducing {};
296332
};
297333

298334
private:

‎python/benchmark/benchmark.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,11 @@ def benchmark_fast_quant() -> None:
5757
print(f'Torch quantization time for {num_runs} runs: {time_torch:.6f} seconds')
5858
print(f'Fast quantization time for {num_runs} runs: {time_fast:.6f} seconds')
5959

60-
labels = ['torch.ao.quantization.fx._decomposed.quantize_per_tensor', 'torch.quantize_per_tensor', 'piquant.quantize_torch']
60+
labels = [
61+
'Torch FX Quantize',
62+
'Torch Quantize',
63+
'piquant',
64+
]
6165
times = [time_torch_fx, time_torch, time_fast]
6266

6367
plt.figure(figsize=(6, 4))

‎python/setup.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ def __init__(self, message):
1818

1919

2020
class CMakeBuildExtension(Extension):
21-
def __init__(self, name, root_dir = ''):
21+
def __init__(self, name, root_dir=''):
2222
super().__init__(name, sources=[])
2323
self.root_dir = os.path.abspath(root_dir)
2424

@@ -46,7 +46,7 @@ def build_extension(self, ext):
4646
os.makedirs(self.build_temp)
4747

4848
cmake_args = [
49-
'-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' + os.path.abspath(os.path.join(self.build_lib, "piquant")),
49+
'-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' + os.path.abspath(os.path.join(self.build_lib, 'piquant')),
5050
'-DCMAKE_BUILD_TYPE=Release',
5151
]
5252
build_args = [
@@ -62,7 +62,7 @@ def build_extension(self, ext):
6262
name='pypiquant',
6363
author='Mario Sieg',
6464
author_email='mario@primeintellect.ai',
65-
packages=['piquant'],
65+
packages=['piquant', 'piquant._compat'],
6666
package_dir={'': 'src'}, # tell setuptools packages are under src/
6767
package_data={
6868
'piquant': ['libquant.so', 'libquant.dylib', 'libquant.dll'],

‎python/src/piquant/__init__.py‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,5 @@
33
__email__ = 'mario.sieg.64@gmail.com'
44
__author_email__ = 'mario.sieg.64@gmail.com'
55

6-
from ._quant import *
7-
from ._torch import *
8-
from ._numpy import *
6+
from ._core import *
7+
from ._compat import *

0 commit comments

Comments
 (0)