Uint2 improvements and fixes

Author: MarioSieg

python/benchmark/benchmark.py

@@ -19,6 +19,7 @@
 QUANT_DTYPES_TO_BENCH: list[torch.dtype] = [
     torch.quint8,
     torch.quint4x2,
+    torch.quint2x4
 ]
 
 def quantize_torch(t: torch.Tensor, scale: float, zp: int, dtype: torch.dtype) -> torch.tensor:
@@ -63,7 +64,11 @@ def _bench_piquant() -> None:
         dq_piquant = piquant.torch.dequantize(results_piquant[i], scale=scale, zero_point=zp, dtype=torch.float32)
         assert dq_torch.numel() == dq_piquant.numel()
         assert dq_torch.dtype == dq_piquant.dtype
-        assert torch.allclose(dq_torch, dq_piquant, atol=1e-1)
+        if not torch.allclose(dq_torch, dq_piquant, atol=1e-1):
+            print(f"Results differ for dtype {torch_d} at run {i}")
+            for j in range(dq_torch.numel()):
+                if not torch.isclose(dq_torch[j], dq_piquant[j], atol=1e-1):
+                    print(f"  Index {j}: torch={dq_torch[j]}, piquant={dq_piquant[j]}")
     print(f'{dtype_labels[-1]:<10} | torch: {torch_time:.6f}s | piquant: {piquant_time:.6f}s')
 
 

python/benchmark/throughput_avg.py

@@ -9,7 +9,7 @@
 TOTAL_GIB = 4
 ITERATIONS = 10
 SRC_DTYPE = torch.bfloat16
-Q_DTYPE = torch.quint4x2
+Q_DTYPE = torch.quint2x4
 OUT_DTYPE = torch.bfloat16
 
 B_PER_ELEM = torch.tensor([], dtype=SRC_DTYPE).element_size()

python/quant_benchmark.png

@@ -69,11 +69,20 @@ static auto PIQUANT_HOT dequant_uint2(
             static_assert(ReduceOp == reduce_op::set || ReduceOp == reduce_op::add, "Invalid reduce operation");
         }
     }
-    if (numel & 3) { /* Handle 1-, 2- or 3-value tail */
-        auto p {x[i>>2].bits};
-        if (numel & 1) o[i] = dequant_step<In, Out>(scale, zp, In{p & 3});
-        if (numel & 2) o[i+1] = dequant_step<In, Out>(scale, zp, In{p>>2 & 3});
-        if (numel & 3) o[i+((numel & 3) == 3 ? 2 : 0)] = dequant_step<In, Out>(scale, zp, In{p>>4 & 3});
+    auto p {x[i>>2].bits};
+    switch (numel&3) {
+        case 1:
+            o[i] = dequant_step<In, Out>(scale, zp, p&3);
+            break;
+        case 2:
+            o[i] = dequant_step<In, Out>(scale, zp, p&3);
+            o[i+1] = dequant_step<In, Out>(scale, zp, (p>>2)&3);
+            break;
+        case 3:
+            o[i] = dequant_step<In, Out>(scale, zp, p&3);
+            o[i+1] = dequant_step<In, Out>(scale, zp, (p>>2)&3);
+            o[i+2] = dequant_step<In, Out>(scale, zp, (p>>4)&3);
+            break;
     }
 }
 

src/kernels/dequantize.inl

@@ -29,7 +29,7 @@ using namespace piquant;
             data_in.resize(numel); \
             quantized.resize(numel_out); \
             std::ranges::generate(data_in, [&] { return dist(gen); }); \
-            piquant::context ctx {std::max(1u, 4u)}; \
+            piquant::context ctx {4}; \
             auto [scale, zero_point] {ctx.compute_quant_config_from_data(data_in, dtype_traits<to>::type_code)}; \
             ctx.quantize_generic<ti, to>(data_in, quantized, scale, zero_point, piquant::round_mode::rnd); \
             std::vector<ti> dequantized {}; \

test/dequant.cpp

@@ -121,6 +121,67 @@ template <const std::uint8_t IDX, typename T>
         } \
     }
 
+#define test_quant_int2(ti, to, rnd, is_stochastic, is_signed) \
+    TEST(quantize, quantize_##ti##_to_##to##_##rnd) { \
+        std::mt19937 gen {0x9032002}; \
+        std::uniform_real_distribution<fp32_t> dist {-1.0, 1.0}; \
+        \
+        for (std::size_t n {}; n < iters; ++n) { \
+            std::cout << "Iteration " << n << std::endl; \
+            fp32_t scale {std::uniform_real_distribution<fp32_t>{0.1, 1.0}(gen)}; \
+            std::size_t numel {std::uniform_int_distribution<std::size_t>{5000, 1'5000}(gen)}; \
+            std::size_t numel_out {std::is_same_v<uint2_t, to> ? (numel+3)>>2 : numel}; \
+            std::int32_t zero_point {std::is_same_v<uint2_t, to> ? std::uniform_int_distribution<std::int32_t>{-1, 2}(gen) : \
+                    std::uniform_int_distribution<std::int32_t>{-128, 127}(gen)}; \
+            \
+            std::vector<ti> data_in {}; \
+            std::vector<to> data_out_naive {}; \
+            std::vector<to> data_out {}; \
+            data_in.resize(numel); \
+            data_out.resize(numel_out); \
+            data_out_naive.resize(numel_out); \
+            std::ranges::generate(data_in, [&] { return dist(gen); }); \
+            quantize_naive<ti, to, piquant::round_mode::rnd>(data_in, data_out_naive, scale, zero_point); \
+            piquant::context ctx {std::max(1u, std::thread::hardware_concurrency())}; \
+            ctx.quantize_generic<ti, to>(data_in, data_out, scale, zero_point, piquant::round_mode::rnd); \
+            for (std::size_t i {}; i < numel_out; ++i) { \
+                bool eq {eq = data_out[i] == data_out_naive[i]}; \
+                std::int32_t a {unpack_nibble<0>(data_out[i], is_signed)}; \
+                std::int32_t b {unpack_nibble<1>(data_out[i], is_signed)}; \
+                std::int32_t a_naive {unpack_nibble<0>(data_out_naive[i], is_signed)}; \
+                std::int32_t b_naive {unpack_nibble<1>(data_out_naive[i], is_signed)}; \
+                if (is_stochastic) { \
+                    eq |= std::abs(a - a_naive) <= stochastic_epsilon; \
+                    eq |= std::abs(b - b_naive) <= stochastic_epsilon; \
+                } else { \
+                    eq = eq && (a == a_naive) && (b == b_naive); \
+                } \
+                if (!eq) { \
+                    std::cout << "Mismatch at index " << i << ": " << "(" << a << ", " << b << ") != (" << a_naive << ", " << b_naive << ") -> " << (std::int32_t)(data_out[i].bits) << " != " << (std::int32_t)(data_out_naive[i].bits) << std::endl; \
+                    std::cout << "Input: " << static_cast<fp32_t>(data_in[i]) << std::endl; \
+                    std::cout << "Data in: ["; \
+                    for (std::size_t j {}; j < numel; ++j) { \
+                        std::cout << static_cast<fp32_t>(data_in[j]) << " "; \
+                    } \
+                    std::cout << "]" << std::endl; \
+                    std::cout << "Data out (f): ["; \
+                    for (std::size_t j {}; j < numel_out; ++j) { \
+                        std::cout << static_cast<std::int32_t>(data_out[j].bits) << " "; \
+                    } \
+                    std::cout << "]" << std::endl; \
+                    std::cout << "Data out (n): ["; \
+                    for (std::size_t j {}; j < numel_out; ++j) { \
+                        std::cout << static_cast<std::int32_t>(data_out_naive[j].bits) << " "; \
+                    } \
+                    std::cout << "]" << std::endl; \
+                    std::cout << "Num el out: " << numel_out << std::endl; \
+                    std::cout << "Num el in: " << numel << std::endl; \
+                    ASSERT_TRUE(eq); \
+                } \
+            } \
+        } \
+    }
+
 test_quant_int4(fp32_t, uint4_t, nearest, false, false)
 test_quant_int4(fp32_t, uint4_t, stochastic, true, false)
 test_quant_int4(bfp16_t, uint4_t, nearest, false, false)