uni-quant-cuda 0.2.0__tar.gz

uni_quant_cuda-0.2.0/LICENSE

@@ -0,0 +1,9 @@
+MIT License
+
+Copyright (c) 2026 Jakub Grula
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

uni_quant_cuda-0.2.0/MANIFEST.in

@@ -0,0 +1,3 @@
+include README.md
+include LICENSE
+include *.cpp

uni_quant_cuda-0.2.0/PKG-INFO

@@ -0,0 +1,40 @@
+Metadata-Version: 2.4
+Name: uni-quant-cuda
+Version: 0.2.0
+Summary: Uni-Quant: CUDA-accelerated quantization/dequantization for Keras and XGBoost models
+Author-email: Jakub Grula <ramsters110@gmail.com>
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: POSIX :: Linux
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Dynamic: license-file
+
+Uni-Quant
+========
+
+Small library to quantize/dequantize Keras and XGBoost models using PyTorch CUDA kernels.
+
+Notes
+- This package compiles CUDA kernels at runtime using `torch.utils.cpp_extension.load_inline`.
+- Installing and using the CUDA compilation requires a compatible PyTorch build and CUDA toolkit on the target machine.
+
+Dependencies are listed in `requirements.txt` and synchronized with `pyproject.toml`.
+
+Quick publish test
+
+Build a source/wheel and check locally:
+
+```
+python -m pip install --upgrade build twine
+python -m build
+python -m twine check dist/*
+```
+
+Upload (example):
+
+```
+python -m twine upload dist/*
+```

uni_quant_cuda-0.2.0/README.md

@@ -0,0 +1,26 @@
+Uni-Quant
+========
+
+Small library to quantize/dequantize Keras and XGBoost models using PyTorch CUDA kernels.
+
+Notes
+- This package compiles CUDA kernels at runtime using `torch.utils.cpp_extension.load_inline`.
+- Installing and using the CUDA compilation requires a compatible PyTorch build and CUDA toolkit on the target machine.
+
+Dependencies are listed in `requirements.txt` and synchronized with `pyproject.toml`.
+
+Quick publish test
+
+Build a source/wheel and check locally:
+
+```
+python -m pip install --upgrade build twine
+python -m build
+python -m twine check dist/*
+```
+
+Upload (example):
+
+```
+python -m twine upload dist/*
+```

uni_quant_cuda-0.2.0/dkernel.cpp

@@ -0,0 +1,357 @@
+#include <torch/types.h>
+#include <c10/cuda/CUDAException.h>
+#include <stdio.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#define CHECK_CPU(x) TORCH_CHECK(!x.is_cuda(), #x " must be a CPU tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_U8(x) TORCH_CHECK(x.scalar_type() == torch::kUInt8, #x " must be uint8")
+#define CHECK_INPUT_CPU_HEX(x) CHECK_CPU(x); CHECK_CONTIGUOUS(x); CHECK_U8(x)
+
+static inline int64_t cdiv_i64(int64_t a, int64_t b) { return (a + b - 1) / b; }
+
+__device__ __forceinline__ uint8_t hex_nibble(uint8_t c) {
+	if (c >= '0' && c <= '9') return (uint8_t)(c - '0');
+	if (c >= 'a' && c <= 'f') return (uint8_t)(10 + (c - 'a'));
+	return (uint8_t)(10 + (c - 'A'));
+}
+
+__device__ __forceinline__ uint8_t hex_byte(const uint8_t* s2) {
+	return (uint8_t)((hex_nibble(s2[0]) << 4) | hex_nibble(s2[1]));
+}
+
+__device__ __forceinline__ float read_be_f32_from_hex8(const uint8_t* s8) {
+	uint32_t b0 = (uint32_t)hex_byte(s8 + 0);
+	uint32_t b1 = (uint32_t)hex_byte(s8 + 2);
+	uint32_t b2 = (uint32_t)hex_byte(s8 + 4);
+	uint32_t b3 = (uint32_t)hex_byte(s8 + 6);
+	uint32_t u = (b0 << 24) | (b1 << 16) | (b2 << 8) | b3;
+	return __uint_as_float(u);
+}
+
+__global__ void dequant_hex_kernel(
+	int mode,
+	const uint8_t* hex,
+	float* out_a,
+	float* out_b,
+	int64_t d1, int64_t d2, int64_t d3, int64_t d4,
+	int64_t pack_size,
+	int quant_size,
+	int half_point,
+	int balanced,
+	int literal,
+	int64_t total_elems
+) {
+	int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+	if (idx >= total_elems) return;
+
+	int64_t batches = 0;
+	int64_t rem = 0;
+
+	int64_t batch_shift = (int64_t)(pack_size * (quant_size / 4.0f));
+	int64_t rec_full = 8 + batch_shift;
+
+	int64_t rem_shift = 0;
+	int64_t rec_rem = 0;
+
+	int64_t elems_weights = 0;
+
+	int64_t rows = 0;
+
+	if (mode == 0 || mode == 4) {
+		batches = d2 / pack_size;
+		rem = d2 % pack_size;
+		rem_shift = (quant_size == 4) ? (int64_t)((rem + (rem & 1)) * (quant_size / 4)) : (int64_t)(rem * (quant_size / 4));
+		rec_rem = rem ? (8 + rem_shift) : 0;
+		rows = d1;
+		elems_weights = d1 * d2;
+	} else if (mode == 1) {
+		batches = d3 / pack_size;
+		rem = d3 % pack_size;
+		rem_shift = (quant_size == 4) ? (int64_t)((rem + (rem & 1)) * (quant_size / 4)) : (int64_t)(rem * (quant_size / 4));
+		rec_rem = rem ? (8 + rem_shift) : 0;
+		rows = d1 * d2;
+		elems_weights = d1 * d2 * d3;
+	} else if (mode == 2) {
+		batches = d4 / pack_size;
+		rem = d4 % pack_size;
+		rem_shift = (quant_size == 4) ? (int64_t)((rem + (rem & 1)) * (quant_size / 4)) : (int64_t)(rem * (quant_size / 4));
+		rec_rem = rem ? (8 + rem_shift) : 0;
+		rows = d1 * d2 * d3;
+		elems_weights = d1 * d2 * d3 * d4;
+	} else {
+		batches = d1 / pack_size;
+		rem = d1 % pack_size;
+		rem_shift = (quant_size == 4) ? (int64_t)((rem + (rem & 1)) * (quant_size / 4)) : (int64_t)(rem * (quant_size / 4));
+		rec_rem = rem ? (8 + rem_shift) : 0;
+		rows = 1;
+		
+		elems_weights = d1;
+	}
+
+	int is_second = 0;
+	int is_bias = 0;
+	int64_t pos = idx;
+
+	int seg = 0;
+	if (mode == 5) {
+		seg = (int)(idx / d1);
+		pos = idx - (int64_t)seg * d1;
+		is_second = 0;
+		is_bias = 0;
+	} else if (mode == 4) {
+		is_bias = 0;
+	} else if (mode == 3) {
+		if (pos >= elems_weights) {
+			is_second = 1;
+			pos -= elems_weights;
+		}
+		is_bias = is_second;
+	} else {
+		if (pos >= elems_weights) {
+			is_bias = 1;
+			pos -= elems_weights;
+		}
+	}
+
+	int64_t row = 0;
+	int64_t col = 0;
+
+	if (mode == 0 || mode == 4) {
+		if (!is_bias) {
+			row = pos / d2;
+			col = pos - row * d2;
+		} else {
+			row = rows;
+			col = pos;
+		}
+	} else if (mode == 1) {
+		if (!is_bias) {
+			row = pos / d3;
+			col = pos - row * d3;
+		} else {
+			row = rows;
+			col = pos;
+		}
+	} else if (mode == 2) {
+		if (!is_bias) {
+			row = pos / d4;
+			col = pos - row * d4;
+		} else {
+			row = rows;
+			col = pos;
+		}
+	} else {
+		row = 0;
+		col = pos;
+	}
+
+	int64_t row_hex = batches * rec_full + (rem ? rec_rem : 0);
+	int64_t base_hex = 0;
+
+	if (mode == 5) {
+		base_hex = (int64_t)seg * row_hex;
+	} else if (mode == 3) {
+		base_hex = is_second ? (row_hex) : 0;
+	} else {
+		base_hex = row * row_hex;
+		if (is_bias) base_hex = rows * row_hex;
+	}
+
+	int64_t blk = col / pack_size;
+	int64_t within = col - blk * pack_size;
+
+	int64_t rec_hex_off = 0;
+	int64_t payload_hex_off = 0;
+	float scale = 0.0f;
+
+	if (blk < batches) {
+		rec_hex_off = base_hex + blk * rec_full;
+		scale = read_be_f32_from_hex8(hex + rec_hex_off);
+		payload_hex_off = rec_hex_off + 8;
+		if (quant_size == 8) {
+			int q = (int)hex_byte(hex + payload_hex_off + within * 2);
+			int n = q - (balanced ? half_point : 0);
+			float v = literal ? (float)n : ((float)n * scale);
+			out_a[idx] = v;
+			if ((mode == 3 || is_bias) && mode != 5) out_b[pos] = v;
+		} else {
+			int64_t byte_i = within >> 1;
+			uint8_t b = hex_byte(hex + payload_hex_off + byte_i * 2);
+			int q = (within & 1) ? (int)(b & 0x0F) : (int)((b >> 4) & 0x0F);
+			if ((within & 1) && q == 0) {
+				out_a[idx] = 0.0f;
+				if ((mode == 3 || is_bias) && mode != 5) out_b[pos] = 0.0f;
+			} else {
+				int n = q - (balanced ? half_point : 0);
+				float v = literal ? (float)n : ((float)n * scale);
+				out_a[idx] = v;
+				if ((mode == 3 || is_bias) && mode != 5) out_b[pos] = v;
+			}
+		}
+	} else {
+		if (!rem) {
+			out_a[idx] = 0.0f;	
+			if ((mode == 3 || is_bias) && mode != 5) out_b[pos] = 0.0f;
+			return;
+		}
+		int64_t rblk = blk - batches;
+		rec_hex_off = base_hex + batches * rec_full + rblk * rec_rem;
+		scale = read_be_f32_from_hex8(hex + rec_hex_off);
+		payload_hex_off = rec_hex_off + 8;
+
+		if (quant_size == 8) {
+			int q = (int)hex_byte(hex + payload_hex_off + within * 2);
+			int n = q - (balanced ? half_point : 0);
+			float v = literal ? (float)n : ((float)n * scale);
+			out_a[idx] = v;
+			if ((mode == 3 || is_bias) && mode != 5) out_b[pos] = v;
+		} else {
+			int64_t byte_i = within >> 1;
+			uint8_t b = hex_byte(hex + payload_hex_off + byte_i * 2);
+			int q = (within & 1) ? (int)(b & 0x0F) : (int)((b >> 4) & 0x0F);
+			if ((within & 1) && q == 0) {
+				out_a[idx] = 0.0f;
+				if ((mode == 3 || is_bias) && mode != 5) out_b[pos] = 0.0f;
+			} else {
+				int n = q - (balanced ? half_point : 0);
+				float v = literal ? (float)n : ((float)n * scale);
+				out_a[idx] = v;
+				if ((mode == 3 || is_bias) && mode != 5) out_b[pos] = v;
+			}
+		}
+	}
+}
+
+static std::vector<torch::Tensor> dequantize_hex_impl(
+	int mode,
+	torch::Tensor hex_cpu,
+	int64_t d1, int64_t d2, int64_t d3, int64_t d4,
+	int64_t pack_size,
+	int quant_size,
+	bool balanced,
+	bool literal
+) {
+	CHECK_INPUT_CPU_HEX(hex_cpu);
+	TORCH_CHECK(pack_size > 0, "pack_size must be > 0");
+	TORCH_CHECK(quant_size == 4 || quant_size == 8, "quant_size must be 4 or 8");
+	int half_point = (1 << quant_size) / 2;
+
+	int64_t total = 0;
+	if (mode == 0) {
+		TORCH_CHECK(d2 >= pack_size, "d2 must be >= pack_size");
+		total = d1 * d2 + d2;
+	} else if (mode == 1) {
+		TORCH_CHECK(d3 >= pack_size, "d3 must be >= pack_size");
+		total = d1 * d2 * d3 + d3;
+	} else if (mode == 2) {
+		TORCH_CHECK(d4 >= pack_size, "d4 must be >= pack_size");
+		total = d1 * d2 * d3 * d4 + d4;
+	} else if (mode == 4) {
+		TORCH_CHECK(d2 >= pack_size, "d2 must be >= pack_size");
+		total = d1 * d2;
+	} else if (mode == 5) {
+		TORCH_CHECK(d1 >= pack_size, "d1 must be >= pack_size");
+		total = 4 * d1;
+	} else {
+		TORCH_CHECK(d1 >= pack_size, "d1 must be >= pack_size");
+		total = d1 + d1;
+	}
+
+	auto hex_gpu = hex_cpu.to(torch::kCUDA);
+	auto out_a_gpu = torch::empty({total}, torch::TensorOptions().device(torch::kCUDA).dtype(torch::kFloat32));
+	torch::Tensor out_b_gpu;
+
+	if (mode == 0) out_b_gpu = torch::empty({d2}, torch::TensorOptions().device(torch::kCUDA).dtype(torch::kFloat32));
+	else if (mode == 1) out_b_gpu = torch::empty({d3}, torch::TensorOptions().device(torch::kCUDA).dtype(torch::kFloat32));
+	else if (mode == 2) out_b_gpu = torch::empty({d4}, torch::TensorOptions().device(torch::kCUDA).dtype(torch::kFloat32));
+	else if (mode == 3) out_b_gpu = torch::empty({d1}, torch::TensorOptions().device(torch::kCUDA).dtype(torch::kFloat32));
+	else out_b_gpu = torch::empty({1}, torch::TensorOptions().device(torch::kCUDA).dtype(torch::kFloat32));
+
+	int threads = 256;
+	int grid = (int)cdiv_i64(total, threads);
+
+	dequant_hex_kernel<<<grid, threads>>>(
+		mode,
+		(const uint8_t*)hex_gpu.data_ptr<uint8_t>(),
+		(float*)out_a_gpu.data_ptr<float>(),
+		(float*)out_b_gpu.data_ptr<float>(),
+		d1, d2, d3, d4,
+		pack_size,
+		quant_size,
+		half_point,
+		(int)balanced,
+		(int)literal,
+		total
+	);
+	C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+	if (mode == 0) {
+		auto w = out_a_gpu.narrow(0, 0, d1 * d2).view({d1, d2}).to(torch::kCPU);
+		auto b = out_b_gpu.to(torch::kCPU);
+		return { w, b };
+	}
+	if (mode == 1) {
+		auto w = out_a_gpu.narrow(0, 0, d1 * d2 * d3).view({d1, d2, d3}).to(torch::kCPU);
+		auto b = out_b_gpu.to(torch::kCPU);
+		return { w, b };
+	}
+	if (mode == 2) {
+		auto w = out_a_gpu.narrow(0, 0, d1 * d2 * d3 * d4).view({d1, d2, d3, d4}).to(torch::kCPU);
+		auto b = out_b_gpu.to(torch::kCPU);
+		return { w, b };
+	}
+	if (mode == 4) {
+		return { out_a_gpu.to(torch::kCPU) };
+	}
+	if (mode == 5) {
+		auto flat = out_a_gpu.to(torch::kCPU);
+		auto w = flat.narrow(0, 0 * d1, d1).contiguous();
+		auto b = flat.narrow(0, 1 * d1, d1).contiguous();
+		auto rm = flat.narrow(0, 2 * d1, d1).contiguous();
+		auto rv = flat.narrow(0, 3 * d1, d1).contiguous();
+		return { w, b, rm, rv };
+	}
+	auto a = out_a_gpu.narrow(0, 0, d1).to(torch::kCPU);
+	auto b = out_b_gpu.to(torch::kCPU);
+	return { a, b };
+}
+
+std::vector<torch::Tensor> dequantize_dense_hex(torch::Tensor hex_cpu, int64_t d1, int64_t d2, int64_t pack_size, int quant_size, bool balanced, bool literal) {
+	return dequantize_hex_impl(0, hex_cpu, d1, d2, 0, 0, pack_size, quant_size, balanced, literal);
+}
+
+std::vector<torch::Tensor> dequantize_conv1d_hex(torch::Tensor hex_cpu, int64_t d1, int64_t d2, int64_t d3, int64_t pack_size, int quant_size, bool balanced, bool literal) {
+	return dequantize_hex_impl(1, hex_cpu, d1, d2, d3, 0, pack_size, quant_size, balanced, literal);
+}
+
+std::vector<torch::Tensor> dequantize_conv2d_hex(torch::Tensor hex_cpu, int64_t d1, int64_t d2, int64_t d3, int64_t d4, int64_t pack_size, int quant_size, bool balanced, bool literal) {
+	return dequantize_hex_impl(2, hex_cpu, d1, d2, d3, d4, pack_size, quant_size, balanced, literal);
+}
+
+std::vector<torch::Tensor> dequantize_gru_hex(torch::Tensor hex_cpu, int64_t d1, int64_t units, int64_t biases, int64_t pack_size, int quant_size, bool balanced, bool literal) {
+	int64_t d2 = 3 * units;
+	int64_t rows = d1 + units + biases;
+
+	auto tmp = dequantize_hex_impl(4, hex_cpu, rows, d2, 0, 0, pack_size, quant_size, balanced, literal);
+	auto flat = tmp[0];
+
+	auto full = flat.view({rows, d2});
+	auto w_in = full.narrow(0, 0, d1).contiguous();
+	auto w_rec = full.narrow(0, d1, units).contiguous();
+	auto b2d = full.narrow(0, d1 + units, biases).contiguous();
+	if (biases == 1) {
+		auto b1d = b2d.view({d2}).contiguous();
+		return { w_in, w_rec, b1d };
+	}
+	return { w_in, w_rec, b2d };
+}
+
+std::vector<torch::Tensor> dequantize_layernorm_hex(torch::Tensor hex_cpu, int64_t d1, int64_t pack_size, int quant_size, bool balanced, bool literal) {
+	return dequantize_hex_impl(3, hex_cpu, d1, 0, 0, 0, pack_size, quant_size, balanced, literal);
+}
+
+std::vector<torch::Tensor> dequantize_batchnorm_hex(torch::Tensor hex_cpu, int64_t d1, int64_t pack_size, int quant_size, bool balanced, bool literal) {
+	return dequantize_hex_impl(5, hex_cpu, d1, 0, 0, 0, pack_size, quant_size, balanced, literal);
+}

uni_quant_cuda-0.2.0/pyproject.toml

@@ -0,0 +1,21 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "uni-quant-cuda"
+version = "0.2.0"
+description = "Uni-Quant: CUDA-accelerated quantization/dequantization for Keras and XGBoost models"
+readme = "README.md"
+license = { text = "MIT" }
+authors = [ { name = "Jakub Grula", email = "ramsters110@gmail.com" } ]
+requires-python = ">=3.11"
+dynamic = ["dependencies"]
+
+classifiers = [
+	"Programming Language :: Python :: 3",
+	"License :: OSI Approved :: MIT License",
+	"Operating System :: POSIX :: Linux"
+]
+
+## No [tool.setuptools] section needed for flat layout

uni_quant_cuda-0.2.0/qkernel.cpp

@@ -0,0 +1,238 @@
+#include <torch/types.h>
+#include <c10/cuda/CUDAException.h>
+#include <stdio.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#define CHECK_CPU(x) TORCH_CHECK(!x.is_cuda(), #x " must be a CPU tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_F32(x) TORCH_CHECK(x.scalar_type() == torch::kFloat32, #x " must be float32")
+#define CHECK_INPUT_CPU_F32(x) CHECK_CPU(x); CHECK_CONTIGUOUS(x); CHECK_F32(x)
+
+static inline int64_t cdiv_i64(int64_t a, int64_t b) { return (a + b - 1) / b; }
+
+__device__ __forceinline__ uint32_t bswap32(uint32_t v) {
+	return ((v & 0x000000FFu) << 24) | ((v & 0x0000FF00u) << 8) | ((v & 0x00FF0000u) >> 8) | ((v & 0xFF000000u) >> 24);
+}
+
+__device__ __forceinline__ void write_be_f32(uint8_t* dst, float v) {
+	uint32_t u = __float_as_uint(v);
+	u = bswap32(u);
+	dst[0] = (uint8_t)(u & 0xFF);
+	dst[1] = (uint8_t)((u >> 8) & 0xFF);
+	dst[2] = (uint8_t)((u >> 16) & 0xFF);
+	dst[3] = (uint8_t)((u >> 24) & 0xFF);
+}
+
+__device__ __forceinline__ int clamp_int(int v, int lo, int hi) {
+	return v < lo ? lo : (v > hi ? hi : v);
+}
+
+__global__ void quant_pack_kernel(
+	int mode,
+	const float* w,
+	int64_t d0, int64_t d1, int64_t d2, int64_t d3,
+	int64_t pack_size,
+	int quant_size,
+	int half_point,
+	int64_t blocks_per_inner,
+	int64_t total_blocks,
+	int64_t stride,
+	int64_t rem,
+	int64_t rem_payload,
+	int64_t rem_stride,
+	int64_t row_bytes,
+	uint8_t* out
+) {
+	int64_t bid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+	if (bid >= total_blocks) return;
+
+	int64_t row_idx = bid / blocks_per_inner;
+	int64_t blk_in_row = bid % blocks_per_inner;
+	int is_rem_block = (rem > 0 && blk_in_row == (blocks_per_inner - (rem > 0 ? 1 : 0))) ? 1 : 0;
+	int64_t actual_pack = is_rem_block ? rem : pack_size;
+
+	int64_t base = 0;
+
+	if (mode == 0) {
+		base = bid * pack_size;
+	} else if (mode == 1) {
+		int64_t i = bid / blocks_per_inner;
+		int64_t jb = bid - i * blocks_per_inner;
+		int64_t start = jb * pack_size;
+		base = i * d1 + start;
+	} else if (mode == 2) {
+		int64_t tmp = bid;
+		int64_t kb = tmp % blocks_per_inner;
+		tmp /= blocks_per_inner;
+		int64_t j = tmp % d1;
+		int64_t i = tmp / d1;
+		int64_t start = kb * pack_size;
+		base = (i * d1 + j) * d2 + start;
+	} else {
+		int64_t tmp = bid;
+		int64_t lb = tmp % blocks_per_inner;
+		tmp /= blocks_per_inner;
+		int64_t k = tmp % d2;
+		tmp /= d2;
+		int64_t j = tmp % d1;
+		int64_t i = tmp / d1;
+		int64_t start = lb * pack_size;
+		base = (((i * d1 + j) * d2 + k) * d3) + start;
+	}
+
+	const float* src = w + base;
+	int64_t out_off = row_idx * row_bytes;
+	if (is_rem_block) {
+		out_off += (blocks_per_inner - 1) * stride;
+	} else {
+		out_off += blk_in_row * stride;
+	}
+	uint8_t* dst = out + out_off;
+
+	float max_abs = 0.0f;
+	for (int64_t t = 0; t < actual_pack; ++t) {
+		float a = fabsf(src[t]);
+		if (a > max_abs) max_abs = a;
+	}
+
+	float scale = max_abs / (float)(half_point - 1);
+	write_be_f32(dst, scale);
+	dst += 4;
+
+	int lo = -(half_point - 1);
+	int hi = (half_point - 1);
+
+	if (quant_size == 8) {
+		for (int64_t t = 0; t < actual_pack; ++t) {
+			float x = src[t];
+			float qf;
+			if (scale == 0.0f) qf = (x > 0.0f) ? (float)hi : ((x < 0.0f) ? (float)lo : 0.0f);
+			else qf = nearbyintf(x / scale);
+			int q = clamp_int((int)qf, lo, hi) + half_point;
+			dst[t] = (uint8_t)q;
+		}
+	} else {
+		int64_t o = 0;
+		int64_t padded = actual_pack + (actual_pack & 1);
+		for (int64_t t = 0; t < padded; t += 2) {
+			float x0 = (t < actual_pack) ? src[t] : 0.0f;
+			float qf0;
+			if (scale == 0.0f) qf0 = (x0 > 0.0f) ? (float)hi : ((x0 < 0.0f) ? (float)lo : 0.0f);
+			else qf0 = nearbyintf(x0 / scale);
+			int q0 = clamp_int((int)qf0, lo, hi) + half_point;
+
+			float x1 = (t + 1 < actual_pack) ? src[t + 1] : 0.0f;
+			float qf1;
+			if (scale == 0.0f) qf1 = (x1 > 0.0f) ? (float)hi : ((x1 < 0.0f) ? (float)lo : 0.0f);
+			else qf1 = nearbyintf(x1 / scale);
+			int q1 = clamp_int((int)qf1, lo, hi) + half_point;
+
+			dst[o++] = (uint8_t)(((q0 & 0x0F) << 4) | (q1 & 0x0F));
+		}
+	}
+}
+
+static torch::Tensor quantize_pack_impl(int mode, torch::Tensor w_cpu, int64_t pack_size, int quant_size) {
+	CHECK_INPUT_CPU_F32(w_cpu);
+	TORCH_CHECK(pack_size > 0, "pack_size must be > 0");
+	TORCH_CHECK(quant_size == 4 || quant_size == 8, "quant_size must be 4 or 8");
+	if (quant_size == 4) TORCH_CHECK((pack_size & 1) == 0, "pack_size must be even for quant_size=4");
+
+	int half_point = (1 << quant_size) / 2;
+
+	int64_t d0 = 0, d1 = 0, d2 = 0, d3 = 0;
+	int64_t blocks_per_inner = 0;
+	int64_t total_blocks = 0;
+
+	if (mode == 0) {
+		TORCH_CHECK(w_cpu.dim() == 1, "w_cpu must be 1D");
+		d0 = w_cpu.size(0);
+		blocks_per_inner = d0 / pack_size;
+		if (d0 % pack_size != 0) blocks_per_inner += 1;
+		total_blocks = blocks_per_inner;
+	} else if (mode == 1) {
+		TORCH_CHECK(w_cpu.dim() == 2, "w_cpu must be 2D");
+		d0 = w_cpu.size(0);
+		d1 = w_cpu.size(1);
+		blocks_per_inner = d1 / pack_size;
+		if (d1 % pack_size != 0) blocks_per_inner += 1;
+		total_blocks = d0 * blocks_per_inner;
+	} else if (mode == 2) {
+		TORCH_CHECK(w_cpu.dim() == 3, "w_cpu must be 3D");
+		d0 = w_cpu.size(0);
+		d1 = w_cpu.size(1);
+		d2 = w_cpu.size(2);
+		blocks_per_inner = d2 / pack_size;
+		if (d2 % pack_size != 0) blocks_per_inner += 1;
+		total_blocks = d0 * d1 * blocks_per_inner;
+	} else {
+		TORCH_CHECK(w_cpu.dim() == 4, "w_cpu must be 4D");
+		d0 = w_cpu.size(0);
+		d1 = w_cpu.size(1);
+		d2 = w_cpu.size(2);
+		d3 = w_cpu.size(3);
+		blocks_per_inner = d3 / pack_size;
+		if (d3 % pack_size != 0) blocks_per_inner += 1;
+		total_blocks = d0 * d1 * d2 * blocks_per_inner;
+	}
+
+	int64_t payload = (quant_size == 4) ? (pack_size / 2) : pack_size;
+	int64_t stride = 4 + payload;
+
+	int64_t inner_dim = (mode == 0) ? d0 : ((mode == 1) ? d1 : ((mode == 2) ? d2 : d3));
+	int64_t rem = inner_dim % pack_size;
+	int64_t rem_payload = 0, rem_stride = 0;
+	if (rem > 0) {
+		int64_t rem_padded = rem + (rem & 1);
+		rem_payload = (quant_size == 4) ? (rem_padded / 2) : rem;
+		rem_stride = 4 + rem_payload;
+	}
+
+	int64_t num_rows = total_blocks / blocks_per_inner;
+	int64_t full_blocks = inner_dim / pack_size;
+	int64_t row_bytes = full_blocks * stride + (rem > 0 ? rem_stride : 0);
+	int64_t total_bytes = num_rows * row_bytes;
+
+	auto w = w_cpu.to(torch::kCUDA);
+	auto out = torch::zeros({total_bytes}, torch::TensorOptions().device(torch::kCUDA).dtype(torch::kUInt8));
+
+	int threads = 256;
+	int grid = (int)cdiv_i64(total_blocks, threads);
+
+	quant_pack_kernel<<<grid, threads>>>(
+		mode,
+		(const float*)w.data_ptr<float>(),
+		d0, d1, d2, d3,
+		pack_size,
+		quant_size,
+		half_point,
+		blocks_per_inner,
+		total_blocks,
+		stride,
+		rem,
+		rem_payload,
+		rem_stride,
+		row_bytes,
+		(uint8_t*)out.data_ptr<uint8_t>()
+	);
+	C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+	return out.to(torch::kCPU);
+}
+
+torch::Tensor quantize_pack_1d(torch::Tensor w_cpu, int64_t pack_size, int quant_size) {
+	return quantize_pack_impl(0, w_cpu, pack_size, quant_size);
+}
+
+torch::Tensor quantize_pack_2d(torch::Tensor w_cpu, int64_t pack_size, int quant_size) {
+	return quantize_pack_impl(1, w_cpu, pack_size, quant_size);
+}
+
+torch::Tensor quantize_pack_3d(torch::Tensor w_cpu, int64_t pack_size, int quant_size) {
+	return quantize_pack_impl(2, w_cpu, pack_size, quant_size);
+}
+
+torch::Tensor quantize_pack_4d(torch::Tensor w_cpu, int64_t pack_size, int quant_size) {
+	return quantize_pack_impl(3, w_cpu, pack_size, quant_size);
+}

uni_quant_cuda-0.2.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

uni_quant_cuda-0.2.0/uni_quant_cuda.egg-info/PKG-INFO

@@ -0,0 +1,40 @@
+Metadata-Version: 2.4
+Name: uni-quant-cuda
+Version: 0.2.0
+Summary: Uni-Quant: CUDA-accelerated quantization/dequantization for Keras and XGBoost models
+Author-email: Jakub Grula <ramsters110@gmail.com>
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: POSIX :: Linux
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Dynamic: license-file
+
+Uni-Quant
+========
+
+Small library to quantize/dequantize Keras and XGBoost models using PyTorch CUDA kernels.
+
+Notes
+- This package compiles CUDA kernels at runtime using `torch.utils.cpp_extension.load_inline`.
+- Installing and using the CUDA compilation requires a compatible PyTorch build and CUDA toolkit on the target machine.
+
+Dependencies are listed in `requirements.txt` and synchronized with `pyproject.toml`.
+
+Quick publish test
+
+Build a source/wheel and check locally:
+
+```
+python -m pip install --upgrade build twine
+python -m build
+python -m twine check dist/*
+```
+
+Upload (example):
+
+```
+python -m twine upload dist/*
+```

uni_quant_cuda-0.2.0/uni_quant_cuda.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+LICENSE
+MANIFEST.in
+README.md
+dkernel.cpp
+pyproject.toml
+qkernel.cpp
+uni_quant_cuda.egg-info/PKG-INFO
+uni_quant_cuda.egg-info/SOURCES.txt
+uni_quant_cuda.egg-info/dependency_links.txt
+uni_quant_cuda.egg-info/top_level.txt

uni_quant_cuda-0.2.0/uni_quant_cuda.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

uni_quant_cuda-0.2.0/uni_quant_cuda.egg-info/top_level.txt

@@ -0,0 +1 @@
+LaTeX