Enable register-codebook W4A4 decode on Ada GPUs and add INT8 output heads

.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

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
@@ -298,7 +318,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 +335,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

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.

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<class T, int RowsPerWarp>
 __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);
@@ -2478,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<float>(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<float>(x[row * k + j]), scale));
+    out[row * k + j] = static_cast<int8_t>(fminf(fmaxf(v, -127.f), 127.f));
+  }
+}
+
+template <class T, int RowsPerWarp>
+__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<const uint32_t *>(w + static_cast<size_t>(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<const uint32_t *>(x + static_cast<size_t>(row + r) * k)[i];
+        sums[r] = __dp4a(static_cast<int>(av), static_cast<int>(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<float>(sums[r]), scale);
+      if (bias) value = __fadd_rn(value, static_cast<float>(bias[col]));
+      out[static_cast<size_t>(row + r) * n + col] = static_cast<T>(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<<<rows, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
+      out.data_ptr<int8_t>(), scales.data_ptr<float>(), x.data_ptr<c10::BFloat16>(), 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<uintptr_t>(w.data_ptr()) & 3) == 0 && (reinterpret_cast<uintptr_t>(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<T, R><<<grid, 128, 0, stream>>>(                                        \
+      out.data_ptr<T>(), x.data_ptr<int8_t>(), x_scales.data_ptr<float>(), w.data_ptr<int8_t>(),          \
+      w_scales.data_ptr<float>(), has_bias ? bias.data_ptr<T>() : 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();
+}

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,41 @@ 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