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/_bitnet_cpu_cpu_bd669ef.abi3.so +3 -0
build/torch211-cxx11-cpu-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cpu-x86_64-linux/bitnet_cpu/__init__.py +26 -0
build/torch211-cxx11-cpu-x86_64-linux/metadata.json +30 -0
build/torch212-cxx11-cpu-x86_64-linux/__init__.py +229 -0
build/torch212-cxx11-cpu-x86_64-linux/_bitnet_cpu_cpu_bd669ef.abi3.so +3 -0
build/torch212-cxx11-cpu-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cpu-x86_64-linux/bitnet_cpu/__init__.py +26 -0
build/torch212-cxx11-cpu-x86_64-linux/metadata.json +30 -0
build/torch213-cxx11-cpu-x86_64-linux/__init__.py +229 -0
build/torch213-cxx11-cpu-x86_64-linux/_bitnet_cpu_cpu_bd669ef.abi3.so +3 -0
build/torch213-cxx11-cpu-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cpu-x86_64-linux/bitnet_cpu/__init__.py +26 -0
build/torch213-cxx11-cpu-x86_64-linux/metadata.json +30 -0

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

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

+"""bitnet-cpu: ternary x INT8 GEMM for BitNet b1.58 (W1.58 A8) on CPUs.
+Same 2-bit weight packing and API as phanerozoic/bitnet-tc (the CUDA member
+of the stack); AVX-512 VNNI / AVX-VNNI / AVX2 paths selected at runtime,
+with a portable scalar fallback.
+High-level API:
+    pack_weights(W)              ternary {-1,0,+1} int8 [N,K] -> packed uint8 [N,K//4]
+    quantize_activation(x)       bf16/f32 [..,K] -> (int8, per-row bf16 scale)
+    bitnet_gemm(...)             int8 activations x packed ternary weights -> bf16
+    bitnet_gemv_fused(...)       bf16/f32 activations (fused quantize) -> bf16, M<16
+    bitnet_linear(x, w, sw)      one-shot forward, auto-dispatches by M
+    BitLinear                    nn.Module wrapper
+    BitLinearKernel              kernelize layer for transformers BitNet modules
+"""
+from typing import Optional
+import torch
+import torch.nn as nn
+from ._ops import ops
+_OK_DTYPES = (torch.bfloat16, torch.float32)
+def pack_weights(W: torch.Tensor) -> torch.Tensor:
+    """Pack ternary weights {-1, 0, +1} into 2-bit codes, 4 per byte.
+    Encoding: -1 -> 1, 0 -> 2, +1 -> 3 (the kernel decodes via byte - 2).
+    Identical to the phanerozoic/bitnet-tc and Microsoft bitnet.cpp packing.
+    Args:
+        W: [N, K] int8 with values in {-1, 0, +1}.
+    Returns:
+        [N, K // 4] uint8.
+    """
+    assert W.dtype == torch.int8, f"W must be int8, got {W.dtype}"
+    assert W.dim() == 2, "W must be 2D"
+    N, K = W.shape
+    assert K % 4 == 0, f"K={K} must be a multiple of 4"
+    assert ((W >= -1) & (W <= 1)).all(), "W values must be in {-1, 0, +1}"
+    codes = (W.to(torch.int32) + 2).to(torch.uint8).view(N, K // 4, 4)
+    packed = (
+        codes[..., 0]
+        | (codes[..., 1] << 2)
+        | (codes[..., 2] << 4)
+        | (codes[..., 3] << 6)
+    ).contiguous()
+    return packed
+def quantize_activation(x: torch.Tensor, eps: float = 1e-5):
+    """Per-token absmax INT8 quantization.
+    Args:
+        x: [..., K] bf16 or f32 tensor.
+    Returns:
+        (x_int8 [M, K], scale [M] bf16) where M is the flattened leading dim.
+    """
+    x_flat = x.reshape(-1, x.shape[-1]).contiguous()
+    M, K = x_flat.shape
+    if x_flat.dtype in _OK_DTYPES and x_flat.device.type == "cpu":
+        x_int8 = torch.empty((M, K), dtype=torch.int8)
+        scale = torch.empty((M,), dtype=torch.bfloat16)
+        ops.quantize_act(x_int8, scale, x_flat)
+        return x_int8, scale
+    absmax = x_flat.abs().amax(dim=-1, keepdim=True).clamp(min=eps)
+    scale = absmax / 127.0
+    x_int8 = (x_flat / scale).round().clamp(-127, 127).to(torch.int8)
+    return x_int8, scale.squeeze(-1).to(torch.bfloat16)
+def bitnet_gemm(
+    x_int8: torch.Tensor,
+    w_packed: torch.Tensor,
+    scale_act: torch.Tensor,
+    scale_wt: torch.Tensor,
+) -> torch.Tensor:
+    """Ternary x INT8 GEMM. Returns [M, N] bf16."""
+    M, K = x_int8.shape
+    N = w_packed.shape[0]
+    out = torch.empty((M, N), dtype=torch.bfloat16)
+    ops.bitnet_gemm(out, x_int8, w_packed, scale_act, scale_wt, None)
+    return out
+def bitnet_gemv_fused(
+    x: torch.Tensor,
+    w_packed: torch.Tensor,
+    scale_wt: torch.Tensor,
+) -> torch.Tensor:
+    """Fused quantize + ternary GEMV for M < 16. Returns [M, N] bf16."""
+    M, K = x.shape
+    N = w_packed.shape[0]
+    out = torch.empty((M, N), dtype=torch.bfloat16)
+    ops.bitnet_gemv_fused(out, x, w_packed, scale_wt)
+    return out
+def bitnet_linear(
+    x: torch.Tensor,
+    w_packed: torch.Tensor,
+    scale_wt: torch.Tensor,
+) -> torch.Tensor:
+    """One-shot BitLinear forward. Auto-dispatches fused (M<16) vs split path.
+    Args:
+        x: [..., K] bf16 or f32.
+        w_packed: [N, K//4] uint8 packed ternary weights.
+        scale_wt: [N] bf16 per-output weight scale.
+    Returns:
+        [..., N] bf16.
+    """
+    assert x.dtype in _OK_DTYPES, f"x must be bf16 or f32, got {x.dtype}"
+    orig = x.shape
+    K = orig[-1]
+    x_flat = x.reshape(-1, K).contiguous()
+    M = x_flat.shape[0]
+    N = w_packed.shape[0]
+    if M < 16:
+        y = bitnet_gemv_fused(x_flat, w_packed, scale_wt)
+    else:
+        x_int8, scale_act = quantize_activation(x_flat)
+        y = bitnet_gemm(x_int8, w_packed, scale_act, scale_wt)
+    return y.view(*orig[:-1], N)
+class BitLinear(nn.Module):
+    """Inference-only BitLinear: ternary weights, per-token INT8 activations."""
+    def __init__(self, in_features: int, out_features: int, bias: bool = False):
+        super().__init__()
+        assert in_features % 32 == 0, f"in_features={in_features} must be % 32"
+        self.in_features = in_features
+        self.out_features = out_features
+        self.register_buffer(
+            "w_packed", torch.zeros(out_features, in_features // 4, dtype=torch.uint8)
+        )
+        self.register_buffer(
+            "scale_wt", torch.ones(out_features, dtype=torch.bfloat16)
+        )
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_features, dtype=torch.bfloat16))
+        else:
+            self.bias = None
+    @classmethod
+    def from_dense(cls, lin: nn.Linear) -> "BitLinear":
+        """Quantize a dense Linear via absmean ternarization (sanity-test only;
+        a real BitNet model is trained with QAT, not post-hoc quantized)."""
+        out_features, in_features = lin.weight.shape
+        bl = cls(in_features, out_features, bias=lin.bias is not None)
+        with torch.no_grad():
+            W = lin.weight.detach().float()
+            gamma = W.abs().mean(dim=-1, keepdim=True).clamp(min=1e-8)
+            Wq = (W / gamma).round().clamp(-1, 1).to(torch.int8)
+            bl.w_packed.copy_(pack_weights(Wq.cpu()))
+            bl.scale_wt.copy_(gamma.squeeze(-1).to(torch.bfloat16))
+            if bl.bias is not None:
+                bl.bias.copy_(lin.bias.detach().to(torch.bfloat16))
+        return bl
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = bitnet_linear(x, self.w_packed, self.scale_wt)
+        if self.bias is not None:
+            y = y + self.bias
+        return y
+class BitLinearKernel(nn.Module):
+    """kernelize layer for the transformers BitNet linear modules.
+    kernelize binds this forward onto the host BitLinear / AutoBitLinear module,
+    so self is that module and the forward reads its attributes (weight,
+    weight_scale, rms_norm, bias). The ternary weight is converted to the
+    kernel's 2-bit layout once and cached; each call runs per-token INT8
+    quantization and the VNNI/AVX2 ternary GEMM.
+    """
+    can_torch_compile: bool = False
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        dt = input.dtype
+        rms = getattr(self, "rms_norm", None)
+        if rms is not None:
+            input = rms(input)
+        if not hasattr(self, "_bitnet_cpu_packed"):
+            w = self.weight
+            if w.dtype == torch.uint8:
+                # transformers stores ternary as int8 viewed through uint8 (255 -> -1).
+                tern = w.contiguous().view(torch.int8)
+                if tern.shape[0] != self.out_features:        # bit-packed (out//4, in)
+                    from transformers.integrations.bitnet import unpack_weights
+                    tern = unpack_weights(w, dtype=torch.bfloat16)
+                tern = tern.round().clamp(-1, 1).to(torch.int8)
+            else:
+                g = w.detach().float().abs().mean().clamp(min=1e-5)
+                tern = (w.detach().float() / g).round().clamp(-1, 1).to(torch.int8)
+            self._bitnet_cpu_n = tern.shape[0]
+            self._bitnet_cpu_packed = pack_weights(tern.contiguous().cpu())
+            ws = getattr(self, "weight_scale", None)
+            if ws is not None:
+                s = ws.detach().float().reshape(-1)
+                s = s.expand(self._bitnet_cpu_n) if s.numel() == 1 else s
+            else:
+                s = w.detach().float().abs().mean().clamp(min=1e-5).reshape(1).expand(self._bitnet_cpu_n)
+            self._bitnet_cpu_sw = s.to(torch.bfloat16).contiguous().cpu()
+        shp = input.shape
+        x = input.reshape(-1, shp[-1])
+        if x.dtype not in _OK_DTYPES:
+            x = x.to(torch.float32)
+        x = x.contiguous()
+        out = bitnet_linear(x, self._bitnet_cpu_packed, self._bitnet_cpu_sw)
+        if getattr(self, "bias", None) is not None:
+            out = out + self.bias
+        return out.reshape(*shp[:-1], self._bitnet_cpu_n).to(dt)
+__all__ = [
+    "pack_weights",
+    "quantize_activation",
+    "bitnet_gemm",
+    "bitnet_gemv_fused",
+    "bitnet_linear",
+    "BitLinear",
+    "BitLinearKernel",
+]

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

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

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

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

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

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

build/torch211-cxx11-cpu-x86_64-linux/bitnet_cpu/__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-cpu-x86_64-linux/metadata.json ADDED Viewed

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

+{
+  "name": "bitnet-cpu",
+  "id": "_bitnet_cpu_cpu_bd669ef",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cpu"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "ipCv02rtmMOwCw75aWDraJag/on6FtbuuKkYLxWUgiA=",
+      "_bitnet_cpu_cpu_bd669ef.abi3.so": "8dgB20e/LOvOHaRhfcC7blFYMNB0hYdABhKy6H3jBUI=",
+      "_ops.py": "CpD3fRDWoqe63dtPPstdFtlIPKs13yPdbOGpg4aH2lk=",
+      "bitnet_cpu/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "2c516fc6678570c013319adeb64605fc8e8d5394",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "bd669ef68e49c5873a7810015ce6a9503fed5394",
+      "dirty": false
+    }
+  }
+}

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

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

+"""bitnet-cpu: ternary x INT8 GEMM for BitNet b1.58 (W1.58 A8) on CPUs.
+Same 2-bit weight packing and API as phanerozoic/bitnet-tc (the CUDA member
+of the stack); AVX-512 VNNI / AVX-VNNI / AVX2 paths selected at runtime,
+with a portable scalar fallback.
+High-level API:
+    pack_weights(W)              ternary {-1,0,+1} int8 [N,K] -> packed uint8 [N,K//4]
+    quantize_activation(x)       bf16/f32 [..,K] -> (int8, per-row bf16 scale)
+    bitnet_gemm(...)             int8 activations x packed ternary weights -> bf16
+    bitnet_gemv_fused(...)       bf16/f32 activations (fused quantize) -> bf16, M<16
+    bitnet_linear(x, w, sw)      one-shot forward, auto-dispatches by M
+    BitLinear                    nn.Module wrapper
+    BitLinearKernel              kernelize layer for transformers BitNet modules
+"""
+from typing import Optional
+import torch
+import torch.nn as nn
+from ._ops import ops
+_OK_DTYPES = (torch.bfloat16, torch.float32)
+def pack_weights(W: torch.Tensor) -> torch.Tensor:
+    """Pack ternary weights {-1, 0, +1} into 2-bit codes, 4 per byte.
+    Encoding: -1 -> 1, 0 -> 2, +1 -> 3 (the kernel decodes via byte - 2).
+    Identical to the phanerozoic/bitnet-tc and Microsoft bitnet.cpp packing.
+    Args:
+        W: [N, K] int8 with values in {-1, 0, +1}.
+    Returns:
+        [N, K // 4] uint8.
+    """
+    assert W.dtype == torch.int8, f"W must be int8, got {W.dtype}"
+    assert W.dim() == 2, "W must be 2D"
+    N, K = W.shape
+    assert K % 4 == 0, f"K={K} must be a multiple of 4"
+    assert ((W >= -1) & (W <= 1)).all(), "W values must be in {-1, 0, +1}"
+    codes = (W.to(torch.int32) + 2).to(torch.uint8).view(N, K // 4, 4)
+    packed = (
+        codes[..., 0]
+        | (codes[..., 1] << 2)
+        | (codes[..., 2] << 4)
+        | (codes[..., 3] << 6)
+    ).contiguous()
+    return packed
+def quantize_activation(x: torch.Tensor, eps: float = 1e-5):
+    """Per-token absmax INT8 quantization.
+    Args:
+        x: [..., K] bf16 or f32 tensor.
+    Returns:
+        (x_int8 [M, K], scale [M] bf16) where M is the flattened leading dim.
+    """
+    x_flat = x.reshape(-1, x.shape[-1]).contiguous()
+    M, K = x_flat.shape
+    if x_flat.dtype in _OK_DTYPES and x_flat.device.type == "cpu":
+        x_int8 = torch.empty((M, K), dtype=torch.int8)
+        scale = torch.empty((M,), dtype=torch.bfloat16)
+        ops.quantize_act(x_int8, scale, x_flat)
+        return x_int8, scale
+    absmax = x_flat.abs().amax(dim=-1, keepdim=True).clamp(min=eps)
+    scale = absmax / 127.0
+    x_int8 = (x_flat / scale).round().clamp(-127, 127).to(torch.int8)
+    return x_int8, scale.squeeze(-1).to(torch.bfloat16)
+def bitnet_gemm(
+    x_int8: torch.Tensor,
+    w_packed: torch.Tensor,
+    scale_act: torch.Tensor,
+    scale_wt: torch.Tensor,
+) -> torch.Tensor:
+    """Ternary x INT8 GEMM. Returns [M, N] bf16."""
+    M, K = x_int8.shape
+    N = w_packed.shape[0]
+    out = torch.empty((M, N), dtype=torch.bfloat16)
+    ops.bitnet_gemm(out, x_int8, w_packed, scale_act, scale_wt, None)
+    return out
+def bitnet_gemv_fused(
+    x: torch.Tensor,
+    w_packed: torch.Tensor,
+    scale_wt: torch.Tensor,
+) -> torch.Tensor:
+    """Fused quantize + ternary GEMV for M < 16. Returns [M, N] bf16."""
+    M, K = x.shape
+    N = w_packed.shape[0]
+    out = torch.empty((M, N), dtype=torch.bfloat16)
+    ops.bitnet_gemv_fused(out, x, w_packed, scale_wt)
+    return out
+def bitnet_linear(
+    x: torch.Tensor,
+    w_packed: torch.Tensor,
+    scale_wt: torch.Tensor,
+) -> torch.Tensor:
+    """One-shot BitLinear forward. Auto-dispatches fused (M<16) vs split path.
+    Args:
+        x: [..., K] bf16 or f32.
+        w_packed: [N, K//4] uint8 packed ternary weights.
+        scale_wt: [N] bf16 per-output weight scale.
+    Returns:
+        [..., N] bf16.
+    """
+    assert x.dtype in _OK_DTYPES, f"x must be bf16 or f32, got {x.dtype}"
+    orig = x.shape
+    K = orig[-1]
+    x_flat = x.reshape(-1, K).contiguous()
+    M = x_flat.shape[0]
+    N = w_packed.shape[0]
+    if M < 16:
+        y = bitnet_gemv_fused(x_flat, w_packed, scale_wt)
+    else:
+        x_int8, scale_act = quantize_activation(x_flat)
+        y = bitnet_gemm(x_int8, w_packed, scale_act, scale_wt)
+    return y.view(*orig[:-1], N)
+class BitLinear(nn.Module):
+    """Inference-only BitLinear: ternary weights, per-token INT8 activations."""
+    def __init__(self, in_features: int, out_features: int, bias: bool = False):
+        super().__init__()
+        assert in_features % 32 == 0, f"in_features={in_features} must be % 32"
+        self.in_features = in_features
+        self.out_features = out_features
+        self.register_buffer(
+            "w_packed", torch.zeros(out_features, in_features // 4, dtype=torch.uint8)
+        )
+        self.register_buffer(
+            "scale_wt", torch.ones(out_features, dtype=torch.bfloat16)
+        )
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_features, dtype=torch.bfloat16))
+        else:
+            self.bias = None
+    @classmethod
+    def from_dense(cls, lin: nn.Linear) -> "BitLinear":
+        """Quantize a dense Linear via absmean ternarization (sanity-test only;
+        a real BitNet model is trained with QAT, not post-hoc quantized)."""
+        out_features, in_features = lin.weight.shape
+        bl = cls(in_features, out_features, bias=lin.bias is not None)
+        with torch.no_grad():
+            W = lin.weight.detach().float()
+            gamma = W.abs().mean(dim=-1, keepdim=True).clamp(min=1e-8)
+            Wq = (W / gamma).round().clamp(-1, 1).to(torch.int8)
+            bl.w_packed.copy_(pack_weights(Wq.cpu()))
+            bl.scale_wt.copy_(gamma.squeeze(-1).to(torch.bfloat16))
+            if bl.bias is not None:
+                bl.bias.copy_(lin.bias.detach().to(torch.bfloat16))
+        return bl
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = bitnet_linear(x, self.w_packed, self.scale_wt)
+        if self.bias is not None:
+            y = y + self.bias
+        return y
+class BitLinearKernel(nn.Module):
+    """kernelize layer for the transformers BitNet linear modules.
+    kernelize binds this forward onto the host BitLinear / AutoBitLinear module,
+    so self is that module and the forward reads its attributes (weight,
+    weight_scale, rms_norm, bias). The ternary weight is converted to the
+    kernel's 2-bit layout once and cached; each call runs per-token INT8
+    quantization and the VNNI/AVX2 ternary GEMM.
+    """
+    can_torch_compile: bool = False
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        dt = input.dtype
+        rms = getattr(self, "rms_norm", None)
+        if rms is not None:
+            input = rms(input)
+        if not hasattr(self, "_bitnet_cpu_packed"):
+            w = self.weight
+            if w.dtype == torch.uint8:
+                # transformers stores ternary as int8 viewed through uint8 (255 -> -1).
+                tern = w.contiguous().view(torch.int8)
+                if tern.shape[0] != self.out_features:        # bit-packed (out//4, in)
+                    from transformers.integrations.bitnet import unpack_weights
+                    tern = unpack_weights(w, dtype=torch.bfloat16)
+                tern = tern.round().clamp(-1, 1).to(torch.int8)
+            else:
+                g = w.detach().float().abs().mean().clamp(min=1e-5)
+                tern = (w.detach().float() / g).round().clamp(-1, 1).to(torch.int8)
+            self._bitnet_cpu_n = tern.shape[0]
+            self._bitnet_cpu_packed = pack_weights(tern.contiguous().cpu())
+            ws = getattr(self, "weight_scale", None)
+            if ws is not None:
+                s = ws.detach().float().reshape(-1)
+                s = s.expand(self._bitnet_cpu_n) if s.numel() == 1 else s
+            else:
+                s = w.detach().float().abs().mean().clamp(min=1e-5).reshape(1).expand(self._bitnet_cpu_n)
+            self._bitnet_cpu_sw = s.to(torch.bfloat16).contiguous().cpu()
+        shp = input.shape
+        x = input.reshape(-1, shp[-1])
+        if x.dtype not in _OK_DTYPES:
+            x = x.to(torch.float32)
+        x = x.contiguous()
+        out = bitnet_linear(x, self._bitnet_cpu_packed, self._bitnet_cpu_sw)
+        if getattr(self, "bias", None) is not None:
+            out = out + self.bias
+        return out.reshape(*shp[:-1], self._bitnet_cpu_n).to(dt)
+__all__ = [
+    "pack_weights",
+    "quantize_activation",
+    "bitnet_gemm",
+    "bitnet_gemv_fused",
+    "bitnet_linear",
+    "BitLinear",
+    "BitLinearKernel",
+]

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:785e2eaef191ee927b27634d64b92b960ae004d20e2c287cad1ec758e03b50e0
+size 91392

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

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

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

build/torch212-cxx11-cpu-x86_64-linux/bitnet_cpu/__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/metadata.json ADDED Viewed

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

+{
+  "name": "bitnet-cpu",
+  "id": "_bitnet_cpu_cpu_bd669ef",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cpu"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "ipCv02rtmMOwCw75aWDraJag/on6FtbuuKkYLxWUgiA=",
+      "_bitnet_cpu_cpu_bd669ef.abi3.so": "eF4urvGR7pJ7J2NNZLkrlgrgBNIOLCh8rR7HWOA7UOA=",
+      "_ops.py": "CpD3fRDWoqe63dtPPstdFtlIPKs13yPdbOGpg4aH2lk=",
+      "bitnet_cpu/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "2c516fc6678570c013319adeb64605fc8e8d5394",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "bd669ef68e49c5873a7810015ce6a9503fed5394",
+      "dirty": false
+    }
+  }
+}

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

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

+"""bitnet-cpu: ternary x INT8 GEMM for BitNet b1.58 (W1.58 A8) on CPUs.
+Same 2-bit weight packing and API as phanerozoic/bitnet-tc (the CUDA member
+of the stack); AVX-512 VNNI / AVX-VNNI / AVX2 paths selected at runtime,
+with a portable scalar fallback.
+High-level API:
+    pack_weights(W)              ternary {-1,0,+1} int8 [N,K] -> packed uint8 [N,K//4]
+    quantize_activation(x)       bf16/f32 [..,K] -> (int8, per-row bf16 scale)
+    bitnet_gemm(...)             int8 activations x packed ternary weights -> bf16
+    bitnet_gemv_fused(...)       bf16/f32 activations (fused quantize) -> bf16, M<16
+    bitnet_linear(x, w, sw)      one-shot forward, auto-dispatches by M
+    BitLinear                    nn.Module wrapper
+    BitLinearKernel              kernelize layer for transformers BitNet modules
+"""
+from typing import Optional
+import torch
+import torch.nn as nn
+from ._ops import ops
+_OK_DTYPES = (torch.bfloat16, torch.float32)
+def pack_weights(W: torch.Tensor) -> torch.Tensor:
+    """Pack ternary weights {-1, 0, +1} into 2-bit codes, 4 per byte.
+    Encoding: -1 -> 1, 0 -> 2, +1 -> 3 (the kernel decodes via byte - 2).
+    Identical to the phanerozoic/bitnet-tc and Microsoft bitnet.cpp packing.
+    Args:
+        W: [N, K] int8 with values in {-1, 0, +1}.
+    Returns:
+        [N, K // 4] uint8.
+    """
+    assert W.dtype == torch.int8, f"W must be int8, got {W.dtype}"
+    assert W.dim() == 2, "W must be 2D"
+    N, K = W.shape
+    assert K % 4 == 0, f"K={K} must be a multiple of 4"
+    assert ((W >= -1) & (W <= 1)).all(), "W values must be in {-1, 0, +1}"
+    codes = (W.to(torch.int32) + 2).to(torch.uint8).view(N, K // 4, 4)
+    packed = (
+        codes[..., 0]
+        | (codes[..., 1] << 2)
+        | (codes[..., 2] << 4)
+        | (codes[..., 3] << 6)
+    ).contiguous()
+    return packed
+def quantize_activation(x: torch.Tensor, eps: float = 1e-5):
+    """Per-token absmax INT8 quantization.
+    Args:
+        x: [..., K] bf16 or f32 tensor.
+    Returns:
+        (x_int8 [M, K], scale [M] bf16) where M is the flattened leading dim.
+    """
+    x_flat = x.reshape(-1, x.shape[-1]).contiguous()
+    M, K = x_flat.shape
+    if x_flat.dtype in _OK_DTYPES and x_flat.device.type == "cpu":
+        x_int8 = torch.empty((M, K), dtype=torch.int8)
+        scale = torch.empty((M,), dtype=torch.bfloat16)
+        ops.quantize_act(x_int8, scale, x_flat)
+        return x_int8, scale
+    absmax = x_flat.abs().amax(dim=-1, keepdim=True).clamp(min=eps)
+    scale = absmax / 127.0
+    x_int8 = (x_flat / scale).round().clamp(-127, 127).to(torch.int8)
+    return x_int8, scale.squeeze(-1).to(torch.bfloat16)
+def bitnet_gemm(
+    x_int8: torch.Tensor,
+    w_packed: torch.Tensor,
+    scale_act: torch.Tensor,
+    scale_wt: torch.Tensor,
+) -> torch.Tensor:
+    """Ternary x INT8 GEMM. Returns [M, N] bf16."""
+    M, K = x_int8.shape
+    N = w_packed.shape[0]
+    out = torch.empty((M, N), dtype=torch.bfloat16)
+    ops.bitnet_gemm(out, x_int8, w_packed, scale_act, scale_wt, None)
+    return out
+def bitnet_gemv_fused(
+    x: torch.Tensor,
+    w_packed: torch.Tensor,
+    scale_wt: torch.Tensor,
+) -> torch.Tensor:
+    """Fused quantize + ternary GEMV for M < 16. Returns [M, N] bf16."""
+    M, K = x.shape
+    N = w_packed.shape[0]
+    out = torch.empty((M, N), dtype=torch.bfloat16)
+    ops.bitnet_gemv_fused(out, x, w_packed, scale_wt)
+    return out
+def bitnet_linear(
+    x: torch.Tensor,
+    w_packed: torch.Tensor,
+    scale_wt: torch.Tensor,
+) -> torch.Tensor:
+    """One-shot BitLinear forward. Auto-dispatches fused (M<16) vs split path.
+    Args:
+        x: [..., K] bf16 or f32.
+        w_packed: [N, K//4] uint8 packed ternary weights.
+        scale_wt: [N] bf16 per-output weight scale.
+    Returns:
+        [..., N] bf16.
+    """
+    assert x.dtype in _OK_DTYPES, f"x must be bf16 or f32, got {x.dtype}"
+    orig = x.shape
+    K = orig[-1]
+    x_flat = x.reshape(-1, K).contiguous()
+    M = x_flat.shape[0]
+    N = w_packed.shape[0]
+    if M < 16:
+        y = bitnet_gemv_fused(x_flat, w_packed, scale_wt)
+    else:
+        x_int8, scale_act = quantize_activation(x_flat)
+        y = bitnet_gemm(x_int8, w_packed, scale_act, scale_wt)
+    return y.view(*orig[:-1], N)
+class BitLinear(nn.Module):
+    """Inference-only BitLinear: ternary weights, per-token INT8 activations."""
+    def __init__(self, in_features: int, out_features: int, bias: bool = False):
+        super().__init__()
+        assert in_features % 32 == 0, f"in_features={in_features} must be % 32"
+        self.in_features = in_features
+        self.out_features = out_features
+        self.register_buffer(
+            "w_packed", torch.zeros(out_features, in_features // 4, dtype=torch.uint8)
+        )
+        self.register_buffer(
+            "scale_wt", torch.ones(out_features, dtype=torch.bfloat16)
+        )
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_features, dtype=torch.bfloat16))
+        else:
+            self.bias = None
+    @classmethod
+    def from_dense(cls, lin: nn.Linear) -> "BitLinear":
+        """Quantize a dense Linear via absmean ternarization (sanity-test only;
+        a real BitNet model is trained with QAT, not post-hoc quantized)."""
+        out_features, in_features = lin.weight.shape
+        bl = cls(in_features, out_features, bias=lin.bias is not None)
+        with torch.no_grad():
+            W = lin.weight.detach().float()
+            gamma = W.abs().mean(dim=-1, keepdim=True).clamp(min=1e-8)
+            Wq = (W / gamma).round().clamp(-1, 1).to(torch.int8)
+            bl.w_packed.copy_(pack_weights(Wq.cpu()))
+            bl.scale_wt.copy_(gamma.squeeze(-1).to(torch.bfloat16))
+            if bl.bias is not None:
+                bl.bias.copy_(lin.bias.detach().to(torch.bfloat16))
+        return bl
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = bitnet_linear(x, self.w_packed, self.scale_wt)
+        if self.bias is not None:
+            y = y + self.bias
+        return y
+class BitLinearKernel(nn.Module):
+    """kernelize layer for the transformers BitNet linear modules.
+    kernelize binds this forward onto the host BitLinear / AutoBitLinear module,
+    so self is that module and the forward reads its attributes (weight,
+    weight_scale, rms_norm, bias). The ternary weight is converted to the
+    kernel's 2-bit layout once and cached; each call runs per-token INT8
+    quantization and the VNNI/AVX2 ternary GEMM.
+    """
+    can_torch_compile: bool = False
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        dt = input.dtype
+        rms = getattr(self, "rms_norm", None)
+        if rms is not None:
+            input = rms(input)
+        if not hasattr(self, "_bitnet_cpu_packed"):
+            w = self.weight
+            if w.dtype == torch.uint8:
+                # transformers stores ternary as int8 viewed through uint8 (255 -> -1).
+                tern = w.contiguous().view(torch.int8)
+                if tern.shape[0] != self.out_features:        # bit-packed (out//4, in)
+                    from transformers.integrations.bitnet import unpack_weights
+                    tern = unpack_weights(w, dtype=torch.bfloat16)
+                tern = tern.round().clamp(-1, 1).to(torch.int8)
+            else:
+                g = w.detach().float().abs().mean().clamp(min=1e-5)
+                tern = (w.detach().float() / g).round().clamp(-1, 1).to(torch.int8)
+            self._bitnet_cpu_n = tern.shape[0]
+            self._bitnet_cpu_packed = pack_weights(tern.contiguous().cpu())
+            ws = getattr(self, "weight_scale", None)
+            if ws is not None:
+                s = ws.detach().float().reshape(-1)
+                s = s.expand(self._bitnet_cpu_n) if s.numel() == 1 else s
+            else:
+                s = w.detach().float().abs().mean().clamp(min=1e-5).reshape(1).expand(self._bitnet_cpu_n)
+            self._bitnet_cpu_sw = s.to(torch.bfloat16).contiguous().cpu()
+        shp = input.shape
+        x = input.reshape(-1, shp[-1])
+        if x.dtype not in _OK_DTYPES:
+            x = x.to(torch.float32)
+        x = x.contiguous()
+        out = bitnet_linear(x, self._bitnet_cpu_packed, self._bitnet_cpu_sw)
+        if getattr(self, "bias", None) is not None:
+            out = out + self.bias
+        return out.reshape(*shp[:-1], self._bitnet_cpu_n).to(dt)
+__all__ = [
+    "pack_weights",
+    "quantize_activation",
+    "bitnet_gemm",
+    "bitnet_gemv_fused",
+    "bitnet_linear",
+    "BitLinear",
+    "BitLinearKernel",
+]

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

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

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

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

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

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

build/torch213-cxx11-cpu-x86_64-linux/bitnet_cpu/__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/metadata.json ADDED Viewed

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

+{
+  "name": "bitnet-cpu",
+  "id": "_bitnet_cpu_cpu_bd669ef",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cpu"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "ipCv02rtmMOwCw75aWDraJag/on6FtbuuKkYLxWUgiA=",
+      "_bitnet_cpu_cpu_bd669ef.abi3.so": "wPL2vHO6y2LmRfrVLk9h+2gOT3GKisvOonNRxFltBp4=",
+      "_ops.py": "CpD3fRDWoqe63dtPPstdFtlIPKs13yPdbOGpg4aH2lk=",
+      "bitnet_cpu/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "2c516fc6678570c013319adeb64605fc8e8d5394",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "bd669ef68e49c5873a7810015ce6a9503fed5394",
+      "dirty": false
+    }
+  }
+}