Works

[ngc92]

a working configuration

[Copilot]

Pull request overview

This PR consolidates CUDA extension loading into a new po2mm.kernels module, extends the GEMV path to handle larger batch sizes by chunking, and introduces a baseline quantized GEMM (TN) kernel with accompanying tests.

Changes:

• Add po2_gemm CUDA kernel + PyTorch binding and expose it via po2mm.kernels / po2mm.__init__.
• Update GEMV binding to launch in chunks of 8 vectors (and extend kernel dispatch to n==8).
• Add new GEMV/GEMM quality tests, including outlier-heavy activation inputs.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
tests/test_gemv.py	New parametrized GEMV quality tests (incl. outlier input distribution).
tests/test_gemm.py	New GEMM quality tests and an “eye” sanity test.
src/po2mm/utils.cuh	Moves scaling decode helpers into a shared header and generalizes LUT decode helper.
src/po2mm/po2_kernels.cuh	Declares the new GEMM launch function.
src/po2mm/po2_gemv.py	Removes the old per-op extension loader (replaced by kernels.py).
src/po2mm/po2_gemv_kernel.cu	Extends dispatch to n==8 and relies on shared scaling helpers.
src/po2mm/po2_gemm_kernel.cu	Adds baseline GEMM(TN) kernel implementation and launcher.
src/po2mm/po2_binding.cpp	Chunks GEMV launches by 8 and adds a po2_gemm binding.
src/po2mm/linear.py	Switches GEMV import to po2mm.kernels and removes manual chunking.
src/po2mm/kernels.py	New unified extension loader and torch.library.custom_op wrappers for GEMV/GEMM.
src/po2mm/init.py	Re-exports po2_gemv and newly added po2_gemm.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.