diff --git a/src/po2mm/__init__.py b/src/po2mm/__init__.py index 976d030..eba792e 100644 --- a/src/po2mm/__init__.py +++ b/src/po2mm/__init__.py @@ -1 +1 @@ -from .po2_gemv import po2_gemv +from .kernels import po2_gemv, po2_gemm diff --git a/src/po2mm/kernels.py b/src/po2mm/kernels.py new file mode 100644 index 0000000..7429fbe --- /dev/null +++ b/src/po2mm/kernels.py @@ -0,0 +1,50 @@ +from pathlib import Path +from torch.utils.cpp_extension import load +import torch + +_here = Path(__file__).parent + +_ext = load( + name="po2_gemv_ext", + sources=[ + str(_here / "po2_binding.cpp"), + str(_here / "po2_gemv_kernel.cu"), + str(_here / "po2_gemm_kernel.cu"), + ], + extra_include_paths=[str(_here)], + extra_cuda_cflags=[ + "-O3", + "--use_fast_math", + "-std=c++17", + "-expt-relaxed-constexpr", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "-U__CUDA_NO_BFLOAT162_OPERATORS__", + ], + extra_cflags=["-O3", "-std=c++17"], + verbose=False, +) + +@torch.library.custom_op("po2::gemv", mutates_args=()) +def po2_gemv(vecs: torch.Tensor, values: torch.Tensor, scales: torch.Tensor, LUTs: torch.Tensor) -> torch.Tensor: + return _ext.po2_gemv(vecs, values, scales, LUTs) + +@po2_gemv.register_fake +def po2_gemv_fake(vecs, values, scales, LUTs): + #assert values.dim == 2, values.shape + if vecs.dim() == 2: + return torch.empty((vecs.shape[0], values.shape[0]), device=vecs.device, dtype=vecs.dtype) + else: + return torch.empty((values.shape[0],), device=vecs.device, dtype=vecs.dtype) + + +@torch.library.custom_op("po2::gemm", mutates_args=()) +def po2_gemm(x: torch.Tensor, values: torch.Tensor, scales: torch.Tensor, LUTs: torch.Tensor) -> torch.Tensor: + return _ext.po2_gemm(x, values, scales, LUTs) + +@po2_gemm.register_fake +def po2_gemm_fake(x, values, scales, LUTs): + return torch.empty((x.shape[0], values.shape[0]), device=x.device, dtype=x.dtype) diff --git a/src/po2mm/linear.py b/src/po2mm/linear.py index 81e6921..f26d1a3 100644 --- a/src/po2mm/linear.py +++ b/src/po2mm/linear.py @@ -2,7 +2,7 @@ import torch from torch import nn -from .po2_gemv import po2_gemv +from .kernels import po2_gemv from .quant import Po2QuantConfig, quantize_po2, ScaleType @@ -34,17 +34,7 @@ def from_linear(cls, def _forward_real(self, x: torch.Tensor) -> torch.Tensor: leading = x.shape[:-1] vecs = x.reshape(-1, self.in_features).contiguous() - N = vecs.shape[0] - - if N <= 7: - out = po2_gemv(vecs, self.weight_values, self.weight_scales, self.weight_lut.flatten()) - else: - chunks = vecs.split(7, dim=0) - out = torch.cat([ - po2_gemv(chunk, self.weight_values, self.weight_scales, self.weight_lut.flatten()) - for chunk in chunks - ], dim=0) - + out = po2_gemv(vecs, self.weight_values, self.weight_scales, self.weight_lut.flatten()) out = out.reshape(*leading, self.out_features).to(x.dtype) if self.bias is not None: out = out + self.bias.to(out.dtype) diff --git a/src/po2mm/po2_binding.cpp b/src/po2mm/po2_binding.cpp index a77df6b..aa991fe 100644 --- a/src/po2mm/po2_binding.cpp +++ b/src/po2mm/po2_binding.cpp @@ -28,6 +28,7 @@ torch::Tensor po2_gemv( TORCH_CHECK(LUTs.dim() == 1, "LUTs must be 1D, got shape ", LUTs.sizes()); TORCH_CHECK(LUTs.size(0) == 32, "LUTs must have length 32"); + TORCH_CHECK(vecs.dim() == 2 || vecs.dim() == 1, "vecs must be 1D or 2D, got shape ", vecs.sizes()); const int m = values.size(0); const int n = vecs.dim() == 2 ? vecs.size(0) : 1; @@ -47,9 +48,70 @@ torch::Tensor po2_gemv( auto out = vecs.dim() == 2 ? torch::empty({n, m}, vecs.options()) : torch::empty({m}, vecs.options()); - launch_po2_gemv( + for (int s = 0; s < n; s += 8) { + int t = std::min(s + 8, n); + launch_po2_gemv( + reinterpret_cast(out.data_ptr()) + s * out.stride(0), + reinterpret_cast(vecs.data_ptr()) + s * vecs.stride(0), + reinterpret_cast(values.data_ptr()), + reinterpret_cast(scales.data_ptr()), + reinterpret_cast(LUTs.data_ptr()), + sf_type, group_size, + m, t - s, k, + at::cuda::getCurrentCUDAStream() + ); + } + + + return out; +} + +torch::Tensor po2_gemm( + torch::Tensor x, + torch::Tensor values, + torch::Tensor scales, + torch::Tensor LUTs +) { + TORCH_CHECK(x.is_cuda() && values.is_cuda() && scales.is_cuda() && LUTs.is_cuda(), + "all tensors must be on CUDA"); + TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be bfloat16"); + TORCH_CHECK(LUTs.scalar_type() == torch::kBFloat16, "LUTs must be bfloat16"); + TORCH_CHECK(values.scalar_type() == torch::kUInt32, "values must be uint32"); + + TORCH_CHECK(x.is_contiguous() && values.is_contiguous() + && scales.is_contiguous() && LUTs.is_contiguous(), + "all tensors must be contiguous"); + + TORCH_CHECK(values.dim() == 2, + "values must be 2D, got shape ", values.sizes()); + TORCH_CHECK(scales.dim() == 2, + "scales must be 2D, got shape ", scales.sizes()); + TORCH_CHECK(LUTs.dim() == 1, + "LUTs must be 1D, got shape ", LUTs.sizes()); + TORCH_CHECK(LUTs.size(0) == 32, "LUTs must have length 32"); + TORCH_CHECK(x.dim() == 2, "x must be 2D, got shape ", x.sizes()); + + const int m = values.size(0); + const int n = x.size(0); + const int k = x.size(1); + + TORCH_CHECK(k % 16 == 0, "k must be a multiple of 16"); + TORCH_CHECK(values.size(0) == m, "values.shape[0] must match n"); + TORCH_CHECK(values.size(1) == k / 8, "values.shape[1] must be k/8"); + TORCH_CHECK(scales.size(0) == m, "scales.shape[0] must match n"); + int group_size = k / scales.size(1); + TORCH_CHECK(group_size == 16 || group_size == 32, "group size must be 16 or 32"); + TORCH_CHECK(scales.size(1) * group_size == k, "k must be divisible by scales.shape[1]"); + + TORCH_CHECK(scales.scalar_type() == torch::kFloat8_e4m3fn || scales.scalar_type() == torch::kUInt8, + "scales must be float8_e4m3fn (E4M3) or uint8 (E7M0)"); + SFType sf_type = (scales.scalar_type() == torch::kFloat8_e4m3fn) ? SFType::G1E4M3 : SFType::G1E7M0; + + auto out = torch::empty({n, m}, x.options()); + + launch_po2_gemm_tn_baseline( reinterpret_cast(out.data_ptr()), - reinterpret_cast(vecs.data_ptr()), + reinterpret_cast(x.data_ptr()), reinterpret_cast(values.data_ptr()), reinterpret_cast(scales.data_ptr()), reinterpret_cast(LUTs.data_ptr()), @@ -64,4 +126,6 @@ torch::Tensor po2_gemv( PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("po2_gemv", &po2_gemv, "Power-of-2 quantized GEMV", py::arg("vecs"), py::arg("values"), py::arg("scales"), py::arg("LUTs")); + m.def("po2_gemm", &po2_gemm, "Power-of-2 quantized GEMM", + py::arg("x"), py::arg("values"), py::arg("scales"), py::arg("LUTs")); } \ No newline at end of file diff --git a/src/po2mm/po2_gemm_kernel.cu b/src/po2mm/po2_gemm_kernel.cu new file mode 100644 index 0000000..0ffcd67 --- /dev/null +++ b/src/po2mm/po2_gemm_kernel.cu @@ -0,0 +1,135 @@ +// Copyright (c) 2026, IST Austria, developed by Erik Schultheis +// SPDX-License-Identifier: Apache-2.0 +// + +#include +#include +#include "vec.cuh" +#include "utils.cuh" +#include "po2_kernels.cuh" + +template +__global__ __launch_bounds__(64, 6) void gemm_tn_baseline(nv_bfloat16* __restrict__ out, const nv_bfloat16* __restrict__ x, const unsigned* __restrict__ values, + const std::byte* __restrict__ scales, const nv_bfloat16* __restrict__ LUTs, int m, int n, int k) { + using bf16x8_t = GenericVector; + static_assert(GroupSize % 8 == 0, "GroupSize must divisible by 8"); + + int i = blockIdx.x * 64; + int j = blockIdx.y * 64; + if (i >= m || j >= n) return; + + __shared__ nv_bfloat16 LUTs_shared[32]; + __shared__ nv_bfloat16 x_smem[64 * 16]; + __shared__ nv_bfloat16 decoded[64][16]; + __shared__ float scales_smem[64]; + if (threadIdx.x < 32) { + LUTs_shared[threadIdx.x] = LUTs[threadIdx.x]; + } + __syncthreads(); + + scales += i * (k / GroupSize); + values += i * (k / 8); + + bf16x8_t acc[8] = { + bf16x8_t::zeros(), bf16x8_t::zeros(), bf16x8_t::zeros(), bf16x8_t::zeros(), + bf16x8_t::zeros(), bf16x8_t::zeros(), bf16x8_t::zeros(), bf16x8_t::zeros() + }; + + auto x_swizzle = [](int p, int q) { + int pf = p % 8; + int pg = p / 8; + return pg * 8 + q * 64 + pf * 128; + }; + + for (int s = 0; s < k / 16; ++s) { + + // fetch x + for (int u = threadIdx.x / 4; u < 32; u += 16) { + int t = ((threadIdx.x / 2) % 2) * 32; + if (j + u + t < n) { + bf16x8_t v = bf16x8_t::load(x + 16 * s + (j + u + t) * k + 8 * (threadIdx.x % 2)); + v.store(x_smem + x_swizzle(u + t, threadIdx.x % 2)); + } + } + + // decode values + for (int t = 0; t < 64; t += 32) { + unsigned val = values[2 * s + (threadIdx.x / 2 + t) * (k / 8) + threadIdx.x % 2]; + std::byte scale = scales[s / (GroupSize / 16) + (threadIdx.x / 2 + t) * (k / GroupSize)]; + auto [sf, group] = Scaler::decode_scale(scale); + bf16x8_t w; + lookup(w, val, group, LUTs_shared); + w.store(&decoded[threadIdx.x / 2 + t][(threadIdx.x % 2) * 8]); + scales_smem[threadIdx.x / 2 + t] = sf; // double write + } + + __syncthreads(); + + for (int u = 0; u < 8; ++u) { + bf16x8_t a1 = bf16x8_t::load(x_smem + x_swizzle(8*(threadIdx.x%8) + u, 0)); + bf16x8_t a2 = bf16x8_t::load(x_smem + x_swizzle(8*(threadIdx.x%8) + u, 1)); + for (int v = 0; v < 8; ++v) { + float inner = 0.f; + bf16x8_t b = bf16x8_t::load(&decoded[8*(threadIdx.x/8) + v][0]); + for (int w = 0; w < 8; ++w) { + inner = fma(a1[w], b[w], inner); + } + + b = bf16x8_t::load(&decoded[8*(threadIdx.x/8) + v][8]); + for (int w = 0; w < 8; ++w) { + inner = fma(a2[w], b[w], inner); + } + acc[u][v] += inner * scales_smem[8*(threadIdx.x/8) + v]; + } + } + __syncthreads(); + } + + for (int u = 0; u < 8; ++u) { + if (j + u + (threadIdx.x%8)*8 < n) { + acc[u].store(out + (j + u + (threadIdx.x%8)*8) * m + i + (threadIdx.x / 8) * 8); + } + } +} + +void launch_po2_gemm_tn_baseline( + nv_bfloat16* out, + const nv_bfloat16* x, + const unsigned* values, + const std::byte* scales, + const nv_bfloat16* LUTs, + SFType sf_type, int group_size, + int m, int n, int k, + cudaStream_t stream +) { + if (k % 16 != 0) { + std::fprintf(stderr, "po2_gemm_tn_baseline: k (%d) must be a multiple of 16\n", k); + std::abort(); + } + if (m % 64 != 0) { + std::fprintf(stderr, "po2_gemm_tn_baseline: m (%d) must be a multiple of 64\n", m); + std::abort(); + } + if (group_size != 16 && group_size != 32) { + std::fprintf(stderr, "po2_gemm_tn_baseline: group_size (%d) must be 16 or 32\n", group_size); + std::abort(); + } + + dim3 grid((m + 63) / 64, (n + 63) / 64); + dim3 block(64); + + if (sf_type == SFType::G1E4M3 && group_size == 16) { + gemm_tn_baseline<<>>(out, x, values, scales, LUTs, m, n, k); + } else if (sf_type == SFType::G1E7M0 && group_size == 16) { + gemm_tn_baseline<<>>(out, x, values, scales, LUTs, m, n, k); + } else if (sf_type == SFType::G1E4M3 && group_size == 32) { + gemm_tn_baseline<<>>(out, x, values, scales, LUTs, m, n, k); + } else if (sf_type == SFType::G1E7M0 && group_size == 32) { + gemm_tn_baseline<<>>(out, x, values, scales, LUTs, m, n, k); + } else { + std::fprintf(stderr, "po2_gemm_tn_baseline: unsupported sf_type (%d)\n", static_cast(sf_type)); + std::abort(); + } + + CUDA_CHECK(cudaGetLastError()); +} \ No newline at end of file diff --git a/src/po2mm/po2_gemv.py b/src/po2mm/po2_gemv.py deleted file mode 100644 index a57d728..0000000 --- a/src/po2mm/po2_gemv.py +++ /dev/null @@ -1,29 +0,0 @@ -from pathlib import Path -from torch.utils.cpp_extension import load - -_here = Path(__file__).parent - -_ext = load( - name="po2_gemv_ext", - sources=[ - str(_here / "po2_binding.cpp"), - str(_here / "po2_gemv_kernel.cu"), - ], - extra_include_paths=[str(_here)], - extra_cuda_cflags=[ - "-O3", - "--use_fast_math", - "-std=c++17", - "-expt-relaxed-constexpr", - "-U__CUDA_NO_HALF_OPERATORS__", - "-U__CUDA_NO_HALF_CONVERSIONS__", - "-U__CUDA_NO_HALF2_OPERATORS__", - "-U__CUDA_NO_BFLOAT16_OPERATORS__", - "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", - "-U__CUDA_NO_BFLOAT162_OPERATORS__", - ], - extra_cflags=["-O3", "-std=c++17"], - verbose=False, -) - -po2_gemv = _ext.po2_gemv diff --git a/src/po2mm/po2_gemv_kernel.cu b/src/po2mm/po2_gemv_kernel.cu index 3508b92..17ea6fc 100644 --- a/src/po2mm/po2_gemv_kernel.cu +++ b/src/po2mm/po2_gemv_kernel.cu @@ -9,36 +9,6 @@ #include #include - -using ScaleResult = std::pair; - -struct E4M3Scaling -{ - static __device__ __forceinline__ ScaleResult decode_scale(std::byte raw) - { - std::uint8_t scale_part = static_cast(raw) & 0b0111'1111u; - bool group = (static_cast(raw) & 0b1000'0000u) != 0; - __nv_fp8_e4m3 factor; - factor.__x = scale_part; - return {static_cast(factor), group}; - } -}; - - -struct E7M0Scaling -{ - static __device__ __forceinline__ ScaleResult decode_scale(std::byte raw) - { - // scale_part encodes 2^(scale_part - 64); float32 exponent field = scale_part + 63 - // since (scale_part + 63) - 127 = scale_part - 64. - std::uint8_t scale_part = static_cast(raw) & 0b0111'1111u; - bool group = (static_cast(raw) & 0b1000'0000u) != 0; - float factor = __uint_as_float((static_cast(scale_part) + 63u) << 23u); - return {factor, group}; - } -}; - - template __global__ void po2_gemv_generic( nv_bfloat16* __restrict__ out, // height x N @@ -260,6 +230,8 @@ void launch_po2_gemv( launch_po2_gemv_dispatch<6>(out, vecs, values, scales, LUTs, sf_type, group_size, m, k, stream); } else if (n == 7) { launch_po2_gemv_dispatch<7>(out, vecs, values, scales, LUTs, sf_type, group_size, m, k, stream); + } else if (n == 8) { + launch_po2_gemv_dispatch<8>(out, vecs, values, scales, LUTs, sf_type, group_size, m, k, stream); } else { std::fprintf(stderr, "po2_gemm: n (%d) must be in [1, 7]\n", n); std::abort(); diff --git a/src/po2mm/po2_kernels.cuh b/src/po2mm/po2_kernels.cuh index bdfb928..9271ae9 100644 --- a/src/po2mm/po2_kernels.cuh +++ b/src/po2mm/po2_kernels.cuh @@ -18,3 +18,14 @@ void launch_po2_gemv( int m, int n, int k, cudaStream_t stream = nullptr ); + +void launch_po2_gemm_tn_baseline( + nv_bfloat16* out, + const nv_bfloat16* x, + const unsigned* values, + const std::byte* scales, + const nv_bfloat16* LUTs, + SFType sf_type, int group_size, + int m, int n, int k, + cudaStream_t stream +); \ No newline at end of file diff --git a/src/po2mm/utils.cuh b/src/po2mm/utils.cuh index 4afc834..9819204 100644 --- a/src/po2mm/utils.cuh +++ b/src/po2mm/utils.cuh @@ -4,6 +4,9 @@ #pragma once +#include +#include + #define CUDA_CHECK(expr) \ do { \ @@ -20,7 +23,8 @@ __device__ __forceinline__ nv_bfloat16 lookup(unsigned idx, unsigned shift, bool return LUTs[((idx >> shift) & 0x0f) + (group ? 0x10 : 0)]; } -__device__ __forceinline__ void lookup(nv_bfloat16* tgt, unsigned src, bool group, const nv_bfloat16* lut) { +template +__device__ __forceinline__ void lookup(TargetType tgt, unsigned src, bool group, const nv_bfloat16* lut) { tgt[0] = lookup(src, 28, group, lut); tgt[1] = lookup(src, 24, group, lut); tgt[2] = lookup(src, 20, group, lut); @@ -42,3 +46,31 @@ __device__ __forceinline__ float fma(nv_bfloat16 a, nv_bfloat16 b, float c) { return __fmaf_rn(static_cast(a), static_cast(b), c); #endif } + +using ScaleResult = std::pair; + +struct E4M3Scaling +{ + static __device__ __forceinline__ ScaleResult decode_scale(std::byte raw) + { + std::uint8_t scale_part = static_cast(raw) & 0b0111'1111u; + bool group = (static_cast(raw) & 0b1000'0000u) != 0; + __nv_fp8_e4m3 factor; + factor.__x = scale_part; + return {static_cast(factor), group}; + } +}; + + +struct E7M0Scaling +{ + static __device__ __forceinline__ ScaleResult decode_scale(std::byte raw) + { + // scale_part encodes 2^(scale_part - 64); float32 exponent field = scale_part + 63 + // since (scale_part + 63) - 127 = scale_part - 64. + std::uint8_t scale_part = static_cast(raw) & 0b0111'1111u; + bool group = (static_cast(raw) & 0b1000'0000u) != 0; + float factor = __uint_as_float((static_cast(scale_part) + 63u) << 23u); + return {factor, group}; + } +}; diff --git a/tests/test_gemm.py b/tests/test_gemm.py new file mode 100644 index 0000000..6f574de --- /dev/null +++ b/tests/test_gemm.py @@ -0,0 +1,114 @@ +import pytest +import torch + +from po2mm.quant import ( + LutType, + ScaleType, + Po2QuantConfig, + _make_lut, + get_luts, + scale_quantize_e4m3, + scale_quantize_e7m0, + quantize_to_simple_lut, + dequantize_from_simple_lut, + quantize_simple_po2, + dequantize_simple_po2, + quantize_po2, + dequantize_po2, +) +from po2mm.kernels import po2_gemm + + +def _inputs_randn(n, in_features, device): + """Current test — ideal Gaussian, unrealistically clean.""" + return torch.randn(n, in_features, device=device).bfloat16() + + +def _inputs_with_outliers(n, in_features, device): + """ + Simulate post-LayerNorm activations: + - Most channels near N(0,1) + - ~1% of channels are persistent outliers (magnitude 10-100x) + These outlier channels are the main source of quantization error in LLMs. + """ + x = torch.randn(n, in_features, device=device) + n_outliers = max(1, in_features // 64) + outlier_channels = torch.randperm(in_features, device=device)[:n_outliers] + x[:, outlier_channels] *= torch.empty(n_outliers, device=device).uniform_(10, 100) + return x.bfloat16() + + +INPUT_FACTORIES = [ + pytest.param(_inputs_randn, id="randn"), + pytest.param(_inputs_with_outliers, id="outliers"), +] + +@pytest.mark.parametrize("size", [64]) +@pytest.mark.parametrize("lut_types,scale_type", [ + pytest.param([LutType.INT4, LutType.NF4], ScaleType.E4M3, id="int4+nf4/e4m3"), + pytest.param([LutType.INT4, LutType.NF4], ScaleType.E7M0, id="int4+nf4/e7m0"), + pytest.param([LutType.FP4, LutType.NF4], ScaleType.E4M3, id="fp4+nf4/e4m3"), +]) +@pytest.mark.parametrize("n", [64]) +@pytest.mark.parametrize("input_factory", INPUT_FACTORIES) +@pytest.mark.parametrize("group_size", [16]) +def test_po2_eye(size, lut_types, scale_type, n, input_factory, group_size): + """Quantized linear output should be close to the reference fp32 output.""" + torch.manual_seed(42) + out_features, in_features = size, size + matrix = torch.eye(out_features, in_features, device="cuda") + + config = Po2QuantConfig(lut_types=lut_types, scale_type=scale_type, group_size=group_size) + luts = config.get_luts().to("cuda") + weight_values, weight_scales, weight_lut = quantize_po2(matrix, luts, config.scale_type, config.group_size) + rdq = dequantize_po2(weight_values, weight_scales, weight_lut, config.scale_type) + x = input_factory(n, in_features, device="cuda") + + ref = (x.float() @ rdq.T.float()).bfloat16() + out = po2_gemm(x, weight_values, weight_scales, weight_lut.flatten()) + + assert out.shape == ref.shape + from test_quant import _assert_roundtrip_quality + _assert_roundtrip_quality(ref, out, snr_threshold_db=15, label=f"{size} {lut_types}/{scale_type}") + assert torch.nn.functional.cosine_similarity(ref.flatten(), out.flatten(), dim=0).item() > 0.98 + rel_mse = ((ref - out)**2).sum() / (ref**2).sum() + assert rel_mse < 0.025 + + +@pytest.mark.parametrize("shape", [ + (64, 64), + (128, 64), + (64, 128), + (128, 256), + (1024, 1024), + (768, 3072), +]) +@pytest.mark.parametrize("lut_types,scale_type", [ + pytest.param([LutType.INT4, LutType.NF4], ScaleType.E4M3, id="int4+nf4/e4m3"), + pytest.param([LutType.INT4, LutType.NF4], ScaleType.E7M0, id="int4+nf4/e7m0"), + pytest.param([LutType.FP4, LutType.NF4], ScaleType.E4M3, id="fp4+nf4/e4m3"), +]) +@pytest.mark.parametrize("n", [5, 64, 256]) +@pytest.mark.parametrize("input_factory", INPUT_FACTORIES) +@pytest.mark.parametrize("group_size", [16]) +def test_po2_gemm(shape, lut_types, scale_type, n, input_factory, group_size): + """Quantized linear output should be close to the reference fp32 output.""" + torch.manual_seed(42) + out_features, in_features = shape + matrix = torch.randn(out_features, in_features, device="cuda") + + config = Po2QuantConfig(lut_types=lut_types, scale_type=scale_type, group_size=group_size) + luts = config.get_luts().to("cuda") + weight_values, weight_scales, weight_lut = quantize_po2(matrix, luts, config.scale_type, config.group_size) + rdq = dequantize_po2(weight_values, weight_scales, weight_lut, config.scale_type) + x = input_factory(n, in_features, device="cuda") + + ref = (x.float() @ rdq.T.float()).bfloat16() + out = po2_gemm(x, weight_values, weight_scales, weight_lut.flatten()) + + assert out.shape == ref.shape + from test_quant import _assert_roundtrip_quality + _assert_roundtrip_quality(ref, out, snr_threshold_db=15, label=f"{shape} {lut_types}/{scale_type}") + assert torch.nn.functional.cosine_similarity(ref.flatten(), out.flatten(), dim=0).item() > 0.98 + rel_mse = ((ref - out)**2).sum() / (ref**2).sum() + assert rel_mse < 0.025 diff --git a/tests/test_gemv.py b/tests/test_gemv.py new file mode 100644 index 0000000..fb24f0b --- /dev/null +++ b/tests/test_gemv.py @@ -0,0 +1,81 @@ +import pytest +import torch + +from po2mm.quant import ( + LutType, + ScaleType, + Po2QuantConfig, + _make_lut, + get_luts, + scale_quantize_e4m3, + scale_quantize_e7m0, + quantize_to_simple_lut, + dequantize_from_simple_lut, + quantize_simple_po2, + dequantize_simple_po2, + quantize_po2, + dequantize_po2, +) +from po2mm.kernels import po2_gemv + + +def _inputs_randn(n, in_features, device): + """Current test — ideal Gaussian, unrealistically clean.""" + return torch.randn(n, in_features, device=device).bfloat16() + + +def _inputs_with_outliers(n, in_features, device): + """ + Simulate post-LayerNorm activations: + - Most channels near N(0,1) + - ~1% of channels are persistent outliers (magnitude 10-100x) + These outlier channels are the main source of quantization error in LLMs. + """ + x = torch.randn(n, in_features, device=device) + n_outliers = max(1, in_features // 64) + outlier_channels = torch.randperm(in_features, device=device)[:n_outliers] + x[:, outlier_channels] *= torch.empty(n_outliers, device=device).uniform_(10, 100) + return x.bfloat16() + + +INPUT_FACTORIES = [ + pytest.param(_inputs_randn, id="randn"), + pytest.param(_inputs_with_outliers, id="outliers"), +] + +@pytest.mark.parametrize("shape", [ + (64, 32), + (128, 256), + (32, 64), + (1024, 1024), + (768, 3072), +]) +@pytest.mark.parametrize("lut_types,scale_type", [ + pytest.param([LutType.INT4, LutType.NF4], ScaleType.E4M3, id="int4+nf4/e4m3"), + pytest.param([LutType.INT4, LutType.NF4], ScaleType.E7M0, id="int4+nf4/e7m0"), + pytest.param([LutType.FP4, LutType.NF4], ScaleType.E4M3, id="fp4+nf4/e4m3"), +]) +@pytest.mark.parametrize("n", [1, 2, 3, 4, 5, 6, 7, 12]) +@pytest.mark.parametrize("input_factory", INPUT_FACTORIES) +@pytest.mark.parametrize("group_size", [16, 32]) +def test_po2_gemv(shape, lut_types, scale_type, n, input_factory, group_size): + """Quantized linear output should be close to the reference fp32 output.""" + torch.manual_seed(42) + out_features, in_features = shape + matrix = torch.randn(out_features, in_features, device="cuda") + + config = Po2QuantConfig(lut_types=lut_types, scale_type=scale_type, group_size=group_size) + luts = config.get_luts().to("cuda") + weight_values, weight_scales, weight_lut = quantize_po2(matrix, luts, config.scale_type, config.group_size) + rdq = dequantize_po2(weight_values, weight_scales, weight_lut, config.scale_type) + x = input_factory(n, in_features, device="cuda") + + ref = (x.float() @ rdq.T.float()).bfloat16() + out = po2_gemv(x, weight_values, weight_scales, weight_lut.flatten()) + + assert out.shape == ref.shape + from test_quant import _assert_roundtrip_quality + _assert_roundtrip_quality(ref, out, snr_threshold_db=15, label=f"{shape} {lut_types}/{scale_type}") + assert torch.nn.functional.cosine_similarity(ref.flatten(), out.flatten(), dim=0).item() > 0.98 + rel_mse = ((ref - out)**2).sum() / (ref**2).sum() + assert rel_mse < 0.025