import ast import os import re import struct import sys import numpy as np from dataclasses import dataclass import torch from torch.utils import _pytree as pytree from optimum.quanto import QModuleMixin, register_qmodule from optimum.quanto.tensor.qtensor import QTensor from optimum.quanto.tensor.qtype import qtype as _quanto_qtype, qtypes as _quanto_qtypes from collections import OrderedDict from mmgp.offload import QEmbedding as BaseQEmbedding HANDLER_NAME = "gguf" try: import gguf except Exception: gguf = None _GGUF_QTYPE_NAME = "gguf" if _GGUF_QTYPE_NAME not in _quanto_qtypes: _quanto_qtypes[_GGUF_QTYPE_NAME] = _quanto_qtype( _GGUF_QTYPE_NAME, is_floating_point=False, bits=6, dtype=torch.uint8, qmin=-32.0, qmax=31.0, ) _GGUF_QTYPE = _quanto_qtypes[_GGUF_QTYPE_NAME] _GGUF_DEFAULT_DTYPE = None _GGUF_LABEL_CACHE = {} _GGUF_METADATA_CACHE = {} _GGUF_INDEX_CACHE = {} _GGUF_RUNTIME_LOGGED = set() _GGUF_CUDA_KERNELS_ENV = "WGP_GGUF_LLAMACPP_CUDA" _GGUF_CUDA_KERNELS_ENABLED_CACHE = None _GGUF_CUDA_MODULE = None _GGUF_CUDA_LOAD_ERROR = None _GGUF_NATIVE_BYTE_ORDER = "<" if sys.byteorder == "little" else ">" _GGUF_ORIG_SHAPE_PREFIX = "comfy.gguf.orig_shape." _GGUF_FAST_METADATA_SCALARS = { None if gguf is None else gguf.GGUFValueType.UINT8: np.uint8, None if gguf is None else gguf.GGUFValueType.INT8: np.int8, None if gguf is None else gguf.GGUFValueType.UINT16: np.uint16, None if gguf is None else gguf.GGUFValueType.INT16: np.int16, None if gguf is None else gguf.GGUFValueType.UINT32: np.uint32, None if gguf is None else gguf.GGUFValueType.INT32: np.int32, None if gguf is None else gguf.GGUFValueType.FLOAT32: np.float32, None if gguf is None else gguf.GGUFValueType.BOOL: np.bool_, None if gguf is None else gguf.GGUFValueType.UINT64: np.uint64, None if gguf is None else gguf.GGUFValueType.INT64: np.int64, None if gguf is None else gguf.GGUFValueType.FLOAT64: np.float64, } _GGUF_FAST_METADATA_STRUCT_FORMATS = { None if gguf is None else gguf.GGUFValueType.UINT8: "B", None if gguf is None else gguf.GGUFValueType.INT8: "b", None if gguf is None else gguf.GGUFValueType.UINT16: "H", None if gguf is None else gguf.GGUFValueType.INT16: "h", None if gguf is None else gguf.GGUFValueType.UINT32: "I", None if gguf is None else gguf.GGUFValueType.INT32: "i", None if gguf is None else gguf.GGUFValueType.FLOAT32: "f", None if gguf is None else gguf.GGUFValueType.BOOL: "?", None if gguf is None else gguf.GGUFValueType.UINT64: "Q", None if gguf is None else gguf.GGUFValueType.INT64: "q", None if gguf is None else gguf.GGUFValueType.FLOAT64: "d", } _GGUF_TYPED_TENSOR_DTYPES = { None if gguf is None else gguf.GGMLQuantizationType.F16: np.float16, None if gguf is None else gguf.GGMLQuantizationType.F32: np.float32, None if gguf is None else gguf.GGMLQuantizationType.F64: np.float64, None if gguf is None else gguf.GGMLQuantizationType.I8: np.int8, None if gguf is None else gguf.GGMLQuantizationType.I16: np.int16, None if gguf is None else gguf.GGMLQuantizationType.I32: np.int32, None if gguf is None else gguf.GGMLQuantizationType.I64: np.int64, } @dataclass(frozen=True) class _GGUFTensorInfo: name: str tensor_type: object raw_shape: tuple[int, ...] data_offset: int n_elements: int n_bytes: int @dataclass(frozen=True) class _GGUFParsedIndex: byte_order: str data_alignment: int data_offset: int tensor_infos: tuple[_GGUFTensorInfo, ...] config: object | None orig_shapes: tuple[tuple[str, tuple[int, ...]], ...] def _normalize_gguf_path(file_path): try: return os.path.normcase(os.path.abspath(file_path)) except Exception: return str(file_path).lower() def normalize(file_path): return _normalize_gguf_path(file_path) def _set_default_dtype_from_loader(dtype): global _GGUF_DEFAULT_DTYPE if dtype is None: return _GGUF_DEFAULT_DTYPE = dtype def _resolve_default_dtype(dtype, fallback=None): if dtype is None: return _GGUF_DEFAULT_DTYPE or fallback return dtype def _gguf_log_once(key, message): if key in _GGUF_RUNTIME_LOGGED: return _GGUF_RUNTIME_LOGGED.add(key) print(message) def _gguf_cuda_env_enabled(): raw = str(os.environ.get(_GGUF_CUDA_KERNELS_ENV, "1")).strip().lower() return raw in ("1", "true", "yes", "y", "on") def _validate_gguf_cuda_module(module): import numpy as np if gguf is None: raise RuntimeError("gguf package is unavailable") if not torch.cuda.is_available(): raise RuntimeError("CUDA is not available") if not module.may_support_linear_qtype_name("Q4_0"): raise RuntimeError("Q4_0 linear is not supported by the installed GGUF CUDA package") weight = np.linspace(-1.0, 1.0, 32, dtype=np.float32).reshape(1, 32) quantized_weight = gguf.quants.quantize(weight, gguf.GGMLQuantizationType.Q4_0) raw_weight = torch.from_numpy(quantized_weight).contiguous().cuda() input_tensor = torch.linspace(-0.5, 0.5, 64, dtype=torch.float32, device="cuda").reshape(2, 32) fast_out = module.linear(raw_weight, "Q4_0", [1, 32], input_tensor, None, torch.float32) dense_weight = torch.from_numpy(gguf.quants.dequantize(quantized_weight, gguf.GGMLQuantizationType.Q4_0)).to(device=input_tensor.device, dtype=torch.float32) ref_out = torch.nn.functional.linear(input_tensor, dense_weight, None) torch.cuda.synchronize() if fast_out.shape != ref_out.shape: raise RuntimeError(f"Q4 self-test shape mismatch: {tuple(fast_out.shape)} vs {tuple(ref_out.shape)}") if not torch.allclose(fast_out, ref_out, atol=2e-2, rtol=2e-2): max_diff = float((fast_out - ref_out).abs().max().item()) raise RuntimeError(f"Q4 self-test output mismatch (max diff {max_diff:.6f})") def _probe_gguf_cuda_runtime(force=False): global _GGUF_CUDA_KERNELS_ENABLED_CACHE, _GGUF_CUDA_MODULE, _GGUF_CUDA_LOAD_ERROR if not force and not _gguf_cuda_env_enabled(): _GGUF_CUDA_KERNELS_ENABLED_CACHE = False _GGUF_CUDA_MODULE = None _GGUF_CUDA_LOAD_ERROR = None return False try: import llamacpp_gguf_cuda as gguf_cuda_module _validate_gguf_cuda_module(gguf_cuda_module) except Exception as exc: _GGUF_CUDA_KERNELS_ENABLED_CACHE = False _GGUF_CUDA_MODULE = None _GGUF_CUDA_LOAD_ERROR = exc _gguf_log_once("gguf_cuda_probe_failed", f"[GGUF][llama.cpp CUDA] kernels unavailable, using fallback") return False _GGUF_CUDA_KERNELS_ENABLED_CACHE = True _GGUF_CUDA_MODULE = gguf_cuda_module _GGUF_CUDA_LOAD_ERROR = None _gguf_log_once("gguf_cuda_probe_ok", "[GGUF][llama.cpp CUDA] kernels available.") return True def _gguf_cuda_kernels_enabled(): global _GGUF_CUDA_KERNELS_ENABLED_CACHE if _GGUF_CUDA_KERNELS_ENABLED_CACHE is not None: return _GGUF_CUDA_KERNELS_ENABLED_CACHE return _probe_gguf_cuda_runtime() def _gguf_cuda_module(): if not _gguf_cuda_kernels_enabled(): return None return _GGUF_CUDA_MODULE def set_gguf_cuda_kernels_enabled(enabled=None): global _GGUF_CUDA_KERNELS_ENABLED_CACHE, _GGUF_CUDA_MODULE, _GGUF_CUDA_LOAD_ERROR if enabled is False: _GGUF_CUDA_KERNELS_ENABLED_CACHE = False _GGUF_CUDA_MODULE = None _GGUF_CUDA_LOAD_ERROR = None return _GGUF_CUDA_KERNELS_ENABLED_CACHE = None _GGUF_CUDA_MODULE = None _GGUF_CUDA_LOAD_ERROR = None if enabled is None: return _probe_gguf_cuda_runtime() return _probe_gguf_cuda_runtime(force=True) _probe_gguf_cuda_runtime() def _gguf_read_array(data, offset, dtype, byte_order): dtype = np.dtype(dtype).newbyteorder(byte_order) value = np.frombuffer(data, dtype=dtype, count=1, offset=offset) return value, offset + int(value.nbytes) def _gguf_read_string_fast(data, offset, byte_order): str_len_arr, offset = _gguf_read_array(data, offset, np.uint64, byte_order) str_len = int(str_len_arr[0]) raw = bytes(data[offset : offset + str_len]) return raw.decode("utf-8"), offset + str_len def _gguf_decode_value_fast(data, offset, raw_type, byte_order): value_type = gguf.GGUFValueType(int(raw_type)) scalar_dtype = _GGUF_FAST_METADATA_SCALARS.get(value_type) if scalar_dtype is not None: value_arr, offset = _gguf_read_array(data, offset, scalar_dtype, byte_order) value = value_arr[0] return (bool(value) if value_type == gguf.GGUFValueType.BOOL else value.item()), offset if value_type == gguf.GGUFValueType.STRING: return _gguf_read_string_fast(data, offset, byte_order) if value_type == gguf.GGUFValueType.ARRAY: item_type_arr, offset = _gguf_read_array(data, offset, np.uint32, byte_order) item_count_arr, offset = _gguf_read_array(data, offset, np.uint64, byte_order) item_type = int(item_type_arr[0]) item_count = int(item_count_arr[0]) items = [] for _ in range(item_count): item, offset = _gguf_decode_value_fast(data, offset, item_type, byte_order) items.append(item) return items, offset raise ValueError(f"Unsupported GGUF metadata field type: {value_type}") def _gguf_stream_read_exact(reader, size): data = reader.read(int(size)) if len(data) != int(size): raise EOFError("Unexpected end of GGUF metadata header.") return data def _gguf_stream_unpack(reader, fmt): return struct.unpack(fmt, _gguf_stream_read_exact(reader, struct.calcsize(fmt))) def _gguf_read_string_stream(reader, byte_order): str_len = int(_gguf_stream_unpack(reader, byte_order + "Q")[0]) return _gguf_stream_read_exact(reader, str_len).decode("utf-8") def _gguf_decode_value_stream(reader, raw_type, byte_order): value_type = gguf.GGUFValueType(int(raw_type)) scalar_fmt = _GGUF_FAST_METADATA_STRUCT_FORMATS.get(value_type) if scalar_fmt is not None: return _gguf_stream_unpack(reader, byte_order + scalar_fmt)[0] if value_type == gguf.GGUFValueType.STRING: return _gguf_read_string_stream(reader, byte_order) if value_type == gguf.GGUFValueType.ARRAY: item_type = int(_gguf_stream_unpack(reader, byte_order + "I")[0]) item_count = int(_gguf_stream_unpack(reader, byte_order + "Q")[0]) return [_gguf_decode_value_stream(reader, item_type, byte_order) for _ in range(item_count)] raise ValueError(f"Unsupported GGUF metadata field type: {value_type}") def _gguf_stream_skip(reader, size): if int(size) <= 0: return reader.seek(int(size), os.SEEK_CUR) def _gguf_skip_value_stream(reader, raw_type, byte_order): value_type = gguf.GGUFValueType(int(raw_type)) scalar_fmt = _GGUF_FAST_METADATA_STRUCT_FORMATS.get(value_type) if scalar_fmt is not None: _gguf_stream_skip(reader, struct.calcsize(byte_order + scalar_fmt)) return if value_type == gguf.GGUFValueType.STRING: str_len = int(_gguf_stream_unpack(reader, byte_order + "Q")[0]) _gguf_stream_skip(reader, str_len) return if value_type == gguf.GGUFValueType.ARRAY: item_type = int(_gguf_stream_unpack(reader, byte_order + "I")[0]) item_count = int(_gguf_stream_unpack(reader, byte_order + "Q")[0]) item_value_type = gguf.GGUFValueType(item_type) item_scalar_fmt = _GGUF_FAST_METADATA_STRUCT_FORMATS.get(item_value_type) if item_scalar_fmt is not None: _gguf_stream_skip(reader, struct.calcsize(byte_order + item_scalar_fmt) * item_count) return if item_value_type == gguf.GGUFValueType.STRING: for _ in range(item_count): item_len = int(_gguf_stream_unpack(reader, byte_order + "Q")[0]) _gguf_stream_skip(reader, item_len) return for _ in range(item_count): _gguf_skip_value_stream(reader, item_type, byte_order) return raise ValueError(f"Unsupported GGUF metadata field type: {value_type}") def _gguf_quant_byte_shape(shape, tensor_type): if len(shape) == 0: return tuple(shape) block_size, type_size = gguf.GGML_QUANT_SIZES[tensor_type] last_dim = int(shape[-1]) if block_size <= 0 or last_dim % block_size != 0: raise ValueError(f"Invalid GGUF tensor shape {tuple(shape)} for tensor type {getattr(tensor_type, 'name', tensor_type)}") return tuple(int(dim) for dim in shape[:-1]) + (last_dim // block_size * type_size,) def _gguf_cache_identity(file_path): try: stat = os.stat(file_path) except OSError: return None return (int(stat.st_size), int(getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1e9)))) def _gguf_parse_index(file_path): if gguf is None: raise RuntimeError("GGUF support requires the 'gguf' package.") with open(file_path, "rb") as reader: magic = int(_gguf_stream_unpack(reader, " None: try: from mmgp import quant_router except Exception: return quant_router.register_handler(__name__) quant_router.register_file_extension("gguf", sys.modules[__name__]) ensure_gguf_handler_registered() def get_file_metadata(file_path): if gguf is None: raise RuntimeError("GGUF support requires the 'gguf' package.") cache_key = _normalize_gguf_path(file_path) cached = _GGUF_METADATA_CACHE.get(cache_key) if cached is not None: state_dict, metadata = cached return OrderedDict(state_dict), dict(metadata) metadata = _get_lightweight_gguf_metadata(file_path) result = (OrderedDict(), metadata) _GGUF_METADATA_CACHE[cache_key] = (OrderedDict(result[0]), dict(result[1])) return OrderedDict(result[0]), dict(result[1]) def _filter_state_dict_basic(state_dict, base_model_prefix, keep_prefix=False): new_state_dict = {} start = -1 if keep_prefix: for k, v in state_dict.items(): if k.startswith(base_model_prefix): new_state_dict[k] = v else: for k, v in state_dict.items(): if k.startswith(base_model_prefix): new_start = len(base_model_prefix) else: pos = k.find("." + base_model_prefix) if pos < 0: continue new_start = pos + len(base_model_prefix) + 1 if start != -1 and start != new_start: new_state_dict = state_dict break start = new_start new_state_dict[k[start:]] = v return new_state_dict def _gguf_resolve_prefix(tensor_names, prefixes): for prefix in prefixes: if any(name.startswith(prefix) for name in tensor_names): return prefix return None def load_gguf_state_dict( file_path, filters=None, keep_prefixes=False, writable_tensors=True, verboseLevel=1, default_dtype=None, pin_to_memory=False, ): if gguf is None: raise RuntimeError("GGUF support requires the 'gguf' package.") if pin_to_memory: raise Exception("Pinning to memory while loading GGUF files is not supported") import warnings def _cast_plain_tensor(torch_tensor, tensor_type): if tensor_type == gguf.GGMLQuantizationType.F16: if torch_tensor.dtype in (torch.uint8, torch.uint16): torch_tensor = torch_tensor.view(torch.float16) elif torch_tensor.dtype != torch.float16: torch_tensor = torch_tensor.to(torch.float16) elif tensor_type == gguf.GGMLQuantizationType.BF16: if torch_tensor.dtype in (torch.uint8, torch.uint16): torch_tensor = torch_tensor.view(torch.bfloat16) elif torch_tensor.dtype != torch.bfloat16: torch_tensor = torch_tensor.to(torch.bfloat16) elif tensor_type == gguf.GGMLQuantizationType.F32: if torch_tensor.dtype in (torch.uint8, torch.uint16, torch.uint32): torch_tensor = torch_tensor.view(torch.float32) elif torch_tensor.dtype != torch.float32: torch_tensor = torch_tensor.to(torch.float32) return torch_tensor def _tensor_type_from_dtype(dtype): if dtype == torch.float16: return gguf.GGMLQuantizationType.F16 if dtype == torch.bfloat16: return gguf.GGMLQuantizationType.BF16 if dtype == torch.float32: return gguf.GGMLQuantizationType.F32 return None parsed = _gguf_get_index(file_path) orig_shapes = dict(parsed.orig_shapes) mmap_data = np.memmap(file_path, mode="r") if verboseLevel >= 2: try: from mmgp import safetensors2 safetensors2.verboseLevel = verboseLevel tracker = safetensors2.MmapTracker(file_path) tracker.register(mmap_data, 0, 0, int(mmap-_data.nbytes)) except Exception: tracker = None tensor_names = [tensor.name for tensor in parsed.tensor_infos] prefix = _gguf_resolve_prefix(tensor_names, ("model.diffusion_model.", "diffusion_model.")) state_dict = {} qtype_counts = {} for tensor in parsed.tensor_infos: name = tensor.name if prefix and name.startswith(prefix): name = name[len(prefix):] with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="The given NumPy array is not writable") torch_tensor = torch.from_numpy(_gguf_open_tensor_numpy(mmap_data, parsed, tensor)) shape = orig_shapes.get(tensor.name, None) if shape is None: shape = torch.Size(tuple(int(v) for v in reversed(tensor.raw_shape))) else: shape = torch.Size(tuple(int(v) for v in shape)) if tensor.tensor_type in ( gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16, gguf.GGMLQuantizationType.BF16, ): torch_tensor = _cast_plain_tensor(torch_tensor, tensor.tensor_type) torch_tensor = torch_tensor.view(*shape) wrapped = GGUFSourceTensor.wrap(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape) if name.endswith(".bias"): wrapped._gguf_bias_orig_dtype = wrapped.dtype wrapped._gguf_bias_orig_tensor_type = tensor.tensor_type state_dict[name] = wrapped type_name = getattr(tensor.tensor_type, "name", str(tensor.tensor_type)) qtype_counts[type_name] = qtype_counts.get(type_name, 0) + 1 if verboseLevel >= 2 and qtype_counts: print("GGUF qtypes: " + ", ".join(f"{k} ({v})" for k, v in qtype_counts.items())) if filters is not None: if not isinstance(filters, list): filters = [filters] new_sd = {} for one_filter in filters: new_sd.update(_filter_state_dict_basic(state_dict, one_filter, keep_prefixes)) state_dict = new_sd if state_dict: for name, bias in list(state_dict.items()): if not name.endswith(".bias") or not torch.is_tensor(bias): continue weight_name = name[:-5] + ".weight" weight = state_dict.get(weight_name) target_dtype = None weight_type = getattr(weight, "tensor_type", None) if torch.is_tensor(weight) else None if torch.is_tensor(weight) and weight_type is None: target_dtype = weight.dtype if weight.dtype.is_floating_point else default_dtype elif weight_type in ( gguf.GGMLQuantizationType.F16, gguf.GGMLQuantizationType.BF16, gguf.GGMLQuantizationType.F32, ): target_dtype = weight.dtype else: target_dtype = default_dtype if target_dtype is None or bias.dtype == target_dtype: continue casted = bias.to(target_dtype) if isinstance(casted, GGUFSourceTensor): casted._gguf_bias_orig_dtype = getattr(bias, "_gguf_bias_orig_dtype", bias.dtype) casted._gguf_bias_orig_tensor_type = getattr( bias, "_gguf_bias_orig_tensor_type", getattr(bias, "tensor_type", None) ) new_tensor_type = _tensor_type_from_dtype(target_dtype) if new_tensor_type is not None: casted.tensor_type = new_tensor_type casted.tensor_shape = getattr(bias, "tensor_shape", casted.shape) state_dict[name] = casted return state_dict, None, None def load_state_dict(*args, **kwargs): return load_gguf_state_dict(*args, **kwargs) class GGUFSourceTensor(torch.Tensor): @staticmethod def wrap(tensor, *, tensor_type, tensor_shape): wrapped = tensor.as_subclass(GGUFSourceTensor) wrapped.tensor_type = tensor_type wrapped.tensor_shape = tensor_shape return wrapped def to(self, *args, **kwargs): new = super().to(*args, **kwargs) new.tensor_type = getattr(self, "tensor_type", None) new.tensor_shape = getattr(self, "tensor_shape", new.shape) return new def clone(self, *args, **kwargs): cloned = super().clone(*args, **kwargs).as_subclass(GGUFSourceTensor) cloned.tensor_type = getattr(self, "tensor_type", None) cloned.tensor_shape = getattr(self, "tensor_shape", cloned.shape) return cloned def detach(self, *args, **kwargs): detached = super().detach(*args, **kwargs).as_subclass(GGUFSourceTensor) detached.tensor_type = getattr(self, "tensor_type", None) detached.tensor_shape = getattr(self, "tensor_shape", detached.shape) return detached def get_quantized_subtensors(self): return [("data", self)] def set_quantized_subtensors(self, sub_tensors): if isinstance(sub_tensors, dict): data = sub_tensors.get("data") else: data = dict(sub_tensors).get("data") if data is None or data is self: return torch.utils.swap_tensors(self, data) def materialize_source_tensor(tensor: torch.Tensor) -> torch.Tensor: if not isinstance(tensor, GGUFSourceTensor): return tensor with torch._C.DisableTorchFunctionSubclass(): plain = torch.empty_strided( tuple(tensor.shape), tuple(tensor.stride()), dtype=tensor.dtype, device=tensor.device, ) plain.copy_(tensor) return plain def materialize_module_source_tensors(module: torch.nn.Module) -> int: converted = 0 for submodule in module.modules(): for name, param in list(submodule._parameters.items()): if param is None or not isinstance(param, GGUFSourceTensor): continue submodule._parameters[name] = torch.nn.Parameter( materialize_source_tensor(param), requires_grad=param.requires_grad, ) converted += 1 for name, buf in list(submodule._buffers.items()): if buf is None or not isinstance(buf, GGUFSourceTensor): continue submodule._buffers[name] = materialize_source_tensor(buf) converted += 1 return converted def _split_gguf_tensor(src, *, dim, split_sizes, context): if not torch.is_tensor(src): return None tensor_type = getattr(src, "tensor_type", None) if tensor_type is None: return None tensor_shape = getattr(src, "tensor_shape", None) or src.shape total = sum(split_sizes) if dim >= len(tensor_shape) or tensor_shape[dim] != total: return None chunks = torch.split(src, split_sizes, dim=dim) out = [] for chunk, size in zip(chunks, split_sizes): new_shape = list(tensor_shape) new_shape[dim] = size wrapped = GGUFSourceTensor.wrap(chunk, tensor_type=tensor_type, tensor_shape=tuple(new_shape)) if hasattr(src, "_gguf_bias_orig_dtype"): wrapped._gguf_bias_orig_dtype = getattr(src, "_gguf_bias_orig_dtype") if hasattr(src, "_gguf_bias_orig_tensor_type"): wrapped._gguf_bias_orig_tensor_type = getattr(src, "_gguf_bias_orig_tensor_type") out.append(wrapped) return out def split_fused_weights(state_dict, fused_split_map, quantization_map=None, allowed_bases=None, default_dtype=None, verboseLevel=1): from mmgp import offload return offload.sd_split_linear( state_dict, fused_split_map, split_fields={"weight": 0, "bias": 0}, split_handlers={"weight": _split_gguf_tensor, "bias": _split_gguf_tensor}, verboseLevel=verboseLevel, allowed_bases=allowed_bases, return_split_bases=True, ) def _is_gguf_qtype(qtype_obj): if gguf is None: return False if qtype_obj is None: return False return qtype_obj not in ( gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16, gguf.GGMLQuantizationType.BF16, ) def _gguf_qtype_name(qtype_obj): if qtype_obj is None: return None return getattr(qtype_obj, "name", None) or str(qtype_obj) def _try_llamacpp_cuda_linear(weight_tensor, input_tensor, bias, target_dtype): if not _gguf_cuda_kernels_enabled(): return None if not torch.is_tensor(input_tensor) or input_tensor.device.type != "cuda": return None raw = getattr(weight_tensor, "_data", None) if not torch.is_tensor(raw) or raw.device.type != "cuda" or not raw.is_contiguous(): return None qtype_name = _gguf_qtype_name(getattr(weight_tensor, "_tensor_type", None)) try: gguf_llamacpp_cuda = _gguf_cuda_module() if gguf_llamacpp_cuda is None or not gguf_llamacpp_cuda.may_support_linear_qtype_name(qtype_name): return None fast_out = gguf_llamacpp_cuda.linear(raw, qtype_name, tuple(getattr(weight_tensor, "_tensor_shape", weight_tensor.shape)), input_tensor, bias, target_dtype) if fast_out is not None: _gguf_log_once(f"llamacpp_cuda_linear_active_{qtype_name}", f"[GGUF][llama.cpp CUDA] linear fast path active for {qtype_name}.") return fast_out except Exception as exc: _gguf_log_once(f"llamacpp_cuda_linear_{qtype_name}", f"[GGUF][llama.cpp CUDA] linear GGUF CUDA kernels failed for {qtype_name}, using fallback: {exc}") return None def _try_llamacpp_cuda_embedding(weight_tensor, index_tensor, target_dtype): if not _gguf_cuda_kernels_enabled(): return None if not torch.is_tensor(index_tensor) or index_tensor.device.type != "cuda": return None raw = getattr(weight_tensor, "_data", None) if not torch.is_tensor(raw) or raw.device.type != "cuda" or not raw.is_contiguous(): return None qtype_name = _gguf_qtype_name(getattr(weight_tensor, "_tensor_type", None)) try: gguf_llamacpp_cuda = _gguf_cuda_module() if gguf_llamacpp_cuda is None or not gguf_llamacpp_cuda.may_support_embedding_qtype_name(qtype_name): return None fast_out = gguf_llamacpp_cuda.embedding(raw, qtype_name, tuple(getattr(weight_tensor, "_tensor_shape", weight_tensor.shape)), index_tensor, target_dtype) if fast_out is not None: _gguf_log_once(f"llamacpp_cuda_embedding_active_{qtype_name}", f"[GGUF][llama.cpp CUDA] embedding fast path active for {qtype_name}.") return fast_out except Exception as exc: _gguf_log_once(f"llamacpp_cuda_embedding_{qtype_name}", f"[GGUF][llama.cpp CUDA] embedding GGUF CUDA kernels failed for {qtype_name}, using fallback: {exc}") return None def _may_try_llamacpp_cuda_linear(weight_tensor, input_tensor): if not _gguf_cuda_kernels_enabled(): return False if not torch.is_tensor(input_tensor) or input_tensor.device.type != "cuda": return False raw = getattr(weight_tensor, "_data", None) if not torch.is_tensor(raw) or raw.device.type != "cuda" or not raw.is_contiguous(): return False qtype_name = _gguf_qtype_name(getattr(weight_tensor, "_tensor_type", None)) try: gguf_llamacpp_cuda = _gguf_cuda_module() return gguf_llamacpp_cuda is not None and gguf_llamacpp_cuda.may_support_linear_qtype_name(qtype_name) except Exception: return False def _may_try_llamacpp_cuda_embedding(weight_tensor, index_tensor): if not _gguf_cuda_kernels_enabled(): return False if not torch.is_tensor(index_tensor) or index_tensor.device.type != "cuda": return False raw = getattr(weight_tensor, "_data", None) if not torch.is_tensor(raw) or raw.device.type != "cuda" or not raw.is_contiguous(): return False qtype_name = _gguf_qtype_name(getattr(weight_tensor, "_tensor_type", None)) try: gguf_llamacpp_cuda = _gguf_cuda_module() return gguf_llamacpp_cuda is not None and gguf_llamacpp_cuda.may_support_embedding_qtype_name(qtype_name) except Exception: return False def _guess_variant_from_filename(filename): base = os.path.basename(str(filename)) match = re.search(r"(?i)(?:^|[_-])(Q\d+_K|Q\d+_\d|Q\d+|IQ\d+_\w+)(?:$|[_.-])", base) if match: return match.group(1).upper() return None def detect_gguf_quantization_variant(file_path, verboseLevel=1): if gguf is None: return None try: parsed = _gguf_get_index(file_path) except Exception: return None counts = {} for tensor in parsed.tensor_infos: qtype = getattr(tensor, "tensor_type", None) if qtype in ( gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16, gguf.GGMLQuantizationType.BF16, ): continue name = _gguf_qtype_name(qtype) if not name: continue counts[name] = counts.get(name, 0) + 1 if not counts: return None return max(counts, key=counts.get) def detect_quantization_kind_for_file(file_path, verboseLevel=1): if not file_path or str(file_path).lower().endswith(".gguf") is False: return None if gguf is None: return None return "gguf" def detect_quantization_label_from_filename(filename, verboseLevel=1): if not filename or str(filename).lower().endswith(".gguf") is False: return "" key = _normalize_gguf_path(filename) cached = _GGUF_LABEL_CACHE.get(key) if cached: return cached variant = _guess_variant_from_filename(filename) if not variant and os.path.isfile(filename): variant = detect_gguf_quantization_variant(filename, verboseLevel=verboseLevel) if variant: label = f"GGUF-{variant}" else: label = "GGUF" _GGUF_LABEL_CACHE[key] = label return label def _gguf_qfallback(callable, *args, **kwargs): args, kwargs = pytree.tree_map_only(GGUFWeightTensor, lambda x: x.dequantize(), (args, kwargs or {})) return callable(*args, **kwargs) def _reshape_scale(scale, weight): if scale.ndim == 0 or scale.numel() == 1: return scale if scale.ndim == 1 and scale.shape[0] == weight.shape[0]: return scale.view(weight.shape[0], *([1] * (weight.ndim - 1))) return scale def _gguf_dequantize_tensor(raw, qtype_obj, oshape, dtype=None): if gguf is None: raise RuntimeError("gguf package is required to dequantize GGUF weights.") if qtype_obj in ( gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16, gguf.GGMLQuantizationType.BF16, ): out = raw.view(*oshape) return out.to(dtype) if dtype is not None else out if qtype_obj not in _DEQUANTIZE_FUNCTIONS: out = gguf.quants.dequantize(raw.cpu().numpy(), qtype_obj) out = torch.from_numpy(out) return out.to(dtype) if dtype is not None else out block_size, type_size = gguf.GGML_QUANT_SIZES[qtype_obj] dequantize_blocks = _DEQUANTIZE_FUNCTIONS[qtype_obj] rows = raw.reshape((-1, raw.shape[-1])).view(torch.uint8) n_blocks = rows.numel() // type_size blocks = rows.reshape((n_blocks, type_size)) blocks = dequantize_blocks(blocks, block_size, type_size, dtype) return blocks.reshape(oshape) def _maybe_cast_bias(bias, target_dtype): if bias is None or not torch.is_tensor(bias) or target_dtype is None: return bias if bias.dtype == target_dtype: return bias if isinstance(bias, GGUFSourceTensor): tensor_type = getattr(bias, "tensor_type", None) tensor_shape = getattr(bias, "tensor_shape", bias.shape) if _is_gguf_qtype(tensor_type): return _gguf_dequantize_tensor(bias, tensor_type, tensor_shape, dtype=target_dtype) return bias.to(target_dtype) def _to_uint32(x): x = x.view(torch.uint8).to(torch.int32) return (x[:, 0] | x[:, 1] << 8 | x[:, 2] << 16 | x[:, 3] << 24).unsqueeze(1) def _to_uint16(x): x = x.view(torch.uint8).to(torch.int32) return (x[:, 0] | x[:, 1] << 8).unsqueeze(1) def _const_like(ref, values, dtype): device = ref.device if torch.is_tensor(ref) else None count = len(values) if count == 0: return torch.empty((0,), device=device, dtype=dtype) if count == 1: return torch.full((1,), values[0], device=device, dtype=dtype) step = values[1] - values[0] if all(values[idx] - values[idx - 1] == step for idx in range(1, count)): end = values[0] + step * count return torch.arange(values[0], end, step, device=device, dtype=dtype) raise ValueError("Unsupported constant pattern for GGUF dequantization.") def _split_block_dims(blocks, *args): n_max = blocks.shape[1] dims = list(args) + [n_max - sum(args)] return torch.split(blocks, dims, dim=1) def _dequantize_blocks_Q8_0(blocks, block_size, type_size, dtype=None): d, x = _split_block_dims(blocks, 2) d = d.view(torch.float16).to(dtype) x = x.view(torch.int8) return d * x def _dequantize_blocks_Q5_1(blocks, block_size, type_size, dtype=None): n_blocks = blocks.shape[0] d, m, qh, qs = _split_block_dims(blocks, 2, 2, 4) d = d.view(torch.float16).to(dtype) m = m.view(torch.float16).to(dtype) qh = _to_uint32(qh) qh = qh.reshape((n_blocks, 1)) >> torch.arange(32, device=d.device, dtype=torch.int32).reshape(1, 32) ql = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> _const_like(d, [0, 4], torch.uint8).reshape(1, 1, 2, 1) qh = (qh & 1).to(torch.uint8) ql = (ql & 0x0F).reshape((n_blocks, -1)) qs = ql | (qh << 4) return d * qs + m def _dequantize_blocks_Q5_0(blocks, block_size, type_size, dtype=None): n_blocks = blocks.shape[0] d, qh, qs = _split_block_dims(blocks, 2, 4) d = d.view(torch.float16).to(dtype) qh = _to_uint32(qh) qh = qh.reshape(n_blocks, 1) >> torch.arange(32, device=d.device, dtype=torch.int32).reshape(1, 32) ql = qs.reshape(n_blocks, -1, 1, block_size // 2) >> _const_like(d, [0, 4], torch.uint8).reshape(1, 1, 2, 1) qh = (qh & 1).to(torch.uint8) ql = (ql & 0x0F).reshape(n_blocks, -1) qs = (ql | (qh << 4)).to(torch.int8) - 16 return d * qs def _dequantize_blocks_Q4_1(blocks, block_size, type_size, dtype=None): n_blocks = blocks.shape[0] d, m, qs = _split_block_dims(blocks, 2, 2) d = d.view(torch.float16).to(dtype) m = m.view(torch.float16).to(dtype) qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> _const_like(d, [0, 4], torch.uint8).reshape(1, 1, 2, 1) qs = (qs & 0x0F).reshape((n_blocks, -1)) return d * qs + m def _dequantize_blocks_Q4_0(blocks, block_size, type_size, dtype=None): n_blocks = blocks.shape[0] d, qs = _split_block_dims(blocks, 2) d = d.view(torch.float16).to(dtype) qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> _const_like(d, [0, 4], torch.uint8).reshape(1, 1, 2, 1) qs = (qs & 0x0F).reshape((n_blocks, -1)) qs = qs.to(torch.int8) - 8 return d * qs QK_K = 256 K_SCALE_SIZE = 12 def _get_scale_min(scales): n_blocks = scales.shape[0] scales = scales.view(torch.uint8) scales = scales.reshape((n_blocks, 3, 4)) d, m, m_d = torch.split(scales, scales.shape[-2] // 3, dim=-2) sc = torch.cat([d & 0x3F, (m_d & 0x0F) | ((d >> 2) & 0x30)], dim=-1) mn = torch.cat([m & 0x3F, (m_d >> 4) | ((m >> 2) & 0x30)], dim=-1) return sc.reshape((n_blocks, 8)), mn.reshape((n_blocks, 8)) def _dequantize_blocks_Q6_K(blocks, block_size, type_size, dtype=None): n_blocks = blocks.shape[0] ql, qh, scales, d = _split_block_dims(blocks, QK_K // 2, QK_K // 4, QK_K // 16) scales = scales.view(torch.int8).to(dtype) d = d.view(torch.float16).to(dtype) d = (d * scales).reshape((n_blocks, QK_K // 16, 1)) ql = ql.reshape((n_blocks, -1, 1, 64)) >> _const_like(d, [0, 4], torch.uint8).reshape((1, 1, 2, 1)) ql = (ql & 0x0F).reshape((n_blocks, -1, 32)) qh = qh.reshape((n_blocks, -1, 1, 32)) >> _const_like(d, [0, 2, 4, 6], torch.uint8).reshape((1, 1, 4, 1)) qh = (qh & 0x03).reshape((n_blocks, -1, 32)) q = (ql | (qh << 4)).to(torch.int8) - 32 q = q.reshape((n_blocks, QK_K // 16, -1)) return (d * q).reshape((n_blocks, QK_K)) def _dequantize_blocks_Q5_K(blocks, block_size, type_size, dtype=None): n_blocks = blocks.shape[0] d, dmin, scales, qh, qs = _split_block_dims(blocks, 2, 2, K_SCALE_SIZE, QK_K // 8) d = d.view(torch.float16).to(dtype) dmin = dmin.view(torch.float16).to(dtype) sc, m = _get_scale_min(scales) d = (d * sc).reshape((n_blocks, -1, 1)) dm = (dmin * m).reshape((n_blocks, -1, 1)) ql = qs.reshape((n_blocks, -1, 1, 32)) >> _const_like(d, [0, 4], torch.uint8).reshape((1, 1, 2, 1)) qh = qh.reshape((n_blocks, -1, 1, 32)) >> _const_like(d, list(range(8)), torch.uint8).reshape((1, 1, 8, 1)) ql = (ql & 0x0F).reshape((n_blocks, -1, 32)) qh = (qh & 0x01).reshape((n_blocks, -1, 32)) q = ql | (qh << 4) return (d * q - dm).reshape((n_blocks, QK_K)) def _dequantize_blocks_Q4_K(blocks, block_size, type_size, dtype=None): n_blocks = blocks.shape[0] d, dmin, scales, qs = _split_block_dims(blocks, 2, 2, K_SCALE_SIZE) d = d.view(torch.float16).to(dtype) dmin = dmin.view(torch.float16).to(dtype) sc, m = _get_scale_min(scales) d = (d * sc).reshape((n_blocks, -1, 1)) dm = (dmin * m).reshape((n_blocks, -1, 1)) qs = qs.reshape((n_blocks, -1, 1, 32)) >> _const_like(d, [0, 4], torch.uint8).reshape((1, 1, 2, 1)) qs = (qs & 0x0F).reshape((n_blocks, -1, 32)) return (d * qs - dm).reshape((n_blocks, QK_K)) def _dequantize_blocks_Q3_K(blocks, block_size, type_size, dtype=None): n_blocks = blocks.shape[0] hmask, qs, scales, d = _split_block_dims(blocks, QK_K // 8, QK_K // 4, 12) d = d.view(torch.float16).to(dtype) lscales, hscales = scales[:, :8], scales[:, 8:] lscales = lscales.reshape((n_blocks, 1, 8)) >> _const_like(d, [0, 4], torch.uint8).reshape((1, 2, 1)) lscales = lscales.reshape((n_blocks, 16)) hscales = hscales.reshape((n_blocks, 1, 4)) >> _const_like(d, [0, 2, 4, 6], torch.uint8).reshape((1, 4, 1)) hscales = hscales.reshape((n_blocks, 16)) scales = (lscales & 0x0F) | ((hscales & 0x03) << 4) scales = (scales.to(torch.int8) - 32) dl = (d * scales).reshape((n_blocks, 16, 1)) ql = qs.reshape((n_blocks, -1, 1, 32)) >> _const_like(d, [0, 2, 4, 6], torch.uint8).reshape((1, 1, 4, 1)) qh = hmask.reshape(n_blocks, -1, 1, 32) >> _const_like(d, list(range(8)), torch.uint8).reshape((1, 1, 8, 1)) ql = ql.reshape((n_blocks, 16, QK_K // 16)) & 3 qh = (qh.reshape((n_blocks, 16, QK_K // 16)) & 1) ^ 1 q = (ql.to(torch.int8) - (qh << 2).to(torch.int8)) return (dl * q).reshape((n_blocks, QK_K)) def _dequantize_blocks_Q2_K(blocks, block_size, type_size, dtype=None): n_blocks = blocks.shape[0] scales, qs, d, dmin = _split_block_dims(blocks, QK_K // 16, QK_K // 4, 2) d = d.view(torch.float16).to(dtype) dmin = dmin.view(torch.float16).to(dtype) dl = (d * (scales & 0xF)).reshape((n_blocks, QK_K // 16, 1)) ml = (dmin * (scales >> 4)).reshape((n_blocks, QK_K // 16, 1)) shift = _const_like(d, [0, 2, 4, 6], torch.uint8).reshape((1, 1, 4, 1)) qs = (qs.reshape((n_blocks, -1, 1, 32)) >> shift) & 3 qs = qs.reshape((n_blocks, QK_K // 16, 16)) qs = dl * qs - ml return qs.reshape((n_blocks, -1)) if gguf is not None: _DEQUANTIZE_FUNCTIONS = { gguf.GGMLQuantizationType.Q8_0: _dequantize_blocks_Q8_0, gguf.GGMLQuantizationType.Q5_1: _dequantize_blocks_Q5_1, gguf.GGMLQuantizationType.Q5_0: _dequantize_blocks_Q5_0, gguf.GGMLQuantizationType.Q4_1: _dequantize_blocks_Q4_1, gguf.GGMLQuantizationType.Q4_0: _dequantize_blocks_Q4_0, gguf.GGMLQuantizationType.Q6_K: _dequantize_blocks_Q6_K, gguf.GGMLQuantizationType.Q5_K: _dequantize_blocks_Q5_K, gguf.GGMLQuantizationType.Q4_K: _dequantize_blocks_Q4_K, gguf.GGMLQuantizationType.Q3_K: _dequantize_blocks_Q3_K, gguf.GGMLQuantizationType.Q2_K: _dequantize_blocks_Q2_K, } else: _DEQUANTIZE_FUNCTIONS = {} class GGUFWeightTensor(QTensor): @staticmethod def create(raw_tensor, size, stride, dtype, device=None, requires_grad=False, tensor_type=None, tensor_shape=None): if tensor_type is None: tensor_type = getattr(raw_tensor, "tensor_type", None) if tensor_shape is None: tensor_shape = getattr(raw_tensor, "tensor_shape", None) or size device = raw_tensor.device if device is None else device if raw_tensor.device != device: raw_tensor = raw_tensor.to(device) return GGUFWeightTensor( qtype=_GGUF_QTYPE, axis=0, size=size, stride=stride, raw=raw_tensor, tensor_type=tensor_type, tensor_shape=tensor_shape, dtype=dtype, requires_grad=requires_grad, ) @staticmethod def __new__(cls, qtype, axis, size, stride, raw, tensor_type, tensor_shape, dtype, requires_grad=False): return torch.Tensor._make_wrapper_subclass( cls, size, strides=stride, dtype=dtype, device=raw.device, requires_grad=requires_grad, ) def __init__(self, qtype, axis, size, stride, raw, tensor_type, tensor_shape, dtype, requires_grad=False): super().__init__(qtype, axis) self._data = raw self._tensor_type = tensor_type self._tensor_shape = torch.Size(tensor_shape) self._gguf_default_dtype = dtype def __repr__(self): cls_name = self.__class__.__name__ try: shape = tuple(self.shape) except Exception: shape = "" try: dtype = str(self.dtype).replace("torch.", "") except Exception: dtype = "" try: device = str(self.device) except Exception: device = "" qtype = getattr(self, "_qtype", None) qtype_name = getattr(qtype, "name", None) or str(qtype) if qtype is not None else "" tensor_type = _gguf_qtype_name(getattr(self, "_tensor_type", None)) or "" return ( f"{cls_name}(shape={shape}, dtype={dtype}, device={device}, " f"qtype={qtype_name}, tensor_type={tensor_type})" ) __str__ = __repr__ def dequantize(self, dtype=None, device=None): if dtype is None: dtype = self.dtype if device is None: device = self.device raw = self._data if self._data.device == device else self._data.to(device) return _gguf_dequantize_tensor(raw, self._tensor_type, self._tensor_shape, dtype=dtype) def linear(self, input, bias=None): if torch.is_tensor(input): target_dtype = _resolve_default_dtype(self._gguf_default_dtype, fallback=input.dtype) target_device = input.device else: target_dtype = _resolve_default_dtype(self._gguf_default_dtype, fallback=self.dtype) target_device = self.device if _may_try_llamacpp_cuda_linear(self, input): fast_out = _try_llamacpp_cuda_linear(self, input, bias, target_dtype) if fast_out is not None: return fast_out weight = self.dequantize(dtype=target_dtype, device=target_device) if torch.is_tensor(input) and input.dtype != weight.dtype: input = input.to(weight.dtype) bias = _maybe_cast_bias(bias, weight.dtype) return torch.nn.functional.linear(input, weight, bias) def get_quantized_subtensors(self): return [("data", self._data)] def set_quantized_subtensors(self, sub_tensors): if isinstance(sub_tensors, dict): sub_map = sub_tensors else: sub_map = {name: tensor for name, tensor in sub_tensors} data = sub_map.get("data", None) if data is not None: old_data = self._data if torch.is_tensor(old_data): try: torch.utils.swap_tensors(old_data, data) self._data = old_data except Exception: self._data = data else: self._data = data if hasattr(self, "_ggml_raw_cpu"): self._ggml_raw_cpu = None def __tensor_flatten__(self): inner_tensors = ["_data"] meta = { "qtype": self._qtype.name, "axis": str(self._axis), "size": str(list(self.size())), "stride": str(list(self.stride())), "dtype": str(self.dtype), "tensor_type": _gguf_qtype_name(self._tensor_type) or "", "tensor_shape": str(list(self._tensor_shape)), } return inner_tensors, meta @staticmethod def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride): qtype = _quanto_qtypes[meta["qtype"]] axis = ast.literal_eval(meta["axis"]) size = ast.literal_eval(meta["size"]) stride = ast.literal_eval(meta["stride"]) dtype_str = meta.get("dtype", "torch.float16") if dtype_str.startswith("torch."): dtype_name = dtype_str.split(".", 1)[1] dtype = getattr(torch, dtype_name, torch.float16) else: dtype = getattr(torch, dtype_str, torch.float16) tensor_shape = ast.literal_eval(meta.get("tensor_shape", str(list(size)))) tensor_type = None if gguf is not None and meta.get("tensor_type"): tensor_type = getattr(gguf.GGMLQuantizationType, meta["tensor_type"], None) return GGUFWeightTensor( qtype=qtype, axis=axis, size=size, stride=stride, raw=inner_tensors["_data"], tensor_type=tensor_type, tensor_shape=tensor_shape, dtype=dtype, ) @classmethod def __torch_function__(cls, func, types, args=(), kwargs=None): kwargs = kwargs or {} if func is torch.nn.functional.linear: input = args[0] if len(args) > 0 else kwargs.get("input", None) weight = args[1] if len(args) > 1 else kwargs.get("weight", None) bias = args[2] if len(args) > 2 else kwargs.get("bias", None) if isinstance(weight, GGUFWeightTensor): return weight.linear(input, bias=bias) with torch._C.DisableTorchFunctionSubclass(): return func(*args, **kwargs) @classmethod def __torch_dispatch__(cls, op, types, args, kwargs=None): op = op.overloadpacket kwargs = kwargs or {} if op is torch.ops.aten.linear: input = args[0] weight = args[1] bias = args[2] if len(args) > 2 else None if isinstance(weight, GGUFWeightTensor): return weight.linear(input, bias=bias) if op is torch.ops.aten.detach: t = args[0] return GGUFWeightTensor.create( raw_tensor=op(t._data), size=t.size(), stride=t.stride(), dtype=t.dtype, device=t.device, requires_grad=t.requires_grad, tensor_type=getattr(t, "_tensor_type", None), tensor_shape=getattr(t, "_tensor_shape", None), ) if op in (torch.ops.aten._to_copy, torch.ops.aten.to): t = args[0] dtype = kwargs.pop("dtype", t.dtype) if kwargs else t.dtype device = kwargs.pop("device", t.device) if kwargs else t.device if dtype != t.dtype: return t.dequantize(dtype=dtype, device=device) out_data = op(t._data, device=device, **(kwargs or {})) return GGUFWeightTensor.create( raw_tensor=out_data, size=t.size(), stride=t.stride(), dtype=t.dtype, device=device, requires_grad=t.requires_grad, tensor_type=getattr(t, "_tensor_type", None), tensor_shape=getattr(t, "_tensor_shape", None), ) return _gguf_qfallback(op, *args, **(kwargs or {})) class QLinearGGUF(QModuleMixin, torch.nn.Linear): def __init__( self, in_features, out_features, bias=True, device=None, dtype=None, weights=None, activations=None, optimizer=None, quantize_input=True, ): super().__init__( in_features, out_features, bias=bias, device=device, dtype=dtype, weights=weights, activations=activations, optimizer=optimizer, quantize_input=quantize_input, ) self._gguf_default_dtype = dtype @classmethod def qcreate(cls, module, weights, activations=None, optimizer=None, device=None): if torch.is_tensor(module.weight) and module.weight.dtype.is_floating_point: weight_dtype = module.weight.dtype elif torch.is_tensor(getattr(module, "bias", None)) and module.bias.dtype.is_floating_point: weight_dtype = module.bias.dtype else: weight_dtype = torch.float16 return cls( module.in_features, module.out_features, module.bias is not None, device=device, dtype=weight_dtype, weights=weights, activations=activations, optimizer=optimizer, quantize_input=True, ) def set_default_dtype(self, dtype): self._gguf_default_dtype = dtype @property def qweight(self): if self.weight_qtype == _GGUF_QTYPE: return self.weight return super().qweight def forward(self, input: torch.Tensor) -> torch.Tensor: qweight = self.qweight if isinstance(qweight, GGUFWeightTensor): return qweight.linear(input, bias=self.bias) return torch.nn.functional.linear(input, qweight, bias=self.bias) def _load_from_state_dict( self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs ): if self.weight_qtype != _GGUF_QTYPE: return super()._load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs ) weight_key = prefix + "weight" bias_key = prefix + "bias" input_scale_key = prefix + "input_scale" output_scale_key = prefix + "output_scale" weight_raw = state_dict.pop(weight_key, None) bias = state_dict.pop(bias_key, None) input_scale = state_dict.pop(input_scale_key, None) output_scale = state_dict.pop(output_scale_key, None) if weight_raw is None: missing_keys.append(weight_key) target_dtype = _resolve_default_dtype(self._gguf_default_dtype, fallback=self.weight.dtype) if weight_raw is not None: gguf_weight = GGUFWeightTensor.create( raw_tensor=weight_raw, size=self.weight.size(), stride=self.weight.stride(), dtype=target_dtype, device=weight_raw.device, requires_grad=False, ) self.weight = torch.nn.Parameter(gguf_weight, requires_grad=False) if bias is not None: self.bias = torch.nn.Parameter(bias, requires_grad=False) if torch.is_tensor(weight_raw): scale_device = weight_raw.device elif torch.is_tensor(self.weight): scale_device = self.weight.device elif torch.is_tensor(bias): scale_device = bias.device else: scale_device = torch.device("cpu") if input_scale is not None: self.input_scale = input_scale.to(scale_device) else: if not hasattr(self, "input_scale") or self.input_scale.is_meta: scale_dtype = self.input_scale.dtype if hasattr(self, "input_scale") else torch.float32 self.input_scale = torch.ones((), dtype=scale_dtype, device=scale_device) if output_scale is not None: self.output_scale = output_scale.to(scale_device) else: if not hasattr(self, "output_scale") or self.output_scale.is_meta: scale_dtype = self.output_scale.dtype if hasattr(self, "output_scale") else torch.float32 self.output_scale = torch.ones((), dtype=scale_dtype, device=scale_device) return @register_qmodule(torch.nn.Conv1d) class QConv1dGGUF(QModuleMixin, torch.nn.Conv1d): def __init__( self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode="zeros", device=None, dtype=None, weights=None, activations=None, optimizer=None, quantize_input=True, ): super().__init__( in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias, padding_mode=padding_mode, device=device, dtype=dtype, weights=weights, activations=activations, optimizer=optimizer, quantize_input=quantize_input, ) self._gguf_default_dtype = dtype @classmethod def qcreate(cls, module, weights, activations=None, optimizer=None, device=None): if torch.is_tensor(module.weight) and module.weight.dtype.is_floating_point: weight_dtype = module.weight.dtype elif torch.is_tensor(getattr(module, "bias", None)) and module.bias.dtype.is_floating_point: weight_dtype = module.bias.dtype else: weight_dtype = torch.float16 return cls( module.in_channels, module.out_channels, module.kernel_size, stride=module.stride, padding=module.padding, dilation=module.dilation, groups=module.groups, bias=module.bias is not None, padding_mode=module.padding_mode, device=device, dtype=weight_dtype, weights=weights, activations=activations, optimizer=optimizer, quantize_input=True, ) def set_default_dtype(self, dtype): self._gguf_default_dtype = dtype @property def qweight(self): if self.weight_qtype == _GGUF_QTYPE: return self.weight return super().qweight def forward(self, input: torch.Tensor) -> torch.Tensor: qweight = self.qweight if isinstance(qweight, GGUFWeightTensor): target_dtype = _resolve_default_dtype(self._gguf_default_dtype, fallback=input.dtype) weight = qweight.dequantize(dtype=target_dtype, device=input.device) bias = _maybe_cast_bias(self.bias, weight.dtype) if input.dtype != weight.dtype: input = input.to(weight.dtype) return torch.nn.functional.conv1d( input, weight, bias=bias, stride=self.stride, padding=self.padding, dilation=self.dilation, groups=self.groups, ) return torch.nn.functional.conv1d( input, qweight, bias=self.bias, stride=self.stride, padding=self.padding, dilation=self.dilation, groups=self.groups, ) def _load_from_state_dict( self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs ): if self.weight_qtype != _GGUF_QTYPE: return super()._load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs ) weight_key = prefix + "weight" bias_key = prefix + "bias" input_scale_key = prefix + "input_scale" output_scale_key = prefix + "output_scale" weight_raw = state_dict.pop(weight_key, None) bias = state_dict.pop(bias_key, None) input_scale = state_dict.pop(input_scale_key, None) output_scale = state_dict.pop(output_scale_key, None) if weight_raw is None: missing_keys.append(weight_key) target_dtype = _resolve_default_dtype(self._gguf_default_dtype, fallback=self.weight.dtype) if weight_raw is not None: gguf_weight = GGUFWeightTensor.create( raw_tensor=weight_raw, size=self.weight.size(), stride=self.weight.stride(), dtype=target_dtype, device=weight_raw.device, requires_grad=False, ) self.weight = torch.nn.Parameter(gguf_weight, requires_grad=False) if bias is not None: self.bias = torch.nn.Parameter(bias, requires_grad=False) if torch.is_tensor(weight_raw): scale_device = weight_raw.device elif torch.is_tensor(self.weight): scale_device = self.weight.device elif torch.is_tensor(bias): scale_device = bias.device else: scale_device = torch.device("cpu") if input_scale is not None: self.input_scale = input_scale.to(scale_device) elif not hasattr(self, "input_scale") or self.input_scale.is_meta: scale_dtype = self.input_scale.dtype if hasattr(self, "input_scale") else torch.float32 self.input_scale = torch.ones((), dtype=scale_dtype, device=scale_device) if output_scale is not None: self.output_scale = output_scale.to(scale_device) elif not hasattr(self, "output_scale") or self.output_scale.is_meta: scale_dtype = self.output_scale.dtype if hasattr(self, "output_scale") else torch.float32 self.output_scale = torch.ones((), dtype=scale_dtype, device=scale_device) return @register_qmodule(torch.nn.Embedding) class QEmbedding(BaseQEmbedding): def __init__( self, num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False, device=None, dtype=None, weights=None, activations=None, optimizer=None, quantize_input=True, ): super().__init__( num_embeddings, embedding_dim, padding_idx=padding_idx, max_norm=max_norm, norm_type=norm_type, scale_grad_by_freq=scale_grad_by_freq, sparse=sparse, device=device, dtype=dtype, weights=weights, activations=activations, optimizer=optimizer, quantize_input=quantize_input, ) self._gguf_default_dtype = dtype def set_default_dtype(self, dtype): self._gguf_default_dtype = dtype @property def qweight(self): if self.weight_qtype == _GGUF_QTYPE: return self.weight return super().qweight def forward(self, input: torch.Tensor) -> torch.Tensor: qweight = self.qweight if isinstance(qweight, GGUFWeightTensor): target_dtype = _resolve_default_dtype(self._gguf_default_dtype, fallback=torch.float16) if self.max_norm is None and not self.scale_grad_by_freq and not self.sparse: if _may_try_llamacpp_cuda_embedding(qweight, input): fast_out = _try_llamacpp_cuda_embedding(qweight, input, target_dtype) if fast_out is not None: return fast_out weight = qweight.dequantize(dtype=target_dtype, device=input.device) return torch.nn.functional.embedding( input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse, ) return super().forward(input) def _load_from_state_dict( self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs ): if self.weight_qtype != _GGUF_QTYPE: return super()._load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs ) weight_key = prefix + "weight" weight_raw = state_dict.pop(weight_key, None) if weight_raw is None: missing_keys.append(weight_key) return target_dtype = _resolve_default_dtype(self._gguf_default_dtype, fallback=self.weight.dtype) gguf_weight = GGUFWeightTensor.create( raw_tensor=weight_raw, size=self.weight.size(), stride=self.weight.stride(), dtype=target_dtype, device=weight_raw.device, requires_grad=False, ) self.weight = torch.nn.Parameter(gguf_weight, requires_grad=False) scale_device = weight_raw.device if torch.is_tensor(weight_raw) else torch.device("cpu") if not hasattr(self, "input_scale") or self.input_scale is None or self.input_scale.is_meta: self.input_scale = torch.ones((), dtype=torch.float32, device=scale_device) else: self.input_scale = self.input_scale.to(scale_device) if not hasattr(self, "output_scale") or self.output_scale is None or self.output_scale.is_meta: self.output_scale = torch.ones((), dtype=torch.float32, device=scale_device) else: self.output_scale = self.output_scale.to(scale_device) def _collect_gguf_specs(state_dict): specs = [] for key, tensor in state_dict.items(): if not key.endswith(".weight"): continue if not _is_gguf_qtype(getattr(tensor, "tensor_type", None)): continue specs.append({"name": key[:-7], "tensor": tensor}) return specs def detect(state_dict, verboseLevel=1): if gguf is None: return {"matched": False, "kind": "none", "details": {"error": "gguf not installed"}} specs = _collect_gguf_specs(state_dict) if not specs: return {"matched": False, "kind": "none", "details": {}} names = [spec["name"] for spec in specs][:8] return {"matched": True, "kind": "gguf", "details": {"count": len(specs), "names": names}} def convert_to_quanto(state_dict, default_dtype, verboseLevel=1, detection=None): if gguf is None: return {"state_dict": state_dict, "quant_map": {}} if detection is not None and not detection.get("matched", False): return {"state_dict": state_dict, "quant_map": {}} specs = _collect_gguf_specs(state_dict) if not specs: return {"state_dict": state_dict, "quant_map": {}} _set_default_dtype_from_loader(default_dtype) quant_map = {spec["name"]: {"weights": _GGUF_QTYPE_NAME, "activations": "none"} for spec in specs} return {"state_dict": state_dict, "quant_map": quant_map} def apply_pre_quantization(model, state_dict, quantization_map, default_dtype=None, verboseLevel=1): if default_dtype is None or model is None or not quantization_map: return quantization_map or {}, [] _set_default_dtype_from_loader(default_dtype) quantized = set(quantization_map.keys()) for name, module in model.named_modules(): if name in quantized and isinstance(module, torch.nn.Linear): module._router_default_dtype = default_dtype return quantization_map or {}, []