Upload 12 files

__main__.py

@@ -14,6 +14,7 @@ Examples::
 from __future__ import annotations
 
 import argparse
+from pathlib import Path
 
 import numpy as np
 
@@ -25,6 +26,7 @@ from .gguf import (
     is_q2k_eligible,
     lookup_imatrix,
     read_imatrix,
+    write_quantized_gguf,
 )
 from .pipeline import Q2KQuantizer
 from .reference import reference_solve, synthetic_cases
@@ -84,6 +86,29 @@ def _run_gguf(args, cfg: HPCConfig) -> None:
         return
 
     quantizer = Q2KQuantizer(cfg)
+    wants_gguf = bool(args.output and str(args.output).lower().endswith(".gguf"))
+    if wants_gguf:
+        if not args.all_q2k and not args.tensor:
+            args.all_q2k = True
+        if args.rows:
+            raise SystemExit("--rows is only valid for .npz sampling output")
+        names = {args.tensor} if args.tensor else None
+        stats = write_quantized_gguf(
+            args.gguf,
+            args.output,
+            quantizer,
+            imat=imat,
+            quantize_all=args.all_q2k,
+            tensor_names=names,
+            limit=args.limit,
+            verbose=True,
+        )
+        print(
+            f"wrote {args.output}: {int(stats['quantized'])}/{int(stats['tensors'])} tensors Q2_K, "
+            f"rmse={stats['rmse']:.6e}"
+        )
+        return
+
     if args.all_q2k:
         payload: dict[str, np.ndarray] = {}
         done = 0
@@ -191,7 +216,7 @@ def main() -> None:
     parser.add_argument("--importance", help="optional .npy importance array")
     parser.add_argument("--imatrix-gguf", help="llama.cpp iMatrix .gguf file")
     parser.add_argument("--imatrix", help="alias for --imatrix-gguf (GGUF or legacy .dat)")
-    parser.add_argument("--output", help="output .npz")
+    parser.add_argument("--output", help="output .gguf (full model) or .npz (block arrays)")
     parser.add_argument("--demo", action="store_true", help="quantize synthetic weights")
     parser.add_argument("--blocks", type=int, default=64, help="demo block count")
     parser.add_argument("--std", type=float, default=0.02, help="demo weight std")
@@ -212,6 +237,13 @@ def main() -> None:
         return
 
     if args.gguf:
+        if (
+            not args.output
+            and not args.list_tensors
+            and (args.all_q2k or args.tensor)
+        ):
+            stem = Path(args.gguf).stem
+            args.output = str(Path(args.gguf).with_name(f"{stem}-Q2K.gguf"))
         _run_gguf(args, cfg)
         return
 

gguf.py

@@ -20,7 +20,7 @@ from pathlib import Path
 
 import numpy as np
 
-from .q2k import QK_K, dequantize_params_batch, unpack_codes
+from .q2k import QK_K, blocks_to_bytes, dequantize_params_batch, unpack_codes
 
 GGUF_MAGIC = 0x46554747
 DEFAULT_ALIGNMENT = 32
@@ -69,6 +69,14 @@ BLOCK_BYTES = {
     GGML_TYPE_Q8_0: 34,
     GGML_TYPE_Q2_K: 84,
 }
+SUPPORTED_READ_TYPES = {
+    GGML_TYPE_F32,
+    GGML_TYPE_F16,
+    GGML_TYPE_BF16,
+    GGML_TYPE_Q4_0,
+    GGML_TYPE_Q8_0,
+    GGML_TYPE_Q2_K,
+}
 
 
 @dataclass(slots=True)
@@ -154,17 +162,26 @@ class GGUFReader:
     def __init__(self, path: str | Path):
         self.path = Path(path)
         self.metadata: dict[str, object] = {}
+        self.kv_pairs: list[tuple[str, int, bytes]] = []
         self.tensors: list[TensorInfo] = []
         with self.path.open("rb") as f:
             magic, version = struct.unpack("<II", f.read(8))
             if magic != GGUF_MAGIC:
                 raise ValueError(f"{self.path} is not a GGUF file")
             self.version = version
+            self.kv_start = f.tell()
             n_tensors, n_kv = struct.unpack("<QQ", f.read(16))
             for _ in range(n_kv):
                 key = _read_string(f)
                 vtype = struct.unpack("<I", f.read(4))[0]
+                value_start = f.tell()
                 self.metadata[key] = _read_value(f, vtype)
+                value_end = f.tell()
+                f.seek(value_start)
+                raw = f.read(value_end - value_start)
+                f.seek(value_end)
+                self.kv_pairs.append((key, vtype, raw))
+            self.kv_end = f.tell()
             for _ in range(n_tensors):
                 name = _read_string(f)
                 n_dims = struct.unpack("<I", f.read(4))[0]
@@ -172,7 +189,9 @@ class GGUFReader:
                 ttype = struct.unpack("<I", f.read(4))[0]
                 offset = struct.unpack("<Q", f.read(8))[0]
                 self.tensors.append(TensorInfo(name, dims, ttype, offset))
-            self.data_start = _align(f.tell(), int(self.metadata.get("general.alignment", DEFAULT_ALIGNMENT)))
+            self.info_end = f.tell()
+            self.alignment = int(self.metadata.get("general.alignment", DEFAULT_ALIGNMENT))
+            self.data_start = _align(self.info_end, self.alignment)
         self._by_name = {ti.name: ti for ti in self.tensors}
 
     def info(self, name: str) -> TensorInfo:
@@ -184,14 +203,30 @@ class GGUFReader:
     def list_tensors(self) -> list[TensorInfo]:
         return list(self.tensors)
 
-    def read_tensor(self, name: str) -> tuple[np.ndarray, int]:
+    def read_tensor_raw(self, name: str) -> bytes:
         ti = self.info(name)
         with self.path.open("rb") as f:
             f.seek(self.data_start + ti.offset)
             raw = f.read(ti.n_bytes)
         if len(raw) != ti.n_bytes:
             raise ValueError(f"{name}: truncated tensor data")
-        return _dequantize(raw, ti), ti.row_width
+        return raw
+
+    def copy_tensor_raw_to(self, name: str, out, chunk: int = 1 << 20) -> None:
+        ti = self.info(name)
+        remaining = ti.n_bytes
+        with self.path.open("rb") as f:
+            f.seek(self.data_start + ti.offset)
+            while remaining > 0:
+                data = f.read(min(chunk, remaining))
+                if not data:
+                    raise ValueError(f"{name}: truncated tensor data")
+                out.write(data)
+                remaining -= len(data)
+
+    def read_tensor(self, name: str) -> tuple[np.ndarray, int]:
+        raw = self.read_tensor_raw(name)
+        return _dequantize(raw, self.info(name)), self.info(name).row_width
 
 
 def _dequantize(raw: bytes, ti: TensorInfo) -> np.ndarray:
@@ -324,3 +359,140 @@ def is_q2k_eligible(name: str, dims: tuple[int, ...]) -> bool:
     for d in dims:
         n_elements *= int(d)
     return n_elements >= QK_K and n_elements % QK_K == 0
+
+
+def _write_string(f, value: str) -> None:
+    data = value.encode("utf-8")
+    f.write(struct.pack("<Q", len(data)))
+    f.write(data)
+
+
+def write_quantized_gguf(
+    input_path: str | Path,
+    output_path: str | Path,
+    quantizer,
+    imat: dict[str, np.ndarray] | None = None,
+    quantize_all: bool = True,
+    tensor_names: set[str] | None = None,
+    limit: int = 0,
+    verbose: bool = True,
+) -> dict[str, float]:
+    """Quantize a source GGUF with the V2 engine and emit a valid GGUF.
+
+    Eligible tensors become Q2_K; everything else is copied byte-for-byte,
+    including the metadata/tokenizer section.  Returns aggregate statistics.
+    """
+    reader = GGUFReader(input_path)
+    alignment = reader.alignment
+    names = set(tensor_names) if tensor_names else None
+    if names is not None:
+        missing = names - set(reader._by_name)
+        if missing:
+            raise KeyError(f"tensors not found: {sorted(missing)}")
+
+    plan: list[bool] = []
+    quantized_count = 0
+    for ti in reader.tensors:
+        if names is not None:
+            want = ti.name in names
+            compatible = (
+                len(ti.dims) >= 2
+                and ti.dims[0] % QK_K == 0
+                and ti.n_elements % QK_K == 0
+            )
+        else:
+            want = quantize_all and is_q2k_eligible(ti.name, ti.dims)
+            compatible = True
+        if want and ti.type not in SUPPORTED_READ_TYPES:
+            if verbose:
+                print(f"  copy (unsupported source {ti.type_name}): {ti.name}")
+            want = False
+        if want and not compatible:
+            if verbose:
+                print(f"  copy (row width not Q2_K aligned): {ti.name}")
+            want = False
+        if want and limit and quantized_count >= limit:
+            want = False
+        if want:
+            quantized_count += 1
+        plan.append(want)
+
+    out_infos = []
+    data_offset = 0
+    for ti, want in zip(reader.tensors, plan):
+        if want:
+            n_blocks = ti.n_elements // QK_K
+            out_type = GGML_TYPE_Q2_K
+            out_size = n_blocks * 84
+        else:
+            out_type = ti.type
+            out_size = ti.n_bytes
+        out_infos.append((ti, out_type, out_size, data_offset))
+        data_offset = _align(data_offset + out_size, alignment)
+
+    kv_pairs = []
+    for key, vtype, raw in reader.kv_pairs:
+        if quantized_count and key == "general.file_type" and vtype == 4:
+            raw = struct.pack("<I", 10)
+        elif quantized_count and key == "general.quantization_version" and vtype == 4:
+            raw = struct.pack("<I", 2)
+        kv_pairs.append((key, vtype, raw))
+
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    total_rmse = 0.0
+    total_quantized = 0
+    total_elements = 0
+    with output_path.open("wb") as fout:
+        fout.write(struct.pack("<IIQQ", GGUF_MAGIC, reader.version, len(out_infos), len(kv_pairs)))
+        for key, vtype, raw in kv_pairs:
+            _write_string(fout, key)
+            fout.write(struct.pack("<I", vtype))
+            fout.write(raw)
+        for ti, out_type, out_size, offset in out_infos:
+            _write_string(fout, ti.name)
+            fout.write(struct.pack("<I", len(ti.dims)))
+            for d in ti.dims:
+                fout.write(struct.pack("<Q", int(d)))
+            fout.write(struct.pack("<I", out_type))
+            fout.write(struct.pack("<Q", offset))
+
+        data_start = _align(fout.tell(), alignment)
+        fout.write(b"\x00" * (data_start - fout.tell()))
+
+        for idx, (ti, out_type, out_size, offset) in enumerate(out_infos):
+            target = data_start + offset
+            if target > fout.tell():
+                fout.write(b"\x00" * (target - fout.tell()))
+            if plan[idx]:
+                x, row = reader.read_tensor(ti.name)
+                imp = None
+                if imat is not None:
+                    raw_imp = lookup_imatrix(imat, ti.name)
+                    if raw_imp is not None:
+                        imp = expand_importance(raw_imp, x.size, row)
+                res = quantizer.quantize(x, importance=imp, row_width=row)
+                payload = blocks_to_bytes(res.scales, res.qs, res.dm)
+                if len(payload) != out_size:
+                    raise RuntimeError(f"{ti.name}: Q2_K payload size mismatch")
+                fout.write(payload)
+                total_rmse += res.stats.rmse * res.stats.rmse * x.size
+                total_elements += x.size
+                total_quantized += 1
+                if verbose:
+                    print(
+                        f"  Q2_K {ti.name:48s} row={row:6d} "
+                        f"rmse={res.stats.rmse:.6e} imatrix={'yes' if imp is not None else 'no'}"
+                    )
+            else:
+                reader.copy_tensor_raw_to(ti.name, fout)
+                if verbose and idx % 25 == 0:
+                    print(f"  copy {ti.name}")
+
+    rmse = float(np.sqrt(total_rmse / total_elements)) if total_elements else 0.0
+    return {
+        "tensors": float(len(out_infos)),
+        "quantized": float(total_quantized),
+        "elements": float(total_elements),
+        "rmse": rmse,
+    }

q2k.py

@@ -169,6 +169,20 @@ def blocks_to_arrays(blocks: list[Q2KBlock]) -> tuple[np.ndarray, np.ndarray, np
     return scales, qs, dm
 
 
+def blocks_to_bytes(scales: np.ndarray, qs: np.ndarray, dm: np.ndarray) -> bytes:
+    """Serialize Q2_K block arrays to GGUF's 84-byte block layout."""
+    scales = np.asarray(scales, dtype=np.uint8)
+    qs = np.asarray(qs, dtype=np.uint8)
+    dm = np.asarray(dm, dtype=np.uint16)
+    n = scales.shape[0]
+    out = np.empty((n, 84), dtype=np.uint8)
+    out[:, :16] = scales
+    out[:, 16:80] = qs
+    out[:, 80:82] = np.ascontiguousarray(dm[:, 0]).view(np.uint8).reshape(n, 2)
+    out[:, 82:84] = np.ascontiguousarray(dm[:, 1]).view(np.uint8).reshape(n, 2)
+    return out.tobytes()
+
+
 def dequantize_blocks(blocks: list[Q2KBlock]) -> np.ndarray:
     if not blocks:
         return np.empty(0, dtype=np.float32)