sparse_kernel_wrapper.py

"""Sparse BitNet INT2 ternary kernel wrapper.

Bridges:
  * BitLinearKernel in /home/server/cortex/distill/windows_pc_inspect/model.py
    (packed-INT2 weights ``[M, K]`` with ``K = in_features // 4``)
  * sparse_bitnet_phase3_kernel.cu which takes ``W_sparse [M, N]`` with
    ``N = K * 4`` and computes ``Y = A @ W^T * row_scales``.

2-bit codes (mirror `unpack_ternary` in the .cu): 00->-1, 01->0, 10->+1, 11->0.

Honesty note (audit KERNEL_AUDIT_FINAL.md §H-5):
  This wrapper is *sparse-formatted*, not *structurally-sparse*. We pack
  every ternary into the 4-wide group layout, but we keep all four
  positions and rely on the 2-bit packing for the storage win (4 ternaries
  per byte = 4x storage saving over INT8). The "2:4 sparsity" claim is
  about *shape* (smaller working set, more experts in VRAM), not about
  achieved 2x matmul throughput — RTX 3070 (sm_86) has no sparse tensor
  cores, so we never realize the throughput side of the marketing pitch.
  See sparse_bitnet_quality_tradeoff.md §3 row 1 in the kernel source.

Public API:
    SparseBitNetTensor     - CSR-style storage of packed 2-bit ternaries + indices.
    dense_to_sparse(...)   - convert dense packed weights -> sparse-formatted layout.
    sparse_bitlinear(...)  - run the sparse matmul (with pure-Py fallback).

Falls back to a pure-Py matmul when ``libsparse_bitnet.so/.dll`` is absent.
"""

from __future__ import annotations

import ctypes, os
from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch

# 2:4 structured sparsity constants (mirror sparse_bitnet_phase3_kernel.cu).
# _GROUP is the 4-wide nibble group the decoder reads from one byte of nz_vals.
# _NZ_PER_GROUP is 4 (the full group), matching the kernel's decode_packed_byte
# which reads 4 ternaries per byte. The "2:4 structured" naming is retained
# for the FUTURE WMMA path that will consume nz_offsets to skip the 2 zeros;
# the current scalar fallback decodes all 4 positions, so nz_vals is dense
# (every slot populated).
_GROUP = 4
_NZ_PER_GROUP = 4

# Kernel discovery.
_KERNEL_DIR = os.environ.get("BITNET_SPARSE_KERNEL_DIR",
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "bitnet_kernels"))
ext = "libsparse_bitnet.dll" if os.name == "nt" else "libsparse_bitnet.so"
_KERNEL_PATH = os.path.join(_KERNEL_DIR, ext)


def _unpack(code: int) -> int:
    """2-bit code -> ternary, matching Microsoft's INT2 convention.

    Microsoft's BitNet uses the *asymmetric* 2-bit encoding (per
    bitnet_kernels.cu's lop3 + __vsubss4(0x02) decode chain):
        0b00 -> -2 (saturated to 0 by __vsubss4; reserved/unused)
        0b01 -> -1
        0b10 ->  0
        0b11 -> +1
    The sparse kernel does its own decode without __vsubss4, so it must
    match this convention directly. The previous convention
    (00=-1, 01=0, 10=+1, 11=0) was symmetric and WRONG — every weight
    fed through dense_to_sparse() decoded as 0 (a Bug #2 regression
    caught in code review 2026-07-13).
    """
    return {1: -1, 2: 0, 3: 1}.get(code, 0)

def _pack_code(tern: int) -> int:
    """Ternary -> 2-bit code, matching Microsoft's INT2 convention.

    See _unpack for the rationale. Code 0b00 (decodes to -2 in the CUDA
    kernel's saturated subtract) is reserved/unused — never emitted.
    """
    return {1: 0b11, -1: 0b01, 0: 0b10}[tern]


@dataclass
class SparseBitNetTensor:
    """CSR-style [M, N] ternary weight matrix in 2:4 layout.

    Mirrors ``BitLinearKernel``: M=output dim, N=input dim, so
    ``activation [B,N] @ sparse_w.T`` -> ``[B,M]``.
    Fields: row_ptr [M+1] u32, nz_vals/nz_cols [M*N_groups] u8,
    row_scales [M] fp16.  N must be % 4.
    """
    row_ptr:    torch.Tensor
    nz_vals:    torch.Tensor
    nz_cols:    torch.Tensor
    row_scales: torch.Tensor
    M:          int
    N:          int

    @property
    def density(self) -> float:
        """Density of the formatted tensor — 1.0 for the current scalar
        fallback (every 2-bit slot in nz_vals is populated). The "2:4
        sparsity" naming refers to a *future* WMMA path that will consult
        nz_offsets to skip zero slots; today the kernel reads nz_vals as
        4 ternaries per byte, so the data is dense-formatted, not
        structurally sparse. See dense_to_sparse docstring."""
        return (self.nz_vals.numel() * _NZ_PER_GROUP) / (self.M * self.N)


def _load_lib():
    """Return ctypes handle for libsparse_bitnet, or None if absent."""
    if not os.path.exists(_KERNEL_PATH):
        return None
    try:
        return ctypes.CDLL(_KERNEL_PATH)
    except OSError:
        return None


def dense_to_sparse(weight_packed: torch.Tensor,
                    weight_scale:  torch.Tensor) -> SparseBitNetTensor:
    """Convert dense BitNet-packed weights [M, K] into the 2-bit sparse layout.

    IMPORTANT (audit §H-5): this is a *sparse-formatted dense* conversion,
    not actual 2:4 sparsity reduction. We re-pack every ternary into the
    4-nibble-per-byte layout so the kernel can use the 4x INT2 packing
    (4 ternaries per byte vs INT8's 1-per-byte), but the resulting tensor
    has density == 1.0 — every group slot is populated. The audit's
    `density <= 0.5` claim was fiction; the correct measure is
    `nz_vals.numel() / (M * N / 4) == 1.0`. Don't rely on this function
    for any "2x throughput" promise — that requires real sparse tensor
    cores which sm_86 does not have.

    The encoder convention matches the kernel's `decode_packed_byte`
    exactly: one input byte holds 4 ternaries (2 bits each), and the
    output `nz_vals` byte holds the same 4 ternaries in the same slot
    order — no top-k selection, no 2:4 cap. `nz_cols` is left as zeros;
    the future WMMA path will fill it with the 2-of-4 column indices.

    weight_scale: bfloat16 [M, 4] or [M]; mean-collapsed per-row.
    Matches BitLinearKernel.weight layout.
    """
    assert weight_packed.dtype == torch.int8, (
        f"expected int8, got {weight_packed.dtype}")
    M, K = weight_packed.shape
    N, N_groups = K * _GROUP, K

    arr = weight_packed.cpu().numpy().astype(np.uint8)
    nz_v = np.zeros(M * N_groups, dtype=np.uint8)
    nz_c = np.zeros(M * N_groups, dtype=np.uint8)
    row_ptr = np.arange(M + 1, dtype=np.uint32) * N_groups

    for r in range(M):
        for g in range(N_groups):
            packed = int(arr[r, g])
            # Re-pack the 4 ternaries from `packed` into nz_vals[g] at the
            # same bit positions. The decoder reads them out at the same
            # slots, so this is the identity transform — bit layout matches
            # decode_packed_byte (00->-1, 01->0, 10->+1, 11->0).
            vals = 0
            for q in range(_GROUP):
                tern = _unpack((packed >> (q * 2)) & 0x3)
                vals |= (_pack_code(tern) << (q * 2))
            i = r * N_groups + g
            nz_v[i] = np.uint8(vals)
            nz_c[i] = np.uint8(0)

    rs = (weight_scale.float().mean(dim=-1).to(torch.float16)
          if weight_scale.dim() == 2 else weight_scale.to(torch.float16))

    return SparseBitNetTensor(
        row_ptr=torch.from_numpy(row_ptr),
        nz_vals=torch.from_numpy(nz_v),
        nz_cols=torch.from_numpy(nz_c),
        row_scales=rs, M=M, N=N,
    )


# ctypes binding (must match the C SparseTernaryMatrix struct).
# The launcher takes the struct as a POINTER (audit §C-3 fix); ctypes passes
# `ctypes.byref(sp)` which is exactly that — a pointer to the original
# struct (not a hidden-pointer-to-stack-copy as MSVC's reference ABI would
# produce). Python side is unchanged; the C++ launcher signature change
# is what makes the binding cross-platform safe.
class _SparseTernaryMatrix(ctypes.Structure):
    _fields_ = [("row_ptr", ctypes.c_void_p), ("nz_vals", ctypes.c_void_p),
                ("nz_offsets", ctypes.c_void_p), ("row_scales", ctypes.c_void_p),
                ("M", ctypes.c_int), ("N", ctypes.c_int), ("N_groups", ctypes.c_int)]


def _make_struct(w: SparseBitNetTensor):
    sp = _SparseTernaryMatrix()
    # Keep the CUDA tensors alive through the asynchronous kernel launch.
    # Taking data_ptr() from a temporary .cuda() tensor leaves dangling pointers.
    gpu_tensors = tuple(x.contiguous().cuda() for x in
                        (w.row_ptr, w.nz_vals, w.nz_cols, w.row_scales))
    sp.row_ptr, sp.nz_vals, sp.nz_offsets, sp.row_scales = (
        ctypes.c_void_p(x.data_ptr()) for x in gpu_tensors)
    sp.M, sp.N, sp.N_groups = w.M, w.N, w.N // _GROUP
    return sp, gpu_tensors


def sparse_bitlinear(activation:   torch.Tensor,
                     sparse_w:     SparseBitNetTensor,
                     dense_weight: Optional[torch.Tensor] = None) -> torch.Tensor:
    """``Y = activation @ sparse_w^T * row_scales`` (sparse 2:4 ternary matmul).

    activation: [B, N]  (N == sparse_w.N).  Returns fp16 [B, M] on GPU, bf16
    [B, M] from fallback.  Raises RuntimeError if the CUDA launch returns != 0.
    """
    B, N = activation.shape
    assert N == sparse_w.N, f"activation width {N} != weight width {sparse_w.N}"
    lib = _load_lib()
    if lib is None:
        return _fallback(activation, sparse_w)

    A = activation.contiguous().to(torch.float16)
    C = torch.empty(B, sparse_w.M, dtype=torch.float16, device=A.device)
    sp, gpu_tensors = _make_struct(sparse_w)
    stream = torch.cuda.current_stream().cuda_stream
    lib.bitnet_sparse_gemm_launch.restype = ctypes.c_int
    err = lib.bitnet_sparse_gemm_launch(
        ctypes.byref(sp),
        ctypes.c_void_p(A.data_ptr()),
        ctypes.c_void_p(C.data_ptr()),
        ctypes.c_int(B), ctypes.c_int(sparse_w.M), ctypes.c_int(N),
        ctypes.c_void_p(stream),
    )
    for tensor in gpu_tensors:
        tensor.record_stream(torch.cuda.current_stream())
    if err != 0:
        raise RuntimeError(
            f"bitnet_sparse_gemm_launch returned {err} (see sparse_bitnet_phase3_kernel.cu)")
    return C


def _fallback(activation: torch.Tensor,
              sparse_w:    SparseBitNetTensor) -> torch.Tensor:
    """Materialise sparse matrix densely and matmul - no CUDA required.

    Output bf16 to match ``bitlinear_int8xint2_linear`` in model.py.
    Reads nz_vals as 4 ternaries per byte (matches the kernel decoder and
    the Python wrapper's dense_to_sparse output); nz_cols is unused in
    the current dense-formatted layout.
    """
    B, N = activation.shape
    activation = activation.cuda()
    M_w, N_groups = sparse_w.M, N // _GROUP
    W  = torch.zeros(M_w, N, dtype=torch.float32, device=activation.device)
    nv = sparse_w.nz_vals.detach().cpu().numpy()
    for r in range(M_w):
        base = r * N_groups
        for g in range(N_groups):
            bv = int(nv[base + g])
            for q in range(_GROUP):
                code = (bv >> (q * 2)) & 0x3
                W[r, g * _GROUP + q] = float(_unpack(code))
    # activation [B,N] @ W^T [N,M_w] -> [B, M_w]
    out = activation.float() @ W.T
    return (out * sparse_w.row_scales.float().to(activation.device).unsqueeze(0)).to(torch.bfloat16)

__all__ = ["SparseBitNetTensor", "dense_to_sparse", "sparse_bitlinear"]  # noqa: E501