Uploaded using `kernel-builder`.

Browse files

Files changed (15) hide show

build/torch211-cxx11-cpu-x86_64-linux/__init__.py +229 -0
build/torch211-cxx11-cpu-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cpu-x86_64-linux/_quant_matmul_cpu_18a2909.abi3.so +3 -0
build/torch211-cxx11-cpu-x86_64-linux/metadata.json +30 -0
build/torch211-cxx11-cpu-x86_64-linux/quant_matmul/__init__.py +26 -0
build/torch212-cxx11-cpu-x86_64-linux/__init__.py +229 -0
build/torch212-cxx11-cpu-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cpu-x86_64-linux/_quant_matmul_cpu_18a2909.abi3.so +3 -0
build/torch212-cxx11-cpu-x86_64-linux/metadata.json +30 -0
build/torch212-cxx11-cpu-x86_64-linux/quant_matmul/__init__.py +26 -0
build/torch213-cxx11-cpu-x86_64-linux/__init__.py +229 -0
build/torch213-cxx11-cpu-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cpu-x86_64-linux/_quant_matmul_cpu_18a2909.abi3.so +3 -0
build/torch213-cxx11-cpu-x86_64-linux/metadata.json +30 -0
build/torch213-cxx11-cpu-x86_64-linux/quant_matmul/__init__.py +26 -0

build/torch211-cxx11-cpu-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,229 @@

+"""quant-matmul: W4A8 / W8A8 / FP16-storage GEMV-GEMM for CPUs, aarch64-first.
+Int4 weights in GPTQ/AWQ-compatible group quantization run against per-token
+INT8 activations on NEON (SDOT where the CPU has dotprod); an int8 W8A8 path
+and an fp16-storage GEMV round out the pack. x86-64 runs a portable scalar
+path (correct, not tuned).
+High-level API:
+    pack_w4(intw, scales, zeros, group_size)   native packing from raw nibbles
+    from_gptq(qweight, qzeros, scales, g_idx)  GPTQ checkpoint tensors -> native
+    from_awq(qweight, qzeros, scales)          AWQ checkpoint tensors -> native
+    w4a8_linear(x, packed, scales, zeros, group_size, perm)
+    w8a8_linear(x, w_int8, scale_wt)
+    fp16_gemv(x, w_f16)
+    QuantLinear                                nn.Module wrapper
+"""
+from typing import Optional, Tuple
+import torch
+import torch.nn as nn
+from ._ops import ops
+_OK_DTYPES = (torch.bfloat16, torch.float32)
+def pack_w4(intw: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor,
+            group_size: int):
+    """Pack int4 nibbles into the native layout.
+    Args:
+        intw: [N, K] integer tensor with values 0..15.
+        scales: [N, G] per-group scales (G = K // group_size).
+        zeros: [N, G] per-group zero points, already in subtract-this form.
+    Returns:
+        (packed uint8 [N, K/2], scales f32 [N, G], zeros f32 [N, G])
+    """
+    assert intw.dim() == 2, "intw must be [N, K]"
+    N, K = intw.shape
+    assert K % group_size == 0 and group_size % 32 == 0
+    w = intw.to(torch.uint8)
+    assert int(w.max()) <= 15, "nibbles must be 0..15"
+    packed = (w[:, 0::2] | (w[:, 1::2] << 4)).contiguous()
+    G = K // group_size
+    s = scales.to(torch.float32).reshape(N, G).contiguous()
+    z = zeros.to(torch.float32).reshape(N, G).contiguous()
+    return packed, s, z
+def dequant_w4(packed: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor,
+               group_size: int) -> torch.Tensor:
+    """Reference dequantization to f32 [N, K] (tests and debugging)."""
+    N = packed.shape[0]
+    K = packed.shape[1] * 2
+    nib = torch.empty(N, K, dtype=torch.float32)
+    nib[:, 0::2] = (packed & 0xF).float()
+    nib[:, 1::2] = (packed >> 4).float()
+    G = K // group_size
+    z = zeros.reshape(N, G, 1)
+    s = scales.reshape(N, G, 1)
+    return ((nib.reshape(N, G, group_size) - z) * s).reshape(N, K)
+_GPTQ_SHIFTS = torch.arange(0, 32, 4, dtype=torch.int32)
+def _unpack_int32_nibbles(t: torch.Tensor, dim_factor: int = 8) -> torch.Tensor:
+    """int32 [..., P] -> uint8 [..., P*8], nibble j of each word at index 8p+j."""
+    u = ((t.unsqueeze(-1) >> _GPTQ_SHIFTS) & 0xF).to(torch.uint8)
+    return u.reshape(*t.shape[:-1], t.shape[-1] * dim_factor)
+def from_gptq(qweight: torch.Tensor, qzeros: torch.Tensor, scales: torch.Tensor,
+              g_idx: Optional[torch.Tensor] = None, zero_offset: int = 1):
+    """GPTQ checkpoint tensors -> native packing.
+    qweight int32 [K/8, N]; qzeros int32 [G, N/8]; scales [G, N].
+    zero_offset follows the AutoGPTQ storage convention (stored zero + 1 is
+    subtracted at dequant); pass 0 for checkpoints without it.
+    Returns (packed, scales, zeros, group_size, perm) where perm is the
+    activation permutation for act-order checkpoints (None otherwise).
+    """
+    K = qweight.shape[0] * 8
+    N = qweight.shape[1]
+    G = scales.shape[0]
+    assert K % G == 0
+    group_size = K // G
+    # nibble j of qweight word [i, n] decodes k = 8i + j.
+    intw = _unpack_int32_nibbles(qweight.t().contiguous())           # [N, K]
+    z = (_unpack_int32_nibbles(qzeros).to(torch.float32) + float(zero_offset))
+    z = z.t().contiguous()                                           # [N, G]
+    s = scales.t().to(torch.float32).contiguous()                    # [N, G]
+    perm = None
+    if g_idx is not None:
+        trivial = torch.arange(K, dtype=g_idx.dtype) // group_size
+        if not torch.equal(g_idx.cpu(), trivial):
+            perm = torch.argsort(g_idx.cpu(), stable=True)
+            intw = intw[:, perm]
+    packed, s, z = pack_w4(intw, s, z, group_size)
+    return packed, s, z, group_size, perm
+_AWQ_ORDER = [0, 2, 4, 6, 1, 3, 5, 7]
+def from_awq(qweight: torch.Tensor, qzeros: torch.Tensor, scales: torch.Tensor,
+             zero_offset: int = 0):
+    """AWQ (AutoAWQ GEMM) checkpoint tensors -> native packing.
+    qweight int32 [K, N/8] with nibble order [0,2,4,6,1,3,5,7]; qzeros int32
+    [G, N/8] same order; scales [G, N].
+    Returns (packed, scales, zeros, group_size, None).
+    """
+    K = qweight.shape[0]
+    G = scales.shape[0]
+    N = scales.shape[1]
+    group_size = K // G
+    inv = torch.empty(8, dtype=torch.int64)
+    for j, o in enumerate(_AWQ_ORDER):
+        inv[o] = j
+    u = _unpack_int32_nibbles(qweight)                               # [K, N] j-order
+    u = u.reshape(K, N // 8, 8)[:, :, inv].reshape(K, N)
+    z = _unpack_int32_nibbles(qzeros).to(torch.float32)              # [G, N] j-order
+    z = z.reshape(G, N // 8, 8)[:, :, inv].reshape(G, N) + float(zero_offset)
+    packed, s, z = pack_w4(u.t().contiguous(), scales.t(), z.t(), group_size)
+    return packed, s, z, group_size, None
+def quantize_activation(x: torch.Tensor):
+    """Per-token absmax INT8: [..., K] bf16/f32 -> (int8 [M, K], f32 scale [M])."""
+    x_flat = x.reshape(-1, x.shape[-1]).contiguous()
+    M, K = x_flat.shape
+    x_int8 = torch.empty((M, K), dtype=torch.int8)
+    scale = torch.empty((M,), dtype=torch.float32)
+    ops.qm_quantize_act(x_int8, scale, x_flat)
+    return x_int8, scale
+def w4a8_linear(x: torch.Tensor, packed: torch.Tensor, scales: torch.Tensor,
+                zeros: torch.Tensor, group_size: int,
+                perm: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """W4A8 forward: [..., K] bf16/f32 -> [..., N], dtype follows x."""
+    assert x.dtype in _OK_DTYPES, f"x must be bf16 or f32, got {x.dtype}"
+    orig = x.shape
+    x_flat = x.reshape(-1, orig[-1])
+    if perm is not None:
+        x_flat = x_flat[:, perm]
+    x_flat = x_flat.contiguous()
+    M = x_flat.shape[0]
+    N = packed.shape[0]
+    out = torch.empty((M, N), dtype=x.dtype)
+    if M < 16:
+        ops.qm_w4a8_gemv_fused(out, x_flat, packed, scales, zeros, group_size)
+    else:
+        q, sa = quantize_activation(x_flat)
+        ops.qm_w4a8_gemm(out, q, packed, scales, zeros, sa, group_size)
+    return out.view(*orig[:-1], N)
+def w8a8_linear(x: torch.Tensor, w_int8: torch.Tensor,
+                scale_wt: torch.Tensor) -> torch.Tensor:
+    """W8A8 forward: per-token int8 activations x int8 weights, per-row scale."""
+    assert x.dtype in _OK_DTYPES
+    orig = x.shape
+    q, sa = quantize_activation(x.reshape(-1, orig[-1]))
+    N = w_int8.shape[0]
+    out = torch.empty((q.shape[0], N), dtype=x.dtype)
+    ops.qm_w8a8_gemm(out, q, w_int8, sa, scale_wt)
+    return out.view(*orig[:-1], N)
+def fp16_gemv(x: torch.Tensor, w_f16: torch.Tensor) -> torch.Tensor:
+    """fp16-storage GEMV: f16 weights widened to f32 in registers."""
+    orig = x.shape
+    x_flat = x.reshape(-1, orig[-1]).to(torch.float32).contiguous()
+    N = w_f16.shape[0]
+    out = torch.empty((x_flat.shape[0], N), dtype=torch.float32)
+    ops.qm_fp16_gemv(out, x_flat, w_f16)
+    return out.view(*orig[:-1], N).to(x.dtype)
+class QuantLinear(nn.Module):
+    """Inference-only W4A8 linear over the native packing."""
+    def __init__(self, packed, scales, zeros, group_size, perm=None, bias=None):
+        super().__init__()
+        self.register_buffer("packed", packed)
+        self.register_buffer("scales", scales)
+        self.register_buffer("zeros", zeros)
+        self.group_size = int(group_size)
+        if perm is not None:
+            self.register_buffer("perm", perm)
+        else:
+            self.perm = None
+        if bias is not None:
+            self.register_buffer("bias", bias)
+        else:
+            self.bias = None
+    @classmethod
+    def from_gptq_tensors(cls, qweight, qzeros, scales, g_idx=None, bias=None,
+                          zero_offset=1):
+        packed, s, z, gs, perm = from_gptq(qweight, qzeros, scales, g_idx, zero_offset)
+        return cls(packed, s, z, gs, perm, bias)
+    @classmethod
+    def from_awq_tensors(cls, qweight, qzeros, scales, bias=None, zero_offset=0):
+        packed, s, z, gs, perm = from_awq(qweight, qzeros, scales, zero_offset)
+        return cls(packed, s, z, gs, perm, bias)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = w4a8_linear(x, self.packed, self.scales, self.zeros, self.group_size,
+                        self.perm)
+        if self.bias is not None:
+            y = y + self.bias
+        return y
+__all__ = [
+    "pack_w4",
+    "dequant_w4",
+    "from_gptq",
+    "from_awq",
+    "quantize_activation",
+    "w4a8_linear",
+    "w8a8_linear",
+    "fp16_gemv",
+    "QuantLinear",
+]

build/torch211-cxx11-cpu-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _quant_matmul_cpu_18a2909
+ops = torch.ops._quant_matmul_cpu_18a2909
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_quant_matmul_cpu_18a2909::{op_name}"

build/torch211-cxx11-cpu-x86_64-linux/_quant_matmul_cpu_18a2909.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:cb00d2b295b018acc9a42048bf276c3d9c5dfe99077ac267e6c45f774bbbf403
+size 143032

build/torch211-cxx11-cpu-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,30 @@

+{
+  "name": "quant-matmul",
+  "id": "_quant_matmul_cpu_18a2909",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cpu"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "FTTRodgwtEJbDPRdBoeZdVWC5sDL6EpzX8St+v6bUyI=",
+      "_ops.py": "7XG0r+IuC4ujAoIFN6lqK2qfTCKug4asMH3GvteKiuw=",
+      "_quant_matmul_cpu_18a2909.abi3.so": "ywDSspWwGKzJpCBIvydsPZxd/pkHesJn5sRfd0u79AM=",
+      "quant_matmul/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "18a2909d5970949d8ec3915c4062d6c3c05466e8",
+      "dirty": false
+    }
+  }
+}

build/torch211-cxx11-cpu-x86_64-linux/quant_matmul/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch212-cxx11-cpu-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,229 @@

+"""quant-matmul: W4A8 / W8A8 / FP16-storage GEMV-GEMM for CPUs, aarch64-first.
+Int4 weights in GPTQ/AWQ-compatible group quantization run against per-token
+INT8 activations on NEON (SDOT where the CPU has dotprod); an int8 W8A8 path
+and an fp16-storage GEMV round out the pack. x86-64 runs a portable scalar
+path (correct, not tuned).
+High-level API:
+    pack_w4(intw, scales, zeros, group_size)   native packing from raw nibbles
+    from_gptq(qweight, qzeros, scales, g_idx)  GPTQ checkpoint tensors -> native
+    from_awq(qweight, qzeros, scales)          AWQ checkpoint tensors -> native
+    w4a8_linear(x, packed, scales, zeros, group_size, perm)
+    w8a8_linear(x, w_int8, scale_wt)
+    fp16_gemv(x, w_f16)
+    QuantLinear                                nn.Module wrapper
+"""
+from typing import Optional, Tuple
+import torch
+import torch.nn as nn
+from ._ops import ops
+_OK_DTYPES = (torch.bfloat16, torch.float32)
+def pack_w4(intw: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor,
+            group_size: int):
+    """Pack int4 nibbles into the native layout.
+    Args:
+        intw: [N, K] integer tensor with values 0..15.
+        scales: [N, G] per-group scales (G = K // group_size).
+        zeros: [N, G] per-group zero points, already in subtract-this form.
+    Returns:
+        (packed uint8 [N, K/2], scales f32 [N, G], zeros f32 [N, G])
+    """
+    assert intw.dim() == 2, "intw must be [N, K]"
+    N, K = intw.shape
+    assert K % group_size == 0 and group_size % 32 == 0
+    w = intw.to(torch.uint8)
+    assert int(w.max()) <= 15, "nibbles must be 0..15"
+    packed = (w[:, 0::2] | (w[:, 1::2] << 4)).contiguous()
+    G = K // group_size
+    s = scales.to(torch.float32).reshape(N, G).contiguous()
+    z = zeros.to(torch.float32).reshape(N, G).contiguous()
+    return packed, s, z
+def dequant_w4(packed: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor,
+               group_size: int) -> torch.Tensor:
+    """Reference dequantization to f32 [N, K] (tests and debugging)."""
+    N = packed.shape[0]
+    K = packed.shape[1] * 2
+    nib = torch.empty(N, K, dtype=torch.float32)
+    nib[:, 0::2] = (packed & 0xF).float()
+    nib[:, 1::2] = (packed >> 4).float()
+    G = K // group_size
+    z = zeros.reshape(N, G, 1)
+    s = scales.reshape(N, G, 1)
+    return ((nib.reshape(N, G, group_size) - z) * s).reshape(N, K)
+_GPTQ_SHIFTS = torch.arange(0, 32, 4, dtype=torch.int32)
+def _unpack_int32_nibbles(t: torch.Tensor, dim_factor: int = 8) -> torch.Tensor:
+    """int32 [..., P] -> uint8 [..., P*8], nibble j of each word at index 8p+j."""
+    u = ((t.unsqueeze(-1) >> _GPTQ_SHIFTS) & 0xF).to(torch.uint8)
+    return u.reshape(*t.shape[:-1], t.shape[-1] * dim_factor)
+def from_gptq(qweight: torch.Tensor, qzeros: torch.Tensor, scales: torch.Tensor,
+              g_idx: Optional[torch.Tensor] = None, zero_offset: int = 1):
+    """GPTQ checkpoint tensors -> native packing.
+    qweight int32 [K/8, N]; qzeros int32 [G, N/8]; scales [G, N].
+    zero_offset follows the AutoGPTQ storage convention (stored zero + 1 is
+    subtracted at dequant); pass 0 for checkpoints without it.
+    Returns (packed, scales, zeros, group_size, perm) where perm is the
+    activation permutation for act-order checkpoints (None otherwise).
+    """
+    K = qweight.shape[0] * 8
+    N = qweight.shape[1]
+    G = scales.shape[0]
+    assert K % G == 0
+    group_size = K // G
+    # nibble j of qweight word [i, n] decodes k = 8i + j.
+    intw = _unpack_int32_nibbles(qweight.t().contiguous())           # [N, K]
+    z = (_unpack_int32_nibbles(qzeros).to(torch.float32) + float(zero_offset))
+    z = z.t().contiguous()                                           # [N, G]
+    s = scales.t().to(torch.float32).contiguous()                    # [N, G]
+    perm = None
+    if g_idx is not None:
+        trivial = torch.arange(K, dtype=g_idx.dtype) // group_size
+        if not torch.equal(g_idx.cpu(), trivial):
+            perm = torch.argsort(g_idx.cpu(), stable=True)
+            intw = intw[:, perm]
+    packed, s, z = pack_w4(intw, s, z, group_size)
+    return packed, s, z, group_size, perm
+_AWQ_ORDER = [0, 2, 4, 6, 1, 3, 5, 7]
+def from_awq(qweight: torch.Tensor, qzeros: torch.Tensor, scales: torch.Tensor,
+             zero_offset: int = 0):
+    """AWQ (AutoAWQ GEMM) checkpoint tensors -> native packing.
+    qweight int32 [K, N/8] with nibble order [0,2,4,6,1,3,5,7]; qzeros int32
+    [G, N/8] same order; scales [G, N].
+    Returns (packed, scales, zeros, group_size, None).
+    """
+    K = qweight.shape[0]
+    G = scales.shape[0]
+    N = scales.shape[1]
+    group_size = K // G
+    inv = torch.empty(8, dtype=torch.int64)
+    for j, o in enumerate(_AWQ_ORDER):
+        inv[o] = j
+    u = _unpack_int32_nibbles(qweight)                               # [K, N] j-order
+    u = u.reshape(K, N // 8, 8)[:, :, inv].reshape(K, N)
+    z = _unpack_int32_nibbles(qzeros).to(torch.float32)              # [G, N] j-order
+    z = z.reshape(G, N // 8, 8)[:, :, inv].reshape(G, N) + float(zero_offset)
+    packed, s, z = pack_w4(u.t().contiguous(), scales.t(), z.t(), group_size)
+    return packed, s, z, group_size, None
+def quantize_activation(x: torch.Tensor):
+    """Per-token absmax INT8: [..., K] bf16/f32 -> (int8 [M, K], f32 scale [M])."""
+    x_flat = x.reshape(-1, x.shape[-1]).contiguous()
+    M, K = x_flat.shape
+    x_int8 = torch.empty((M, K), dtype=torch.int8)
+    scale = torch.empty((M,), dtype=torch.float32)
+    ops.qm_quantize_act(x_int8, scale, x_flat)
+    return x_int8, scale
+def w4a8_linear(x: torch.Tensor, packed: torch.Tensor, scales: torch.Tensor,
+                zeros: torch.Tensor, group_size: int,
+                perm: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """W4A8 forward: [..., K] bf16/f32 -> [..., N], dtype follows x."""
+    assert x.dtype in _OK_DTYPES, f"x must be bf16 or f32, got {x.dtype}"
+    orig = x.shape
+    x_flat = x.reshape(-1, orig[-1])
+    if perm is not None:
+        x_flat = x_flat[:, perm]
+    x_flat = x_flat.contiguous()
+    M = x_flat.shape[0]
+    N = packed.shape[0]
+    out = torch.empty((M, N), dtype=x.dtype)
+    if M < 16:
+        ops.qm_w4a8_gemv_fused(out, x_flat, packed, scales, zeros, group_size)
+    else:
+        q, sa = quantize_activation(x_flat)
+        ops.qm_w4a8_gemm(out, q, packed, scales, zeros, sa, group_size)
+    return out.view(*orig[:-1], N)
+def w8a8_linear(x: torch.Tensor, w_int8: torch.Tensor,
+                scale_wt: torch.Tensor) -> torch.Tensor:
+    """W8A8 forward: per-token int8 activations x int8 weights, per-row scale."""
+    assert x.dtype in _OK_DTYPES
+    orig = x.shape
+    q, sa = quantize_activation(x.reshape(-1, orig[-1]))
+    N = w_int8.shape[0]
+    out = torch.empty((q.shape[0], N), dtype=x.dtype)
+    ops.qm_w8a8_gemm(out, q, w_int8, sa, scale_wt)
+    return out.view(*orig[:-1], N)
+def fp16_gemv(x: torch.Tensor, w_f16: torch.Tensor) -> torch.Tensor:
+    """fp16-storage GEMV: f16 weights widened to f32 in registers."""
+    orig = x.shape
+    x_flat = x.reshape(-1, orig[-1]).to(torch.float32).contiguous()
+    N = w_f16.shape[0]
+    out = torch.empty((x_flat.shape[0], N), dtype=torch.float32)
+    ops.qm_fp16_gemv(out, x_flat, w_f16)
+    return out.view(*orig[:-1], N).to(x.dtype)
+class QuantLinear(nn.Module):
+    """Inference-only W4A8 linear over the native packing."""
+    def __init__(self, packed, scales, zeros, group_size, perm=None, bias=None):
+        super().__init__()
+        self.register_buffer("packed", packed)
+        self.register_buffer("scales", scales)
+        self.register_buffer("zeros", zeros)
+        self.group_size = int(group_size)
+        if perm is not None:
+            self.register_buffer("perm", perm)
+        else:
+            self.perm = None
+        if bias is not None:
+            self.register_buffer("bias", bias)
+        else:
+            self.bias = None
+    @classmethod
+    def from_gptq_tensors(cls, qweight, qzeros, scales, g_idx=None, bias=None,
+                          zero_offset=1):
+        packed, s, z, gs, perm = from_gptq(qweight, qzeros, scales, g_idx, zero_offset)
+        return cls(packed, s, z, gs, perm, bias)
+    @classmethod
+    def from_awq_tensors(cls, qweight, qzeros, scales, bias=None, zero_offset=0):
+        packed, s, z, gs, perm = from_awq(qweight, qzeros, scales, zero_offset)
+        return cls(packed, s, z, gs, perm, bias)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = w4a8_linear(x, self.packed, self.scales, self.zeros, self.group_size,
+                        self.perm)
+        if self.bias is not None:
+            y = y + self.bias
+        return y
+__all__ = [
+    "pack_w4",
+    "dequant_w4",
+    "from_gptq",
+    "from_awq",
+    "quantize_activation",
+    "w4a8_linear",
+    "w8a8_linear",
+    "fp16_gemv",
+    "QuantLinear",
+]

build/torch212-cxx11-cpu-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _quant_matmul_cpu_18a2909
+ops = torch.ops._quant_matmul_cpu_18a2909
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_quant_matmul_cpu_18a2909::{op_name}"

build/torch212-cxx11-cpu-x86_64-linux/_quant_matmul_cpu_18a2909.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:4b7aab09304bfdb2a5e7ab767a859f4d2f21d545b2bb3d00944e5c26a291d65b
+size 147104

build/torch212-cxx11-cpu-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,30 @@

+{
+  "name": "quant-matmul",
+  "id": "_quant_matmul_cpu_18a2909",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cpu"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "FTTRodgwtEJbDPRdBoeZdVWC5sDL6EpzX8St+v6bUyI=",
+      "_ops.py": "7XG0r+IuC4ujAoIFN6lqK2qfTCKug4asMH3GvteKiuw=",
+      "_quant_matmul_cpu_18a2909.abi3.so": "S3qrCTBL/bKl56t2eoWfTS8h1UWyuz0AlE5cJqKR1ls=",
+      "quant_matmul/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "18a2909d5970949d8ec3915c4062d6c3c05466e8",
+      "dirty": false
+    }
+  }
+}

build/torch212-cxx11-cpu-x86_64-linux/quant_matmul/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch213-cxx11-cpu-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,229 @@

+"""quant-matmul: W4A8 / W8A8 / FP16-storage GEMV-GEMM for CPUs, aarch64-first.
+Int4 weights in GPTQ/AWQ-compatible group quantization run against per-token
+INT8 activations on NEON (SDOT where the CPU has dotprod); an int8 W8A8 path
+and an fp16-storage GEMV round out the pack. x86-64 runs a portable scalar
+path (correct, not tuned).
+High-level API:
+    pack_w4(intw, scales, zeros, group_size)   native packing from raw nibbles
+    from_gptq(qweight, qzeros, scales, g_idx)  GPTQ checkpoint tensors -> native
+    from_awq(qweight, qzeros, scales)          AWQ checkpoint tensors -> native
+    w4a8_linear(x, packed, scales, zeros, group_size, perm)
+    w8a8_linear(x, w_int8, scale_wt)
+    fp16_gemv(x, w_f16)
+    QuantLinear                                nn.Module wrapper
+"""
+from typing import Optional, Tuple
+import torch
+import torch.nn as nn
+from ._ops import ops
+_OK_DTYPES = (torch.bfloat16, torch.float32)
+def pack_w4(intw: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor,
+            group_size: int):
+    """Pack int4 nibbles into the native layout.
+    Args:
+        intw: [N, K] integer tensor with values 0..15.
+        scales: [N, G] per-group scales (G = K // group_size).
+        zeros: [N, G] per-group zero points, already in subtract-this form.
+    Returns:
+        (packed uint8 [N, K/2], scales f32 [N, G], zeros f32 [N, G])
+    """
+    assert intw.dim() == 2, "intw must be [N, K]"
+    N, K = intw.shape
+    assert K % group_size == 0 and group_size % 32 == 0
+    w = intw.to(torch.uint8)
+    assert int(w.max()) <= 15, "nibbles must be 0..15"
+    packed = (w[:, 0::2] | (w[:, 1::2] << 4)).contiguous()
+    G = K // group_size
+    s = scales.to(torch.float32).reshape(N, G).contiguous()
+    z = zeros.to(torch.float32).reshape(N, G).contiguous()
+    return packed, s, z
+def dequant_w4(packed: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor,
+               group_size: int) -> torch.Tensor:
+    """Reference dequantization to f32 [N, K] (tests and debugging)."""
+    N = packed.shape[0]
+    K = packed.shape[1] * 2
+    nib = torch.empty(N, K, dtype=torch.float32)
+    nib[:, 0::2] = (packed & 0xF).float()
+    nib[:, 1::2] = (packed >> 4).float()
+    G = K // group_size
+    z = zeros.reshape(N, G, 1)
+    s = scales.reshape(N, G, 1)
+    return ((nib.reshape(N, G, group_size) - z) * s).reshape(N, K)
+_GPTQ_SHIFTS = torch.arange(0, 32, 4, dtype=torch.int32)
+def _unpack_int32_nibbles(t: torch.Tensor, dim_factor: int = 8) -> torch.Tensor:
+    """int32 [..., P] -> uint8 [..., P*8], nibble j of each word at index 8p+j."""
+    u = ((t.unsqueeze(-1) >> _GPTQ_SHIFTS) & 0xF).to(torch.uint8)
+    return u.reshape(*t.shape[:-1], t.shape[-1] * dim_factor)
+def from_gptq(qweight: torch.Tensor, qzeros: torch.Tensor, scales: torch.Tensor,
+              g_idx: Optional[torch.Tensor] = None, zero_offset: int = 1):
+    """GPTQ checkpoint tensors -> native packing.
+    qweight int32 [K/8, N]; qzeros int32 [G, N/8]; scales [G, N].
+    zero_offset follows the AutoGPTQ storage convention (stored zero + 1 is
+    subtracted at dequant); pass 0 for checkpoints without it.
+    Returns (packed, scales, zeros, group_size, perm) where perm is the
+    activation permutation for act-order checkpoints (None otherwise).
+    """
+    K = qweight.shape[0] * 8
+    N = qweight.shape[1]
+    G = scales.shape[0]
+    assert K % G == 0
+    group_size = K // G
+    # nibble j of qweight word [i, n] decodes k = 8i + j.
+    intw = _unpack_int32_nibbles(qweight.t().contiguous())           # [N, K]
+    z = (_unpack_int32_nibbles(qzeros).to(torch.float32) + float(zero_offset))
+    z = z.t().contiguous()                                           # [N, G]
+    s = scales.t().to(torch.float32).contiguous()                    # [N, G]
+    perm = None
+    if g_idx is not None:
+        trivial = torch.arange(K, dtype=g_idx.dtype) // group_size
+        if not torch.equal(g_idx.cpu(), trivial):
+            perm = torch.argsort(g_idx.cpu(), stable=True)
+            intw = intw[:, perm]
+    packed, s, z = pack_w4(intw, s, z, group_size)
+    return packed, s, z, group_size, perm
+_AWQ_ORDER = [0, 2, 4, 6, 1, 3, 5, 7]
+def from_awq(qweight: torch.Tensor, qzeros: torch.Tensor, scales: torch.Tensor,
+             zero_offset: int = 0):
+    """AWQ (AutoAWQ GEMM) checkpoint tensors -> native packing.
+    qweight int32 [K, N/8] with nibble order [0,2,4,6,1,3,5,7]; qzeros int32
+    [G, N/8] same order; scales [G, N].
+    Returns (packed, scales, zeros, group_size, None).
+    """
+    K = qweight.shape[0]
+    G = scales.shape[0]
+    N = scales.shape[1]
+    group_size = K // G
+    inv = torch.empty(8, dtype=torch.int64)
+    for j, o in enumerate(_AWQ_ORDER):
+        inv[o] = j
+    u = _unpack_int32_nibbles(qweight)                               # [K, N] j-order
+    u = u.reshape(K, N // 8, 8)[:, :, inv].reshape(K, N)
+    z = _unpack_int32_nibbles(qzeros).to(torch.float32)              # [G, N] j-order
+    z = z.reshape(G, N // 8, 8)[:, :, inv].reshape(G, N) + float(zero_offset)
+    packed, s, z = pack_w4(u.t().contiguous(), scales.t(), z.t(), group_size)
+    return packed, s, z, group_size, None
+def quantize_activation(x: torch.Tensor):
+    """Per-token absmax INT8: [..., K] bf16/f32 -> (int8 [M, K], f32 scale [M])."""
+    x_flat = x.reshape(-1, x.shape[-1]).contiguous()
+    M, K = x_flat.shape
+    x_int8 = torch.empty((M, K), dtype=torch.int8)
+    scale = torch.empty((M,), dtype=torch.float32)
+    ops.qm_quantize_act(x_int8, scale, x_flat)
+    return x_int8, scale
+def w4a8_linear(x: torch.Tensor, packed: torch.Tensor, scales: torch.Tensor,
+                zeros: torch.Tensor, group_size: int,
+                perm: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """W4A8 forward: [..., K] bf16/f32 -> [..., N], dtype follows x."""
+    assert x.dtype in _OK_DTYPES, f"x must be bf16 or f32, got {x.dtype}"
+    orig = x.shape
+    x_flat = x.reshape(-1, orig[-1])
+    if perm is not None:
+        x_flat = x_flat[:, perm]
+    x_flat = x_flat.contiguous()
+    M = x_flat.shape[0]
+    N = packed.shape[0]
+    out = torch.empty((M, N), dtype=x.dtype)
+    if M < 16:
+        ops.qm_w4a8_gemv_fused(out, x_flat, packed, scales, zeros, group_size)
+    else:
+        q, sa = quantize_activation(x_flat)
+        ops.qm_w4a8_gemm(out, q, packed, scales, zeros, sa, group_size)
+    return out.view(*orig[:-1], N)
+def w8a8_linear(x: torch.Tensor, w_int8: torch.Tensor,
+                scale_wt: torch.Tensor) -> torch.Tensor:
+    """W8A8 forward: per-token int8 activations x int8 weights, per-row scale."""
+    assert x.dtype in _OK_DTYPES
+    orig = x.shape
+    q, sa = quantize_activation(x.reshape(-1, orig[-1]))
+    N = w_int8.shape[0]
+    out = torch.empty((q.shape[0], N), dtype=x.dtype)
+    ops.qm_w8a8_gemm(out, q, w_int8, sa, scale_wt)
+    return out.view(*orig[:-1], N)
+def fp16_gemv(x: torch.Tensor, w_f16: torch.Tensor) -> torch.Tensor:
+    """fp16-storage GEMV: f16 weights widened to f32 in registers."""
+    orig = x.shape
+    x_flat = x.reshape(-1, orig[-1]).to(torch.float32).contiguous()
+    N = w_f16.shape[0]
+    out = torch.empty((x_flat.shape[0], N), dtype=torch.float32)
+    ops.qm_fp16_gemv(out, x_flat, w_f16)
+    return out.view(*orig[:-1], N).to(x.dtype)
+class QuantLinear(nn.Module):
+    """Inference-only W4A8 linear over the native packing."""
+    def __init__(self, packed, scales, zeros, group_size, perm=None, bias=None):
+        super().__init__()
+        self.register_buffer("packed", packed)
+        self.register_buffer("scales", scales)
+        self.register_buffer("zeros", zeros)
+        self.group_size = int(group_size)
+        if perm is not None:
+            self.register_buffer("perm", perm)
+        else:
+            self.perm = None
+        if bias is not None:
+            self.register_buffer("bias", bias)
+        else:
+            self.bias = None
+    @classmethod
+    def from_gptq_tensors(cls, qweight, qzeros, scales, g_idx=None, bias=None,
+                          zero_offset=1):
+        packed, s, z, gs, perm = from_gptq(qweight, qzeros, scales, g_idx, zero_offset)
+        return cls(packed, s, z, gs, perm, bias)
+    @classmethod
+    def from_awq_tensors(cls, qweight, qzeros, scales, bias=None, zero_offset=0):
+        packed, s, z, gs, perm = from_awq(qweight, qzeros, scales, zero_offset)
+        return cls(packed, s, z, gs, perm, bias)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = w4a8_linear(x, self.packed, self.scales, self.zeros, self.group_size,
+                        self.perm)
+        if self.bias is not None:
+            y = y + self.bias
+        return y
+__all__ = [
+    "pack_w4",
+    "dequant_w4",
+    "from_gptq",
+    "from_awq",
+    "quantize_activation",
+    "w4a8_linear",
+    "w8a8_linear",
+    "fp16_gemv",
+    "QuantLinear",
+]

build/torch213-cxx11-cpu-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _quant_matmul_cpu_18a2909
+ops = torch.ops._quant_matmul_cpu_18a2909
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_quant_matmul_cpu_18a2909::{op_name}"

build/torch213-cxx11-cpu-x86_64-linux/_quant_matmul_cpu_18a2909.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:3c1e74f4c3e314bc8cfaf7d735b16b7176a13f20735d91b9b4843b264dc54cba
+size 146952

build/torch213-cxx11-cpu-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,30 @@

+{
+  "name": "quant-matmul",
+  "id": "_quant_matmul_cpu_18a2909",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cpu"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "FTTRodgwtEJbDPRdBoeZdVWC5sDL6EpzX8St+v6bUyI=",
+      "_ops.py": "7XG0r+IuC4ujAoIFN6lqK2qfTCKug4asMH3GvteKiuw=",
+      "_quant_matmul_cpu_18a2909.abi3.so": "PB509MPjFLyM+vfXNbFrcXahPyBzXZG5tIQ7Jk3FTLo=",
+      "quant_matmul/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "18a2909d5970949d8ec3915c4062d6c3c05466e8",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cpu-x86_64-linux/quant_matmul/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))