File size: 4,177 Bytes
e04ed2d c453fcd e04ed2d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """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
""")
|