drbh HF Staff commited on 16 days ago

Commit

0247a3b

verified ·

1 Parent(s): 69f84f8

Uploaded using `kernel-builder`.

Browse files

Files changed (48) hide show

build/torch211-cxx11-cu126-x86_64-linux/__init__.py +43 -0
build/torch211-cxx11-cu126-x86_64-linux/_gemmaq_cuda_6820a9e.abi3.so +3 -0
build/torch211-cxx11-cu126-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu126-x86_64-linux/gemmaq/__init__.py +26 -0
build/torch211-cxx11-cu126-x86_64-linux/layer.py +81 -0
build/torch211-cxx11-cu126-x86_64-linux/layers.py +48 -0
build/torch211-cxx11-cu126-x86_64-linux/metadata.json +25 -0
build/torch211-cxx11-cu126-x86_64-linux/quant.py +67 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +43 -0
build/torch211-cxx11-cu128-x86_64-linux/_gemmaq_cuda_6820a9e.abi3.so +3 -0
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu128-x86_64-linux/gemmaq/__init__.py +26 -0
build/torch211-cxx11-cu128-x86_64-linux/layer.py +81 -0
build/torch211-cxx11-cu128-x86_64-linux/layers.py +48 -0
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +25 -0
build/torch211-cxx11-cu128-x86_64-linux/quant.py +67 -0
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +43 -0
build/torch211-cxx11-cu130-x86_64-linux/_gemmaq_cuda_6820a9e.abi3.so +3 -0
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu130-x86_64-linux/gemmaq/__init__.py +26 -0
build/torch211-cxx11-cu130-x86_64-linux/layer.py +81 -0
build/torch211-cxx11-cu130-x86_64-linux/layers.py +48 -0
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +25 -0
build/torch211-cxx11-cu130-x86_64-linux/quant.py +67 -0
build/torch212-cxx11-cu126-x86_64-linux/__init__.py +43 -0
build/torch212-cxx11-cu126-x86_64-linux/_gemmaq_cuda_6820a9e.abi3.so +3 -0
build/torch212-cxx11-cu126-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu126-x86_64-linux/gemmaq/__init__.py +26 -0
build/torch212-cxx11-cu126-x86_64-linux/layer.py +81 -0
build/torch212-cxx11-cu126-x86_64-linux/layers.py +48 -0
build/torch212-cxx11-cu126-x86_64-linux/metadata.json +25 -0
build/torch212-cxx11-cu126-x86_64-linux/quant.py +67 -0
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +43 -0
build/torch212-cxx11-cu130-x86_64-linux/_gemmaq_cuda_6820a9e.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu130-x86_64-linux/gemmaq/__init__.py +26 -0
build/torch212-cxx11-cu130-x86_64-linux/layer.py +81 -0
build/torch212-cxx11-cu130-x86_64-linux/layers.py +48 -0
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +25 -0
build/torch212-cxx11-cu130-x86_64-linux/quant.py +67 -0
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +43 -0
build/torch212-cxx11-cu132-x86_64-linux/_gemmaq_cuda_6820a9e.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu132-x86_64-linux/gemmaq/__init__.py +26 -0
build/torch212-cxx11-cu132-x86_64-linux/layer.py +81 -0
build/torch212-cxx11-cu132-x86_64-linux/layers.py +48 -0
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +25 -0
build/torch212-cxx11-cu132-x86_64-linux/quant.py +67 -0

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

	@@ -0,0 +1,43 @@

+"""gemmaq: W4A16 quantized GEMV kernel + quantization utilities."""
+from . import quant
+from .quant import (
+    dequantize_w4a16_asym,
+    gemv_reference,
+    quantize_w4a16_asym,
+)
+# The compiled op (._ops) only exists in a built package; importing the source
+# tree (for quant utilities) must not fail.
+try:
+    from ._ops import ops
+except ImportError:
+    ops = None
+# Model-side layer + swap utility, and the hub-consumable kernel layers
+# (gemmaq.layers.W4A16Linear, resolved by kernels.kernelize via LayerRepository).
+try:
+    from . import layers
+    from .layer import W4A16Linear, quantize_model
+except Exception:  # torch.library API or ops unavailable
+    layers = None
+    W4A16Linear = quantize_model = None
+def w4a16_gemv(x, qweight, scales, zeros, group_size=128):
+    """y = x @ dequant(W).T using the compiled W4A16 GEMV kernel."""
+    if ops is None:
+        raise RuntimeError("gemmaq compiled ops not available (build with `make dev`)")
+    return ops.w4a16_gemv(x, qweight, scales, zeros, group_size)
+__all__ = [
+    "ops",
+    "w4a16_gemv",
+    "quant",
+    "quantize_w4a16_asym",
+    "dequantize_w4a16_asym",
+    "gemv_reference",
+    "W4A16Linear",
+    "quantize_model",
+    "layers",
+]

build/torch211-cxx11-cu126-x86_64-linux/_gemmaq_cuda_6820a9e.abi3.so ADDED Viewed

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

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

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

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

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

build/torch211-cxx11-cu126-x86_64-linux/gemmaq/__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/torch211-cxx11-cu126-x86_64-linux/layer.py ADDED Viewed

	@@ -0,0 +1,81 @@

+"""Model-side W4A16 layer + model-swap utility.
+`quantize_model` replaces the target `nn.Linear`s with `W4A16Linear` (which packs
+the weights to int4). `W4A16Linear` is decorated with
+`@use_kernel_forward_from_hub("W4A16Linear")`, so its pure-torch forward is the
+*fallback*; calling `kernels.kernelize(model, ...)` with a mapping to the
+`gemmaq.layers.W4A16Linear` kernel layer swaps in the compiled GEMV. The model is
+therefore runnable with or without the compiled kernel.
+"""
+import torch
+import torch.nn as nn
+from .quant import dequantize_w4a16_asym, quantize_w4a16_asym
+try:
+    from kernels import use_kernel_forward_from_hub
+except ImportError:  # kernels not installed: no-op decorator (forward still works)
+    def use_kernel_forward_from_hub(_name):
+        return lambda cls: cls
+# 7 linear types per decoder layer (lm_head deferred to Phase 2).
+DEFAULT_TARGETS = ("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj")
+_SKIP = ("vision", "audio", "multi_modal", "mm_")
+@use_kernel_forward_from_hub("W4A16Linear")
+class W4A16Linear(nn.Module):
+    """Weight-only int4 linear. Holds the quantized buffers; the forward here is
+    the pure-torch fallback (dequant + mm). kernelize() swaps in the kernel."""
+    def __init__(self, in_features, out_features, bias=False, group_size=128, device=None):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.group_size = group_size
+        self.register_buffer("qweight", torch.empty(out_features, in_features // 8, dtype=torch.int32, device=device))
+        self.register_buffer("scales", torch.empty(out_features, in_features // group_size, dtype=torch.float16, device=device))
+        self.register_buffer("zeros", torch.empty(out_features, in_features // group_size, dtype=torch.uint8, device=device))
+        self.register_buffer("bias", torch.empty(out_features, device=device) if bias else None)
+    @classmethod
+    def from_linear(cls, lin: nn.Linear, group_size=128):
+        qw, sc, zp = quantize_w4a16_asym(lin.weight.data, group_size)
+        m = cls(lin.in_features, lin.out_features, bias=lin.bias is not None,
+                group_size=group_size, device=lin.weight.device)
+        m.qweight.copy_(qw)
+        m.scales.copy_(sc)
+        m.zeros.copy_(zp)
+        if lin.bias is not None:
+            m.bias.copy_(lin.bias.data.to(m.bias.dtype))
+        return m
+    def forward(self, x):
+        w = dequantize_w4a16_asym(self.qweight, self.scales, self.zeros, self.group_size).to(x.dtype)
+        out = torch.matmul(x, w.t())
+        if self.bias is not None:
+            out = out + self.bias.to(out.dtype)
+        return out
+def quantize_model(model, group_size=128, targets=DEFAULT_TARGETS):
+    """In-place swap of the target nn.Linear layers with W4A16Linear. Returns count.
+    After this, call `kernels.kernelize(model, mode=...)` with a mapping for
+    `"W4A16Linear"` to activate the compiled kernel (otherwise the torch fallback runs).
+    """
+    to_swap = []
+    for name, mod in model.named_modules():
+        if (isinstance(mod, nn.Linear)
+                and name.rsplit(".", 1)[-1] in targets
+                and not any(s in name for s in _SKIP)
+                and mod.in_features % group_size == 0):
+            to_swap.append((name, mod))
+    for name, mod in to_swap:
+        parent = model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else model
+        child = name.rsplit(".", 1)[-1]
+        setattr(parent, child, W4A16Linear.from_linear(mod, group_size))
+        del mod
+    torch.cuda.empty_cache()
+    return len(to_swap)

build/torch211-cxx11-cu126-x86_64-linux/layers.py ADDED Viewed

	@@ -0,0 +1,48 @@

+"""Hub-consumable kernel layers, resolved by `kernels.kernelize` via a
+LayerRepository (`layer_name="W4A16Linear"`).
+Per the kernels layer contract these layers are STATELESS: kernelize binds only
+`forward` onto an existing module instance, so `self` carries the quantized
+buffers (qweight/scales/zeros/group_size/bias) from gemmaq.layer.W4A16Linear.
+The class must not define __init__ or extra members (only `forward` plus the
+`can_torch_compile`/`has_backward` flags).
+"""
+import torch
+import torch.nn as nn
+from .quant import dequantize_w4a16_asym
+try:
+    from ._ops import ops as _ops
+except ImportError:  # importable from source tree without the build
+    _ops = None
+def _w4a16_fake(x, qweight, scales, zeros, group_size):
+    return x.new_empty((*x.shape[:-1], qweight.shape[0]))
+# Register a meta/fake impl so the compiled op shape-propagates under torch.compile.
+if _ops is not None:
+    try:
+        torch.library.register_fake(_ops.w4a16_gemv.default.name(), _w4a16_fake)
+    except Exception:
+        pass  # already registered
+class W4A16Linear(nn.Module):
+    """Optimized W4A16 forward: compiled GEMV for decode (M=1), dequant+mm for
+    prefill (M>1, where the GEMV kernel loses to cuBLAS GEMM)."""
+    can_torch_compile = True
+    def forward(self, x):
+        m = x.numel() // x.shape[-1]
+        if m == 1:
+            out = _ops.w4a16_gemv(x, self.qweight, self.scales, self.zeros, self.group_size)
+        else:
+            w = dequantize_w4a16_asym(self.qweight, self.scales, self.zeros, self.group_size).to(x.dtype)
+            out = torch.matmul(x, w.t())
+        if self.bias is not None:
+            out = out + self.bias.to(out.dtype)
+        return out

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

	@@ -0,0 +1,25 @@

+{
+  "name": "gemmaq",
+  "id": "_gemmaq_cuda_6820a9e",
+  "version": 1,
+  "license": "MIT",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.9"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "nufuk/+eZjEX7g9AJLBII5Nr7alASHAUDgNRJTdMs1c=",
+      "_gemmaq_cuda_6820a9e.abi3.so": "+Tqi9Z395m1Q7ugS9evkzW9rP4+WMEo4Pq9Ida+W/pA=",
+      "_ops.py": "VrnRWJF6J13HL5IPqZyNmcgDlxYow62J+G2ZD6PU8Ac=",
+      "gemmaq/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
+      "layer.py": "Eml+uGNCtnO3dtrmIte8vtO+BteLT6wYfGUKvbbuz6E=",
+      "layers.py": "SM4DMpETjS/NXe7dMdjLR+r9NTcuO0gaqW26oaDKC5E=",
+      "quant.py": "OeosAKwPsCax7g8GfotC7FBbI/p/pVIg3tNoAFilQzw="
+    }
+  }
+}

build/torch211-cxx11-cu126-x86_64-linux/quant.py ADDED Viewed

	@@ -0,0 +1,67 @@

+"""W4A16 weight-only quantization: asymmetric, group-wise, RTN.
+Layout (must match csrc/w4a16_gemv.cu):
+  weight W: [N, K]  (PyTorch nn.Linear weight; y = x @ W.T)
+  quantize along K in groups of `group_size` (default 128).
+    scale[n,g] = (wmax-wmin)/15,  zero[n,g] = round(-wmin/scale)  in [0,15]
+    q[n,k]     = clamp(round(W/scale)+zero, 0, 15)                 uint4
+    w_hat      = (q - zero) * scale
+  qweight: int32[N, K//8], 8 nibbles packed along K per word
+           (nibble i -> k = 8*j + i, at bits [4i, 4i+4)).
+  scales:  fp16 [N, K//group_size]
+  zeros:   uint8[N, K//group_size]
+"""
+import torch
+PACK = 8  # int4 values per int32 word
+def pack_int4(q: torch.Tensor) -> torch.Tensor:
+    """q: [N, K] int (values 0..15) -> packed int32 [N, K//8]."""
+    N, K = q.shape
+    assert K % PACK == 0, f"K={K} must be divisible by {PACK}"
+    q = q.to(torch.int32).reshape(N, K // PACK, PACK)
+    packed = torch.zeros(N, K // PACK, dtype=torch.int32, device=q.device)
+    for i in range(PACK):
+        packed |= (q[:, :, i] & 0xF) << (4 * i)
+    return packed
+def unpack_int4(qweight: torch.Tensor) -> torch.Tensor:
+    """packed int32 [N, K//8] -> [N, K] int32 values 0..15."""
+    N, Kp = qweight.shape
+    out = torch.empty(N, Kp, PACK, dtype=torch.int32, device=qweight.device)
+    for i in range(PACK):
+        out[:, :, i] = (qweight >> (4 * i)) & 0xF  # & 0xF undoes arithmetic-shift sign bits
+    return out.reshape(N, Kp * PACK)
+def quantize_w4a16_asym(weight: torch.Tensor, group_size: int = 128):
+    """RTN asymmetric int4. Returns (qweight int32, scales fp16, zeros uint8)."""
+    N, K = weight.shape
+    G = group_size
+    assert K % G == 0, f"K={K} must be divisible by group_size={G}"
+    w = weight.detach().float().reshape(N, K // G, G)
+    wmin = w.amin(dim=2)
+    wmax = w.amax(dim=2)
+    scale = ((wmax - wmin) / 15.0).clamp(min=1e-8)              # [N, K//G]
+    zero = torch.round(-wmin / scale).clamp(0, 15)             # [N, K//G]
+    q = torch.round(w / scale.unsqueeze(2) + zero.unsqueeze(2)).clamp(0, 15)
+    qweight = pack_int4(q.reshape(N, K).to(torch.int32))
+    return qweight, scale.to(torch.float16), zero.to(torch.uint8)
+def dequantize_w4a16_asym(qweight, scales, zeros, group_size: int = 128) -> torch.Tensor:
+    """Golden reference: reconstruct w_hat [N, K] float from the packed form."""
+    q = unpack_int4(qweight).float()                          # [N, K]
+    N, K = q.shape
+    G = group_size
+    q = q.reshape(N, K // G, G)
+    w = (q - zeros.float().unsqueeze(2)) * scales.float().unsqueeze(2)
+    return w.reshape(N, K)
+def gemv_reference(x, qweight, scales, zeros, group_size: int = 128) -> torch.Tensor:
+    """Reference W4A16 linear (no bias): y = x @ dequant(W).T."""
+    w_hat = dequantize_w4a16_asym(qweight, scales, zeros, group_size).to(x.dtype)
+    return x @ w_hat.t()

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

	@@ -0,0 +1,43 @@

+"""gemmaq: W4A16 quantized GEMV kernel + quantization utilities."""
+from . import quant
+from .quant import (
+    dequantize_w4a16_asym,
+    gemv_reference,
+    quantize_w4a16_asym,
+)
+# The compiled op (._ops) only exists in a built package; importing the source
+# tree (for quant utilities) must not fail.
+try:
+    from ._ops import ops
+except ImportError:
+    ops = None
+# Model-side layer + swap utility, and the hub-consumable kernel layers
+# (gemmaq.layers.W4A16Linear, resolved by kernels.kernelize via LayerRepository).
+try:
+    from . import layers
+    from .layer import W4A16Linear, quantize_model
+except Exception:  # torch.library API or ops unavailable
+    layers = None
+    W4A16Linear = quantize_model = None
+def w4a16_gemv(x, qweight, scales, zeros, group_size=128):
+    """y = x @ dequant(W).T using the compiled W4A16 GEMV kernel."""
+    if ops is None:
+        raise RuntimeError("gemmaq compiled ops not available (build with `make dev`)")
+    return ops.w4a16_gemv(x, qweight, scales, zeros, group_size)
+__all__ = [
+    "ops",
+    "w4a16_gemv",
+    "quant",
+    "quantize_w4a16_asym",
+    "dequantize_w4a16_asym",
+    "gemv_reference",
+    "W4A16Linear",
+    "quantize_model",
+    "layers",
+]

build/torch211-cxx11-cu128-x86_64-linux/_gemmaq_cuda_6820a9e.abi3.so ADDED Viewed

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

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

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

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

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

build/torch211-cxx11-cu128-x86_64-linux/gemmaq/__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/torch211-cxx11-cu128-x86_64-linux/layer.py ADDED Viewed

	@@ -0,0 +1,81 @@

+"""Model-side W4A16 layer + model-swap utility.
+`quantize_model` replaces the target `nn.Linear`s with `W4A16Linear` (which packs
+the weights to int4). `W4A16Linear` is decorated with
+`@use_kernel_forward_from_hub("W4A16Linear")`, so its pure-torch forward is the
+*fallback*; calling `kernels.kernelize(model, ...)` with a mapping to the
+`gemmaq.layers.W4A16Linear` kernel layer swaps in the compiled GEMV. The model is
+therefore runnable with or without the compiled kernel.
+"""
+import torch
+import torch.nn as nn
+from .quant import dequantize_w4a16_asym, quantize_w4a16_asym
+try:
+    from kernels import use_kernel_forward_from_hub
+except ImportError:  # kernels not installed: no-op decorator (forward still works)
+    def use_kernel_forward_from_hub(_name):
+        return lambda cls: cls
+# 7 linear types per decoder layer (lm_head deferred to Phase 2).
+DEFAULT_TARGETS = ("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj")
+_SKIP = ("vision", "audio", "multi_modal", "mm_")
+@use_kernel_forward_from_hub("W4A16Linear")
+class W4A16Linear(nn.Module):
+    """Weight-only int4 linear. Holds the quantized buffers; the forward here is
+    the pure-torch fallback (dequant + mm). kernelize() swaps in the kernel."""
+    def __init__(self, in_features, out_features, bias=False, group_size=128, device=None):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.group_size = group_size
+        self.register_buffer("qweight", torch.empty(out_features, in_features // 8, dtype=torch.int32, device=device))
+        self.register_buffer("scales", torch.empty(out_features, in_features // group_size, dtype=torch.float16, device=device))
+        self.register_buffer("zeros", torch.empty(out_features, in_features // group_size, dtype=torch.uint8, device=device))
+        self.register_buffer("bias", torch.empty(out_features, device=device) if bias else None)
+    @classmethod
+    def from_linear(cls, lin: nn.Linear, group_size=128):
+        qw, sc, zp = quantize_w4a16_asym(lin.weight.data, group_size)
+        m = cls(lin.in_features, lin.out_features, bias=lin.bias is not None,
+                group_size=group_size, device=lin.weight.device)
+        m.qweight.copy_(qw)
+        m.scales.copy_(sc)
+        m.zeros.copy_(zp)
+        if lin.bias is not None:
+            m.bias.copy_(lin.bias.data.to(m.bias.dtype))
+        return m
+    def forward(self, x):
+        w = dequantize_w4a16_asym(self.qweight, self.scales, self.zeros, self.group_size).to(x.dtype)
+        out = torch.matmul(x, w.t())
+        if self.bias is not None:
+            out = out + self.bias.to(out.dtype)
+        return out
+def quantize_model(model, group_size=128, targets=DEFAULT_TARGETS):
+    """In-place swap of the target nn.Linear layers with W4A16Linear. Returns count.
+    After this, call `kernels.kernelize(model, mode=...)` with a mapping for
+    `"W4A16Linear"` to activate the compiled kernel (otherwise the torch fallback runs).
+    """
+    to_swap = []
+    for name, mod in model.named_modules():
+        if (isinstance(mod, nn.Linear)
+                and name.rsplit(".", 1)[-1] in targets
+                and not any(s in name for s in _SKIP)
+                and mod.in_features % group_size == 0):
+            to_swap.append((name, mod))
+    for name, mod in to_swap:
+        parent = model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else model
+        child = name.rsplit(".", 1)[-1]
+        setattr(parent, child, W4A16Linear.from_linear(mod, group_size))
+        del mod
+    torch.cuda.empty_cache()
+    return len(to_swap)

build/torch211-cxx11-cu128-x86_64-linux/layers.py ADDED Viewed

	@@ -0,0 +1,48 @@

+"""Hub-consumable kernel layers, resolved by `kernels.kernelize` via a
+LayerRepository (`layer_name="W4A16Linear"`).
+Per the kernels layer contract these layers are STATELESS: kernelize binds only
+`forward` onto an existing module instance, so `self` carries the quantized
+buffers (qweight/scales/zeros/group_size/bias) from gemmaq.layer.W4A16Linear.
+The class must not define __init__ or extra members (only `forward` plus the
+`can_torch_compile`/`has_backward` flags).
+"""
+import torch
+import torch.nn as nn
+from .quant import dequantize_w4a16_asym
+try:
+    from ._ops import ops as _ops
+except ImportError:  # importable from source tree without the build
+    _ops = None
+def _w4a16_fake(x, qweight, scales, zeros, group_size):
+    return x.new_empty((*x.shape[:-1], qweight.shape[0]))
+# Register a meta/fake impl so the compiled op shape-propagates under torch.compile.
+if _ops is not None:
+    try:
+        torch.library.register_fake(_ops.w4a16_gemv.default.name(), _w4a16_fake)
+    except Exception:
+        pass  # already registered
+class W4A16Linear(nn.Module):
+    """Optimized W4A16 forward: compiled GEMV for decode (M=1), dequant+mm for
+    prefill (M>1, where the GEMV kernel loses to cuBLAS GEMM)."""
+    can_torch_compile = True
+    def forward(self, x):
+        m = x.numel() // x.shape[-1]
+        if m == 1:
+            out = _ops.w4a16_gemv(x, self.qweight, self.scales, self.zeros, self.group_size)
+        else:
+            w = dequantize_w4a16_asym(self.qweight, self.scales, self.zeros, self.group_size).to(x.dtype)
+            out = torch.matmul(x, w.t())
+        if self.bias is not None:
+            out = out + self.bias.to(out.dtype)
+        return out

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

	@@ -0,0 +1,25 @@

+{
+  "name": "gemmaq",
+  "id": "_gemmaq_cuda_6820a9e",
+  "version": 1,
+  "license": "MIT",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.9"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "nufuk/+eZjEX7g9AJLBII5Nr7alASHAUDgNRJTdMs1c=",
+      "_gemmaq_cuda_6820a9e.abi3.so": "rQqnG2nWULB2znF1awO/ieKpa0dMsuG22M0/vWQhw50=",
+      "_ops.py": "VrnRWJF6J13HL5IPqZyNmcgDlxYow62J+G2ZD6PU8Ac=",
+      "gemmaq/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
+      "layer.py": "Eml+uGNCtnO3dtrmIte8vtO+BteLT6wYfGUKvbbuz6E=",
+      "layers.py": "SM4DMpETjS/NXe7dMdjLR+r9NTcuO0gaqW26oaDKC5E=",
+      "quant.py": "OeosAKwPsCax7g8GfotC7FBbI/p/pVIg3tNoAFilQzw="
+    }
+  }
+}

build/torch211-cxx11-cu128-x86_64-linux/quant.py ADDED Viewed

	@@ -0,0 +1,67 @@

+"""W4A16 weight-only quantization: asymmetric, group-wise, RTN.
+Layout (must match csrc/w4a16_gemv.cu):
+  weight W: [N, K]  (PyTorch nn.Linear weight; y = x @ W.T)
+  quantize along K in groups of `group_size` (default 128).
+    scale[n,g] = (wmax-wmin)/15,  zero[n,g] = round(-wmin/scale)  in [0,15]
+    q[n,k]     = clamp(round(W/scale)+zero, 0, 15)                 uint4
+    w_hat      = (q - zero) * scale
+  qweight: int32[N, K//8], 8 nibbles packed along K per word
+           (nibble i -> k = 8*j + i, at bits [4i, 4i+4)).
+  scales:  fp16 [N, K//group_size]
+  zeros:   uint8[N, K//group_size]
+"""
+import torch
+PACK = 8  # int4 values per int32 word
+def pack_int4(q: torch.Tensor) -> torch.Tensor:
+    """q: [N, K] int (values 0..15) -> packed int32 [N, K//8]."""
+    N, K = q.shape
+    assert K % PACK == 0, f"K={K} must be divisible by {PACK}"
+    q = q.to(torch.int32).reshape(N, K // PACK, PACK)
+    packed = torch.zeros(N, K // PACK, dtype=torch.int32, device=q.device)
+    for i in range(PACK):
+        packed |= (q[:, :, i] & 0xF) << (4 * i)
+    return packed
+def unpack_int4(qweight: torch.Tensor) -> torch.Tensor:
+    """packed int32 [N, K//8] -> [N, K] int32 values 0..15."""
+    N, Kp = qweight.shape
+    out = torch.empty(N, Kp, PACK, dtype=torch.int32, device=qweight.device)
+    for i in range(PACK):
+        out[:, :, i] = (qweight >> (4 * i)) & 0xF  # & 0xF undoes arithmetic-shift sign bits
+    return out.reshape(N, Kp * PACK)
+def quantize_w4a16_asym(weight: torch.Tensor, group_size: int = 128):
+    """RTN asymmetric int4. Returns (qweight int32, scales fp16, zeros uint8)."""
+    N, K = weight.shape
+    G = group_size
+    assert K % G == 0, f"K={K} must be divisible by group_size={G}"
+    w = weight.detach().float().reshape(N, K // G, G)
+    wmin = w.amin(dim=2)
+    wmax = w.amax(dim=2)
+    scale = ((wmax - wmin) / 15.0).clamp(min=1e-8)              # [N, K//G]
+    zero = torch.round(-wmin / scale).clamp(0, 15)             # [N, K//G]
+    q = torch.round(w / scale.unsqueeze(2) + zero.unsqueeze(2)).clamp(0, 15)
+    qweight = pack_int4(q.reshape(N, K).to(torch.int32))
+    return qweight, scale.to(torch.float16), zero.to(torch.uint8)
+def dequantize_w4a16_asym(qweight, scales, zeros, group_size: int = 128) -> torch.Tensor:
+    """Golden reference: reconstruct w_hat [N, K] float from the packed form."""
+    q = unpack_int4(qweight).float()                          # [N, K]
+    N, K = q.shape
+    G = group_size
+    q = q.reshape(N, K // G, G)
+    w = (q - zeros.float().unsqueeze(2)) * scales.float().unsqueeze(2)
+    return w.reshape(N, K)
+def gemv_reference(x, qweight, scales, zeros, group_size: int = 128) -> torch.Tensor:
+    """Reference W4A16 linear (no bias): y = x @ dequant(W).T."""
+    w_hat = dequantize_w4a16_asym(qweight, scales, zeros, group_size).to(x.dtype)
+    return x @ w_hat.t()

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

	@@ -0,0 +1,43 @@

+"""gemmaq: W4A16 quantized GEMV kernel + quantization utilities."""
+from . import quant
+from .quant import (
+    dequantize_w4a16_asym,
+    gemv_reference,
+    quantize_w4a16_asym,
+)
+# The compiled op (._ops) only exists in a built package; importing the source
+# tree (for quant utilities) must not fail.
+try:
+    from ._ops import ops
+except ImportError:
+    ops = None
+# Model-side layer + swap utility, and the hub-consumable kernel layers
+# (gemmaq.layers.W4A16Linear, resolved by kernels.kernelize via LayerRepository).
+try:
+    from . import layers
+    from .layer import W4A16Linear, quantize_model
+except Exception:  # torch.library API or ops unavailable
+    layers = None
+    W4A16Linear = quantize_model = None
+def w4a16_gemv(x, qweight, scales, zeros, group_size=128):
+    """y = x @ dequant(W).T using the compiled W4A16 GEMV kernel."""
+    if ops is None:
+        raise RuntimeError("gemmaq compiled ops not available (build with `make dev`)")
+    return ops.w4a16_gemv(x, qweight, scales, zeros, group_size)
+__all__ = [
+    "ops",
+    "w4a16_gemv",
+    "quant",
+    "quantize_w4a16_asym",
+    "dequantize_w4a16_asym",
+    "gemv_reference",
+    "W4A16Linear",
+    "quantize_model",
+    "layers",
+]

build/torch211-cxx11-cu130-x86_64-linux/_gemmaq_cuda_6820a9e.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:8b72b9a5fd737934030f9e7ce029be3338060e4cc3fea70849b0fd6b349a8deb
+size 126760

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

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

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

build/torch211-cxx11-cu130-x86_64-linux/gemmaq/__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/torch211-cxx11-cu130-x86_64-linux/layer.py ADDED Viewed

	@@ -0,0 +1,81 @@

+"""Model-side W4A16 layer + model-swap utility.
+`quantize_model` replaces the target `nn.Linear`s with `W4A16Linear` (which packs
+the weights to int4). `W4A16Linear` is decorated with
+`@use_kernel_forward_from_hub("W4A16Linear")`, so its pure-torch forward is the
+*fallback*; calling `kernels.kernelize(model, ...)` with a mapping to the
+`gemmaq.layers.W4A16Linear` kernel layer swaps in the compiled GEMV. The model is
+therefore runnable with or without the compiled kernel.
+"""
+import torch
+import torch.nn as nn
+from .quant import dequantize_w4a16_asym, quantize_w4a16_asym
+try:
+    from kernels import use_kernel_forward_from_hub
+except ImportError:  # kernels not installed: no-op decorator (forward still works)
+    def use_kernel_forward_from_hub(_name):
+        return lambda cls: cls
+# 7 linear types per decoder layer (lm_head deferred to Phase 2).
+DEFAULT_TARGETS = ("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj")
+_SKIP = ("vision", "audio", "multi_modal", "mm_")
+@use_kernel_forward_from_hub("W4A16Linear")
+class W4A16Linear(nn.Module):
+    """Weight-only int4 linear. Holds the quantized buffers; the forward here is
+    the pure-torch fallback (dequant + mm). kernelize() swaps in the kernel."""
+    def __init__(self, in_features, out_features, bias=False, group_size=128, device=None):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.group_size = group_size
+        self.register_buffer("qweight", torch.empty(out_features, in_features // 8, dtype=torch.int32, device=device))
+        self.register_buffer("scales", torch.empty(out_features, in_features // group_size, dtype=torch.float16, device=device))
+        self.register_buffer("zeros", torch.empty(out_features, in_features // group_size, dtype=torch.uint8, device=device))
+        self.register_buffer("bias", torch.empty(out_features, device=device) if bias else None)
+    @classmethod
+    def from_linear(cls, lin: nn.Linear, group_size=128):
+        qw, sc, zp = quantize_w4a16_asym(lin.weight.data, group_size)
+        m = cls(lin.in_features, lin.out_features, bias=lin.bias is not None,
+                group_size=group_size, device=lin.weight.device)
+        m.qweight.copy_(qw)
+        m.scales.copy_(sc)
+        m.zeros.copy_(zp)
+        if lin.bias is not None:
+            m.bias.copy_(lin.bias.data.to(m.bias.dtype))
+        return m
+    def forward(self, x):
+        w = dequantize_w4a16_asym(self.qweight, self.scales, self.zeros, self.group_size).to(x.dtype)
+        out = torch.matmul(x, w.t())
+        if self.bias is not None:
+            out = out + self.bias.to(out.dtype)
+        return out
+def quantize_model(model, group_size=128, targets=DEFAULT_TARGETS):
+    """In-place swap of the target nn.Linear layers with W4A16Linear. Returns count.
+    After this, call `kernels.kernelize(model, mode=...)` with a mapping for
+    `"W4A16Linear"` to activate the compiled kernel (otherwise the torch fallback runs).
+    """
+    to_swap = []
+    for name, mod in model.named_modules():
+        if (isinstance(mod, nn.Linear)
+                and name.rsplit(".", 1)[-1] in targets
+                and not any(s in name for s in _SKIP)
+                and mod.in_features % group_size == 0):
+            to_swap.append((name, mod))
+    for name, mod in to_swap:
+        parent = model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else model
+        child = name.rsplit(".", 1)[-1]
+        setattr(parent, child, W4A16Linear.from_linear(mod, group_size))
+        del mod
+    torch.cuda.empty_cache()
+    return len(to_swap)

build/torch211-cxx11-cu130-x86_64-linux/layers.py ADDED Viewed

	@@ -0,0 +1,48 @@

+"""Hub-consumable kernel layers, resolved by `kernels.kernelize` via a
+LayerRepository (`layer_name="W4A16Linear"`).
+Per the kernels layer contract these layers are STATELESS: kernelize binds only
+`forward` onto an existing module instance, so `self` carries the quantized
+buffers (qweight/scales/zeros/group_size/bias) from gemmaq.layer.W4A16Linear.
+The class must not define __init__ or extra members (only `forward` plus the
+`can_torch_compile`/`has_backward` flags).
+"""
+import torch
+import torch.nn as nn
+from .quant import dequantize_w4a16_asym
+try:
+    from ._ops import ops as _ops
+except ImportError:  # importable from source tree without the build
+    _ops = None
+def _w4a16_fake(x, qweight, scales, zeros, group_size):
+    return x.new_empty((*x.shape[:-1], qweight.shape[0]))
+# Register a meta/fake impl so the compiled op shape-propagates under torch.compile.
+if _ops is not None:
+    try:
+        torch.library.register_fake(_ops.w4a16_gemv.default.name(), _w4a16_fake)
+    except Exception:
+        pass  # already registered
+class W4A16Linear(nn.Module):
+    """Optimized W4A16 forward: compiled GEMV for decode (M=1), dequant+mm for
+    prefill (M>1, where the GEMV kernel loses to cuBLAS GEMM)."""
+    can_torch_compile = True
+    def forward(self, x):
+        m = x.numel() // x.shape[-1]
+        if m == 1:
+            out = _ops.w4a16_gemv(x, self.qweight, self.scales, self.zeros, self.group_size)
+        else:
+            w = dequantize_w4a16_asym(self.qweight, self.scales, self.zeros, self.group_size).to(x.dtype)
+            out = torch.matmul(x, w.t())
+        if self.bias is not None:
+            out = out + self.bias.to(out.dtype)
+        return out

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

	@@ -0,0 +1,25 @@

+{
+  "name": "gemmaq",
+  "id": "_gemmaq_cuda_6820a9e",
+  "version": 1,
+  "license": "MIT",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.9"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "nufuk/+eZjEX7g9AJLBII5Nr7alASHAUDgNRJTdMs1c=",
+      "_gemmaq_cuda_6820a9e.abi3.so": "i3K5pf1zeTQDD5584Cm+MzgGDkzD/qcISbD9azSajes=",
+      "_ops.py": "VrnRWJF6J13HL5IPqZyNmcgDlxYow62J+G2ZD6PU8Ac=",
+      "gemmaq/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
+      "layer.py": "Eml+uGNCtnO3dtrmIte8vtO+BteLT6wYfGUKvbbuz6E=",
+      "layers.py": "SM4DMpETjS/NXe7dMdjLR+r9NTcuO0gaqW26oaDKC5E=",
+      "quant.py": "OeosAKwPsCax7g8GfotC7FBbI/p/pVIg3tNoAFilQzw="
+    }
+  }
+}

build/torch211-cxx11-cu130-x86_64-linux/quant.py ADDED Viewed

	@@ -0,0 +1,67 @@

+"""W4A16 weight-only quantization: asymmetric, group-wise, RTN.
+Layout (must match csrc/w4a16_gemv.cu):
+  weight W: [N, K]  (PyTorch nn.Linear weight; y = x @ W.T)
+  quantize along K in groups of `group_size` (default 128).
+    scale[n,g] = (wmax-wmin)/15,  zero[n,g] = round(-wmin/scale)  in [0,15]
+    q[n,k]     = clamp(round(W/scale)+zero, 0, 15)                 uint4
+    w_hat      = (q - zero) * scale
+  qweight: int32[N, K//8], 8 nibbles packed along K per word
+           (nibble i -> k = 8*j + i, at bits [4i, 4i+4)).
+  scales:  fp16 [N, K//group_size]
+  zeros:   uint8[N, K//group_size]
+"""
+import torch
+PACK = 8  # int4 values per int32 word
+def pack_int4(q: torch.Tensor) -> torch.Tensor:
+    """q: [N, K] int (values 0..15) -> packed int32 [N, K//8]."""
+    N, K = q.shape
+    assert K % PACK == 0, f"K={K} must be divisible by {PACK}"
+    q = q.to(torch.int32).reshape(N, K // PACK, PACK)
+    packed = torch.zeros(N, K // PACK, dtype=torch.int32, device=q.device)
+    for i in range(PACK):
+        packed |= (q[:, :, i] & 0xF) << (4 * i)
+    return packed
+def unpack_int4(qweight: torch.Tensor) -> torch.Tensor:
+    """packed int32 [N, K//8] -> [N, K] int32 values 0..15."""
+    N, Kp = qweight.shape
+    out = torch.empty(N, Kp, PACK, dtype=torch.int32, device=qweight.device)
+    for i in range(PACK):
+        out[:, :, i] = (qweight >> (4 * i)) & 0xF  # & 0xF undoes arithmetic-shift sign bits
+    return out.reshape(N, Kp * PACK)
+def quantize_w4a16_asym(weight: torch.Tensor, group_size: int = 128):
+    """RTN asymmetric int4. Returns (qweight int32, scales fp16, zeros uint8)."""
+    N, K = weight.shape
+    G = group_size
+    assert K % G == 0, f"K={K} must be divisible by group_size={G}"
+    w = weight.detach().float().reshape(N, K // G, G)
+    wmin = w.amin(dim=2)
+    wmax = w.amax(dim=2)
+    scale = ((wmax - wmin) / 15.0).clamp(min=1e-8)              # [N, K//G]
+    zero = torch.round(-wmin / scale).clamp(0, 15)             # [N, K//G]
+    q = torch.round(w / scale.unsqueeze(2) + zero.unsqueeze(2)).clamp(0, 15)
+    qweight = pack_int4(q.reshape(N, K).to(torch.int32))
+    return qweight, scale.to(torch.float16), zero.to(torch.uint8)
+def dequantize_w4a16_asym(qweight, scales, zeros, group_size: int = 128) -> torch.Tensor:
+    """Golden reference: reconstruct w_hat [N, K] float from the packed form."""
+    q = unpack_int4(qweight).float()                          # [N, K]
+    N, K = q.shape
+    G = group_size
+    q = q.reshape(N, K // G, G)
+    w = (q - zeros.float().unsqueeze(2)) * scales.float().unsqueeze(2)
+    return w.reshape(N, K)
+def gemv_reference(x, qweight, scales, zeros, group_size: int = 128) -> torch.Tensor:
+    """Reference W4A16 linear (no bias): y = x @ dequant(W).T."""
+    w_hat = dequantize_w4a16_asym(qweight, scales, zeros, group_size).to(x.dtype)
+    return x @ w_hat.t()

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

	@@ -0,0 +1,43 @@

+"""gemmaq: W4A16 quantized GEMV kernel + quantization utilities."""
+from . import quant
+from .quant import (
+    dequantize_w4a16_asym,
+    gemv_reference,
+    quantize_w4a16_asym,
+)
+# The compiled op (._ops) only exists in a built package; importing the source
+# tree (for quant utilities) must not fail.
+try:
+    from ._ops import ops
+except ImportError:
+    ops = None
+# Model-side layer + swap utility, and the hub-consumable kernel layers
+# (gemmaq.layers.W4A16Linear, resolved by kernels.kernelize via LayerRepository).
+try:
+    from . import layers
+    from .layer import W4A16Linear, quantize_model
+except Exception:  # torch.library API or ops unavailable
+    layers = None
+    W4A16Linear = quantize_model = None
+def w4a16_gemv(x, qweight, scales, zeros, group_size=128):
+    """y = x @ dequant(W).T using the compiled W4A16 GEMV kernel."""
+    if ops is None:
+        raise RuntimeError("gemmaq compiled ops not available (build with `make dev`)")
+    return ops.w4a16_gemv(x, qweight, scales, zeros, group_size)
+__all__ = [
+    "ops",
+    "w4a16_gemv",
+    "quant",
+    "quantize_w4a16_asym",
+    "dequantize_w4a16_asym",
+    "gemv_reference",
+    "W4A16Linear",
+    "quantize_model",
+    "layers",
+]

build/torch212-cxx11-cu126-x86_64-linux/_gemmaq_cuda_6820a9e.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:0002d3886f6511b21dda2bfde912bb3dfac82bd3a3aa07e03a85b0decf439666
+size 110616

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

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

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

build/torch212-cxx11-cu126-x86_64-linux/gemmaq/__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-cu126-x86_64-linux/layer.py ADDED Viewed

	@@ -0,0 +1,81 @@

+"""Model-side W4A16 layer + model-swap utility.
+`quantize_model` replaces the target `nn.Linear`s with `W4A16Linear` (which packs
+the weights to int4). `W4A16Linear` is decorated with
+`@use_kernel_forward_from_hub("W4A16Linear")`, so its pure-torch forward is the
+*fallback*; calling `kernels.kernelize(model, ...)` with a mapping to the
+`gemmaq.layers.W4A16Linear` kernel layer swaps in the compiled GEMV. The model is
+therefore runnable with or without the compiled kernel.
+"""
+import torch
+import torch.nn as nn
+from .quant import dequantize_w4a16_asym, quantize_w4a16_asym
+try:
+    from kernels import use_kernel_forward_from_hub
+except ImportError:  # kernels not installed: no-op decorator (forward still works)
+    def use_kernel_forward_from_hub(_name):
+        return lambda cls: cls
+# 7 linear types per decoder layer (lm_head deferred to Phase 2).
+DEFAULT_TARGETS = ("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj")
+_SKIP = ("vision", "audio", "multi_modal", "mm_")
+@use_kernel_forward_from_hub("W4A16Linear")
+class W4A16Linear(nn.Module):
+    """Weight-only int4 linear. Holds the quantized buffers; the forward here is
+    the pure-torch fallback (dequant + mm). kernelize() swaps in the kernel."""
+    def __init__(self, in_features, out_features, bias=False, group_size=128, device=None):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.group_size = group_size
+        self.register_buffer("qweight", torch.empty(out_features, in_features // 8, dtype=torch.int32, device=device))
+        self.register_buffer("scales", torch.empty(out_features, in_features // group_size, dtype=torch.float16, device=device))
+        self.register_buffer("zeros", torch.empty(out_features, in_features // group_size, dtype=torch.uint8, device=device))
+        self.register_buffer("bias", torch.empty(out_features, device=device) if bias else None)
+    @classmethod
+    def from_linear(cls, lin: nn.Linear, group_size=128):
+        qw, sc, zp = quantize_w4a16_asym(lin.weight.data, group_size)
+        m = cls(lin.in_features, lin.out_features, bias=lin.bias is not None,
+                group_size=group_size, device=lin.weight.device)
+        m.qweight.copy_(qw)
+        m.scales.copy_(sc)
+        m.zeros.copy_(zp)
+        if lin.bias is not None:
+            m.bias.copy_(lin.bias.data.to(m.bias.dtype))
+        return m
+    def forward(self, x):
+        w = dequantize_w4a16_asym(self.qweight, self.scales, self.zeros, self.group_size).to(x.dtype)
+        out = torch.matmul(x, w.t())
+        if self.bias is not None:
+            out = out + self.bias.to(out.dtype)
+        return out
+def quantize_model(model, group_size=128, targets=DEFAULT_TARGETS):
+    """In-place swap of the target nn.Linear layers with W4A16Linear. Returns count.
+    After this, call `kernels.kernelize(model, mode=...)` with a mapping for
+    `"W4A16Linear"` to activate the compiled kernel (otherwise the torch fallback runs).
+    """
+    to_swap = []
+    for name, mod in model.named_modules():
+        if (isinstance(mod, nn.Linear)
+                and name.rsplit(".", 1)[-1] in targets
+                and not any(s in name for s in _SKIP)
+                and mod.in_features % group_size == 0):
+            to_swap.append((name, mod))
+    for name, mod in to_swap:
+        parent = model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else model
+        child = name.rsplit(".", 1)[-1]
+        setattr(parent, child, W4A16Linear.from_linear(mod, group_size))
+        del mod
+    torch.cuda.empty_cache()
+    return len(to_swap)

build/torch212-cxx11-cu126-x86_64-linux/layers.py ADDED Viewed

	@@ -0,0 +1,48 @@

+"""Hub-consumable kernel layers, resolved by `kernels.kernelize` via a
+LayerRepository (`layer_name="W4A16Linear"`).
+Per the kernels layer contract these layers are STATELESS: kernelize binds only
+`forward` onto an existing module instance, so `self` carries the quantized
+buffers (qweight/scales/zeros/group_size/bias) from gemmaq.layer.W4A16Linear.
+The class must not define __init__ or extra members (only `forward` plus the
+`can_torch_compile`/`has_backward` flags).
+"""
+import torch
+import torch.nn as nn
+from .quant import dequantize_w4a16_asym
+try:
+    from ._ops import ops as _ops
+except ImportError:  # importable from source tree without the build
+    _ops = None
+def _w4a16_fake(x, qweight, scales, zeros, group_size):
+    return x.new_empty((*x.shape[:-1], qweight.shape[0]))
+# Register a meta/fake impl so the compiled op shape-propagates under torch.compile.
+if _ops is not None:
+    try:
+        torch.library.register_fake(_ops.w4a16_gemv.default.name(), _w4a16_fake)
+    except Exception:
+        pass  # already registered
+class W4A16Linear(nn.Module):
+    """Optimized W4A16 forward: compiled GEMV for decode (M=1), dequant+mm for
+    prefill (M>1, where the GEMV kernel loses to cuBLAS GEMM)."""
+    can_torch_compile = True
+    def forward(self, x):
+        m = x.numel() // x.shape[-1]
+        if m == 1:
+            out = _ops.w4a16_gemv(x, self.qweight, self.scales, self.zeros, self.group_size)
+        else:
+            w = dequantize_w4a16_asym(self.qweight, self.scales, self.zeros, self.group_size).to(x.dtype)
+            out = torch.matmul(x, w.t())
+        if self.bias is not None:
+            out = out + self.bias.to(out.dtype)
+        return out

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

	@@ -0,0 +1,25 @@

+{
+  "name": "gemmaq",
+  "id": "_gemmaq_cuda_6820a9e",
+  "version": 1,
+  "license": "MIT",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.9"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "nufuk/+eZjEX7g9AJLBII5Nr7alASHAUDgNRJTdMs1c=",
+      "_gemmaq_cuda_6820a9e.abi3.so": "AALTiG9lEbId2iv96RK7PfrIK9OjqgfgOoWw3s9DlmY=",
+      "_ops.py": "VrnRWJF6J13HL5IPqZyNmcgDlxYow62J+G2ZD6PU8Ac=",
+      "gemmaq/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
+      "layer.py": "Eml+uGNCtnO3dtrmIte8vtO+BteLT6wYfGUKvbbuz6E=",
+      "layers.py": "SM4DMpETjS/NXe7dMdjLR+r9NTcuO0gaqW26oaDKC5E=",
+      "quant.py": "OeosAKwPsCax7g8GfotC7FBbI/p/pVIg3tNoAFilQzw="
+    }
+  }
+}

build/torch212-cxx11-cu126-x86_64-linux/quant.py ADDED Viewed

	@@ -0,0 +1,67 @@

+"""W4A16 weight-only quantization: asymmetric, group-wise, RTN.
+Layout (must match csrc/w4a16_gemv.cu):
+  weight W: [N, K]  (PyTorch nn.Linear weight; y = x @ W.T)
+  quantize along K in groups of `group_size` (default 128).
+    scale[n,g] = (wmax-wmin)/15,  zero[n,g] = round(-wmin/scale)  in [0,15]
+    q[n,k]     = clamp(round(W/scale)+zero, 0, 15)                 uint4
+    w_hat      = (q - zero) * scale
+  qweight: int32[N, K//8], 8 nibbles packed along K per word
+           (nibble i -> k = 8*j + i, at bits [4i, 4i+4)).
+  scales:  fp16 [N, K//group_size]
+  zeros:   uint8[N, K//group_size]
+"""
+import torch
+PACK = 8  # int4 values per int32 word
+def pack_int4(q: torch.Tensor) -> torch.Tensor:
+    """q: [N, K] int (values 0..15) -> packed int32 [N, K//8]."""
+    N, K = q.shape
+    assert K % PACK == 0, f"K={K} must be divisible by {PACK}"
+    q = q.to(torch.int32).reshape(N, K // PACK, PACK)
+    packed = torch.zeros(N, K // PACK, dtype=torch.int32, device=q.device)
+    for i in range(PACK):
+        packed |= (q[:, :, i] & 0xF) << (4 * i)
+    return packed
+def unpack_int4(qweight: torch.Tensor) -> torch.Tensor:
+    """packed int32 [N, K//8] -> [N, K] int32 values 0..15."""
+    N, Kp = qweight.shape
+    out = torch.empty(N, Kp, PACK, dtype=torch.int32, device=qweight.device)
+    for i in range(PACK):
+        out[:, :, i] = (qweight >> (4 * i)) & 0xF  # & 0xF undoes arithmetic-shift sign bits
+    return out.reshape(N, Kp * PACK)
+def quantize_w4a16_asym(weight: torch.Tensor, group_size: int = 128):
+    """RTN asymmetric int4. Returns (qweight int32, scales fp16, zeros uint8)."""
+    N, K = weight.shape
+    G = group_size
+    assert K % G == 0, f"K={K} must be divisible by group_size={G}"
+    w = weight.detach().float().reshape(N, K // G, G)
+    wmin = w.amin(dim=2)
+    wmax = w.amax(dim=2)
+    scale = ((wmax - wmin) / 15.0).clamp(min=1e-8)              # [N, K//G]
+    zero = torch.round(-wmin / scale).clamp(0, 15)             # [N, K//G]
+    q = torch.round(w / scale.unsqueeze(2) + zero.unsqueeze(2)).clamp(0, 15)
+    qweight = pack_int4(q.reshape(N, K).to(torch.int32))
+    return qweight, scale.to(torch.float16), zero.to(torch.uint8)
+def dequantize_w4a16_asym(qweight, scales, zeros, group_size: int = 128) -> torch.Tensor:
+    """Golden reference: reconstruct w_hat [N, K] float from the packed form."""
+    q = unpack_int4(qweight).float()                          # [N, K]
+    N, K = q.shape
+    G = group_size
+    q = q.reshape(N, K // G, G)
+    w = (q - zeros.float().unsqueeze(2)) * scales.float().unsqueeze(2)
+    return w.reshape(N, K)
+def gemv_reference(x, qweight, scales, zeros, group_size: int = 128) -> torch.Tensor:
+    """Reference W4A16 linear (no bias): y = x @ dequant(W).T."""
+    w_hat = dequantize_w4a16_asym(qweight, scales, zeros, group_size).to(x.dtype)
+    return x @ w_hat.t()

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

	@@ -0,0 +1,43 @@

+"""gemmaq: W4A16 quantized GEMV kernel + quantization utilities."""
+from . import quant
+from .quant import (
+    dequantize_w4a16_asym,
+    gemv_reference,
+    quantize_w4a16_asym,
+)
+# The compiled op (._ops) only exists in a built package; importing the source
+# tree (for quant utilities) must not fail.
+try:
+    from ._ops import ops
+except ImportError:
+    ops = None
+# Model-side layer + swap utility, and the hub-consumable kernel layers
+# (gemmaq.layers.W4A16Linear, resolved by kernels.kernelize via LayerRepository).
+try:
+    from . import layers
+    from .layer import W4A16Linear, quantize_model
+except Exception:  # torch.library API or ops unavailable
+    layers = None
+    W4A16Linear = quantize_model = None
+def w4a16_gemv(x, qweight, scales, zeros, group_size=128):
+    """y = x @ dequant(W).T using the compiled W4A16 GEMV kernel."""
+    if ops is None:
+        raise RuntimeError("gemmaq compiled ops not available (build with `make dev`)")
+    return ops.w4a16_gemv(x, qweight, scales, zeros, group_size)
+__all__ = [
+    "ops",
+    "w4a16_gemv",
+    "quant",
+    "quantize_w4a16_asym",
+    "dequantize_w4a16_asym",
+    "gemv_reference",
+    "W4A16Linear",
+    "quantize_model",
+    "layers",
+]

build/torch212-cxx11-cu130-x86_64-linux/_gemmaq_cuda_6820a9e.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:374f596b66f0e821feb268c44b787d480ed24fa14a0d859001c8157a452f3503
+size 126768

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

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

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

build/torch212-cxx11-cu130-x86_64-linux/gemmaq/__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-cu130-x86_64-linux/layer.py ADDED Viewed

	@@ -0,0 +1,81 @@

+"""Model-side W4A16 layer + model-swap utility.
+`quantize_model` replaces the target `nn.Linear`s with `W4A16Linear` (which packs
+the weights to int4). `W4A16Linear` is decorated with
+`@use_kernel_forward_from_hub("W4A16Linear")`, so its pure-torch forward is the
+*fallback*; calling `kernels.kernelize(model, ...)` with a mapping to the
+`gemmaq.layers.W4A16Linear` kernel layer swaps in the compiled GEMV. The model is
+therefore runnable with or without the compiled kernel.
+"""
+import torch
+import torch.nn as nn
+from .quant import dequantize_w4a16_asym, quantize_w4a16_asym
+try:
+    from kernels import use_kernel_forward_from_hub
+except ImportError:  # kernels not installed: no-op decorator (forward still works)
+    def use_kernel_forward_from_hub(_name):
+        return lambda cls: cls
+# 7 linear types per decoder layer (lm_head deferred to Phase 2).
+DEFAULT_TARGETS = ("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj")
+_SKIP = ("vision", "audio", "multi_modal", "mm_")
+@use_kernel_forward_from_hub("W4A16Linear")
+class W4A16Linear(nn.Module):
+    """Weight-only int4 linear. Holds the quantized buffers; the forward here is
+    the pure-torch fallback (dequant + mm). kernelize() swaps in the kernel."""
+    def __init__(self, in_features, out_features, bias=False, group_size=128, device=None):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.group_size = group_size
+        self.register_buffer("qweight", torch.empty(out_features, in_features // 8, dtype=torch.int32, device=device))
+        self.register_buffer("scales", torch.empty(out_features, in_features // group_size, dtype=torch.float16, device=device))
+        self.register_buffer("zeros", torch.empty(out_features, in_features // group_size, dtype=torch.uint8, device=device))
+        self.register_buffer("bias", torch.empty(out_features, device=device) if bias else None)
+    @classmethod
+    def from_linear(cls, lin: nn.Linear, group_size=128):
+        qw, sc, zp = quantize_w4a16_asym(lin.weight.data, group_size)
+        m = cls(lin.in_features, lin.out_features, bias=lin.bias is not None,
+                group_size=group_size, device=lin.weight.device)
+        m.qweight.copy_(qw)
+        m.scales.copy_(sc)
+        m.zeros.copy_(zp)
+        if lin.bias is not None:
+            m.bias.copy_(lin.bias.data.to(m.bias.dtype))
+        return m
+    def forward(self, x):
+        w = dequantize_w4a16_asym(self.qweight, self.scales, self.zeros, self.group_size).to(x.dtype)
+        out = torch.matmul(x, w.t())
+        if self.bias is not None:
+            out = out + self.bias.to(out.dtype)
+        return out
+def quantize_model(model, group_size=128, targets=DEFAULT_TARGETS):
+    """In-place swap of the target nn.Linear layers with W4A16Linear. Returns count.
+    After this, call `kernels.kernelize(model, mode=...)` with a mapping for
+    `"W4A16Linear"` to activate the compiled kernel (otherwise the torch fallback runs).
+    """
+    to_swap = []
+    for name, mod in model.named_modules():
+        if (isinstance(mod, nn.Linear)
+                and name.rsplit(".", 1)[-1] in targets
+                and not any(s in name for s in _SKIP)
+                and mod.in_features % group_size == 0):
+            to_swap.append((name, mod))
+    for name, mod in to_swap:
+        parent = model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else model
+        child = name.rsplit(".", 1)[-1]
+        setattr(parent, child, W4A16Linear.from_linear(mod, group_size))
+        del mod
+    torch.cuda.empty_cache()
+    return len(to_swap)

build/torch212-cxx11-cu130-x86_64-linux/layers.py ADDED Viewed

	@@ -0,0 +1,48 @@

+"""Hub-consumable kernel layers, resolved by `kernels.kernelize` via a
+LayerRepository (`layer_name="W4A16Linear"`).
+Per the kernels layer contract these layers are STATELESS: kernelize binds only
+`forward` onto an existing module instance, so `self` carries the quantized
+buffers (qweight/scales/zeros/group_size/bias) from gemmaq.layer.W4A16Linear.
+The class must not define __init__ or extra members (only `forward` plus the
+`can_torch_compile`/`has_backward` flags).
+"""
+import torch
+import torch.nn as nn
+from .quant import dequantize_w4a16_asym
+try:
+    from ._ops import ops as _ops
+except ImportError:  # importable from source tree without the build
+    _ops = None
+def _w4a16_fake(x, qweight, scales, zeros, group_size):
+    return x.new_empty((*x.shape[:-1], qweight.shape[0]))
+# Register a meta/fake impl so the compiled op shape-propagates under torch.compile.
+if _ops is not None:
+    try:
+        torch.library.register_fake(_ops.w4a16_gemv.default.name(), _w4a16_fake)
+    except Exception:
+        pass  # already registered
+class W4A16Linear(nn.Module):
+    """Optimized W4A16 forward: compiled GEMV for decode (M=1), dequant+mm for
+    prefill (M>1, where the GEMV kernel loses to cuBLAS GEMM)."""
+    can_torch_compile = True
+    def forward(self, x):
+        m = x.numel() // x.shape[-1]
+        if m == 1:
+            out = _ops.w4a16_gemv(x, self.qweight, self.scales, self.zeros, self.group_size)
+        else:
+            w = dequantize_w4a16_asym(self.qweight, self.scales, self.zeros, self.group_size).to(x.dtype)
+            out = torch.matmul(x, w.t())
+        if self.bias is not None:
+            out = out + self.bias.to(out.dtype)
+        return out

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

	@@ -0,0 +1,25 @@

+{
+  "name": "gemmaq",
+  "id": "_gemmaq_cuda_6820a9e",
+  "version": 1,
+  "license": "MIT",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.9"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "nufuk/+eZjEX7g9AJLBII5Nr7alASHAUDgNRJTdMs1c=",
+      "_gemmaq_cuda_6820a9e.abi3.so": "N09Za2bw6CH+smjES3h9SA7ST6FKDYWQAcgVekUvNQM=",
+      "_ops.py": "VrnRWJF6J13HL5IPqZyNmcgDlxYow62J+G2ZD6PU8Ac=",
+      "gemmaq/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
+      "layer.py": "Eml+uGNCtnO3dtrmIte8vtO+BteLT6wYfGUKvbbuz6E=",
+      "layers.py": "SM4DMpETjS/NXe7dMdjLR+r9NTcuO0gaqW26oaDKC5E=",
+      "quant.py": "OeosAKwPsCax7g8GfotC7FBbI/p/pVIg3tNoAFilQzw="
+    }
+  }
+}

build/torch212-cxx11-cu130-x86_64-linux/quant.py ADDED Viewed

	@@ -0,0 +1,67 @@

+"""W4A16 weight-only quantization: asymmetric, group-wise, RTN.
+Layout (must match csrc/w4a16_gemv.cu):
+  weight W: [N, K]  (PyTorch nn.Linear weight; y = x @ W.T)
+  quantize along K in groups of `group_size` (default 128).
+    scale[n,g] = (wmax-wmin)/15,  zero[n,g] = round(-wmin/scale)  in [0,15]
+    q[n,k]     = clamp(round(W/scale)+zero, 0, 15)                 uint4
+    w_hat      = (q - zero) * scale
+  qweight: int32[N, K//8], 8 nibbles packed along K per word
+           (nibble i -> k = 8*j + i, at bits [4i, 4i+4)).
+  scales:  fp16 [N, K//group_size]
+  zeros:   uint8[N, K//group_size]
+"""
+import torch
+PACK = 8  # int4 values per int32 word
+def pack_int4(q: torch.Tensor) -> torch.Tensor:
+    """q: [N, K] int (values 0..15) -> packed int32 [N, K//8]."""
+    N, K = q.shape
+    assert K % PACK == 0, f"K={K} must be divisible by {PACK}"
+    q = q.to(torch.int32).reshape(N, K // PACK, PACK)
+    packed = torch.zeros(N, K // PACK, dtype=torch.int32, device=q.device)
+    for i in range(PACK):
+        packed |= (q[:, :, i] & 0xF) << (4 * i)
+    return packed
+def unpack_int4(qweight: torch.Tensor) -> torch.Tensor:
+    """packed int32 [N, K//8] -> [N, K] int32 values 0..15."""
+    N, Kp = qweight.shape
+    out = torch.empty(N, Kp, PACK, dtype=torch.int32, device=qweight.device)
+    for i in range(PACK):
+        out[:, :, i] = (qweight >> (4 * i)) & 0xF  # & 0xF undoes arithmetic-shift sign bits
+    return out.reshape(N, Kp * PACK)
+def quantize_w4a16_asym(weight: torch.Tensor, group_size: int = 128):
+    """RTN asymmetric int4. Returns (qweight int32, scales fp16, zeros uint8)."""
+    N, K = weight.shape
+    G = group_size
+    assert K % G == 0, f"K={K} must be divisible by group_size={G}"
+    w = weight.detach().float().reshape(N, K // G, G)
+    wmin = w.amin(dim=2)
+    wmax = w.amax(dim=2)
+    scale = ((wmax - wmin) / 15.0).clamp(min=1e-8)              # [N, K//G]
+    zero = torch.round(-wmin / scale).clamp(0, 15)             # [N, K//G]
+    q = torch.round(w / scale.unsqueeze(2) + zero.unsqueeze(2)).clamp(0, 15)
+    qweight = pack_int4(q.reshape(N, K).to(torch.int32))
+    return qweight, scale.to(torch.float16), zero.to(torch.uint8)
+def dequantize_w4a16_asym(qweight, scales, zeros, group_size: int = 128) -> torch.Tensor:
+    """Golden reference: reconstruct w_hat [N, K] float from the packed form."""
+    q = unpack_int4(qweight).float()                          # [N, K]
+    N, K = q.shape
+    G = group_size
+    q = q.reshape(N, K // G, G)
+    w = (q - zeros.float().unsqueeze(2)) * scales.float().unsqueeze(2)
+    return w.reshape(N, K)
+def gemv_reference(x, qweight, scales, zeros, group_size: int = 128) -> torch.Tensor:
+    """Reference W4A16 linear (no bias): y = x @ dequant(W).T."""
+    w_hat = dequantize_w4a16_asym(qweight, scales, zeros, group_size).to(x.dtype)
+    return x @ w_hat.t()

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

	@@ -0,0 +1,43 @@

+"""gemmaq: W4A16 quantized GEMV kernel + quantization utilities."""
+from . import quant
+from .quant import (
+    dequantize_w4a16_asym,
+    gemv_reference,
+    quantize_w4a16_asym,
+)
+# The compiled op (._ops) only exists in a built package; importing the source
+# tree (for quant utilities) must not fail.
+try:
+    from ._ops import ops
+except ImportError:
+    ops = None
+# Model-side layer + swap utility, and the hub-consumable kernel layers
+# (gemmaq.layers.W4A16Linear, resolved by kernels.kernelize via LayerRepository).
+try:
+    from . import layers
+    from .layer import W4A16Linear, quantize_model
+except Exception:  # torch.library API or ops unavailable
+    layers = None
+    W4A16Linear = quantize_model = None
+def w4a16_gemv(x, qweight, scales, zeros, group_size=128):
+    """y = x @ dequant(W).T using the compiled W4A16 GEMV kernel."""
+    if ops is None:
+        raise RuntimeError("gemmaq compiled ops not available (build with `make dev`)")
+    return ops.w4a16_gemv(x, qweight, scales, zeros, group_size)
+__all__ = [
+    "ops",
+    "w4a16_gemv",
+    "quant",
+    "quantize_w4a16_asym",
+    "dequantize_w4a16_asym",
+    "gemv_reference",
+    "W4A16Linear",
+    "quantize_model",
+    "layers",
+]

build/torch212-cxx11-cu132-x86_64-linux/_gemmaq_cuda_6820a9e.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:383b0e8b74e09e214c812e32e8b6569ec9a5dc481162f2197190170aa71f7cd6
+size 136616

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

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

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

build/torch212-cxx11-cu132-x86_64-linux/gemmaq/__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-cu132-x86_64-linux/layer.py ADDED Viewed

	@@ -0,0 +1,81 @@

+"""Model-side W4A16 layer + model-swap utility.
+`quantize_model` replaces the target `nn.Linear`s with `W4A16Linear` (which packs
+the weights to int4). `W4A16Linear` is decorated with
+`@use_kernel_forward_from_hub("W4A16Linear")`, so its pure-torch forward is the
+*fallback*; calling `kernels.kernelize(model, ...)` with a mapping to the
+`gemmaq.layers.W4A16Linear` kernel layer swaps in the compiled GEMV. The model is
+therefore runnable with or without the compiled kernel.
+"""
+import torch
+import torch.nn as nn
+from .quant import dequantize_w4a16_asym, quantize_w4a16_asym
+try:
+    from kernels import use_kernel_forward_from_hub
+except ImportError:  # kernels not installed: no-op decorator (forward still works)
+    def use_kernel_forward_from_hub(_name):
+        return lambda cls: cls
+# 7 linear types per decoder layer (lm_head deferred to Phase 2).
+DEFAULT_TARGETS = ("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj")
+_SKIP = ("vision", "audio", "multi_modal", "mm_")
+@use_kernel_forward_from_hub("W4A16Linear")
+class W4A16Linear(nn.Module):
+    """Weight-only int4 linear. Holds the quantized buffers; the forward here is
+    the pure-torch fallback (dequant + mm). kernelize() swaps in the kernel."""
+    def __init__(self, in_features, out_features, bias=False, group_size=128, device=None):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.group_size = group_size
+        self.register_buffer("qweight", torch.empty(out_features, in_features // 8, dtype=torch.int32, device=device))
+        self.register_buffer("scales", torch.empty(out_features, in_features // group_size, dtype=torch.float16, device=device))
+        self.register_buffer("zeros", torch.empty(out_features, in_features // group_size, dtype=torch.uint8, device=device))
+        self.register_buffer("bias", torch.empty(out_features, device=device) if bias else None)
+    @classmethod
+    def from_linear(cls, lin: nn.Linear, group_size=128):
+        qw, sc, zp = quantize_w4a16_asym(lin.weight.data, group_size)
+        m = cls(lin.in_features, lin.out_features, bias=lin.bias is not None,
+                group_size=group_size, device=lin.weight.device)
+        m.qweight.copy_(qw)
+        m.scales.copy_(sc)
+        m.zeros.copy_(zp)
+        if lin.bias is not None:
+            m.bias.copy_(lin.bias.data.to(m.bias.dtype))
+        return m
+    def forward(self, x):
+        w = dequantize_w4a16_asym(self.qweight, self.scales, self.zeros, self.group_size).to(x.dtype)
+        out = torch.matmul(x, w.t())
+        if self.bias is not None:
+            out = out + self.bias.to(out.dtype)
+        return out
+def quantize_model(model, group_size=128, targets=DEFAULT_TARGETS):
+    """In-place swap of the target nn.Linear layers with W4A16Linear. Returns count.
+    After this, call `kernels.kernelize(model, mode=...)` with a mapping for
+    `"W4A16Linear"` to activate the compiled kernel (otherwise the torch fallback runs).
+    """
+    to_swap = []
+    for name, mod in model.named_modules():
+        if (isinstance(mod, nn.Linear)
+                and name.rsplit(".", 1)[-1] in targets
+                and not any(s in name for s in _SKIP)
+                and mod.in_features % group_size == 0):
+            to_swap.append((name, mod))
+    for name, mod in to_swap:
+        parent = model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else model
+        child = name.rsplit(".", 1)[-1]
+        setattr(parent, child, W4A16Linear.from_linear(mod, group_size))
+        del mod
+    torch.cuda.empty_cache()
+    return len(to_swap)

build/torch212-cxx11-cu132-x86_64-linux/layers.py ADDED Viewed

	@@ -0,0 +1,48 @@

+"""Hub-consumable kernel layers, resolved by `kernels.kernelize` via a
+LayerRepository (`layer_name="W4A16Linear"`).
+Per the kernels layer contract these layers are STATELESS: kernelize binds only
+`forward` onto an existing module instance, so `self` carries the quantized
+buffers (qweight/scales/zeros/group_size/bias) from gemmaq.layer.W4A16Linear.
+The class must not define __init__ or extra members (only `forward` plus the
+`can_torch_compile`/`has_backward` flags).
+"""
+import torch
+import torch.nn as nn
+from .quant import dequantize_w4a16_asym
+try:
+    from ._ops import ops as _ops
+except ImportError:  # importable from source tree without the build
+    _ops = None
+def _w4a16_fake(x, qweight, scales, zeros, group_size):
+    return x.new_empty((*x.shape[:-1], qweight.shape[0]))
+# Register a meta/fake impl so the compiled op shape-propagates under torch.compile.
+if _ops is not None:
+    try:
+        torch.library.register_fake(_ops.w4a16_gemv.default.name(), _w4a16_fake)
+    except Exception:
+        pass  # already registered
+class W4A16Linear(nn.Module):
+    """Optimized W4A16 forward: compiled GEMV for decode (M=1), dequant+mm for
+    prefill (M>1, where the GEMV kernel loses to cuBLAS GEMM)."""
+    can_torch_compile = True
+    def forward(self, x):
+        m = x.numel() // x.shape[-1]
+        if m == 1:
+            out = _ops.w4a16_gemv(x, self.qweight, self.scales, self.zeros, self.group_size)
+        else:
+            w = dequantize_w4a16_asym(self.qweight, self.scales, self.zeros, self.group_size).to(x.dtype)
+            out = torch.matmul(x, w.t())
+        if self.bias is not None:
+            out = out + self.bias.to(out.dtype)
+        return out

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

	@@ -0,0 +1,25 @@

+{
+  "name": "gemmaq",
+  "id": "_gemmaq_cuda_6820a9e",
+  "version": 1,
+  "license": "MIT",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.9"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "nufuk/+eZjEX7g9AJLBII5Nr7alASHAUDgNRJTdMs1c=",
+      "_gemmaq_cuda_6820a9e.abi3.so": "ODsOi3TgniFMgS4y6LZWnsml3EgRYvIZcZAXCqcffNY=",
+      "_ops.py": "VrnRWJF6J13HL5IPqZyNmcgDlxYow62J+G2ZD6PU8Ac=",
+      "gemmaq/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
+      "layer.py": "Eml+uGNCtnO3dtrmIte8vtO+BteLT6wYfGUKvbbuz6E=",
+      "layers.py": "SM4DMpETjS/NXe7dMdjLR+r9NTcuO0gaqW26oaDKC5E=",
+      "quant.py": "OeosAKwPsCax7g8GfotC7FBbI/p/pVIg3tNoAFilQzw="
+    }
+  }
+}

build/torch212-cxx11-cu132-x86_64-linux/quant.py ADDED Viewed

	@@ -0,0 +1,67 @@

+"""W4A16 weight-only quantization: asymmetric, group-wise, RTN.
+Layout (must match csrc/w4a16_gemv.cu):
+  weight W: [N, K]  (PyTorch nn.Linear weight; y = x @ W.T)
+  quantize along K in groups of `group_size` (default 128).
+    scale[n,g] = (wmax-wmin)/15,  zero[n,g] = round(-wmin/scale)  in [0,15]
+    q[n,k]     = clamp(round(W/scale)+zero, 0, 15)                 uint4
+    w_hat      = (q - zero) * scale
+  qweight: int32[N, K//8], 8 nibbles packed along K per word
+           (nibble i -> k = 8*j + i, at bits [4i, 4i+4)).
+  scales:  fp16 [N, K//group_size]
+  zeros:   uint8[N, K//group_size]
+"""
+import torch
+PACK = 8  # int4 values per int32 word
+def pack_int4(q: torch.Tensor) -> torch.Tensor:
+    """q: [N, K] int (values 0..15) -> packed int32 [N, K//8]."""
+    N, K = q.shape
+    assert K % PACK == 0, f"K={K} must be divisible by {PACK}"
+    q = q.to(torch.int32).reshape(N, K // PACK, PACK)
+    packed = torch.zeros(N, K // PACK, dtype=torch.int32, device=q.device)
+    for i in range(PACK):
+        packed |= (q[:, :, i] & 0xF) << (4 * i)
+    return packed
+def unpack_int4(qweight: torch.Tensor) -> torch.Tensor:
+    """packed int32 [N, K//8] -> [N, K] int32 values 0..15."""
+    N, Kp = qweight.shape
+    out = torch.empty(N, Kp, PACK, dtype=torch.int32, device=qweight.device)
+    for i in range(PACK):
+        out[:, :, i] = (qweight >> (4 * i)) & 0xF  # & 0xF undoes arithmetic-shift sign bits
+    return out.reshape(N, Kp * PACK)
+def quantize_w4a16_asym(weight: torch.Tensor, group_size: int = 128):
+    """RTN asymmetric int4. Returns (qweight int32, scales fp16, zeros uint8)."""
+    N, K = weight.shape
+    G = group_size
+    assert K % G == 0, f"K={K} must be divisible by group_size={G}"
+    w = weight.detach().float().reshape(N, K // G, G)
+    wmin = w.amin(dim=2)
+    wmax = w.amax(dim=2)
+    scale = ((wmax - wmin) / 15.0).clamp(min=1e-8)              # [N, K//G]
+    zero = torch.round(-wmin / scale).clamp(0, 15)             # [N, K//G]
+    q = torch.round(w / scale.unsqueeze(2) + zero.unsqueeze(2)).clamp(0, 15)
+    qweight = pack_int4(q.reshape(N, K).to(torch.int32))
+    return qweight, scale.to(torch.float16), zero.to(torch.uint8)
+def dequantize_w4a16_asym(qweight, scales, zeros, group_size: int = 128) -> torch.Tensor:
+    """Golden reference: reconstruct w_hat [N, K] float from the packed form."""
+    q = unpack_int4(qweight).float()                          # [N, K]
+    N, K = q.shape
+    G = group_size
+    q = q.reshape(N, K // G, G)
+    w = (q - zeros.float().unsqueeze(2)) * scales.float().unsqueeze(2)
+    return w.reshape(N, K)
+def gemv_reference(x, qweight, scales, zeros, group_size: int = 128) -> torch.Tensor:
+    """Reference W4A16 linear (no bias): y = x @ dequant(W).T."""
+    w_hat = dequantize_w4a16_asym(qweight, scales, zeros, group_size).to(x.dtype)
+    return x @ w_hat.t()