Sync the source snapshot with iamwavecut/OrbitQuant@2d586f189ab8

Tracked kernel package sources mirrored 1:1 from native-kernels/orbitquant-packed-matmul at 2d586f189ab8c44ca10685a64e71aaab10c61430. Adds the CPU backend (AVX2/AVX-512/NEON + stable torch ABI), the CUDA W4A4 int8 runtime ops with the cold-L2 pipelined mma64 path, and the Metal skinny-batch GEMV; no generated build files.

CARD.md

@@ -28,29 +28,90 @@ Inputs:
 
 - `x`: contiguous or reshapeable tensor with shape `[..., in_features]`.
 - `packed_weight_indices`: `uint8` low-bit packed row-major codebook indices.
-- `row_norms`: `float32` row norms with shape `[out_features]`.
+- `row_norms`: row norms with shape `[out_features]`; CUDA consumes the
+  artifact's `bfloat16` values directly, while Metal uses `float32` internally.
 - `centroids`: `float32` Lloyd-Max centroids with shape `[2**bits]`.
 - `bias`: optional projection bias.
 
 `x` may be `float32`, `float16`, or `bfloat16`. The output has shape
 `[..., out_features]` and the same dtype as `x`.
 
+The CUDA package also exports the operations used by OrbitQuant's W4A4 runtime:
+
+- `quantize_activations_int8`: token norm, RPBH/FWHT, nearest-codebook
+  assignment, and INT8-surrogate output in one native launch.
+- `quantize_activations_packed_w4`: the same activation path with packed 4-bit
+  output for the direct packed matmul fallback.
+- `matmul_packed_w4a4_int8`: direct packed A4/W4 CUDA MMA with fused token norm,
+  row norm, surrogate scales, and bias epilogue.
+
+On CUDA compute capability 8.0 or newer, OrbitQuant normally combines
+`quantize_activations_int8` with chunked packed-weight decode and Torch's
+CUTLASS-backed INT8 matmul. The direct packed MMA operation remains available
+when that path is unsupported. Neither path materializes a complete BF16/FP16
+weight matrix.
+
 ## Build And Test
 
 ```bash
-nix --option sandbox relaxed run .#build-and-copy -L
-nix --option sandbox relaxed run .#ci-test -L
+nix --option sandbox relaxed --option max-jobs 1 --option cores 8 \
+  run .#build-and-copy -L
+nix --option sandbox relaxed --option max-jobs 1 --option cores 8 \
+  run .#ci-test -L
 ```
 
 The build produces ABI3 Hugging Face Kernels artifacts under `build/` for the
 supported backend variants on the current platform. On macOS, `sandbox relaxed`
-or enabled Nix sandboxing is required by `kernel-builder`.
+or enabled Nix sandboxing is required by `kernel-builder`. The commands build
+local files only; they do not upload to Kernel Hub.
+
+For a faster CUDA-only development build on a machine with a matching Torch and
+CUDA toolchain:
+
+```bash
+cargo install --git https://github.com/huggingface/kernels hf-kernel-builder
+kernel-builder check-config .
+kernel-builder create-pyproject -f .
+TORCH_CUDA_ARCH_LIST="8.9" CUDACXX=/usr/local/cuda/bin/nvcc \
+  python setup.py build_kernel
+```
+
+For a local Metal build compatible with macOS 15 and newer:
+
+```bash
+cargo install --git https://github.com/huggingface/kernels hf-kernel-builder
+kernel-builder check-config .
+kernel-builder create-pyproject -f .
+MACOSX_DEPLOYMENT_TARGET=15.0 \
+  CMAKE_ARGS="-DCMAKE_OSX_DEPLOYMENT_TARGET=15.0" \
+  python setup.py build_kernel
+kernel-builder check-abi --macos 15.0 --python-abi 3.9 .
+```
+
+This generated project is for local testing and must not be committed or
+distributed without a successful `kernel-builder check-abi`. Use the Nix build
+for redistributable variants.
 
 For direct local imports, add the matching `build/torch*-<backend>-<platform>`
 directory to `PYTHONPATH`; the `torch*` variant must match the runtime PyTorch
-version. For Hugging Face `kernels` local loading, set `LOCAL_KERNELS` to the
-same built variant directory containing `metadata.json`, not to the source
-package root:
+version:
+
+```bash
+export PYTHONPATH="/path/to/build/torch212-metal-aarch64-darwin:$PYTHONPATH"
+python -c "import orbitquant_packed_matmul; print(orbitquant_packed_matmul)"
+```
+
+For PyTorch 2.9 CUDA inference, set the allocator before starting Python when
+minimum reserved memory is important:
+
+```bash
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python generate.py
+```
+
+OrbitQuant detects that importable package before trying any Hub loader. For
+Hugging Face `kernels` local loading instead, set `LOCAL_KERNELS` to the same
+built variant directory containing `metadata.json`, not to the source package
+root:
 
 ```bash
 export LOCAL_KERNELS="WaveCut/orbitquant-packed-matmul=/path/to/build/torch212-metal-aarch64-darwin"
@@ -77,9 +138,15 @@ python benchmarks/benchmark.py \
   --iters 20
 ```
 
+`--rows` accepts a comma-separated sweep (the default covers decode-bound
+small batches and GEMM-bound large batches), and `--dtype` selects the
+activation dtype explicitly. Headline timings are hot-loop medians.
+
 The script prints JSON with `packed_seconds_per_iter`,
 `predequantized_f_linear_seconds_per_iter`,
-`dequantize_then_f_linear_seconds_per_iter`,
+`dequantize_then_f_linear_seconds_per_iter` (all hot-loop medians), the
+per-path `*_hot_mean_seconds`, `*_hot_median_seconds`, and
+`*_hot_p95_seconds` distributions,
 `packed_vs_predequantized_f_linear_speedup`,
 `packed_vs_dequantize_then_f_linear_speedup`, compatibility aliases
 `reference_seconds_per_iter` and `packed_vs_reference_speedup`, and
@@ -91,3 +158,31 @@ It also reports storage accounting for the packed weight path:
 `packed_weight_path_vs_materialized_weight_ratio`. These values describe only
 the weight-side storage used by this operator; they are not end-to-end model
 VRAM measurements.
+
+### Metal reference results
+
+Measured on an Apple M2 Max with Torch 2.12.1, FP16 activations, W4 packed
+weights, transformer-scale shapes (hot-loop medians over 30 iterations; each
+iteration synchronizes, so sub-millisecond rows include the MPS submit
+floor):
+
+| Shape (rows x in x out) | Packed Metal | Resident FP16 `F.linear` | Materialize + `F.linear` | Packed vs resident | Packed vs materialize |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| 1 x 3072 x 3072 | 0.330 ms | 0.200 ms | 1.791 ms | 0.60x | 5.42x |
+| 4 x 3072 x 3072 | 0.377 ms | 0.215 ms | 1.768 ms | 0.57x | 4.69x |
+| 32 x 3072 x 3072 | 0.367 ms | 0.312 ms | 1.700 ms | 0.85x | 4.63x |
+| 512 x 3072 x 3072 | 1.355 ms | 1.077 ms | 2.331 ms | 0.79x | 1.72x |
+| 512 x 3072 x 12288 | 5.138 ms | 3.907 ms | 9.312 ms | 0.76x | 1.81x |
+| 1 x 3072 x 12288 | 0.417 ms | 0.383 ms | 6.202 ms | 0.92x | 14.88x |
+
+Batches of at most four rows dispatch a skinny-batch GEMV (one simdgroup per
+output column, decoded weight segments reused across the batch); larger
+batches dispatch the simdgroup-matrix tiles. The packed weight payload, row
+norms, and centroids occupy about 25% of the materialized FP16 weight size at
+W4. The resident reference excludes weight materialization time and retains
+the complete FP16 matrix in memory — it is the throughput ceiling for a
+kernel that decodes weights on the fly, not a like-for-like memory
+configuration.
+
+End-to-end FLUX.2 Klein 9B measurements and the SDNQ comparison are recorded in
+[`docs/flux2-klein-9b-sdnq-vs-orbitquant.md`](../../docs/flux2-klein-9b-sdnq-vs-orbitquant.md).

benchmarks/benchmark.py

@@ -2,6 +2,9 @@ from __future__ import annotations
 
 import argparse
 import json
+import os
+import platform
+import statistics
 import time
 
 import torch
@@ -28,48 +31,57 @@ def _synchronize(device: str) -> None:
         torch.mps.synchronize()
 
 
-def _time_call(device: str, iters: int, fn) -> float:
+def _time_call(device: str, fn) -> float:
     _synchronize(device)
-    start = time.perf_counter()
-    for _ in range(iters):
-        fn()
+    start = time.perf_counter_ns()
+    fn()
     _synchronize(device)
-    return (time.perf_counter() - start) / iters
+    return (time.perf_counter_ns() - start) / 1_000_000_000
 
 
-def main() -> None:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--device", choices=["cuda", "mps"], default="cuda")
-    parser.add_argument("--bits", type=int, default=4)
-    parser.add_argument("--rows", type=int, default=4096)
-    parser.add_argument("--in-features", type=int, default=3072)
-    parser.add_argument("--out-features", type=int, default=3072)
-    parser.add_argument("--iters", type=int, default=20)
-    parser.add_argument("--warmup", type=int, default=3)
-    parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument("--with-bias", action="store_true")
-    args = parser.parse_args()
+def _time_distribution(device: str, iters: int, fn) -> dict[str, float]:
+    samples = []
+    for _ in range(iters):
+        samples.append(_time_call(device, fn))
+    samples.sort()
+    return {
+        "mean": statistics.fmean(samples),
+        "median": statistics.median(samples),
+        "p95": samples[min(len(samples) - 1, int(len(samples) * 0.95))],
+    }
 
-    torch.manual_seed(args.seed)
-    dtype = torch.float16 if args.device == "mps" else torch.bfloat16
-    x = torch.randn(args.rows, args.in_features, device=args.device, dtype=dtype)
-    indices = torch.randint(
-        0,
-        2**args.bits,
-        (args.out_features, args.in_features),
-        dtype=torch.uint8,
-    )
-    packed = _pack(indices, args.bits).to(args.device)
-    row_norms = torch.ones(args.out_features, device=args.device)
-    centroids = torch.linspace(-1.0, 1.0, 2**args.bits, device=args.device)
+
+def _parse_rows(raw: str) -> list[int]:
+    rows_values = []
+    for chunk in raw.split(","):
+        chunk = chunk.strip()
+        if not chunk:
+            continue
+        rows = int(chunk)
+        if rows <= 0:
+            raise argparse.ArgumentTypeError("--rows values must be positive")
+        rows_values.append(rows)
+    if not rows_values:
+        raise argparse.ArgumentTypeError("--rows must list at least one row count")
+    return rows_values
+
+
+_DTYPES = {
+    "float32": torch.float32,
+    "float16": torch.float16,
+    "bfloat16": torch.bfloat16,
+}
+
+
+def _benchmark_rows(args, rows: int, dtype: torch.dtype, weights) -> dict:
+    packed, row_norms, centroids, indices_device = weights
+    x = torch.randn(rows, args.in_features, device=args.device, dtype=dtype)
     bias = (
         torch.randn(args.out_features, device=args.device, dtype=dtype)
         if args.with_bias
         else None
     )
 
-    indices_device = indices.long().to(args.device)
-
     def materialize_reference_weight() -> torch.Tensor:
         return (row_norms[:, None] * centroids[indices_device]).to(dtype)
 
@@ -98,41 +110,79 @@ def main() -> None:
     def dequantize_then_linear_call() -> torch.Tensor:
         return torch.nn.functional.linear(x, materialize_reference_weight(), bias)
 
+    packed_first_call_seconds = _time_call(args.device, packed_call)
+    predequantized_first_call_seconds = _time_call(args.device, predequantized_linear_call)
+    dequantize_then_first_call_seconds = _time_call(args.device, dequantize_then_linear_call)
+
     for _ in range(args.warmup):
         packed_call()
         predequantized_linear_call()
         dequantize_then_linear_call()
-    packed_seconds = _time_call(args.device, args.iters, packed_call)
-    predequantized_linear_seconds = _time_call(
+    packed_distribution = _time_distribution(args.device, args.iters, packed_call)
+    predequantized_distribution = _time_distribution(
         args.device,
         args.iters,
         predequantized_linear_call,
     )
-    dequantize_then_linear_seconds = _time_call(
+    dequantize_then_distribution = _time_distribution(
         args.device,
         args.iters,
         dequantize_then_linear_call,
     )
+    # Headline numbers are hot-loop medians; the mean is retained alongside the
+    # median/p95 so noisy first-iteration outliers cannot skew comparisons.
+    packed_seconds = packed_distribution["median"]
+    predequantized_linear_seconds = predequantized_distribution["median"]
+    dequantize_then_linear_seconds = dequantize_then_distribution["median"]
 
     packed_output = packed_call()
     reference_output = predequantized_linear_call()
     _synchronize(args.device)
-    max_abs_error = (packed_output.float() - reference_output.float()).abs().max().item()
+    error = packed_output.float() - reference_output.float()
+    max_abs_error = error.abs().max().item()
+    rmse = error.square().mean().sqrt().item()
+    reference_rms = reference_output.float().square().mean().sqrt().item()
+    relative_rmse = rmse / max(reference_rms, 1e-12)
 
-    payload = {
+    return {
         "device": args.device,
-        "device_name": torch.cuda.get_device_name(0) if args.device == "cuda" else "mps",
+        "device_name": (
+            torch.cuda.get_device_name(0)
+            if args.device == "cuda"
+            else "mps"
+            if args.device == "mps"
+            else f"{platform.processor() or platform.machine()} "
+            f"({torch.backends.cpu.get_cpu_capability()})"
+        ),
         "dtype": str(dtype).replace("torch.", ""),
         "bits": args.bits,
-        "rows": args.rows,
+        "rows": rows,
         "in_features": args.in_features,
         "out_features": args.out_features,
         "iters": args.iters,
         "warmup": args.warmup,
+        "threads": (
+            os.environ.get("ORBITQUANT_CPU_THREADS", "runtime default")
+            if args.device == "cpu"
+            else None
+        ),
+        "torch_threads": torch.get_num_threads() if args.device == "cpu" else None,
         "with_bias": args.with_bias,
         "packed_seconds_per_iter": packed_seconds,
+        "packed_first_call_seconds": packed_first_call_seconds,
+        "packed_hot_mean_seconds": packed_distribution["mean"],
+        "packed_hot_median_seconds": packed_distribution["median"],
+        "packed_hot_p95_seconds": packed_distribution["p95"],
         "predequantized_f_linear_seconds_per_iter": predequantized_linear_seconds,
+        "predequantized_first_call_seconds": predequantized_first_call_seconds,
+        "predequantized_hot_mean_seconds": predequantized_distribution["mean"],
+        "predequantized_hot_median_seconds": predequantized_distribution["median"],
+        "predequantized_hot_p95_seconds": predequantized_distribution["p95"],
         "dequantize_then_f_linear_seconds_per_iter": dequantize_then_linear_seconds,
+        "dequantize_then_first_call_seconds": dequantize_then_first_call_seconds,
+        "dequantize_then_hot_mean_seconds": dequantize_then_distribution["mean"],
+        "dequantize_then_hot_median_seconds": dequantize_then_distribution["median"],
+        "dequantize_then_hot_p95_seconds": dequantize_then_distribution["p95"],
         "packed_weight_indices_bytes": packed_weight_indices_bytes,
         "row_norms_bytes": row_norms_bytes,
         "centroid_bytes": centroid_bytes,
@@ -155,6 +205,9 @@ def main() -> None:
         if packed_seconds > 0
         else None,
         "max_abs_error": max_abs_error,
+        "rmse": rmse,
+        "relative_rmse": relative_rmse,
+        "timing_headline": "hot-loop median seconds per iteration",
         "reference": (
             "predequantized PyTorch F.linear over a materialized dequantized "
             "weight matrix"
@@ -163,7 +216,64 @@ def main() -> None:
             "materialize the dequantized weight matrix, then call PyTorch F.linear"
         ),
     }
-    print(json.dumps(payload, indent=2, sort_keys=True))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", choices=["cpu", "cuda", "mps"], default="cuda")
+    parser.add_argument("--bits", type=int, default=4)
+    parser.add_argument(
+        "--rows",
+        type=_parse_rows,
+        default=[1, 8, 512, 4096],
+        help="comma-separated row counts to sweep (default covers decode-bound "
+        "small batches and GEMM-bound large batches)",
+    )
+    parser.add_argument(
+        "--dtype",
+        choices=["auto", *sorted(_DTYPES)],
+        default="auto",
+        help="activation dtype; auto picks float16 on mps and bfloat16 elsewhere",
+    )
+    parser.add_argument("--in-features", type=int, default=3072)
+    parser.add_argument("--out-features", type=int, default=3072)
+    parser.add_argument("--iters", type=int, default=20)
+    parser.add_argument("--warmup", type=int, default=3)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--threads", type=int, default=0)
+    parser.add_argument("--with-bias", action="store_true")
+    args = parser.parse_args()
+
+    if args.threads < 0:
+        parser.error("--threads must be non-negative")
+    if args.iters <= 0 or args.warmup < 0:
+        parser.error("--iters must be positive and --warmup must be non-negative")
+    if args.device == "cpu" and args.threads > 0:
+        os.environ["ORBITQUANT_CPU_THREADS"] = str(args.threads)
+        torch.set_num_threads(args.threads)
+
+    torch.manual_seed(args.seed)
+    if args.dtype == "auto":
+        dtype = torch.float16 if args.device == "mps" else torch.bfloat16
+    else:
+        dtype = _DTYPES[args.dtype]
+    indices = torch.randint(
+        0,
+        2**args.bits,
+        (args.out_features, args.in_features),
+        dtype=torch.uint8,
+    )
+    packed = _pack(indices, args.bits).to(args.device)
+    row_norms = torch.ones(args.out_features, device=args.device, dtype=torch.bfloat16)
+    centroids = torch.linspace(-1.0, 1.0, 2**args.bits, device=args.device)
+    indices_device = indices.long().to(args.device)
+    weights = (packed, row_norms, centroids, indices_device)
+
+    payloads = [_benchmark_rows(args, rows, dtype, weights) for rows in args.rows]
+    if len(payloads) == 1:
+        print(json.dumps(payloads[0], indent=2, sort_keys=True))
+    else:
+        print(json.dumps(payloads, indent=2, sort_keys=True))
 
 
 if __name__ == "__main__":

build.toml

@@ -3,7 +3,7 @@ name = "orbitquant-packed-matmul"
 version = 1
 edition = 5
 license = "Apache-2.0"
-backends = ["cuda", "metal"]
+backends = ["cpu", "cuda", "metal"]
 upstream = "https://github.com/iamwavecut/OrbitQuant"
 source = "https://huggingface.co/WaveCut/orbitquant-packed-matmul"
 
@@ -16,6 +16,39 @@ src = [
   "torch-ext/torch_binding.h",
 ]
 
+[torch.stable-abi]
+cpu = "2.11"
+
+[kernel.packed_matmul_cpu]
+backend = "cpu"
+depends = ["torch"]
+include = ["orbitquant_packed_matmul_cpu"]
+src = [
+  "orbitquant_packed_matmul_cpu/cpu_isa.cpp",
+  "orbitquant_packed_matmul_cpu/cpu_kernel_args.h",
+  "orbitquant_packed_matmul_cpu/cpu_pool.cpp",
+  "orbitquant_packed_matmul_cpu/cpu_pool.h",
+  "orbitquant_packed_matmul_cpu/cpu_threads.cpp",
+  "orbitquant_packed_matmul_cpu/cpu_threads.h",
+  "orbitquant_packed_matmul_cpu/packed_adaln_cpu.cpp",
+  "orbitquant_packed_matmul_cpu/packed_matmul_cpu.cpp",
+  "orbitquant_packed_matmul_cpu/packed_matmul_cpu.h",
+  "orbitquant_packed_matmul_cpu/packed_matmul_scalar.cpp",
+  "orbitquant_packed_matmul_cpu/packed_matmul_neon.cpp",
+  "orbitquant_packed_matmul_cpu/packed_matmul_x86_avx512.cpp",
+  "orbitquant_packed_matmul_cpu/quantize_activations_cpu.cpp",
+]
+
+[kernel.packed_matmul_cpu_x86_avx2]
+backend = "cpu"
+depends = ["torch"]
+include = ["orbitquant_packed_matmul_cpu"]
+cxx-flags = ["$<$<CXX_COMPILER_ID:MSVC>:/arch:AVX2>"]
+src = [
+  "orbitquant_packed_matmul_cpu/cpu_msvc_avx2.cpp",
+  "orbitquant_packed_matmul_cpu/packed_matmul_x86.cpp",
+]
+
 [kernel.packed_matmul_cuda]
 backend = "cuda"
 depends = ["torch"]

orbitquant_packed_matmul_cpu/cpu_isa.cpp

@@ -0,0 +1,42 @@
+#include "packed_matmul_cpu.h"
+
+#if defined(__x86_64__) || defined(_M_X64)
+#if defined(_MSC_VER)
+#include <intrin.h>
+#pragma intrinsic(_xgetbv)
+#endif
+#endif
+
+namespace orbitquant::cpu {
+namespace {
+
+bool runtime_has_avx2_fma_f16c() {
+#if defined(_MSC_VER) && defined(_M_X64)
+  int registers[4]{};
+  __cpuid(registers, 1);
+  const bool osxsave = (registers[2] & (1 << 27)) != 0;
+  const bool avx = (registers[2] & (1 << 28)) != 0;
+  const bool fma = (registers[2] & (1 << 12)) != 0;
+  const bool f16c = (registers[2] & (1 << 29)) != 0;
+  if (!osxsave || !avx || !fma || !f16c || (_xgetbv(0) & 0x6) != 0x6) {
+    return false;
+  }
+  __cpuidex(registers, 7, 0);
+  return (registers[1] & (1 << 5)) != 0;
+#elif defined(__x86_64__)
+  __builtin_cpu_init();
+  return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma") &&
+      __builtin_cpu_supports("f16c");
+#else
+  return false;
+#endif
+}
+
+}  // namespace
+
+bool packed_matmul_x86_avx2_available() {
+  static const bool available = runtime_has_avx2_fma_f16c();
+  return available;
+}
+
+}  // namespace orbitquant::cpu

orbitquant_packed_matmul_cpu/cpu_kernel_args.h

@@ -0,0 +1,75 @@
+#pragma once
+
+#include "packed_matmul_cpu.h"
+
+#include <cstdint>
+
+namespace orbitquant::cpu {
+
+enum class ActivationIsa : std::uint8_t {
+  Portable,
+  Neon,
+  Avx2,
+  Avx512,
+};
+
+struct ActivationArgs {
+  void *out;
+  void const *x;
+  // Exactly one of the permutation pointers is set (int64 checkpoints and the
+  // int32 hot-path constants are both accepted).
+  std::int64_t const *permutation;
+  std::int32_t const *permutation_i32;
+  std::int8_t const *signs;
+  float const *centroids;
+  float const *boundaries;
+  ScalarKind scalar_kind;
+  ActivationIsa isa;
+  std::int64_t rows;
+  std::int64_t dim;
+  std::int64_t boundary_count;
+  std::int64_t block_size;
+  float eps;
+  float inv_sqrt_block;
+};
+
+void activation_fwht_msvc_avx2(float *values, std::int64_t block_size);
+
+float activation_squared_norm_msvc_avx2(
+    void const *data,
+    ScalarKind scalar_kind,
+    std::int64_t offset,
+    std::int64_t dim);
+
+void activation_quantize_lookup_msvc_avx2(
+    ActivationArgs const &args,
+    float const *scratch,
+    std::int64_t output_offset,
+    float norm);
+
+struct AdalnArgs {
+  void *out;
+  void const *x;
+  std::uint8_t const *packed_weight;
+  float const *scales;
+  float const *bias;
+  bool has_bias;
+  std::int64_t rows;
+  std::int64_t out_features;
+  std::int64_t in_features;
+  std::int64_t group_size;
+  std::int64_t num_groups;
+  std::int64_t padded_in_features;
+};
+
+using AdalnRangeFn = void (*)(
+    AdalnArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end);
+
+void packed_adaln_msvc_avx2_range(
+    AdalnArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end);
+
+}  // namespace orbitquant::cpu

orbitquant_packed_matmul_cpu/cpu_msvc_avx2.cpp

@@ -0,0 +1,431 @@
+#include "cpu_kernel_args.h"
+
+#if defined(_MSC_VER) && defined(_M_X64)
+#include <immintrin.h>
+
+#include <torch/headeronly/util/BFloat16.h>
+#include <torch/headeronly/util/Half.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <type_traits>
+
+namespace orbitquant::cpu {
+namespace {
+
+template <typename scalar_t>
+inline float load_scalar(void const *data, std::int64_t offset) {
+  return static_cast<float>(static_cast<scalar_t const *>(data)[offset]);
+}
+
+template <>
+inline float load_scalar<float>(void const *data, std::int64_t offset) {
+  return static_cast<float const *>(data)[offset];
+}
+
+template <typename scalar_t>
+inline void store_scalar(void *data, std::int64_t offset, float value) {
+  static_cast<scalar_t *>(data)[offset] = scalar_t(value);
+}
+
+template <>
+inline void store_scalar<float>(void *data, std::int64_t offset, float value) {
+  static_cast<float *>(data)[offset] = value;
+}
+
+template <typename scalar_t>
+float squared_norm_avx2(
+    void const *data,
+    std::int64_t offset,
+    std::int64_t dim) {
+  __m256 accumulator = _mm256_setzero_ps();
+  std::int64_t index = 0;
+  for (; index + 8 <= dim; index += 8) {
+    __m256 values;
+    if constexpr (std::is_same_v<scalar_t, float>) {
+      values = _mm256_loadu_ps(
+          static_cast<float const *>(data) + offset + index);
+    } else if constexpr (std::is_same_v<scalar_t, c10::Half>) {
+      const auto *source =
+          static_cast<std::uint16_t const *>(data) + offset + index;
+      values = _mm256_cvtph_ps(
+          _mm_loadu_si128(reinterpret_cast<__m128i const *>(source)));
+    } else {
+      const auto *source =
+          static_cast<std::uint16_t const *>(data) + offset + index;
+      const __m128i packed =
+          _mm_loadu_si128(reinterpret_cast<__m128i const *>(source));
+      values = _mm256_castsi256_ps(
+          _mm256_slli_epi32(_mm256_cvtepu16_epi32(packed), 16));
+    }
+    accumulator = _mm256_fmadd_ps(values, values, accumulator);
+  }
+  const __m128 halves = _mm_add_ps(
+      _mm256_castps256_ps128(accumulator),
+      _mm256_extractf128_ps(accumulator, 1));
+  const __m128 pairs = _mm_hadd_ps(halves, halves);
+  float result = _mm_cvtss_f32(_mm_hadd_ps(pairs, pairs));
+  for (; index < dim; ++index) {
+    const float value = load_scalar<scalar_t>(data, offset + index);
+    result += value * value;
+  }
+  return result;
+}
+
+inline std::int64_t nearest_centroid(
+    float value,
+    float const *boundaries,
+    std::int64_t boundary_count) {
+  std::int64_t low = 0;
+  std::int64_t high = boundary_count;
+  while (low < high) {
+    const std::int64_t middle = low + (high - low) / 2;
+    if (value <= boundaries[middle]) {
+      high = middle;
+    } else {
+      low = middle + 1;
+    }
+  }
+  return low;
+}
+
+inline __m128i float8_to_bfloat8(__m256 values) {
+  const __m256i bits = _mm256_castps_si256(values);
+  const __m256i absolute_bits =
+      _mm256_and_si256(bits, _mm256_set1_epi32(0x7fffffff));
+  const __m256i nan_mask =
+      _mm256_cmpgt_epi32(absolute_bits, _mm256_set1_epi32(0x7f800000));
+  const __m256i rounding = _mm256_add_epi32(
+      _mm256_set1_epi32(0x7fff),
+      _mm256_and_si256(_mm256_srli_epi32(bits, 16), _mm256_set1_epi32(1)));
+  const __m256i upper = _mm256_srli_epi32(
+      _mm256_add_epi32(bits, rounding),
+      16);
+  __m128i packed = _mm_packus_epi32(
+      _mm256_castsi256_si128(upper),
+      _mm256_extracti128_si256(upper, 1));
+  const __m128i packed_nan_mask = _mm_packs_epi32(
+      _mm256_castsi256_si128(nan_mask),
+      _mm256_extracti128_si256(nan_mask, 1));
+  return _mm_blendv_epi8(
+      packed,
+      _mm_set1_epi16(0x7fc0),
+      packed_nan_mask);
+}
+
+template <typename scalar_t>
+void quantize_lookup_avx2(
+    ActivationArgs const &args,
+    float const *scratch,
+    std::int64_t output_offset,
+    float norm) {
+  const __m256 inverse_sqrt_block = _mm256_set1_ps(args.inv_sqrt_block);
+  const __m256 output_norm = _mm256_set1_ps(norm);
+  const __m256i ones = _mm256_set1_epi32(1);
+  std::int64_t index = 0;
+  for (; index + 8 <= args.dim; index += 8) {
+    const __m256 direction = _mm256_mul_ps(
+        _mm256_loadu_ps(scratch + index), inverse_sqrt_block);
+    __m256i centroid_indices = _mm256_setzero_si256();
+    for (std::int64_t boundary = 0; boundary < args.boundary_count; ++boundary) {
+      const __m256 comparison = _mm256_cmp_ps(
+          direction,
+          _mm256_set1_ps(args.boundaries[boundary]),
+          _CMP_NLE_UQ);
+      centroid_indices = _mm256_add_epi32(
+          centroid_indices,
+          _mm256_and_si256(_mm256_castps_si256(comparison), ones));
+    }
+    const __m256 output = _mm256_mul_ps(
+        _mm256_i32gather_ps(args.centroids, centroid_indices, 4),
+        output_norm);
+    if constexpr (std::is_same_v<scalar_t, float>) {
+      _mm256_storeu_ps(
+          static_cast<float *>(args.out) + output_offset + index,
+          output);
+    } else if constexpr (std::is_same_v<scalar_t, c10::Half>) {
+      _mm_storeu_si128(
+          reinterpret_cast<__m128i *>(
+              static_cast<std::uint16_t *>(args.out) + output_offset + index),
+          _mm256_cvtps_ph(output, _MM_FROUND_TO_NEAREST_INT));
+    } else {
+      _mm_storeu_si128(
+          reinterpret_cast<__m128i *>(
+              static_cast<std::uint16_t *>(args.out) + output_offset + index),
+          float8_to_bfloat8(output));
+    }
+  }
+  for (; index < args.dim; ++index) {
+    const float direction = scratch[index] * args.inv_sqrt_block;
+    const std::int64_t centroid_index = nearest_centroid(
+        direction, args.boundaries, args.boundary_count);
+    store_scalar<scalar_t>(
+        args.out,
+        output_offset + index,
+        args.centroids[centroid_index] * norm);
+  }
+}
+
+inline float load_bfloat(void const *data, std::int64_t offset) {
+  return static_cast<float>(
+      static_cast<c10::BFloat16 const *>(data)[offset]);
+}
+
+inline void store_bfloat(void *data, std::int64_t offset, float value) {
+  static_cast<c10::BFloat16 *>(data)[offset] = c10::BFloat16(value);
+}
+
+inline std::uint8_t unpack_adaln_index(
+    std::uint8_t const *packed,
+    std::int64_t flat_index) {
+  const std::uint8_t byte = packed[flat_index / 2];
+  return (flat_index & 1) == 0 ? byte & 15u : (byte >> 4) & 15u;
+}
+
+inline float dequantized_adaln_value(std::uint8_t index, float scale) {
+  return static_cast<float>(
+      c10::BFloat16((static_cast<int>(index) - 8) * scale));
+}
+
+inline void fill_group_lut(float *lut, float scale) {
+  for (int index = 0; index < 16; ++index) {
+    lut[index] = dequantized_adaln_value(
+        static_cast<std::uint8_t>(index), scale);
+  }
+}
+
+void packed_adaln_scalar_range(
+    AdalnArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) {
+    for (std::int64_t row = 0; row < args.rows; ++row) {
+      float accumulator = 0.0f;
+      const std::int64_t input_offset = row * args.in_features;
+      for (std::int64_t k = 0; k < args.in_features; ++k) {
+        const std::int64_t weight_offset =
+            out_col * args.padded_in_features + k;
+        const std::uint8_t index =
+            unpack_adaln_index(args.packed_weight, weight_offset);
+        const float scale =
+            args.scales[out_col * args.num_groups + k / args.group_size];
+        accumulator += load_bfloat(args.x, input_offset + k) *
+            dequantized_adaln_value(index, scale);
+      }
+      if (args.has_bias) {
+        accumulator += args.bias[out_col];
+      }
+      store_bfloat(
+          args.out,
+          row * args.out_features + out_col,
+          accumulator);
+    }
+  }
+}
+
+inline __m256 load_bfloat8(void const *data, std::int64_t offset) {
+  const auto *source =
+      static_cast<std::uint16_t const *>(data) + offset;
+  const __m128i packed =
+      _mm_loadu_si128(reinterpret_cast<__m128i const *>(source));
+  return _mm256_castsi256_ps(
+      _mm256_slli_epi32(_mm256_cvtepu16_epi32(packed), 16));
+}
+
+inline float horizontal_sum(__m256 value) {
+  const __m128 halves = _mm_add_ps(
+      _mm256_castps256_ps128(value),
+      _mm256_extractf128_ps(value, 1));
+  const __m128 pairs = _mm_hadd_ps(halves, halves);
+  return _mm_cvtss_f32(_mm_hadd_ps(pairs, pairs));
+}
+
+template <int row_tile>
+inline void packed_adaln_avx2_rows(
+    AdalnArgs const &args,
+    std::uint8_t const *packed_row,
+    float const *scale_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  __m256 accumulators[row_tile];
+  float tails[row_tile]{};
+  for (int row = 0; row < row_tile; ++row) {
+    accumulators[row] = _mm256_setzero_ps();
+  }
+  const __m128i nibble_mask = _mm_set1_epi8(15);
+  const __m256i low_table_limit = _mm256_set1_epi32(7);
+
+  for (std::int64_t group = 0; group < args.num_groups; ++group) {
+    float lut[16];
+    fill_group_lut(lut, scale_row[group]);
+    const __m256 lut_low = _mm256_loadu_ps(lut);
+    const __m256 lut_high = _mm256_loadu_ps(lut + 8);
+    const std::int64_t group_start = group * args.group_size;
+    const std::int64_t group_end = std::min(
+        args.in_features,
+        group_start + args.group_size);
+    std::int64_t k = group_start;
+    for (; k + 8 <= group_end; k += 8) {
+      std::int32_t packed;
+      std::memcpy(&packed, packed_row + k / 2, sizeof(packed));
+      const __m128i bytes = _mm_cvtsi32_si128(packed);
+      const __m128i low = _mm_and_si128(bytes, nibble_mask);
+      const __m128i high = _mm_and_si128(
+          _mm_srli_epi16(bytes, 4), nibble_mask);
+      const __m256i indices =
+          _mm256_cvtepu8_epi32(_mm_unpacklo_epi8(low, high));
+      const __m256 weight = _mm256_blendv_ps(
+          _mm256_permutevar8x32_ps(lut_low, indices),
+          _mm256_permutevar8x32_ps(lut_high, indices),
+          _mm256_castsi256_ps(
+              _mm256_cmpgt_epi32(indices, low_table_limit)));
+      for (int row = 0; row < row_tile; ++row) {
+        const std::int64_t input_offset =
+            (row_start + row) * args.in_features + k;
+        accumulators[row] = _mm256_fmadd_ps(
+            load_bfloat8(args.x, input_offset),
+            weight,
+            accumulators[row]);
+      }
+    }
+    for (; k < group_end; ++k) {
+      const float weight = lut[unpack_adaln_index(packed_row, k)];
+      for (int row = 0; row < row_tile; ++row) {
+        tails[row] += load_bfloat(
+            args.x,
+            (row_start + row) * args.in_features + k) * weight;
+      }
+    }
+  }
+
+  for (int row = 0; row < row_tile; ++row) {
+    float accumulator = horizontal_sum(accumulators[row]) + tails[row];
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_bfloat(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+template <
+    void (*rows8)(AdalnArgs const &, std::uint8_t const *, float const *, std::int64_t, std::int64_t),
+    void (*rows4)(AdalnArgs const &, std::uint8_t const *, float const *, std::int64_t, std::int64_t),
+    void (*rows3)(AdalnArgs const &, std::uint8_t const *, float const *, std::int64_t, std::int64_t),
+    void (*rows2)(AdalnArgs const &, std::uint8_t const *, float const *, std::int64_t, std::int64_t),
+    void (*rows1)(AdalnArgs const &, std::uint8_t const *, float const *, std::int64_t, std::int64_t)>
+void packed_adaln_tiled_range(
+    AdalnArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  const std::int64_t packed_row_bytes = args.padded_in_features / 2;
+  for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) {
+    const auto *packed_row =
+        args.packed_weight + out_col * packed_row_bytes;
+    const auto *scale_row =
+        args.scales + out_col * args.num_groups;
+    std::int64_t row = 0;
+    for (; row + 8 <= args.rows; row += 8) {
+      rows8(args, packed_row, scale_row, out_col, row);
+    }
+    if (row + 4 <= args.rows) {
+      rows4(args, packed_row, scale_row, out_col, row);
+      row += 4;
+    }
+    switch (args.rows - row) {
+      case 3:
+        rows3(args, packed_row, scale_row, out_col, row);
+        break;
+      case 2:
+        rows2(args, packed_row, scale_row, out_col, row);
+        break;
+      case 1:
+        rows1(args, packed_row, scale_row, out_col, row);
+        break;
+      default:
+        break;
+    }
+  }
+}
+
+}  // namespace
+
+void activation_fwht_msvc_avx2(float *values, std::int64_t block_size) {
+  for (std::int64_t half = 1; half < block_size; half *= 2) {
+    for (std::int64_t base = 0; base < block_size; base += 2 * half) {
+      std::int64_t offset = 0;
+      for (; offset + 8 <= half; offset += 8) {
+        const __m256 left = _mm256_loadu_ps(values + base + offset);
+        const __m256 right =
+            _mm256_loadu_ps(values + base + half + offset);
+        _mm256_storeu_ps(values + base + offset, _mm256_add_ps(left, right));
+        _mm256_storeu_ps(
+            values + base + half + offset,
+            _mm256_sub_ps(left, right));
+      }
+      for (; offset < half; ++offset) {
+        const float left = values[base + offset];
+        const float right = values[base + half + offset];
+        values[base + offset] = left + right;
+        values[base + half + offset] = left - right;
+      }
+    }
+  }
+}
+
+float activation_squared_norm_msvc_avx2(
+    void const *data,
+    ScalarKind scalar_kind,
+    std::int64_t offset,
+    std::int64_t dim) {
+  switch (scalar_kind) {
+    case ScalarKind::Float32:
+      return squared_norm_avx2<float>(data, offset, dim);
+    case ScalarKind::Float16:
+      return squared_norm_avx2<c10::Half>(data, offset, dim);
+    case ScalarKind::BFloat16:
+      return squared_norm_avx2<c10::BFloat16>(data, offset, dim);
+  }
+  return 0.0f;
+}
+
+void activation_quantize_lookup_msvc_avx2(
+    ActivationArgs const &args,
+    float const *scratch,
+    std::int64_t output_offset,
+    float norm) {
+  switch (args.scalar_kind) {
+    case ScalarKind::Float32:
+      quantize_lookup_avx2<float>(args, scratch, output_offset, norm);
+      return;
+    case ScalarKind::Float16:
+      quantize_lookup_avx2<c10::Half>(args, scratch, output_offset, norm);
+      return;
+    case ScalarKind::BFloat16:
+      quantize_lookup_avx2<c10::BFloat16>(
+          args, scratch, output_offset, norm);
+      return;
+  }
+}
+
+void packed_adaln_msvc_avx2_range(
+    AdalnArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  if (args.group_size % 8 != 0 || args.padded_in_features % 2 != 0) {
+    packed_adaln_scalar_range(args, out_start, out_end);
+    return;
+  }
+  packed_adaln_tiled_range<
+      packed_adaln_avx2_rows<8>,
+      packed_adaln_avx2_rows<4>,
+      packed_adaln_avx2_rows<3>,
+      packed_adaln_avx2_rows<2>,
+      packed_adaln_avx2_rows<1>>(args, out_start, out_end);
+}
+
+}  // namespace orbitquant::cpu
+#endif

orbitquant_packed_matmul_cpu/cpu_pool.cpp

@@ -0,0 +1,126 @@
+#include "cpu_pool.h"
+#include "cpu_threads.h"
+
+#include <atomic>
+#include <condition_variable>
+#include <cstddef>
+#include <mutex>
+#include <thread>
+
+namespace orbitquant::cpu {
+namespace {
+
+struct Job {
+  std::vector<std::pair<std::int64_t, std::int64_t>> const *ranges;
+  std::function<void(std::int64_t, std::int64_t)> const *fn;
+  std::atomic<std::size_t> next_range{0};
+  std::atomic<int> pending_workers{0};
+};
+
+class WorkerPool {
+ public:
+  static WorkerPool &instance() {
+    // Intentionally leaked: joining workers during static destruction can
+    // deadlock interpreter teardown, so the pool lives for the process.
+    static WorkerPool *pool = new WorkerPool(requested_threads() - 1);
+    return *pool;
+  }
+
+  int worker_count() const {
+    return worker_count_;
+  }
+
+  void run(Job &job) {
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      job.pending_workers.store(worker_count_, std::memory_order_relaxed);
+      job_ = &job;
+      ++generation_;
+    }
+    wake_cv_.notify_all();
+    drain(job);
+    std::unique_lock<std::mutex> lock(mutex_);
+    done_cv_.wait(lock, [&job] {
+      return job.pending_workers.load(std::memory_order_acquire) == 0;
+    });
+    job_ = nullptr;
+  }
+
+ private:
+  explicit WorkerPool(int requested_workers)
+      : worker_count_(requested_workers > 0 ? requested_workers : 0) {
+    for (int worker = 0; worker < worker_count_; ++worker) {
+      std::thread([this] { worker_loop(); }).detach();
+    }
+  }
+
+  static void drain(Job &job) {
+    const std::size_t range_count = job.ranges->size();
+    for (;;) {
+      const std::size_t index =
+          job.next_range.fetch_add(1, std::memory_order_relaxed);
+      if (index >= range_count) {
+        return;
+      }
+      auto const &range = (*job.ranges)[index];
+      (*job.fn)(range.first, range.second);
+    }
+  }
+
+  void worker_loop() {
+    std::uint64_t seen_generation = 0;
+    for (;;) {
+      Job *job = nullptr;
+      {
+        std::unique_lock<std::mutex> lock(mutex_);
+        wake_cv_.wait(lock, [this, seen_generation] {
+          return generation_ != seen_generation;
+        });
+        seen_generation = generation_;
+        job = job_;
+      }
+      if (job == nullptr) {
+        continue;
+      }
+      drain(*job);
+      if (job->pending_workers.fetch_sub(1, std::memory_order_acq_rel) == 1) {
+        std::lock_guard<std::mutex> lock(mutex_);
+        done_cv_.notify_all();
+      }
+    }
+  }
+
+  const int worker_count_;
+  std::mutex mutex_;
+  std::condition_variable wake_cv_;
+  std::condition_variable done_cv_;
+  Job *job_ = nullptr;
+  std::uint64_t generation_ = 0;
+};
+
+}  // namespace
+
+void run_ranges(
+    std::vector<std::pair<std::int64_t, std::int64_t>> const &ranges,
+    std::function<void(std::int64_t, std::int64_t)> const &fn) {
+  if (ranges.empty()) {
+    return;
+  }
+  if (ranges.size() == 1) {
+    fn(ranges.front().first, ranges.front().second);
+    return;
+  }
+  WorkerPool &pool = WorkerPool::instance();
+  if (pool.worker_count() == 0) {
+    for (auto const &range : ranges) {
+      fn(range.first, range.second);
+    }
+    return;
+  }
+  Job job;
+  job.ranges = &ranges;
+  job.fn = &fn;
+  pool.run(job);
+}
+
+}  // namespace orbitquant::cpu

orbitquant_packed_matmul_cpu/cpu_pool.h

@@ -0,0 +1,18 @@
+#pragma once
+
+#include <cstdint>
+#include <functional>
+#include <utility>
+#include <vector>
+
+namespace orbitquant::cpu {
+
+// Executes fn over each [start, end) range using a lazily created persistent
+// worker pool plus the calling thread. Ranges are disjoint, so results stay
+// deterministic regardless of which thread runs which range. A single range
+// (or an empty pool) runs inline on the caller.
+void run_ranges(
+    std::vector<std::pair<std::int64_t, std::int64_t>> const &ranges,
+    std::function<void(std::int64_t, std::int64_t)> const &fn);
+
+}  // namespace orbitquant::cpu

orbitquant_packed_matmul_cpu/cpu_threads.cpp

@@ -0,0 +1,103 @@
+#include "cpu_threads.h"
+
+#include <algorithm>
+#include <cstdlib>
+#include <fstream>
+#include <set>
+#include <string>
+#include <thread>
+#include <utility>
+
+#if defined(__APPLE__)
+#include <sys/sysctl.h>
+#endif
+
+#if defined(__linux__)
+#include <sched.h>
+#endif
+
+namespace orbitquant::cpu {
+namespace {
+
+int environment_thread_count() {
+  char const *value = std::getenv("ORBITQUANT_CPU_THREADS");
+  if (value == nullptr) {
+    value = std::getenv("OMP_NUM_THREADS");
+  }
+  if (value == nullptr) {
+    return 0;
+  }
+  char *end = nullptr;
+  const long parsed = std::strtol(value, &end, 10);
+  if (end == value || parsed <= 0) {
+    return 0;
+  }
+  return static_cast<int>(std::min<long>(parsed, 64));
+}
+
+#if defined(__linux__)
+int affinity_physical_core_count() {
+  cpu_set_t affinity;
+  CPU_ZERO(&affinity);
+  if (sched_getaffinity(0, sizeof(affinity), &affinity) != 0) {
+    return 0;
+  }
+
+  int logical_cpus = 0;
+  std::set<std::pair<int, int>> physical_cores;
+  for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) {
+    if (!CPU_ISSET(cpu, &affinity)) {
+      continue;
+    }
+    ++logical_cpus;
+    int package = -1;
+    int core = -1;
+    std::ifstream package_file(
+        "/sys/devices/system/cpu/cpu" + std::to_string(cpu) +
+        "/topology/physical_package_id");
+    std::ifstream core_file(
+        "/sys/devices/system/cpu/cpu" + std::to_string(cpu) +
+        "/topology/core_id");
+    if (package_file >> package && core_file >> core) {
+      physical_cores.emplace(package, core);
+    }
+  }
+  return physical_cores.empty()
+      ? logical_cpus
+      : static_cast<int>(physical_cores.size());
+}
+#endif
+
+int default_thread_count() {
+#if defined(__APPLE__)
+  int performance_cores = 0;
+  std::size_t size = sizeof(performance_cores);
+  if (sysctlbyname(
+          "hw.perflevel0.physicalcpu",
+          &performance_cores,
+          &size,
+          nullptr,
+          0) == 0 &&
+      performance_cores > 0) {
+    return std::min(performance_cores, 64);
+  }
+#endif
+#if defined(__linux__)
+  const int affinity_cores = affinity_physical_core_count();
+  if (affinity_cores > 0) {
+    return std::min(affinity_cores, 64);
+  }
+#endif
+  const unsigned hardware = std::thread::hardware_concurrency();
+  return static_cast<int>(hardware == 0 ? 1 : std::min<unsigned>(hardware, 64));
+}
+
+}  // namespace
+
+int requested_threads() {
+  const int environment = environment_thread_count();
+  static const int default_threads = default_thread_count();
+  return environment > 0 ? environment : default_threads;
+}
+
+}  // namespace orbitquant::cpu

orbitquant_packed_matmul_cpu/cpu_threads.h

@@ -0,0 +1,7 @@
+#pragma once
+
+namespace orbitquant::cpu {
+
+int requested_threads();
+
+}  // namespace orbitquant::cpu

orbitquant_packed_matmul_cpu/packed_adaln_cpu.cpp

@@ -0,0 +1,878 @@
+#include "cpu_pool.h"
+#include "cpu_threads.h"
+#include "cpu_kernel_args.h"
+#include "packed_matmul_cpu.h"
+#include "../torch-ext/torch_binding.h"
+
+#include <torch/headeronly/core/DeviceType.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/BFloat16.h>
+
+#include <algorithm>
+#include <cstdlib>
+#include <cstdint>
+#include <cstring>
+#include <utility>
+#include <vector>
+
+#if defined(__aarch64__) || defined(_M_ARM64)
+#include <arm_neon.h>
+#endif
+
+#if (defined(__x86_64__) || defined(_M_X64)) && !defined(_MSC_VER)
+#include <immintrin.h>
+#define ORBITQUANT_TARGET_AVX2 __attribute__((target("avx2,fma,f16c")))
+#define ORBITQUANT_TARGET_AVX512 \
+  __attribute__((target("avx512f,avx512dq,avx512bw,avx512vl,fma,f16c")))
+#define ORBITQUANT_TARGET_AVX512_BF16 \
+  __attribute__((target( \
+      "avx512f,avx512dq,avx512bw,avx512vl,avx512bf16,fma,f16c")))
+#define ORBITQUANT_HAS_AVX512_BF16_INTRINSICS 1
+#define ORBITQUANT_NOINLINE __attribute__((noinline))
+#endif
+
+namespace {
+
+using orbitquant::cpu::AdalnArgs;
+using orbitquant::cpu::AdalnRangeFn;
+
+inline float load_bfloat(
+    void const *data,
+    std::int64_t offset) {
+  return static_cast<float>(
+      static_cast<c10::BFloat16 const *>(data)[offset]);
+}
+
+inline void store_bfloat(
+    void *data,
+    std::int64_t offset,
+    float value) {
+  static_cast<c10::BFloat16 *>(data)[offset] = c10::BFloat16(value);
+}
+
+inline std::uint8_t unpack_adaln_index(
+    std::uint8_t const *packed,
+    std::int64_t flat_index) {
+  const std::uint8_t byte = packed[flat_index / 2];
+  return (flat_index & 1) == 0 ? byte & 15u : (byte >> 4) & 15u;
+}
+
+inline float dequantized_adaln_value(
+    std::uint8_t index,
+    float scale) {
+  return static_cast<float>(
+      c10::BFloat16((static_cast<int>(index) - 8) * scale));
+}
+
+inline void fill_group_lut(float *lut, float scale) {
+  for (int index = 0; index < 16; ++index) {
+    lut[index] = dequantized_adaln_value(
+        static_cast<std::uint8_t>(index),
+        scale);
+  }
+}
+
+void packed_adaln_scalar_range(
+    AdalnArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) {
+    for (std::int64_t row = 0; row < args.rows; ++row) {
+      float accumulator = 0.0f;
+      const std::int64_t input_offset = row * args.in_features;
+      for (std::int64_t k = 0; k < args.in_features; ++k) {
+        const std::int64_t weight_offset =
+            out_col * args.padded_in_features + k;
+        const std::uint8_t index =
+            unpack_adaln_index(args.packed_weight, weight_offset);
+        const float scale =
+            args.scales[out_col * args.num_groups + k / args.group_size];
+        accumulator += load_bfloat(args.x, input_offset + k) *
+            dequantized_adaln_value(index, scale);
+      }
+      if (args.has_bias) {
+        accumulator += args.bias[out_col];
+      }
+      store_bfloat(
+          args.out,
+          row * args.out_features + out_col,
+          accumulator);
+    }
+  }
+}
+
+#if defined(__aarch64__) || defined(_M_ARM64)
+inline float32x4_t load_bfloat4(
+    void const *data,
+    std::int64_t offset) {
+  const uint16x4_t raw = vld1_u16(
+      static_cast<std::uint16_t const *>(data) + offset);
+  return vreinterpretq_f32_u32(vshlq_n_u32(vmovl_u16(raw), 16));
+}
+
+template <int row_tile>
+inline void packed_adaln_neon_rows(
+    AdalnArgs const &args,
+    std::uint8_t const *packed_row,
+    float const *scale_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  float32x4_t accumulators[row_tile];
+  float tails[row_tile]{};
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    accumulators[row] = vdupq_n_f32(0.0f);
+  }
+
+  for (std::int64_t group = 0; group < args.num_groups; ++group) {
+    float lut[16];
+    fill_group_lut(lut, scale_row[group]);
+    const std::int64_t group_start = group * args.group_size;
+    const std::int64_t group_end = std::min(
+        args.in_features,
+        group_start + args.group_size);
+    std::int64_t k = group_start;
+    for (; k + 4 <= group_end; k += 4) {
+      float weights[4];
+#pragma clang loop unroll(full)
+      for (int lane = 0; lane < 4; ++lane) {
+        weights[lane] = lut[unpack_adaln_index(packed_row, k + lane)];
+      }
+      const float32x4_t weight = vld1q_f32(weights);
+#pragma clang loop unroll(full)
+      for (int row = 0; row < row_tile; ++row) {
+        const std::int64_t input_offset =
+            (row_start + row) * args.in_features + k;
+        accumulators[row] = vfmaq_f32(
+            accumulators[row],
+            load_bfloat4(args.x, input_offset),
+            weight);
+      }
+    }
+    for (; k < group_end; ++k) {
+      const float weight = lut[unpack_adaln_index(packed_row, k)];
+#pragma clang loop unroll(full)
+      for (int row = 0; row < row_tile; ++row) {
+        tails[row] += load_bfloat(
+            args.x,
+            (row_start + row) * args.in_features + k) * weight;
+      }
+    }
+  }
+
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    float accumulator = vaddvq_f32(accumulators[row]) + tails[row];
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_bfloat(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+void packed_adaln_neon_range(
+    AdalnArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  if (args.group_size % 4 != 0 || args.padded_in_features % 2 != 0) {
+    packed_adaln_scalar_range(args, out_start, out_end);
+    return;
+  }
+  constexpr int kPrimaryRowTile = 8;
+  const std::int64_t packed_row_bytes = args.padded_in_features / 2;
+  for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) {
+    const auto *packed_row =
+        args.packed_weight + out_col * packed_row_bytes;
+    const auto *scale_row =
+        args.scales + out_col * args.num_groups;
+    std::int64_t row = 0;
+    for (; row + kPrimaryRowTile <= args.rows; row += kPrimaryRowTile) {
+      packed_adaln_neon_rows<kPrimaryRowTile>(
+          args, packed_row, scale_row, out_col, row);
+    }
+    if (row + 4 <= args.rows) {
+      packed_adaln_neon_rows<4>(
+          args, packed_row, scale_row, out_col, row);
+      row += 4;
+    }
+    switch (args.rows - row) {
+      case 3:
+        packed_adaln_neon_rows<3>(
+            args, packed_row, scale_row, out_col, row);
+        break;
+      case 2:
+        packed_adaln_neon_rows<2>(
+            args, packed_row, scale_row, out_col, row);
+        break;
+      case 1:
+        packed_adaln_neon_rows<1>(
+            args, packed_row, scale_row, out_col, row);
+        break;
+      default:
+        break;
+    }
+  }
+}
+#else
+void packed_adaln_neon_range(
+    AdalnArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  packed_adaln_scalar_range(args, out_start, out_end);
+}
+#endif
+
+#if (defined(__x86_64__) || defined(_M_X64)) && !defined(_MSC_VER)
+ORBITQUANT_TARGET_AVX2 inline __m256 load_bfloat8(
+    void const *data,
+    std::int64_t offset) {
+  const auto *source =
+      static_cast<std::uint16_t const *>(data) + offset;
+  const __m128i packed =
+      _mm_loadu_si128(reinterpret_cast<__m128i const *>(source));
+  return _mm256_castsi256_ps(
+      _mm256_slli_epi32(_mm256_cvtepu16_epi32(packed), 16));
+}
+
+ORBITQUANT_TARGET_AVX512 inline __m512 load_bfloat16(
+    void const *data,
+    std::int64_t offset) {
+  const auto *source =
+      static_cast<std::uint16_t const *>(data) + offset;
+  const __m256i packed =
+      _mm256_loadu_si256(reinterpret_cast<__m256i const *>(source));
+  return _mm512_castsi512_ps(
+      _mm512_slli_epi32(_mm512_cvtepu16_epi32(packed), 16));
+}
+
+ORBITQUANT_TARGET_AVX2 inline float horizontal_sum(__m256 value) {
+  const __m128 halves = _mm_add_ps(
+      _mm256_castps256_ps128(value),
+      _mm256_extractf128_ps(value, 1));
+  const __m128 pairs = _mm_hadd_ps(halves, halves);
+  return _mm_cvtss_f32(_mm_hadd_ps(pairs, pairs));
+}
+
+ORBITQUANT_TARGET_AVX512 inline float horizontal_sum(__m512 value) {
+  return _mm512_reduce_add_ps(value);
+}
+
+template <int row_tile>
+ORBITQUANT_TARGET_AVX2 inline void packed_adaln_avx2_rows(
+    AdalnArgs const &args,
+    std::uint8_t const *packed_row,
+    float const *scale_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  __m256 accumulators[row_tile];
+  float tails[row_tile]{};
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    accumulators[row] = _mm256_setzero_ps();
+  }
+  const __m128i nibble_mask = _mm_set1_epi8(15);
+  const __m256i low_table_limit = _mm256_set1_epi32(7);
+
+  for (std::int64_t group = 0; group < args.num_groups; ++group) {
+    float lut[16];
+    fill_group_lut(lut, scale_row[group]);
+    const __m256 lut_low = _mm256_loadu_ps(lut);
+    const __m256 lut_high = _mm256_loadu_ps(lut + 8);
+    const std::int64_t group_start = group * args.group_size;
+    const std::int64_t group_end = std::min(
+        args.in_features,
+        group_start + args.group_size);
+    std::int64_t k = group_start;
+    for (; k + 8 <= group_end; k += 8) {
+      std::int32_t packed;
+      std::memcpy(&packed, packed_row + k / 2, sizeof(packed));
+      const __m128i bytes = _mm_cvtsi32_si128(packed);
+      const __m128i low = _mm_and_si128(bytes, nibble_mask);
+      const __m128i high = _mm_and_si128(
+          _mm_srli_epi16(bytes, 4),
+          nibble_mask);
+      const __m256i indices =
+          _mm256_cvtepu8_epi32(_mm_unpacklo_epi8(low, high));
+      const __m256 weight = _mm256_blendv_ps(
+          _mm256_permutevar8x32_ps(lut_low, indices),
+          _mm256_permutevar8x32_ps(lut_high, indices),
+          _mm256_castsi256_ps(
+              _mm256_cmpgt_epi32(indices, low_table_limit)));
+#pragma clang loop unroll(full)
+      for (int row = 0; row < row_tile; ++row) {
+        const std::int64_t input_offset =
+            (row_start + row) * args.in_features + k;
+        accumulators[row] = _mm256_fmadd_ps(
+            load_bfloat8(args.x, input_offset),
+            weight,
+            accumulators[row]);
+      }
+    }
+    for (; k < group_end; ++k) {
+      const float weight = lut[unpack_adaln_index(packed_row, k)];
+#pragma clang loop unroll(full)
+      for (int row = 0; row < row_tile; ++row) {
+        tails[row] += load_bfloat(
+            args.x,
+            (row_start + row) * args.in_features + k) * weight;
+      }
+    }
+  }
+
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    float accumulator = horizontal_sum(accumulators[row]) + tails[row];
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_bfloat(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+template <int row_tile>
+ORBITQUANT_TARGET_AVX512 inline void packed_adaln_avx512_rows(
+    AdalnArgs const &args,
+    std::uint8_t const *packed_row,
+    float const *scale_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  __m512 accumulators[row_tile];
+  float tails[row_tile]{};
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    accumulators[row] = _mm512_setzero_ps();
+  }
+  const __m128i nibble_mask = _mm_set1_epi8(15);
+
+  for (std::int64_t group = 0; group < args.num_groups; ++group) {
+    float lut[16];
+    fill_group_lut(lut, scale_row[group]);
+    const __m512 centroid_lut = _mm512_loadu_ps(lut);
+    const std::int64_t group_start = group * args.group_size;
+    const std::int64_t group_end = std::min(
+        args.in_features,
+        group_start + args.group_size);
+    std::int64_t k = group_start;
+    for (; k + 16 <= group_end; k += 16) {
+      std::int64_t packed;
+      std::memcpy(&packed, packed_row + k / 2, sizeof(packed));
+      const __m128i bytes = _mm_cvtsi64_si128(packed);
+      const __m128i low = _mm_and_si128(bytes, nibble_mask);
+      const __m128i high = _mm_and_si128(
+          _mm_srli_epi16(bytes, 4),
+          nibble_mask);
+      const __m512i indices =
+          _mm512_cvtepu8_epi32(_mm_unpacklo_epi8(low, high));
+      const __m512 weight =
+          _mm512_permutexvar_ps(indices, centroid_lut);
+#pragma clang loop unroll(full)
+      for (int row = 0; row < row_tile; ++row) {
+        const std::int64_t input_offset =
+            (row_start + row) * args.in_features + k;
+        accumulators[row] = _mm512_fmadd_ps(
+            load_bfloat16(args.x, input_offset),
+            weight,
+            accumulators[row]);
+      }
+    }
+    for (; k < group_end; ++k) {
+      const float weight = lut[unpack_adaln_index(packed_row, k)];
+#pragma clang loop unroll(full)
+      for (int row = 0; row < row_tile; ++row) {
+        tails[row] += load_bfloat(
+            args.x,
+            (row_start + row) * args.in_features + k) * weight;
+      }
+    }
+  }
+
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    float accumulator = horizontal_sum(accumulators[row]) + tails[row];
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_bfloat(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+#if defined(ORBITQUANT_HAS_AVX512_BF16_INTRINSICS)
+template <bool aligned_k>
+ORBITQUANT_TARGET_AVX512_BF16 inline __m512bh load_bfloat32(
+    std::uint16_t const *data,
+    __mmask32 valid_mask) {
+  if constexpr (aligned_k) {
+    return (__m512bh)_mm512_loadu_si512(data);
+  }
+  return (__m512bh)_mm512_maskz_loadu_epi16(valid_mask, data);
+}
+
+template <int row_tile, bool aligned_k>
+ORBITQUANT_TARGET_AVX512_BF16 inline void packed_adaln_avx512_bf16_rows(
+    AdalnArgs const &args,
+    std::uint8_t const *packed_row,
+    float const *scale_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  static_assert(row_tile >= 1 && row_tile <= 8);
+  __m512 accumulator0 = _mm512_setzero_ps();
+  __m512 accumulator1 = _mm512_setzero_ps();
+  __m512 accumulator2 = _mm512_setzero_ps();
+  __m512 accumulator3 = _mm512_setzero_ps();
+  __m512 accumulator4 = _mm512_setzero_ps();
+  __m512 accumulator5 = _mm512_setzero_ps();
+  __m512 accumulator6 = _mm512_setzero_ps();
+  __m512 accumulator7 = _mm512_setzero_ps();
+  const __m128i nibble_mask = _mm_set1_epi8(15);
+  const __m512 signed_codes = _mm512_setr_ps(
+      -8.0f, -7.0f, -6.0f, -5.0f,
+      -4.0f, -3.0f, -2.0f, -1.0f,
+      0.0f, 1.0f, 2.0f, 3.0f,
+      4.0f, 5.0f, 6.0f, 7.0f);
+  const auto *input_base =
+      static_cast<std::uint16_t const *>(args.x) +
+      row_start * args.in_features;
+  const std::int64_t input_stride = args.in_features;
+
+  for (std::int64_t group = 0; group < args.num_groups; ++group) {
+    const __m512 scaled_codes = _mm512_mul_ps(
+        signed_codes,
+        _mm512_set1_ps(scale_row[group]));
+    const __m512i lut_words = (__m512i)_mm512_cvtne2ps_pbh(
+        scaled_codes,
+        scaled_codes);
+    const std::int64_t group_start = group * args.group_size;
+    for (std::int64_t offset = 0; offset < args.group_size; offset += 32) {
+      const std::int64_t k = group_start + offset;
+      __mmask32 valid_mask = static_cast<__mmask32>(0xffffffffu);
+      if constexpr (!aligned_k) {
+        const std::int64_t valid = std::max<std::int64_t>(
+            0,
+            std::min<std::int64_t>(32, args.in_features - k));
+        if (valid == 0) {
+          continue;
+        }
+        valid_mask = valid == 32
+            ? static_cast<__mmask32>(0xffffffffu)
+            : static_cast<__mmask32>((1u << valid) - 1u);
+      }
+      const __m128i bytes = _mm_loadu_si128(
+          reinterpret_cast<__m128i const *>(packed_row + k / 2));
+      const __m128i low = _mm_and_si128(bytes, nibble_mask);
+      const __m128i high = _mm_and_si128(
+          _mm_srli_epi16(bytes, 4),
+          nibble_mask);
+      const __m256i packed_indices = _mm256_set_m128i(
+          _mm_unpackhi_epi8(low, high),
+          _mm_unpacklo_epi8(low, high));
+      const __m512i indices =
+          _mm512_cvtepu8_epi16(packed_indices);
+      const __m512bh weights = (__m512bh)_mm512_permutexvar_epi16(
+          indices,
+          lut_words);
+      accumulator0 = _mm512_dpbf16_ps(
+          accumulator0,
+          load_bfloat32<aligned_k>(input_base + k, valid_mask),
+          weights);
+      if constexpr (row_tile > 1) {
+        accumulator1 = _mm512_dpbf16_ps(
+            accumulator1,
+            load_bfloat32<aligned_k>(
+                input_base + input_stride + k,
+                valid_mask),
+            weights);
+      }
+      if constexpr (row_tile > 2) {
+        accumulator2 = _mm512_dpbf16_ps(
+            accumulator2,
+            load_bfloat32<aligned_k>(
+                input_base + 2 * input_stride + k,
+                valid_mask),
+            weights);
+      }
+      if constexpr (row_tile > 3) {
+        accumulator3 = _mm512_dpbf16_ps(
+            accumulator3,
+            load_bfloat32<aligned_k>(
+                input_base + 3 * input_stride + k,
+                valid_mask),
+            weights);
+      }
+      if constexpr (row_tile > 4) {
+        accumulator4 = _mm512_dpbf16_ps(
+            accumulator4,
+            load_bfloat32<aligned_k>(
+                input_base + 4 * input_stride + k,
+                valid_mask),
+            weights);
+      }
+      if constexpr (row_tile > 5) {
+        accumulator5 = _mm512_dpbf16_ps(
+            accumulator5,
+            load_bfloat32<aligned_k>(
+                input_base + 5 * input_stride + k,
+                valid_mask),
+            weights);
+      }
+      if constexpr (row_tile > 6) {
+        accumulator6 = _mm512_dpbf16_ps(
+            accumulator6,
+            load_bfloat32<aligned_k>(
+                input_base + 6 * input_stride + k,
+                valid_mask),
+            weights);
+      }
+      if constexpr (row_tile > 7) {
+        accumulator7 = _mm512_dpbf16_ps(
+            accumulator7,
+            load_bfloat32<aligned_k>(
+                input_base + 7 * input_stride + k,
+                valid_mask),
+            weights);
+      }
+    }
+  }
+
+  const float bias = args.has_bias ? args.bias[out_col] : 0.0f;
+  const std::int64_t output_offset =
+      row_start * args.out_features + out_col;
+  store_bfloat(
+      args.out,
+      output_offset,
+      horizontal_sum(accumulator0) + bias);
+  if constexpr (row_tile > 1) {
+    store_bfloat(
+        args.out,
+        output_offset + args.out_features,
+        horizontal_sum(accumulator1) + bias);
+  }
+  if constexpr (row_tile > 2) {
+    store_bfloat(
+        args.out,
+        output_offset + 2 * args.out_features,
+        horizontal_sum(accumulator2) + bias);
+  }
+  if constexpr (row_tile > 3) {
+    store_bfloat(
+        args.out,
+        output_offset + 3 * args.out_features,
+        horizontal_sum(accumulator3) + bias);
+  }
+  if constexpr (row_tile > 4) {
+    store_bfloat(
+        args.out,
+        output_offset + 4 * args.out_features,
+        horizontal_sum(accumulator4) + bias);
+  }
+  if constexpr (row_tile > 5) {
+    store_bfloat(
+        args.out,
+        output_offset + 5 * args.out_features,
+        horizontal_sum(accumulator5) + bias);
+  }
+  if constexpr (row_tile > 6) {
+    store_bfloat(
+        args.out,
+        output_offset + 6 * args.out_features,
+        horizontal_sum(accumulator6) + bias);
+  }
+  if constexpr (row_tile > 7) {
+    store_bfloat(
+        args.out,
+        output_offset + 7 * args.out_features,
+        horizontal_sum(accumulator7) + bias);
+  }
+}
+
+bool adaln_avx512_bf16_available() {
+  __builtin_cpu_init();
+  return orbitquant::cpu::packed_matmul_x86_avx512_available() &&
+      __builtin_cpu_supports("avx512bf16");
+}
+#else
+bool adaln_avx512_bf16_available() {
+  return false;
+}
+#endif
+
+template <
+    void (*rows8)(AdalnArgs const &, std::uint8_t const *, float const *, std::int64_t, std::int64_t),
+    void (*rows4)(AdalnArgs const &, std::uint8_t const *, float const *, std::int64_t, std::int64_t),
+    void (*rows3)(AdalnArgs const &, std::uint8_t const *, float const *, std::int64_t, std::int64_t),
+    void (*rows2)(AdalnArgs const &, std::uint8_t const *, float const *, std::int64_t, std::int64_t),
+    void (*rows1)(AdalnArgs const &, std::uint8_t const *, float const *, std::int64_t, std::int64_t)>
+void packed_adaln_x86_tiled_range(
+    AdalnArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  const std::int64_t packed_row_bytes = args.padded_in_features / 2;
+  for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) {
+    const auto *packed_row =
+        args.packed_weight + out_col * packed_row_bytes;
+    const auto *scale_row =
+        args.scales + out_col * args.num_groups;
+    std::int64_t row = 0;
+    for (; row + 8 <= args.rows; row += 8) {
+      rows8(args, packed_row, scale_row, out_col, row);
+    }
+    if (row + 4 <= args.rows) {
+      rows4(args, packed_row, scale_row, out_col, row);
+      row += 4;
+    }
+    switch (args.rows - row) {
+      case 3:
+        rows3(args, packed_row, scale_row, out_col, row);
+        break;
+      case 2:
+        rows2(args, packed_row, scale_row, out_col, row);
+        break;
+      case 1:
+        rows1(args, packed_row, scale_row, out_col, row);
+        break;
+      default:
+        break;
+    }
+  }
+}
+
+void packed_adaln_avx2_range(
+    AdalnArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  if (args.group_size % 8 != 0 || args.padded_in_features % 2 != 0) {
+    packed_adaln_scalar_range(args, out_start, out_end);
+    return;
+  }
+  packed_adaln_x86_tiled_range<
+      packed_adaln_avx2_rows<8>,
+      packed_adaln_avx2_rows<4>,
+      packed_adaln_avx2_rows<3>,
+      packed_adaln_avx2_rows<2>,
+      packed_adaln_avx2_rows<1>>(args, out_start, out_end);
+}
+
+void packed_adaln_avx512_range(
+    AdalnArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+#if defined(ORBITQUANT_HAS_AVX512_BF16_INTRINSICS)
+  if (args.group_size % 32 == 0 &&
+      args.padded_in_features % 2 == 0 &&
+      adaln_avx512_bf16_available()) {
+    if (args.in_features % args.group_size == 0) {
+      packed_adaln_x86_tiled_range<
+          packed_adaln_avx512_bf16_rows<8, true>,
+          packed_adaln_avx512_bf16_rows<4, true>,
+          packed_adaln_avx512_bf16_rows<3, true>,
+          packed_adaln_avx512_bf16_rows<2, true>,
+          packed_adaln_avx512_bf16_rows<1, true>>(args, out_start, out_end);
+      return;
+    }
+    packed_adaln_x86_tiled_range<
+        packed_adaln_avx512_bf16_rows<8, false>,
+        packed_adaln_avx512_bf16_rows<4, false>,
+        packed_adaln_avx512_bf16_rows<3, false>,
+        packed_adaln_avx512_bf16_rows<2, false>,
+        packed_adaln_avx512_bf16_rows<1, false>>(args, out_start, out_end);
+    return;
+  }
+#endif
+  if (args.group_size % 16 != 0 || args.padded_in_features % 2 != 0) {
+    packed_adaln_scalar_range(args, out_start, out_end);
+    return;
+  }
+  packed_adaln_x86_tiled_range<
+      packed_adaln_avx512_rows<8>,
+      packed_adaln_avx512_rows<4>,
+      packed_adaln_avx512_rows<3>,
+      packed_adaln_avx512_rows<2>,
+      packed_adaln_avx512_rows<1>>(args, out_start, out_end);
+}
+#else
+void packed_adaln_avx2_range(
+    AdalnArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  packed_adaln_scalar_range(args, out_start, out_end);
+}
+
+void packed_adaln_avx512_range(
+    AdalnArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  packed_adaln_scalar_range(args, out_start, out_end);
+}
+#endif
+
+AdalnRangeFn select_adaln_kernel() {
+  const char *requested = std::getenv("ORBITQUANT_CPU_ISA");
+  if (requested == nullptr || std::strcmp(requested, "auto") == 0) {
+    if (orbitquant::cpu::packed_matmul_x86_avx512_available()) {
+      return packed_adaln_avx512_range;
+    }
+    if (orbitquant::cpu::packed_matmul_x86_avx2_available()) {
+#if defined(_MSC_VER) && defined(_M_X64)
+      return orbitquant::cpu::packed_adaln_msvc_avx2_range;
+#else
+      return packed_adaln_avx2_range;
+#endif
+    }
+    if (orbitquant::cpu::packed_matmul_neon_available()) {
+      return packed_adaln_neon_range;
+    }
+    return packed_adaln_scalar_range;
+  }
+  if (std::strcmp(requested, "scalar") == 0) {
+    return packed_adaln_scalar_range;
+  }
+  if (std::strcmp(requested, "avx2") == 0) {
+    STD_TORCH_CHECK(
+        orbitquant::cpu::packed_matmul_x86_avx2_available(),
+        "ORBITQUANT_CPU_ISA=avx2 requested AVX2/FMA/F16C on an unsupported CPU");
+#if defined(_MSC_VER) && defined(_M_X64)
+    return orbitquant::cpu::packed_adaln_msvc_avx2_range;
+#else
+    return packed_adaln_avx2_range;
+#endif
+  }
+  if (std::strcmp(requested, "avx512") == 0) {
+    STD_TORCH_CHECK(
+        orbitquant::cpu::packed_matmul_x86_avx512_available(),
+        "ORBITQUANT_CPU_ISA=avx512 requested AVX-512F/DQ/BW/VL on an unsupported CPU");
+    return packed_adaln_avx512_range;
+  }
+  if (std::strcmp(requested, "neon") == 0) {
+    STD_TORCH_CHECK(
+        orbitquant::cpu::packed_matmul_neon_available(),
+        "ORBITQUANT_CPU_ISA=neon requested NEON on an unsupported CPU");
+    return packed_adaln_neon_range;
+  }
+  STD_TORCH_CHECK(
+      false,
+      "ORBITQUANT_CPU_ISA must be auto, scalar, avx2, avx512, or neon");
+  return packed_adaln_scalar_range;
+}
+
+void parallel_packed_adaln(
+    AdalnArgs const &args,
+    AdalnRangeFn function) {
+  const std::int64_t arithmetic =
+      args.rows * args.out_features * args.in_features;
+  const int max_threads = orbitquant::cpu::requested_threads();
+  const int threads = arithmetic < 1'000'000
+      ? 1
+      : std::max<int>(
+            1,
+            std::min<std::int64_t>(
+                max_threads,
+                args.out_features / 16));
+  if (threads == 1) {
+    function(args, 0, args.out_features);
+    return;
+  }
+
+  std::vector<std::pair<std::int64_t, std::int64_t>> ranges;
+  ranges.reserve(threads);
+  const std::int64_t columns_per_thread =
+      (args.out_features + threads - 1) / threads;
+  for (int thread = 0; thread < threads; ++thread) {
+    const std::int64_t start = thread * columns_per_thread;
+    const std::int64_t end = std::min(
+        args.out_features,
+        start + columns_per_thread);
+    if (start >= end) {
+      break;
+    }
+    ranges.emplace_back(start, end);
+  }
+  orbitquant::cpu::run_ranges(
+      ranges,
+      [&args, function](std::int64_t start, std::int64_t end) {
+        function(args, start, end);
+      });
+}
+
+}  // namespace
+
+void matmul_packed_adaln_int4_cpu(
+    OrbitQuantTensor &out,
+    OrbitQuantTensor const &x,
+    OrbitQuantTensor const &packed_weight,
+    OrbitQuantTensor const &scales,
+    OrbitQuantTensor const &bias,
+    bool has_bias,
+    int64_t out_features,
+    int64_t in_features,
+    int64_t group_size) {
+  using torch::headeronly::DeviceType;
+  using torch::headeronly::ScalarType;
+
+  STD_TORCH_CHECK(x.device().type() == DeviceType::CPU, "x must be a CPU tensor");
+  STD_TORCH_CHECK(out.device().type() == DeviceType::CPU, "out must be a CPU tensor");
+  STD_TORCH_CHECK(
+      packed_weight.device().type() == DeviceType::CPU,
+      "packed_weight must be a CPU tensor");
+  STD_TORCH_CHECK(
+      scales.device().type() == DeviceType::CPU,
+      "scales must be a CPU tensor");
+  STD_TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
+  STD_TORCH_CHECK(out.is_contiguous(), "out must be contiguous");
+  STD_TORCH_CHECK(packed_weight.is_contiguous(), "packed_weight must be contiguous");
+  STD_TORCH_CHECK(scales.is_contiguous(), "scales must be contiguous");
+  STD_TORCH_CHECK(x.scalar_type() == ScalarType::BFloat16, "x must be bfloat16");
+  STD_TORCH_CHECK(out.scalar_type() == ScalarType::BFloat16, "out must be bfloat16");
+  STD_TORCH_CHECK(
+      packed_weight.scalar_type() == ScalarType::Byte,
+      "packed_weight must be uint8");
+  STD_TORCH_CHECK(scales.scalar_type() == ScalarType::Float, "scales must be float32");
+  STD_TORCH_CHECK(x.dim() == 2, "x must be rank 2");
+  STD_TORCH_CHECK(out.dim() == 2, "out must be rank 2");
+  STD_TORCH_CHECK(group_size > 0, "group_size must be positive");
+  STD_TORCH_CHECK(x.size(1) == in_features, "x has an unexpected input dimension");
+  STD_TORCH_CHECK(out.size(0) == x.size(0), "out has an unexpected row count");
+  STD_TORCH_CHECK(out.size(1) == out_features, "out has an unexpected output dimension");
+  const int64_t num_groups = (in_features + group_size - 1) / group_size;
+  const int64_t padded_in_features = num_groups * group_size;
+  const int64_t packed_values = out_features * padded_in_features;
+  STD_TORCH_CHECK(
+      packed_weight.numel() >= (packed_values + 1) / 2,
+      "packed_weight is too short");
+  STD_TORCH_CHECK(
+      scales.numel() == out_features * num_groups,
+      "scales must match out_features and num_groups");
+  if (has_bias) {
+    STD_TORCH_CHECK(bias.device().type() == DeviceType::CPU, "bias must be a CPU tensor");
+    STD_TORCH_CHECK(bias.is_contiguous(), "bias must be contiguous");
+    STD_TORCH_CHECK(bias.scalar_type() == ScalarType::Float, "bias must be float32");
+    STD_TORCH_CHECK(bias.numel() == out_features, "bias must match out_features");
+  }
+  if (x.numel() == 0 || out_features == 0) {
+    return;
+  }
+
+  const AdalnArgs args{
+      out.mutable_data_ptr(),
+      x.const_data_ptr(),
+      packed_weight.const_data_ptr<std::uint8_t>(),
+      scales.const_data_ptr<float>(),
+      has_bias ? bias.const_data_ptr<float>() : nullptr,
+      has_bias,
+      x.size(0),
+      out_features,
+      in_features,
+      group_size,
+      num_groups,
+      padded_in_features,
+  };
+  parallel_packed_adaln(args, select_adaln_kernel());
+}

orbitquant_packed_matmul_cpu/packed_matmul_cpu.cpp

@@ -0,0 +1,201 @@
+#include "cpu_pool.h"
+#include "cpu_threads.h"
+#include "packed_matmul_cpu.h"
+#include "../torch-ext/torch_binding.h"
+
+#include <torch/headeronly/core/DeviceType.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/macros/Macros.h>
+
+#include <algorithm>
+#include <cstdlib>
+#include <cstring>
+#include <utility>
+#include <vector>
+
+namespace {
+
+orbitquant::cpu::ScalarKind scalar_kind(OrbitQuantTensor const &tensor) {
+  using torch::headeronly::ScalarType;
+  switch (tensor.scalar_type()) {
+    case ScalarType::Float:
+      return orbitquant::cpu::ScalarKind::Float32;
+    case ScalarType::Half:
+      return orbitquant::cpu::ScalarKind::Float16;
+    case ScalarType::BFloat16:
+      return orbitquant::cpu::ScalarKind::BFloat16;
+    default:
+      STD_TORCH_CHECK(
+          false,
+          "CPU packed matmul supports float32, float16, and bfloat16 inputs");
+  }
+  return orbitquant::cpu::ScalarKind::Float32;
+}
+
+void parallel_packed_matmul(
+    orbitquant::cpu::PackedMatmulArgs const &args,
+    orbitquant::cpu::PackedMatmulRangeFn function) {
+  const std::int64_t arithmetic =
+      args.rows * args.out_features * args.in_features;
+  const int max_threads = orbitquant::cpu::requested_threads();
+  const int threads = arithmetic < 1'000'000
+      ? 1
+      : std::max<int>(
+            1,
+            std::min<std::int64_t>(max_threads, args.out_features / 16));
+  if (threads == 1) {
+    function(args, 0, args.out_features);
+    return;
+  }
+
+  std::vector<std::pair<std::int64_t, std::int64_t>> ranges;
+  ranges.reserve(threads);
+  const std::int64_t columns_per_thread =
+      (args.out_features + threads - 1) / threads;
+  for (int thread = 0; thread < threads; ++thread) {
+    const std::int64_t start = thread * columns_per_thread;
+    const std::int64_t end = std::min(
+        args.out_features,
+        start + columns_per_thread);
+    if (start >= end) {
+      break;
+    }
+    ranges.emplace_back(start, end);
+  }
+  orbitquant::cpu::run_ranges(
+      ranges,
+      [&args, function](std::int64_t start, std::int64_t end) {
+        function(args, start, end);
+      });
+}
+
+orbitquant::cpu::PackedMatmulRangeFn select_packed_matmul() {
+  const char *requested = std::getenv("ORBITQUANT_CPU_ISA");
+  if (requested == nullptr || std::strcmp(requested, "auto") == 0) {
+    if (orbitquant::cpu::packed_matmul_x86_avx512_available()) {
+      return orbitquant::cpu::packed_matmul_x86_avx512_range;
+    }
+    if (orbitquant::cpu::packed_matmul_x86_avx2_available()) {
+      return orbitquant::cpu::packed_matmul_x86_avx2_range;
+    }
+    if (orbitquant::cpu::packed_matmul_neon_available()) {
+      return orbitquant::cpu::packed_matmul_neon_range;
+    }
+    return orbitquant::cpu::packed_matmul_scalar_range;
+  }
+  if (std::strcmp(requested, "scalar") == 0) {
+    return orbitquant::cpu::packed_matmul_scalar_range;
+  }
+  if (std::strcmp(requested, "avx2") == 0) {
+    STD_TORCH_CHECK(
+        orbitquant::cpu::packed_matmul_x86_avx2_available(),
+        "ORBITQUANT_CPU_ISA=avx2 requested AVX2/FMA/F16C on an unsupported CPU");
+    return orbitquant::cpu::packed_matmul_x86_avx2_range;
+  }
+  if (std::strcmp(requested, "avx512") == 0) {
+    STD_TORCH_CHECK(
+        orbitquant::cpu::packed_matmul_x86_avx512_available(),
+        "ORBITQUANT_CPU_ISA=avx512 requested AVX-512F/DQ/BW/VL on an unsupported CPU");
+    return orbitquant::cpu::packed_matmul_x86_avx512_range;
+  }
+  if (std::strcmp(requested, "neon") == 0) {
+    STD_TORCH_CHECK(
+        orbitquant::cpu::packed_matmul_neon_available(),
+        "ORBITQUANT_CPU_ISA=neon requested NEON on an unsupported CPU");
+    return orbitquant::cpu::packed_matmul_neon_range;
+  }
+  STD_TORCH_CHECK(
+      false,
+      "ORBITQUANT_CPU_ISA must be auto, scalar, avx2, avx512, or neon");
+  return orbitquant::cpu::packed_matmul_scalar_range;
+}
+
+}  // namespace
+
+void matmul_packed_weight(
+    OrbitQuantTensor &out,
+    OrbitQuantTensor const &x,
+    OrbitQuantTensor const &packed_weight_indices,
+    OrbitQuantTensor const &row_norms,
+    OrbitQuantTensor const &centroids,
+    OrbitQuantTensor const &bias,
+    bool has_bias,
+    int64_t bits,
+    int64_t out_features,
+    int64_t in_features,
+    int64_t block_m,
+    int64_t block_n,
+    int64_t block_k) {
+  using torch::headeronly::DeviceType;
+  using torch::headeronly::ScalarType;
+
+  STD_TORCH_CHECK(x.device().type() == DeviceType::CPU, "x must be a CPU tensor");
+  STD_TORCH_CHECK(out.device().type() == DeviceType::CPU, "out must be a CPU tensor");
+  STD_TORCH_CHECK(
+      packed_weight_indices.device().type() == DeviceType::CPU,
+      "packed weights must be CPU tensors");
+  STD_TORCH_CHECK(
+      row_norms.device().type() == DeviceType::CPU,
+      "row norms must be CPU tensors");
+  STD_TORCH_CHECK(
+      centroids.device().type() == DeviceType::CPU,
+      "centroids must be CPU tensors");
+  STD_TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
+  STD_TORCH_CHECK(out.is_contiguous(), "out must be contiguous");
+  STD_TORCH_CHECK(
+      packed_weight_indices.is_contiguous(), "packed weights must be contiguous");
+  STD_TORCH_CHECK(row_norms.is_contiguous(), "row norms must be contiguous");
+  STD_TORCH_CHECK(centroids.is_contiguous(), "centroids must be contiguous");
+  STD_TORCH_CHECK(
+      packed_weight_indices.scalar_type() == ScalarType::Byte,
+      "packed weights must be uint8");
+  STD_TORCH_CHECK(
+      row_norms.scalar_type() == ScalarType::Float,
+      "row_norms must be float32");
+  STD_TORCH_CHECK(
+      centroids.scalar_type() == ScalarType::Float,
+      "centroids must be float32");
+  STD_TORCH_CHECK(out.scalar_type() == x.scalar_type(), "out dtype must match x dtype");
+  STD_TORCH_CHECK(x.dim() == 2, "x must be rank 2");
+  STD_TORCH_CHECK(out.dim() == 2, "out must be rank 2");
+  STD_TORCH_CHECK(bits > 0 && bits <= 8, "bits must be in [1, 8]");
+  STD_TORCH_CHECK(block_m > 0 && block_n > 0 && block_k > 0, "tile sizes must be positive");
+  STD_TORCH_CHECK(x.size(1) == in_features, "x has an unexpected input dimension");
+  STD_TORCH_CHECK(out.size(0) == x.size(0), "out has an unexpected row count");
+  STD_TORCH_CHECK(out.size(1) == out_features, "out has an unexpected output dimension");
+  STD_TORCH_CHECK(row_norms.numel() == out_features, "row_norms must match out_features");
+  STD_TORCH_CHECK(centroids.numel() >= (1LL << bits), "centroids are too short");
+  const int64_t packed_bytes = (out_features * in_features * bits + 7) / 8;
+  STD_TORCH_CHECK(
+      packed_weight_indices.numel() >= packed_bytes,
+      "packed weights are too short");
+  if (has_bias) {
+    STD_TORCH_CHECK(
+        bias.device().type() == DeviceType::CPU,
+        "bias must be a CPU tensor");
+    STD_TORCH_CHECK(bias.is_contiguous(), "bias must be contiguous");
+    STD_TORCH_CHECK(
+        bias.scalar_type() == ScalarType::Float,
+        "bias must be float32");
+    STD_TORCH_CHECK(bias.numel() == out_features, "bias must match out_features");
+  }
+  if (x.numel() == 0 || out_features == 0) {
+    return;
+  }
+
+  orbitquant::cpu::PackedMatmulArgs args{
+      out.mutable_data_ptr(),
+      x.const_data_ptr(),
+      packed_weight_indices.const_data_ptr<std::uint8_t>(),
+      row_norms.const_data_ptr<float>(),
+      centroids.const_data_ptr<float>(),
+      has_bias ? bias.const_data_ptr<float>() : nullptr,
+      has_bias,
+      scalar_kind(x),
+      x.size(0),
+      out_features,
+      in_features,
+      bits,
+  };
+  parallel_packed_matmul(args, select_packed_matmul());
+}

orbitquant_packed_matmul_cpu/packed_matmul_cpu.h

@@ -0,0 +1,59 @@
+#pragma once
+
+#include <cstdint>
+
+namespace orbitquant::cpu {
+
+enum class ScalarKind : std::uint8_t {
+  Float32,
+  Float16,
+  BFloat16,
+};
+
+struct PackedMatmulArgs {
+  void *out;
+  void const *x;
+  std::uint8_t const *packed_weight_indices;
+  float const *row_norms;
+  float const *centroids;
+  float const *bias;
+  bool has_bias;
+  ScalarKind scalar_kind;
+  std::int64_t rows;
+  std::int64_t out_features;
+  std::int64_t in_features;
+  std::int64_t bits;
+};
+
+using PackedMatmulRangeFn = void (*)(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end);
+
+void packed_matmul_scalar_range(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end);
+
+bool packed_matmul_neon_available();
+
+void packed_matmul_neon_range(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end);
+
+bool packed_matmul_x86_avx2_available();
+
+void packed_matmul_x86_avx2_range(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end);
+
+bool packed_matmul_x86_avx512_available();
+
+void packed_matmul_x86_avx512_range(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end);
+
+}  // namespace orbitquant::cpu

orbitquant_packed_matmul_cpu/packed_matmul_neon.cpp

@@ -0,0 +1,484 @@
+#include "packed_matmul_cpu.h"
+
+#if defined(__aarch64__) || defined(_M_ARM64)
+#include <arm_neon.h>
+
+#include <torch/headeronly/util/BFloat16.h>
+#include <torch/headeronly/util/Half.h>
+
+#include <type_traits>
+#include <vector>
+
+namespace orbitquant::cpu {
+namespace {
+
+inline float horizontal_sum(float32x4_t value) {
+#if defined(__aarch64__)
+  return vaddvq_f32(value);
+#else
+  const float32x2_t pair = vadd_f32(vget_low_f32(value), vget_high_f32(value));
+  return vget_lane_f32(vpadd_f32(pair, pair), 0);
+#endif
+}
+
+inline float32x4_t load_float4(void const *data, std::int64_t offset) {
+  return vld1q_f32(static_cast<float const *>(data) + offset);
+}
+
+inline float32x4_t load_half4(void const *data, std::int64_t offset) {
+  const auto *source = reinterpret_cast<float16_t const *>(
+      static_cast<std::uint16_t const *>(data) + offset);
+  return vcvt_f32_f16(vld1_f16(source));
+}
+
+inline float32x4_t load_bfloat4(void const *data, std::int64_t offset) {
+  const uint16x4_t raw = vld1_u16(
+      static_cast<std::uint16_t const *>(data) + offset);
+  return vreinterpretq_f32_u32(vshlq_n_u32(vmovl_u16(raw), 16));
+}
+
+template <typename scalar_t>
+inline void store_value(void *data, std::int64_t offset, float value) {
+  static_cast<scalar_t *>(data)[offset] = scalar_t(value);
+}
+
+template <>
+inline void store_value<float>(void *data, std::int64_t offset, float value) {
+  static_cast<float *>(data)[offset] = value;
+}
+
+template <
+    typename scalar_t,
+    float32x4_t (*load4)(void const *, std::int64_t),
+    int row_tile>
+inline void packed_matmul_neon_w4_rows(
+    PackedMatmulArgs const &args,
+    std::uint8_t const *packed_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  float32x4_t accumulator0[row_tile];
+  float32x4_t accumulator1[row_tile];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    accumulator0[row] = vdupq_n_f32(0.0f);
+    accumulator1[row] = vdupq_n_f32(0.0f);
+  }
+
+  std::int64_t k = 0;
+  for (; k + 8 <= args.in_features; k += 8) {
+    const std::int64_t byte_offset = k / 2;
+    float weights0[4];
+    float weights1[4];
+#pragma clang loop unroll(full)
+    for (int pair = 0; pair < 4; ++pair) {
+      const std::uint8_t packed = packed_row[byte_offset + pair];
+      const int value = pair * 2;
+      if (value < 4) {
+        weights0[value] = args.centroids[packed & 15u];
+        weights0[value + 1] = args.centroids[(packed >> 4) & 15u];
+      } else {
+        weights1[value - 4] = args.centroids[packed & 15u];
+        weights1[value - 3] = args.centroids[(packed >> 4) & 15u];
+      }
+    }
+    const float32x4_t weight0 = vld1q_f32(weights0);
+    const float32x4_t weight1 = vld1q_f32(weights1);
+#pragma clang loop unroll(full)
+    for (int row = 0; row < row_tile; ++row) {
+      const std::int64_t input_offset =
+          (row_start + row) * args.in_features + k;
+      accumulator0[row] =
+          vfmaq_f32(accumulator0[row], load4(args.x, input_offset), weight0);
+      accumulator1[row] = vfmaq_f32(
+          accumulator1[row], load4(args.x, input_offset + 4), weight1);
+    }
+  }
+
+  const float row_norm = args.row_norms[out_col];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    const std::int64_t input_row_offset =
+        (row_start + row) * args.in_features;
+    float accumulator =
+        horizontal_sum(vaddq_f32(accumulator0[row], accumulator1[row]));
+    for (std::int64_t tail = k; tail < args.in_features; ++tail) {
+      const std::uint8_t packed = packed_row[tail / 2];
+      const std::uint8_t index =
+          (tail & 1) == 0 ? packed & 15u : (packed >> 4) & 15u;
+      if constexpr (std::is_same_v<scalar_t, float>) {
+        accumulator +=
+            static_cast<float const *>(args.x)[input_row_offset + tail] *
+            args.centroids[index];
+      } else {
+        accumulator += static_cast<float>(
+                           static_cast<scalar_t const *>(
+                               args.x)[input_row_offset + tail]) *
+            args.centroids[index];
+      }
+    }
+    accumulator *= row_norm;
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_value<scalar_t>(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+template <int Bits>
+inline std::uint32_t unpack_index_generic(
+    std::uint8_t const *packed_row,
+    std::int64_t value_index) {
+  const std::int64_t bit_start = value_index * Bits;
+  const std::int64_t byte_index = bit_start >> 3;
+  const unsigned bit_offset = static_cast<unsigned>(bit_start & 7);
+  std::uint32_t raw = packed_row[byte_index];
+  if (bit_offset + static_cast<unsigned>(Bits) > 8) {
+    raw |= static_cast<std::uint32_t>(packed_row[byte_index + 1]) << 8;
+  }
+  return (raw >> bit_offset) & ((1u << Bits) - 1u);
+}
+
+template <
+    typename scalar_t,
+    float32x4_t (*load4)(void const *, std::int64_t),
+    int Bits,
+    int row_tile>
+inline void packed_matmul_neon_lowbit_rows(
+    PackedMatmulArgs const &args,
+    std::uint8_t const *packed_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  float32x4_t accumulator0[row_tile];
+  float32x4_t accumulator1[row_tile];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    accumulator0[row] = vdupq_n_f32(0.0f);
+    accumulator1[row] = vdupq_n_f32(0.0f);
+  }
+
+  std::int64_t k = 0;
+  for (; k + 8 <= args.in_features; k += 8) {
+    float weights0[4];
+    float weights1[4];
+#pragma clang loop unroll(full)
+    for (int value = 0; value < 4; ++value) {
+      weights0[value] =
+          args.centroids[unpack_index_generic<Bits>(packed_row, k + value)];
+      weights1[value] =
+          args.centroids[unpack_index_generic<Bits>(packed_row, k + 4 + value)];
+    }
+    const float32x4_t weight0 = vld1q_f32(weights0);
+    const float32x4_t weight1 = vld1q_f32(weights1);
+#pragma clang loop unroll(full)
+    for (int row = 0; row < row_tile; ++row) {
+      const std::int64_t input_offset =
+          (row_start + row) * args.in_features + k;
+      accumulator0[row] =
+          vfmaq_f32(accumulator0[row], load4(args.x, input_offset), weight0);
+      accumulator1[row] = vfmaq_f32(
+          accumulator1[row], load4(args.x, input_offset + 4), weight1);
+    }
+  }
+
+  const float row_norm = args.row_norms[out_col];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    const std::int64_t input_row_offset =
+        (row_start + row) * args.in_features;
+    float accumulator =
+        horizontal_sum(vaddq_f32(accumulator0[row], accumulator1[row]));
+    for (std::int64_t tail = k; tail < args.in_features; ++tail) {
+      const std::uint32_t index =
+          unpack_index_generic<Bits>(packed_row, tail);
+      if constexpr (std::is_same_v<scalar_t, float>) {
+        accumulator +=
+            static_cast<float const *>(args.x)[input_row_offset + tail] *
+            args.centroids[index];
+      } else {
+        accumulator += static_cast<float>(
+                           static_cast<scalar_t const *>(
+                               args.x)[input_row_offset + tail]) *
+            args.centroids[index];
+      }
+    }
+    accumulator *= row_norm;
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_value<scalar_t>(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+template <typename scalar_t, float32x4_t (*load4)(void const *, std::int64_t), int row_tile>
+inline void packed_matmul_neon_buffered_rows(
+    PackedMatmulArgs const &args,
+    float const *decoded_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  float32x4_t accumulator0[row_tile];
+  float32x4_t accumulator1[row_tile];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    accumulator0[row] = vdupq_n_f32(0.0f);
+    accumulator1[row] = vdupq_n_f32(0.0f);
+  }
+
+  std::int64_t k = 0;
+  for (; k + 8 <= args.in_features; k += 8) {
+    const float32x4_t weight0 = vld1q_f32(decoded_row + k);
+    const float32x4_t weight1 = vld1q_f32(decoded_row + k + 4);
+#pragma clang loop unroll(full)
+    for (int row = 0; row < row_tile; ++row) {
+      const std::int64_t input_offset =
+          (row_start + row) * args.in_features + k;
+      accumulator0[row] =
+          vfmaq_f32(accumulator0[row], load4(args.x, input_offset), weight0);
+      accumulator1[row] = vfmaq_f32(
+          accumulator1[row], load4(args.x, input_offset + 4), weight1);
+    }
+  }
+
+  const float row_norm = args.row_norms[out_col];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    const std::int64_t input_row_offset =
+        (row_start + row) * args.in_features;
+    float accumulator =
+        horizontal_sum(vaddq_f32(accumulator0[row], accumulator1[row]));
+    for (std::int64_t tail = k; tail < args.in_features; ++tail) {
+      if constexpr (std::is_same_v<scalar_t, float>) {
+        accumulator +=
+            static_cast<float const *>(args.x)[input_row_offset + tail] *
+            decoded_row[tail];
+      } else {
+        accumulator += static_cast<float>(
+                           static_cast<scalar_t const *>(
+                               args.x)[input_row_offset + tail]) *
+            decoded_row[tail];
+      }
+    }
+    accumulator *= row_norm;
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_value<scalar_t>(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+template <
+    typename scalar_t,
+    float32x4_t (*load4)(void const *, std::int64_t),
+    int Bits>
+void packed_matmul_neon_lowbit_typed(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  const std::int64_t packed_row_bytes = args.in_features * Bits / 8;
+  // The NEON decode goes through scalar table lookups, so expanding the
+  // column once per >=16-row call amortizes the costliest stage.
+  const bool use_decoded_buffer = args.rows >= 16;
+  thread_local std::vector<float> decoded_row_storage;
+  if (use_decoded_buffer &&
+      decoded_row_storage.size() < static_cast<std::size_t>(args.in_features)) {
+    decoded_row_storage.resize(static_cast<std::size_t>(args.in_features));
+  }
+  for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) {
+    const auto *packed_row =
+        args.packed_weight_indices + out_col * packed_row_bytes;
+    if (use_decoded_buffer) {
+      float *decoded_row = decoded_row_storage.data();
+      for (std::int64_t k = 0; k < args.in_features; ++k) {
+        decoded_row[k] =
+            args.centroids[unpack_index_generic<Bits>(packed_row, k)];
+      }
+      std::int64_t row = 0;
+      for (; row + 8 <= args.rows; row += 8) {
+        packed_matmul_neon_buffered_rows<scalar_t, load4, 8>(
+            args, decoded_row, out_col, row);
+      }
+      for (; row + 4 <= args.rows; row += 4) {
+        packed_matmul_neon_buffered_rows<scalar_t, load4, 4>(
+            args, decoded_row, out_col, row);
+      }
+      switch (args.rows - row) {
+        case 3:
+          packed_matmul_neon_buffered_rows<scalar_t, load4, 3>(
+              args, decoded_row, out_col, row);
+          break;
+        case 2:
+          packed_matmul_neon_buffered_rows<scalar_t, load4, 2>(
+              args, decoded_row, out_col, row);
+          break;
+        case 1:
+          packed_matmul_neon_buffered_rows<scalar_t, load4, 1>(
+              args, decoded_row, out_col, row);
+          break;
+        default:
+          break;
+      }
+      continue;
+    }
+    std::int64_t row = 0;
+    for (; row + 8 <= args.rows; row += 8) {
+      packed_matmul_neon_lowbit_rows<scalar_t, load4, Bits, 8>(
+          args, packed_row, out_col, row);
+    }
+    for (; row + 4 <= args.rows; row += 4) {
+      packed_matmul_neon_lowbit_rows<scalar_t, load4, Bits, 4>(
+          args, packed_row, out_col, row);
+    }
+    switch (args.rows - row) {
+      case 3:
+        packed_matmul_neon_lowbit_rows<scalar_t, load4, Bits, 3>(
+            args, packed_row, out_col, row);
+        break;
+      case 2:
+        packed_matmul_neon_lowbit_rows<scalar_t, load4, Bits, 2>(
+            args, packed_row, out_col, row);
+        break;
+      case 1:
+        packed_matmul_neon_lowbit_rows<scalar_t, load4, Bits, 1>(
+            args, packed_row, out_col, row);
+        break;
+      default:
+        break;
+    }
+  }
+}
+
+template <int Bits>
+void packed_matmul_neon_lowbit_dispatch(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  switch (args.scalar_kind) {
+    case ScalarKind::Float32:
+      packed_matmul_neon_lowbit_typed<float, load_float4, Bits>(
+          args, out_start, out_end);
+      return;
+    case ScalarKind::Float16:
+      packed_matmul_neon_lowbit_typed<c10::Half, load_half4, Bits>(
+          args, out_start, out_end);
+      return;
+    case ScalarKind::BFloat16:
+      packed_matmul_neon_lowbit_typed<c10::BFloat16, load_bfloat4, Bits>(
+          args, out_start, out_end);
+      return;
+  }
+}
+
+template <typename scalar_t, float32x4_t (*load4)(void const *, std::int64_t)>
+void packed_matmul_neon_w4_typed(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  constexpr int kPrimaryRowTile = 8;
+  const std::int64_t packed_row_bytes = args.in_features / 2;
+  for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) {
+    const auto *packed_row =
+        args.packed_weight_indices + out_col * packed_row_bytes;
+    std::int64_t row = 0;
+    if constexpr (kPrimaryRowTile == 8) {
+      for (; row + 8 <= args.rows; row += 8) {
+        packed_matmul_neon_w4_rows<scalar_t, load4, 8>(
+            args, packed_row, out_col, row);
+      }
+    }
+    for (; row + 4 <= args.rows; row += 4) {
+      packed_matmul_neon_w4_rows<scalar_t, load4, 4>(
+          args, packed_row, out_col, row);
+    }
+    switch (args.rows - row) {
+      case 3:
+        packed_matmul_neon_w4_rows<scalar_t, load4, 3>(
+            args, packed_row, out_col, row);
+        break;
+      case 2:
+        packed_matmul_neon_w4_rows<scalar_t, load4, 2>(
+            args, packed_row, out_col, row);
+        break;
+      case 1:
+        packed_matmul_neon_w4_rows<scalar_t, load4, 1>(
+            args, packed_row, out_col, row);
+        break;
+      default:
+        break;
+    }
+  }
+}
+
+}  // namespace
+
+bool packed_matmul_neon_available() {
+  return true;
+}
+
+void packed_matmul_neon_range(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  if (args.bits == 4 && args.in_features % 2 == 0 && args.rows >= 16) {
+    // The buffered generic path decodes each column once, which beats the
+    // per-tile scalar decode as soon as several row tiles reuse it.
+    packed_matmul_neon_lowbit_dispatch<4>(args, out_start, out_end);
+    return;
+  }
+  if (args.bits == 2 && args.in_features % 4 == 0) {
+    packed_matmul_neon_lowbit_dispatch<2>(args, out_start, out_end);
+    return;
+  }
+  if (args.bits == 3 && args.in_features % 8 == 0) {
+    packed_matmul_neon_lowbit_dispatch<3>(args, out_start, out_end);
+    return;
+  }
+  if (args.bits == 6 && args.in_features % 4 == 0) {
+    packed_matmul_neon_lowbit_dispatch<6>(args, out_start, out_end);
+    return;
+  }
+  if (args.bits != 4 || args.in_features % 2 != 0) {
+    packed_matmul_scalar_range(args, out_start, out_end);
+    return;
+  }
+  switch (args.scalar_kind) {
+    case ScalarKind::Float32:
+      packed_matmul_neon_w4_typed<float, load_float4>(args, out_start, out_end);
+      return;
+    case ScalarKind::Float16:
+      packed_matmul_neon_w4_typed<c10::Half, load_half4>(args, out_start, out_end);
+      return;
+    case ScalarKind::BFloat16:
+      packed_matmul_neon_w4_typed<c10::BFloat16, load_bfloat4>(
+          args, out_start, out_end);
+      return;
+  }
+}
+
+}  // namespace orbitquant::cpu
+
+#else
+
+namespace orbitquant::cpu {
+
+bool packed_matmul_neon_available() {
+  return false;
+}
+
+void packed_matmul_neon_range(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  packed_matmul_scalar_range(args, out_start, out_end);
+}
+
+}  // namespace orbitquant::cpu
+
+#endif

orbitquant_packed_matmul_cpu/packed_matmul_scalar.cpp

@@ -0,0 +1,95 @@
+#include "packed_matmul_cpu.h"
+
+#include <torch/headeronly/util/BFloat16.h>
+#include <torch/headeronly/util/Half.h>
+
+#include <cstddef>
+
+namespace orbitquant::cpu {
+namespace {
+
+template <typename scalar_t>
+inline float load_scalar(void const *data, std::int64_t offset) {
+  return static_cast<float>(static_cast<scalar_t const *>(data)[offset]);
+}
+
+template <typename scalar_t>
+inline void store_scalar(void *data, std::int64_t offset, float value) {
+  static_cast<scalar_t *>(data)[offset] = scalar_t(value);
+}
+
+template <>
+inline float load_scalar<float>(void const *data, std::int64_t offset) {
+  return static_cast<float const *>(data)[offset];
+}
+
+template <>
+inline void store_scalar<float>(void *data, std::int64_t offset, float value) {
+  static_cast<float *>(data)[offset] = value;
+}
+
+inline std::uint32_t unpack_index(
+    std::uint8_t const *packed,
+    std::int64_t value_offset,
+    std::int64_t bits) {
+  const std::int64_t bit_start = value_offset * bits;
+  const std::int64_t byte_index = bit_start >> 3;
+  const unsigned bit_offset = static_cast<unsigned>(bit_start & 7);
+  std::uint32_t raw = packed[byte_index];
+  if (bit_offset + static_cast<unsigned>(bits) > 8) {
+    raw |= static_cast<std::uint32_t>(packed[byte_index + 1]) << 8;
+  }
+  return (raw >> bit_offset) & ((1u << static_cast<unsigned>(bits)) - 1u);
+}
+
+template <typename scalar_t>
+void packed_matmul_scalar_typed(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) {
+    const float row_norm = args.row_norms[out_col];
+    const std::int64_t weight_row_offset = out_col * args.in_features;
+    for (std::int64_t row = 0; row < args.rows; ++row) {
+      const std::int64_t input_row_offset = row * args.in_features;
+      float accumulator = 0.0f;
+      for (std::int64_t k = 0; k < args.in_features; ++k) {
+        const std::uint32_t index = unpack_index(
+            args.packed_weight_indices,
+            weight_row_offset + k,
+            args.bits);
+        accumulator += load_scalar<scalar_t>(args.x, input_row_offset + k) *
+            args.centroids[index];
+      }
+      accumulator *= row_norm;
+      if (args.has_bias) {
+        accumulator += args.bias[out_col];
+      }
+      store_scalar<scalar_t>(
+          args.out,
+          row * args.out_features + out_col,
+          accumulator);
+    }
+  }
+}
+
+}  // namespace
+
+void packed_matmul_scalar_range(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  switch (args.scalar_kind) {
+    case ScalarKind::Float32:
+      packed_matmul_scalar_typed<float>(args, out_start, out_end);
+      return;
+    case ScalarKind::Float16:
+      packed_matmul_scalar_typed<c10::Half>(args, out_start, out_end);
+      return;
+    case ScalarKind::BFloat16:
+      packed_matmul_scalar_typed<c10::BFloat16>(args, out_start, out_end);
+      return;
+  }
+}
+
+}  // namespace orbitquant::cpu

orbitquant_packed_matmul_cpu/packed_matmul_x86.cpp

@@ -0,0 +1,616 @@
+#include "packed_matmul_cpu.h"
+
+#if defined(__x86_64__) || defined(_M_X64)
+#include <immintrin.h>
+#if defined(_MSC_VER)
+#include <intrin.h>
+#else
+#include <cpuid.h>
+#endif
+
+#include <torch/headeronly/util/BFloat16.h>
+#include <torch/headeronly/util/Half.h>
+
+#include <cstring>
+#include <cstdint>
+#include <type_traits>
+#include <vector>
+
+#if defined(_MSC_VER)
+#define ORBITQUANT_TARGET_AVX2
+#define ORBITQUANT_NOINLINE __declspec(noinline)
+#else
+#define ORBITQUANT_TARGET_AVX2 __attribute__((target("avx2,fma,f16c")))
+#define ORBITQUANT_NOINLINE __attribute__((noinline))
+#endif
+
+namespace orbitquant::cpu {
+namespace {
+
+ORBITQUANT_TARGET_AVX2 inline float horizontal_sum(__m256 value) {
+  const __m128 halves =
+      _mm_add_ps(_mm256_castps256_ps128(value), _mm256_extractf128_ps(value, 1));
+  const __m128 pairs = _mm_hadd_ps(halves, halves);
+  return _mm_cvtss_f32(_mm_hadd_ps(pairs, pairs));
+}
+
+ORBITQUANT_TARGET_AVX2 inline __m256 load_float8(
+    void const *data,
+    std::int64_t offset) {
+  return _mm256_loadu_ps(static_cast<float const *>(data) + offset);
+}
+
+ORBITQUANT_TARGET_AVX2 inline __m256 load_half8(
+    void const *data,
+    std::int64_t offset) {
+  const auto *source = static_cast<std::uint16_t const *>(data) + offset;
+  const __m128i packed =
+      _mm_loadu_si128(reinterpret_cast<__m128i const *>(source));
+  return _mm256_cvtph_ps(packed);
+}
+
+ORBITQUANT_TARGET_AVX2 inline __m256 load_bfloat8(
+    void const *data,
+    std::int64_t offset) {
+  const auto *source = static_cast<std::uint16_t const *>(data) + offset;
+  const __m128i packed =
+      _mm_loadu_si128(reinterpret_cast<__m128i const *>(source));
+  const __m256i widened = _mm256_cvtepu16_epi32(packed);
+  return _mm256_castsi256_ps(_mm256_slli_epi32(widened, 16));
+}
+
+template <typename scalar_t>
+inline void store_value(void *data, std::int64_t offset, float value) {
+  static_cast<scalar_t *>(data)[offset] = scalar_t(value);
+}
+
+template <>
+inline void store_value<float>(void *data, std::int64_t offset, float value) {
+  static_cast<float *>(data)[offset] = value;
+}
+
+template <
+    typename scalar_t,
+    __m256 (*load8)(void const *, std::int64_t),
+    int row_tile>
+ORBITQUANT_TARGET_AVX2 inline void packed_matmul_avx2_w4_rows(
+    PackedMatmulArgs const &args,
+    std::uint8_t const *packed_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  __m256 accumulators[row_tile];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    accumulators[row] = _mm256_setzero_ps();
+  }
+  const __m256 centroid_lut_low = _mm256_loadu_ps(args.centroids);
+  const __m256 centroid_lut_high = _mm256_loadu_ps(args.centroids + 8);
+  const __m128i nibble_mask = _mm_set1_epi8(15);
+  const __m256i low_table_limit = _mm256_set1_epi32(7);
+
+  std::int64_t k = 0;
+  for (; k + 8 <= args.in_features; k += 8) {
+    const std::int64_t byte_offset = k / 2;
+    std::int32_t packed;
+    std::memcpy(&packed, packed_row + byte_offset, sizeof(packed));
+    const __m128i bytes = _mm_cvtsi32_si128(packed);
+    const __m128i low = _mm_and_si128(bytes, nibble_mask);
+    const __m128i high = _mm_and_si128(
+        _mm_srli_epi16(bytes, 4),
+        nibble_mask);
+    const __m256i indices =
+        _mm256_cvtepu8_epi32(_mm_unpacklo_epi8(low, high));
+    const __m256 low_weights =
+        _mm256_permutevar8x32_ps(centroid_lut_low, indices);
+    const __m256 high_weights =
+        _mm256_permutevar8x32_ps(centroid_lut_high, indices);
+    const __m256 weight = _mm256_blendv_ps(
+        low_weights,
+        high_weights,
+        _mm256_castsi256_ps(_mm256_cmpgt_epi32(indices, low_table_limit)));
+#pragma clang loop unroll(full)
+    for (int row = 0; row < row_tile; ++row) {
+      const std::int64_t input_offset =
+          (row_start + row) * args.in_features + k;
+      accumulators[row] = _mm256_fmadd_ps(
+          load8(args.x, input_offset),
+          weight,
+          accumulators[row]);
+    }
+  }
+
+  const float row_norm = args.row_norms[out_col];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    const std::int64_t input_row_offset =
+        (row_start + row) * args.in_features;
+    float accumulator = horizontal_sum(accumulators[row]);
+    for (std::int64_t tail = k; tail < args.in_features; ++tail) {
+      const std::uint8_t packed = packed_row[tail / 2];
+      const std::uint8_t index =
+          (tail & 1) == 0 ? packed & 15u : (packed >> 4) & 15u;
+      if constexpr (std::is_same_v<scalar_t, float>) {
+        accumulator +=
+            static_cast<float const *>(args.x)[input_row_offset + tail] *
+            args.centroids[index];
+      } else {
+        accumulator += static_cast<float>(
+                           static_cast<scalar_t const *>(
+                               args.x)[input_row_offset + tail]) *
+            args.centroids[index];
+      }
+    }
+    accumulator *= row_norm;
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_value<scalar_t>(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+template <int Bits>
+inline std::uint32_t unpack_index_generic(
+    std::uint8_t const *packed_row,
+    std::int64_t value_index) {
+  const std::int64_t bit_start = value_index * Bits;
+  const std::int64_t byte_index = bit_start >> 3;
+  const unsigned bit_offset = static_cast<unsigned>(bit_start & 7);
+  std::uint32_t raw = packed_row[byte_index];
+  if (bit_offset + static_cast<unsigned>(Bits) > 8) {
+    raw |= static_cast<std::uint32_t>(packed_row[byte_index + 1]) << 8;
+  }
+  return (raw >> bit_offset) & ((1u << Bits) - 1u);
+}
+
+// Each decoder turns 8 consecutive packed indices into 8 fp32 centroid
+// values. Rows are byte-aligned because dispatch requires in_features % 4 == 0.
+struct W2Avx2Decoder {
+  static constexpr int kBits = 2;
+  struct Tables {
+    __m256 lut;
+  };
+
+  ORBITQUANT_TARGET_AVX2 static inline Tables load_tables(
+      float const *centroids) {
+    const __m128 lut4 = _mm_loadu_ps(centroids);
+    return Tables{_mm256_set_m128(lut4, lut4)};
+  }
+
+  ORBITQUANT_TARGET_AVX2 static inline __m256 decode(
+      std::uint8_t const *packed_row,
+      std::int64_t k,
+      Tables const &tables) {
+    std::uint16_t packed_bits;
+    std::memcpy(&packed_bits, packed_row + (k >> 2), sizeof(packed_bits));
+    const __m128i bytes = _mm_cvtsi32_si128(packed_bits);
+    const __m128i replicated = _mm_shuffle_epi8(
+        bytes,
+        _mm_setr_epi8(0, 0, 0, 0, 1, 1, 1, 1, -1, -1, -1, -1, -1, -1, -1, -1));
+    const __m256i widened = _mm256_cvtepu8_epi32(replicated);
+    const __m256i shifts = _mm256_setr_epi32(0, 2, 4, 6, 0, 2, 4, 6);
+    const __m256i indices = _mm256_and_si256(
+        _mm256_srlv_epi32(widened, shifts),
+        _mm256_set1_epi32(3));
+    return _mm256_permutevar8x32_ps(tables.lut, indices);
+  }
+};
+
+struct W6Avx2Decoder {
+  static constexpr int kBits = 6;
+  struct Tables {
+    float const *centroids;
+  };
+
+  ORBITQUANT_TARGET_AVX2 static inline Tables load_tables(
+      float const *centroids) {
+    return Tables{centroids};
+  }
+
+  ORBITQUANT_TARGET_AVX2 static inline __m256 decode(
+      std::uint8_t const *packed_row,
+      std::int64_t k,
+      Tables const &tables) {
+    std::uint64_t raw_bits = 0;
+    std::memcpy(&raw_bits, packed_row + (k * 6 >> 3), 6);
+    const __m128i raw =
+        _mm_cvtsi64_si128(static_cast<long long>(raw_bits));
+    const __m128i windows = _mm_shuffle_epi8(
+        raw,
+        _mm_setr_epi8(0, 1, 0, 1, 1, 2, 2, 3, 3, 4, 3, 4, 4, 5, 5, 6));
+    const __m256i widened = _mm256_cvtepu16_epi32(windows);
+    const __m256i shifts = _mm256_setr_epi32(0, 6, 4, 2, 0, 6, 4, 2);
+    const __m256i indices = _mm256_and_si256(
+        _mm256_srlv_epi32(widened, shifts),
+        _mm256_set1_epi32(63));
+    return _mm256_i32gather_ps(tables.centroids, indices, 4);
+  }
+};
+
+template <
+    typename scalar_t,
+    __m256 (*load8)(void const *, std::int64_t),
+    typename decoder_t,
+    int row_tile>
+ORBITQUANT_TARGET_AVX2 inline void packed_matmul_avx2_lowbit_rows(
+    PackedMatmulArgs const &args,
+    std::uint8_t const *packed_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  __m256 accumulators[row_tile];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    accumulators[row] = _mm256_setzero_ps();
+  }
+  const typename decoder_t::Tables tables =
+      decoder_t::load_tables(args.centroids);
+
+  std::int64_t k = 0;
+  for (; k + 8 <= args.in_features; k += 8) {
+    const __m256 weight = decoder_t::decode(packed_row, k, tables);
+#pragma clang loop unroll(full)
+    for (int row = 0; row < row_tile; ++row) {
+      const std::int64_t input_offset =
+          (row_start + row) * args.in_features + k;
+      accumulators[row] = _mm256_fmadd_ps(
+          load8(args.x, input_offset),
+          weight,
+          accumulators[row]);
+    }
+  }
+
+  const float row_norm = args.row_norms[out_col];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    const std::int64_t input_row_offset =
+        (row_start + row) * args.in_features;
+    float accumulator = horizontal_sum(accumulators[row]);
+    for (std::int64_t tail = k; tail < args.in_features; ++tail) {
+      const std::uint32_t index =
+          unpack_index_generic<decoder_t::kBits>(packed_row, tail);
+      if constexpr (std::is_same_v<scalar_t, float>) {
+        accumulator +=
+            static_cast<float const *>(args.x)[input_row_offset + tail] *
+            args.centroids[index];
+      } else {
+        accumulator += static_cast<float>(
+                           static_cast<scalar_t const *>(
+                               args.x)[input_row_offset + tail]) *
+            args.centroids[index];
+      }
+    }
+    accumulator *= row_norm;
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_value<scalar_t>(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+template <typename scalar_t, __m256 (*load8)(void const *, std::int64_t), int row_tile>
+ORBITQUANT_TARGET_AVX2 inline void packed_matmul_avx2_buffered_rows(
+    PackedMatmulArgs const &args,
+    float const *decoded_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  __m256 accumulators[row_tile];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    accumulators[row] = _mm256_setzero_ps();
+  }
+
+  std::int64_t k = 0;
+  for (; k + 8 <= args.in_features; k += 8) {
+    const __m256 weight = _mm256_loadu_ps(decoded_row + k);
+#pragma clang loop unroll(full)
+    for (int row = 0; row < row_tile; ++row) {
+      const std::int64_t input_offset =
+          (row_start + row) * args.in_features + k;
+      accumulators[row] = _mm256_fmadd_ps(
+          load8(args.x, input_offset),
+          weight,
+          accumulators[row]);
+    }
+  }
+
+  const float row_norm = args.row_norms[out_col];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    const std::int64_t input_row_offset =
+        (row_start + row) * args.in_features;
+    float accumulator = horizontal_sum(accumulators[row]);
+    for (std::int64_t tail = k; tail < args.in_features; ++tail) {
+      if constexpr (std::is_same_v<scalar_t, float>) {
+        accumulator +=
+            static_cast<float const *>(args.x)[input_row_offset + tail] *
+            decoded_row[tail];
+      } else {
+        accumulator += static_cast<float>(
+                           static_cast<scalar_t const *>(
+                               args.x)[input_row_offset + tail]) *
+            decoded_row[tail];
+      }
+    }
+    accumulator *= row_norm;
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_value<scalar_t>(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+template <
+    typename scalar_t,
+    __m256 (*load8)(void const *, std::int64_t),
+    typename decoder_t>
+ORBITQUANT_TARGET_AVX2 ORBITQUANT_NOINLINE void
+packed_matmul_avx2_lowbit_typed(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  constexpr int kPrimaryRowTile = 8;
+  const std::int64_t packed_row_bytes =
+      args.in_features * decoder_t::kBits / 8;
+  // Two or more row tiles amortize the packed decode: expand the column once
+  // into a per-thread scratch row and stream plain FMA tiles from it.
+  const bool use_decoded_buffer = args.rows >= 16;
+  thread_local std::vector<float> decoded_row_storage;
+  if (use_decoded_buffer &&
+      decoded_row_storage.size() < static_cast<std::size_t>(args.in_features)) {
+    decoded_row_storage.resize(static_cast<std::size_t>(args.in_features));
+  }
+  for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) {
+    const auto *packed_row =
+        args.packed_weight_indices + out_col * packed_row_bytes;
+    if (use_decoded_buffer) {
+      float *decoded_row = decoded_row_storage.data();
+      const typename decoder_t::Tables tables =
+          decoder_t::load_tables(args.centroids);
+      std::int64_t k = 0;
+      for (; k + 8 <= args.in_features; k += 8) {
+        _mm256_storeu_ps(decoded_row + k, decoder_t::decode(packed_row, k, tables));
+      }
+      for (; k < args.in_features; ++k) {
+        decoded_row[k] =
+            args.centroids[unpack_index_generic<decoder_t::kBits>(packed_row, k)];
+      }
+      std::int64_t row = 0;
+      for (; row + kPrimaryRowTile <= args.rows; row += kPrimaryRowTile) {
+        packed_matmul_avx2_buffered_rows<scalar_t, load8, 8>(
+            args, decoded_row, out_col, row);
+      }
+      if (row + 4 <= args.rows) {
+        packed_matmul_avx2_buffered_rows<scalar_t, load8, 4>(
+            args, decoded_row, out_col, row);
+        row += 4;
+      }
+      switch (args.rows - row) {
+        case 3:
+          packed_matmul_avx2_buffered_rows<scalar_t, load8, 3>(
+              args, decoded_row, out_col, row);
+          break;
+        case 2:
+          packed_matmul_avx2_buffered_rows<scalar_t, load8, 2>(
+              args, decoded_row, out_col, row);
+          break;
+        case 1:
+          packed_matmul_avx2_buffered_rows<scalar_t, load8, 1>(
+              args, decoded_row, out_col, row);
+          break;
+        default:
+          break;
+      }
+      continue;
+    }
+    std::int64_t row = 0;
+    for (; row + kPrimaryRowTile <= args.rows; row += kPrimaryRowTile) {
+      packed_matmul_avx2_lowbit_rows<scalar_t, load8, decoder_t, 8>(
+          args, packed_row, out_col, row);
+    }
+    if (row + 4 <= args.rows) {
+      packed_matmul_avx2_lowbit_rows<scalar_t, load8, decoder_t, 4>(
+          args, packed_row, out_col, row);
+      row += 4;
+    }
+    switch (args.rows - row) {
+      case 3:
+        packed_matmul_avx2_lowbit_rows<scalar_t, load8, decoder_t, 3>(
+            args, packed_row, out_col, row);
+        break;
+      case 2:
+        packed_matmul_avx2_lowbit_rows<scalar_t, load8, decoder_t, 2>(
+            args, packed_row, out_col, row);
+        break;
+      case 1:
+        packed_matmul_avx2_lowbit_rows<scalar_t, load8, decoder_t, 1>(
+            args, packed_row, out_col, row);
+        break;
+      default:
+        break;
+    }
+  }
+}
+
+template <typename decoder_t>
+void packed_matmul_avx2_lowbit_dispatch(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  switch (args.scalar_kind) {
+    case ScalarKind::Float32:
+      packed_matmul_avx2_lowbit_typed<float, load_float8, decoder_t>(
+          args, out_start, out_end);
+      return;
+    case ScalarKind::Float16:
+      packed_matmul_avx2_lowbit_typed<c10::Half, load_half8, decoder_t>(
+          args, out_start, out_end);
+      return;
+    case ScalarKind::BFloat16:
+      packed_matmul_avx2_lowbit_typed<c10::BFloat16, load_bfloat8, decoder_t>(
+          args, out_start, out_end);
+      return;
+  }
+}
+
+template <typename scalar_t>
+bool use_verified_amd_cezanne_row_tile(PackedMatmulArgs const &args) {
+  if constexpr (!std::is_same_v<scalar_t, c10::BFloat16>) {
+    return false;
+  }
+  const bool tuned_dimension = args.in_features == 1536 ||
+      args.in_features == 1920 || args.in_features == 3072;
+  if (args.rows < 16 || !tuned_dimension) {
+    return false;
+  }
+  static const bool verified_cpu = [] {
+    unsigned int eax = 0;
+    unsigned int ebx = 0;
+    unsigned int ecx = 0;
+    unsigned int edx = 0;
+#if defined(_MSC_VER)
+    int registers[4]{};
+    __cpuid(registers, 0);
+    eax = static_cast<unsigned int>(registers[0]);
+    ebx = static_cast<unsigned int>(registers[1]);
+    ecx = static_cast<unsigned int>(registers[2]);
+    edx = static_cast<unsigned int>(registers[3]);
+    if (ebx != 0x68747541u || edx != 0x69746e65u ||
+        ecx != 0x444d4163u) {
+      return false;
+    }
+    __cpuid(registers, 1);
+    eax = static_cast<unsigned int>(registers[0]);
+#else
+    // CPUID vendor registers spell "AuthenticAMD" in EBX, EDX, ECX order.
+    if (!__get_cpuid(0, &eax, &ebx, &ecx, &edx) ||
+        ebx != 0x68747541u || edx != 0x69746e65u || ecx != 0x444d4163u ||
+        !__get_cpuid(1, &eax, &ebx, &ecx, &edx)) {
+      return false;
+    }
+#endif
+    const unsigned int base_family = (eax >> 8) & 0xfu;
+    const unsigned int family = base_family == 0xfu
+        ? base_family + ((eax >> 20) & 0xffu)
+        : base_family;
+    // The 16-row tile was measured on Zen 3 (Ryzen 5 5600G); apply it to the
+    // whole AVX2-only Zen 3 family (19h without AVX-512) instead of pinning
+    // the one benchmarked model.
+    return family == 0x19u && !packed_matmul_x86_avx512_available();
+  }();
+  return verified_cpu;
+}
+
+template <
+    typename scalar_t,
+    __m256 (*load8)(void const *, std::int64_t),
+    int primary_row_tile>
+ORBITQUANT_TARGET_AVX2 ORBITQUANT_NOINLINE void packed_matmul_avx2_w4_typed(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  static_assert(primary_row_tile == 8 || primary_row_tile == 16);
+  const std::int64_t packed_row_bytes = args.in_features / 2;
+  for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) {
+    const auto *packed_row =
+        args.packed_weight_indices + out_col * packed_row_bytes;
+    std::int64_t row = 0;
+    for (; row + primary_row_tile <= args.rows; row += primary_row_tile) {
+      packed_matmul_avx2_w4_rows<scalar_t, load8, primary_row_tile>(
+          args, packed_row, out_col, row);
+    }
+    if (row + 8 <= args.rows) {
+      packed_matmul_avx2_w4_rows<scalar_t, load8, 8>(
+          args, packed_row, out_col, row);
+      row += 8;
+    }
+    if (row + 4 <= args.rows) {
+      packed_matmul_avx2_w4_rows<scalar_t, load8, 4>(
+          args, packed_row, out_col, row);
+      row += 4;
+    }
+    switch (args.rows - row) {
+      case 3:
+        packed_matmul_avx2_w4_rows<scalar_t, load8, 3>(
+            args, packed_row, out_col, row);
+        break;
+      case 2:
+        packed_matmul_avx2_w4_rows<scalar_t, load8, 2>(
+            args, packed_row, out_col, row);
+        break;
+      case 1:
+        packed_matmul_avx2_w4_rows<scalar_t, load8, 1>(
+            args, packed_row, out_col, row);
+        break;
+      default:
+        break;
+    }
+  }
+}
+
+}  // namespace
+
+void packed_matmul_x86_avx2_range(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  if (!packed_matmul_x86_avx2_available()) {
+    packed_matmul_scalar_range(args, out_start, out_end);
+    return;
+  }
+  if (args.bits == 2 && args.in_features % 4 == 0) {
+    packed_matmul_avx2_lowbit_dispatch<W2Avx2Decoder>(args, out_start, out_end);
+    return;
+  }
+  if (args.bits == 6 && args.in_features % 4 == 0) {
+    packed_matmul_avx2_lowbit_dispatch<W6Avx2Decoder>(args, out_start, out_end);
+    return;
+  }
+  if (args.bits != 4 || args.in_features % 2 != 0) {
+    packed_matmul_scalar_range(args, out_start, out_end);
+    return;
+  }
+  switch (args.scalar_kind) {
+    case ScalarKind::Float32:
+      packed_matmul_avx2_w4_typed<float, load_float8, 8>(
+          args, out_start, out_end);
+      return;
+    case ScalarKind::Float16:
+      packed_matmul_avx2_w4_typed<c10::Half, load_half8, 8>(
+          args, out_start, out_end);
+      return;
+    case ScalarKind::BFloat16:
+      if (use_verified_amd_cezanne_row_tile<c10::BFloat16>(args)) {
+        packed_matmul_avx2_w4_typed<c10::BFloat16, load_bfloat8, 16>(
+            args, out_start, out_end);
+      } else {
+        packed_matmul_avx2_w4_typed<c10::BFloat16, load_bfloat8, 8>(
+            args, out_start, out_end);
+      }
+      return;
+  }
+}
+
+}  // namespace orbitquant::cpu
+
+#else
+
+namespace orbitquant::cpu {
+
+void packed_matmul_x86_avx2_range(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  packed_matmul_scalar_range(args, out_start, out_end);
+}
+
+}  // namespace orbitquant::cpu
+
+#endif

orbitquant_packed_matmul_cpu/packed_matmul_x86_avx512.cpp

@@ -0,0 +1,851 @@
+#include "packed_matmul_cpu.h"
+
+#if (defined(__x86_64__) || defined(_M_X64)) && !defined(_MSC_VER)
+#include <cpuid.h>
+#include <immintrin.h>
+
+#include <torch/headeronly/util/BFloat16.h>
+#include <torch/headeronly/util/Half.h>
+
+#include <cstring>
+#include <cstdint>
+#include <type_traits>
+#include <vector>
+
+#define ORBITQUANT_TARGET_AVX512 \
+  __attribute__((target("avx512f,avx512dq,avx512bw,avx512vl,fma,f16c")))
+#define ORBITQUANT_TARGET_AVX512_BF16 \
+  __attribute__((target( \
+      "avx512f,avx512dq,avx512bw,avx512vl,avx512bf16,fma,f16c")))
+#define ORBITQUANT_HAS_AVX512_BF16_INTRINSICS 1
+#define ORBITQUANT_NOINLINE __attribute__((noinline))
+#define ORBITQUANT_ALWAYS_INLINE __attribute__((always_inline))
+
+namespace orbitquant::cpu {
+namespace {
+
+ORBITQUANT_TARGET_AVX512 inline float horizontal_sum(__m512 value) {
+  return _mm512_reduce_add_ps(value);
+}
+
+ORBITQUANT_TARGET_AVX512 inline __m512 load_float16(
+    void const *data,
+    std::int64_t offset) {
+  return _mm512_loadu_ps(static_cast<float const *>(data) + offset);
+}
+
+ORBITQUANT_TARGET_AVX512 inline __m512 load_half16(
+    void const *data,
+    std::int64_t offset) {
+  const auto *source = static_cast<std::uint16_t const *>(data) + offset;
+  const __m256i packed =
+      _mm256_loadu_si256(reinterpret_cast<__m256i const *>(source));
+  return _mm512_cvtph_ps(packed);
+}
+
+ORBITQUANT_TARGET_AVX512 inline __m512 load_bfloat16(
+    void const *data,
+    std::int64_t offset) {
+  const auto *source = static_cast<std::uint16_t const *>(data) + offset;
+  const __m256i packed =
+      _mm256_loadu_si256(reinterpret_cast<__m256i const *>(source));
+  const __m512i widened = _mm512_cvtepu16_epi32(packed);
+  return _mm512_castsi512_ps(_mm512_slli_epi32(widened, 16));
+}
+
+template <typename scalar_t>
+inline void store_value(void *data, std::int64_t offset, float value) {
+  static_cast<scalar_t *>(data)[offset] = scalar_t(value);
+}
+
+template <>
+inline void store_value<float>(void *data, std::int64_t offset, float value) {
+  static_cast<float *>(data)[offset] = value;
+}
+
+template <
+    typename scalar_t,
+    __m512 (*load16)(void const *, std::int64_t),
+    int row_tile>
+ORBITQUANT_TARGET_AVX512 inline void packed_matmul_avx512_w4_rows(
+    PackedMatmulArgs const &args,
+    std::uint8_t const *packed_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  __m512 accumulators[row_tile];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    accumulators[row] = _mm512_setzero_ps();
+  }
+  const __m512 centroid_lut = _mm512_loadu_ps(args.centroids);
+  const __m128i nibble_mask = _mm_set1_epi8(15);
+
+  std::int64_t k = 0;
+  for (; k + 16 <= args.in_features; k += 16) {
+    const std::int64_t byte_offset = k / 2;
+    std::int64_t packed;
+    std::memcpy(&packed, packed_row + byte_offset, sizeof(packed));
+    const __m128i bytes = _mm_cvtsi64_si128(packed);
+    const __m128i low = _mm_and_si128(bytes, nibble_mask);
+    const __m128i high = _mm_and_si128(
+        _mm_srli_epi16(bytes, 4),
+        nibble_mask);
+    const __m512i indices =
+        _mm512_cvtepu8_epi32(_mm_unpacklo_epi8(low, high));
+    const __m512 weight = _mm512_permutexvar_ps(indices, centroid_lut);
+#pragma clang loop unroll(full)
+    for (int row = 0; row < row_tile; ++row) {
+      const std::int64_t input_offset =
+          (row_start + row) * args.in_features + k;
+      accumulators[row] = _mm512_fmadd_ps(
+          load16(args.x, input_offset),
+          weight,
+          accumulators[row]);
+    }
+  }
+
+  const float row_norm = args.row_norms[out_col];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    const std::int64_t input_row_offset =
+        (row_start + row) * args.in_features;
+    float accumulator = horizontal_sum(accumulators[row]);
+    for (std::int64_t tail = k; tail < args.in_features; ++tail) {
+      const std::uint8_t packed = packed_row[tail / 2];
+      const std::uint8_t index =
+          (tail & 1) == 0 ? packed & 15u : (packed >> 4) & 15u;
+      if constexpr (std::is_same_v<scalar_t, float>) {
+        accumulator +=
+            static_cast<float const *>(args.x)[input_row_offset + tail] *
+            args.centroids[index];
+      } else {
+        accumulator += static_cast<float>(
+                           static_cast<scalar_t const *>(
+                               args.x)[input_row_offset + tail]) *
+            args.centroids[index];
+      }
+    }
+    accumulator *= row_norm;
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_value<scalar_t>(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+template <typename scalar_t, __m512 (*load16)(void const *, std::int64_t)>
+ORBITQUANT_TARGET_AVX512 ORBITQUANT_NOINLINE void packed_matmul_avx512_w4_typed(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  constexpr int kPrimaryRowTile = 8;
+  const std::int64_t packed_row_bytes = args.in_features / 2;
+  for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) {
+    const auto *packed_row =
+        args.packed_weight_indices + out_col * packed_row_bytes;
+    std::int64_t row = 0;
+    for (; row + kPrimaryRowTile <= args.rows; row += kPrimaryRowTile) {
+      packed_matmul_avx512_w4_rows<scalar_t, load16, kPrimaryRowTile>(
+          args, packed_row, out_col, row);
+    }
+    if (row + 8 <= args.rows) {
+      packed_matmul_avx512_w4_rows<scalar_t, load16, 8>(
+          args, packed_row, out_col, row);
+      row += 8;
+    }
+    if (row + 4 <= args.rows) {
+      packed_matmul_avx512_w4_rows<scalar_t, load16, 4>(
+          args, packed_row, out_col, row);
+      row += 4;
+    }
+    switch (args.rows - row) {
+      case 3:
+        packed_matmul_avx512_w4_rows<scalar_t, load16, 3>(
+            args, packed_row, out_col, row);
+        break;
+      case 2:
+        packed_matmul_avx512_w4_rows<scalar_t, load16, 2>(
+            args, packed_row, out_col, row);
+        break;
+      case 1:
+        packed_matmul_avx512_w4_rows<scalar_t, load16, 1>(
+            args, packed_row, out_col, row);
+        break;
+      default:
+        break;
+    }
+  }
+}
+
+template <int Bits>
+inline std::uint32_t unpack_index_generic(
+    std::uint8_t const *packed_row,
+    std::int64_t value_index) {
+  const std::int64_t bit_start = value_index * Bits;
+  const std::int64_t byte_index = bit_start >> 3;
+  const unsigned bit_offset = static_cast<unsigned>(bit_start & 7);
+  std::uint32_t raw = packed_row[byte_index];
+  if (bit_offset + static_cast<unsigned>(Bits) > 8) {
+    raw |= static_cast<std::uint32_t>(packed_row[byte_index + 1]) << 8;
+  }
+  return (raw >> bit_offset) & ((1u << Bits) - 1u);
+}
+
+// Each decoder turns 16 consecutive packed indices into 16 fp32 centroid
+// values. Rows are byte-aligned because dispatch requires in_features
+// divisibility (W2/W6: in % 4 == 0, W3: in % 8 == 0).
+struct W2Avx512Decoder {
+  static constexpr int kBits = 2;
+  struct Tables {
+    __m512 lut;
+  };
+
+  ORBITQUANT_TARGET_AVX512 static inline Tables load_tables(
+      float const *centroids) {
+    return Tables{_mm512_broadcast_f32x4(_mm_loadu_ps(centroids))};
+  }
+
+  ORBITQUANT_TARGET_AVX512 static inline __m512 decode(
+      std::uint8_t const *packed_row,
+      std::int64_t k,
+      Tables const &tables) {
+    std::uint32_t packed_bits;
+    std::memcpy(&packed_bits, packed_row + (k >> 2), sizeof(packed_bits));
+    const __m128i bytes = _mm_cvtsi32_si128(static_cast<int>(packed_bits));
+    const __m128i replicated = _mm_shuffle_epi8(
+        bytes,
+        _mm_setr_epi8(0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3));
+    const __m512i widened = _mm512_cvtepu8_epi32(replicated);
+    const __m512i shifts = _mm512_set_epi32(
+        6, 4, 2, 0, 6, 4, 2, 0, 6, 4, 2, 0, 6, 4, 2, 0);
+    const __m512i indices = _mm512_and_si512(
+        _mm512_srlv_epi32(widened, shifts),
+        _mm512_set1_epi32(3));
+    return _mm512_permutexvar_ps(indices, tables.lut);
+  }
+};
+
+struct W3Avx512Decoder {
+  static constexpr int kBits = 3;
+  struct Tables {
+    __m512 lut;
+  };
+
+  ORBITQUANT_TARGET_AVX512 static inline Tables load_tables(
+      float const *centroids) {
+    return Tables{_mm512_broadcast_f32x8(_mm256_loadu_ps(centroids))};
+  }
+
+  ORBITQUANT_TARGET_AVX512 static inline __m512 decode(
+      std::uint8_t const *packed_row,
+      std::int64_t k,
+      Tables const &tables) {
+    std::uint64_t raw_bits = 0;
+    std::memcpy(&raw_bits, packed_row + (k * 3 >> 3), 6);
+    const __m128i raw = _mm_cvtsi64_si128(static_cast<long long>(raw_bits));
+    const __m128i window_low = _mm_shuffle_epi8(
+        raw,
+        _mm_setr_epi8(0, 1, 0, 1, 0, 1, 1, 2, 1, 2, 1, 2, 2, 3, 2, 3));
+    const __m128i window_high = _mm_shuffle_epi8(
+        raw,
+        _mm_setr_epi8(3, 4, 3, 4, 3, 4, 4, 5, 4, 5, 4, 5, 5, 6, 5, 6));
+    const __m512i widened = _mm512_cvtepu16_epi32(
+        _mm256_set_m128i(window_high, window_low));
+    const __m512i shifts = _mm512_set_epi32(
+        5, 2, 7, 4, 1, 6, 3, 0, 5, 2, 7, 4, 1, 6, 3, 0);
+    const __m512i indices = _mm512_and_si512(
+        _mm512_srlv_epi32(widened, shifts),
+        _mm512_set1_epi32(7));
+    return _mm512_permutexvar_ps(indices, tables.lut);
+  }
+};
+
+struct W6Avx512Decoder {
+  static constexpr int kBits = 6;
+  struct Tables {
+    __m512 lut0;
+    __m512 lut1;
+    __m512 lut2;
+    __m512 lut3;
+  };
+
+  ORBITQUANT_TARGET_AVX512 static inline Tables load_tables(
+      float const *centroids) {
+    return Tables{
+        _mm512_loadu_ps(centroids),
+        _mm512_loadu_ps(centroids + 16),
+        _mm512_loadu_ps(centroids + 32),
+        _mm512_loadu_ps(centroids + 48)};
+  }
+
+  ORBITQUANT_TARGET_AVX512 static inline __m512 decode(
+      std::uint8_t const *packed_row,
+      std::int64_t k,
+      Tables const &tables) {
+    std::uint64_t low_bytes;
+    std::uint32_t high_bytes;
+    std::uint8_t const *source = packed_row + (k * 6 >> 3);
+    std::memcpy(&low_bytes, source, sizeof(low_bytes));
+    std::memcpy(&high_bytes, source + 8, sizeof(high_bytes));
+    const __m128i raw = _mm_set_epi64x(
+        static_cast<long long>(high_bytes),
+        static_cast<long long>(low_bytes));
+    const __m128i window_low = _mm_shuffle_epi8(
+        raw,
+        _mm_setr_epi8(0, 1, 0, 1, 1, 2, 2, 3, 3, 4, 3, 4, 4, 5, 5, 6));
+    const __m128i window_high = _mm_shuffle_epi8(
+        raw,
+        _mm_setr_epi8(6, 7, 6, 7, 7, 8, 8, 9, 9, 10, 9, 10, 10, 11, 11, 12));
+    const __m512i widened = _mm512_cvtepu16_epi32(
+        _mm256_set_m128i(window_high, window_low));
+    const __m512i shifts = _mm512_set_epi32(
+        2, 4, 6, 0, 2, 4, 6, 0, 2, 4, 6, 0, 2, 4, 6, 0);
+    const __m512i indices = _mm512_and_si512(
+        _mm512_srlv_epi32(widened, shifts),
+        _mm512_set1_epi32(63));
+    const __m512 low_pair = _mm512_permutex2var_ps(tables.lut0, indices, tables.lut1);
+    const __m512 high_pair = _mm512_permutex2var_ps(tables.lut2, indices, tables.lut3);
+    const __mmask16 use_high =
+        _mm512_test_epi32_mask(indices, _mm512_set1_epi32(32));
+    return _mm512_mask_blend_ps(use_high, low_pair, high_pair);
+  }
+};
+
+template <
+    typename scalar_t,
+    __m512 (*load16)(void const *, std::int64_t),
+    typename decoder_t,
+    int row_tile>
+ORBITQUANT_TARGET_AVX512 inline void packed_matmul_avx512_lowbit_rows(
+    PackedMatmulArgs const &args,
+    std::uint8_t const *packed_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  __m512 accumulators[row_tile];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    accumulators[row] = _mm512_setzero_ps();
+  }
+  const typename decoder_t::Tables tables =
+      decoder_t::load_tables(args.centroids);
+
+  std::int64_t k = 0;
+  for (; k + 16 <= args.in_features; k += 16) {
+    const __m512 weight = decoder_t::decode(packed_row, k, tables);
+#pragma clang loop unroll(full)
+    for (int row = 0; row < row_tile; ++row) {
+      const std::int64_t input_offset =
+          (row_start + row) * args.in_features + k;
+      accumulators[row] = _mm512_fmadd_ps(
+          load16(args.x, input_offset),
+          weight,
+          accumulators[row]);
+    }
+  }
+
+  const float row_norm = args.row_norms[out_col];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    const std::int64_t input_row_offset =
+        (row_start + row) * args.in_features;
+    float accumulator = horizontal_sum(accumulators[row]);
+    for (std::int64_t tail = k; tail < args.in_features; ++tail) {
+      const std::uint32_t index =
+          unpack_index_generic<decoder_t::kBits>(packed_row, tail);
+      if constexpr (std::is_same_v<scalar_t, float>) {
+        accumulator +=
+            static_cast<float const *>(args.x)[input_row_offset + tail] *
+            args.centroids[index];
+      } else {
+        accumulator += static_cast<float>(
+                           static_cast<scalar_t const *>(
+                               args.x)[input_row_offset + tail]) *
+            args.centroids[index];
+      }
+    }
+    accumulator *= row_norm;
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_value<scalar_t>(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+template <typename scalar_t, __m512 (*load16)(void const *, std::int64_t), int row_tile>
+ORBITQUANT_TARGET_AVX512 inline void packed_matmul_avx512_buffered_rows(
+    PackedMatmulArgs const &args,
+    float const *decoded_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  __m512 accumulators[row_tile];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    accumulators[row] = _mm512_setzero_ps();
+  }
+
+  std::int64_t k = 0;
+  for (; k + 16 <= args.in_features; k += 16) {
+    const __m512 weight = _mm512_loadu_ps(decoded_row + k);
+#pragma clang loop unroll(full)
+    for (int row = 0; row < row_tile; ++row) {
+      const std::int64_t input_offset =
+          (row_start + row) * args.in_features + k;
+      accumulators[row] = _mm512_fmadd_ps(
+          load16(args.x, input_offset),
+          weight,
+          accumulators[row]);
+    }
+  }
+
+  const float row_norm = args.row_norms[out_col];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    const std::int64_t input_row_offset =
+        (row_start + row) * args.in_features;
+    float accumulator = horizontal_sum(accumulators[row]);
+    for (std::int64_t tail = k; tail < args.in_features; ++tail) {
+      if constexpr (std::is_same_v<scalar_t, float>) {
+        accumulator +=
+            static_cast<float const *>(args.x)[input_row_offset + tail] *
+            decoded_row[tail];
+      } else {
+        accumulator += static_cast<float>(
+                           static_cast<scalar_t const *>(
+                               args.x)[input_row_offset + tail]) *
+            decoded_row[tail];
+      }
+    }
+    accumulator *= row_norm;
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_value<scalar_t>(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+template <
+    typename scalar_t,
+    __m512 (*load16)(void const *, std::int64_t),
+    typename decoder_t>
+ORBITQUANT_TARGET_AVX512 ORBITQUANT_NOINLINE void
+packed_matmul_avx512_lowbit_typed(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  constexpr int kPrimaryRowTile = 8;
+  const std::int64_t packed_row_bytes =
+      args.in_features * decoder_t::kBits / 8;
+  // Two or more row tiles amortize the packed decode: expand the column once
+  // into a per-thread scratch row and stream plain FMA tiles from it.
+  const bool use_decoded_buffer = args.rows >= 16;
+  thread_local std::vector<float> decoded_row_storage;
+  if (use_decoded_buffer &&
+      decoded_row_storage.size() < static_cast<std::size_t>(args.in_features)) {
+    decoded_row_storage.resize(static_cast<std::size_t>(args.in_features));
+  }
+  for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) {
+    const auto *packed_row =
+        args.packed_weight_indices + out_col * packed_row_bytes;
+    if (use_decoded_buffer) {
+      float *decoded_row = decoded_row_storage.data();
+      const typename decoder_t::Tables tables =
+          decoder_t::load_tables(args.centroids);
+      std::int64_t k = 0;
+      for (; k + 16 <= args.in_features; k += 16) {
+        _mm512_storeu_ps(decoded_row + k, decoder_t::decode(packed_row, k, tables));
+      }
+      for (; k < args.in_features; ++k) {
+        decoded_row[k] =
+            args.centroids[unpack_index_generic<decoder_t::kBits>(packed_row, k)];
+      }
+      std::int64_t row = 0;
+      for (; row + kPrimaryRowTile <= args.rows; row += kPrimaryRowTile) {
+        packed_matmul_avx512_buffered_rows<scalar_t, load16, 8>(
+            args, decoded_row, out_col, row);
+      }
+      if (row + 4 <= args.rows) {
+        packed_matmul_avx512_buffered_rows<scalar_t, load16, 4>(
+            args, decoded_row, out_col, row);
+        row += 4;
+      }
+      switch (args.rows - row) {
+        case 3:
+          packed_matmul_avx512_buffered_rows<scalar_t, load16, 3>(
+              args, decoded_row, out_col, row);
+          break;
+        case 2:
+          packed_matmul_avx512_buffered_rows<scalar_t, load16, 2>(
+              args, decoded_row, out_col, row);
+          break;
+        case 1:
+          packed_matmul_avx512_buffered_rows<scalar_t, load16, 1>(
+              args, decoded_row, out_col, row);
+          break;
+        default:
+          break;
+      }
+      continue;
+    }
+    std::int64_t row = 0;
+    for (; row + kPrimaryRowTile <= args.rows; row += kPrimaryRowTile) {
+      packed_matmul_avx512_lowbit_rows<scalar_t, load16, decoder_t, 8>(
+          args, packed_row, out_col, row);
+    }
+    if (row + 4 <= args.rows) {
+      packed_matmul_avx512_lowbit_rows<scalar_t, load16, decoder_t, 4>(
+          args, packed_row, out_col, row);
+      row += 4;
+    }
+    switch (args.rows - row) {
+      case 3:
+        packed_matmul_avx512_lowbit_rows<scalar_t, load16, decoder_t, 3>(
+            args, packed_row, out_col, row);
+        break;
+      case 2:
+        packed_matmul_avx512_lowbit_rows<scalar_t, load16, decoder_t, 2>(
+            args, packed_row, out_col, row);
+        break;
+      case 1:
+        packed_matmul_avx512_lowbit_rows<scalar_t, load16, decoder_t, 1>(
+            args, packed_row, out_col, row);
+        break;
+      default:
+        break;
+    }
+  }
+}
+
+template <typename decoder_t>
+ORBITQUANT_TARGET_AVX512 void packed_matmul_avx512_lowbit_dispatch(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  switch (args.scalar_kind) {
+    case ScalarKind::Float32:
+      packed_matmul_avx512_lowbit_typed<float, load_float16, decoder_t>(
+          args, out_start, out_end);
+      return;
+    case ScalarKind::Float16:
+      packed_matmul_avx512_lowbit_typed<c10::Half, load_half16, decoder_t>(
+          args, out_start, out_end);
+      return;
+    case ScalarKind::BFloat16:
+      packed_matmul_avx512_lowbit_typed<c10::BFloat16, load_bfloat16, decoder_t>(
+          args, out_start, out_end);
+      return;
+  }
+}
+
+#if defined(ORBITQUANT_HAS_AVX512_BF16_INTRINSICS)
+template <int row_tile>
+ORBITQUANT_TARGET_AVX512_BF16 ORBITQUANT_ALWAYS_INLINE inline void
+accumulate_bf16_w4_chunk(
+    PackedMatmulArgs const &args,
+    std::uint8_t const *packed_row,
+    std::int64_t row_start,
+    std::int64_t k,
+    __m512i centroid_lut_words,
+    __m128i nibble_mask,
+    __m512 (&accumulators)[row_tile]) {
+  const std::int64_t byte_offset = k / 2;
+  const __m128i bytes = _mm_loadu_si128(
+      reinterpret_cast<__m128i const *>(packed_row + byte_offset));
+  const __m128i low = _mm_and_si128(bytes, nibble_mask);
+  const __m128i high = _mm_and_si128(
+      _mm_srli_epi16(bytes, 4),
+      nibble_mask);
+  const __m256i packed_indices = _mm256_set_m128i(
+      _mm_unpackhi_epi8(low, high),
+      _mm_unpacklo_epi8(low, high));
+  const __m512i indices = _mm512_cvtepu8_epi16(packed_indices);
+  const __m512bh weights = (__m512bh)_mm512_permutexvar_epi16(
+      indices,
+      centroid_lut_words);
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    const std::int64_t input_offset =
+        (row_start + row) * args.in_features + k;
+    const __m512bh activations = (__m512bh)_mm512_loadu_si512(
+        static_cast<std::uint16_t const *>(args.x) + input_offset);
+    accumulators[row] =
+        _mm512_dpbf16_ps(accumulators[row], activations, weights);
+  }
+}
+
+template <int row_tile, bool aligned_k, bool unroll_k>
+ORBITQUANT_TARGET_AVX512_BF16 inline void packed_matmul_avx512_bf16_w4_rows(
+    PackedMatmulArgs const &args,
+    std::uint8_t const *packed_row,
+    std::int64_t out_col,
+    std::int64_t row_start) {
+  __m512 accumulators[row_tile];
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    accumulators[row] = _mm512_setzero_ps();
+  }
+  const float row_norm = args.row_norms[out_col];
+  const __m512 centroid_lut = _mm512_mul_ps(
+      _mm512_loadu_ps(args.centroids),
+      _mm512_set1_ps(row_norm));
+  const __m512i centroid_lut_words = (__m512i)_mm512_cvtne2ps_pbh(
+      centroid_lut,
+      centroid_lut);
+  const __m128i nibble_mask = _mm_set1_epi8(15);
+
+  std::int64_t k = 0;
+  if constexpr (unroll_k) {
+#if defined(__clang__)
+#pragma clang loop unroll_count(2)
+#elif defined(__GNUC__)
+#pragma GCC unroll 2
+#endif
+    for (; k + 32 <= args.in_features; k += 32) {
+      accumulate_bf16_w4_chunk<row_tile>(
+          args,
+          packed_row,
+          row_start,
+          k,
+          centroid_lut_words,
+          nibble_mask,
+          accumulators);
+    }
+  } else {
+    for (; k + 32 <= args.in_features; k += 32) {
+      accumulate_bf16_w4_chunk<row_tile>(
+          args,
+          packed_row,
+          row_start,
+          k,
+          centroid_lut_words,
+          nibble_mask,
+          accumulators);
+    }
+  }
+
+#pragma clang loop unroll(full)
+  for (int row = 0; row < row_tile; ++row) {
+    const std::int64_t input_row_offset =
+        (row_start + row) * args.in_features;
+    float accumulator = horizontal_sum(accumulators[row]);
+    if constexpr (!aligned_k) {
+      for (std::int64_t tail = k; tail < args.in_features; ++tail) {
+        const std::uint8_t packed = packed_row[tail / 2];
+        const std::uint8_t index =
+            (tail & 1) == 0 ? packed & 15u : (packed >> 4) & 15u;
+        const float activation = static_cast<float>(
+            static_cast<c10::BFloat16 const *>(args.x)[input_row_offset + tail]);
+        const float weight = static_cast<float>(
+            c10::BFloat16(args.centroids[index] * row_norm));
+        accumulator += activation * weight;
+      }
+    }
+    if (args.has_bias) {
+      accumulator += args.bias[out_col];
+    }
+    store_value<c10::BFloat16>(
+        args.out,
+        (row_start + row) * args.out_features + out_col,
+        accumulator);
+  }
+}
+
+template <int primary_row_tile, bool aligned_k, bool unroll_k>
+ORBITQUANT_TARGET_AVX512_BF16 ORBITQUANT_NOINLINE void
+packed_matmul_avx512_bf16_w4_typed(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  static_assert(primary_row_tile == 4 || primary_row_tile == 8);
+  const std::int64_t packed_row_bytes = args.in_features / 2;
+  for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) {
+    const auto *packed_row =
+        args.packed_weight_indices + out_col * packed_row_bytes;
+    std::int64_t row = 0;
+    for (; row + primary_row_tile <= args.rows; row += primary_row_tile) {
+      packed_matmul_avx512_bf16_w4_rows<primary_row_tile, aligned_k, unroll_k>(
+          args, packed_row, out_col, row);
+    }
+    if (row + 4 <= args.rows) {
+      packed_matmul_avx512_bf16_w4_rows<4, aligned_k, unroll_k>(
+          args, packed_row, out_col, row);
+      row += 4;
+    }
+    switch (args.rows - row) {
+      case 3:
+        packed_matmul_avx512_bf16_w4_rows<3, aligned_k, unroll_k>(
+            args, packed_row, out_col, row);
+        break;
+      case 2:
+        packed_matmul_avx512_bf16_w4_rows<2, aligned_k, unroll_k>(
+            args, packed_row, out_col, row);
+        break;
+      case 1:
+        packed_matmul_avx512_bf16_w4_rows<1, aligned_k, unroll_k>(
+            args, packed_row, out_col, row);
+        break;
+      default:
+        break;
+    }
+  }
+}
+#endif
+
+bool runtime_has_avx512() {
+#if defined(_MSC_VER)
+  int registers[4]{};
+  __cpuid(registers, 1);
+  const bool osxsave = (registers[2] & (1 << 27)) != 0;
+  const bool avx = (registers[2] & (1 << 28)) != 0;
+  const bool fma = (registers[2] & (1 << 12)) != 0;
+  const bool f16c = (registers[2] & (1 << 29)) != 0;
+  if (!osxsave || !avx || !fma || !f16c || (_xgetbv(0) & 0xE6) != 0xE6) {
+    return false;
+  }
+  __cpuidex(registers, 7, 0);
+  constexpr unsigned required_ebx =
+      (1u << 16) | (1u << 17) | (1u << 30) | (1u << 31);
+  return (static_cast<unsigned>(registers[1]) & required_ebx) == required_ebx;
+#else
+  __builtin_cpu_init();
+  return __builtin_cpu_supports("avx512f") &&
+      __builtin_cpu_supports("avx512dq") &&
+      __builtin_cpu_supports("avx512bw") &&
+      __builtin_cpu_supports("avx512vl") &&
+      __builtin_cpu_supports("fma") && __builtin_cpu_supports("f16c");
+#endif
+}
+
+bool runtime_has_avx512_bf16() {
+#if defined(ORBITQUANT_HAS_AVX512_BF16_INTRINSICS)
+  __builtin_cpu_init();
+  return runtime_has_avx512() && __builtin_cpu_supports("avx512bf16");
+#else
+  return false;
+#endif
+}
+
+bool runtime_has_verified_amd_bf16_tuning() {
+  static const bool available = [] {
+    unsigned int eax = 0;
+    unsigned int ebx = 0;
+    unsigned int ecx = 0;
+    unsigned int edx = 0;
+    // CPUID vendor registers spell "AuthenticAMD" in EBX, EDX, ECX order.
+    if (!__get_cpuid(0, &eax, &ebx, &ecx, &edx) ||
+        ebx != 0x68747541u || edx != 0x69746e65u || ecx != 0x444d4163u ||
+        !__get_cpuid(1, &eax, &ebx, &ecx, &edx)) {
+      return false;
+    }
+    const unsigned int base_family = (eax >> 8) & 0xfu;
+    const unsigned int family = base_family == 0xfu
+        ? base_family + ((eax >> 20) & 0xffu)
+        : base_family;
+    // The four-row unrolled tile was measured on Zen 4 (EPYC 4564P); apply it
+    // to the AMD families that ship AVX-512 BF16 (Zen 4 = 19h, Zen 5 = 1Ah)
+    // instead of pinning the one benchmarked model.
+    return (family == 0x19u || family == 0x1au) && runtime_has_avx512_bf16();
+  }();
+  return available;
+}
+
+}  // namespace
+
+bool packed_matmul_x86_avx512_available() {
+  static const bool available = runtime_has_avx512();
+  return available;
+}
+
+void packed_matmul_x86_avx512_range(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  if (!packed_matmul_x86_avx512_available()) {
+    packed_matmul_scalar_range(args, out_start, out_end);
+    return;
+  }
+  if (args.bits == 2 && args.in_features % 4 == 0) {
+    packed_matmul_avx512_lowbit_dispatch<W2Avx512Decoder>(
+        args, out_start, out_end);
+    return;
+  }
+  if (args.bits == 3 && args.in_features % 8 == 0) {
+    packed_matmul_avx512_lowbit_dispatch<W3Avx512Decoder>(
+        args, out_start, out_end);
+    return;
+  }
+  if (args.bits == 6 && args.in_features % 4 == 0) {
+    packed_matmul_avx512_lowbit_dispatch<W6Avx512Decoder>(
+        args, out_start, out_end);
+    return;
+  }
+  if (args.bits != 4 || args.in_features % 2 != 0) {
+    packed_matmul_scalar_range(args, out_start, out_end);
+    return;
+  }
+  switch (args.scalar_kind) {
+    case ScalarKind::Float32:
+      packed_matmul_avx512_w4_typed<float, load_float16>(
+          args, out_start, out_end);
+      return;
+    case ScalarKind::Float16:
+      packed_matmul_avx512_w4_typed<c10::Half, load_half16>(
+          args, out_start, out_end);
+      return;
+    case ScalarKind::BFloat16:
+#if defined(ORBITQUANT_HAS_AVX512_BF16_INTRINSICS)
+      if (runtime_has_avx512_bf16()) {
+        const bool tuned_dimension = args.in_features == 1536 ||
+            args.in_features == 1920 || args.in_features == 3072;
+        const bool use_tuned_shape =
+            runtime_has_verified_amd_bf16_tuning() && args.rows >= 16 &&
+            tuned_dimension;
+        if (args.in_features % 32 == 0) {
+          if (use_tuned_shape) {
+            packed_matmul_avx512_bf16_w4_typed<4, true, true>(
+                args, out_start, out_end);
+          } else {
+            packed_matmul_avx512_bf16_w4_typed<8, true, false>(
+                args, out_start, out_end);
+          }
+        } else {
+          packed_matmul_avx512_bf16_w4_typed<8, false, false>(
+              args, out_start, out_end);
+        }
+        return;
+      }
+#endif
+      packed_matmul_avx512_w4_typed<c10::BFloat16, load_bfloat16>(
+          args, out_start, out_end);
+      return;
+  }
+}
+
+}  // namespace orbitquant::cpu
+
+#else
+
+namespace orbitquant::cpu {
+
+bool packed_matmul_x86_avx512_available() {
+  return false;
+}
+
+void packed_matmul_x86_avx512_range(
+    PackedMatmulArgs const &args,
+    std::int64_t out_start,
+    std::int64_t out_end) {
+  packed_matmul_scalar_range(args, out_start, out_end);
+}
+
+}  // namespace orbitquant::cpu
+
+#endif

orbitquant_packed_matmul_cpu/quantize_activations_cpu.cpp

@@ -0,0 +1,867 @@
+#include "cpu_pool.h"
+#include "cpu_threads.h"
+#include "cpu_kernel_args.h"
+#include "packed_matmul_cpu.h"
+#include "../torch-ext/torch_binding.h"
+
+#include <torch/headeronly/core/DeviceType.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/BFloat16.h>
+#include <torch/headeronly/util/Half.h>
+
+#include <algorithm>
+#include <array>
+#include <cmath>
+#include <cstdlib>
+#include <cstdint>
+#include <cstring>
+#include <mutex>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+#if defined(__aarch64__) || defined(_M_ARM64)
+#include <arm_neon.h>
+#endif
+
+#if (defined(__x86_64__) || defined(_M_X64)) && !defined(_MSC_VER)
+#include <immintrin.h>
+#define ORBITQUANT_TARGET_AVX2 __attribute__((target("avx2,fma,f16c")))
+#define ORBITQUANT_TARGET_AVX512 \
+  __attribute__((target("avx512f,avx512dq,avx512bw,avx512vl,fma,f16c")))
+#endif
+
+namespace {
+
+using orbitquant::cpu::ActivationArgs;
+using orbitquant::cpu::ActivationIsa;
+
+ActivationIsa select_activation_isa() {
+  const char *requested = std::getenv("ORBITQUANT_CPU_ISA");
+  if (requested == nullptr || std::strcmp(requested, "auto") == 0) {
+    if (orbitquant::cpu::packed_matmul_x86_avx512_available()) {
+      return ActivationIsa::Avx512;
+    }
+    if (orbitquant::cpu::packed_matmul_x86_avx2_available()) {
+      return ActivationIsa::Avx2;
+    }
+    if (orbitquant::cpu::packed_matmul_neon_available()) {
+      return ActivationIsa::Neon;
+    }
+    return ActivationIsa::Portable;
+  }
+  if (std::strcmp(requested, "scalar") == 0) {
+    return ActivationIsa::Portable;
+  }
+  if (std::strcmp(requested, "avx2") == 0) {
+    STD_TORCH_CHECK(
+        orbitquant::cpu::packed_matmul_x86_avx2_available(),
+        "ORBITQUANT_CPU_ISA=avx2 requested AVX2/FMA/F16C on an unsupported CPU");
+    return ActivationIsa::Avx2;
+  }
+  if (std::strcmp(requested, "avx512") == 0) {
+    STD_TORCH_CHECK(
+        orbitquant::cpu::packed_matmul_x86_avx512_available(),
+        "ORBITQUANT_CPU_ISA=avx512 requested AVX-512F/DQ/BW/VL on an unsupported CPU");
+    return ActivationIsa::Avx512;
+  }
+  if (std::strcmp(requested, "neon") == 0) {
+    STD_TORCH_CHECK(
+        orbitquant::cpu::packed_matmul_neon_available(),
+        "ORBITQUANT_CPU_ISA=neon requested NEON on an unsupported CPU");
+    return ActivationIsa::Neon;
+  }
+  STD_TORCH_CHECK(
+      false,
+      "ORBITQUANT_CPU_ISA must be auto, scalar, avx2, avx512, or neon");
+  return ActivationIsa::Portable;
+}
+
+orbitquant::cpu::ScalarKind activation_scalar_kind(
+    OrbitQuantTensor const &tensor) {
+  using torch::headeronly::ScalarType;
+  switch (tensor.scalar_type()) {
+    case ScalarType::Float:
+      return orbitquant::cpu::ScalarKind::Float32;
+    case ScalarType::Half:
+      return orbitquant::cpu::ScalarKind::Float16;
+    case ScalarType::BFloat16:
+      return orbitquant::cpu::ScalarKind::BFloat16;
+    default:
+      STD_TORCH_CHECK(
+          false,
+          "CPU activation quantization supports float32, float16, and bfloat16 inputs");
+  }
+  return orbitquant::cpu::ScalarKind::Float32;
+}
+
+template <typename scalar_t>
+inline float load_scalar(void const *data, std::int64_t offset) {
+  return static_cast<float>(static_cast<scalar_t const *>(data)[offset]);
+}
+
+template <>
+inline float load_scalar<float>(void const *data, std::int64_t offset) {
+  return static_cast<float const *>(data)[offset];
+}
+
+template <typename scalar_t>
+inline void store_scalar(void *data, std::int64_t offset, float value) {
+  static_cast<scalar_t *>(data)[offset] = scalar_t(value);
+}
+
+template <>
+inline void store_scalar<float>(void *data, std::int64_t offset, float value) {
+  static_cast<float *>(data)[offset] = value;
+}
+
+inline void fwht_block_portable(float *values, std::int64_t block_size) {
+  for (std::int64_t half = 1; half < block_size; half *= 2) {
+    for (std::int64_t base = 0; base < block_size; base += 2 * half) {
+      for (std::int64_t offset = 0; offset < half; ++offset) {
+        const float left = values[base + offset];
+        const float right = values[base + half + offset];
+        values[base + offset] = left + right;
+        values[base + half + offset] = left - right;
+      }
+    }
+  }
+}
+
+#if defined(__aarch64__) || defined(_M_ARM64)
+inline void fwht_block_neon(float *values, std::int64_t block_size) {
+  if (block_size < 4) {
+    fwht_block_portable(values, block_size);
+    return;
+  }
+  const uint32x4_t odd_lanes = {0u, 0xFFFFFFFFu, 0u, 0xFFFFFFFFu};
+  const uint32x4_t high_lanes = {0u, 0u, 0xFFFFFFFFu, 0xFFFFFFFFu};
+  for (std::int64_t base = 0; base < block_size; base += 4) {
+    float32x4_t value = vld1q_f32(values + base);
+    float32x4_t swapped = vrev64q_f32(value);
+    value = vbslq_f32(
+        odd_lanes, vsubq_f32(swapped, value), vaddq_f32(value, swapped));
+    swapped = vextq_f32(value, value, 2);
+    value = vbslq_f32(
+        high_lanes, vsubq_f32(swapped, value), vaddq_f32(value, swapped));
+    vst1q_f32(values + base, value);
+  }
+  for (std::int64_t half = 4; half < block_size; half *= 2) {
+    for (std::int64_t base = 0; base < block_size; base += 2 * half) {
+      for (std::int64_t offset = 0; offset + 4 <= half; offset += 4) {
+        const float32x4_t left = vld1q_f32(values + base + offset);
+        const float32x4_t right = vld1q_f32(values + base + half + offset);
+        vst1q_f32(values + base + offset, vaddq_f32(left, right));
+        vst1q_f32(values + base + half + offset, vsubq_f32(left, right));
+      }
+    }
+  }
+}
+#endif
+
+#if (defined(__x86_64__) || defined(_M_X64)) && !defined(_MSC_VER)
+ORBITQUANT_TARGET_AVX2 void fwht_block_avx2(
+    float *values,
+    std::int64_t block_size) {
+  if (block_size < 8) {
+    fwht_block_portable(values, block_size);
+    return;
+  }
+  for (std::int64_t base = 0; base < block_size; base += 8) {
+    __m256 value = _mm256_loadu_ps(values + base);
+    __m256 swapped = _mm256_permute_ps(value, 0xB1);
+    value = _mm256_blend_ps(
+        _mm256_add_ps(value, swapped), _mm256_sub_ps(swapped, value), 0xAA);
+    swapped = _mm256_permute_ps(value, 0x4E);
+    value = _mm256_blend_ps(
+        _mm256_add_ps(value, swapped), _mm256_sub_ps(swapped, value), 0xCC);
+    swapped = _mm256_permute2f128_ps(value, value, 0x01);
+    value = _mm256_blend_ps(
+        _mm256_add_ps(value, swapped), _mm256_sub_ps(swapped, value), 0xF0);
+    _mm256_storeu_ps(values + base, value);
+  }
+  for (std::int64_t half = 8; half < block_size; half *= 2) {
+    for (std::int64_t base = 0; base < block_size; base += 2 * half) {
+      for (std::int64_t offset = 0; offset + 8 <= half; offset += 8) {
+        const __m256 left = _mm256_loadu_ps(values + base + offset);
+        const __m256 right =
+            _mm256_loadu_ps(values + base + half + offset);
+        _mm256_storeu_ps(values + base + offset, _mm256_add_ps(left, right));
+        _mm256_storeu_ps(
+            values + base + half + offset,
+            _mm256_sub_ps(left, right));
+      }
+    }
+  }
+}
+
+ORBITQUANT_TARGET_AVX512 void fwht_block_avx512(
+    float *values,
+    std::int64_t block_size) {
+  if (block_size < 16) {
+    fwht_block_portable(values, block_size);
+    return;
+  }
+  // Butterfly widths 1-8 stay inside one 16-float register: swap the paired
+  // lanes with shuffles and blend the +/- results, one pass over memory
+  // instead of four.
+  for (std::int64_t base = 0; base < block_size; base += 16) {
+    __m512 value = _mm512_loadu_ps(values + base);
+    __m512 swapped = _mm512_permute_ps(value, 0xB1);
+    value = _mm512_mask_blend_ps(
+        0xAAAA, _mm512_add_ps(value, swapped), _mm512_sub_ps(swapped, value));
+    swapped = _mm512_permute_ps(value, 0x4E);
+    value = _mm512_mask_blend_ps(
+        0xCCCC, _mm512_add_ps(value, swapped), _mm512_sub_ps(swapped, value));
+    swapped = _mm512_shuffle_f32x4(value, value, 0xB1);
+    value = _mm512_mask_blend_ps(
+        0xF0F0, _mm512_add_ps(value, swapped), _mm512_sub_ps(swapped, value));
+    swapped = _mm512_shuffle_f32x4(value, value, 0x4E);
+    value = _mm512_mask_blend_ps(
+        0xFF00, _mm512_add_ps(value, swapped), _mm512_sub_ps(swapped, value));
+    _mm512_storeu_ps(values + base, value);
+  }
+  for (std::int64_t half = 16; half < block_size; half *= 2) {
+    for (std::int64_t base = 0; base < block_size; base += 2 * half) {
+      for (std::int64_t offset = 0; offset + 16 <= half; offset += 16) {
+        const __m512 left = _mm512_loadu_ps(values + base + offset);
+        const __m512 right =
+            _mm512_loadu_ps(values + base + half + offset);
+        _mm512_storeu_ps(values + base + offset, _mm512_add_ps(left, right));
+        _mm512_storeu_ps(
+            values + base + half + offset,
+            _mm512_sub_ps(left, right));
+      }
+    }
+  }
+}
+#endif
+
+inline void fwht_block(
+    float *values,
+    std::int64_t block_size,
+    ActivationIsa isa) {
+#if defined(__x86_64__) || defined(_M_X64)
+  if (isa == ActivationIsa::Avx512) {
+#if defined(_MSC_VER)
+    orbitquant::cpu::activation_fwht_msvc_avx2(values, block_size);
+#else
+    fwht_block_avx512(values, block_size);
+#endif
+    return;
+  }
+  if (isa == ActivationIsa::Avx2) {
+#if defined(_MSC_VER)
+    orbitquant::cpu::activation_fwht_msvc_avx2(values, block_size);
+#else
+    fwht_block_avx2(values, block_size);
+#endif
+    return;
+  }
+#endif
+#if defined(__aarch64__) || defined(_M_ARM64)
+  if (isa == ActivationIsa::Neon) {
+    fwht_block_neon(values, block_size);
+    return;
+  }
+#endif
+  fwht_block_portable(values, block_size);
+}
+
+template <typename scalar_t>
+float squared_norm_scalar(void const *data, std::int64_t offset, std::int64_t dim) {
+  float result = 0.0f;
+  for (std::int64_t index = 0; index < dim; ++index) {
+    const float value = load_scalar<scalar_t>(data, offset + index);
+    result += value * value;
+  }
+  return result;
+}
+
+#if defined(__aarch64__) || defined(_M_ARM64)
+template <typename scalar_t>
+float squared_norm_neon(void const *data, std::int64_t offset, std::int64_t dim) {
+  float32x4_t accumulator = vdupq_n_f32(0.0f);
+  std::int64_t index = 0;
+  for (; index + 4 <= dim; index += 4) {
+    float32x4_t values;
+    if constexpr (std::is_same_v<scalar_t, float>) {
+      values = vld1q_f32(static_cast<float const *>(data) + offset + index);
+    } else if constexpr (std::is_same_v<scalar_t, c10::Half>) {
+      const auto *source = reinterpret_cast<float16_t const *>(
+          static_cast<std::uint16_t const *>(data) + offset + index);
+      values = vcvt_f32_f16(vld1_f16(source));
+    } else {
+      const uint16x4_t raw = vld1_u16(
+          static_cast<std::uint16_t const *>(data) + offset + index);
+      values = vreinterpretq_f32_u32(vshlq_n_u32(vmovl_u16(raw), 16));
+    }
+    accumulator = vfmaq_f32(accumulator, values, values);
+  }
+  float result = vaddvq_f32(accumulator);
+  for (; index < dim; ++index) {
+    const float value = load_scalar<scalar_t>(data, offset + index);
+    result += value * value;
+  }
+  return result;
+}
+#endif
+
+#if (defined(__x86_64__) || defined(_M_X64)) && !defined(_MSC_VER)
+template <typename scalar_t>
+ORBITQUANT_TARGET_AVX2 float squared_norm_avx2(
+    void const *data,
+    std::int64_t offset,
+    std::int64_t dim) {
+  __m256 accumulator = _mm256_setzero_ps();
+  std::int64_t index = 0;
+  for (; index + 8 <= dim; index += 8) {
+    __m256 values;
+    if constexpr (std::is_same_v<scalar_t, float>) {
+      values = _mm256_loadu_ps(
+          static_cast<float const *>(data) + offset + index);
+    } else if constexpr (std::is_same_v<scalar_t, c10::Half>) {
+      const auto *source =
+          static_cast<std::uint16_t const *>(data) + offset + index;
+      values = _mm256_cvtph_ps(
+          _mm_loadu_si128(reinterpret_cast<__m128i const *>(source)));
+    } else {
+      const auto *source =
+          static_cast<std::uint16_t const *>(data) + offset + index;
+      const __m128i packed =
+          _mm_loadu_si128(reinterpret_cast<__m128i const *>(source));
+      values = _mm256_castsi256_ps(
+          _mm256_slli_epi32(_mm256_cvtepu16_epi32(packed), 16));
+    }
+    accumulator = _mm256_fmadd_ps(values, values, accumulator);
+  }
+  const __m128 halves = _mm_add_ps(
+      _mm256_castps256_ps128(accumulator),
+      _mm256_extractf128_ps(accumulator, 1));
+  const __m128 pairs = _mm_hadd_ps(halves, halves);
+  float result = _mm_cvtss_f32(_mm_hadd_ps(pairs, pairs));
+  for (; index < dim; ++index) {
+    const float value = load_scalar<scalar_t>(data, offset + index);
+    result += value * value;
+  }
+  return result;
+}
+
+template <typename scalar_t>
+ORBITQUANT_TARGET_AVX512 float squared_norm_avx512(
+    void const *data,
+    std::int64_t offset,
+    std::int64_t dim) {
+  __m512 accumulator = _mm512_setzero_ps();
+  std::int64_t index = 0;
+  for (; index + 16 <= dim; index += 16) {
+    __m512 values;
+    if constexpr (std::is_same_v<scalar_t, float>) {
+      values = _mm512_loadu_ps(
+          static_cast<float const *>(data) + offset + index);
+    } else if constexpr (std::is_same_v<scalar_t, c10::Half>) {
+      const auto *source =
+          static_cast<std::uint16_t const *>(data) + offset + index;
+      values = _mm512_cvtph_ps(
+          _mm256_loadu_si256(reinterpret_cast<__m256i const *>(source)));
+    } else {
+      const auto *source =
+          static_cast<std::uint16_t const *>(data) + offset + index;
+      const __m256i packed =
+          _mm256_loadu_si256(reinterpret_cast<__m256i const *>(source));
+      values = _mm512_castsi512_ps(
+          _mm512_slli_epi32(_mm512_cvtepu16_epi32(packed), 16));
+    }
+    accumulator = _mm512_fmadd_ps(values, values, accumulator);
+  }
+  float result = _mm512_reduce_add_ps(accumulator);
+  for (; index < dim; ++index) {
+    const float value = load_scalar<scalar_t>(data, offset + index);
+    result += value * value;
+  }
+  return result;
+}
+#endif
+
+template <typename scalar_t>
+float squared_norm(
+    void const *data,
+    std::int64_t offset,
+    std::int64_t dim,
+    ActivationIsa isa) {
+#if defined(__x86_64__) || defined(_M_X64)
+  if (isa == ActivationIsa::Avx512) {
+#if defined(_MSC_VER)
+    constexpr auto scalar_kind = std::is_same_v<scalar_t, float>
+        ? orbitquant::cpu::ScalarKind::Float32
+        : std::is_same_v<scalar_t, c10::Half>
+        ? orbitquant::cpu::ScalarKind::Float16
+        : orbitquant::cpu::ScalarKind::BFloat16;
+    return orbitquant::cpu::activation_squared_norm_msvc_avx2(
+        data, scalar_kind, offset, dim);
+#else
+    return squared_norm_avx512<scalar_t>(data, offset, dim);
+#endif
+  }
+  if (isa == ActivationIsa::Avx2) {
+#if defined(_MSC_VER)
+    constexpr auto scalar_kind = std::is_same_v<scalar_t, float>
+        ? orbitquant::cpu::ScalarKind::Float32
+        : std::is_same_v<scalar_t, c10::Half>
+        ? orbitquant::cpu::ScalarKind::Float16
+        : orbitquant::cpu::ScalarKind::BFloat16;
+    return orbitquant::cpu::activation_squared_norm_msvc_avx2(
+        data, scalar_kind, offset, dim);
+#else
+    return squared_norm_avx2<scalar_t>(data, offset, dim);
+#endif
+  }
+#endif
+#if defined(__aarch64__) || defined(_M_ARM64)
+  if (isa == ActivationIsa::Neon) {
+    return squared_norm_neon<scalar_t>(data, offset, dim);
+  }
+#endif
+  return squared_norm_scalar<scalar_t>(data, offset, dim);
+}
+
+inline std::int64_t nearest_centroid(
+    float value,
+    float const *boundaries,
+    std::int64_t boundary_count) {
+  std::int64_t low = 0;
+  std::int64_t high = boundary_count;
+  while (low < high) {
+    const std::int64_t middle = low + (high - low) / 2;
+    if (value <= boundaries[middle]) {
+      high = middle;
+    } else {
+      low = middle + 1;
+    }
+  }
+  return low;
+}
+
+template <typename scalar_t>
+void quantize_lookup_scalar(
+    ActivationArgs const &args,
+    float const *scratch,
+    std::int64_t output_offset,
+    float norm,
+    std::int64_t start = 0) {
+  for (std::int64_t index = start; index < args.dim; ++index) {
+    const float direction = scratch[index] * args.inv_sqrt_block;
+    const std::int64_t centroid_index =
+        nearest_centroid(direction, args.boundaries, args.boundary_count);
+    store_scalar<scalar_t>(
+        args.out,
+        output_offset + index,
+        args.centroids[centroid_index] * norm);
+  }
+}
+
+#if (defined(__x86_64__) || defined(_M_X64)) && !defined(_MSC_VER)
+ORBITQUANT_TARGET_AVX2 inline __m128i float8_to_bfloat8(__m256 values) {
+  const __m256i bits = _mm256_castps_si256(values);
+  const __m256i absolute_bits =
+      _mm256_and_si256(bits, _mm256_set1_epi32(0x7fffffff));
+  const __m256i nan_mask =
+      _mm256_cmpgt_epi32(absolute_bits, _mm256_set1_epi32(0x7f800000));
+  const __m256i rounding = _mm256_add_epi32(
+      _mm256_set1_epi32(0x7fff),
+      _mm256_and_si256(_mm256_srli_epi32(bits, 16), _mm256_set1_epi32(1)));
+  const __m256i upper = _mm256_srli_epi32(
+      _mm256_add_epi32(bits, rounding),
+      16);
+  __m128i packed = _mm_packus_epi32(
+      _mm256_castsi256_si128(upper),
+      _mm256_extracti128_si256(upper, 1));
+  const __m128i packed_nan_mask = _mm_packs_epi32(
+      _mm256_castsi256_si128(nan_mask),
+      _mm256_extracti128_si256(nan_mask, 1));
+  packed = _mm_blendv_epi8(
+      packed,
+      _mm_set1_epi16(0x7fc0),
+      packed_nan_mask);
+  return packed;
+}
+
+ORBITQUANT_TARGET_AVX512 inline __m256i float16_to_bfloat16(__m512 values) {
+  const __m512i bits = _mm512_castps_si512(values);
+  const __m512i absolute_bits =
+      _mm512_and_si512(bits, _mm512_set1_epi32(0x7fffffff));
+  const __mmask16 nan_mask = _mm512_cmp_epu32_mask(
+      absolute_bits,
+      _mm512_set1_epi32(0x7f800000),
+      _MM_CMPINT_GT);
+  const __m512i rounding = _mm512_add_epi32(
+      _mm512_set1_epi32(0x7fff),
+      _mm512_and_si512(_mm512_srli_epi32(bits, 16), _mm512_set1_epi32(1)));
+  const __m512i upper = _mm512_srli_epi32(
+      _mm512_add_epi32(bits, rounding),
+      16);
+  return _mm256_mask_mov_epi16(
+      _mm512_cvtepi32_epi16(upper),
+      nan_mask,
+      _mm256_set1_epi16(0x7fc0));
+}
+
+template <typename scalar_t>
+ORBITQUANT_TARGET_AVX2 void quantize_lookup_avx2(
+    ActivationArgs const &args,
+    float const *scratch,
+    std::int64_t output_offset,
+    float norm) {
+  const __m256 inverse_sqrt_block = _mm256_set1_ps(args.inv_sqrt_block);
+  const __m256 output_norm = _mm256_set1_ps(norm);
+  const __m256i ones = _mm256_set1_epi32(1);
+  std::int64_t index = 0;
+  for (; index + 8 <= args.dim; index += 8) {
+    const __m256 direction = _mm256_mul_ps(
+        _mm256_loadu_ps(scratch + index), inverse_sqrt_block);
+    __m256i centroid_indices = _mm256_setzero_si256();
+    for (std::int64_t boundary = 0; boundary < args.boundary_count; ++boundary) {
+      const __m256 comparison = _mm256_cmp_ps(
+          direction,
+          _mm256_set1_ps(args.boundaries[boundary]),
+          _CMP_NLE_UQ);
+      centroid_indices = _mm256_add_epi32(
+          centroid_indices,
+          _mm256_and_si256(_mm256_castps_si256(comparison), ones));
+    }
+    const __m256 output = _mm256_mul_ps(
+        _mm256_i32gather_ps(args.centroids, centroid_indices, 4), output_norm);
+    if constexpr (std::is_same_v<scalar_t, float>) {
+      _mm256_storeu_ps(
+          static_cast<float *>(args.out) + output_offset + index,
+          output);
+    } else if constexpr (std::is_same_v<scalar_t, c10::Half>) {
+      _mm_storeu_si128(
+          reinterpret_cast<__m128i *>(
+              static_cast<std::uint16_t *>(args.out) + output_offset + index),
+          _mm256_cvtps_ph(output, _MM_FROUND_TO_NEAREST_INT));
+    } else {
+      _mm_storeu_si128(
+          reinterpret_cast<__m128i *>(
+              static_cast<std::uint16_t *>(args.out) + output_offset + index),
+          float8_to_bfloat8(output));
+    }
+  }
+  quantize_lookup_scalar<scalar_t>(args, scratch, output_offset, norm, index);
+}
+
+template <typename scalar_t>
+ORBITQUANT_TARGET_AVX512 void quantize_lookup_avx512(
+    ActivationArgs const &args,
+    float const *scratch,
+    std::int64_t output_offset,
+    float norm) {
+  const __m512 inverse_sqrt_block = _mm512_set1_ps(args.inv_sqrt_block);
+  const __m512 output_norm = _mm512_set1_ps(norm);
+  const __m512i ones = _mm512_set1_epi32(1);
+  std::int64_t index = 0;
+  for (; index + 16 <= args.dim; index += 16) {
+    const __m512 direction = _mm512_mul_ps(
+        _mm512_loadu_ps(scratch + index), inverse_sqrt_block);
+    __m512i centroid_indices = _mm512_setzero_si512();
+    for (std::int64_t boundary = 0; boundary < args.boundary_count; ++boundary) {
+      const __mmask16 comparison = _mm512_cmp_ps_mask(
+          direction,
+          _mm512_set1_ps(args.boundaries[boundary]),
+          _CMP_NLE_UQ);
+      centroid_indices = _mm512_mask_add_epi32(
+          centroid_indices,
+          comparison,
+          centroid_indices,
+          ones);
+    }
+    const __m512 output = _mm512_mul_ps(
+        _mm512_i32gather_ps(centroid_indices, args.centroids, 4),
+        output_norm);
+    if constexpr (std::is_same_v<scalar_t, float>) {
+      _mm512_storeu_ps(
+          static_cast<float *>(args.out) + output_offset + index,
+          output);
+    } else if constexpr (std::is_same_v<scalar_t, c10::Half>) {
+      _mm256_storeu_si256(
+          reinterpret_cast<__m256i *>(
+              static_cast<std::uint16_t *>(args.out) + output_offset + index),
+          _mm512_cvtps_ph(output, _MM_FROUND_TO_NEAREST_INT));
+    } else {
+      _mm256_storeu_si256(
+          reinterpret_cast<__m256i *>(
+              static_cast<std::uint16_t *>(args.out) + output_offset + index),
+          float16_to_bfloat16(output));
+    }
+  }
+  quantize_lookup_scalar<scalar_t>(args, scratch, output_offset, norm, index);
+}
+#endif
+
+template <typename scalar_t>
+void quantize_lookup(
+    ActivationArgs const &args,
+    float const *scratch,
+    std::int64_t output_offset,
+    float norm,
+    ActivationIsa isa) {
+#if defined(__x86_64__) || defined(_M_X64)
+  if (isa == ActivationIsa::Avx512) {
+#if defined(_MSC_VER)
+    orbitquant::cpu::activation_quantize_lookup_msvc_avx2(
+        args, scratch, output_offset, norm);
+#else
+    quantize_lookup_avx512<scalar_t>(args, scratch, output_offset, norm);
+#endif
+    return;
+  }
+  if (isa == ActivationIsa::Avx2) {
+#if defined(_MSC_VER)
+    orbitquant::cpu::activation_quantize_lookup_msvc_avx2(
+        args, scratch, output_offset, norm);
+#else
+    quantize_lookup_avx2<scalar_t>(args, scratch, output_offset, norm);
+#endif
+    return;
+  }
+#endif
+  quantize_lookup_scalar<scalar_t>(args, scratch, output_offset, norm);
+}
+
+template <typename scalar_t, typename index_t>
+void quantize_activation_range(
+    ActivationArgs const &args,
+    index_t const *permutation,
+    std::int64_t row_start,
+    std::int64_t row_end) {
+  thread_local std::vector<float> scratch;
+  if (scratch.size() < static_cast<std::size_t>(args.dim)) {
+    scratch.resize(static_cast<std::size_t>(args.dim));
+  }
+  for (std::int64_t row = row_start; row < row_end; ++row) {
+    const std::int64_t input_offset = row * args.dim;
+    const float norm_squared =
+        squared_norm<scalar_t>(args.x, input_offset, args.dim, args.isa);
+    const float norm = std::sqrt(norm_squared);
+    const float inverse_norm = 1.0f / (norm + args.eps);
+    for (std::int64_t index = 0; index < args.dim; ++index) {
+      const float value = load_scalar<scalar_t>(
+          args.x,
+          input_offset + static_cast<std::int64_t>(permutation[index]));
+      scratch[index] = value * static_cast<float>(args.signs[index]) * inverse_norm;
+    }
+    for (std::int64_t block = 0; block < args.dim; block += args.block_size) {
+      fwht_block(scratch.data() + block, args.block_size, args.isa);
+    }
+    quantize_lookup<scalar_t>(
+        args,
+        scratch.data(),
+        input_offset,
+        norm,
+        args.isa);
+  }
+}
+
+template <typename index_t>
+void quantize_activation_dispatch_indexed(
+    ActivationArgs const &args,
+    index_t const *permutation,
+    std::int64_t row_start,
+    std::int64_t row_end) {
+  switch (args.scalar_kind) {
+    case orbitquant::cpu::ScalarKind::Float32:
+      quantize_activation_range<float>(args, permutation, row_start, row_end);
+      return;
+    case orbitquant::cpu::ScalarKind::Float16:
+      quantize_activation_range<c10::Half>(args, permutation, row_start, row_end);
+      return;
+    case orbitquant::cpu::ScalarKind::BFloat16:
+      quantize_activation_range<c10::BFloat16>(args, permutation, row_start, row_end);
+      return;
+  }
+}
+
+void quantize_activation_dispatch(
+    ActivationArgs const &args,
+    std::int64_t row_start,
+    std::int64_t row_end) {
+  if (args.permutation_i32 != nullptr) {
+    quantize_activation_dispatch_indexed(
+        args, args.permutation_i32, row_start, row_end);
+    return;
+  }
+  quantize_activation_dispatch_indexed(args, args.permutation, row_start, row_end);
+}
+
+void parallel_quantize_activations(ActivationArgs const &args) {
+  const std::int64_t values = args.rows * args.dim;
+  const int threads = values < 16'384
+      ? 1
+      : std::max<int>(
+            1,
+            std::min<std::int64_t>(
+                orbitquant::cpu::requested_threads(),
+                args.rows));
+  if (threads == 1) {
+    quantize_activation_dispatch(args, 0, args.rows);
+    return;
+  }
+
+  std::vector<std::pair<std::int64_t, std::int64_t>> ranges;
+  ranges.reserve(threads);
+  const std::int64_t rows_per_thread = (args.rows + threads - 1) / threads;
+  for (int thread = 0; thread < threads; ++thread) {
+    const std::int64_t start = thread * rows_per_thread;
+    const std::int64_t end = std::min(args.rows, start + rows_per_thread);
+    if (start >= end) {
+      break;
+    }
+    ranges.emplace_back(start, end);
+  }
+  orbitquant::cpu::run_ranges(
+      ranges,
+      [&args](std::int64_t start, std::int64_t end) {
+        quantize_activation_dispatch(args, start, end);
+      });
+}
+
+}  // namespace
+
+void quantize_activations_cpu(
+    OrbitQuantTensor &out,
+    OrbitQuantTensor const &x,
+    OrbitQuantTensor const &permutation,
+    OrbitQuantTensor const &signs,
+    OrbitQuantTensor const &centroids,
+    OrbitQuantTensor const &boundaries,
+    double eps,
+    double inv_sqrt_block,
+    int64_t block_size) {
+  using torch::headeronly::DeviceType;
+  using torch::headeronly::ScalarType;
+
+  STD_TORCH_CHECK(x.device().type() == DeviceType::CPU, "x must be a CPU tensor");
+  STD_TORCH_CHECK(out.device().type() == DeviceType::CPU, "out must be a CPU tensor");
+  STD_TORCH_CHECK(
+      permutation.device().type() == DeviceType::CPU,
+      "permutation must be a CPU tensor");
+  STD_TORCH_CHECK(signs.device().type() == DeviceType::CPU, "signs must be a CPU tensor");
+  STD_TORCH_CHECK(
+      centroids.device().type() == DeviceType::CPU,
+      "centroids must be a CPU tensor");
+  STD_TORCH_CHECK(
+      boundaries.device().type() == DeviceType::CPU,
+      "boundaries must be a CPU tensor");
+  STD_TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
+  STD_TORCH_CHECK(out.is_contiguous(), "out must be contiguous");
+  STD_TORCH_CHECK(permutation.is_contiguous(), "permutation must be contiguous");
+  STD_TORCH_CHECK(signs.is_contiguous(), "signs must be contiguous");
+  STD_TORCH_CHECK(centroids.is_contiguous(), "centroids must be contiguous");
+  STD_TORCH_CHECK(boundaries.is_contiguous(), "boundaries must be contiguous");
+  STD_TORCH_CHECK(out.scalar_type() == x.scalar_type(), "out dtype must match x dtype");
+  STD_TORCH_CHECK(
+      permutation.scalar_type() == ScalarType::Long ||
+          permutation.scalar_type() == ScalarType::Int,
+      "permutation must be int32 or int64");
+  STD_TORCH_CHECK(signs.scalar_type() == ScalarType::Char, "signs must be int8");
+  STD_TORCH_CHECK(centroids.scalar_type() == ScalarType::Float, "centroids must be float32");
+  STD_TORCH_CHECK(boundaries.scalar_type() == ScalarType::Float, "boundaries must be float32");
+  STD_TORCH_CHECK(x.dim() == 2, "x must be rank 2");
+  STD_TORCH_CHECK(out.dim() == 2, "out must be rank 2");
+  STD_TORCH_CHECK(out.size(0) == x.size(0), "out row count must match x");
+  STD_TORCH_CHECK(out.size(1) == x.size(1), "out dimension must match x");
+  const int64_t dim = x.size(1);
+  STD_TORCH_CHECK(permutation.numel() == dim, "permutation must match the input dimension");
+  STD_TORCH_CHECK(signs.numel() == dim, "signs must match the input dimension");
+  STD_TORCH_CHECK(
+      centroids.numel() == boundaries.numel() + 1,
+      "centroids must contain exactly one more value than boundaries");
+  STD_TORCH_CHECK(centroids.numel() >= 2, "at least two centroids are required");
+  STD_TORCH_CHECK(block_size > 0, "block_size must be positive");
+  STD_TORCH_CHECK(
+      (block_size & (block_size - 1)) == 0,
+      "block_size must be a power of two");
+  STD_TORCH_CHECK(dim % block_size == 0, "block_size must divide the input dimension");
+  STD_TORCH_CHECK(eps >= 0.0, "eps must be non-negative");
+  STD_TORCH_CHECK(inv_sqrt_block > 0.0, "inv_sqrt_block must be positive");
+
+  const bool int32_permutation = permutation.scalar_type() == ScalarType::Int;
+  const std::int64_t *permutation_values =
+      int32_permutation ? nullptr : permutation.const_data_ptr<std::int64_t>();
+  const std::int32_t *permutation_values_i32 =
+      int32_permutation ? permutation.const_data_ptr<std::int32_t>() : nullptr;
+  const auto *sign_values = signs.const_data_ptr<std::int8_t>();
+  const auto permutation_at = [&](std::int64_t index) -> std::int64_t {
+    return int32_permutation
+        ? static_cast<std::int64_t>(permutation_values_i32[index])
+        : permutation_values[index];
+  };
+  // The permutation/sign buffers are immutable module constants, so cache a
+  // fingerprint of validated buffers instead of re-scanning them per forward.
+  struct ValidatedEntry {
+    void const *permutation;
+    void const *signs;
+    std::int64_t dim;
+    std::int64_t head;
+    std::int64_t tail;
+  };
+  static std::mutex validated_mutex;
+  static std::array<ValidatedEntry, 16> validated_entries{};
+  static std::size_t validated_cursor = 0;
+  const ValidatedEntry candidate{
+      int32_permutation ? static_cast<void const *>(permutation_values_i32)
+                        : static_cast<void const *>(permutation_values),
+      sign_values,
+      dim,
+      permutation_at(0),
+      permutation_at(dim - 1)};
+  bool already_validated = false;
+  {
+    std::lock_guard<std::mutex> lock(validated_mutex);
+    for (auto const &entry : validated_entries) {
+      if (entry.permutation == candidate.permutation &&
+          entry.signs == candidate.signs && entry.dim == candidate.dim &&
+          entry.head == candidate.head && entry.tail == candidate.tail) {
+        already_validated = true;
+        break;
+      }
+    }
+  }
+  if (!already_validated) {
+    for (int64_t index = 0; index < dim; ++index) {
+      const std::int64_t source_index = permutation_at(index);
+      STD_TORCH_CHECK(
+          source_index >= 0 && source_index < dim,
+          "permutation contains an out-of-range index");
+      STD_TORCH_CHECK(
+          sign_values[index] == -1 || sign_values[index] == 1,
+          "signs must contain only -1 and 1");
+    }
+    std::lock_guard<std::mutex> lock(validated_mutex);
+    validated_entries[validated_cursor % validated_entries.size()] = candidate;
+    ++validated_cursor;
+  }
+  if (x.numel() == 0) {
+    return;
+  }
+
+  const ActivationArgs args{
+      out.mutable_data_ptr(),
+      x.const_data_ptr(),
+      permutation_values,
+      permutation_values_i32,
+      sign_values,
+      centroids.const_data_ptr<float>(),
+      boundaries.const_data_ptr<float>(),
+      activation_scalar_kind(x),
+      select_activation_isa(),
+      x.size(0),
+      dim,
+      boundaries.numel(),
+      block_size,
+      static_cast<float>(eps),
+      static_cast<float>(inv_sqrt_block),
+  };
+  parallel_quantize_activations(args);
+}

orbitquant_packed_matmul_cuda/packed_matmul.cu

@@ -11,9 +11,34 @@
 
 #include <algorithm>
 #include <cstdint>
+#include <cstdlib>
 
 using namespace nvcuda;
 
+namespace {
+
+// Escape hatch (and A/B benchmarking toggle) for the cp.async-pipelined mma64
+// path: set ORBITQUANT_MMA64_DISABLE_PIPELINE=1 to force the legacy kernel.
+inline bool orbitquant_mma64_pipeline_disabled() {
+  static const bool disabled = []() {
+    char const *value = std::getenv("ORBITQUANT_MMA64_DISABLE_PIPELINE");
+    return value != nullptr && value[0] != '\0' && value[0] != '0';
+  }();
+  return disabled;
+}
+
+// ORBITQUANT_MMA64_FORCE_PIPELINE=1 extends the pipeline to every eligible
+// width (W4/W6) for re-benchmarking on other GPUs.
+inline bool orbitquant_mma64_pipeline_forced() {
+  static const bool forced = []() {
+    char const *value = std::getenv("ORBITQUANT_MMA64_FORCE_PIPELINE");
+    return value != nullptr && value[0] != '\0' && value[0] != '0';
+  }();
+  return forced;
+}
+
+}  // namespace
+
 __device__ __forceinline__ uint32_t unpack_lowbit_index(
     uint8_t const *__restrict__ packed_weight_indices,
     int64_t value_offset,
@@ -44,18 +69,943 @@ __device__ __forceinline__ uint32_t unpack_lowbit_index_const(
   return (raw >> bit_offset) & mask;
 }
 
+template <typename T>
+__device__ __forceinline__ T orbitquant_mma_from_float(float value);
+
+template <>
+__device__ __forceinline__ half orbitquant_mma_from_float<half>(float value) {
+  return __float2half(value);
+}
+
+template <>
+__device__ __forceinline__ __nv_bfloat16 orbitquant_mma_from_float<__nv_bfloat16>(
+    float value) {
+  return __float2bfloat16(value);
+}
+
+template <int Bits, int Values, typename mma_t>
+__device__ __forceinline__ void decode_mma64_weight_segment(
+    mma_t *__restrict__ destination,
+    uint8_t const *__restrict__ packed_weight_indices,
+    c10::BFloat16 const *__restrict__ row_norms,
+    float const *__restrict__ centroids,
+    int64_t global_col,
+    int64_t global_k,
+    int64_t in_features,
+    bool valid_segment) {
+  constexpr uint32_t mask = (1u << Bits) - 1u;
+  constexpr int word_count = (Values * Bits + 31) / 32;
+  uint32_t packed_words[word_count] = {};
+  float norm = 0.0f;
+  if (valid_segment) {
+    const int64_t value_offset = global_col * in_features + global_k;
+    const int64_t byte_index = (value_offset * Bits) >> 3;
+    auto const *words = reinterpret_cast<uint32_t const *>(
+        packed_weight_indices + byte_index);
+#pragma unroll
+    for (int word = 0; word < word_count; ++word) {
+      packed_words[word] = words[word];
+    }
+    norm = static_cast<float>(row_norms[global_col]);
+  }
+
+#pragma unroll
+  for (int index_offset = 0; index_offset < Values; ++index_offset) {
+    const int bit_start = index_offset * Bits;
+    const int word_index = bit_start >> 5;
+    const int shift = bit_start & 31;
+    uint32_t raw = packed_words[word_index] >> shift;
+    if (shift + Bits > 32) {
+      raw |= packed_words[word_index + 1] << (32 - shift);
+    }
+    const uint32_t codebook_index = raw & mask;
+    const float value = valid_segment ? norm * centroids[codebook_index] : 0.0f;
+    destination[index_offset] = orbitquant_mma_from_float<mma_t>(value);
+  }
+}
+
+template <int Bits, int Values, typename mma_t>
+__device__ __forceinline__ void decode_mma64_weight_segment_from_stage(
+    mma_t *__restrict__ destination,
+    uint8_t const *__restrict__ staged_bytes,
+    float const *__restrict__ centroids,
+    float norm,
+    bool valid_segment) {
+  constexpr uint32_t mask = (1u << Bits) - 1u;
+  constexpr int word_count = (Values * Bits + 31) / 32;
+  uint32_t packed_words[word_count] = {};
+  if (valid_segment) {
+    auto const *words = reinterpret_cast<uint32_t const *>(staged_bytes);
+#pragma unroll
+    for (int word = 0; word < word_count; ++word) {
+      packed_words[word] = words[word];
+    }
+  }
+
+#pragma unroll
+  for (int index_offset = 0; index_offset < Values; ++index_offset) {
+    const int bit_start = index_offset * Bits;
+    const int word_index = bit_start >> 5;
+    const int shift = bit_start & 31;
+    uint32_t raw = packed_words[word_index] >> shift;
+    if (shift + Bits > 32) {
+      raw |= packed_words[word_index + 1] << (32 - shift);
+    }
+    const uint32_t codebook_index = raw & mask;
+    const float value = valid_segment ? norm * centroids[codebook_index] : 0.0f;
+    destination[index_offset] = orbitquant_mma_from_float<mma_t>(value);
+  }
+}
+
+// WMMA fragments (and the bf16 variants in particular) only exist on sm80+.
+// Multi-architecture builds still instantiate these templates for older
+// targets, so the bodies compile to empty stubs there; the host dispatch
+// requires compute capability >= 8 before launching either kernel.
+template <typename storage_t, typename mma_t, int Bits>
+__global__ void orbitquant_packed_matmul_mma64_kernel(
+    storage_t *__restrict__ out,
+    storage_t const *__restrict__ x,
+    uint8_t const *__restrict__ packed_weight_indices,
+    c10::BFloat16 const *__restrict__ row_norms,
+    float const *__restrict__ centroids,
+    storage_t const *__restrict__ bias,
+    bool has_bias,
+    int64_t rows,
+    int64_t out_features,
+    int64_t in_features) {
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
+  constexpr int tile_m = 128;
+  constexpr int tile_n = 128;
+  constexpr int tile_k = 64;
+  constexpr int padded_k = 72;
+  constexpr int warps_per_block = 8;
+  constexpr int warp_tile = 16;
+  constexpr int col_tiles = tile_n / warp_tile;
+  constexpr int x_vector_values = 8;
+  constexpr int x_vectors_per_row = tile_k / x_vector_values;
+  constexpr int weight_segment_values = Bits == 3 ? 32 : 16;
+  constexpr int weight_segments_per_row = tile_k / weight_segment_values;
+  static_assert(sizeof(storage_t) == sizeof(mma_t));
+
+  __shared__ mma_t x_tile[tile_m * padded_k];
+  __shared__ mma_t weight_tile[tile_n * padded_k];
+  __shared__ float accumulator_tile[warps_per_block * warp_tile * warp_tile];
+
+  const int warp_id = threadIdx.x / warpSize;
+  const int lane = threadIdx.x & (warpSize - 1);
+  const int64_t block_row = int64_t(blockIdx.y) * tile_m;
+  const int64_t block_col = int64_t(blockIdx.x) * tile_n;
+
+  wmma::fragment<wmma::accumulator, warp_tile, warp_tile, warp_tile, float>
+      accumulators[col_tiles];
+#pragma unroll
+  for (int col_tile = 0; col_tile < col_tiles; ++col_tile) {
+    wmma::fill_fragment(accumulators[col_tile], 0.0f);
+  }
+
+  for (int64_t k_start = 0; k_start < in_features; k_start += tile_k) {
+    constexpr int x_vector_tasks = tile_m * x_vectors_per_row;
+    for (int task = threadIdx.x; task < x_vector_tasks; task += blockDim.x) {
+      const int local_row = task / x_vectors_per_row;
+      const int local_vector = task - local_row * x_vectors_per_row;
+      const int local_k = local_vector * x_vector_values;
+      const int64_t global_row = block_row + local_row;
+      auto *destination = reinterpret_cast<uint4 *>(
+          x_tile + local_row * padded_k + local_k);
+      if (global_row < rows) {
+        auto const *source = reinterpret_cast<uint4 const *>(
+            x + global_row * in_features + k_start + local_k);
+        *destination = *source;
+      } else {
+        *destination = make_uint4(0, 0, 0, 0);
+      }
+    }
+
+    constexpr int weight_tasks = tile_n * weight_segments_per_row;
+    for (int weight_task = threadIdx.x; weight_task < weight_tasks;
+         weight_task += blockDim.x) {
+      const int local_col = weight_task / weight_segments_per_row;
+      const int local_segment =
+          weight_task - local_col * weight_segments_per_row;
+      const int local_k = local_segment * weight_segment_values;
+      const int64_t global_col = block_col + local_col;
+      decode_mma64_weight_segment<Bits, weight_segment_values>(
+          weight_tile + local_col * padded_k + local_k,
+          packed_weight_indices,
+          row_norms,
+          centroids,
+          global_col,
+          k_start + local_k,
+          in_features,
+          global_col < out_features);
+    }
+    __syncthreads();
+
+#pragma unroll
+    for (int local_k = 0; local_k < tile_k; local_k += warp_tile) {
+      wmma::fragment<wmma::matrix_a, warp_tile, warp_tile, warp_tile, mma_t,
+                     wmma::row_major>
+          lhs;
+      wmma::load_matrix_sync(
+          lhs,
+          x_tile + warp_id * warp_tile * padded_k + local_k,
+          padded_k);
+#pragma unroll
+      for (int col_tile = 0; col_tile < col_tiles; ++col_tile) {
+        wmma::fragment<wmma::matrix_b, warp_tile, warp_tile, warp_tile, mma_t,
+                       wmma::col_major>
+            rhs;
+        wmma::load_matrix_sync(
+            rhs,
+            weight_tile + col_tile * warp_tile * padded_k + local_k,
+            padded_k);
+        wmma::mma_sync(
+            accumulators[col_tile], lhs, rhs, accumulators[col_tile]);
+      }
+    }
+    __syncthreads();
+  }
+
+  float *warp_accumulator = accumulator_tile + warp_id * warp_tile * warp_tile;
+#pragma unroll
+  for (int col_tile = 0; col_tile < col_tiles; ++col_tile) {
+    wmma::store_matrix_sync(
+        warp_accumulator,
+        accumulators[col_tile],
+        warp_tile,
+        wmma::mem_row_major);
+    __syncwarp();
+    for (int offset = lane; offset < warp_tile * warp_tile; offset += warpSize) {
+      const int local_row = offset / warp_tile;
+      const int local_col = offset - local_row * warp_tile;
+      const int64_t global_row = block_row + warp_id * warp_tile + local_row;
+      const int64_t global_col = block_col + col_tile * warp_tile + local_col;
+      if (global_row < rows && global_col < out_features) {
+        float value = warp_accumulator[offset];
+        if (has_bias) {
+          value += static_cast<float>(bias[global_col]);
+        }
+        out[global_row * out_features + global_col] =
+            static_cast<storage_t>(value);
+      }
+    }
+    __syncwarp();
+  }
+#endif  // __CUDA_ARCH__ >= 800
+}
+
+__device__ __forceinline__ void copy_async_16(
+    void *__restrict__ destination,
+    void const *__restrict__ source) {
+#if __CUDA_ARCH__ >= 800
+  const uint32_t shared_address =
+      static_cast<uint32_t>(__cvta_generic_to_shared(destination));
+  asm volatile(
+      "cp.async.ca.shared.global [%0], [%1], 16;\n" : : "r"(shared_address),
+      "l"(source));
+#else
+  *reinterpret_cast<uint4 *>(destination) =
+      *reinterpret_cast<uint4 const *>(source);
+#endif
+}
+
+__device__ __forceinline__ void commit_async_copies() {
+#if __CUDA_ARCH__ >= 800
+  asm volatile("cp.async.commit_group;\n" : :);
+#endif
+}
+
+__device__ __forceinline__ void wait_for_async_copies() {
+#if __CUDA_ARCH__ >= 800
+  asm volatile("cp.async.wait_group 0;\n" : :);
+#endif
+}
+
+// cp.async-pipelined variant of the mma64 kernel (sm80+, byte-aligned k-tile
+// strides, i.e. every supported width except W3). Interior blocks double-buffer
+// the X tile and the packed weight bytes so global loads overlap the MMA work;
+// edge blocks keep the guarded synchronous path.
+template <typename storage_t, typename mma_t, int Bits>
+__global__ void orbitquant_packed_matmul_mma64_pipelined_kernel(
+    storage_t *__restrict__ out,
+    storage_t const *__restrict__ x,
+    uint8_t const *__restrict__ packed_weight_indices,
+    c10::BFloat16 const *__restrict__ row_norms,
+    float const *__restrict__ centroids,
+    storage_t const *__restrict__ bias,
+    bool has_bias,
+    int64_t rows,
+    int64_t out_features,
+    int64_t in_features) {
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
+  constexpr int tile_m = 128;
+  constexpr int tile_n = 128;
+  constexpr int tile_k = 64;
+  constexpr int padded_k = 72;
+  constexpr int warps_per_block = 8;
+  constexpr int warp_tile = 16;
+  constexpr int col_tiles = tile_n / warp_tile;
+  constexpr int x_vector_values = 8;
+  constexpr int x_vectors_per_row = tile_k / x_vector_values;
+  constexpr int seg_stride = tile_k * Bits / 8;
+  constexpr int weight_segment_values = 16;
+  constexpr int weight_segments_per_row = tile_k / weight_segment_values;
+  constexpr int segment_bytes = weight_segment_values * Bits / 8;
+  static_assert(sizeof(storage_t) == sizeof(mma_t));
+  static_assert(Bits == 2 || Bits == 4 || Bits == 6);
+
+  extern __shared__ __align__(16) uint8_t dynamic_shared[];
+  mma_t *x_tiles = reinterpret_cast<mma_t *>(dynamic_shared);
+  mma_t *weight_tile = x_tiles + 2 * tile_m * padded_k;
+  float *accumulator_tile =
+      reinterpret_cast<float *>(weight_tile + tile_n * padded_k);
+  uint8_t *packed_stage = reinterpret_cast<uint8_t *>(
+      accumulator_tile + warps_per_block * warp_tile * warp_tile);
+
+  const int warp_id = threadIdx.x / warpSize;
+  const int lane = threadIdx.x & (warpSize - 1);
+  const int64_t block_row = int64_t(blockIdx.y) * tile_m;
+  const int64_t block_col = int64_t(blockIdx.x) * tile_n;
+  const int64_t packed_row_bytes = in_features * Bits / 8;
+
+  wmma::fragment<wmma::accumulator, warp_tile, warp_tile, warp_tile, float>
+      accumulators[col_tiles];
+#pragma unroll
+  for (int col_tile = 0; col_tile < col_tiles; ++col_tile) {
+    wmma::fill_fragment(accumulators[col_tile], 0.0f);
+  }
+
+  const bool interior =
+      block_row + tile_m <= rows && block_col + tile_n <= out_features;
+
+  if (interior) {
+    constexpr int x_copies = tile_m * x_vectors_per_row;
+    constexpr int w_copies = tile_n * seg_stride / 16;
+    auto stage_tile = [&](int buffer, int64_t k_start) {
+      mma_t *x_destination = x_tiles + buffer * tile_m * padded_k;
+      for (int task = threadIdx.x; task < x_copies; task += blockDim.x) {
+        const int local_row = task / x_vectors_per_row;
+        const int local_vector = task - local_row * x_vectors_per_row;
+        copy_async_16(
+            x_destination + local_row * padded_k + local_vector * x_vector_values,
+            x + (block_row + local_row) * in_features + k_start +
+                local_vector * x_vector_values);
+      }
+      uint8_t *stage_destination = packed_stage + buffer * tile_n * seg_stride;
+      for (int task = threadIdx.x; task < w_copies; task += blockDim.x) {
+        const int byte_offset = task * 16;
+        const int local_col = byte_offset / seg_stride;
+        const int seg_byte = byte_offset - local_col * seg_stride;
+        copy_async_16(
+            stage_destination + byte_offset,
+            packed_weight_indices + (block_col + local_col) * packed_row_bytes +
+                (k_start * Bits) / 8 + seg_byte);
+      }
+      commit_async_copies();
+    };
+
+    stage_tile(0, 0);
+    wait_for_async_copies();
+    __syncthreads();
+
+    for (int64_t k_start = 0; k_start < in_features; k_start += tile_k) {
+      const int buffer = static_cast<int>((k_start / tile_k) & 1);
+      const bool has_next = k_start + tile_k < in_features;
+      if (has_next) {
+        stage_tile(buffer ^ 1, k_start + tile_k);
+      }
+
+      uint8_t const *stage_source = packed_stage + buffer * tile_n * seg_stride;
+      constexpr int weight_tasks = tile_n * weight_segments_per_row;
+      for (int task = threadIdx.x; task < weight_tasks; task += blockDim.x) {
+        const int local_col = task / weight_segments_per_row;
+        const int local_segment = task - local_col * weight_segments_per_row;
+        decode_mma64_weight_segment_from_stage<Bits, weight_segment_values>(
+            weight_tile + local_col * padded_k +
+                local_segment * weight_segment_values,
+            stage_source + local_col * seg_stride + local_segment * segment_bytes,
+            centroids,
+            static_cast<float>(row_norms[block_col + local_col]),
+            true);
+      }
+      __syncthreads();
+
+      mma_t const *x_source = x_tiles + buffer * tile_m * padded_k;
+#pragma unroll
+      for (int local_k = 0; local_k < tile_k; local_k += warp_tile) {
+        wmma::fragment<wmma::matrix_a, warp_tile, warp_tile, warp_tile, mma_t,
+                       wmma::row_major>
+            lhs;
+        wmma::load_matrix_sync(
+            lhs,
+            x_source + warp_id * warp_tile * padded_k + local_k,
+            padded_k);
+#pragma unroll
+        for (int col_tile = 0; col_tile < col_tiles; ++col_tile) {
+          wmma::fragment<wmma::matrix_b, warp_tile, warp_tile, warp_tile, mma_t,
+                         wmma::col_major>
+              rhs;
+          wmma::load_matrix_sync(
+              rhs,
+              weight_tile + col_tile * warp_tile * padded_k + local_k,
+              padded_k);
+          wmma::mma_sync(
+              accumulators[col_tile], lhs, rhs, accumulators[col_tile]);
+        }
+      }
+      __syncthreads();
+      if (has_next) {
+        wait_for_async_copies();
+        __syncthreads();
+      }
+    }
+  } else {
+    mma_t *x_tile = x_tiles;
+    for (int64_t k_start = 0; k_start < in_features; k_start += tile_k) {
+      constexpr int x_vector_tasks = tile_m * x_vectors_per_row;
+      for (int task = threadIdx.x; task < x_vector_tasks; task += blockDim.x) {
+        const int local_row = task / x_vectors_per_row;
+        const int local_vector = task - local_row * x_vectors_per_row;
+        const int local_k = local_vector * x_vector_values;
+        const int64_t global_row = block_row + local_row;
+        auto *destination = reinterpret_cast<uint4 *>(
+            x_tile + local_row * padded_k + local_k);
+        if (global_row < rows) {
+          auto const *source = reinterpret_cast<uint4 const *>(
+              x + global_row * in_features + k_start + local_k);
+          *destination = *source;
+        } else {
+          *destination = make_uint4(0, 0, 0, 0);
+        }
+      }
+
+      constexpr int weight_tasks = tile_n * weight_segments_per_row;
+      for (int weight_task = threadIdx.x; weight_task < weight_tasks;
+           weight_task += blockDim.x) {
+        const int local_col = weight_task / weight_segments_per_row;
+        const int local_segment =
+            weight_task - local_col * weight_segments_per_row;
+        const int local_k = local_segment * weight_segment_values;
+        const int64_t global_col = block_col + local_col;
+        decode_mma64_weight_segment<Bits, weight_segment_values>(
+            weight_tile + local_col * padded_k + local_k,
+            packed_weight_indices,
+            row_norms,
+            centroids,
+            global_col,
+            k_start + local_k,
+            in_features,
+            global_col < out_features);
+      }
+      __syncthreads();
+
+#pragma unroll
+      for (int local_k = 0; local_k < tile_k; local_k += warp_tile) {
+        wmma::fragment<wmma::matrix_a, warp_tile, warp_tile, warp_tile, mma_t,
+                       wmma::row_major>
+            lhs;
+        wmma::load_matrix_sync(
+            lhs,
+            x_tile + warp_id * warp_tile * padded_k + local_k,
+            padded_k);
+#pragma unroll
+        for (int col_tile = 0; col_tile < col_tiles; ++col_tile) {
+          wmma::fragment<wmma::matrix_b, warp_tile, warp_tile, warp_tile, mma_t,
+                         wmma::col_major>
+              rhs;
+          wmma::load_matrix_sync(
+              rhs,
+              weight_tile + col_tile * warp_tile * padded_k + local_k,
+              padded_k);
+          wmma::mma_sync(
+              accumulators[col_tile], lhs, rhs, accumulators[col_tile]);
+        }
+      }
+      __syncthreads();
+    }
+  }
+
+  float *warp_accumulator = accumulator_tile + warp_id * warp_tile * warp_tile;
+#pragma unroll
+  for (int col_tile = 0; col_tile < col_tiles; ++col_tile) {
+    wmma::store_matrix_sync(
+        warp_accumulator,
+        accumulators[col_tile],
+        warp_tile,
+        wmma::mem_row_major);
+    __syncwarp();
+    for (int offset = lane; offset < warp_tile * warp_tile; offset += warpSize) {
+      const int local_row = offset / warp_tile;
+      const int local_col = offset - local_row * warp_tile;
+      const int64_t global_row = block_row + warp_id * warp_tile + local_row;
+      const int64_t global_col = block_col + col_tile * warp_tile + local_col;
+      if (global_row < rows && global_col < out_features) {
+        float value = warp_accumulator[offset];
+        if (has_bias) {
+          value += static_cast<float>(bias[global_col]);
+        }
+        out[global_row * out_features + global_col] =
+            static_cast<storage_t>(value);
+      }
+    }
+    __syncwarp();
+  }
+#endif  // __CUDA_ARCH__ >= 800
+}
+
+__device__ __forceinline__ uint8_t orbitquant_bucketize_w4(
+    float value,
+    float const *__restrict__ boundaries) {
+  int index = value > boundaries[7] ? 8 : 0;
+  index += value > boundaries[index + 3] ? 4 : 0;
+  index += value > boundaries[index + 1] ? 2 : 0;
+  index += value > boundaries[index] ? 1 : 0;
+  return static_cast<uint8_t>(index);
+}
+
+template <typename storage_t, typename index_t, int Dim>
+__global__ void orbitquant_rpbh_quantize_pack_w4_kernel(
+    uint8_t *__restrict__ packed_out,
+    float *__restrict__ norms_out,
+    storage_t const *__restrict__ x,
+    index_t const *__restrict__ permutation,
+    int8_t const *__restrict__ signs,
+    float const *__restrict__ boundaries,
+    float eps,
+    float inv_sqrt_block,
+    int64_t rows) {
+  extern __shared__ __align__(16) float shared[];
+  float *values = shared;
+  float *reduction = values + Dim;
+  float *boundary_table = reduction + blockDim.x;
+  const int tid = threadIdx.x;
+  const int64_t row = blockIdx.x;
+
+  float squared_sum = 0.0f;
+  for (int col = tid; col < Dim; col += blockDim.x) {
+    const int64_t source_col = permutation[col];
+    const float value =
+        static_cast<float>(x[row * Dim + source_col]) *
+        static_cast<float>(signs[col]);
+    values[col] = value;
+    squared_sum = fmaf(value, value, squared_sum);
+  }
+  reduction[tid] = squared_sum;
+  if (tid < 15) {
+    boundary_table[tid] = boundaries[tid];
+  }
+  __syncthreads();
+
+  for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+    if (tid < stride) {
+      reduction[tid] += reduction[tid + stride];
+    }
+    __syncthreads();
+  }
+  const float norm = sqrtf(reduction[0]);
+  if (tid == 0) {
+    norms_out[row] = norm;
+  }
+  const float inv_norm = 1.0f / (norm + eps);
+  for (int col = tid; col < Dim; col += blockDim.x) {
+    values[col] *= inv_norm;
+  }
+  __syncthreads();
+
+#pragma unroll
+  for (int butterfly_width = 1; butterfly_width < Dim;
+       butterfly_width <<= 1) {
+    constexpr int butterflies = Dim / 2;
+    for (int butterfly = tid; butterfly < butterflies;
+         butterfly += blockDim.x) {
+      const int group = butterfly / butterfly_width;
+      const int offset = butterfly - group * butterfly_width;
+      const int left = group * (butterfly_width * 2) + offset;
+      const int right = left + butterfly_width;
+      const float lhs = values[left];
+      const float rhs = values[right];
+      values[left] = lhs + rhs;
+      values[right] = lhs - rhs;
+    }
+    __syncthreads();
+  }
+
+  constexpr int packed_dim = Dim / 2;
+  for (int byte_col = tid; byte_col < packed_dim; byte_col += blockDim.x) {
+    const float low_value = values[byte_col * 2] * inv_sqrt_block;
+    const float high_value = values[byte_col * 2 + 1] * inv_sqrt_block;
+    const uint8_t low = orbitquant_bucketize_w4(low_value, boundary_table);
+    const uint8_t high = orbitquant_bucketize_w4(high_value, boundary_table);
+    packed_out[row * packed_dim + byte_col] =
+        static_cast<uint8_t>(low | (high << 4));
+  }
+}
+
+template <typename storage_t, typename index_t, int Dim, int OrbitBlock>
+__global__ void orbitquant_rpbh_quantize_int8_kernel(
+    int8_t *__restrict__ int8_out,
+    float *__restrict__ norms_out,
+    storage_t const *__restrict__ x,
+    index_t const *__restrict__ permutation,
+    int8_t const *__restrict__ signs,
+    float const *__restrict__ boundaries,
+    int8_t const *__restrict__ codes,
+    float eps,
+    float inv_sqrt_block,
+    int64_t rows) {
+  static_assert(Dim % OrbitBlock == 0);
+  extern __shared__ __align__(16) float shared[];
+  float *values = shared;
+  float *reduction = values + Dim;
+  float *boundary_table = reduction + blockDim.x;
+  int8_t *code_table = reinterpret_cast<int8_t *>(boundary_table + 15);
+  const int tid = threadIdx.x;
+  const int64_t row = blockIdx.x;
+
+  float squared_sum = 0.0f;
+  for (int col = tid; col < Dim; col += blockDim.x) {
+    const int64_t source_col = permutation[col];
+    const float value =
+        static_cast<float>(x[row * Dim + source_col]) *
+        static_cast<float>(signs[col]);
+    values[col] = value;
+    squared_sum = fmaf(value, value, squared_sum);
+  }
+  reduction[tid] = squared_sum;
+  if (tid < 15) {
+    boundary_table[tid] = boundaries[tid];
+  }
+  if (tid < 16) {
+    code_table[tid] = codes[tid];
+  }
+  __syncthreads();
+
+  for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+    if (tid < stride) {
+      reduction[tid] += reduction[tid + stride];
+    }
+    __syncthreads();
+  }
+  const float norm = sqrtf(reduction[0]);
+  if (tid == 0) {
+    norms_out[row] = norm;
+  }
+  const float inv_norm = 1.0f / (norm + eps);
+  for (int col = tid; col < Dim; col += blockDim.x) {
+    values[col] *= inv_norm;
+  }
+  __syncthreads();
+
+#pragma unroll
+  for (int butterfly_width = 1; butterfly_width < OrbitBlock;
+       butterfly_width <<= 1) {
+    constexpr int butterflies = Dim / 2;
+    constexpr int butterflies_per_block = OrbitBlock / 2;
+    for (int butterfly = tid; butterfly < butterflies;
+         butterfly += blockDim.x) {
+      const int orbit_block = butterfly / butterflies_per_block;
+      const int local_butterfly =
+          butterfly - orbit_block * butterflies_per_block;
+      const int group = local_butterfly / butterfly_width;
+      const int offset = local_butterfly - group * butterfly_width;
+      const int left = orbit_block * OrbitBlock +
+                       group * (butterfly_width * 2) + offset;
+      const int right = left + butterfly_width;
+      const float lhs = values[left];
+      const float rhs = values[right];
+      values[left] = lhs + rhs;
+      values[right] = lhs - rhs;
+    }
+    __syncthreads();
+  }
+
+  for (int col = tid; col < Dim; col += blockDim.x) {
+    const float value = values[col] * inv_sqrt_block;
+    const uint8_t index = orbitquant_bucketize_w4(value, boundary_table);
+    int8_out[row * Dim + col] = code_table[index];
+  }
+}
+
+template <
+    typename storage_t,
+    int TileM,
+    int TileN,
+    bool AsyncPacked,
+    bool KMajorWeight>
+__global__ void orbitquant_packed_w4a4_int8_mma_kernel(
+    storage_t *__restrict__ out,
+    uint8_t const *__restrict__ packed_activations,
+    uint8_t const *__restrict__ packed_weight_indices,
+    float const *__restrict__ token_norms,
+    c10::BFloat16 const *__restrict__ row_norms,
+    int8_t const *__restrict__ activation_codes,
+    int8_t const *__restrict__ weight_codes,
+    storage_t const *__restrict__ bias,
+    bool has_bias,
+    float activation_scale,
+    float weight_scale,
+    int64_t rows,
+    int64_t out_features,
+    int64_t in_features) {
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
+  constexpr int tile_k = 64;
+  constexpr int packed_tile_k = tile_k / 2;
+  constexpr int padded_k = 80;
+  constexpr int warp_tile = 16;
+  constexpr int warp_rows = TileM / warp_tile;
+  constexpr int col_tiles_per_warp = 8;
+  constexpr int warp_col_groups = TileN / (col_tiles_per_warp * warp_tile);
+  constexpr int warps_per_block = warp_rows * warp_col_groups;
+  static_assert(TileM == 128 || TileM == 256);
+  static_assert(TileN == 128 || TileN == 256);
+  static_assert(warps_per_block == 8 || warps_per_block == 16);
+
+  extern __shared__ __align__(16) uint8_t shared_memory[];
+  int8_t *activation_tile = reinterpret_cast<int8_t *>(shared_memory);
+  int8_t *weight_tile = activation_tile + TileM * padded_k;
+  int32_t *accumulator_tile = reinterpret_cast<int32_t *>(
+      weight_tile + TileN * padded_k);
+  int8_t *code_tables = reinterpret_cast<int8_t *>(
+      accumulator_tile + warps_per_block * warp_tile * warp_tile);
+  uint8_t *packed_activation_stage =
+      reinterpret_cast<uint8_t *>(code_tables + 32);
+  uint8_t *packed_weight_stage =
+      packed_activation_stage + (AsyncPacked ? TileM * packed_tile_k : 0);
+
+  const int warp_id = threadIdx.x / warpSize;
+  const int lane = threadIdx.x & (warpSize - 1);
+  const int warp_row = warp_id % warp_rows;
+  const int warp_col_group = warp_id / warp_rows;
+  const int64_t block_row = int64_t(blockIdx.y) * TileM;
+  const int64_t block_col = int64_t(blockIdx.x) * TileN;
+  const int64_t packed_row_stride = in_features / 2;
+
+  if (threadIdx.x < 16) {
+    code_tables[threadIdx.x] = activation_codes[threadIdx.x];
+    code_tables[16 + threadIdx.x] = weight_codes[threadIdx.x];
+  }
+
+  wmma::fragment<wmma::accumulator, warp_tile, warp_tile, warp_tile, int>
+      accumulators[col_tiles_per_warp];
+#pragma unroll
+  for (int col_tile = 0; col_tile < col_tiles_per_warp; ++col_tile) {
+    wmma::fill_fragment(accumulators[col_tile], 0);
+  }
+  __syncthreads();
+
+  const bool use_async =
+      AsyncPacked && block_row + TileM <= rows &&
+      block_col + TileN <= out_features && out_features % 16 == 0;
+  if constexpr (AsyncPacked) {
+    if (use_async) {
+      constexpr int activation_vectors = TileM * packed_tile_k / 16;
+      for (int vector = threadIdx.x; vector < activation_vectors;
+           vector += blockDim.x) {
+        const int byte_offset = vector * 16;
+        const int local_row = byte_offset / packed_tile_k;
+        const int local_k_byte = byte_offset - local_row * packed_tile_k;
+        copy_async_16(
+            packed_activation_stage + byte_offset,
+            packed_activations +
+                (block_row + local_row) * packed_row_stride + local_k_byte);
+      }
+      constexpr int weight_vectors = packed_tile_k * TileN / 16;
+      for (int vector = threadIdx.x; vector < weight_vectors;
+           vector += blockDim.x) {
+        const int byte_offset = vector * 16;
+        if constexpr (KMajorWeight) {
+          const int local_k_byte = byte_offset / TileN;
+          const int local_col = byte_offset - local_k_byte * TileN;
+          copy_async_16(
+              packed_weight_stage + byte_offset,
+              packed_weight_indices + local_k_byte * out_features + block_col +
+                  local_col);
+        } else {
+          const int local_col = byte_offset / packed_tile_k;
+          const int local_k_byte = byte_offset - local_col * packed_tile_k;
+          copy_async_16(
+              packed_weight_stage + byte_offset,
+              packed_weight_indices +
+                  (block_col + local_col) * packed_row_stride + local_k_byte);
+        }
+      }
+      commit_async_copies();
+      wait_for_async_copies();
+      __syncthreads();
+    }
+  }
+
+  for (int64_t k_start = 0; k_start < in_features; k_start += tile_k) {
+    constexpr int activation_tasks = TileM * packed_tile_k;
+    for (int task = threadIdx.x; task < activation_tasks; task += blockDim.x) {
+      const int local_row = task / packed_tile_k;
+      const int local_k_byte = task - local_row * packed_tile_k;
+      const int64_t global_row = block_row + local_row;
+      uint8_t packed = 0;
+      if (use_async) {
+        packed = packed_activation_stage[task];
+      } else if (global_row < rows) {
+        packed = packed_activations[
+            global_row * packed_row_stride + k_start / 2 + local_k_byte];
+      }
+      const int destination = local_row * padded_k + local_k_byte * 2;
+      activation_tile[destination] = code_tables[packed & 15u];
+      activation_tile[destination + 1] = code_tables[packed >> 4];
+    }
+
+    constexpr int weight_tasks = packed_tile_k * TileN;
+    for (int task = threadIdx.x; task < weight_tasks; task += blockDim.x) {
+      const int local_k_byte = task / TileN;
+      const int local_col = task - local_k_byte * TileN;
+      const int64_t global_col = block_col + local_col;
+      uint8_t packed = 0;
+      if (use_async) {
+        if constexpr (KMajorWeight) {
+          packed = packed_weight_stage[task];
+        } else {
+          packed = packed_weight_stage[local_col * packed_tile_k + local_k_byte];
+        }
+      } else if (global_col < out_features) {
+        if constexpr (KMajorWeight) {
+          packed = packed_weight_indices[
+              (k_start / 2 + local_k_byte) * out_features + global_col];
+        } else {
+          packed = packed_weight_indices[
+              global_col * packed_row_stride + k_start / 2 + local_k_byte];
+        }
+      }
+      const int destination = local_col * padded_k + local_k_byte * 2;
+      weight_tile[destination] = code_tables[16 + (packed & 15u)];
+      weight_tile[destination + 1] = code_tables[16 + (packed >> 4)];
+    }
+    __syncthreads();
+
+    const bool has_next_tile = k_start + tile_k < in_features;
+    if constexpr (AsyncPacked) {
+      if (use_async && has_next_tile) {
+        const int64_t next_k_byte = (k_start + tile_k) / 2;
+        constexpr int activation_vectors = TileM * packed_tile_k / 16;
+        for (int vector = threadIdx.x; vector < activation_vectors;
+             vector += blockDim.x) {
+          const int byte_offset = vector * 16;
+          const int local_row = byte_offset / packed_tile_k;
+          const int local_k_byte = byte_offset - local_row * packed_tile_k;
+          copy_async_16(
+              packed_activation_stage + byte_offset,
+              packed_activations +
+                  (block_row + local_row) * packed_row_stride + next_k_byte +
+                  local_k_byte);
+        }
+        constexpr int weight_vectors = packed_tile_k * TileN / 16;
+        for (int vector = threadIdx.x; vector < weight_vectors;
+             vector += blockDim.x) {
+          const int byte_offset = vector * 16;
+          if constexpr (KMajorWeight) {
+            const int local_k_byte = byte_offset / TileN;
+            const int local_col = byte_offset - local_k_byte * TileN;
+            copy_async_16(
+                packed_weight_stage + byte_offset,
+                packed_weight_indices +
+                    (next_k_byte + local_k_byte) * out_features + block_col +
+                    local_col);
+          } else {
+            const int local_col = byte_offset / packed_tile_k;
+            const int local_k_byte = byte_offset - local_col * packed_tile_k;
+            copy_async_16(
+                packed_weight_stage + byte_offset,
+                packed_weight_indices +
+                    (block_col + local_col) * packed_row_stride + next_k_byte +
+                    local_k_byte);
+          }
+        }
+        commit_async_copies();
+      }
+    }
+
+#pragma unroll
+    for (int local_k = 0; local_k < tile_k; local_k += warp_tile) {
+      wmma::fragment<wmma::matrix_a, warp_tile, warp_tile, warp_tile, signed char,
+                     wmma::row_major>
+          lhs;
+      wmma::load_matrix_sync(
+          lhs,
+          reinterpret_cast<signed char const *>(
+              activation_tile + warp_row * warp_tile * padded_k + local_k),
+          padded_k);
+#pragma unroll
+      for (int col_tile = 0; col_tile < col_tiles_per_warp; ++col_tile) {
+        wmma::fragment<wmma::matrix_b, warp_tile, warp_tile, warp_tile, signed char,
+                       wmma::col_major>
+            rhs;
+        wmma::load_matrix_sync(
+            rhs,
+            reinterpret_cast<signed char const *>(
+                weight_tile +
+                (warp_col_group * col_tiles_per_warp + col_tile) * warp_tile *
+                    padded_k +
+                local_k),
+            padded_k);
+        wmma::mma_sync(
+            accumulators[col_tile], lhs, rhs, accumulators[col_tile]);
+      }
+    }
+    __syncthreads();
+    if constexpr (AsyncPacked) {
+      if (use_async && has_next_tile) {
+        wait_for_async_copies();
+        __syncthreads();
+      }
+    }
+  }
+
+  int32_t *warp_accumulator =
+      accumulator_tile + warp_id * warp_tile * warp_tile;
+  const float surrogate_scale = activation_scale * weight_scale;
+#pragma unroll
+  for (int col_tile = 0; col_tile < col_tiles_per_warp; ++col_tile) {
+    wmma::store_matrix_sync(
+        warp_accumulator,
+        accumulators[col_tile],
+        warp_tile,
+        wmma::mem_row_major);
+    __syncwarp();
+    for (int offset = lane; offset < warp_tile * warp_tile; offset += warpSize) {
+      const int local_row = offset / warp_tile;
+      const int local_col = offset - local_row * warp_tile;
+      const int64_t global_row = block_row + warp_row * warp_tile + local_row;
+      const int64_t global_col =
+          block_col +
+          (warp_col_group * col_tiles_per_warp + col_tile) * warp_tile +
+          local_col;
+      if (global_row < rows && global_col < out_features) {
+        float value = static_cast<float>(warp_accumulator[offset]);
+        value *= token_norms[global_row] *
+            static_cast<float>(row_norms[global_col]) * surrogate_scale;
+        if (has_bias) {
+          value += static_cast<float>(bias[global_col]);
+        }
+        out[global_row * out_features + global_col] =
+            static_cast<storage_t>(value);
+      }
+    }
+    __syncwarp();
+  }
+#endif  // __CUDA_ARCH__ >= 800
+}
+
 template <int Bits>
 __global__ void orbitquant_packed_matmul_wmma_bf16_kernel(
     c10::BFloat16 *__restrict__ out,
     c10::BFloat16 const *__restrict__ x,
     uint8_t const *__restrict__ packed_weight_indices,
-    float const *__restrict__ row_norms,
+    c10::BFloat16 const *__restrict__ row_norms,
     float const *__restrict__ centroids,
-    float const *__restrict__ bias,
+    c10::BFloat16 const *__restrict__ bias,
     bool has_bias,
     int64_t rows,
     int64_t out_features,
     int64_t in_features) {
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
   constexpr int tile = 16;
   constexpr int col_tiles = 4;
   constexpr int warps_per_block = 8;
@@ -102,7 +1052,7 @@ __global__ void orbitquant_packed_matmul_wmma_bf16_kernel(
           const int64_t value_offset = global_col * in_features + global_k;
           const uint32_t index = unpack_lowbit_index_const<Bits>(
               packed_weight_indices, value_offset);
-          value = row_norms[global_col] * centroids[index];
+          value = static_cast<float>(row_norms[global_col]) * centroids[index];
         }
         warp_w_tile[offset] = __float2bfloat16(value);
       }
@@ -131,12 +1081,13 @@ __global__ void orbitquant_packed_matmul_wmma_bf16_kernel(
       if (global_row < rows && global_col < out_features) {
         float value = warp_acc_tile[offset];
         if (has_bias) {
-          value += bias[global_col];
+          value += static_cast<float>(bias[global_col]);
         }
         out[global_row * out_features + global_col] = static_cast<c10::BFloat16>(value);
       }
     }
   }
+#endif  // __CUDA_ARCH__ >= 800
 }
 
 template <int Bits>
@@ -144,9 +1095,9 @@ __global__ void orbitquant_packed_matmul_wmma_half_kernel(
     c10::Half *__restrict__ out,
     c10::Half const *__restrict__ x,
     uint8_t const *__restrict__ packed_weight_indices,
-    float const *__restrict__ row_norms,
+    c10::BFloat16 const *__restrict__ row_norms,
     float const *__restrict__ centroids,
-    float const *__restrict__ bias,
+    c10::Half const *__restrict__ bias,
     bool has_bias,
     int64_t rows,
     int64_t out_features,
@@ -197,7 +1148,7 @@ __global__ void orbitquant_packed_matmul_wmma_half_kernel(
           const int64_t value_offset = global_col * in_features + global_k;
           const uint32_t index = unpack_lowbit_index_const<Bits>(
               packed_weight_indices, value_offset);
-          value = row_norms[global_col] * centroids[index];
+          value = static_cast<float>(row_norms[global_col]) * centroids[index];
         }
         warp_w_tile[offset] = __float2half(value);
       }
@@ -225,7 +1176,7 @@ __global__ void orbitquant_packed_matmul_wmma_half_kernel(
       if (global_row < rows && global_col < out_features) {
         float value = warp_acc_tile[offset];
         if (has_bias) {
-          value += bias[global_col];
+          value += static_cast<float>(bias[global_col]);
         }
         out[global_row * out_features + global_col] = static_cast<c10::Half>(value);
       }
@@ -233,14 +1184,73 @@ __global__ void orbitquant_packed_matmul_wmma_half_kernel(
   }
 }
 
+template <typename scalar_t>
+__global__ void orbitquant_packed_matmul_small_rows_kernel(
+    scalar_t *__restrict__ out,
+    scalar_t const *__restrict__ x,
+    uint8_t const *__restrict__ packed_weight_indices,
+    c10::BFloat16 const *__restrict__ row_norms,
+    float const *__restrict__ centroids,
+    scalar_t const *__restrict__ bias,
+    bool has_bias,
+    int64_t rows,
+    int64_t out_features,
+    int64_t in_features,
+    int64_t bits) {
+  constexpr int channels_per_warp = 4;
+  const int lane = threadIdx.x;
+  const int64_t row = blockIdx.y;
+  const int64_t col_start = int64_t(blockIdx.x) * channels_per_warp;
+  const uint32_t mask = (1u << bits) - 1u;
+  float accumulators[channels_per_warp] = {};
+  float norms[channels_per_warp];
+
+#pragma unroll
+  for (int col_offset = 0; col_offset < channels_per_warp; ++col_offset) {
+    const int64_t col = col_start + col_offset;
+    norms[col_offset] =
+        col < out_features ? static_cast<float>(row_norms[col]) : 0.0f;
+  }
+
+  for (int64_t k = lane; k < in_features; k += warpSize) {
+    const float x_value = static_cast<float>(x[row * in_features + k]);
+#pragma unroll
+    for (int col_offset = 0; col_offset < channels_per_warp; ++col_offset) {
+      const int64_t col = col_start + col_offset;
+      if (col < out_features) {
+        const int64_t value_offset = col * in_features + k;
+        const uint32_t index =
+            unpack_lowbit_index(packed_weight_indices, value_offset, bits, mask);
+        accumulators[col_offset] += x_value * norms[col_offset] * centroids[index];
+      }
+    }
+  }
+
+#pragma unroll
+  for (int col_offset = 0; col_offset < channels_per_warp; ++col_offset) {
+#pragma unroll
+    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
+      accumulators[col_offset] +=
+          __shfl_down_sync(0xffffffffu, accumulators[col_offset], offset);
+    }
+    const int64_t col = col_start + col_offset;
+    if (lane == 0 && col < out_features) {
+      const float value =
+          accumulators[col_offset] +
+          (has_bias ? static_cast<float>(bias[col]) : 0.0f);
+      out[row * out_features + col] = static_cast<scalar_t>(value);
+    }
+  }
+}
+
 template <typename scalar_t>
 __global__ void orbitquant_packed_matmul_tiled_kernel(
     scalar_t *__restrict__ out,
     scalar_t const *__restrict__ x,
     uint8_t const *__restrict__ packed_weight_indices,
-    float const *__restrict__ row_norms,
+    c10::BFloat16 const *__restrict__ row_norms,
     float const *__restrict__ centroids,
-    float const *__restrict__ bias,
+    scalar_t const *__restrict__ bias,
     bool has_bias,
     int64_t rows,
     int64_t out_features,
@@ -260,7 +1270,8 @@ __global__ void orbitquant_packed_matmul_tiled_kernel(
 
   const uint32_t mask = (1u << bits) - 1u;
   const bool output_valid = row < rows && col < out_features;
-  float acc = output_valid && has_bias ? bias[col] : 0.0f;
+  float acc =
+      output_valid && has_bias ? static_cast<float>(bias[col]) : 0.0f;
 
   for (int64_t k_start = 0; k_start < in_features; k_start += block_k) {
     const int64_t x_tile_values = blockDim.y * block_k;
@@ -287,7 +1298,7 @@ __global__ void orbitquant_packed_matmul_tiled_kernel(
         const int64_t value_offset = global_col * in_features + global_k;
         const uint32_t index =
             unpack_lowbit_index(packed_weight_indices, value_offset, bits, mask);
-        value = row_norms[global_col] * centroids[index];
+        value = static_cast<float>(row_norms[global_col]) * centroids[index];
       }
       w_tile[offset] = value;
     }
@@ -331,7 +1342,9 @@ void matmul_packed_weight(
   TORCH_CHECK(row_norms.is_contiguous(), "row norms must be contiguous");
   TORCH_CHECK(centroids.is_contiguous(), "centroids must be contiguous");
   TORCH_CHECK(packed_weight_indices.scalar_type() == torch::kUInt8, "packed weights must be uint8");
-  TORCH_CHECK(row_norms.scalar_type() == torch::kFloat, "row_norms must be float32");
+  TORCH_CHECK(
+      row_norms.scalar_type() == torch::kBFloat16,
+      "CUDA row_norms must be bfloat16");
   TORCH_CHECK(centroids.scalar_type() == torch::kFloat, "centroids must be float32");
   TORCH_CHECK(x.dim() == 2, "x must be rank 2");
   TORCH_CHECK(out.dim() == 2, "out must be rank 2");
@@ -348,7 +1361,7 @@ void matmul_packed_weight(
   if (has_bias) {
     TORCH_CHECK(bias.device().is_cuda(), "bias must be a CUDA tensor");
     TORCH_CHECK(bias.is_contiguous(), "bias must be contiguous");
-    TORCH_CHECK(bias.scalar_type() == torch::kFloat, "bias must be float32");
+    TORCH_CHECK(bias.scalar_type() == x.scalar_type(), "CUDA bias dtype must match x");
     TORCH_CHECK(bias.numel() == out_features, "bias must match out_features");
   }
   if (x.numel() == 0 || out_features == 0) {
@@ -356,8 +1369,9 @@ void matmul_packed_weight(
   }
 
   const int threads_n = static_cast<int>(std::min<int64_t>(std::max<int64_t>(block_n, 1), 64));
-  const int threads_m =
-      static_cast<int>(std::min<int64_t>(std::max<int64_t>(block_m, 1), 1024 / threads_n));
+  const int threads_m = static_cast<int>(std::min<int64_t>(
+      x.size(0),
+      std::min<int64_t>(std::max<int64_t>(block_m, 1), 1024 / threads_n)));
   const int tile_k = static_cast<int>(std::min<int64_t>(std::max<int64_t>(block_k, 1), 128));
   const dim3 block(threads_n, threads_m);
   const dim3 grid(
@@ -367,8 +1381,136 @@ void matmul_packed_weight(
       sizeof(float);
   const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaDeviceProp *mma_properties = at::cuda::getCurrentDeviceProperties();
+
+  if (x.size(0) <= 8) {
+    constexpr int channels_per_warp = 4;
+    const dim3 small_rows_block(32);
+    const dim3 small_rows_grid(
+        (out_features + channels_per_warp - 1) / channels_per_warp,
+        x.size(0));
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::kHalf, at::kBFloat16, x.scalar_type(),
+        "orbitquant_packed_matmul_cuda_small_rows", [&] {
+          orbitquant_packed_matmul_small_rows_kernel<scalar_t>
+              <<<small_rows_grid, small_rows_block, 0, stream>>>(
+                  out.data_ptr<scalar_t>(),
+                  x.data_ptr<scalar_t>(),
+                  packed_weight_indices.data_ptr<uint8_t>(),
+                  row_norms.data_ptr<c10::BFloat16>(),
+                  centroids.data_ptr<float>(),
+                  has_bias ? bias.data_ptr<scalar_t>() : nullptr,
+                  has_bias,
+                  x.size(0),
+                  out_features,
+                  in_features,
+                  bits);
+        });
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+    return;
+  }
 
-  if (x.scalar_type() == at::kBFloat16) {
+  if (x.scalar_type() == at::kBFloat16 && x.size(0) >= 9 &&
+      mma_properties->major >= 8) {
+    if (in_features % 64 == 0 &&
+        (bits == 2 || bits == 3 || bits == 4 || bits == 6)) {
+      constexpr int mma_tile_m = 128;
+      constexpr int mma_tile_n = 128;
+      const dim3 mma_block(256);
+      const dim3 mma_grid(
+          (out_features + mma_tile_n - 1) / mma_tile_n,
+          (x.size(0) + mma_tile_m - 1) / mma_tile_m);
+
+#define ORBITQUANT_LAUNCH_MMA64_PIPELINED(STORAGE_TYPE, MMA_TYPE, BITS_VALUE)   \
+  do {                                                                          \
+    constexpr int kSegStride = 64 * (BITS_VALUE) / 8;                           \
+    const int pipelined_shared_bytes = static_cast<int>(                        \
+        2 * 128 * 72 * sizeof(MMA_TYPE) + 128 * 72 * sizeof(MMA_TYPE) +         \
+        8 * 16 * 16 * sizeof(float) + 2 * 128 * kSegStride);                    \
+    if (pipelined_shared_bytes > mma_properties->sharedMemPerBlock) {           \
+      C10_CUDA_CHECK(cudaFuncSetAttribute(                                      \
+          orbitquant_packed_matmul_mma64_pipelined_kernel<STORAGE_TYPE,         \
+                                                          MMA_TYPE,             \
+                                                          BITS_VALUE>,          \
+          cudaFuncAttributeMaxDynamicSharedMemorySize,                          \
+          pipelined_shared_bytes));                                             \
+    }                                                                           \
+    orbitquant_packed_matmul_mma64_pipelined_kernel<STORAGE_TYPE, MMA_TYPE,     \
+                                                    BITS_VALUE>                 \
+        <<<mma_grid, mma_block, pipelined_shared_bytes, stream>>>(              \
+            reinterpret_cast<STORAGE_TYPE *>(out.data_ptr()),                   \
+            reinterpret_cast<STORAGE_TYPE const *>(x.data_ptr()),               \
+            packed_weight_indices.data_ptr<uint8_t>(),                          \
+            row_norms.data_ptr<c10::BFloat16>(),                                \
+            centroids.data_ptr<float>(),                                        \
+            has_bias ? bias.data_ptr<STORAGE_TYPE>() : nullptr,                  \
+            has_bias,                                                            \
+            x.size(0),                                                           \
+            out_features,                                                        \
+            in_features);                                                        \
+  } while (0)
+// The cp.async pipeline wins when the launch is latency-bound: with cold L2
+// (each layer's weights are evicted between calls in a real model) it is
+// 1.35-1.46x faster for W2/W4/W6 on RTX 4060 Ti, A40, and RTX 4090 whenever
+// the grid fits in one wave (blocks <= SM count), and it regresses up to 15%
+// once the grid oversubscribes the device. Gate on grid size; measured
+// 2026-07 across the three architectures above.
+// ORBITQUANT_MMA64_FORCE_PIPELINE=1 bypasses the grid gate,
+// ORBITQUANT_MMA64_DISABLE_PIPELINE=1 forces the legacy kernel everywhere.
+#define ORBITQUANT_MMA64_USE_PIPELINE(BITS_VALUE)                               \
+  ((orbitquant_mma64_pipeline_forced() ||                                       \
+    static_cast<int64_t>(mma_grid.x) * mma_grid.y <=                            \
+        mma_properties->multiProcessorCount) &&                                 \
+   !orbitquant_mma64_pipeline_disabled() &&                                     \
+   mma_properties->major >= 8 &&                                                \
+   (2 * 128 * 72 * 2 + 128 * 72 * 2 + 8 * 16 * 16 * 4 +                        \
+    2 * 128 * (64 * (BITS_VALUE) / 8)) <=                                       \
+       static_cast<int>(mma_properties->sharedMemPerBlockOptin))
+
+#define ORBITQUANT_LAUNCH_MMA64_BF16(BITS_VALUE)                               \
+  orbitquant_packed_matmul_mma64_kernel<c10::BFloat16, __nv_bfloat16,          \
+                                         BITS_VALUE><<<mma_grid, mma_block, 0,  \
+                                                       stream>>>(               \
+      reinterpret_cast<c10::BFloat16 *>(out.data_ptr()),                         \
+      reinterpret_cast<c10::BFloat16 const *>(x.data_ptr()),                     \
+      packed_weight_indices.data_ptr<uint8_t>(),                                 \
+      row_norms.data_ptr<c10::BFloat16>(),                                       \
+      centroids.data_ptr<float>(),                                               \
+      has_bias ? bias.data_ptr<c10::BFloat16>() : nullptr,                       \
+      has_bias,                                                                  \
+      x.size(0),                                                                 \
+      out_features,                                                              \
+      in_features)
+      switch (bits) {
+        case 2:
+          if (ORBITQUANT_MMA64_USE_PIPELINE(2)) {
+            ORBITQUANT_LAUNCH_MMA64_PIPELINED(c10::BFloat16, __nv_bfloat16, 2);
+          } else {
+            ORBITQUANT_LAUNCH_MMA64_BF16(2);
+          }
+          break;
+        case 3:
+          ORBITQUANT_LAUNCH_MMA64_BF16(3);
+          break;
+        case 4:
+          if (ORBITQUANT_MMA64_USE_PIPELINE(4)) {
+            ORBITQUANT_LAUNCH_MMA64_PIPELINED(c10::BFloat16, __nv_bfloat16, 4);
+          } else {
+            ORBITQUANT_LAUNCH_MMA64_BF16(4);
+          }
+          break;
+        case 6:
+          if (ORBITQUANT_MMA64_USE_PIPELINE(6)) {
+            ORBITQUANT_LAUNCH_MMA64_PIPELINED(c10::BFloat16, __nv_bfloat16, 6);
+          } else {
+            ORBITQUANT_LAUNCH_MMA64_BF16(6);
+          }
+          break;
+      }
+#undef ORBITQUANT_LAUNCH_MMA64_BF16
+      C10_CUDA_KERNEL_LAUNCH_CHECK();
+      return;
+    }
     constexpr int tile = 16;
     constexpr int col_tiles = 4;
     constexpr int rows_per_block = tile * 8;
@@ -382,9 +1524,9 @@ void matmul_packed_weight(
       reinterpret_cast<c10::BFloat16 *>(out.data_ptr()),                                \
       reinterpret_cast<c10::BFloat16 const *>(x.data_ptr()),                            \
       packed_weight_indices.data_ptr<uint8_t>(),                                        \
-      row_norms.data_ptr<float>(),                                                      \
+      row_norms.data_ptr<c10::BFloat16>(),                                              \
       centroids.data_ptr<float>(),                                                      \
-      has_bias ? bias.data_ptr<float>() : nullptr,                                      \
+      has_bias ? bias.data_ptr<c10::BFloat16>() : nullptr,                              \
       has_bias,                                                                         \
       x.size(0),                                                                        \
       out_features,                                                                     \
@@ -420,7 +1562,61 @@ void matmul_packed_weight(
     return;
   }
 
-  if (x.scalar_type() == at::kHalf) {
+  if (x.scalar_type() == at::kHalf && x.size(0) >= 9 &&
+      mma_properties->major >= 8) {
+    if (in_features % 64 == 0 &&
+        (bits == 2 || bits == 3 || bits == 4 || bits == 6)) {
+      constexpr int mma_tile_m = 128;
+      constexpr int mma_tile_n = 128;
+      const dim3 mma_block(256);
+      const dim3 mma_grid(
+          (out_features + mma_tile_n - 1) / mma_tile_n,
+          (x.size(0) + mma_tile_m - 1) / mma_tile_m);
+#define ORBITQUANT_LAUNCH_MMA64_HALF(BITS_VALUE)                                \
+  orbitquant_packed_matmul_mma64_kernel<c10::Half, half, BITS_VALUE>            \
+      <<<mma_grid, mma_block, 0, stream>>>(                                     \
+          reinterpret_cast<c10::Half *>(out.data_ptr()),                         \
+          reinterpret_cast<c10::Half const *>(x.data_ptr()),                     \
+          packed_weight_indices.data_ptr<uint8_t>(),                             \
+          row_norms.data_ptr<c10::BFloat16>(),                                   \
+          centroids.data_ptr<float>(),                                           \
+          has_bias ? bias.data_ptr<c10::Half>() : nullptr,                       \
+          has_bias,                                                              \
+          x.size(0),                                                             \
+          out_features,                                                          \
+          in_features)
+      switch (bits) {
+        case 2:
+          if (ORBITQUANT_MMA64_USE_PIPELINE(2)) {
+            ORBITQUANT_LAUNCH_MMA64_PIPELINED(c10::Half, half, 2);
+          } else {
+            ORBITQUANT_LAUNCH_MMA64_HALF(2);
+          }
+          break;
+        case 3:
+          ORBITQUANT_LAUNCH_MMA64_HALF(3);
+          break;
+        case 4:
+          if (ORBITQUANT_MMA64_USE_PIPELINE(4)) {
+            ORBITQUANT_LAUNCH_MMA64_PIPELINED(c10::Half, half, 4);
+          } else {
+            ORBITQUANT_LAUNCH_MMA64_HALF(4);
+          }
+          break;
+        case 6:
+          if (ORBITQUANT_MMA64_USE_PIPELINE(6)) {
+            ORBITQUANT_LAUNCH_MMA64_PIPELINED(c10::Half, half, 6);
+          } else {
+            ORBITQUANT_LAUNCH_MMA64_HALF(6);
+          }
+          break;
+      }
+#undef ORBITQUANT_LAUNCH_MMA64_HALF
+#undef ORBITQUANT_MMA64_USE_PIPELINE
+#undef ORBITQUANT_LAUNCH_MMA64_PIPELINED
+      C10_CUDA_KERNEL_LAUNCH_CHECK();
+      return;
+    }
     constexpr int tile = 16;
     constexpr int col_tiles = 4;
     constexpr int rows_per_block = tile * 8;
@@ -434,9 +1630,9 @@ void matmul_packed_weight(
       reinterpret_cast<c10::Half *>(out.data_ptr()),                                    \
       reinterpret_cast<c10::Half const *>(x.data_ptr()),                                \
       packed_weight_indices.data_ptr<uint8_t>(),                                        \
-      row_norms.data_ptr<float>(),                                                      \
+      row_norms.data_ptr<c10::BFloat16>(),                                              \
       centroids.data_ptr<float>(),                                                      \
-      has_bias ? bias.data_ptr<float>() : nullptr,                                      \
+      has_bias ? bias.data_ptr<c10::Half>() : nullptr,                                  \
       has_bias,                                                                         \
       x.size(0),                                                                        \
       out_features,                                                                     \
@@ -478,9 +1674,9 @@ void matmul_packed_weight(
             out.data_ptr<scalar_t>(),
             x.data_ptr<scalar_t>(),
             packed_weight_indices.data_ptr<uint8_t>(),
-            row_norms.data_ptr<float>(),
+            row_norms.data_ptr<c10::BFloat16>(),
             centroids.data_ptr<float>(),
-            has_bias ? bias.data_ptr<float>() : nullptr,
+            has_bias ? bias.data_ptr<scalar_t>() : nullptr,
             has_bias,
             x.size(0),
             out_features,
@@ -490,3 +1686,454 @@ void matmul_packed_weight(
       });
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
+
+void matmul_packed_w4a4_int8(
+    torch::Tensor &out,
+    torch::Tensor const &packed_activations,
+    torch::Tensor const &packed_weight_indices,
+    torch::Tensor const &token_norms,
+    torch::Tensor const &row_norms,
+    torch::Tensor const &activation_codes,
+    torch::Tensor const &weight_codes,
+    torch::Tensor const &bias,
+    bool has_bias,
+    double activation_scale,
+    double weight_scale,
+    int64_t out_features,
+    int64_t in_features,
+    int64_t tile_m,
+    int64_t tile_n,
+    bool async_packed,
+    bool weight_k_major) {
+  TORCH_CHECK(out.device().is_cuda(), "out must be a CUDA tensor");
+  TORCH_CHECK(
+      packed_activations.device().is_cuda(),
+      "packed activations must be a CUDA tensor");
+  TORCH_CHECK(
+      packed_weight_indices.device().is_cuda(),
+      "packed weights must be a CUDA tensor");
+  TORCH_CHECK(token_norms.device().is_cuda(), "token norms must be a CUDA tensor");
+  TORCH_CHECK(row_norms.device().is_cuda(), "row norms must be a CUDA tensor");
+  TORCH_CHECK(
+      activation_codes.device().is_cuda(),
+      "activation surrogate codes must be a CUDA tensor");
+  TORCH_CHECK(
+      weight_codes.device().is_cuda(),
+      "weight surrogate codes must be a CUDA tensor");
+  TORCH_CHECK(out.is_contiguous(), "out must be contiguous");
+  TORCH_CHECK(
+      packed_activations.is_contiguous(), "packed activations must be contiguous");
+  TORCH_CHECK(
+      packed_weight_indices.is_contiguous(), "packed weights must be contiguous");
+  TORCH_CHECK(token_norms.is_contiguous(), "token norms must be contiguous");
+  TORCH_CHECK(row_norms.is_contiguous(), "row norms must be contiguous");
+  TORCH_CHECK(
+      activation_codes.is_contiguous(), "activation surrogate codes must be contiguous");
+  TORCH_CHECK(
+      weight_codes.is_contiguous(), "weight surrogate codes must be contiguous");
+  TORCH_CHECK(
+      packed_activations.scalar_type() == torch::kUInt8,
+      "packed activations must be uint8");
+  TORCH_CHECK(
+      packed_weight_indices.scalar_type() == torch::kUInt8,
+      "packed weights must be uint8");
+  TORCH_CHECK(token_norms.scalar_type() == torch::kFloat, "token norms must be float32");
+  TORCH_CHECK(row_norms.scalar_type() == torch::kBFloat16, "row norms must be bfloat16");
+  TORCH_CHECK(
+      activation_codes.scalar_type() == torch::kChar,
+      "activation surrogate codes must be int8");
+  TORCH_CHECK(
+      weight_codes.scalar_type() == torch::kChar,
+      "weight surrogate codes must be int8");
+  TORCH_CHECK(
+      out.scalar_type() == torch::kBFloat16 || out.scalar_type() == torch::kHalf,
+      "packed W4A4 INT8 output must be bfloat16 or float16");
+  TORCH_CHECK(packed_activations.dim() == 2, "packed activations must be rank 2");
+  TORCH_CHECK(out.dim() == 2, "out must be rank 2");
+  TORCH_CHECK(
+      in_features > 0 && in_features % 64 == 0,
+      "in_features must be positive and divisible by 64");
+  TORCH_CHECK(out_features >= 0, "out_features must be non-negative");
+  TORCH_CHECK(
+      (tile_m == 128 && tile_n == 128) ||
+          (tile_m == 256 && tile_n == 128) ||
+          (tile_m == 128 && tile_n == 256),
+      "packed W4A4 INT8 tile must be 128x128, 256x128, or 128x256");
+  TORCH_CHECK(
+      packed_activations.size(1) == in_features / 2,
+      "packed activations have an unexpected input dimension");
+  const int64_t rows = packed_activations.size(0);
+  TORCH_CHECK(out.size(0) == rows, "out has an unexpected row count");
+  TORCH_CHECK(out.size(1) == out_features, "out has an unexpected output dimension");
+  TORCH_CHECK(token_norms.numel() == rows, "token norms must match rows");
+  TORCH_CHECK(row_norms.numel() == out_features, "row norms must match out_features");
+  TORCH_CHECK(activation_codes.numel() == 16, "activation codes must contain 16 values");
+  TORCH_CHECK(weight_codes.numel() == 16, "weight codes must contain 16 values");
+  TORCH_CHECK(
+      packed_weight_indices.numel() == out_features * (in_features / 2),
+      "K-major packed weights have an unexpected size");
+  if (has_bias) {
+    TORCH_CHECK(bias.device().is_cuda(), "bias must be a CUDA tensor");
+    TORCH_CHECK(bias.is_contiguous(), "bias must be contiguous");
+    TORCH_CHECK(bias.scalar_type() == out.scalar_type(), "bias dtype must match out");
+    TORCH_CHECK(bias.numel() == out_features, "bias must match out_features");
+  }
+  if (rows == 0 || out_features == 0) {
+    return;
+  }
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(packed_activations));
+  const cudaDeviceProp *properties = at::cuda::getCurrentDeviceProperties();
+  TORCH_CHECK(
+      properties->major > 7 || (properties->major == 7 && properties->minor >= 5),
+      "packed W4A4 INT8 Tensor Core matmul requires compute capability 7.5+");
+  const int warp_count = static_cast<int>((tile_m / 16) * (tile_n / 128));
+  const dim3 block(warp_count * 32);
+  const dim3 grid(
+      (out_features + tile_n - 1) / tile_n,
+      (rows + tile_m - 1) / tile_m);
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const int shared_bytes = static_cast<int>(
+      tile_m * 80 + tile_n * 80 + warp_count * 16 * 16 * sizeof(int32_t) + 32 +
+      (async_packed ? (tile_m + tile_n) * 32 : 0));
+  TORCH_CHECK(
+      shared_bytes <= properties->sharedMemPerBlockOptin,
+      "packed W4A4 INT8 tile requires ",
+      shared_bytes,
+      " bytes of shared memory, but the device supports ",
+      properties->sharedMemPerBlockOptin);
+
+#define ORBITQUANT_LAUNCH_PACKED_W4A4_INT8(                                      \
+    STORAGE_TYPE, TILE_M, TILE_N, ASYNC_PACKED, K_MAJOR_WEIGHT)                  \
+  if (shared_bytes > properties->sharedMemPerBlock) {                            \
+    C10_CUDA_CHECK(cudaFuncSetAttribute(                                         \
+        orbitquant_packed_w4a4_int8_mma_kernel<                                  \
+            STORAGE_TYPE, TILE_M, TILE_N, ASYNC_PACKED, K_MAJOR_WEIGHT>,         \
+        cudaFuncAttributeMaxDynamicSharedMemorySize,                             \
+        shared_bytes));                                                          \
+  }                                                                              \
+  orbitquant_packed_w4a4_int8_mma_kernel<                                        \
+      STORAGE_TYPE, TILE_M, TILE_N, ASYNC_PACKED, K_MAJOR_WEIGHT>                \
+      <<<grid, block, shared_bytes, stream>>>(                                    \
+      reinterpret_cast<STORAGE_TYPE *>(out.data_ptr()),                          \
+      packed_activations.data_ptr<uint8_t>(),                                    \
+      packed_weight_indices.data_ptr<uint8_t>(),                                 \
+      token_norms.data_ptr<float>(),                                             \
+      row_norms.data_ptr<c10::BFloat16>(),                                       \
+      activation_codes.data_ptr<int8_t>(),                                       \
+      weight_codes.data_ptr<int8_t>(),                                           \
+      has_bias ? bias.data_ptr<STORAGE_TYPE>() : nullptr,                         \
+      has_bias,                                                                  \
+      static_cast<float>(activation_scale),                                      \
+      static_cast<float>(weight_scale),                                          \
+      rows,                                                                       \
+      out_features,                                                               \
+      in_features)
+#define ORBITQUANT_DISPATCH_PACKED_W4A4_TILE(                                    \
+    STORAGE_TYPE, ASYNC_PACKED, K_MAJOR_WEIGHT)                                  \
+  if (tile_m == 256) {                                                           \
+    ORBITQUANT_LAUNCH_PACKED_W4A4_INT8(                                          \
+        STORAGE_TYPE, 256, 128, ASYNC_PACKED, K_MAJOR_WEIGHT);                   \
+  } else if (tile_n == 256) {                                                    \
+    ORBITQUANT_LAUNCH_PACKED_W4A4_INT8(                                          \
+        STORAGE_TYPE, 128, 256, ASYNC_PACKED, K_MAJOR_WEIGHT);                   \
+  } else {                                                                       \
+    ORBITQUANT_LAUNCH_PACKED_W4A4_INT8(                                          \
+        STORAGE_TYPE, 128, 128, ASYNC_PACKED, K_MAJOR_WEIGHT);                   \
+  }
+#define ORBITQUANT_DISPATCH_PACKED_W4A4_LAYOUT(STORAGE_TYPE, ASYNC_PACKED)       \
+  if (weight_k_major) {                                                          \
+    ORBITQUANT_DISPATCH_PACKED_W4A4_TILE(STORAGE_TYPE, ASYNC_PACKED, true);      \
+  } else {                                                                       \
+    ORBITQUANT_DISPATCH_PACKED_W4A4_TILE(STORAGE_TYPE, ASYNC_PACKED, false);     \
+  }
+  if (out.scalar_type() == torch::kBFloat16) {
+    if (async_packed) {
+      ORBITQUANT_DISPATCH_PACKED_W4A4_LAYOUT(c10::BFloat16, true);
+    } else {
+      ORBITQUANT_DISPATCH_PACKED_W4A4_LAYOUT(c10::BFloat16, false);
+    }
+  } else {
+    if (async_packed) {
+      ORBITQUANT_DISPATCH_PACKED_W4A4_LAYOUT(c10::Half, true);
+    } else {
+      ORBITQUANT_DISPATCH_PACKED_W4A4_LAYOUT(c10::Half, false);
+    }
+  }
+#undef ORBITQUANT_DISPATCH_PACKED_W4A4_LAYOUT
+#undef ORBITQUANT_DISPATCH_PACKED_W4A4_TILE
+#undef ORBITQUANT_LAUNCH_PACKED_W4A4_INT8
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+void quantize_activations_packed_w4(
+    torch::Tensor &packed_out,
+    torch::Tensor &norms_out,
+    torch::Tensor const &x,
+    torch::Tensor const &permutation,
+    torch::Tensor const &signs,
+    torch::Tensor const &boundaries,
+    double eps,
+    double inv_sqrt_block,
+    int64_t threads) {
+  TORCH_CHECK(packed_out.device().is_cuda(), "packed_out must be a CUDA tensor");
+  TORCH_CHECK(norms_out.device().is_cuda(), "norms_out must be a CUDA tensor");
+  TORCH_CHECK(x.device().is_cuda(), "x must be a CUDA tensor");
+  TORCH_CHECK(permutation.device().is_cuda(), "permutation must be a CUDA tensor");
+  TORCH_CHECK(signs.device().is_cuda(), "signs must be a CUDA tensor");
+  TORCH_CHECK(boundaries.device().is_cuda(), "boundaries must be a CUDA tensor");
+  TORCH_CHECK(packed_out.is_contiguous(), "packed_out must be contiguous");
+  TORCH_CHECK(norms_out.is_contiguous(), "norms_out must be contiguous");
+  TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
+  TORCH_CHECK(permutation.is_contiguous(), "permutation must be contiguous");
+  TORCH_CHECK(signs.is_contiguous(), "signs must be contiguous");
+  TORCH_CHECK(boundaries.is_contiguous(), "boundaries must be contiguous");
+  TORCH_CHECK(
+      packed_out.scalar_type() == torch::kUInt8,
+      "packed_out must be uint8");
+  TORCH_CHECK(norms_out.scalar_type() == torch::kFloat, "norms_out must be float32");
+  TORCH_CHECK(
+      x.scalar_type() == torch::kBFloat16 || x.scalar_type() == torch::kHalf,
+      "x must be bfloat16 or float16");
+  TORCH_CHECK(
+      permutation.scalar_type() == torch::kLong ||
+          permutation.scalar_type() == torch::kInt,
+      "permutation must be int32 or int64");
+  TORCH_CHECK(signs.scalar_type() == torch::kChar, "signs must be int8");
+  TORCH_CHECK(boundaries.scalar_type() == torch::kFloat, "boundaries must be float32");
+  TORCH_CHECK(x.dim() == 2, "x must be rank 2");
+  TORCH_CHECK(packed_out.dim() == 2, "packed_out must be rank 2");
+  TORCH_CHECK(norms_out.dim() == 1, "norms_out must be rank 1");
+  const int64_t rows = x.size(0);
+  const int64_t dim = x.size(1);
+  TORCH_CHECK(
+      dim == 512 || dim == 1024 || dim == 2048 || dim == 4096 ||
+          dim == 8192 || dim == 16384,
+      "native packed W4 activation quantization supports dimensions "
+      "512, 1024, 2048, 4096, 8192, and 16384");
+  TORCH_CHECK(
+      threads == 128 || threads == 256 || threads == 512,
+      "native packed W4 activation quantization threads must be 128, 256, or 512");
+  TORCH_CHECK(
+      packed_out.size(0) == rows && packed_out.size(1) == dim / 2,
+      "packed_out has an unexpected shape");
+  TORCH_CHECK(norms_out.numel() == rows, "norms_out must match rows");
+  TORCH_CHECK(permutation.numel() == dim, "permutation must match the input dimension");
+  TORCH_CHECK(signs.numel() == dim, "signs must match the input dimension");
+  TORCH_CHECK(boundaries.numel() == 15, "boundaries must contain 15 values");
+  if (rows == 0) {
+    return;
+  }
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+  const cudaDeviceProp *properties = at::cuda::getCurrentDeviceProperties();
+  const int shared_bytes = static_cast<int>((dim + threads + 15) * sizeof(float));
+  TORCH_CHECK(
+      shared_bytes <= properties->sharedMemPerBlockOptin,
+      "native packed W4 activation quantization requires ",
+      shared_bytes,
+      " bytes of shared memory, but the device supports ",
+      properties->sharedMemPerBlockOptin);
+  const dim3 block(static_cast<unsigned int>(threads));
+  const dim3 grid(static_cast<unsigned int>(rows));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+#define ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, DIM_VALUE)     \
+  if (shared_bytes > properties->sharedMemPerBlock) {                           \
+    C10_CUDA_CHECK(cudaFuncSetAttribute(                                        \
+        orbitquant_rpbh_quantize_pack_w4_kernel<STORAGE_TYPE, INDEX_TYPE,       \
+                                                DIM_VALUE>,                     \
+        cudaFuncAttributeMaxDynamicSharedMemorySize,                            \
+        shared_bytes));                                                         \
+  }                                                                             \
+  orbitquant_rpbh_quantize_pack_w4_kernel<STORAGE_TYPE, INDEX_TYPE, DIM_VALUE>  \
+      <<<grid, block, shared_bytes, stream>>>(                                   \
+          packed_out.data_ptr<uint8_t>(),                                       \
+          norms_out.data_ptr<float>(),                                          \
+          reinterpret_cast<STORAGE_TYPE const *>(x.data_ptr()),                 \
+          permutation.data_ptr<INDEX_TYPE>(),                                   \
+          signs.data_ptr<int8_t>(),                                             \
+          boundaries.data_ptr<float>(),                                         \
+          static_cast<float>(eps),                                              \
+          static_cast<float>(inv_sqrt_block),                                   \
+          rows)
+#define ORBITQUANT_DISPATCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE)              \
+  switch (dim) {                                                                \
+    case 512:                                                                   \
+      ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, 512);            \
+      break;                                                                    \
+    case 1024:                                                                  \
+      ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, 1024);           \
+      break;                                                                    \
+    case 2048:                                                                  \
+      ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, 2048);           \
+      break;                                                                    \
+    case 4096:                                                                  \
+      ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, 4096);           \
+      break;                                                                    \
+    case 8192:                                                                  \
+      ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, 8192);           \
+      break;                                                                    \
+    case 16384:                                                                 \
+      ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, 16384);          \
+      break;                                                                    \
+  }
+  const bool int32_permutation = permutation.scalar_type() == torch::kInt;
+  if (x.scalar_type() == torch::kBFloat16) {
+    if (int32_permutation) {
+      ORBITQUANT_DISPATCH_RPBH_PACK_W4(c10::BFloat16, int32_t);
+    } else {
+      ORBITQUANT_DISPATCH_RPBH_PACK_W4(c10::BFloat16, int64_t);
+    }
+  } else {
+    if (int32_permutation) {
+      ORBITQUANT_DISPATCH_RPBH_PACK_W4(c10::Half, int32_t);
+    } else {
+      ORBITQUANT_DISPATCH_RPBH_PACK_W4(c10::Half, int64_t);
+    }
+  }
+#undef ORBITQUANT_DISPATCH_RPBH_PACK_W4
+#undef ORBITQUANT_LAUNCH_RPBH_PACK_W4
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+void quantize_activations_int8(
+    torch::Tensor &int8_out,
+    torch::Tensor &norms_out,
+    torch::Tensor const &x,
+    torch::Tensor const &permutation,
+    torch::Tensor const &signs,
+    torch::Tensor const &boundaries,
+    torch::Tensor const &codes,
+    double eps,
+    double inv_sqrt_block,
+    int64_t threads) {
+  TORCH_CHECK(int8_out.device().is_cuda(), "int8_out must be a CUDA tensor");
+  TORCH_CHECK(norms_out.device().is_cuda(), "norms_out must be a CUDA tensor");
+  TORCH_CHECK(x.device().is_cuda(), "x must be a CUDA tensor");
+  TORCH_CHECK(permutation.device().is_cuda(), "permutation must be a CUDA tensor");
+  TORCH_CHECK(signs.device().is_cuda(), "signs must be a CUDA tensor");
+  TORCH_CHECK(boundaries.device().is_cuda(), "boundaries must be a CUDA tensor");
+  TORCH_CHECK(codes.device().is_cuda(), "codes must be a CUDA tensor");
+  TORCH_CHECK(int8_out.is_contiguous(), "int8_out must be contiguous");
+  TORCH_CHECK(norms_out.is_contiguous(), "norms_out must be contiguous");
+  TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
+  TORCH_CHECK(permutation.is_contiguous(), "permutation must be contiguous");
+  TORCH_CHECK(signs.is_contiguous(), "signs must be contiguous");
+  TORCH_CHECK(boundaries.is_contiguous(), "boundaries must be contiguous");
+  TORCH_CHECK(codes.is_contiguous(), "codes must be contiguous");
+  TORCH_CHECK(int8_out.scalar_type() == torch::kChar, "int8_out must be int8");
+  TORCH_CHECK(norms_out.scalar_type() == torch::kFloat, "norms_out must be float32");
+  TORCH_CHECK(
+      x.scalar_type() == torch::kBFloat16 || x.scalar_type() == torch::kHalf,
+      "x must be bfloat16 or float16");
+  TORCH_CHECK(
+      permutation.scalar_type() == torch::kLong ||
+          permutation.scalar_type() == torch::kInt,
+      "permutation must be int32 or int64");
+  TORCH_CHECK(signs.scalar_type() == torch::kChar, "signs must be int8");
+  TORCH_CHECK(boundaries.scalar_type() == torch::kFloat, "boundaries must be float32");
+  TORCH_CHECK(codes.scalar_type() == torch::kChar, "codes must be int8");
+  TORCH_CHECK(x.dim() == 2, "x must be rank 2");
+  TORCH_CHECK(int8_out.dim() == 2, "int8_out must be rank 2");
+  TORCH_CHECK(norms_out.dim() == 1, "norms_out must be rank 1");
+  const int64_t rows = x.size(0);
+  const int64_t dim = x.size(1);
+  TORCH_CHECK(
+      dim == 512 || dim == 1024 || dim == 2048 || dim == 4096 ||
+          dim == 8192 || dim == 12288 || dim == 16384,
+      "native INT8 activation quantization supports dimensions "
+      "512, 1024, 2048, 4096, 8192, 12288, and 16384");
+  TORCH_CHECK(
+      threads == 128 || threads == 256 || threads == 512,
+      "native INT8 activation quantization threads must be 128, 256, or 512");
+  TORCH_CHECK(
+      int8_out.size(0) == rows && int8_out.size(1) == dim,
+      "int8_out has an unexpected shape");
+  TORCH_CHECK(norms_out.numel() == rows, "norms_out must match rows");
+  TORCH_CHECK(permutation.numel() == dim, "permutation must match the input dimension");
+  TORCH_CHECK(signs.numel() == dim, "signs must match the input dimension");
+  TORCH_CHECK(boundaries.numel() == 15, "boundaries must contain 15 values");
+  TORCH_CHECK(codes.numel() == 16, "codes must contain 16 values");
+  if (rows == 0) {
+    return;
+  }
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+  const cudaDeviceProp *properties = at::cuda::getCurrentDeviceProperties();
+  const int shared_bytes =
+      static_cast<int>((dim + threads + 15) * sizeof(float) + 16);
+  TORCH_CHECK(
+      shared_bytes <= properties->sharedMemPerBlockOptin,
+      "native INT8 activation quantization requires ",
+      shared_bytes,
+      " bytes of shared memory, but the device supports ",
+      properties->sharedMemPerBlockOptin);
+  const dim3 block(static_cast<unsigned int>(threads));
+  const dim3 grid(static_cast<unsigned int>(rows));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+#define ORBITQUANT_LAUNCH_RPBH_INT8(                                           \
+    STORAGE_TYPE, INDEX_TYPE, DIM_VALUE, ORBIT_BLOCK_VALUE)                    \
+  if (shared_bytes > properties->sharedMemPerBlock) {                           \
+    C10_CUDA_CHECK(cudaFuncSetAttribute(                                        \
+        orbitquant_rpbh_quantize_int8_kernel<                                   \
+            STORAGE_TYPE, INDEX_TYPE, DIM_VALUE, ORBIT_BLOCK_VALUE>,            \
+        cudaFuncAttributeMaxDynamicSharedMemorySize,                            \
+        shared_bytes));                                                         \
+  }                                                                             \
+  orbitquant_rpbh_quantize_int8_kernel<                                         \
+      STORAGE_TYPE, INDEX_TYPE, DIM_VALUE, ORBIT_BLOCK_VALUE>                   \
+      <<<grid, block, shared_bytes, stream>>>(                                   \
+          int8_out.data_ptr<int8_t>(),                                          \
+          norms_out.data_ptr<float>(),                                          \
+          reinterpret_cast<STORAGE_TYPE const *>(x.data_ptr()),                 \
+          permutation.data_ptr<INDEX_TYPE>(),                                   \
+          signs.data_ptr<int8_t>(),                                             \
+          boundaries.data_ptr<float>(),                                         \
+          codes.data_ptr<int8_t>(),                                             \
+          static_cast<float>(eps),                                              \
+          static_cast<float>(inv_sqrt_block),                                   \
+          rows)
+#define ORBITQUANT_DISPATCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE)                 \
+  switch (dim) {                                                                \
+    case 512:                                                                   \
+      ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 512, 512);          \
+      break;                                                                    \
+    case 1024:                                                                  \
+      ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 1024, 1024);        \
+      break;                                                                    \
+    case 2048:                                                                  \
+      ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 2048, 2048);        \
+      break;                                                                    \
+    case 4096:                                                                  \
+      ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 4096, 4096);        \
+      break;                                                                    \
+    case 8192:                                                                  \
+      ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 8192, 8192);        \
+      break;                                                                    \
+    case 12288:                                                                 \
+      ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 12288, 4096);       \
+      break;                                                                    \
+    case 16384:                                                                 \
+      ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 16384, 16384);      \
+      break;                                                                    \
+  }
+  const bool int32_permutation = permutation.scalar_type() == torch::kInt;
+  if (x.scalar_type() == torch::kBFloat16) {
+    if (int32_permutation) {
+      ORBITQUANT_DISPATCH_RPBH_INT8(c10::BFloat16, int32_t);
+    } else {
+      ORBITQUANT_DISPATCH_RPBH_INT8(c10::BFloat16, int64_t);
+    }
+  } else {
+    if (int32_permutation) {
+      ORBITQUANT_DISPATCH_RPBH_INT8(c10::Half, int32_t);
+    } else {
+      ORBITQUANT_DISPATCH_RPBH_INT8(c10::Half, int64_t);
+    }
+  }
+#undef ORBITQUANT_DISPATCH_RPBH_INT8
+#undef ORBITQUANT_LAUNCH_RPBH_INT8
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}

orbitquant_packed_matmul_metal/packed_matmul.metal

@@ -11,14 +11,120 @@ struct PackedMatmulParams {
   int has_bias;
 };
 
-inline float bf16_to_float(ushort value) {
-  return as_type<float>(uint(value) << 16);
-}
+struct alignas(16) PackedMMAReadVector {
+  uchar values[16];
+};
+
+struct alignas(1) PackedMMARead3 {
+  uchar values[3];
+};
+
+constant uint packed_gemv_simdgroups = 8;
+constant uint packed_gemv_max_rows = 8;
+
+// Skinny-batch GEMV: one simdgroup owns one output column, the 32 lanes strip
+// the K dimension in 8-value segments (Bits bytes each), and the decoded
+// weights are reused across every row of the batch. The packed column is read
+// with one vector load per segment, which is what makes this path faster than
+// a dequantized F.linear for small row counts: the weight stream is 16/Bits
+// times smaller.
+template <typename scalar_t, uint Bits>
+inline void packed_matmul_gemv_value(
+    device scalar_t *out,
+    device const scalar_t *x,
+    device const uchar *packed_weight_indices,
+    device const float *row_norms,
+    device const float *centroids,
+    device const float *bias,
+    constant PackedMatmulParams &params,
+    threadgroup float *centroid_lut,
+    uint2 group_id,
+    uint thread_index,
+    ushort simdgroup_id,
+    ushort lane) {
+  constexpr uint mask = (1u << Bits) - 1u;
+  if (thread_index < (1u << Bits)) {
+    centroid_lut[thread_index] = centroids[thread_index];
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  const long col = long(group_id.x) * packed_gemv_simdgroups + simdgroup_id;
+  if (col >= params.out_features) {
+    return;
+  }
+  const long rows = min(params.rows, long(packed_gemv_max_rows));
+  device const uchar *column_bytes =
+      packed_weight_indices + (col * params.in_features * Bits) / 8;
+
+  float accumulators[packed_gemv_max_rows] = {};
+
+  // Each lane consumes an 8-value segment (Bits bytes) per step; the 32 lanes
+  // cover 256 values per iteration.
+  const long segments = params.in_features / 8;
+  for (long segment = long(lane); segment < segments; segment += 32) {
+    const long byte_index = segment * Bits;
+    uint low = 0;
+    uint high = 0;
+    if (Bits == 2) {
+      low = uint(*reinterpret_cast<device const ushort *>(
+          column_bytes + byte_index));
+    } else if (Bits == 3) {
+      const PackedMMARead3 bytes =
+          *reinterpret_cast<device const PackedMMARead3 *>(
+              column_bytes + byte_index);
+      low = uint(bytes.values[0]) | (uint(bytes.values[1]) << 8) |
+          (uint(bytes.values[2]) << 16);
+    } else if (Bits == 4) {
+      low = *reinterpret_cast<device const uint *>(column_bytes + byte_index);
+    } else {
+      // Six-byte segments are only 2-byte aligned; read three ushorts.
+      const ushort word0 = *reinterpret_cast<device const ushort *>(
+          column_bytes + byte_index);
+      const ushort word1 = *reinterpret_cast<device const ushort *>(
+          column_bytes + byte_index + 2);
+      const ushort word2 = *reinterpret_cast<device const ushort *>(
+          column_bytes + byte_index + 4);
+      low = uint(word0) | (uint(word1) << 16);
+      high = uint(word2);
+    }
 
-inline ushort float_to_bf16(float value) {
-  const uint bits = as_type<uint>(value);
-  const uint lsb = (bits >> 16) & 1u;
-  return ushort((bits + 0x7fffu + lsb) >> 16);
+    float weights[8];
+#pragma clang loop unroll(full)
+    for (uint idx = 0; idx < 8; ++idx) {
+      uint raw;
+      if (Bits == 6) {
+        const uint bit_start = idx * 6;
+        raw = bit_start < 32
+            ? ((low >> bit_start) |
+               (bit_start > 26 ? (high << (32 - bit_start)) : 0u))
+            : (high >> (bit_start - 32));
+      } else {
+        raw = low >> (idx * Bits);
+      }
+      weights[idx] = centroid_lut[raw & mask];
+    }
+
+    const long k = segment * 8;
+    for (long row = 0; row < rows; ++row) {
+      device const scalar_t *x_row = x + row * params.in_features + k;
+      float partial = 0.0f;
+#pragma clang loop unroll(full)
+      for (uint idx = 0; idx < 8; ++idx) {
+        partial += weights[idx] * float(x_row[idx]);
+      }
+      accumulators[row] += partial;
+    }
+  }
+
+  const float norm = row_norms[col];
+  const float bias_value = params.has_bias != 0 ? bias[col] : 0.0f;
+  for (long row = 0; row < rows; ++row) {
+    const float value = simd_sum(accumulators[row]);
+    if (lane == 0) {
+      out[row * params.out_features + col] =
+          scalar_t(value * norm + bias_value);
+    }
+  }
 }
 
 template <typename scalar_t>
@@ -102,83 +208,214 @@ inline void packed_matmul_tiled_value(
   }
 }
 
-inline void packed_matmul_tiled_value_bfloat16(
-    device ushort *out,
-    device const ushort *x,
+constant uint packed_mma_tile = 32;
+constant uint packed_mma_padded_k = 40;
+
+template <typename scalar_t>
+inline void packed_mma_fragment(
+    thread float2 &output,
+    thread vec<scalar_t, 2> &lhs,
+    thread vec<scalar_t, 2> &rhs,
+    thread float2 &accumulator) {
+  simdgroup_matrix<float, 8, 8> output_matrix;
+  simdgroup_matrix<scalar_t, 8, 8> lhs_matrix;
+  simdgroup_matrix<scalar_t, 8, 8> rhs_matrix;
+  simdgroup_matrix<float, 8, 8> accumulator_matrix;
+  reinterpret_cast<thread vec<scalar_t, 2> &>(lhs_matrix.thread_elements()) = lhs;
+  reinterpret_cast<thread vec<scalar_t, 2> &>(rhs_matrix.thread_elements()) = rhs;
+  reinterpret_cast<thread float2 &>(accumulator_matrix.thread_elements()) = accumulator;
+  simdgroup_multiply_accumulate(
+      output_matrix, lhs_matrix, rhs_matrix, accumulator_matrix);
+  output = reinterpret_cast<thread float2 &>(output_matrix.thread_elements());
+}
+
+template <typename scalar_t, uint Bits>
+inline void decode_packed_mma_weight_segment(
+    threadgroup scalar_t *destination,
+    device const uchar *packed_weight_indices,
+    device const float *row_norms,
+    device const float *centroids,
+    long global_col,
+    long global_k,
+    long in_features) {
+  constexpr uint values = 8;
+  constexpr uint byte_count = Bits;
+  constexpr uint mask = (1u << Bits) - 1u;
+  const long value_offset = global_col * in_features + global_k;
+  const long byte_index = (value_offset * Bits) >> 3;
+  const float norm = row_norms[global_col];
+
+  if (Bits == 2) {
+    const ushort packed = *reinterpret_cast<device const ushort *>(
+        packed_weight_indices + byte_index);
+#pragma clang loop unroll(full)
+    for (uint idx = 0; idx < values; ++idx) {
+      destination[idx] =
+          scalar_t(norm * centroids[(uint(packed) >> (idx * 2)) & 3u]);
+    }
+    return;
+  }
+
+  if (Bits == 3) {
+    const PackedMMARead3 bytes = *reinterpret_cast<device const PackedMMARead3 *>(
+        packed_weight_indices + byte_index);
+    const uint packed =
+        uint(bytes.values[0]) | (uint(bytes.values[1]) << 8) |
+        (uint(bytes.values[2]) << 16);
+#pragma clang loop unroll(full)
+    for (uint idx = 0; idx < values; ++idx) {
+      destination[idx] =
+          scalar_t(norm * centroids[(packed >> (idx * 3)) & 7u]);
+    }
+    return;
+  }
+
+  if (Bits == 4) {
+    const uchar4 packed = *reinterpret_cast<device const uchar4 *>(
+        packed_weight_indices + byte_index);
+#pragma clang loop unroll(full)
+    for (uint idx = 0; idx < 4; ++idx) {
+      const uchar value = packed[idx];
+      destination[idx * 2] = scalar_t(norm * centroids[uint(value) & 15u]);
+      destination[idx * 2 + 1] =
+          scalar_t(norm * centroids[(uint(value) >> 4) & 15u]);
+    }
+    return;
+  }
+
+  uchar packed[byte_count];
+#pragma clang loop unroll(full)
+  for (uint idx = 0; idx < byte_count; ++idx) {
+    packed[idx] = packed_weight_indices[byte_index + idx];
+  }
+#pragma clang loop unroll(full)
+  for (uint idx = 0; idx < values; ++idx) {
+    const uint bit_start = idx * Bits;
+    const uint source_byte = bit_start >> 3;
+    const uint shift = bit_start & 7u;
+    uint raw = uint(packed[source_byte]);
+    if (shift + Bits > 8u) {
+      raw |= uint(packed[source_byte + 1]) << 8;
+    }
+    destination[idx] = scalar_t(norm * centroids[(raw >> shift) & mask]);
+  }
+}
+
+template <typename scalar_t, uint Bits>
+inline void packed_matmul_padded_mma_value(
+    device scalar_t *out,
+    device const scalar_t *x,
     device const uchar *packed_weight_indices,
     device const float *row_norms,
     device const float *centroids,
     device const float *bias,
     constant PackedMatmulParams &params,
-    threadgroup float *shared,
-    uint2 tid,
-    uint2 local_tid,
-    uint2 threads_per_group) {
-  const long col = tid.x;
-  const long row = tid.y;
-  const bool output_valid = row < params.rows && col < params.out_features;
-  const long block_k = params.block_k;
-  threadgroup float *x_tile = shared;
-  threadgroup float *w_tile = shared + long(threads_per_group.y) * block_k;
-  const long local_col = local_tid.x;
-  const long local_row = local_tid.y;
-  const long thread_linear = local_row * long(threads_per_group.x) + local_col;
-  const long thread_count = long(threads_per_group.x) * long(threads_per_group.y);
-
-  const uint mask = (1u << uint(params.bits)) - 1u;
-  float acc = output_valid && params.has_bias != 0 ? bias[col] : 0.0f;
+    threadgroup uchar *shared,
+    uint2 group_id,
+    uint thread_index,
+    ushort simdgroup_id,
+    ushort lane) {
+  threadgroup scalar_t *x_tile = reinterpret_cast<threadgroup scalar_t *>(shared);
+  threadgroup scalar_t *weight_tile = x_tile + packed_mma_tile * packed_mma_padded_k;
+  const long row_start = long(group_id.y) * packed_mma_tile;
+  const long col_start = long(group_id.x) * packed_mma_tile;
+  const ushort load_row = ushort(thread_index / 4);
+  const ushort load_k = ushort(thread_index % 4) * 8;
+  const ushort quad = lane / 4;
+  const ushort fragment_row = (quad & 4) + ((lane / 2) % 4);
+  const ushort fragment_col = (quad & 2) * 2 + (lane % 2) * 2;
+  const ushort simd_row = 8 * (simdgroup_id / 2);
+  const ushort simd_col = 8 * (simdgroup_id % 2);
 
-  for (long k_start = 0; k_start < params.in_features; k_start += block_k) {
-    const long x_tile_values = long(threads_per_group.y) * block_k;
-    for (long offset = thread_linear; offset < x_tile_values; offset += thread_count) {
-      const long tile_row = offset / block_k;
-      const long tile_k = offset - tile_row * block_k;
-      const long global_row = long(tid.y) - local_row + tile_row;
-      const long global_k = k_start + tile_k;
-      float value = 0.0f;
-      if (global_row < params.rows && global_k < params.in_features) {
-        value = bf16_to_float(x[global_row * params.in_features + global_k]);
-      }
-      x_tile[offset] = value;
+  float2 accumulators[2][2];
+#pragma clang loop unroll(full)
+  for (ushort matrix_m = 0; matrix_m < 2; ++matrix_m) {
+#pragma clang loop unroll(full)
+    for (ushort matrix_n = 0; matrix_n < 2; ++matrix_n) {
+      accumulators[matrix_m][matrix_n] = float2(0.0f);
     }
+  }
 
-    const long w_tile_values = block_k * long(threads_per_group.x);
-    for (long offset = thread_linear; offset < w_tile_values; offset += thread_count) {
-      const long tile_k = offset / long(threads_per_group.x);
-      const long tile_col = offset - tile_k * long(threads_per_group.x);
-      const long global_k = k_start + tile_k;
-      const long global_col = long(tid.x) - local_col + tile_col;
-      float value = 0.0f;
-      if (global_col < params.out_features && global_k < params.in_features) {
-        const long value_offset = global_col * params.in_features + global_k;
-        const long bit_start = value_offset * params.bits;
-        const long byte_index = bit_start >> 3;
-        const long bit_offset = bit_start & 7;
-        uint raw = packed_weight_indices[byte_index];
-        if (bit_offset + params.bits > 8) {
-          raw |= uint(packed_weight_indices[byte_index + 1]) << 8;
-        }
-        const uint index = (raw >> uint(bit_offset)) & mask;
-        value = row_norms[global_col] * centroids[index];
+  for (long k_start = 0; k_start < params.in_features; k_start += packed_mma_tile) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    threadgroup scalar_t *x_destination =
+        x_tile + load_row * packed_mma_padded_k + load_k;
+    if (row_start + long(load_row) < params.rows) {
+      *reinterpret_cast<threadgroup PackedMMAReadVector *>(x_destination) =
+          *reinterpret_cast<device const PackedMMAReadVector *>(
+              x + (row_start + long(load_row)) * params.in_features +
+              k_start + long(load_k));
+    } else {
+#pragma clang loop unroll(full)
+      for (ushort offset = 0; offset < 8; ++offset) {
+        x_destination[offset] = scalar_t(0.0f);
       }
-      w_tile[offset] = value;
     }
-
+    decode_packed_mma_weight_segment<scalar_t, Bits>(
+        weight_tile + load_row * packed_mma_padded_k + load_k,
+        packed_weight_indices,
+        row_norms,
+        centroids,
+        col_start + long(load_row),
+        k_start + long(load_k),
+        params.in_features);
     threadgroup_barrier(mem_flags::mem_threadgroup);
 
-    if (output_valid) {
-      for (long tile_k = 0; tile_k < block_k && k_start + tile_k < params.in_features;
-           ++tile_k) {
-        acc += x_tile[local_row * block_k + tile_k] *
-            w_tile[tile_k * long(threads_per_group.x) + local_col];
+#pragma clang loop unroll(full)
+    for (ushort matrix_k = 0; matrix_k < packed_mma_tile; matrix_k += 8) {
+      vec<scalar_t, 2> x_fragments[2];
+      vec<scalar_t, 2> weight_fragments[2];
+      simdgroup_barrier(mem_flags::mem_none);
+#pragma clang loop unroll(full)
+      for (ushort matrix_m = 0; matrix_m < 2; ++matrix_m) {
+        threadgroup const scalar_t *source =
+            x_tile + (simd_row + fragment_row + matrix_m * 16) *
+                packed_mma_padded_k +
+            matrix_k + fragment_col;
+        x_fragments[matrix_m] = vec<scalar_t, 2>(source[0], source[1]);
+      }
+      simdgroup_barrier(mem_flags::mem_none);
+#pragma clang loop unroll(full)
+      for (ushort matrix_n = 0; matrix_n < 2; ++matrix_n) {
+        threadgroup const scalar_t *source =
+            weight_tile + (simd_col + fragment_col + matrix_n * 16) *
+                packed_mma_padded_k +
+            matrix_k + fragment_row;
+        weight_fragments[matrix_n] =
+            vec<scalar_t, 2>(source[0], source[packed_mma_padded_k]);
+      }
+      simdgroup_barrier(mem_flags::mem_none);
+#pragma clang loop unroll(full)
+      for (ushort matrix_m = 0; matrix_m < 2; ++matrix_m) {
+#pragma clang loop unroll(full)
+        for (ushort matrix_n = 0; matrix_n < 2; ++matrix_n) {
+          packed_mma_fragment(
+              accumulators[matrix_m][matrix_n],
+              x_fragments[matrix_m],
+              weight_fragments[matrix_n],
+              accumulators[matrix_m][matrix_n]);
+        }
       }
     }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
   }
 
-  if (output_valid) {
-    out[row * params.out_features + col] = float_to_bf16(acc);
+#pragma clang loop unroll(full)
+  for (ushort matrix_m = 0; matrix_m < 2; ++matrix_m) {
+#pragma clang loop unroll(full)
+    for (ushort matrix_n = 0; matrix_n < 2; ++matrix_n) {
+      const long global_row =
+          row_start + long(simd_row + fragment_row + matrix_m * 16);
+      const long global_col =
+          col_start + long(simd_col + fragment_col + matrix_n * 16);
+      float2 values = accumulators[matrix_m][matrix_n];
+      if (params.has_bias != 0) {
+        values += float2(bias[global_col], bias[global_col + 1]);
+      }
+      if (global_row < params.rows) {
+        out[global_row * params.out_features + global_col] = scalar_t(values[0]);
+        out[global_row * params.out_features + global_col + 1] = scalar_t(values[1]);
+      }
+    }
   }
 }
 
@@ -209,6 +446,74 @@ kernel void packed_matmul_forward_float(
 }
 
 kernel void packed_matmul_forward_half(
+    device half *out [[buffer(0)]],
+    device const half *x [[buffer(1)]],
+    device const uchar *packed_weight_indices [[buffer(2)]],
+    device const float *row_norms [[buffer(3)]],
+    device const float *centroids [[buffer(4)]],
+    device const float *bias [[buffer(5)]],
+    constant PackedMatmulParams &params [[buffer(6)]],
+    threadgroup uchar *shared [[threadgroup(0)]],
+    uint2 group_id [[threadgroup_position_in_grid]],
+    uint thread_index [[thread_index_in_threadgroup]],
+    ushort simdgroup_id [[simdgroup_index_in_threadgroup]],
+    ushort lane [[thread_index_in_simdgroup]]) {
+#define ORBITQUANT_PACKED_MMA_HALF(BITS_VALUE)                              \
+  packed_matmul_padded_mma_value<half, BITS_VALUE>(                         \
+      out, x, packed_weight_indices, row_norms, centroids, bias, params,    \
+      shared, group_id, thread_index, simdgroup_id, lane)
+  switch (params.bits) {
+    case 2:
+      ORBITQUANT_PACKED_MMA_HALF(2);
+      break;
+    case 3:
+      ORBITQUANT_PACKED_MMA_HALF(3);
+      break;
+    case 4:
+      ORBITQUANT_PACKED_MMA_HALF(4);
+      break;
+    case 6:
+      ORBITQUANT_PACKED_MMA_HALF(6);
+      break;
+  }
+#undef ORBITQUANT_PACKED_MMA_HALF
+}
+
+kernel void packed_matmul_forward_bfloat16(
+    device bfloat *out [[buffer(0)]],
+    device const bfloat *x [[buffer(1)]],
+    device const uchar *packed_weight_indices [[buffer(2)]],
+    device const float *row_norms [[buffer(3)]],
+    device const float *centroids [[buffer(4)]],
+    device const float *bias [[buffer(5)]],
+    constant PackedMatmulParams &params [[buffer(6)]],
+    threadgroup uchar *shared [[threadgroup(0)]],
+    uint2 group_id [[threadgroup_position_in_grid]],
+    uint thread_index [[thread_index_in_threadgroup]],
+    ushort simdgroup_id [[simdgroup_index_in_threadgroup]],
+    ushort lane [[thread_index_in_simdgroup]]) {
+#define ORBITQUANT_PACKED_MMA_BFLOAT16(BITS_VALUE)                          \
+  packed_matmul_padded_mma_value<bfloat, BITS_VALUE>(                       \
+      out, x, packed_weight_indices, row_norms, centroids, bias, params,    \
+      shared, group_id, thread_index, simdgroup_id, lane)
+  switch (params.bits) {
+    case 2:
+      ORBITQUANT_PACKED_MMA_BFLOAT16(2);
+      break;
+    case 3:
+      ORBITQUANT_PACKED_MMA_BFLOAT16(3);
+      break;
+    case 4:
+      ORBITQUANT_PACKED_MMA_BFLOAT16(4);
+      break;
+    case 6:
+      ORBITQUANT_PACKED_MMA_BFLOAT16(6);
+      break;
+  }
+#undef ORBITQUANT_PACKED_MMA_BFLOAT16
+}
+
+kernel void packed_matmul_forward_half_scalar(
     device half *out [[buffer(0)]],
     device const half *x [[buffer(1)]],
     device const uchar *packed_weight_indices [[buffer(2)]],
@@ -234,9 +539,9 @@ kernel void packed_matmul_forward_half(
       threads_per_group);
 }
 
-kernel void packed_matmul_forward_bfloat16(
-    device ushort *out [[buffer(0)]],
-    device const ushort *x [[buffer(1)]],
+kernel void packed_matmul_forward_bfloat16_scalar(
+    device bfloat *out [[buffer(0)]],
+    device const bfloat *x [[buffer(1)]],
     device const uchar *packed_weight_indices [[buffer(2)]],
     device const float *row_norms [[buffer(3)]],
     device const float *centroids [[buffer(4)]],
@@ -246,7 +551,7 @@ kernel void packed_matmul_forward_bfloat16(
     uint2 tid [[thread_position_in_grid]],
     uint2 local_tid [[thread_position_in_threadgroup]],
     uint2 threads_per_group [[threads_per_threadgroup]]) {
-  packed_matmul_tiled_value_bfloat16(
+  packed_matmul_tiled_value(
       out,
       x,
       packed_weight_indices,
@@ -259,3 +564,51 @@ kernel void packed_matmul_forward_bfloat16(
       local_tid,
       threads_per_group);
 }
+
+#define ORBITQUANT_PACKED_SMALL_ROWS_KERNEL(NAME, TYPE)                     \
+  kernel void NAME(                                                        \
+      device TYPE *out [[buffer(0)]],                                      \
+      device const TYPE *x [[buffer(1)]],                                  \
+      device const uchar *packed_weight_indices [[buffer(2)]],             \
+      device const float *row_norms [[buffer(3)]],                          \
+      device const float *centroids [[buffer(4)]],                          \
+      device const float *bias [[buffer(5)]],                               \
+      constant PackedMatmulParams &params [[buffer(6)]],                    \
+      uint2 group_id [[threadgroup_position_in_grid]],                      \
+      uint thread_index [[thread_index_in_threadgroup]],                    \
+      ushort simdgroup_id [[simdgroup_index_in_threadgroup]],               \
+      ushort lane [[thread_index_in_simdgroup]]) {                          \
+    threadgroup float centroid_lut[64];                                     \
+    switch (params.bits) {                                                  \
+      case 2:                                                               \
+        packed_matmul_gemv_value<TYPE, 2>(                                  \
+            out, x, packed_weight_indices, row_norms, centroids, bias,      \
+            params, centroid_lut, group_id, thread_index, simdgroup_id,     \
+            lane);                                                          \
+        break;                                                              \
+      case 3:                                                               \
+        packed_matmul_gemv_value<TYPE, 3>(                                  \
+            out, x, packed_weight_indices, row_norms, centroids, bias,      \
+            params, centroid_lut, group_id, thread_index, simdgroup_id,     \
+            lane);                                                          \
+        break;                                                              \
+      case 4:                                                               \
+        packed_matmul_gemv_value<TYPE, 4>(                                  \
+            out, x, packed_weight_indices, row_norms, centroids, bias,      \
+            params, centroid_lut, group_id, thread_index, simdgroup_id,     \
+            lane);                                                          \
+        break;                                                              \
+      case 6:                                                               \
+        packed_matmul_gemv_value<TYPE, 6>(                                  \
+            out, x, packed_weight_indices, row_norms, centroids, bias,      \
+            params, centroid_lut, group_id, thread_index, simdgroup_id,     \
+            lane);                                                          \
+        break;                                                              \
+    }                                                                       \
+  }
+
+ORBITQUANT_PACKED_SMALL_ROWS_KERNEL(packed_matmul_forward_float_small_rows, float)
+ORBITQUANT_PACKED_SMALL_ROWS_KERNEL(packed_matmul_forward_half_small_rows, half)
+ORBITQUANT_PACKED_SMALL_ROWS_KERNEL(packed_matmul_forward_bfloat16_small_rows, bfloat)
+
+#undef ORBITQUANT_PACKED_SMALL_ROWS_KERNEL

orbitquant_packed_matmul_metal/packed_matmul.mm

@@ -33,6 +33,11 @@ struct PackedMatmulPipelineCache {
   id<MTLComputePipelineState> float_pipeline;
   id<MTLComputePipelineState> half_pipeline;
   id<MTLComputePipelineState> bfloat16_pipeline;
+  id<MTLComputePipelineState> half_scalar_pipeline;
+  id<MTLComputePipelineState> bfloat16_scalar_pipeline;
+  id<MTLComputePipelineState> float_small_rows_pipeline;
+  id<MTLComputePipelineState> half_small_rows_pipeline;
+  id<MTLComputePipelineState> bfloat16_small_rows_pipeline;
 };
 
 static id<MTLComputePipelineState> create_pipeline(
@@ -68,6 +73,16 @@ static PackedMatmulPipelineCache &packed_matmul_pipeline_cache() {
           create_pipeline(cache.device, library, "packed_matmul_forward_half");
       cache.bfloat16_pipeline =
           create_pipeline(cache.device, library, "packed_matmul_forward_bfloat16");
+      cache.half_scalar_pipeline =
+          create_pipeline(cache.device, library, "packed_matmul_forward_half_scalar");
+      cache.bfloat16_scalar_pipeline =
+          create_pipeline(cache.device, library, "packed_matmul_forward_bfloat16_scalar");
+      cache.float_small_rows_pipeline =
+          create_pipeline(cache.device, library, "packed_matmul_forward_float_small_rows");
+      cache.half_small_rows_pipeline =
+          create_pipeline(cache.device, library, "packed_matmul_forward_half_small_rows");
+      cache.bfloat16_small_rows_pipeline = create_pipeline(
+          cache.device, library, "packed_matmul_forward_bfloat16_small_rows");
     }
   });
   return cache;
@@ -75,14 +90,25 @@ static PackedMatmulPipelineCache &packed_matmul_pipeline_cache() {
 
 static id<MTLComputePipelineState> select_packed_matmul_pipeline(
     PackedMatmulPipelineCache &cache,
-    c10::ScalarType dtype) {
+    c10::ScalarType dtype,
+    bool use_padded_mma,
+    bool use_small_rows) {
+  if (use_small_rows) {
+    if (dtype == torch::kFloat) {
+      return cache.float_small_rows_pipeline;
+    }
+    if (dtype == torch::kHalf) {
+      return cache.half_small_rows_pipeline;
+    }
+    return cache.bfloat16_small_rows_pipeline;
+  }
   if (dtype == torch::kFloat) {
     return cache.float_pipeline;
   }
   if (dtype == torch::kHalf) {
-    return cache.half_pipeline;
+    return use_padded_mma ? cache.half_pipeline : cache.half_scalar_pipeline;
   }
-  return cache.bfloat16_pipeline;
+  return use_padded_mma ? cache.bfloat16_pipeline : cache.bfloat16_scalar_pipeline;
 }
 
 static void dispatch_packed_matmul_kernel(
@@ -100,9 +126,21 @@ static void dispatch_packed_matmul_kernel(
     int64_t block_n,
     int64_t block_k) {
   @autoreleasepool {
+    // Skinny batches take the GEMV path even when the MMA tiles would fit:
+    // below 5 rows the 32-row MMA tile wastes most of its work, while the
+    // GEMV path reuses each decoded weight segment across the whole batch.
+    // From 5 rows up the measured crossover already favors the MMA tiles.
+    const bool use_small_rows = x.size(0) <= 4 &&
+        (bits == 2 || bits == 3 || bits == 4 || bits == 6) &&
+        in_features % 8 == 0;
+    const bool use_padded_mma = !use_small_rows &&
+        x.scalar_type() != torch::kFloat && x.size(0) > 1 &&
+        out_features % 32 == 0 && in_features % 32 == 0 &&
+        (bits == 2 || bits == 3 || bits == 4 || bits == 6);
     PackedMatmulPipelineCache &cache = packed_matmul_pipeline_cache();
     id<MTLComputePipelineState> pipeline =
-        select_packed_matmul_pipeline(cache, x.scalar_type());
+        select_packed_matmul_pipeline(
+            cache, x.scalar_type(), use_padded_mma, use_small_rows);
 
     id<MTLCommandBuffer> command_buffer = torch::mps::get_command_buffer();
     TORCH_CHECK(command_buffer, "Failed to retrieve MPS command buffer");
@@ -141,14 +179,54 @@ static void dispatch_packed_matmul_kernel(
                  atIndex:5];
       [encoder setBytes:&params length:sizeof(params) atIndex:6];
 
-      const NSUInteger threads_x = std::min<int64_t>(std::max<int64_t>(block_n, 1), 32);
-      const NSUInteger threads_y = std::min<int64_t>(std::max<int64_t>(block_m, 1), 32);
-      const NSUInteger tile_k = std::min<int64_t>(std::max<int64_t>(block_k, 1), 128);
-      const NSUInteger shared_bytes = (threads_y * tile_k + tile_k * threads_x) * sizeof(float);
-      [encoder setThreadgroupMemoryLength:shared_bytes atIndex:0];
-      MTLSize grid_size = MTLSizeMake(out_features, x.size(0), 1);
-      MTLSize threadgroup_size = MTLSizeMake(threads_x, threads_y, 1);
-      [encoder dispatchThreads:grid_size threadsPerThreadgroup:threadgroup_size];
+      if (use_small_rows) {
+        constexpr NSUInteger columns_per_threadgroup = 8;
+        constexpr NSUInteger threads = 8 * 32;
+        MTLSize threadgroups = MTLSizeMake(
+            (out_features + columns_per_threadgroup - 1) /
+                columns_per_threadgroup,
+            1,
+            1);
+        MTLSize threadgroup_size = MTLSizeMake(threads, 1, 1);
+        [encoder dispatchThreadgroups:threadgroups
+            threadsPerThreadgroup:threadgroup_size];
+      } else if (!use_padded_mma) {
+        const NSUInteger threads_x =
+            std::min<int64_t>(std::max<int64_t>(block_n, 1), 32);
+        const NSUInteger threads_y = std::min<int64_t>(
+            x.size(0), std::min<int64_t>(std::max<int64_t>(block_m, 1), 32));
+        const NSUInteger tile_k =
+            std::min<int64_t>(std::max<int64_t>(block_k, 1), 128);
+        const NSUInteger shared_bytes =
+            (threads_y * tile_k + tile_k * threads_x) * sizeof(float);
+        [encoder setThreadgroupMemoryLength:shared_bytes atIndex:0];
+        MTLSize grid_size = MTLSizeMake(out_features, x.size(0), 1);
+        MTLSize threadgroup_size = MTLSizeMake(threads_x, threads_y, 1);
+        [encoder dispatchThreads:grid_size threadsPerThreadgroup:threadgroup_size];
+      } else {
+        constexpr NSUInteger tile_m = 32;
+        constexpr NSUInteger tile_n = 32;
+        constexpr NSUInteger padded_k = 40;
+        constexpr NSUInteger threads = 128;
+        const NSUInteger scalar_bytes = x.element_size();
+        const NSUInteger shared_bytes =
+            2 * tile_m * padded_k * scalar_bytes;
+        TORCH_CHECK(
+            threads <= pipeline.maxTotalThreadsPerThreadgroup,
+            "Metal packed matmul pipeline supports only ",
+            pipeline.maxTotalThreadsPerThreadgroup,
+            " threads per threadgroup, but ",
+            threads,
+            " are required");
+        [encoder setThreadgroupMemoryLength:shared_bytes atIndex:0];
+        MTLSize threadgroups = MTLSizeMake(
+            (out_features + tile_n - 1) / tile_n,
+            (x.size(0) + tile_m - 1) / tile_m,
+            1);
+        MTLSize threadgroup_size = MTLSizeMake(threads, 1, 1);
+        [encoder dispatchThreadgroups:threadgroups
+            threadsPerThreadgroup:threadgroup_size];
+      }
       [encoder endEncoding];
       torch::mps::commit();
     });

scripts/prepare_wheel_project.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Prepare kernel-builder output for a platform wheel build."
+    )
+    parser.add_argument("project", type=Path)
+    parser.add_argument("--version", required=True)
+    parser.add_argument(
+        "--torch-requirement",
+        default="torch>=2.11",
+        help=(
+            "torch dependency for the wheel metadata; non-stable-ABI variants "
+            'must pin the torch minor they were built against (e.g. "torch>=2.9,<2.10")'
+        ),
+    )
+    args = parser.parse_args()
+
+    pyproject = args.project / "pyproject.toml"
+    text = pyproject.read_text(encoding="utf-8")
+
+    old_version = 'version = "0.1.0"'
+    if text.count(old_version) != 1:
+        raise RuntimeError("generated pyproject must contain one stub version")
+    text = text.replace(old_version, f'version = "{args.version}"', 1)
+
+    requires_python = 'requires-python = ">=3.9"'
+    if text.count(requires_python) != 1:
+        raise RuntimeError("generated pyproject must contain one Python requirement")
+    text = text.replace(
+        requires_python,
+        f'{requires_python}\ndependencies = ["{args.torch_requirement}"]',
+        1,
+    )
+    pyproject.write_text(text, encoding="utf-8")
+
+    cmake = args.project / "CMakeLists.txt"
+    cmake_text = cmake.read_text(encoding="utf-8")
+    for required in (False, True):
+        marker = " REQUIRED" if required else ""
+        development = (
+            f"find_package(Python3{marker} COMPONENTS Development "
+            "Development.SABIModule Interpreter)"
+        )
+        if cmake_text.count(development) != 1:
+            raise RuntimeError(
+                "generated CMake must contain one Python development lookup"
+            )
+        cmake_text = cmake_text.replace(
+            development,
+            f"find_package(Python3{marker} COMPONENTS Development.SABIModule Interpreter)",
+            1,
+        )
+    cmake.write_text(cmake_text, encoding="utf-8")
+
+    setup = args.project / "setup.py"
+    setup_text = setup.read_text(encoding="utf-8")
+    shutil_import = "from shutil import which, move\n"
+    if setup_text.count(shutil_import) != 1:
+        raise RuntimeError("generated setup must contain one shutil import")
+    setup_text = setup_text.replace(
+        shutil_import,
+        "from shutil import copy2, move, which\n",
+        1,
+    )
+    ninja_path = 'ninja_executable_path = Path(ninja.BIN_DIR) / "ninja"'
+    if setup_text.count(ninja_path) != 1:
+        raise RuntimeError("generated setup must contain one Ninja executable path")
+    setup_text = setup_text.replace(
+        ninja_path,
+        'ninja_executable_path = Path(ninja.BIN_DIR) / '
+        '("ninja.exe" if os.name == "nt" else "ninja")',
+        1,
+    )
+    for cache_tool in ("sccache", "ccache"):
+        availability = f'return which("{cache_tool}") is not None'
+        if setup_text.count(availability) != 1:
+            raise RuntimeError(
+                f"generated setup must contain one {cache_tool} availability check"
+            )
+        setup_text = setup_text.replace(
+            availability,
+            f'return os.name != "nt" and which("{cache_tool}") is not None',
+            1,
+        )
+    cmake_args_hook = (
+        '    if "CMAKE_ARGS" in os.environ:\n'
+        '        cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") '
+        "if item]\n"
+    )
+    if setup_text.count(cmake_args_hook) != 1:
+        raise RuntimeError("generated setup must contain one CMAKE_ARGS hook")
+    setup_text = setup_text.replace(
+        cmake_args_hook,
+        cmake_args_hook
+        + '    cmake_make_program = os.environ.get("ORBITQUANT_CMAKE_MAKE_PROGRAM")\n'
+        + "    if cmake_make_program:\n"
+        + '        cmake_args.append(f"-DCMAKE_MAKE_PROGRAM:FILEPATH={cmake_make_program}")\n',
+        1,
+    )
+    build_temp = "        build_temp = Path(self.build_temp) / ext.name"
+    if setup_text.count(build_temp) != 1:
+        raise RuntimeError("generated setup must contain one extension build temp")
+    setup_text = setup_text.replace(
+        build_temp,
+        '        build_temp_root = os.environ.get("ORBITQUANT_BUILD_TEMP", '
+        "self.build_temp)\n"
+        "        build_temp = (Path(build_temp_root) / ext.name).resolve()",
+        1,
+    )
+    windows_multi_config = (
+        '        if sys.platform == "win32":\n'
+        "            # Move the dylib one folder up for discovery."
+    )
+    if setup_text.count(windows_multi_config) != 1:
+        raise RuntimeError("generated setup must contain one Windows output move")
+    setup_text = setup_text.replace(
+        windows_multi_config,
+        '        if sys.platform == "win32" and (extdir / cfg).is_dir():\n'
+        "            # Move the dylib one folder up for discovery.",
+        1,
+    )
+    build_call = (
+        "        subprocess.run(\n"
+        '            ["cmake", "--build", str(build_temp), *build_args], '
+        "cwd=build_temp, check=True\n"
+        "        )\n"
+    )
+    if setup_text.count(build_call) != 1:
+        raise RuntimeError("generated setup must contain one wheel CMake build call")
+    setup_text = setup_text.replace(
+        build_call,
+        build_call
+        + "\n"
+        + '        package_name = ext.name.split(".", 1)[0]\n'
+        + "        generated_ops = (\n"
+        + '            Path(ext.sourcedir) / "torch-ext" / package_name / "_ops.py"\n'
+        + "        )\n"
+        + '        copy2(generated_ops, extdir / "_ops.py")\n',
+        1,
+    )
+    zip_safe = "    zip_safe=False,\n"
+    if setup_text.count(zip_safe) != 1:
+        raise RuntimeError("generated setup must contain one zip-safe option")
+    setup_text = setup_text.replace(
+        zip_safe,
+        '    options={"bdist_wheel": {"py_limited_api": "cp39"}},\n' + zip_safe,
+        1,
+    )
+    setup.write_text(setup_text, encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()

tests/test_packed_matmul.py

@@ -1,8 +1,20 @@
 from __future__ import annotations
 
+import platform
+
 import pytest
 import torch
-from orbitquant_packed_matmul import matmul_packed_weight
+from orbitquant_packed_matmul import (
+    matmul_packed_adaln_int4_cpu,
+    matmul_packed_w4a4_int8,
+    matmul_packed_weight,
+    quantize_activations_cpu,
+    quantize_activations_int8,
+    quantize_activations_packed_w4,
+    supports_cpu_activation,
+    supports_cpu_adaln,
+    supports_device,
+)
 
 
 def _pack(values: torch.Tensor, bits: int) -> torch.Tensor:
@@ -19,15 +31,17 @@ def _pack(values: torch.Tensor, bits: int) -> torch.Tensor:
 
 
 def _device() -> str:
-    if torch.cuda.is_available():
+    if supports_device("cuda") and torch.cuda.is_available():
         return "cuda"
-    if torch.backends.mps.is_available():
+    if supports_device("mps") and torch.backends.mps.is_available():
         return "mps"
-    pytest.skip("CUDA or MPS is required")
+    if supports_device("cpu"):
+        return "cpu"
+    pytest.skip("the built variant has no runnable backend")
 
 
 def _mps_device() -> str:
-    if not torch.backends.mps.is_available():
+    if not supports_device("mps") or not torch.backends.mps.is_available():
         pytest.skip("MPS is required")
     return "mps"
 
@@ -41,6 +55,313 @@ def _mps_bfloat16_device() -> str:
     return device
 
 
+def _cuda_device() -> str:
+    if not supports_device("cuda") or not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    return "cuda"
+
+
+def _runnable_cpu_isas() -> list[str]:
+    machine = platform.machine().lower()
+    capability = torch.backends.cpu.get_cpu_capability().upper()
+    isas = ["scalar"]
+    if machine in {"x86_64", "amd64"} and capability in {"AVX2", "AVX512"}:
+        isas.append("avx2")
+    if machine in {"x86_64", "amd64"} and capability == "AVX512" and platform.system() != "Windows":
+        # The MSVC wheel currently ships the separately compiled AVX2 TU; the
+        # AVX-512 implementation uses GCC/Clang per-function target attributes.
+        isas.append("avx512")
+    if machine in {"aarch64", "arm64"}:
+        isas.append("neon")
+    return isas
+
+
+def _row_norms(device: str, out_features: int) -> torch.Tensor:
+    dtype = torch.bfloat16 if device == "cuda" else torch.float32
+    return torch.linspace(0.5, 1.5, out_features, device=device, dtype=dtype)
+
+
+def _fwht_reference(values: torch.Tensor) -> torch.Tensor:
+    output = values.clone()
+    half = 1
+    while half < output.shape[-1]:
+        blocks = output.reshape(*output.shape[:-1], -1, 2 * half)
+        left = blocks[..., :half].clone()
+        right = blocks[..., half:].clone()
+        blocks[..., :half] = left + right
+        blocks[..., half:] = left - right
+        half *= 2
+    return output
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_quantize_activations_cpu_matches_independent_reference(dtype: torch.dtype) -> None:
+    if not supports_cpu_activation():
+        pytest.skip("the built variant has no native CPU activation pipeline")
+    torch.manual_seed(41)
+    dim = 24
+    block_size = 8
+    x = torch.randn(2, 3, dim, dtype=dtype)
+    x[0, 0].zero_()
+    permutation = torch.randperm(dim)
+    signs = torch.randint(0, 2, (dim,), dtype=torch.int8).mul(2).sub(1)
+    centroids = torch.tanh(torch.linspace(-1.7, 1.7, 16))
+    boundaries = (centroids[:-1] + centroids[1:]) / 2
+    eps = 1e-10
+
+    work = x.float()
+    norms = work.norm(dim=-1, keepdim=True)
+    unit = work / (norms + eps)
+    gathered = unit.index_select(-1, permutation) * signs.float()
+    rotated = _fwht_reference(gathered.reshape(2, 3, 3, block_size)) / block_size**0.5
+    rotated = rotated.reshape_as(work)
+    indices = (rotated.unsqueeze(-1) - centroids).abs().argmin(dim=-1)
+    expected = (centroids[indices] * norms).to(dtype)
+
+    actual = quantize_activations_cpu(
+        x,
+        permutation,
+        signs,
+        centroids,
+        boundaries,
+        eps=eps,
+        inv_sqrt_block=block_size**-0.5,
+        block_size=block_size,
+    )
+
+    torch.testing.assert_close(actual, expected, atol=2e-3, rtol=2e-3)
+    assert torch.equal(actual[0, 0], torch.zeros(dim, dtype=dtype))
+
+
+@pytest.mark.kernels_ci
+def test_cpu_runtime_isa_dispatch_matches_scalar_reference(monkeypatch) -> None:
+    if not supports_device("cpu") or not supports_cpu_activation():
+        pytest.skip("the built variant has no complete native CPU pipeline")
+    torch.manual_seed(43)
+    rows = 8
+    in_features = 64
+    out_features = 11
+    x = torch.randn(rows, in_features)
+    indices = torch.randint(0, 16, (out_features, in_features), dtype=torch.uint8)
+    packed = _pack(indices, 4)
+    row_norms = torch.linspace(0.5, 1.5, out_features)
+    centroids = torch.tanh(torch.linspace(-1.7, 1.7, 16))
+    boundaries = (centroids[:-1] + centroids[1:]) / 2
+    bias = torch.randn(out_features)
+    permutation = torch.randperm(in_features)
+    signs = torch.randint(0, 2, (in_features,), dtype=torch.int8).mul(2).sub(1)
+
+    outputs = {}
+    activations = {}
+    for isa in _runnable_cpu_isas():
+        monkeypatch.setenv("ORBITQUANT_CPU_ISA", isa)
+        activations[isa] = quantize_activations_cpu(
+            x,
+            permutation,
+            signs,
+            centroids,
+            boundaries,
+            eps=1e-10,
+            inv_sqrt_block=in_features**-0.5,
+            block_size=in_features,
+        )
+        outputs[isa] = matmul_packed_weight(
+            activations[isa],
+            packed,
+            row_norms,
+            centroids,
+            bits=4,
+            out_features=out_features,
+            in_features=in_features,
+            bias=bias,
+        )
+
+    for isa in _runnable_cpu_isas()[1:]:
+        torch.testing.assert_close(activations[isa], activations["scalar"], atol=2e-6, rtol=2e-6)
+        torch.testing.assert_close(outputs[isa], outputs["scalar"], atol=2e-5, rtol=2e-5)
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.parametrize("bits", [2, 3, 6])
+@pytest.mark.parametrize("in_features", [64, 84])
+def test_cpu_isa_matmul_matches_scalar_for_low_bit_widths(
+    monkeypatch, bits: int, in_features: int
+) -> None:
+    if not supports_device("cpu"):
+        pytest.skip("the built variant has no native CPU backend")
+    torch.manual_seed(47 + bits + in_features)
+    rows = 9
+    out_features = 11
+    levels = 2**bits
+    x = torch.randn(rows, in_features)
+    indices = torch.randint(0, levels, (out_features, in_features), dtype=torch.uint8)
+    packed = _pack(indices, bits)
+    row_norms = torch.linspace(0.5, 1.5, out_features)
+    centroids = torch.tanh(torch.linspace(-1.7, 1.7, levels))
+    bias = torch.randn(out_features)
+
+    outputs = {}
+    for isa in _runnable_cpu_isas():
+        monkeypatch.setenv("ORBITQUANT_CPU_ISA", isa)
+        outputs[isa] = matmul_packed_weight(
+            x,
+            packed,
+            row_norms,
+            centroids,
+            bits=bits,
+            out_features=out_features,
+            in_features=in_features,
+            bias=bias,
+        )
+
+    for isa in _runnable_cpu_isas()[1:]:
+        torch.testing.assert_close(outputs[isa], outputs["scalar"], atol=2e-5, rtol=2e-5)
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.parametrize("rows", [16, 24, 32, 33])
+def test_cpu_avx2_bf16_realistic_row_tiles_match_reference(monkeypatch, rows: int) -> None:
+    if "avx2" not in _runnable_cpu_isas() or not supports_device("cpu"):
+        pytest.skip("the built variant has no runnable AVX2 CPU path")
+    monkeypatch.setenv("ORBITQUANT_CPU_ISA", "avx2")
+    torch.manual_seed(59 + rows)
+    in_features = 1536
+    out_features = 17
+    x = torch.randn(rows, in_features, dtype=torch.bfloat16)
+    indices = torch.randint(0, 16, (out_features, in_features), dtype=torch.uint8)
+    packed = _pack(indices, 4)
+    row_norms = torch.linspace(0.5, 1.5, out_features)
+    centroids = torch.tanh(torch.linspace(-1.7, 1.7, 16))
+
+    expected_weight = (row_norms[:, None] * centroids[indices.long()]).to(torch.bfloat16)
+    expected = torch.nn.functional.linear(x, expected_weight)
+    actual = matmul_packed_weight(
+        x,
+        packed,
+        row_norms,
+        centroids,
+        bits=4,
+        out_features=out_features,
+        in_features=in_features,
+    )
+
+    torch.testing.assert_close(actual, expected, atol=0.25, rtol=3e-2)
+
+
+@pytest.mark.kernels_ci
+def test_cpu_avx512_bf16_realistic_row_tile_matches_reference(monkeypatch) -> None:
+    if "avx512" not in _runnable_cpu_isas() or not supports_device("cpu"):
+        pytest.skip("the built variant has no runnable AVX-512 CPU path")
+    monkeypatch.setenv("ORBITQUANT_CPU_ISA", "avx512")
+    torch.manual_seed(59)
+    rows = 32
+    in_features = 1536
+    out_features = 17
+    x = torch.randn(rows, in_features, dtype=torch.bfloat16)
+    indices = torch.randint(0, 16, (out_features, in_features), dtype=torch.uint8)
+    packed = _pack(indices, 4)
+    row_norms = torch.linspace(0.5, 1.5, out_features)
+    centroids = torch.tanh(torch.linspace(-1.7, 1.7, 16))
+
+    expected_weight = (row_norms[:, None] * centroids[indices.long()]).to(torch.bfloat16)
+    expected = torch.nn.functional.linear(x, expected_weight)
+    actual = matmul_packed_weight(
+        x,
+        packed,
+        row_norms,
+        centroids,
+        bits=4,
+        out_features=out_features,
+        in_features=in_features,
+    )
+
+    torch.testing.assert_close(actual, expected, atol=0.125, rtol=3e-2)
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.parametrize("group_size", [8, 64])
+@pytest.mark.parametrize("with_bias", [False, True])
+def test_matmul_packed_adaln_cpu_matches_independent_bf16_reference(
+    group_size: int,
+    with_bias: bool,
+) -> None:
+    if not supports_cpu_adaln():
+        pytest.skip("the built variant has no native CPU AdaLN kernel")
+    torch.manual_seed(47)
+    in_features = 65
+    out_features = 9
+    num_groups = (in_features + group_size - 1) // group_size
+    padded_in_features = num_groups * group_size
+    indices = torch.randint(
+        0,
+        16,
+        (out_features, num_groups, group_size),
+        dtype=torch.uint8,
+    )
+    indices.reshape(out_features, padded_in_features)[:, in_features:] = 8
+    packed = _pack(indices, 4)
+    scales = torch.rand(out_features, num_groups, dtype=torch.bfloat16).mul(0.1)
+    x = torch.randn(2, 3, in_features, dtype=torch.bfloat16)
+    bias = torch.randn(out_features, dtype=torch.bfloat16) if with_bias else None
+
+    signed = indices.to(torch.int16).sub(8).float()
+    weight = (signed * scales.float()[..., None]).reshape(out_features, padded_in_features)[
+        :, :in_features
+    ]
+    expected = torch.nn.functional.linear(x, weight.to(torch.bfloat16), bias)
+    actual = matmul_packed_adaln_int4_cpu(
+        x,
+        packed,
+        scales,
+        out_features=out_features,
+        in_features=in_features,
+        group_size=group_size,
+        bias=bias,
+    )
+
+    assert actual.dtype == torch.bfloat16
+    assert actual.shape == (2, 3, out_features)
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
+
+
+@pytest.mark.kernels_ci
+def test_cpu_adaln_runtime_isa_dispatch_matches_scalar_reference(monkeypatch) -> None:
+    if not supports_cpu_adaln():
+        pytest.skip("the built variant has no native CPU AdaLN kernel")
+    torch.manual_seed(53)
+    rows = 8
+    in_features = 64
+    out_features = 11
+    group_size = 64
+    indices = torch.randint(
+        0,
+        16,
+        (out_features, 1, group_size),
+        dtype=torch.uint8,
+    )
+    packed = _pack(indices, 4)
+    scales = torch.rand(out_features, 1, dtype=torch.bfloat16).mul(0.1)
+    x = torch.randn(rows, in_features, dtype=torch.bfloat16)
+    bias = torch.randn(out_features, dtype=torch.bfloat16)
+
+    outputs = {}
+    for isa in _runnable_cpu_isas():
+        monkeypatch.setenv("ORBITQUANT_CPU_ISA", isa)
+        outputs[isa] = matmul_packed_adaln_int4_cpu(
+            x,
+            packed,
+            scales,
+            out_features=out_features,
+            in_features=in_features,
+            group_size=group_size,
+            bias=bias,
+        )
+
+    for isa in _runnable_cpu_isas()[1:]:
+        torch.testing.assert_close(outputs[isa], outputs["scalar"], atol=3e-2, rtol=3e-2)
+
+
 @pytest.mark.kernels_ci
 @pytest.mark.parametrize("bits", [2, 3, 4, 6])
 @pytest.mark.parametrize("in_features", [16, 19])
@@ -53,14 +374,11 @@ def test_matmul_packed_weight_matches_dequantized_reference(
     rows = 9
     out_features = 7
     x = torch.randn(rows, in_features, device=device, dtype=dtype)
-    indices = (
-        torch.arange(out_features * in_features, dtype=torch.uint8).reshape(
-            out_features, in_features
-        )
-        % (2**bits)
-    )
+    indices = torch.arange(out_features * in_features, dtype=torch.uint8).reshape(
+        out_features, in_features
+    ) % (2**bits)
     packed = _pack(indices, bits).to(device)
-    row_norms = torch.linspace(0.5, 1.5, out_features, device=device)
+    row_norms = _row_norms(device, out_features)
     centroids = torch.linspace(-1.0, 1.0, 2**bits, device=device)
     bias = torch.randn(out_features, device=device, dtype=dtype) if with_bias else None
 
@@ -87,6 +405,53 @@ def test_matmul_packed_weight_matches_dequantized_reference(
     assert torch.allclose(actual.float().cpu(), expected, atol=2e-2, rtol=2e-2)
 
 
+@pytest.mark.kernels_ci
+@pytest.mark.parametrize("bits", [2, 3, 4, 6])
+@pytest.mark.parametrize("rows", [1, 2, 3, 8, 9, 15])
+@pytest.mark.parametrize(
+    ("in_features", "out_features"),
+    [(32, 32), (37, 29)],
+)
+def test_matmul_packed_weight_short_sequence_matches_reference(
+    bits: int,
+    rows: int,
+    in_features: int,
+    out_features: int,
+) -> None:
+    torch.manual_seed(1000 + bits * 100 + rows * 10 + in_features + out_features)
+    device = _device()
+    dtype = torch.float16 if device == "mps" else torch.bfloat16
+    x = torch.randn(rows, in_features, device=device, dtype=dtype)
+    indices = torch.randint(0, 2**bits, (out_features, in_features), dtype=torch.uint8)
+    packed = _pack(indices, bits).to(device)
+    row_norms = _row_norms(device, out_features)
+    centroids = torch.linspace(-1.0, 1.0, 2**bits, device=device)
+    bias = torch.randn(out_features, device=device, dtype=dtype)
+
+    expected_weight = row_norms.cpu()[:, None] * centroids.cpu()[indices.long()]
+    if device == "cpu":
+        expected = torch.nn.functional.linear(
+            x.cpu(),
+            expected_weight.to(torch.bfloat16),
+            bias.cpu(),
+        ).float()
+    else:
+        expected = torch.nn.functional.linear(x.float().cpu(), expected_weight, bias.float().cpu())
+    actual = matmul_packed_weight(
+        x,
+        packed,
+        row_norms,
+        centroids,
+        bits=bits,
+        out_features=out_features,
+        in_features=in_features,
+        bias=bias,
+    )
+
+    assert actual.shape == (rows, out_features)
+    assert torch.allclose(actual.float().cpu(), expected, atol=3e-2, rtol=3e-2)
+
+
 @pytest.mark.kernels_ci
 def test_matmul_packed_weight_explicit_mps_path_matches_dequantized_reference() -> None:
     device = _mps_device()
@@ -95,14 +460,11 @@ def test_matmul_packed_weight_explicit_mps_path_matches_dequantized_reference()
     in_features = 19
     out_features = 7
     x = torch.randn(rows, in_features, device=device, dtype=torch.float16)
-    indices = (
-        torch.arange(out_features * in_features, dtype=torch.uint8).reshape(
-            out_features, in_features
-        )
-        % (2**bits)
-    )
+    indices = torch.arange(out_features * in_features, dtype=torch.uint8).reshape(
+        out_features, in_features
+    ) % (2**bits)
     packed = _pack(indices, bits).to(device)
-    row_norms = torch.linspace(0.5, 1.5, out_features, device=device)
+    row_norms = _row_norms(device, out_features)
     centroids = torch.linspace(-1.0, 1.0, 2**bits, device=device)
     bias = torch.randn(out_features, device=device, dtype=torch.float16)
 
@@ -136,14 +498,11 @@ def test_matmul_packed_weight_explicit_mps_bfloat16_path_matches_dequantized_ref
     in_features = 19
     out_features = 7
     x = torch.randn(rows, in_features, device=device, dtype=torch.bfloat16)
-    indices = (
-        torch.arange(out_features * in_features, dtype=torch.uint8).reshape(
-            out_features, in_features
-        )
-        % (2**bits)
-    )
+    indices = torch.arange(out_features * in_features, dtype=torch.uint8).reshape(
+        out_features, in_features
+    ) % (2**bits)
     packed = _pack(indices, bits).to(device)
-    row_norms = torch.linspace(0.5, 1.5, out_features, device=device)
+    row_norms = _row_norms(device, out_features)
     centroids = torch.linspace(-1.0, 1.0, 2**bits, device=device)
     bias = torch.randn(out_features, device=device, dtype=torch.bfloat16)
 
@@ -167,3 +526,344 @@ def test_matmul_packed_weight_explicit_mps_bfloat16_path_matches_dequantized_ref
     assert actual.dtype == torch.bfloat16
     assert actual.shape == (rows, out_features)
     assert torch.allclose(actual.float().cpu(), expected, atol=3e-2, rtol=3e-2)
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.parametrize("bits", [2, 3, 4, 6])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_matmul_packed_weight_mps_aligned_mma_path_matches_dequantized_reference(
+    bits: int,
+    dtype: torch.dtype,
+) -> None:
+    device = _mps_bfloat16_device() if dtype == torch.bfloat16 else _mps_device()
+    rows = in_features = out_features = 32
+    x = torch.randn(rows, in_features, device=device, dtype=dtype)
+    indices = torch.arange(out_features * in_features, dtype=torch.uint8).reshape(
+        out_features, in_features
+    ) % (2**bits)
+    packed = _pack(indices, bits).to(device)
+    row_norms = _row_norms(device, out_features)
+    centroids = torch.linspace(-1.0, 1.0, 2**bits, device=device)
+    bias = torch.randn(out_features, device=device, dtype=dtype)
+
+    expected_weight = row_norms.cpu()[:, None] * centroids.cpu()[indices.long()]
+    expected = torch.nn.functional.linear(x.float().cpu(), expected_weight, bias.float().cpu())
+    actual = matmul_packed_weight(
+        x,
+        packed,
+        row_norms,
+        centroids,
+        bits=bits,
+        out_features=out_features,
+        in_features=in_features,
+        bias=bias,
+    )
+
+    tolerance = 3e-2 if dtype == torch.bfloat16 else 2e-2
+    assert actual.dtype == dtype
+    assert torch.allclose(actual.float().cpu(), expected, atol=tolerance, rtol=tolerance)
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.parametrize("bits", [2, 3, 4, 6])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_matmul_packed_weight_cuda_mma64_path_matches_dequantized_reference(
+    bits: int,
+    dtype: torch.dtype,
+) -> None:
+    device = _cuda_device()
+    rows = 65
+    in_features = 64
+    out_features = 70
+    x = torch.randn(rows, in_features, device=device, dtype=dtype)
+    indices = torch.arange(out_features * in_features, dtype=torch.uint8).reshape(
+        out_features, in_features
+    ) % (2**bits)
+    packed = _pack(indices, bits).to(device)
+    row_norms = _row_norms(device, out_features)
+    centroids = torch.linspace(-1.0, 1.0, 2**bits, device=device)
+    bias = torch.randn(out_features, device=device, dtype=dtype)
+
+    expected_weight = (row_norms[:, None] * centroids[indices.long().to(device)]).to(dtype)
+    expected = torch.nn.functional.linear(x, expected_weight, bias)
+    actual = matmul_packed_weight(
+        x,
+        packed,
+        row_norms,
+        centroids,
+        bits=bits,
+        out_features=out_features,
+        in_features=in_features,
+        bias=bias,
+    )
+
+    tolerance = 3e-2 if dtype == torch.bfloat16 else 2e-2
+    assert actual.dtype == dtype
+    assert torch.allclose(actual.float(), expected.float(), atol=tolerance, rtol=tolerance)
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.parametrize(
+    ("rows", "out_features", "tile_m", "tile_n"),
+    [
+        (128, 128, 128, 128),
+        (256, 128, 256, 128),
+        (128, 256, 128, 256),
+        (130, 258, 128, 128),
+    ],
+)
+@pytest.mark.parametrize("weight_k_major", [False, True])
+def test_matmul_packed_w4a4_async_matches_sync_and_float_reference(
+    rows: int,
+    out_features: int,
+    tile_m: int,
+    tile_n: int,
+    weight_k_major: bool,
+) -> None:
+    device = _cuda_device()
+    in_features = 256
+    torch.manual_seed(0)
+    activation_indices = torch.randint(0, 16, (rows, in_features), device=device, dtype=torch.uint8)
+    weight_indices = torch.randint(
+        0, 16, (out_features, in_features), device=device, dtype=torch.uint8
+    )
+    packed_activations = (
+        activation_indices[:, 0::2] | (activation_indices[:, 1::2] << 4)
+    ).contiguous()
+    row_major_weights = (weight_indices[:, 0::2] | (weight_indices[:, 1::2] << 4)).contiguous()
+    packed_weights = row_major_weights.T.contiguous() if weight_k_major else row_major_weights
+    codes = torch.tensor(
+        [-104, -79, -62, -48, -36, -25, -15, -5, 5, 15, 25, 36, 48, 62, 79, 104],
+        device=device,
+        dtype=torch.int8,
+    )
+    token_norms = torch.linspace(0.5, 1.0, rows, device=device)
+    row_norms = torch.linspace(0.5, 1.5, out_features, device=device, dtype=torch.bfloat16)
+    activation_scale = 0.005
+    weight_scale = 0.005
+
+    kwargs = {
+        "activation_scale": activation_scale,
+        "weight_scale": weight_scale,
+        "out_features": out_features,
+        "in_features": in_features,
+        "tile_m": tile_m,
+        "tile_n": tile_n,
+    }
+    sync = matmul_packed_w4a4_int8(
+        packed_activations,
+        packed_weights,
+        token_norms,
+        row_norms,
+        codes,
+        codes,
+        async_packed=False,
+        weight_k_major=weight_k_major,
+        **kwargs,
+    )
+    asynchronous = matmul_packed_w4a4_int8(
+        packed_activations,
+        packed_weights,
+        token_norms,
+        row_norms,
+        codes,
+        codes,
+        async_packed=True,
+        weight_k_major=weight_k_major,
+        **kwargs,
+    )
+
+    assert torch.equal(asynchronous, sync)
+
+    activation_values = codes[activation_indices.long()].float()
+    weight_values = codes[weight_indices.long()].float()
+    reference = activation_values @ weight_values.T
+    reference *= token_norms[:, None]
+    reference *= row_norms.float()[None, :]
+    reference *= activation_scale * weight_scale
+
+    assert torch.allclose(asynchronous.float(), reference, atol=0.25, rtol=1e-2)
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.parametrize(("dim", "threads"), [(512, 128), (4096, 256), (16384, 512)])
+def test_quantize_activations_packed_w4_matches_torch_reference(
+    dim: int,
+    threads: int,
+) -> None:
+    device = _cuda_device()
+    rows = 2
+    x = torch.zeros((rows, dim), device=device, dtype=torch.bfloat16)
+    x[0, 3] = 1
+    x[1, dim - 7] = -1
+    permutation = torch.randperm(dim, device=device)
+    signs = torch.where(
+        torch.arange(dim, device=device) % 2 == 0,
+        torch.ones(dim, device=device, dtype=torch.int8),
+        -torch.ones(dim, device=device, dtype=torch.int8),
+    )
+    boundaries = torch.linspace(-0.2, 0.2, 15, device=device)
+    eps = 1e-12
+    inv_sqrt_block = dim**-0.5
+
+    packed, norms = quantize_activations_packed_w4(
+        x,
+        permutation,
+        signs,
+        boundaries,
+        eps=eps,
+        inv_sqrt_block=inv_sqrt_block,
+        threads=threads,
+    )
+
+    work = x.float()[:, permutation] * signs.float()
+    expected_norms = work.norm(dim=-1)
+    work /= expected_norms[:, None] + eps
+    width = 1
+    while width < dim:
+        blocks = work.reshape(rows, -1, width * 2)
+        left = blocks[..., :width]
+        right = blocks[..., width:]
+        work = torch.cat((left + right, left - right), dim=-1).reshape(rows, dim)
+        width *= 2
+    indices = torch.bucketize(work * inv_sqrt_block, boundaries).to(torch.uint8)
+    expected_packed = (indices[:, 0::2] | (indices[:, 1::2] << 4)).contiguous()
+
+    assert torch.equal(packed, expected_packed)
+    assert torch.allclose(norms, expected_norms, atol=1e-6, rtol=1e-6)
+
+
+@pytest.mark.kernels_ci
+def test_quantize_activations_packed_w4_accepts_int32_permutation() -> None:
+    device = _cuda_device()
+    dim = 512
+    torch.manual_seed(0)
+    x = torch.randn((3, dim), device=device, dtype=torch.bfloat16)
+    permutation = torch.randperm(dim, device=device)
+    signs = torch.where(
+        torch.arange(dim, device=device) % 2 == 0,
+        torch.ones(dim, device=device, dtype=torch.int8),
+        -torch.ones(dim, device=device, dtype=torch.int8),
+    )
+    boundaries = torch.linspace(-0.2, 0.2, 15, device=device)
+
+    packed_int64, norms_int64 = quantize_activations_packed_w4(
+        x,
+        permutation,
+        signs,
+        boundaries,
+        eps=1e-12,
+        inv_sqrt_block=dim**-0.5,
+        threads=256,
+    )
+    packed_int32, norms_int32 = quantize_activations_packed_w4(
+        x,
+        permutation.to(torch.int32),
+        signs,
+        boundaries,
+        eps=1e-12,
+        inv_sqrt_block=dim**-0.5,
+        threads=256,
+    )
+
+    assert torch.equal(packed_int32, packed_int64)
+    assert torch.equal(norms_int32, norms_int64)
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.parametrize(("dim", "threads"), [(512, 128), (4096, 256), (16384, 512)])
+def test_quantize_activations_int8_matches_packed_codes(
+    dim: int,
+    threads: int,
+) -> None:
+    device = _cuda_device()
+    x = torch.randn((2, dim), device=device, dtype=torch.bfloat16)
+    permutation = torch.randperm(dim, device=device)
+    signs = torch.where(
+        torch.arange(dim, device=device) % 2 == 0,
+        torch.ones(dim, device=device, dtype=torch.int8),
+        -torch.ones(dim, device=device, dtype=torch.int8),
+    )
+    boundaries = torch.linspace(-0.2, 0.2, 15, device=device)
+    codes = torch.tensor(
+        [-120, -92, -68, -49, -34, -22, -12, -4, 4, 12, 22, 34, 49, 68, 92, 120],
+        device=device,
+        dtype=torch.int8,
+    )
+    kwargs = {
+        "eps": 1e-12,
+        "inv_sqrt_block": dim**-0.5,
+        "threads": threads,
+    }
+
+    packed, packed_norms = quantize_activations_packed_w4(
+        x,
+        permutation,
+        signs,
+        boundaries,
+        **kwargs,
+    )
+    quantized, int8_norms = quantize_activations_int8(
+        x,
+        permutation,
+        signs,
+        boundaries,
+        codes,
+        **kwargs,
+    )
+
+    indices = torch.empty((2, dim), device=device, dtype=torch.long)
+    indices[:, 0::2] = packed & 15
+    indices[:, 1::2] = packed >> 4
+    expected = codes[indices]
+
+    assert torch.equal(quantized, expected)
+    assert torch.equal(int8_norms, packed_norms)
+
+
+@pytest.mark.kernels_ci
+def test_quantize_activations_int8_matches_blocked_rpbh_reference() -> None:
+    device = _cuda_device()
+    rows = 2
+    dim = 12288
+    block_size = 4096
+    x = torch.randn((rows, dim), device=device, dtype=torch.bfloat16)
+    permutation = torch.randperm(dim, device=device)
+    signs = torch.where(
+        torch.arange(dim, device=device) % 2 == 0,
+        torch.ones(dim, device=device, dtype=torch.int8),
+        -torch.ones(dim, device=device, dtype=torch.int8),
+    )
+    boundaries = torch.linspace(-0.2, 0.2, 15, device=device)
+    codes = torch.tensor(
+        [-120, -92, -68, -49, -34, -22, -12, -4, 4, 12, 22, 34, 49, 68, 92, 120],
+        device=device,
+        dtype=torch.int8,
+    )
+
+    quantized, norms = quantize_activations_int8(
+        x,
+        permutation,
+        signs,
+        boundaries,
+        codes,
+        eps=1e-12,
+        inv_sqrt_block=block_size**-0.5,
+        threads=512,
+    )
+
+    work = x.float()[:, permutation] * signs.float()
+    expected_norms = work.norm(dim=-1)
+    work = (work / expected_norms[:, None]).reshape(rows, -1, block_size)
+    width = 1
+    while width < block_size:
+        blocks = work.reshape(rows, -1, width * 2)
+        left = blocks[..., :width]
+        right = blocks[..., width:]
+        work = torch.cat((left + right, left - right), dim=-1).reshape(rows, -1, block_size)
+        width *= 2
+    indices = torch.bucketize(work.reshape(rows, dim) * (block_size**-0.5), boundaries)
+    expected = codes[indices]
+
+    assert torch.equal(quantized, expected)
+    assert torch.allclose(norms, expected_norms, atol=1e-5, rtol=1e-6)

torch-ext/orbitquant_packed_matmul/__init__.py

@@ -4,7 +4,106 @@ import torch
 
 from ._ops import ops
 
-__all__ = ["matmul_packed_weight"]
+__all__ = [
+    "matmul_packed_w4a4_int8",
+    "matmul_packed_adaln_int4_cpu",
+    "matmul_packed_weight",
+    "quantize_activations_cpu",
+    "quantize_activations_int8",
+    "quantize_activations_packed_w4",
+    "supports_cpu_activation",
+    "supports_cpu_adaln",
+    "supports_device",
+]
+
+
+def supports_device(device_type: str) -> bool:
+    dispatch_keys = {
+        "cpu": "CPU",
+        "cuda": "CUDA",
+        "mps": "MPS",
+    }
+    try:
+        dispatch_key = dispatch_keys[device_type]
+    except KeyError as exc:
+        raise ValueError(f"unknown device type {device_type!r}") from exc
+    return bool(
+        torch._C._dispatch_has_kernel_for_dispatch_key(  # noqa: SLF001
+            ops.matmul_packed_weight._qualified_op_name,  # noqa: SLF001
+            dispatch_key,
+        )
+    )
+
+
+def supports_cpu_activation() -> bool:
+    try:
+        operation = ops.quantize_activations_cpu
+        qualified_name = operation._qualified_op_name  # noqa: SLF001
+    except (AttributeError, RuntimeError):
+        return False
+    return bool(
+        torch._C._dispatch_has_kernel_for_dispatch_key(  # noqa: SLF001
+            qualified_name,
+            "CPU",
+        )
+    )
+
+
+def supports_cpu_adaln() -> bool:
+    try:
+        operation = ops.matmul_packed_adaln_int4_cpu
+        qualified_name = operation._qualified_op_name  # noqa: SLF001
+    except (AttributeError, RuntimeError):
+        return False
+    return bool(
+        torch._C._dispatch_has_kernel_for_dispatch_key(  # noqa: SLF001
+            qualified_name,
+            "CPU",
+        )
+    )
+
+
+def quantize_activations_cpu(
+    x: torch.Tensor,
+    permutation: torch.Tensor,
+    signs: torch.Tensor,
+    centroids: torch.Tensor,
+    boundaries: torch.Tensor,
+    *,
+    eps: float,
+    inv_sqrt_block: float,
+    block_size: int,
+) -> torch.Tensor:
+    if x.device.type != "cpu":
+        raise RuntimeError("native CPU activation quantization requires CPU tensors")
+    if x.dtype not in {torch.float32, torch.float16, torch.bfloat16}:
+        raise ValueError("x must be float32, float16, or bfloat16")
+    if block_size <= 0 or block_size & (block_size - 1):
+        raise ValueError("block_size must be a positive power of two")
+    dim = x.shape[-1]
+    if dim % block_size != 0:
+        raise ValueError("block_size must divide the input dimension")
+
+    original_shape = x.shape
+    values = x.contiguous().reshape(-1, dim)
+    out = torch.empty_like(values)
+    permutation_dtype = torch.int32 if permutation.dtype == torch.int32 else torch.int64
+    permutation_values = permutation.to(device="cpu", dtype=permutation_dtype).contiguous()
+    sign_values = signs.to(device="cpu", dtype=torch.int8).contiguous()
+    centroid_values = centroids.to(device="cpu", dtype=torch.float32).contiguous()
+    boundary_values = boundaries.to(device="cpu", dtype=torch.float32).contiguous()
+    ops.quantize_activations_cpu(
+        out,
+        values,
+        permutation_values,
+        sign_values,
+        centroid_values,
+        boundary_values,
+        eps,
+        inv_sqrt_block,
+        block_size,
+    )
+    return out.reshape(original_shape)
 
 
 def matmul_packed_weight(
@@ -17,9 +116,9 @@ def matmul_packed_weight(
     out_features: int,
     in_features: int,
     bias: torch.Tensor | None = None,
-    block_m: int = 32,
+    block_m: int = 64,
     block_n: int = 64,
-    block_k: int = 64,
+    block_k: int = 128,
 ) -> torch.Tensor:
     if bits <= 0 or bits > 8:
         raise ValueError("bits must be in [1, 8]")
@@ -32,13 +131,15 @@ def matmul_packed_weight(
     x_2d = x.contiguous().reshape(-1, in_features)
     out = torch.empty((x_2d.shape[0], out_features), device=x.device, dtype=x.dtype)
     packed = packed_weight_indices.to(device=x.device, dtype=torch.uint8).contiguous()
-    norms = row_norms.to(device=x.device, dtype=torch.float32).contiguous()
+    auxiliary_dtype = torch.bfloat16 if x.device.type == "cuda" else torch.float32
+    norms = row_norms.to(device=x.device, dtype=auxiliary_dtype).contiguous()
     centroid_values = centroids.to(device=x.device, dtype=torch.float32).contiguous()
     if bias is None:
         bias_values = torch.empty((1,), device=x.device, dtype=torch.float32)
         has_bias = False
     else:
-        bias_values = bias.to(device=x.device, dtype=torch.float32).contiguous()
+        bias_dtype = x.dtype if x.device.type == "cuda" else torch.float32
+        bias_values = bias.to(device=x.device, dtype=bias_dtype).contiguous()
         has_bias = True
 
     ops.matmul_packed_weight(
@@ -57,3 +158,234 @@ def matmul_packed_weight(
         block_k,
     )
     return out.reshape(*original_shape[:-1], out_features)
+
+
+def matmul_packed_adaln_int4_cpu(
+    x: torch.Tensor,
+    packed_weight: torch.Tensor,
+    scales: torch.Tensor,
+    *,
+    out_features: int,
+    in_features: int,
+    group_size: int,
+    bias: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if x.device.type != "cpu":
+        raise RuntimeError("native packed AdaLN requires CPU tensors")
+    if x.shape[-1] != in_features:
+        raise ValueError(f"expected input last dimension {in_features}, got {x.shape[-1]}")
+    if group_size <= 0:
+        raise ValueError("group_size must be positive")
+
+    original_shape = x.shape
+    x_2d = x.to(dtype=torch.bfloat16).contiguous().reshape(-1, in_features)
+    out = torch.empty((x_2d.shape[0], out_features), dtype=torch.bfloat16)
+    packed = packed_weight.to(device="cpu", dtype=torch.uint8).contiguous()
+    scale_values = scales.to(device="cpu", dtype=torch.float32).contiguous()
+    if bias is None:
+        bias_values = torch.empty((1,), dtype=torch.float32)
+        has_bias = False
+    else:
+        bias_values = (
+            bias.to(device="cpu", dtype=torch.bfloat16).to(dtype=torch.float32).contiguous()
+        )
+        has_bias = True
+    ops.matmul_packed_adaln_int4_cpu(
+        out,
+        x_2d,
+        packed,
+        scale_values,
+        bias_values,
+        has_bias,
+        out_features,
+        in_features,
+        group_size,
+    )
+    return out.reshape(*original_shape[:-1], out_features)
+
+
+def matmul_packed_w4a4_int8(
+    packed_activations: torch.Tensor,
+    packed_weight_indices: torch.Tensor,
+    token_norms: torch.Tensor,
+    row_norms: torch.Tensor,
+    activation_codes: torch.Tensor,
+    weight_codes: torch.Tensor,
+    *,
+    activation_scale: float,
+    weight_scale: float,
+    out_features: int,
+    in_features: int,
+    bias: torch.Tensor | None = None,
+    output_dtype: torch.dtype = torch.bfloat16,
+    tile_m: int = 128,
+    tile_n: int = 128,
+    async_packed: bool = False,
+    weight_k_major: bool = False,
+) -> torch.Tensor:
+    if not packed_activations.is_cuda:
+        raise RuntimeError("packed W4A4 INT8 matmul requires CUDA tensors")
+    if in_features <= 0 or in_features % 64 != 0:
+        raise ValueError("in_features must be positive and divisible by 64")
+    if packed_activations.shape[-1] != in_features // 2:
+        raise ValueError(
+            f"expected packed activation last dimension {in_features // 2}, "
+            f"got {packed_activations.shape[-1]}"
+        )
+    if output_dtype not in {torch.float16, torch.bfloat16}:
+        raise ValueError("output_dtype must be float16 or bfloat16")
+    if (tile_m, tile_n) not in {(128, 128), (256, 128), (128, 256)}:
+        raise ValueError("tile must be 128x128, 256x128, or 128x256")
+
+    original_shape = packed_activations.shape
+    activations = (
+        packed_activations.to(dtype=torch.uint8).contiguous().reshape(-1, in_features // 2)
+    )
+    weights = packed_weight_indices.to(device=activations.device, dtype=torch.uint8).contiguous()
+    norms = token_norms.to(device=activations.device, dtype=torch.float32).contiguous()
+    weight_norms = row_norms.to(device=activations.device, dtype=torch.bfloat16).contiguous()
+    activation_code_values = activation_codes.to(
+        device=activations.device, dtype=torch.int8
+    ).contiguous()
+    weight_code_values = weight_codes.to(device=activations.device, dtype=torch.int8).contiguous()
+    out = torch.empty(
+        (activations.shape[0], out_features),
+        device=activations.device,
+        dtype=output_dtype,
+    )
+    if bias is None:
+        bias_values = out
+        has_bias = False
+    else:
+        bias_values = bias.to(device=activations.device, dtype=output_dtype).contiguous()
+        has_bias = True
+
+    ops.matmul_packed_w4a4_int8(
+        out,
+        activations,
+        weights,
+        norms,
+        weight_norms,
+        activation_code_values,
+        weight_code_values,
+        bias_values,
+        has_bias,
+        activation_scale,
+        weight_scale,
+        out_features,
+        in_features,
+        tile_m,
+        tile_n,
+        async_packed,
+        weight_k_major,
+    )
+    return out.reshape(*original_shape[:-1], out_features)
+
+
+def quantize_activations_packed_w4(
+    x: torch.Tensor,
+    permutation: torch.Tensor,
+    signs: torch.Tensor,
+    boundaries: torch.Tensor,
+    *,
+    eps: float = 1e-12,
+    inv_sqrt_block: float,
+    threads: int = 256,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if not x.is_cuda:
+        raise RuntimeError("native packed W4 activation quantization requires CUDA tensors")
+    if x.dtype not in {torch.float16, torch.bfloat16}:
+        raise ValueError("x must be float16 or bfloat16")
+    dim = x.shape[-1]
+    if dim not in {512, 1024, 2048, 4096, 8192, 16384}:
+        raise ValueError(
+            "native packed W4 activation quantization supports dimensions "
+            "512, 1024, 2048, 4096, 8192, and 16384"
+        )
+    if threads not in {128, 256, 512}:
+        raise ValueError("threads must be 128, 256, or 512")
+    if permutation.numel() != dim or signs.numel() != dim:
+        raise ValueError("permutation and signs must match the input dimension")
+    if boundaries.numel() != 15:
+        raise ValueError("boundaries must contain 15 values")
+
+    original_shape = x.shape
+    values = x.contiguous().reshape(-1, dim)
+    packed = torch.empty((values.shape[0], dim // 2), device=x.device, dtype=torch.uint8)
+    norms = torch.empty(values.shape[0], device=x.device, dtype=torch.float32)
+    permutation_dtype = torch.int32 if permutation.dtype == torch.int32 else torch.int64
+    permutation_values = permutation.to(device=x.device, dtype=permutation_dtype).contiguous()
+    sign_values = signs.to(device=x.device, dtype=torch.int8).contiguous()
+    boundary_values = boundaries.to(device=x.device, dtype=torch.float32).contiguous()
+    ops.quantize_activations_packed_w4(
+        packed,
+        norms,
+        values,
+        permutation_values,
+        sign_values,
+        boundary_values,
+        eps,
+        inv_sqrt_block,
+        threads,
+    )
+    return (
+        packed.reshape(*original_shape[:-1], dim // 2),
+        norms.reshape(original_shape[:-1]),
+    )
+
+
+def quantize_activations_int8(
+    x: torch.Tensor,
+    permutation: torch.Tensor,
+    signs: torch.Tensor,
+    boundaries: torch.Tensor,
+    codes: torch.Tensor,
+    *,
+    eps: float = 1e-12,
+    inv_sqrt_block: float,
+    threads: int = 256,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if not x.is_cuda:
+        raise RuntimeError("native INT8 activation quantization requires CUDA tensors")
+    if x.dtype not in {torch.float16, torch.bfloat16}:
+        raise ValueError("x must be float16 or bfloat16")
+    dim = x.shape[-1]
+    if dim not in {512, 1024, 2048, 4096, 8192, 12288, 16384}:
+        raise ValueError(
+            "native INT8 activation quantization supports dimensions "
+            "512, 1024, 2048, 4096, 8192, 12288, and 16384"
+        )
+    if threads not in {128, 256, 512}:
+        raise ValueError("threads must be 128, 256, or 512")
+    if permutation.numel() != dim or signs.numel() != dim:
+        raise ValueError("permutation and signs must match the input dimension")
+    if boundaries.numel() != 15:
+        raise ValueError("boundaries must contain 15 values")
+    if codes.numel() != 16:
+        raise ValueError("codes must contain 16 values")
+
+    original_shape = x.shape
+    values = x.contiguous().reshape(-1, dim)
+    quantized = torch.empty(values.shape, device=x.device, dtype=torch.int8)
+    norms = torch.empty(values.shape[0], device=x.device, dtype=torch.float32)
+    permutation_dtype = torch.int32 if permutation.dtype == torch.int32 else torch.int64
+    permutation_values = permutation.to(device=x.device, dtype=permutation_dtype).contiguous()
+    sign_values = signs.to(device=x.device, dtype=torch.int8).contiguous()
+    boundary_values = boundaries.to(device=x.device, dtype=torch.float32).contiguous()
+    code_values = codes.to(device=x.device, dtype=torch.int8).contiguous()
+    ops.quantize_activations_int8(
+        quantized,
+        norms,
+        values,
+        permutation_values,
+        sign_values,
+        boundary_values,
+        code_values,
+        eps,
+        inv_sqrt_block,
+        threads,
+    )
+    return (
+        quantized.reshape(*original_shape[:-1], dim),
+        norms.reshape(original_shape[:-1]),
+    )

torch-ext/torch_binding.cpp

@@ -1,8 +1,34 @@
-#include <torch/library.h>
-
 #include "registration.h"
 #include "torch_binding.h"
 
+#if defined(CPU_KERNEL)
+#include <torch/csrc/stable/library.h>
+
+STABLE_TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def(
+      "matmul_packed_weight(Tensor! out, Tensor x, Tensor packed_weight_indices, "
+      "Tensor row_norms, Tensor centroids, Tensor bias, bool has_bias, int bits, "
+      "int out_features, int in_features, int block_m, int block_n, int block_k) -> ()");
+  ops.def(
+      "quantize_activations_cpu(Tensor! out, Tensor x, Tensor permutation, "
+      "Tensor signs, Tensor centroids, Tensor boundaries, float eps, "
+      "float inv_sqrt_block, int block_size) -> ()");
+  ops.def(
+      "matmul_packed_adaln_int4_cpu(Tensor! out, Tensor x, Tensor packed_weight, "
+      "Tensor scales, Tensor bias, bool has_bias, int out_features, "
+      "int in_features, int group_size) -> ()");
+}
+
+STABLE_TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CPU, ops) {
+  ops.impl("matmul_packed_weight", TORCH_BOX(&matmul_packed_weight));
+  ops.impl("quantize_activations_cpu", TORCH_BOX(&quantize_activations_cpu));
+  ops.impl(
+      "matmul_packed_adaln_int4_cpu",
+      TORCH_BOX(&matmul_packed_adaln_int4_cpu));
+}
+#else
+#include <torch/library.h>
+
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def(
       "matmul_packed_weight(Tensor! out, Tensor x, Tensor packed_weight_indices, "
@@ -10,9 +36,34 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "int out_features, int in_features, int block_m, int block_n, int block_k) -> ()");
 #if defined(CUDA_KERNEL)
   ops.impl("matmul_packed_weight", torch::kCUDA, &matmul_packed_weight);
+  ops.def(
+      "matmul_packed_w4a4_int8(Tensor! out, Tensor packed_activations, "
+      "Tensor packed_weight_indices, Tensor token_norms, Tensor row_norms, "
+      "Tensor activation_codes, Tensor weight_codes, Tensor bias, bool has_bias, "
+      "float activation_scale, float weight_scale, int out_features, "
+      "int in_features, int tile_m, int tile_n, bool async_packed, "
+      "bool weight_k_major) -> ()");
+  ops.impl("matmul_packed_w4a4_int8", torch::kCUDA, &matmul_packed_w4a4_int8);
+  ops.def(
+      "quantize_activations_packed_w4(Tensor! packed_out, Tensor! norms_out, "
+      "Tensor x, Tensor permutation, Tensor signs, Tensor boundaries, float eps, "
+      "float inv_sqrt_block, int threads) -> ()");
+  ops.impl(
+      "quantize_activations_packed_w4",
+      torch::kCUDA,
+      &quantize_activations_packed_w4);
+  ops.def(
+      "quantize_activations_int8(Tensor! int8_out, Tensor! norms_out, Tensor x, "
+      "Tensor permutation, Tensor signs, Tensor boundaries, Tensor codes, float eps, "
+      "float inv_sqrt_block, int threads) -> ()");
+  ops.impl(
+      "quantize_activations_int8",
+      torch::kCUDA,
+      &quantize_activations_int8);
 #elif defined(METAL_KERNEL)
   ops.impl("matmul_packed_weight", torch::kMPS, &matmul_packed_weight);
 #endif
 }
+#endif
 
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h

@@ -1,14 +1,20 @@
 #pragma once
 
+#if defined(CPU_KERNEL)
+#include <torch/csrc/stable/tensor.h>
+using OrbitQuantTensor = torch::stable::Tensor;
+#else
 #include <torch/torch.h>
+using OrbitQuantTensor = torch::Tensor;
+#endif
 
 void matmul_packed_weight(
-    torch::Tensor &out,
-    torch::Tensor const &x,
-    torch::Tensor const &packed_weight_indices,
-    torch::Tensor const &row_norms,
-    torch::Tensor const &centroids,
-    torch::Tensor const &bias,
+    OrbitQuantTensor &out,
+    OrbitQuantTensor const &x,
+    OrbitQuantTensor const &packed_weight_indices,
+    OrbitQuantTensor const &row_norms,
+    OrbitQuantTensor const &centroids,
+    OrbitQuantTensor const &bias,
     bool has_bias,
     int64_t bits,
     int64_t out_features,
@@ -16,3 +22,71 @@ void matmul_packed_weight(
     int64_t block_m,
     int64_t block_n,
     int64_t block_k);
+
+#if defined(CPU_KERNEL)
+void quantize_activations_cpu(
+    OrbitQuantTensor &out,
+    OrbitQuantTensor const &x,
+    OrbitQuantTensor const &permutation,
+    OrbitQuantTensor const &signs,
+    OrbitQuantTensor const &centroids,
+    OrbitQuantTensor const &boundaries,
+    double eps,
+    double inv_sqrt_block,
+    int64_t block_size);
+
+void matmul_packed_adaln_int4_cpu(
+    OrbitQuantTensor &out,
+    OrbitQuantTensor const &x,
+    OrbitQuantTensor const &packed_weight,
+    OrbitQuantTensor const &scales,
+    OrbitQuantTensor const &bias,
+    bool has_bias,
+    int64_t out_features,
+    int64_t in_features,
+    int64_t group_size);
+#endif
+
+#if defined(CUDA_KERNEL)
+void matmul_packed_w4a4_int8(
+    torch::Tensor &out,
+    torch::Tensor const &packed_activations,
+    torch::Tensor const &packed_weight_indices,
+    torch::Tensor const &token_norms,
+    torch::Tensor const &row_norms,
+    torch::Tensor const &activation_codes,
+    torch::Tensor const &weight_codes,
+    torch::Tensor const &bias,
+    bool has_bias,
+    double activation_scale,
+    double weight_scale,
+    int64_t out_features,
+    int64_t in_features,
+    int64_t tile_m,
+    int64_t tile_n,
+    bool async_packed,
+    bool weight_k_major);
+
+void quantize_activations_packed_w4(
+    torch::Tensor &packed_out,
+    torch::Tensor &norms_out,
+    torch::Tensor const &x,
+    torch::Tensor const &permutation,
+    torch::Tensor const &signs,
+    torch::Tensor const &boundaries,
+    double eps,
+    double inv_sqrt_block,
+    int64_t threads);
+
+void quantize_activations_int8(
+    torch::Tensor &int8_out,
+    torch::Tensor &norms_out,
+    torch::Tensor const &x,
+    torch::Tensor const &permutation,
+    torch::Tensor const &signs,
+    torch::Tensor const &boundaries,
+    torch::Tensor const &codes,
+    double eps,
+    double inv_sqrt_block,
+    int64_t threads);
+#endif