Works

src/po2mm/__init__.py

@@ -1 +1 @@
-from .po2_gemv import po2_gemv
+from .kernels import po2_gemv, po2_gemm

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)

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)

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<nv_bfloat16*>(out.data_ptr()) + s * out.stride(0),
+            reinterpret_cast<const nv_bfloat16*>(vecs.data_ptr()) + s * vecs.stride(0),
+            reinterpret_cast<const unsigned*>(values.data_ptr()),
+            reinterpret_cast<const std::byte*>(scales.data_ptr()),
+            reinterpret_cast<const nv_bfloat16*>(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<nv_bfloat16*>(out.data_ptr()),
-        reinterpret_cast<const nv_bfloat16*>(vecs.data_ptr()),
+        reinterpret_cast<const nv_bfloat16*>(x.data_ptr()),
         reinterpret_cast<const unsigned*>(values.data_ptr()),
         reinterpret_cast<const std::byte*>(scales.data_ptr()),
         reinterpret_cast<const nv_bfloat16*>(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"));
 }

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 <cuda_bf16.h>
+#include <cstdio>
+#include "vec.cuh"
+#include "utils.cuh"
+#include "po2_kernels.cuh"
+
+template<typename Scaler, int GroupSize=16>
+__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<nv_bfloat16, 8>;
+    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<bf16x8_t&>(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<E4M3Scaling, 16><<<grid, block, 0, stream>>>(out, x, values, scales, LUTs, m, n, k);
+    } else if (sf_type == SFType::G1E7M0 && group_size == 16) {
+        gemm_tn_baseline<E7M0Scaling, 16><<<grid, block, 0, stream>>>(out, x, values, scales, LUTs, m, n, k);
+    } else if (sf_type == SFType::G1E4M3 && group_size == 32) {
+        gemm_tn_baseline<E4M3Scaling, 32><<<grid, block, 0, stream>>>(out, x, values, scales, LUTs, m, n, k);
+    } else if (sf_type == SFType::G1E7M0 && group_size == 32) {
+        gemm_tn_baseline<E7M0Scaling, 32><<<grid, block, 0, stream>>>(out, x, values, scales, LUTs, m, n, k);
+    } else {
+        std::fprintf(stderr, "po2_gemm_tn_baseline: unsupported sf_type (%d)\n", static_cast<int>(sf_type));
+        std::abort();
+    }
+
+    CUDA_CHECK(cudaGetLastError());
+}

src/po2mm/po2_gemv_kernel.cu

@@ -9,36 +9,6 @@
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
 
-
-using ScaleResult = std::pair<float, bool>;
-
-struct E4M3Scaling
-{
-    static __device__ __forceinline__ ScaleResult decode_scale(std::byte raw)
-    {
-        std::uint8_t scale_part = static_cast<std::uint8_t>(raw) & 0b0111'1111u;
-        bool group = (static_cast<std::uint8_t>(raw) & 0b1000'0000u) != 0;
-        __nv_fp8_e4m3 factor;
-        factor.__x = scale_part;
-        return {static_cast<float>(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<std::uint8_t>(raw) & 0b0111'1111u;
-        bool group = (static_cast<std::uint8_t>(raw) & 0b1000'0000u) != 0;
-        float factor = __uint_as_float((static_cast<unsigned>(scale_part) + 63u) << 23u);
-        return {factor, group};
-    }
-};
-
-
 template<int N, typename Scaler, int GroupSize=16>
 __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();

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
+);