codebook_init.py

"""Flattened from lib/codebook/__init__.py.

Provides `kdict` (imported by bitshift.py) and registers the quip_lib torch.library
ops for each compiled kernel shape. The registration is pure torch.library metadata;
the actual CUDA impl only runs when qtip_kernels is importable. In manifest
('train-fixW') inference the fused-kernel path is never taken, so this whole module is
effectively just supplying an (empty) `kdict` symbol — but we keep the op registration
so the eval/kernel path still works verbatim if the compiled .so happens to be present.
"""
import torch  # must precede qtip_kernels so libc10 is loaded (the .so links it)

try:
    import qtip_kernels
except Exception as _qtip_err:  # ABI mismatch (torch upgrade) / not built / no CUDA
    qtip_kernels = None
    import warnings
    warnings.warn(
        f"qtip_kernels unavailable ({_qtip_err}); CUDA trellis kernels disabled - "
        f"falling back to torch decode. has_kernel() is forced False, so inference "
        f"runs the (pure-torch) manifest path.")

kernels = [
    (5120, 1, 5120, 4),
    (5120, 1, 5120, 3),
    (5120, 1, 5120, 2),
    (1024, 1, 5120, 4),
    (1024, 1, 5120, 3),
    (1024, 1, 5120, 2),
    (8192, 1, 5120, 4),
    (8192, 1, 5120, 3),
    (8192, 1, 5120, 2),
    (5120, 1, 8192, 4),
    (5120, 1, 8192, 3),
    (5120, 1, 8192, 2),
    (1024, 1, 3072, 4),
    (8192, 1, 3072, 4),
    (3072, 1, 8192, 4),
    (3072, 1, 3072, 4),
    (53248, 1, 16384, 2),
    (53248, 1, 16384, 3),
    (53248, 1, 16384, 4),
    (16384, 1, 53248, 2),
    (16384, 1, 53248, 3),
    (16384, 1, 53248, 4),
    (1024, 1, 16384, 2),
    (1024, 1, 16384, 3),
    (1024, 1, 16384, 4),
    (16384, 1, 16384, 2),
    (16384, 1, 16384, 3),
    (16384, 1, 16384, 4),
    (4096, 1, 14336, 2),
    (4096, 1, 14336, 3),
    (4096, 1, 14336, 4),
    (14336, 1, 4096, 2),
    (14336, 1, 4096, 3),
    (14336, 1, 4096, 4),
    (1024, 1, 4096, 2),
    (1024, 1, 4096, 3),
    (1024, 1, 4096, 4),
    (4096, 1, 4096, 2),
    (4096, 1, 11008, 2),
    (4096, 1, 12288, 2),
    (11008, 1, 4096, 2),
    (12288, 1, 4096, 2),
    (22016, 1, 4096, 2),
    (8192, 1, 8192, 2),
    (10240, 1, 8192, 2),
    (10240, 1, 8192, 3),
    (10240, 1, 8192, 4),
    (57344, 1, 8192, 2),
    (57344, 1, 8192, 3),
    (57344, 1, 8192, 4),
    (8192, 1, 1024, 2),
    (8192, 1, 28672, 2),
    (28672, 1, 8192, 2),
    (1024, 1, 8192, 2),
    (4096, 1, 4096, 3),
    (4096, 1, 11008, 3),
    (4096, 1, 12288, 3),
    (11008, 1, 4096, 3),
    (12288, 1, 4096, 3),
    (22016, 1, 4096, 3),
    (8192, 1, 8192, 3),
    (8192, 1, 1024, 3),
    (8192, 1, 28672, 3),
    (28672, 1, 8192, 3),
    (1024, 1, 8192, 3),
    (4096, 1, 4096, 4),
    (4096, 1, 11008, 4),
    (4096, 1, 12288, 4),
    (11008, 1, 4096, 4),
    (12288, 1, 4096, 4),
    (22016, 1, 4096, 4),
    (8192, 1, 8192, 4),
    (8192, 1, 1024, 4),
    (8192, 1, 28672, 4),
    (28672, 1, 8192, 4),
    (1024, 1, 8192, 4),
    (27648, 1, 5120, 2),
    (27648, 1, 5120, 3),
    (27648, 1, 5120, 4),
    (5120, 1, 27648, 2),
    (5120, 1, 27648, 3),
    (5120, 1, 27648, 4),
]

kdict = {}

for m, n, k, bitrate in kernels:
    name = f"decompress_matvec_qtip_{m}_{n}_{k}_{bitrate}"
    kernel_name = f"qtip_kernels.decompress_matvec_16_9_{bitrate}_1_{m}_{n}_{k}"
    try:
        torch.library.define(
            f"quip_lib::{name}",
            "(Tensor compressed, Tensor x, Tensor codebook) -> Tensor")
    except RuntimeError:
        continue   # op already registered in this process (loading another checkpoint in one kernel)
    exec(f"""\
@torch.library.register_fake("quip_lib::{name}")
def {name}_abstract(
        compressed: torch.Tensor,
        x: torch.Tensor,
        codebook: torch.Tensor) -> torch.Tensor:
    return torch.zeros(1, {m}, dtype=torch.float32, device=x.device)

@torch.library.impl("quip_lib::{name}", "cuda")
def {name}_cuda(
        compressed: torch.Tensor,
        x: torch.Tensor,
        codebook: torch.Tensor) -> torch.Tensor:
    out = torch.zeros(({m}, 1), dtype=torch.float32, device=x.device)
    {kernel_name}(out, compressed.reshape(-1).view(torch.int32), x.to(torch.float16).T, codebook.reshape(-1))
    return out.T
    """)