From 4083601c6dd26bf6001667617e3d5ed39968b577 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Wed, 29 Apr 2026 21:41:14 +0200 Subject: [PATCH 1/7] library binding --- src/po2mm/po2_binding.cpp | 1 + src/po2mm/po2_gemv.py | 13 ++++++++++++- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/src/po2mm/po2_binding.cpp b/src/po2mm/po2_binding.cpp index a77df6b..04a5f0e 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; diff --git a/src/po2mm/po2_gemv.py b/src/po2mm/po2_gemv.py index a57d728..aebe76d 100644 --- a/src/po2mm/po2_gemv.py +++ b/src/po2mm/po2_gemv.py @@ -1,5 +1,6 @@ from pathlib import Path from torch.utils.cpp_extension import load +import torch _here = Path(__file__).parent @@ -26,4 +27,14 @@ verbose=False, ) -po2_gemv = _ext.po2_gemv +@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) From d22aac398013ec6f2d7a60c00fdde5388505578a Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Wed, 29 Apr 2026 22:40:50 +0200 Subject: [PATCH 2/7] loop over input vectors in C++ --- src/po2mm/linear.py | 12 +----------- src/po2mm/po2_binding.cpp | 24 ++++++++++++++---------- 2 files changed, 15 insertions(+), 21 deletions(-) diff --git a/src/po2mm/linear.py b/src/po2mm/linear.py index 81e6921..e6ebeba 100644 --- a/src/po2mm/linear.py +++ b/src/po2mm/linear.py @@ -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 04a5f0e..8c88624 100644 --- a/src/po2mm/po2_binding.cpp +++ b/src/po2mm/po2_binding.cpp @@ -48,16 +48,20 @@ torch::Tensor po2_gemv( auto out = vecs.dim() == 2 ? torch::empty({n, m}, vecs.options()) : torch::empty({m}, vecs.options()); - launch_po2_gemv( - reinterpret_cast(out.data_ptr()), - reinterpret_cast(vecs.data_ptr()), - reinterpret_cast(values.data_ptr()), - reinterpret_cast(scales.data_ptr()), - reinterpret_cast(LUTs.data_ptr()), - sf_type, group_size, - m, n, k, - at::cuda::getCurrentCUDAStream() - ); + for (int s = 0; s < n; s += 7) { + int t = std::min(s + 7, 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; } From 54268b3727eee9cee59097cb28ab229b3bbc93bc Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Thu, 30 Apr 2026 14:45:05 +0200 Subject: [PATCH 3/7] baseline gemm kernel --- src/po2mm/po2_binding.cpp | 59 +++++++++++++++++ src/po2mm/po2_gemm_kernel.cu | 125 +++++++++++++++++++++++++++++++++++ src/po2mm/po2_gemv.py | 10 +++ src/po2mm/po2_gemv_kernel.cu | 30 --------- src/po2mm/po2_kernels.cuh | 11 +++ src/po2mm/utils.cuh | 34 +++++++++- tests/test_gemm.py | 114 ++++++++++++++++++++++++++++++++ tests/test_gemv.py | 81 +++++++++++++++++++++++ 8 files changed, 433 insertions(+), 31 deletions(-) create mode 100644 src/po2mm/po2_gemm_kernel.cu create mode 100644 tests/test_gemm.py create mode 100644 tests/test_gemv.py diff --git a/src/po2mm/po2_binding.cpp b/src/po2mm/po2_binding.cpp index 8c88624..e5dbd77 100644 --- a/src/po2mm/po2_binding.cpp +++ b/src/po2mm/po2_binding.cpp @@ -66,7 +66,66 @@ torch::Tensor po2_gemv( 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(x.data_ptr()), + reinterpret_cast(values.data_ptr()), + reinterpret_cast(scales.data_ptr()), + reinterpret_cast(LUTs.data_ptr()), + sf_type, group_size, + m, n, k, + at::cuda::getCurrentCUDAStream() + ); + + return out; +} + 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..cbe2d53 --- /dev/null +++ b/src/po2mm/po2_gemm_kernel.cu @@ -0,0 +1,125 @@ +// 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__ 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); + x += j * k; + + 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() + }; + + for (int s = 0; s < k / 16; ++s) { + + // fetch x + for (int t = 0; t < 64; t += 32) { + bf16x8_t v = bf16x8_t::load(x + 16 * s + (threadIdx.x/2 + t) * k + 8 * (threadIdx.x % 2)); + v.store(&x_smem[threadIdx.x/2 + t][(threadIdx.x % 2) * 8]); + } + + // 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) { + for (int v = 0; v < 8; ++v) { + float inner = 0.f; + bf16x8_t a = bf16x8_t::load(&x_smem[8*(threadIdx.x%8) + u][0]); + bf16x8_t b = bf16x8_t::load(&decoded[8*(threadIdx.x/8) + v][0]); + for (int w = 0; w < 8; ++w) { + inner = fma(a[w], b[w], inner); + } + a = bf16x8_t::load(&x_smem[8*(threadIdx.x%8) + u][8]); + b = bf16x8_t::load(&decoded[8*(threadIdx.x/8) + v][8]); + for (int w = 0; w < 8; ++w) { + inner = fma(a[w], b[w], inner); + } + acc[u][v] += inner * scales_smem[8*(threadIdx.x/8) + v]; + } + } + __syncthreads(); + } + + for (int u = 0; u < 8; ++u) { + acc[u].store(out + (j + u + (threadIdx.x%8)*8) * m + i + (threadIdx.x / 8) * 8); + //acc[u].store(out + (i + (threadIdx.x/8)*8 + u) * n + j + (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 || n % 64 != 0) { + std::fprintf(stderr, "po2_gemm_tn_baseline: m (%d) and n (%d) must be multiples of 64\n", m, n); + 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 index aebe76d..7429fbe 100644 --- a/src/po2mm/po2_gemv.py +++ b/src/po2mm/po2_gemv.py @@ -9,6 +9,7 @@ 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=[ @@ -38,3 +39,12 @@ def po2_gemv_fake(vecs, values, scales, LUTs): 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/po2_gemv_kernel.cu b/src/po2mm/po2_gemv_kernel.cu index 3508b92..a134bb2 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 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..ea17610 --- /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.po2_gemv 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", [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..4616138 --- /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.po2_gemv 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 From c82b2ac893df7f6e5bd393f42484c57a9f9ba1ea Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Thu, 30 Apr 2026 14:51:52 +0200 Subject: [PATCH 4/7] renaming --- src/po2mm/__init__.py | 2 +- src/po2mm/{po2_gemv.py => kernels.py} | 0 src/po2mm/linear.py | 2 +- tests/test_gemm.py | 2 +- tests/test_gemv.py | 2 +- 5 files changed, 4 insertions(+), 4 deletions(-) rename src/po2mm/{po2_gemv.py => kernels.py} (100%) 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/po2_gemv.py b/src/po2mm/kernels.py similarity index 100% rename from src/po2mm/po2_gemv.py rename to src/po2mm/kernels.py diff --git a/src/po2mm/linear.py b/src/po2mm/linear.py index e6ebeba..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 diff --git a/tests/test_gemm.py b/tests/test_gemm.py index ea17610..16e8afa 100644 --- a/tests/test_gemm.py +++ b/tests/test_gemm.py @@ -16,7 +16,7 @@ quantize_po2, dequantize_po2, ) -from po2mm.po2_gemv import po2_gemm +from po2mm.kernels import po2_gemm def _inputs_randn(n, in_features, device): diff --git a/tests/test_gemv.py b/tests/test_gemv.py index 4616138..fb24f0b 100644 --- a/tests/test_gemv.py +++ b/tests/test_gemv.py @@ -16,7 +16,7 @@ quantize_po2, dequantize_po2, ) -from po2mm.po2_gemv import po2_gemv +from po2mm.kernels import po2_gemv def _inputs_randn(n, in_features, device): From 863f7572c7b8fa42fde9b603a4d360d8cfb8c090 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sat, 2 May 2026 03:02:33 +0200 Subject: [PATCH 5/7] enable gemm for arbitrary n --- src/po2mm/po2_gemm_kernel.cu | 16 +++++++++------- tests/test_gemm.py | 6 +++--- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/src/po2mm/po2_gemm_kernel.cu b/src/po2mm/po2_gemm_kernel.cu index cbe2d53..f0c1550 100644 --- a/src/po2mm/po2_gemm_kernel.cu +++ b/src/po2mm/po2_gemm_kernel.cu @@ -29,7 +29,6 @@ __global__ void gemm_tn_baseline(nv_bfloat16* __restrict__ out, const nv_bfloat1 scales += i * (k / GroupSize); values += i * (k / 8); - x += j * k; bf16x8_t acc[8] = { bf16x8_t::zeros(), bf16x8_t::zeros(), bf16x8_t::zeros(), bf16x8_t::zeros(), @@ -40,8 +39,10 @@ __global__ void gemm_tn_baseline(nv_bfloat16* __restrict__ out, const nv_bfloat1 // fetch x for (int t = 0; t < 64; t += 32) { - bf16x8_t v = bf16x8_t::load(x + 16 * s + (threadIdx.x/2 + t) * k + 8 * (threadIdx.x % 2)); - v.store(&x_smem[threadIdx.x/2 + t][(threadIdx.x % 2) * 8]); + if (j + threadIdx.x/2 + t < n) { + bf16x8_t v = bf16x8_t::load(x + 16 * s + (j + threadIdx.x/2 + t) * k + 8 * (threadIdx.x % 2)); + v.store(&x_smem[threadIdx.x/2 + t][(threadIdx.x % 2) * 8]); + } } // decode values @@ -77,8 +78,9 @@ __global__ void gemm_tn_baseline(nv_bfloat16* __restrict__ out, const nv_bfloat1 } for (int u = 0; u < 8; ++u) { - acc[u].store(out + (j + u + (threadIdx.x%8)*8) * m + i + (threadIdx.x / 8) * 8); - //acc[u].store(out + (i + (threadIdx.x/8)*8 + u) * n + j + (threadIdx.x%8)*8); + if (j + u + (threadIdx.x%8)*8 < n) { + acc[u].store(out + (j + u + (threadIdx.x%8)*8) * m + i + (threadIdx.x / 8) * 8); + } } } @@ -96,8 +98,8 @@ void launch_po2_gemm_tn_baseline( std::fprintf(stderr, "po2_gemm_tn_baseline: k (%d) must be a multiple of 16\n", k); std::abort(); } - if (m % 64 != 0 || n % 64 != 0) { - std::fprintf(stderr, "po2_gemm_tn_baseline: m (%d) and n (%d) must be multiples of 64\n", m, n); + 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) { diff --git a/tests/test_gemm.py b/tests/test_gemm.py index 16e8afa..6f574de 100644 --- a/tests/test_gemm.py +++ b/tests/test_gemm.py @@ -80,15 +80,15 @@ def test_po2_eye(size, lut_types, scale_type, n, input_factory, group_size): (128, 64), (64, 128), (128, 256), - (1024, 1024), - (768, 3072), + (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", [64, 256]) +@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): From fe0b16e47c0cc9b866e1f6ee4656730d9a6efe93 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sat, 2 May 2026 11:57:24 +0200 Subject: [PATCH 6/7] fix shared load bank conflicts --- src/po2mm/po2_gemm_kernel.cu | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/src/po2mm/po2_gemm_kernel.cu b/src/po2mm/po2_gemm_kernel.cu index f0c1550..0ffcd67 100644 --- a/src/po2mm/po2_gemm_kernel.cu +++ b/src/po2mm/po2_gemm_kernel.cu @@ -9,7 +9,7 @@ #include "po2_kernels.cuh" template -__global__ void gemm_tn_baseline(nv_bfloat16* __restrict__ out, const nv_bfloat16* __restrict__ x, const unsigned* __restrict__ values, +__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"); @@ -19,7 +19,7 @@ __global__ void gemm_tn_baseline(nv_bfloat16* __restrict__ out, const nv_bfloat1 if (i >= m || j >= n) return; __shared__ nv_bfloat16 LUTs_shared[32]; - __shared__ nv_bfloat16 x_smem[64][16]; + __shared__ nv_bfloat16 x_smem[64 * 16]; __shared__ nv_bfloat16 decoded[64][16]; __shared__ float scales_smem[64]; if (threadIdx.x < 32) { @@ -35,13 +35,20 @@ __global__ void gemm_tn_baseline(nv_bfloat16* __restrict__ out, const nv_bfloat1 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 t = 0; t < 64; t += 32) { - if (j + threadIdx.x/2 + t < n) { - bf16x8_t v = bf16x8_t::load(x + 16 * s + (j + threadIdx.x/2 + t) * k + 8 * (threadIdx.x % 2)); - v.store(&x_smem[threadIdx.x/2 + t][(threadIdx.x % 2) * 8]); + 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)); } } @@ -59,17 +66,18 @@ __global__ void gemm_tn_baseline(nv_bfloat16* __restrict__ out, const nv_bfloat1 __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 a = bf16x8_t::load(&x_smem[8*(threadIdx.x%8) + u][0]); bf16x8_t b = bf16x8_t::load(&decoded[8*(threadIdx.x/8) + v][0]); for (int w = 0; w < 8; ++w) { - inner = fma(a[w], b[w], inner); + inner = fma(a1[w], b[w], inner); } - a = bf16x8_t::load(&x_smem[8*(threadIdx.x%8) + u][8]); + b = bf16x8_t::load(&decoded[8*(threadIdx.x/8) + v][8]); for (int w = 0; w < 8; ++w) { - inner = fma(a[w], b[w], inner); + inner = fma(a2[w], b[w], inner); } acc[u][v] += inner * scales_smem[8*(threadIdx.x/8) + v]; } From 3416a74fcaa87a9448cf060dec4a5be2640e5516 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sat, 2 May 2026 13:00:18 +0200 Subject: [PATCH 7/7] enable N=8 gemv --- src/po2mm/po2_binding.cpp | 4 ++-- src/po2mm/po2_gemv_kernel.cu | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/po2mm/po2_binding.cpp b/src/po2mm/po2_binding.cpp index e5dbd77..aa991fe 100644 --- a/src/po2mm/po2_binding.cpp +++ b/src/po2mm/po2_binding.cpp @@ -48,8 +48,8 @@ torch::Tensor po2_gemv( auto out = vecs.dim() == 2 ? torch::empty({n, m}, vecs.options()) : torch::empty({m}, vecs.options()); - for (int s = 0; s < n; s += 7) { - int t = std::min(s + 7, n); + 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), diff --git a/src/po2mm/po2_gemv_kernel.cu b/src/po2mm/po2_gemv_kernel.cu index a134bb2..17ea6fc 100644 --- a/src/po2mm/po2_gemv_kernel.cu +++ b/src/po2mm/po2_gemv_kernel.cu @@ -230,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();