File size: 845 Bytes
862cf3c | 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 | # The CUDA trellis kernels live in the qtip_kernels extension. If it can't be
# imported (ABI mismatch after a torch upgrade, not built, no nvcc to rebuild),
# every shape must use the torch decode fallback — force has_kernel False.
try:
import torch # ensure libc10 is loaded before the qtip_kernels .so
import qtip_kernels as _qtip_kernels # noqa: F401
_QTIP_KERNELS_OK = True
except Exception:
_QTIP_KERNELS_OK = False
def has_kernel(decode_mode, L, K, V, tlut_bits, td_x, td_y):
if not _QTIP_KERNELS_OK:
return False
if decode_mode != 'quantlut_sym':
return False
if L != 16:
return False
if V != 2:
return False
if K < 2 or K > 4:
return False
if tlut_bits != 9:
return False
if td_x != 16 or td_y != 16:
return False
return True
|