From 0a62d10bba8490a9e97d29d1c2fac0757ac3d299 Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Sat, 12 Sep 2026 10:30:23 +0200 Subject: [PATCH 1/4] Enable measured register-codebook decode on Ada GPUs --- .../packed_matmul.cu | 5 ++-- .../tests/test_decode_gemv.py | 28 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu index c9c37e5..f3646b4 100644 --- a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu +++ b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu @@ -34,7 +34,7 @@ __device__ __forceinline__ uint32_t orbitquant_decode_four( } // Register lookup removes shared-table initialization and bank conflicts. -// The host restricts this path to measured SM120 decode shapes. +// The host restricts this path to measured SM89/SM120 decode shapes. template __global__ void orbitquant_packed_w4a4_gemv_register(T* out, const uint8_t* x, const uint8_t* w, const float* xn, const c10::BFloat16* wn, const int8_t* ac, @@ -2060,7 +2060,8 @@ void matmul_packed_w4a4_int8( // Keep the existing layout/tile path for prefill, K-major weights and wide K. if (rows <= 8 && in_features <= 16384 && !weight_k_major && !orbitquant_w4a4_gemv_disabled()) { - if (properties->major == 12 && properties->minor == 0 && + if (((properties->major == 12 && properties->minor == 0) || + (properties->major == 8 && properties->minor == 9)) && out_features >= 2048 && in_features >= 1024) { const int group = rows >= 2 ? 2 : 1; const dim3 register_grid((out_features + 3) / 4, (rows + group - 1) / group); diff --git a/native-kernels/orbitquant-packed-matmul/tests/test_decode_gemv.py b/native-kernels/orbitquant-packed-matmul/tests/test_decode_gemv.py index bf31bf4..bfba01a 100644 --- a/native-kernels/orbitquant-packed-matmul/tests/test_decode_gemv.py +++ b/native-kernels/orbitquant-packed-matmul/tests/test_decode_gemv.py @@ -1,5 +1,7 @@ from __future__ import annotations +import json + import pytest import torch from orbitquant_packed_matmul import matmul_packed_w4a4_int8, supports_device @@ -99,3 +101,29 @@ def device_slice(values): graph.replay() torch.cuda.synchronize() torch.testing.assert_close(captured, actual, rtol=0, atol=0) + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("rows", [1, 2]) +def test_register_decode_dispatch_on_validated_devices(rows, tmp_path): + if not torch.cuda.is_available() or not supports_device("cuda"): + pytest.skip("CUDA kernel required") + if torch.cuda.get_device_capability() not in {(8, 9), (12, 0)}: + pytest.skip("Register decode is enabled only on measured GPU architectures") + x = torch.zeros((rows, 512), device="cuda", dtype=torch.uint8) + w = torch.zeros(2048 * 512, device="cuda", dtype=torch.uint8) + xn = torch.ones(rows, device="cuda") + wn = torch.ones(2048, device="cuda", dtype=torch.bfloat16) + codes = torch.arange(-8, 8, device="cuda", dtype=torch.int8) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA]) as prof: + actual = matmul_packed_w4a4_int8(x, w, xn, wn, codes, codes, + activation_scale=1.0, weight_scale=1.0, out_features=2048, + in_features=1024, output_dtype=torch.bfloat16) + torch.cuda.synchronize() + assert torch.all(actual == 65536) + trace = tmp_path / "dispatch.json" + prof.export_chrome_trace(str(trace)) + kernels = [e["name"] for e in json.loads(trace.read_text())["traceEvents"] + if e.get("cat") == "kernel"] + assert any("orbitquant_packed_w4a4_gemv_register" in name for name in kernels), kernels From 249f18e6005ff79b582918944517b9ba0748021a Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Sat, 12 Sep 2026 10:41:51 +0200 Subject: [PATCH 2/4] Prepare OrbitQuant 0.9.8 and native 1.0.5 for Ada decode --- .github/workflows/build-kernels.yml | 12 ++++----- README.md | 5 ++-- .../tests/test_decode_gemv.py | 26 ++++++++++++++----- pyproject.toml | 2 +- src/orbitquant/__init__.py | 2 +- src/orbitquant/kernels/provision.py | 2 +- tests/test_kernel_provision.py | 7 ++--- uv.lock | 2 +- 8 files changed, 36 insertions(+), 22 deletions(-) diff --git a/.github/workflows/build-kernels.yml b/.github/workflows/build-kernels.yml index e70e5ce..23e7d89 100644 --- a/.github/workflows/build-kernels.yml +++ b/.github/workflows/build-kernels.yml @@ -184,7 +184,7 @@ jobs: "$builder" create-pyproject --unique-id ${{ matrix.unique_id }} -f "$kernel_root" .venv/bin/python "$kernel_root/scripts/prepare_wheel_project.py" \ "$kernel_root" \ - --version "1.0.4+${{ matrix.local_tag }}" \ + --version "1.0.5+${{ matrix.local_tag }}" \ --torch-requirement "${{ matrix.torch_requirement }}" - name: Build CUDA wheel @@ -345,7 +345,7 @@ jobs: & .venv\Scripts\python.exe ` "$kernelRoot\scripts\prepare_wheel_project.py" ` $kernelRoot ` - --version "1.0.4+${{ matrix.local_tag }}" ` + --version "1.0.5+${{ matrix.local_tag }}" ` --torch-requirement "${{ matrix.torch_requirement }}" - name: Build CUDA wheel with MSVC @@ -466,7 +466,7 @@ jobs: "$builder" create-pyproject --unique-id ${{ matrix.unique_id }} -f "$kernel_root" .venv/bin/python "$kernel_root/scripts/prepare_wheel_project.py" \ "$kernel_root" \ - --version "1.0.4+${{ matrix.local_tag }}" \ + --version "1.0.5+${{ matrix.local_tag }}" \ --torch-requirement "${{ matrix.torch_requirement }}" - name: Build Metal wheel @@ -553,7 +553,7 @@ jobs: "$builder" create-pyproject --unique-id rel_cpu_darwin -f "$kernel_root" .venv/bin/python "$kernel_root/scripts/prepare_wheel_project.py" \ "$kernel_root" \ - --version "1.0.4" + --version "1.0.5" - name: Build ABI3 CPU wheel working-directory: native-kernels/orbitquant-packed-matmul @@ -641,7 +641,7 @@ jobs: uv run --no-project --with setuptools python \ "$kernel_root/scripts/prepare_wheel_project.py" \ "$kernel_root" \ - --version "1.0.4" + --version "1.0.5" - name: Build manylinux 2.28 ABI3 wheel env: @@ -738,7 +738,7 @@ jobs: & .venv\Scripts\python.exe ` "$kernelRoot\scripts\prepare_wheel_project.py" ` $kernelRoot ` - --version "1.0.4" + --version "1.0.5" - name: Build ABI3 CPU wheel with MSVC shell: pwsh diff --git a/README.md b/README.md index c9037d7..bcb30d2 100644 --- a/README.md +++ b/README.md @@ -298,7 +298,8 @@ pipe = load_quantized_pipeline_from_artifact( ## CUDA Decode -Version 0.9.7 includes native 1.0.4 register-codebook GEMV for SM120 GPUs. +Version 0.9.8 includes native 1.0.5 register-codebook GEMV for SM89 (Ada) +and SM120 GPUs. For row-major W4A4 inputs with one to eight rows, input width 1024–16384, and at least 2048 output features, it decodes arbitrary 16-entry signed INT8 codebooks with register byte permutations. Two rows share each decoded weight @@ -314,7 +315,7 @@ These are warm-generation measurements on one GPU, not a universal gain or a fresh-process memory comparison. The Python package and native kernel are separate installations. Managed -caches are partitioned by the minimum native release (currently 1.0.4), so an +caches are partitioned by the minimum native release (currently 1.0.5), so an upgrade does not silently reuse an older managed binary. Older cache directories are preserved. Downloaded wheels must satisfy that version floor and match the release checksum; a current cached variant remains usable offline. Provision the diff --git a/native-kernels/orbitquant-packed-matmul/tests/test_decode_gemv.py b/native-kernels/orbitquant-packed-matmul/tests/test_decode_gemv.py index bfba01a..dd41b59 100644 --- a/native-kernels/orbitquant-packed-matmul/tests/test_decode_gemv.py +++ b/native-kernels/orbitquant-packed-matmul/tests/test_decode_gemv.py @@ -115,15 +115,27 @@ def test_register_decode_dispatch_on_validated_devices(rows, tmp_path): xn = torch.ones(rows, device="cuda") wn = torch.ones(2048, device="cuda", dtype=torch.bfloat16) codes = torch.arange(-8, 8, device="cuda", dtype=torch.int8) - with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU, - torch.profiler.ProfilerActivity.CUDA]) as prof: - actual = matmul_packed_w4a4_int8(x, w, xn, wn, codes, codes, - activation_scale=1.0, weight_scale=1.0, out_features=2048, - in_features=1024, output_dtype=torch.bfloat16) + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA] + ) as prof: + actual = matmul_packed_w4a4_int8( + x, + w, + xn, + wn, + codes, + codes, + activation_scale=1.0, + weight_scale=1.0, + out_features=2048, + in_features=1024, + output_dtype=torch.bfloat16, + ) torch.cuda.synchronize() assert torch.all(actual == 65536) trace = tmp_path / "dispatch.json" prof.export_chrome_trace(str(trace)) - kernels = [e["name"] for e in json.loads(trace.read_text())["traceEvents"] - if e.get("cat") == "kernel"] + kernels = [ + e["name"] for e in json.loads(trace.read_text())["traceEvents"] if e.get("cat") == "kernel" + ] assert any("orbitquant_packed_w4a4_gemv_register" in name for name in kernels), kernels diff --git a/pyproject.toml b/pyproject.toml index 92d2904..7ce5ccd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "orbitquant" -version = "0.9.7" +version = "0.9.8" description = "Calibration-free OrbitQuant for transformer linear projections" readme = "README.md" requires-python = ">=3.11" diff --git a/src/orbitquant/__init__.py b/src/orbitquant/__init__.py index a9efb68..8ec3111 100644 --- a/src/orbitquant/__init__.py +++ b/src/orbitquant/__init__.py @@ -18,7 +18,7 @@ from orbitquant.quantizer import register_hf_quantizers from orbitquant.recipes import recipe -__version__ = "0.9.7" +__version__ = "0.9.8" register_hf_quantizers() diff --git a/src/orbitquant/kernels/provision.py b/src/orbitquant/kernels/provision.py index 8485dea..279c772 100644 --- a/src/orbitquant/kernels/provision.py +++ b/src/orbitquant/kernels/provision.py @@ -45,7 +45,7 @@ from packaging.version import Version KERNEL_VERSION = 1 -NATIVE_RELEASE_MINIMUM = "1.0.4" +NATIVE_RELEASE_MINIMUM = "1.0.5" KERNEL_PACKAGE_NAME = "orbitquant_packed_matmul" KERNEL_REPO_ID = "WaveCut/orbitquant-packed-matmul" _RELEASE_TAG = f"kernels-v{KERNEL_VERSION}" diff --git a/tests/test_kernel_provision.py b/tests/test_kernel_provision.py index 9fda89b..e8348ab 100644 --- a/tests/test_kernel_provision.py +++ b/tests/test_kernel_provision.py @@ -116,7 +116,7 @@ def test_provision_ignores_cache_without_marker(monkeypatch): assert report.source == "unavailable" -@pytest.mark.parametrize("namespace", ["", "1.0.3"]) +@pytest.mark.parametrize("namespace", ["", "1.0.3", "1.0.4"]) def test_upgrade_preserves_but_does_not_load_older_managed_cache(monkeypatch, namespace): _patch_runtime(monkeypatch, torch_version="2.12.1") old_root = provision.kernels_cache_root() / f"v{provision.KERNEL_VERSION}" / namespace @@ -243,10 +243,11 @@ def test_provision_rejects_checksum_mismatch(monkeypatch): assert not cache_dir.exists() -def test_provision_rejects_older_release_before_download(monkeypatch): +@pytest.mark.parametrize("version", ["1.0.3", "1.0.4"]) +def test_provision_rejects_older_release_before_download(monkeypatch, version): _patch_runtime(monkeypatch, torch_version="2.12.1") variant = provision.cpu_variant_name() - _serve_release(monkeypatch, variant, _wheel_bytes(), version="1.0.3") + _serve_release(monkeypatch, variant, _wheel_bytes(), version=version) original_get = provision._http_get def manifest_only(url, timeout): diff --git a/uv.lock b/uv.lock index 02ddd99..cc5feaf 100644 --- a/uv.lock +++ b/uv.lock @@ -844,7 +844,7 @@ wheels = [ [[package]] name = "orbitquant" -version = "0.9.7" +version = "0.9.8" source = { editable = "." } dependencies = [ { name = "huggingface-hub" }, From f2a2bac1f770aa1172c6e88662bbe58859067bca Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Sat, 12 Sep 2026 19:59:09 +0200 Subject: [PATCH 3/4] Add INT8 row-quantized output heads (Int8RowLinear, native quantize_rows_int8/matmul_int8_rows) Opt-in per-row absmax INT8 for output heads kept outside the W4A4 packed format: Int8RowLinear / quantize_output_heads in Python, plus two CUDA operators in native 1.0.5 (quantize_rows_int8, matmul_int8_rows: one warp per row, DP4A, 1..8 activation rows). The native kernel and the torch._int_mm fallback share the epilogue float(sum) * (x_scale * w_scale) and derive scales with a true IEEE division, so both paths return identical values (verified on RTX 4090: 81 native + 10 Python tests, bit-exact against _int_mm). --- README.md | 20 ++ docs/kernel-audit.md | 28 +++ .../packed_matmul.cu | 124 ++++++++++++ .../tests/test_int8_rows.py | 87 +++++++++ .../orbitquant_packed_matmul/__init__.py | 52 +++++ .../torch-ext/torch_binding.cpp | 6 + .../torch-ext/torch_binding.h | 11 ++ src/orbitquant/__init__.py | 10 + src/orbitquant/int8_head.py | 179 ++++++++++++++++++ tests/test_int8_head.py | 83 ++++++++ 10 files changed, 600 insertions(+) create mode 100644 native-kernels/orbitquant-packed-matmul/tests/test_int8_rows.py create mode 100644 src/orbitquant/int8_head.py create mode 100644 tests/test_int8_head.py diff --git a/README.md b/README.md index bcb30d2..8ef3e7a 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,26 @@ Embeddings, timestep modules, task heads, and common final projections are kept in source precision by default. Every automatic decision is available as a machine-readable inventory before quantization. +## INT8 output heads (opt-in) + +Output heads stay in source precision by default. For memory-bound decode the head can be +the largest per-token read (a 150k-row BF16 vocabulary head is 0.6 GB per step), so +OrbitQuant also offers an explicit per-row INT8 head: + +```python +from orbitquant import quantize_output_heads + +quantize_output_heads(model, names=("lm_head",)) # in-place Int8RowLinear replacement +``` + +`Int8RowLinear` stores per-row absmax INT8 weights plus FP32 scales (2x smaller than BF16), +quantizes activations per token at forward time, and serves 1..8-row decode with the native +`matmul_int8_rows` DP4A kernel (CUDA package 1.0.5+), falling back to `torch._int_mm` for +larger batches or when the native package is unavailable. All paths share the epilogue +`float(sum) * (x_scale * w_scale)`, so results are identical across them. Head quantization +changes logits at the INT8 rounding level (~0.4% relative); measure sampled outputs before +adopting it for a model. + ## Install ```bash diff --git a/docs/kernel-audit.md b/docs/kernel-audit.md index 8662b94..2e5f455 100644 --- a/docs/kernel-audit.md +++ b/docs/kernel-audit.md @@ -820,3 +820,31 @@ All ten complete PCM arrays matched their corresponding archived control, without truncation or non-finite samples. These results measure generation latency after loading; shared-process NVML does not establish fresh-profile memory use, and other architectures have no performance claim from this run. + +## INT8 row-quantized output heads + +Native 1.0.5 adds two CUDA operators for output heads kept outside the W4A4 +packed format: `quantize_rows_int8` (per-token absmax INT8 activations of a +BF16 `[rows, K]` tensor) and `matmul_int8_rows` (one warp per output row, +four rows per block, DP4A accumulation over per-row absmax INT8 weights, +1–8 activation rows, `K` a multiple of 4, BF16 or FP16 output, optional bias). +The epilogue is `float(sum) * (x_scale * w_scale) (+ bias)`, the same formula +the Python `Int8RowLinear` module uses for its cuBLASLt `torch._int_mm` +fallback, and both derive `scale = absmax / 127` with a true IEEE division +(`__fdiv_rn` in the kernel, a same-device tensor divisor in PyTorch: dividing a +CUDA tensor by a Python scalar makes PyTorch multiply by a reciprocal, which +is 1 ulp off in some rows), so both paths return identical values. No dense dequantized copy or +persistent state is added; larger batches keep the INT8 GEMM fallback. + +`orbitquant.Int8RowLinear` / `quantize_output_heads` are opt-in. They change +logits at INT8 rounding level; the YuE2-3B semantic head (32,769 classes) +measured a mean forward KL of 0.0015 over 128 forced positions with top-1 +agreement 0.914, and the ABC head (151,644 classes) 0.0002 with top-1 +agreement 1.0, both against the BF16 head on RTX 4090. + +On RTX 4090 (Torch 2.10.0+cu128) the 151,644 x 2048 head reads 310 MB per +step in 331 us (937 GB/s) versus 649 us for the BF16 cuBLAS GEMV of the same +rows; the padded `torch._int_mm` fallback was about twice slower than the +native kernel for one row. Native package tests for the new operators compare +against `torch._int_mm` bit for bit across 129–151,644 rows, K of 64–4096, +BF16/FP16 outputs, bias, graph replay, and the contract errors. diff --git a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu index f3646b4..d2e1e1f 100644 --- a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu +++ b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu @@ -2479,3 +2479,127 @@ void quantize_activations_int8( #undef ORBITQUANT_LAUNCH_RPBH_INT8 C10_CUDA_KERNEL_LAUNCH_CHECK(); } + +// --------------------------------------------------------------------------- +// INT8 row-quantized GEMV for output heads (per-row INT8 weights, per-token INT8 +// activations, DP4A accumulation). One warp per output row, four rows per block; +// the epilogue mirrors the torch._int_mm reference: float(sum) * (xs[row] * ws[col]). +// --------------------------------------------------------------------------- +namespace { + +__global__ void orbitquant_quantize_rows_int8_kernel( + int8_t *__restrict__ out, float *__restrict__ scales, const c10::BFloat16 *__restrict__ x, + int64_t k) { + const int64_t row = blockIdx.x; + const int t = threadIdx.x; + __shared__ float red[256]; + float m = 0.f; + for (int64_t j = t; j < k; j += 256) m = fmaxf(m, fabsf(static_cast(x[row * k + j]))); + red[t] = m; + __syncthreads(); + for (int s = 128; s > 0; s >>= 1) { + if (t < s) red[t] = fmaxf(red[t], red[t + s]); + __syncthreads(); + } + // Round-to-nearest intrinsics keep host/PyTorch parity regardless of fast-math or FMA contraction. + const float scale = __fdiv_rn(fmaxf(red[0], 1e-12f), 127.0f); + if (t == 0) scales[row] = scale; + for (int64_t j = t; j < k; j += 256) { + const float v = rintf(__fdiv_rn(static_cast(x[row * k + j]), scale)); + out[row * k + j] = static_cast(fminf(fmaxf(v, -127.f), 127.f)); + } +} + +template +__global__ void __launch_bounds__(128) orbitquant_int8_rows_gemv_kernel( + T *__restrict__ out, const int8_t *__restrict__ x, const float *__restrict__ xs, + const int8_t *__restrict__ w, const float *__restrict__ ws, const T *__restrict__ bias, + int n, int k, int rows) { + const int lane = threadIdx.x & 31; + const int col = blockIdx.x * 4 + (threadIdx.x >> 5); + const int row = blockIdx.y * RowsPerWarp; + if (col >= n) return; + const int words = k / 4; + const uint32_t *wp = reinterpret_cast(w + static_cast(col) * k); + int sums[RowsPerWarp] = {}; + for (int i = lane; i < words; i += 32) { + const uint32_t bv = __ldg(wp + i); +#pragma unroll + for (int r = 0; r < RowsPerWarp; ++r) { + if (row + r < rows) { + const uint32_t av = reinterpret_cast(x + static_cast(row + r) * k)[i]; + sums[r] = __dp4a(static_cast(av), static_cast(bv), sums[r]); + } + } + } +#pragma unroll + for (int r = 0; r < RowsPerWarp; ++r) { +#pragma unroll + for (int offset = 16; offset; offset >>= 1) sums[r] += __shfl_down_sync(0xffffffff, sums[r], offset); + if (lane == 0 && row + r < rows) { + // Same operation order as the torch._int_mm reference: (x_scale * w_scale), multiply, add bias. + const float scale = __fmul_rn(xs[row + r], ws[col]); + float value = __fmul_rn(static_cast(sums[r]), scale); + if (bias) value = __fadd_rn(value, static_cast(bias[col])); + out[static_cast(row + r) * n + col] = static_cast(value); + } + } +} + +} // namespace + +void quantize_rows_int8(torch::Tensor &out, torch::Tensor &scales, torch::Tensor const &x) { + TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2, + "quantize_rows_int8 expects contiguous CUDA BF16 [rows, K]"); + const int64_t rows = x.size(0), k = x.size(1); + TORCH_CHECK(k > 0 && k % 4 == 0, "K must be a positive multiple of 4"); + TORCH_CHECK(out.device() == x.device() && out.is_contiguous() && out.scalar_type() == torch::kChar && + out.sizes() == x.sizes(), "invalid INT8 output"); + TORCH_CHECK(scales.device() == x.device() && scales.is_contiguous() && scales.scalar_type() == torch::kFloat && + scales.numel() == rows, "invalid scales"); + const at::cuda::OptionalCUDAGuard guard(device_of(x)); + if (!rows) return; + orbitquant_quantize_rows_int8_kernel<<>>( + out.data_ptr(), scales.data_ptr(), x.data_ptr(), k); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void matmul_int8_rows(torch::Tensor &out, torch::Tensor const &x, torch::Tensor const &x_scales, + torch::Tensor const &w, torch::Tensor const &w_scales, torch::Tensor const &bias, + bool has_bias) { + TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.scalar_type() == torch::kChar && x.dim() == 2, + "matmul_int8_rows expects contiguous CUDA INT8 activations [rows, K]"); + const int rows = x.size(0), k = x.size(1); + TORCH_CHECK(rows >= 1 && rows <= 8, "matmul_int8_rows serves 1..8 rows; use INT8 GEMM for larger batches"); + TORCH_CHECK(k > 0 && k % 4 == 0, "K must be a positive multiple of 4"); + TORCH_CHECK(w.device() == x.device() && w.is_contiguous() && w.scalar_type() == torch::kChar && w.dim() == 2 && + w.size(1) == k, "weights must be contiguous INT8 [N, K]"); + const int n = w.size(0); + TORCH_CHECK(n >= 1, "N must be positive"); + TORCH_CHECK(w_scales.device() == x.device() && w_scales.is_contiguous() && w_scales.scalar_type() == torch::kFloat && + w_scales.numel() == n, "invalid weight scales"); + TORCH_CHECK(x_scales.device() == x.device() && x_scales.is_contiguous() && x_scales.scalar_type() == torch::kFloat && + x_scales.numel() == rows, "invalid activation scales"); + TORCH_CHECK(out.device() == x.device() && out.is_contiguous() && out.dim() == 2 && out.size(0) == rows && + out.size(1) == n && (out.scalar_type() == torch::kBFloat16 || out.scalar_type() == torch::kHalf), + "invalid output"); + TORCH_CHECK((reinterpret_cast(w.data_ptr()) & 3) == 0 && (reinterpret_cast(x.data_ptr()) & 3) == 0, + "INT8 tensors must be 4-byte aligned"); + if (has_bias) + TORCH_CHECK(bias.device() == x.device() && bias.is_contiguous() && bias.scalar_type() == out.scalar_type() && + bias.numel() == n, "invalid bias"); + const at::cuda::OptionalCUDAGuard guard(device_of(x)); + const auto stream = at::cuda::getCurrentCUDAStream(); + const dim3 grid((n + 3) / 4, rows == 1 ? 1 : (rows + 1) / 2); +#define ORBITQUANT_LAUNCH_INT8_ROWS(T, R) \ + orbitquant_int8_rows_gemv_kernel<<>>( \ + out.data_ptr(), x.data_ptr(), x_scales.data_ptr(), w.data_ptr(), \ + w_scales.data_ptr(), has_bias ? bias.data_ptr() : nullptr, n, k, rows) + if (out.scalar_type() == torch::kBFloat16) { + if (rows == 1) { ORBITQUANT_LAUNCH_INT8_ROWS(c10::BFloat16, 1); } else { ORBITQUANT_LAUNCH_INT8_ROWS(c10::BFloat16, 2); } + } else { + if (rows == 1) { ORBITQUANT_LAUNCH_INT8_ROWS(c10::Half, 1); } else { ORBITQUANT_LAUNCH_INT8_ROWS(c10::Half, 2); } + } +#undef ORBITQUANT_LAUNCH_INT8_ROWS + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/native-kernels/orbitquant-packed-matmul/tests/test_int8_rows.py b/native-kernels/orbitquant-packed-matmul/tests/test_int8_rows.py new file mode 100644 index 0000000..8a09642 --- /dev/null +++ b/native-kernels/orbitquant-packed-matmul/tests/test_int8_rows.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import pytest +import torch +from orbitquant_packed_matmul import matmul_int8_rows, quantize_rows_int8, supports_device + + +def _absmax_scale(values): + # tensor divisor => true IEEE division; a Python-scalar divisor makes torch's CUDA kernel + # multiply by a reciprocal, which is 1 ulp off the kernel's __fdiv_rn in some rows + amax = values.float().abs().amax(-1).clamp_min(1e-12) + return amax / amax.new_full((), 127.0) + + +def _int8_rows(weight): + scale = _absmax_scale(weight) + rows = (weight.float() / scale[:, None]).round().clamp(-127, 127).to(torch.int8) + return rows.contiguous(), scale.contiguous() + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("rows", [1, 2, 3, 8]) +@pytest.mark.parametrize( + "n,k", [(129, 64), (2048, 2048), (32769, 2048), (151644, 2048), (777, 4096)] +) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("bias_enabled", [False, True]) +def test_int8_rows_matches_int_mm(rows, n, k, dtype, bias_enabled): + if not torch.cuda.is_available() or not supports_device("cuda"): + pytest.skip("CUDA kernel required") + torch.manual_seed(rows * 31 + n + k) + x = torch.randn(rows, k, device="cuda").bfloat16() + w = torch.randn(n, k, device="cuda").bfloat16() + w8, ws = _int8_rows(w) + x8, xs = quantize_rows_int8(x) + xf = x.float() + xs_ref = _absmax_scale(xf) + torch.testing.assert_close(xs, xs_ref, rtol=0, atol=0) + x8_ref = (xf / xs_ref[:, None]).round().clamp(-127, 127).to(torch.int8) + assert (x8 != x8_ref).sum().item() <= rows * 2 # rintf vs round on exact ties only + pad_rows = 32 + xpad = torch.zeros((pad_rows, k), device="cuda", dtype=torch.int8) + xpad[:rows] = x8 + npad = (-n) % 8 + wpad = w8 + if npad: + wpad = torch.cat([w8, torch.zeros((npad, k), device="cuda", dtype=torch.int8)], 0) + acc = torch._int_mm(xpad, wpad.t())[:rows, :n] + bias = torch.randn(n, device="cuda").to(dtype) if bias_enabled else None + ref = acc.float() * (xs[:, None] * ws[None, :]) + if bias is not None: + ref = ref + bias.float() + ref = ref.to(dtype) + out = matmul_int8_rows(x8, xs, w8, ws, bias=bias, output_dtype=dtype) + torch.testing.assert_close(out, ref, rtol=0, atol=0) + + +@pytest.mark.kernels_ci +def test_int8_rows_graph_capture_and_contract(): + if not torch.cuda.is_available() or not supports_device("cuda"): + pytest.skip("CUDA kernel required") + torch.manual_seed(5) + x = torch.randn(2, 2048, device="cuda").bfloat16() + w8, ws = _int8_rows(torch.randn(4096, 2048, device="cuda").bfloat16()) + + def run(): + x8, xs = quantize_rows_int8(x) + return matmul_int8_rows(x8, xs, w8, ws) + + ref = run() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = run() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(out, ref, rtol=0, atol=0) + nine = torch.zeros(9, 2048, device="cuda", dtype=torch.int8) + with pytest.raises(RuntimeError): + matmul_int8_rows(nine, torch.ones(9, device="cuda"), w8, ws) + with pytest.raises(RuntimeError): + quantize_rows_int8(torch.randn(1, 2046, device="cuda").bfloat16()) diff --git a/native-kernels/orbitquant-packed-matmul/torch-ext/orbitquant_packed_matmul/__init__.py b/native-kernels/orbitquant-packed-matmul/torch-ext/orbitquant_packed_matmul/__init__.py index 2534ecf..ab448ab 100644 --- a/native-kernels/orbitquant-packed-matmul/torch-ext/orbitquant_packed_matmul/__init__.py +++ b/native-kernels/orbitquant-packed-matmul/torch-ext/orbitquant_packed_matmul/__init__.py @@ -389,3 +389,55 @@ def quantize_activations_int8( quantized.reshape(*original_shape[:-1], dim), norms.reshape(original_shape[:-1]), ) + + +def quantize_rows_int8(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Per-row absmax INT8 quantization of a BF16 [rows, K] tensor (K multiple of 4).""" + if not x.is_cuda: + raise RuntimeError("native INT8 row quantization requires CUDA tensors") + if x.dtype != torch.bfloat16: + raise ValueError("x must be bfloat16") + original_shape = x.shape + values = x.contiguous().reshape(-1, original_shape[-1]) + out = torch.empty(values.shape, device=x.device, dtype=torch.int8) + scales = torch.empty(values.shape[0], device=x.device, dtype=torch.float32) + ops.quantize_rows_int8(out, scales, values) + return out.reshape(original_shape), scales.reshape(original_shape[:-1]) + + +def matmul_int8_rows( + x_int8: torch.Tensor, + x_scales: torch.Tensor, + weight_int8: torch.Tensor, + weight_scales: torch.Tensor, + *, + bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """INT8 activations [rows, K] x INT8 row-quantized weights [N, K] -> [rows, N] (rows 1..8). + + The epilogue computes float(sum) * (x_scale * weight_scale) (+ bias), matching the + torch._int_mm based reference bit for bit. + """ + if not x_int8.is_cuda: + raise RuntimeError("native INT8 row GEMV requires CUDA tensors") + if output_dtype not in {torch.bfloat16, torch.float16}: + raise ValueError("output_dtype must be bfloat16 or float16") + original_shape = x_int8.shape + values = x_int8.contiguous().reshape(-1, original_shape[-1]) + n = weight_int8.shape[0] + out = torch.empty((values.shape[0], n), device=x_int8.device, dtype=output_dtype) + if bias is None: + bias_values = out + else: + bias_values = bias.to(device=x_int8.device, dtype=output_dtype).contiguous() + ops.matmul_int8_rows( + out, + values, + x_scales.to(device=x_int8.device, dtype=torch.float32).contiguous().reshape(-1), + weight_int8.contiguous(), + weight_scales.to(device=x_int8.device, dtype=torch.float32).contiguous(), + bias_values, + bias is not None, + ) + return out.reshape(*original_shape[:-1], n) diff --git a/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.cpp b/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.cpp index d9db6c5..5dc90c1 100644 --- a/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.cpp +++ b/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.cpp @@ -60,6 +60,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "quantize_activations_int8", torch::kCUDA, &quantize_activations_int8); + ops.def("quantize_rows_int8(Tensor! out, Tensor! scales, Tensor x) -> ()"); + ops.impl("quantize_rows_int8", torch::kCUDA, &quantize_rows_int8); + ops.def( + "matmul_int8_rows(Tensor! out, Tensor x, Tensor x_scales, Tensor w, Tensor w_scales, " + "Tensor bias, bool has_bias) -> ()"); + ops.impl("matmul_int8_rows", torch::kCUDA, &matmul_int8_rows); #elif defined(METAL_KERNEL) ops.impl("matmul_packed_weight", torch::kMPS, &matmul_packed_weight); #endif diff --git a/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.h b/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.h index 559db93..83df770 100644 --- a/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.h +++ b/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.h @@ -89,4 +89,15 @@ void quantize_activations_int8( double eps, double inv_sqrt_block, int64_t threads); + +void quantize_rows_int8(torch::Tensor &out, torch::Tensor &scales, torch::Tensor const &x); + +void matmul_int8_rows( + torch::Tensor &out, + torch::Tensor const &x, + torch::Tensor const &x_scales, + torch::Tensor const &w, + torch::Tensor const &w_scales, + torch::Tensor const &bias, + bool has_bias); #endif diff --git a/src/orbitquant/__init__.py b/src/orbitquant/__init__.py index 8ec3111..fd7148e 100644 --- a/src/orbitquant/__init__.py +++ b/src/orbitquant/__init__.py @@ -1,6 +1,12 @@ """Calibration-free OrbitQuant for transformer linear projections.""" from orbitquant.config import OrbitQuantConfig +from orbitquant.int8_head import ( + Int8RowLinear, + absmax_int8_scales, + quantize_int8_rows, + quantize_output_heads, +) from orbitquant.layers import OrbitQuantLinear from orbitquant.linear_adapters import register_linear_adapter from orbitquant.modeling import ( @@ -37,4 +43,8 @@ "register_linear_adapter", "register_hf_quantizers", "save_quantized_pipeline_component", + "Int8RowLinear", + "absmax_int8_scales", + "quantize_int8_rows", + "quantize_output_heads", ] diff --git a/src/orbitquant/int8_head.py b/src/orbitquant/int8_head.py new file mode 100644 index 0000000..3396ce0 --- /dev/null +++ b/src/orbitquant/int8_head.py @@ -0,0 +1,179 @@ +"""INT8 row-quantized output heads. + +OrbitQuant keeps output heads (``lm_head`` and similar) in source precision by default, +because their logits feed sampling directly. For memory-bound decode the head is often the +single largest BF16 read per token (a 150k-row vocabulary head is 0.6 GB per step). This +module offers an explicit, opt-in per-row INT8 representation: + +* weights: per-row absmax INT8 (``weight_int8 [N, K]`` + ``scales [N]`` FP32), +* activations: per-token absmax INT8 at forward time, +* decode (1..8 rows): the native ``matmul_int8_rows`` DP4A kernel when the CUDA package is + available, otherwise a cuBLASLt ``torch._int_mm`` fallback with padded rows, +* larger batches: ``torch._int_mm`` (CUDA) or a dequantized float matmul. + +The INT8 GEMM epilogue is ``float(sum) * (x_scale * w_scale)`` in every path, and every path +derives ``scale = absmax / 127`` with a true IEEE division (the native kernel uses +``__fdiv_rn``; the PyTorch paths divide by a same-device tensor, because dividing by a Python +scalar makes PyTorch's CUDA kernel multiply by a reciprocal, which can be 1 ulp off). The native +kernel and the ``_int_mm`` fallback therefore agree bit for bit. +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch +from torch import nn + +__all__ = ["Int8RowLinear", "absmax_int8_scales", "quantize_int8_rows", "quantize_output_heads"] + + +def absmax_int8_scales(values: torch.Tensor) -> torch.Tensor: + """Per-row ``absmax / 127`` in float32 with a true IEEE division (see module docstring).""" + amax = values.to(torch.float32).abs().amax(-1).clamp_min(1e-12) + return amax / amax.new_full((), 127.0) + + +def quantize_int8_rows(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Per-row absmax INT8 quantization: returns ``(weight_int8 [N, K], scales [N] float32)``.""" + if weight.ndim != 2: + raise ValueError("weight must be a rank-2 [out_features, in_features] tensor") + values = weight.detach().to(torch.float32) + scales = absmax_int8_scales(values) + rows = (values / scales[:, None]).round().clamp(-127, 127).to(torch.int8) + return rows.contiguous(), scales.contiguous() + + +def _native_int8_rows(): + try: + from orbitquant.kernels.native_packed_matmul import load_native_packed_matmul_kernel + + kernel = load_native_packed_matmul_kernel() + except Exception: # pragma: no cover - depends on the provisioned binary + return None + if not hasattr(kernel, "matmul_int8_rows") or not hasattr(kernel, "quantize_rows_int8"): + return None + return kernel + + +class Int8RowLinear(nn.Module): + """``nn.Linear`` replacement holding per-row INT8 weights; forward returns ``x.dtype``.""" + + def __init__( + self, + weight_int8: torch.Tensor, + scales: torch.Tensor, + bias: torch.Tensor | None = None, + *, + module_name: str = "", + ) -> None: + super().__init__() + if weight_int8.dtype != torch.int8 or weight_int8.ndim != 2: + raise ValueError("weight_int8 must be an INT8 [out_features, in_features] tensor") + if scales.numel() != weight_int8.shape[0]: + raise ValueError("scales must have one entry per output row") + self.out_features, self.in_features = (int(s) for s in weight_int8.shape) + if self.in_features % 4: + raise ValueError("in_features must be a multiple of 4") + self.register_buffer("weight_int8", weight_int8.contiguous()) + self.register_buffer("scales", scales.to(torch.float32).contiguous()) + if bias is not None: + self.register_buffer("bias", bias.detach().clone()) + else: + self.bias = None + self.module_name = module_name + self.last_path: str | None = None + + @classmethod + def from_linear(cls, linear: nn.Module, *, module_name: str = "") -> Int8RowLinear: + weight = getattr(linear, "weight", None) + if weight is None or weight.ndim != 2: + raise TypeError("from_linear expects a module with a rank-2 weight (nn.Linear layout)") + rows, scales = quantize_int8_rows(weight) + bias = getattr(linear, "bias", None) + return cls(rows, scales, None if bias is None else bias.detach(), module_name=module_name) + + def extra_repr(self) -> str: + return ( + f"in_features={self.in_features}, out_features={self.out_features}, " + f"bias={self.bias is not None}" + ) + + def _int_mm_pad(self) -> tuple[torch.Tensor, int]: + pad = (-self.out_features) % 8 + if not pad: + return self.weight_int8, 0 + cached = getattr(self, "_padded_weight", None) + if cached is None or cached.device != self.weight_int8.device: + zeros = torch.zeros( + (pad, self.in_features), dtype=torch.int8, device=self.weight_int8.device + ) + cached = torch.cat([self.weight_int8, zeros], 0) + self._padded_weight = cached + return cached, pad + + def forward(self, x: torch.Tensor) -> torch.Tensor: + lead = x.shape[:-1] + x2 = x.reshape(-1, self.in_features) + rows = x2.shape[0] + out_dtype = x.dtype if x.dtype in {torch.bfloat16, torch.float16} else torch.bfloat16 + if x.device.type == "cuda": + xf = x2.to(torch.float32) + x_scales = absmax_int8_scales(xf) + kernel = None + if 1 <= rows <= 8: + kernel = _native_int8_rows() + if kernel is not None and x2.dtype == torch.bfloat16: + x_int8, x_scales = kernel.quantize_rows_int8(x2.contiguous()) + bias = None if self.bias is None else self.bias.to(out_dtype) + out = kernel.matmul_int8_rows( + x_int8, + x_scales, + self.weight_int8, + self.scales, + bias=bias, + output_dtype=out_dtype, + ) + self.last_path = "native_int8_rows" + return out.reshape(*lead, self.out_features) + x_int8 = (xf / x_scales[:, None]).round().clamp(-127, 127).to(torch.int8) + if callable(getattr(torch, "_int_mm", None)): + weight, pad = self._int_mm_pad() + pad_rows = max(32, rows) + if pad_rows != rows: + zeros = torch.zeros( + (pad_rows - rows, self.in_features), dtype=torch.int8, device=x.device + ) + x_int8 = torch.cat([x_int8, zeros], 0) + acc = torch._int_mm(x_int8, weight.t())[:rows, : self.out_features] + out = acc.to(torch.float32) * (x_scales[:, None] * self.scales[None, :]) + self.last_path = "int_mm" + else: # pragma: no cover - very old torch + dequantized = self.weight_int8.float() * self.scales[:, None] + out = (x_int8.float() * x_scales[:, None]) @ dequantized.t() + self.last_path = "float_fallback" + else: + xf = x2.to(torch.float32) + x_scales = absmax_int8_scales(xf) + x_int8 = (xf / x_scales[:, None]).round().clamp(-127, 127) + acc = x_int8 @ (self.weight_int8.to(torch.float32)).t() + out = acc * (x_scales[:, None] * self.scales[None, :]) + self.last_path = "cpu_float" + if self.bias is not None: + out = out + self.bias.to(torch.float32) + return out.to(out_dtype).reshape(*lead, self.out_features) + + +def quantize_output_heads( + model: nn.Module, names: Iterable[str] = ("lm_head",) +) -> dict[str, Int8RowLinear]: + """Replace the named ``nn.Linear`` heads of ``model`` by :class:`Int8RowLinear` in place.""" + replaced: dict[str, Int8RowLinear] = {} + for name in names: + parent_name, _, leaf = name.rpartition(".") + parent = model.get_submodule(parent_name) if parent_name else model + original = getattr(parent, leaf) + head = Int8RowLinear.from_linear(original, module_name=name) + setattr(parent, leaf, head) + replaced[name] = head + return replaced diff --git a/tests/test_int8_head.py b/tests/test_int8_head.py new file mode 100644 index 0000000..7d323bb --- /dev/null +++ b/tests/test_int8_head.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import pytest +import torch +from torch import nn + +from orbitquant import ( + Int8RowLinear, + absmax_int8_scales, + quantize_int8_rows, + quantize_output_heads, +) + + +def test_absmax_int8_scales_is_true_division(): + values = torch.randn(64, 2048) + scales = absmax_int8_scales(values) + ref = (values.abs().amax(-1).clamp_min(1e-12).double() / 127.0).float() + torch.testing.assert_close(scales, ref, rtol=0, atol=0) + assert scales.dtype == torch.float32 and scales.shape == (64,) + + +def test_quantize_int8_rows_shapes_and_range(): + weight = torch.randn(37, 64) + rows, scales = quantize_int8_rows(weight) + assert rows.dtype == torch.int8 and rows.shape == (37, 64) and scales.shape == (37,) + assert rows.abs().max().item() <= 127 + recon = rows.float() * scales[:, None] + step = (weight.abs().amax(-1) / 127.0).max().item() + assert (recon - weight).abs().max().item() <= step * 0.5 + 1e-6 + + +@pytest.mark.parametrize("rows", [1, 3, 40]) +def test_int8_head_cpu_close_to_linear(rows): + torch.manual_seed(1) + linear = nn.Linear(64, 200) + head = Int8RowLinear.from_linear(linear, module_name="lm_head") + x = torch.randn(rows, 64) + ref = linear(x) + out = head(x) + assert out.dtype == torch.bfloat16 and out.shape == ref.shape + err = (out.float() - ref).abs().max().item() + assert err < 0.05 * ref.abs().max().item() + assert head.last_path == "cpu_float" + + +def test_quantize_output_heads_replaces_in_place(): + class Tiny(nn.Module): + def __init__(self): + super().__init__() + self.body = nn.Linear(16, 16) + self.lm_head = nn.Linear(16, 50, bias=False) + + model = Tiny() + replaced = quantize_output_heads(model, ("lm_head",)) + assert isinstance(model.lm_head, Int8RowLinear) and "lm_head" in replaced + assert model.lm_head.bias is None and model.lm_head.out_features == 50 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("rows", [1, 2, 8, 33]) +def test_int8_head_cuda_paths_agree(rows): + torch.manual_seed(7) + linear = nn.Linear(2048, 4104, bias=True).cuda().bfloat16() + head = Int8RowLinear.from_linear(linear, module_name="lm_head").cuda() + x = torch.randn(rows, 2048, device="cuda").bfloat16() + out = head(x) + assert out.dtype == torch.bfloat16 and out.shape == (rows, 4104) + # explicit _int_mm reference with the same epilogue formula + xf = x.float() + xs = absmax_int8_scales(xf) + x8 = (xf / xs[:, None]).round().clamp(-127, 127).to(torch.int8) + pad = torch.zeros((max(32, rows), 2048), device="cuda", dtype=torch.int8) + pad[:rows] = x8 + acc = torch._int_mm(pad, head.weight_int8.t())[:rows] + scaled = acc.float() * (xs[:, None] * head.scales[None, :]) + head.bias.float() + ref = scaled.to(torch.bfloat16) + if head.last_path == "native_int8_rows": + torch.testing.assert_close(out, ref, rtol=0, atol=0) + else: + torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-2) + ref_linear = linear(x).float() + assert (out.float() - ref_linear).abs().max().item() < 0.05 * ref_linear.abs().max().item() From e5cb66afc979fef607b239b06639c15700cc4bf7 Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:03:18 +0200 Subject: [PATCH 4/4] Extend the kernel-builder guard for the INT8 row operators The CUDA source now carries 14 kernel launches with launch checks (the two INT8 row operators added two); the guard also names both new kernels. --- tests/test_ci.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_ci.py b/tests/test_ci.py index 002bb97..741ca7f 100644 --- a/tests/test_ci.py +++ b/tests/test_ci.py @@ -339,10 +339,12 @@ def test_native_packed_matmul_kernel_package_stays_kernel_builder_abi3_compliant package_root / "orbitquant_packed_matmul_cuda/packed_matmul.cu" ).read_text(encoding="utf-8") assert "#include " in cuda_source - assert cuda_source.count("C10_CUDA_KERNEL_LAUNCH_CHECK();") == 12 + assert cuda_source.count("C10_CUDA_KERNEL_LAUNCH_CHECK();") == 14 assert "orbitquant_packed_matmul_small_rows_kernel" in cuda_source assert "orbitquant_rpbh_quantize_pack_w4_kernel" in cuda_source assert "orbitquant_rpbh_quantize_int8_kernel" in cuda_source + assert "orbitquant_quantize_rows_int8_kernel" in cuda_source + assert "orbitquant_int8_rows_gemv_kernel" in cuda_source package_api = ( package_root / "torch-ext/orbitquant_packed_matmul/__init__.py"