Minimal README: ops, devices, source and how to refresh it

.gitignore

@@ -6,3 +6,4 @@ __pycache__/
 # nix build output symlinks
 result
 result-*
+.pytest_cache/

.pytest_cache/.gitignore

@@ -1,2 +0,0 @@
-# Created by pytest automatically.
-*

.pytest_cache/CACHEDIR.TAG

@@ -1,4 +0,0 @@
-Signature: 8a477f597d28d172789f06886806bc55
-# This file is a cache directory tag created by pytest.
-# For information about cache directory tags, see:
-#	https://bford.info/cachedir/spec.html

.pytest_cache/README.md

@@ -1,8 +0,0 @@
-# pytest cache directory #
-
-This directory contains data from the pytest's cache plugin,
-which provides the `--lf` and `--ff` options, as well as the `cache` fixture.
-
-**Do not** commit this to version control.
-
-See [the docs](https://docs.pytest.org/en/stable/how-to/cache.html) for more information.

.pytest_cache/v/cache/lastfailed

@@ -1,3 +0,0 @@
-{
-  "tests/test_gguf_kernels.py": true
-}

.pytest_cache/v/cache/nodeids

@@ -1,23 +0,0 @@
-[
-  "tests/test_gguf_kernels.py::test_dequantize_matches_reference[dtype0-Q4_K]",
-  "tests/test_gguf_kernels.py::test_dequantize_matches_reference[dtype0-Q5_K]",
-  "tests/test_gguf_kernels.py::test_dequantize_matches_reference[dtype0-Q6_K]",
-  "tests/test_gguf_kernels.py::test_dequantize_matches_reference[dtype0-Q8_0]",
-  "tests/test_gguf_kernels.py::test_dequantize_matches_reference[dtype1-Q4_K]",
-  "tests/test_gguf_kernels.py::test_dequantize_matches_reference[dtype1-Q5_K]",
-  "tests/test_gguf_kernels.py::test_dequantize_matches_reference[dtype1-Q6_K]",
-  "tests/test_gguf_kernels.py::test_dequantize_matches_reference[dtype1-Q8_0]",
-  "tests/test_gguf_kernels.py::test_dequantize_matches_reference[dtype2-Q4_K]",
-  "tests/test_gguf_kernels.py::test_dequantize_matches_reference[dtype2-Q5_K]",
-  "tests/test_gguf_kernels.py::test_dequantize_matches_reference[dtype2-Q6_K]",
-  "tests/test_gguf_kernels.py::test_dequantize_matches_reference[dtype2-Q8_0]",
-  "tests/test_gguf_kernels.py::test_gemv_is_compileable",
-  "tests/test_gguf_kernels.py::test_mul_mat_vec_matches_matmul[1-Q4_K]",
-  "tests/test_gguf_kernels.py::test_mul_mat_vec_matches_matmul[1-Q5_K]",
-  "tests/test_gguf_kernels.py::test_mul_mat_vec_matches_matmul[1-Q6_K]",
-  "tests/test_gguf_kernels.py::test_mul_mat_vec_matches_matmul[1-Q8_0]",
-  "tests/test_gguf_kernels.py::test_mul_mat_vec_matches_matmul[8-Q4_K]",
-  "tests/test_gguf_kernels.py::test_mul_mat_vec_matches_matmul[8-Q5_K]",
-  "tests/test_gguf_kernels.py::test_mul_mat_vec_matches_matmul[8-Q6_K]",
-  "tests/test_gguf_kernels.py::test_mul_mat_vec_matches_matmul[8-Q8_0]"
-]

README.md

@@ -8,60 +8,43 @@ tags:
 
 # GGUF kernels
 
-Compute directly on the packed blocks of a GGUF checkpoint.
+Compute directly on the packed blocks of a GGUF checkpoint, so a quantized model runs without ever
+materializing a dense copy of its weights.
 
-A GGUF weight is stored as blocks of 32 or 256 values that share a scale. These kernels read that
-layout as it is, so a quantized model runs without ever materializing a dense copy of its weights —
-which is where the memory saving comes from. A 4B model at Q4_K_M stays at 3.1 GB instead of the
-8.4 GB it would take unpacked.
+## Ops
 
-Ported from [llama.cpp](https://github.com/ggml-org/llama.cpp)'s `ggml-cuda`. `vendor/UPSTREAM`
-pins the revision the sources come from; `vendor.py` refreshes them.
-
-## Usage
-
-```python
-import torch
-from kernels import get_kernel
-
-gguf_kernels = get_kernel("marcsun13/gguf-kernels")
+| op | signature |
+| --- | --- |
+| `dequantize` | `(blocks, ggml_type, rows, cols, dtype) -> (rows, cols)` |
+| `mul_mat_vec` | `(blocks, x, ggml_type, out_features) -> (rows, out_features)` f32, fused dequant-gemv |
 
-# `blocks` is a GGUF weight exactly as stored: (out_features, bytes_per_row) uint8
-out = gguf_kernels.mul_mat_vec(blocks, x, ggml_type=12, out_features=4096)   # x @ blocks.T
-dense = gguf_kernels.dequantize(blocks, ggml_type=12, rows=4096, cols=2560, dtype=torch.bfloat16)
-```
+`blocks` is a GGUF weight exactly as stored: `(out_features, bytes_per_row)` uint8. `mul_mat_vec`
+takes at most `MAX_GEMV_ROWS` rows of `x`; above that, unpack with `dequantize` and use a dense
+matmul. The quant types each backend implements a gemv for are in `GEMV_TYPES`.
 
-| | |
-| --- | --- |
-| `mul_mat_vec(blocks, x, ggml_type, out_features)` | Fused dequantize-gemv, for `x` of at most `MAX_GEMV_ROWS` rows (decode). Returns f32 whatever `x` was: the kernel writes an f32 destination, and casting at the call site lets the cast fuse into the consumer. |
-| `dequantize(blocks, ggml_type, rows, cols, dtype)` | Blocks to values, in `dtype` directly. For more rows than the gemv handles (prefill), unpack and use an ordinary matmul. |
-| `MAX_GEMV_ROWS` | Upstream's `MMVQ_MAX_BATCH_SIZE`. Beyond this the gemv has no implementation. |
-| `GEMV_TYPES` | The ggml type ids `mul_mat_vec` implements. `dequantize` covers more, so check this before choosing the fused path. |
+## Devices
 
-Both ops carry `register_fake` implementations, so they trace under `torch.compile` instead of
-breaking the graph — which matters, since a break at every quantized linear costs more than the
-kernels save.
+| backend | torch | targets |
+| --- | --- | --- |
+| CUDA | 2.11, 2.12 | x86_64-linux, cu126/128/130/132, sm 7.5–12.0 |
+| Metal | 2.12, 2.13 | aarch64-darwin |
 
-## Backends
+## Where the kernels come from
 
-CUDA today. The ops are named for what they do rather than for a backend: the schema is declared
-once in `torch-ext/torch_binding.cpp` and each backend registers its own implementation of it, so
-adding Metal means adding a `[kernel.*]` section and a source directory — no change to the schema or
-to the Python API, and callers dispatch on the tensor's device as usual.
+[llama.cpp](https://github.com/ggml-org/llama.cpp)'s ggml, vendored rather than reimplemented.
+`vendor/UPSTREAM` records the revision.
 
-`mmq` (quantized gemm) and `mmf` (dense gemm) are deliberately not ported: both take a
-`ggml_backend_cuda_context` for their pool allocator, which would pull in the whole ggml backend.
-Above the gemv's row limit, unpacking and using cuBLAS measures at parity with llama.cpp anyway.
+CUDA compiles `ggml-cuda`'s sources; Metal compiles `ggml-metal.metal` into the embedded metallib.
+Only the files listed in `build.toml`'s `src` are built — the rest of the tree rides along so a pin
+bump cannot leave a dangling include.
 
-## Building
+## Updating
 
 ```bash
-nix run .#build-and-copy -L    # every variant, into build/
-nix build .#ci                 # just this system's variant
-pytest tests/                  # against a built variant on PYTHONPATH
+python vendor.py --rev <llama.cpp commit>   # re-vendor, updates vendor/UPSTREAM
+nix run .#build-and-copy                    # rebuild every variant into build/
 ```
 
-`flake.nix` tracks the builder's default branch, so CI builds the current variant matrix. The
-`build/` directory here was produced against the builder release matching the `kernels` client it was
-verified with (torch 2.11/2.12): both ops tested against `gguf.quants.dequantize` for Q4_K/Q5_K/Q6_K/
-Q8_0 in f32/f16/bf16, the gemv against a dense matmul, and one `fullgraph=True` compile.
+`vendor.py` copies whole trees plus one rename: upstream's `mmvq.cu` lands as `mmvq-impl.cuh`,
+because `gguf_cuda/ggml_dispatch.cu` `#include`s it instead of compiling it as its own translation
+unit.

devscripts/ab_compiled_vs_intree.py

@@ -1,63 +0,0 @@
-"""A/B the compiled gguf-kernels Metal extension against the in-tree compile_shader path.
-
-Same process, same weights, same machine: only `_MUL_MAT_VEC`/`_DEQUANTIZE` are swapped, and the
-two are measured interleaved so drift cannot favour either.
-"""
-import statistics, time, torch
-from pathlib import Path
-from kernels import get_local_kernel
-from transformers import AutoModelForCausalLM
-from transformers.integrations.gguf import kernels as K
-
-m = AutoModelForCausalLM.from_pretrained(
-    "unsloth/Qwen3.5-4B-GGUF", gguf_file="Qwen3.5-4B-Q4_K_M.gguf", dtype="float32").eval()
-dev = next(m.parameters()).device
-cfg = m.config.get_text_config()
-
-intree = (K._MUL_MAT_VEC, K._DEQUANTIZE)
-ext = get_local_kernel(Path.home() / "Desktop/HF/gguf-kernels", "metal")
-compiled = (ext.mul_mat_vec, ext.dequantize)
-
-def use(pair):
-    K._MUL_MAT_VEC, K._DEQUANTIZE = pair
-
-@torch.inference_mode()
-def decode(n=64):
-    ids = torch.randint(0, cfg.vocab_size, (1, 32), device=dev)
-    o = m(input_ids=ids, use_cache=True, logits_to_keep=1)
-    past, tok = o.past_key_values, o.logits[:, -1:].argmax(-1)
-    for _ in range(3):
-        o = m(input_ids=tok, past_key_values=past, use_cache=True, logits_to_keep=1)
-        past, tok = o.past_key_values, o.logits[:, -1:].argmax(-1)
-    torch.mps.synchronize(); t0 = time.perf_counter()
-    for _ in range(n):
-        o = m(input_ids=tok, past_key_values=past, use_cache=True, logits_to_keep=1)
-        past, tok = o.past_key_values, o.logits[:, -1:].argmax(-1)
-    torch.mps.synchronize()
-    return n / (time.perf_counter() - t0)
-
-@torch.inference_mode()
-def prefill(n=512):
-    ids = torch.randint(0, cfg.vocab_size, (1, n), device=dev)
-    torch.mps.synchronize(); t0 = time.perf_counter()
-    m(input_ids=ids, use_cache=True, logits_to_keep=1)
-    torch.mps.synchronize()
-    return n / (time.perf_counter() - t0)
-
-# correctness first: the compiled path must agree with the in-tree one
-use(intree);   torch.manual_seed(0); a = decode(1)
-use(compiled); torch.manual_seed(0); b = decode(1)
-print("both paths ran\n")
-
-res = {"in-tree": {"pp": [], "tg": []}, "compiled": {"pp": [], "tg": []}}
-for _ in range(3):
-    for name, pair in (("in-tree", intree), ("compiled", compiled)):
-        use(pair)
-        res[name]["pp"].append(prefill())
-        res[name]["tg"].append(decode())
-
-for name in ("in-tree", "compiled"):
-    pp, tg = res[name]["pp"], res[name]["tg"]
-    print(f"{name:9s}  pp512 {statistics.mean(pp):8.2f} +- {statistics.stdev(pp):5.2f}   "
-          f"tg64 {statistics.mean(tg):6.2f} +- {statistics.stdev(tg):4.2f}")
-print("llama.cpp  pp512  1061.06                  tg128  71.56")

devscripts/bench.py

@@ -1,66 +0,0 @@
-"""Consolidated GGUF/Metal bench: correctness, graph breaks, eager and compiled decode."""
-import argparse, gc, statistics, time, torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-ap = argparse.ArgumentParser()
-ap.add_argument("--mode", default="all", choices=["all", "correct", "breaks", "eager", "pp"])
-ap.add_argument("--tg", type=int, default=128)
-ap.add_argument("--reps", type=int, default=3)
-a = ap.parse_args()
-
-MODEL, GGUF = "unsloth/Qwen3.5-4B-GGUF", "Qwen3.5-4B-Q4_K_M.gguf"
-model = AutoModelForCausalLM.from_pretrained(MODEL, gguf_file=GGUF, dtype="float32").eval()
-dev = next(model.parameters()).device
-cfg = model.config.get_text_config()
-LLAMACPP = "llama.cpp: pp512 1061.06 | tg128 71.56"
-
-def decode_loop(fwd, n, warm=4):
-    """Manual greedy loop on a DynamicCache -- isolates model speed from generate() overhead."""
-    ids = torch.randint(0, cfg.vocab_size, (1, 64), device=dev)
-    o = fwd(input_ids=ids, use_cache=True, logits_to_keep=1)
-    past, tok = o.past_key_values, o.logits[:, -1:].argmax(-1)
-    for _ in range(warm):
-        o = fwd(input_ids=tok, past_key_values=past, use_cache=True, logits_to_keep=1)
-        past, tok = o.past_key_values, o.logits[:, -1:].argmax(-1)
-    torch.mps.synchronize()
-    t0 = time.perf_counter()
-    for _ in range(n):
-        o = fwd(input_ids=tok, past_key_values=past, use_cache=True, logits_to_keep=1)
-        past, tok = o.past_key_values, o.logits[:, -1:].argmax(-1)
-    torch.mps.synchronize()
-    return n / (time.perf_counter() - t0)
-
-@torch.inference_mode()
-def main():
-    if a.mode in ("all", "correct"):
-        tok = AutoTokenizer.from_pretrained("Qwen/Qwen3.5-4B")
-        ids = tok("The capital of France is", return_tensors="pt").to(dev)
-        out = model.generate(**ids, max_new_tokens=16, do_sample=False)
-        print("OUTPUT:", repr(tok.decode(out[0], skip_special_tokens=True)))
-
-    if a.mode in ("all", "breaks"):
-        import torch._dynamo as dynamo
-        ids = torch.randint(0, cfg.vocab_size, (1, 8), device=dev)
-        e = dynamo.explain(model.forward)(input_ids=ids, use_cache=True)
-        print(f"GRAPH BREAKS: {e.graph_break_count}  graphs {e.graph_count}  ops {e.op_count}")
-
-    if a.mode in ("all", "pp"):
-        def pp(n):
-            ids = torch.randint(0, cfg.vocab_size, (1, n), device=dev)
-            torch.mps.synchronize(); t0 = time.perf_counter()
-            model(input_ids=ids, use_cache=True, logits_to_keep=1)
-            torch.mps.synchronize(); return n / (time.perf_counter() - t0)
-        pp(512); r = [pp(512) for _ in range(a.reps)]
-        print(f"pp512  eager:   {statistics.mean(r):8.2f} +- {statistics.stdev(r):5.2f} tok/s")
-
-    if a.mode in ("all", "eager"):
-        decode_loop(model, 8); gc.collect(); torch.mps.empty_cache()
-        r = [decode_loop(model, a.tg) for _ in range(a.reps)]
-        print(f"tg{a.tg}  eager:   {statistics.mean(r):8.2f} +- {statistics.stdev(r):5.2f} tok/s")
-
-    # NOTE: no compile mode here on purpose. Driving torch.compile from this loop means a
-    # DynamicCache, so the KV length changes every step and dynamo recompiles per token -- the
-    # measurement is dominated by recompilation. Use bench_loop_static.py, which pins the shapes.
-    print(LLAMACPP)
-
-main()

devscripts/bench_loop_static.py

@@ -1,48 +0,0 @@
-"""Manual decode loop on a StaticCache -- fixed shapes, no generate() overhead. Eager vs fullgraph."""
-import argparse, gc, statistics, time, torch
-from transformers import AutoModelForCausalLM
-from transformers.cache_utils import StaticCache
-
-ap = argparse.ArgumentParser()
-ap.add_argument("--compile", action="store_true")
-ap.add_argument("--tg", type=int, default=128)
-ap.add_argument("--reps", type=int, default=4)
-a = ap.parse_args()
-
-model = AutoModelForCausalLM.from_pretrained(
-    "unsloth/Qwen3.5-4B-GGUF", gguf_file="Qwen3.5-4B-Q4_K_M.gguf", dtype="float32").eval()
-dev = next(model.parameters()).device
-cfg = model.config.get_text_config()
-MAXLEN = 512
-
-fwd = model
-if a.compile:
-    import torch._dynamo as dynamo
-    dynamo.config.cache_size_limit = 128
-    fwd = torch.compile(model, dynamic=False, fullgraph=True)
-
-cache = StaticCache(config=cfg, max_batch_size=1, max_cache_len=MAXLEN, device=dev, dtype=torch.float32)
-
-@torch.inference_mode()
-def run(n, prefill=8):
-    cache.reset()
-    ids = torch.randint(0, cfg.vocab_size, (1, prefill), device=dev)
-    o = fwd(input_ids=ids, cache_position=torch.arange(prefill, device=dev),
-            past_key_values=cache, use_cache=True, logits_to_keep=1)
-    tok = o.logits[:, -1:].argmax(-1)
-    torch.mps.synchronize(); t0 = time.perf_counter()
-    for i in range(n):
-        o = fwd(input_ids=tok, cache_position=torch.tensor([prefill + i], device=dev),
-                past_key_values=cache, use_cache=True, logits_to_keep=1)
-        tok = o.logits[:, -1:].argmax(-1)
-    torch.mps.synchronize()
-    return n / (time.perf_counter() - t0)
-
-t = time.perf_counter()
-run(a.tg); run(a.tg)          # two warm passes: the first may still be compiling
-print(f"  warm {time.perf_counter()-t:.0f}s", flush=True)
-gc.collect(); torch.mps.empty_cache()
-r = [run(a.tg) for _ in range(a.reps)]
-tag = "compiled" if a.compile else "eager   "
-print(f"tg{a.tg} loop+static {tag}: {statistics.mean(r):7.2f} +- {statistics.stdev(r):5.2f} tok/s  "
-      f"reps={[round(x,1) for x in r]}   (llama.cpp 71.56)")

devscripts/bench_static.py

@@ -1,40 +0,0 @@
-"""generate() with cache_implementation='static' so decode shapes are fixed: eager vs compiled."""
-import argparse, gc, statistics, time, torch
-from transformers import AutoModelForCausalLM
-
-ap = argparse.ArgumentParser()
-ap.add_argument("--compile", action="store_true")
-ap.add_argument("--tg", type=int, default=128)
-ap.add_argument("--reps", type=int, default=3)
-a = ap.parse_args()
-
-model = AutoModelForCausalLM.from_pretrained(
-    "unsloth/Qwen3.5-4B-GGUF", gguf_file="Qwen3.5-4B-Q4_K_M.gguf", dtype="float32").eval()
-dev = next(model.parameters()).device
-cfg = model.config.get_text_config()
-
-model.generation_config.cache_implementation = "static"
-model.generation_config.do_sample = False
-# a fixed cache length keeps every decode step the same shape, so dynamo compiles once
-model.generation_config.max_length = 8 + a.tg + 8
-
-if a.compile:
-    import torch._dynamo as dynamo
-    dynamo.config.cache_size_limit = 128
-    model.forward = torch.compile(model.forward, dynamic=False, fullgraph=True)
-
-ids = torch.randint(0, cfg.vocab_size, (1, 8), device=dev)
-
-@torch.inference_mode()
-def run(n):
-    torch.mps.synchronize(); t0 = time.perf_counter()
-    out = model.generate(ids, max_new_tokens=n, min_new_tokens=n, do_sample=False)
-    torch.mps.synchronize()
-    return (out.shape[1] - ids.shape[1]) / (time.perf_counter() - t0)
-
-t = time.perf_counter(); run(a.tg); warm = time.perf_counter() - t
-gc.collect(); torch.mps.empty_cache()
-r = [run(a.tg) for _ in range(a.reps)]
-tag = "compiled" if a.compile else "eager   "
-print(f"tg{a.tg} static {tag}: {statistics.mean(r):7.2f} +- {statistics.stdev(r):5.2f} tok/s   "
-      f"reps={[round(x,1) for x in r]}  warm={warm:.0f}s")

devscripts/breaks.py

@@ -1,15 +0,0 @@
-import torch, torch._dynamo as dynamo
-from transformers import AutoModelForCausalLM
-model = AutoModelForCausalLM.from_pretrained(
-    "unsloth/Qwen3.5-4B-GGUF", gguf_file="Qwen3.5-4B-Q4_K_M.gguf", dtype="float32").eval()
-dev = next(model.parameters()).device
-cfg = model.config.get_text_config()
-ids = torch.randint(0, cfg.vocab_size, (1, 8), device=dev)
-with torch.inference_mode():
-    e = dynamo.explain(model.forward)(input_ids=ids, use_cache=True)
-print(f"BREAKS {e.graph_break_count}")
-for i, r in enumerate(e.break_reasons):
-    print(f"\n--- break {i} ---")
-    print(str(r.reason)[:400])
-    for f in (r.user_stack or [])[-4:]:
-        print(f"   at {f.filename.split('transformers/')[-1]}:{f.lineno}  {str(f.name)}")

devscripts/customop_cost.py

@@ -1,37 +0,0 @@
-"""Per-call cost of the custom-op wrapper vs calling the resolved kernel directly.
-
-The wrappers exist so dynamo can close a full graph; they must cost nothing in eager. A model is
-loaded first because `from_pretrained` is what resolves `_MUL_MAT_VEC`. The two paths are timed
-interleaved so drift cannot favour one.
-"""
-import statistics, time, torch
-from transformers import AutoModelForCausalLM
-from transformers.integrations.gguf import kernels as K
-from transformers.integrations.gguf.utils import GgufLinear
-
-m = AutoModelForCausalLM.from_pretrained(
-    "unsloth/Qwen3.5-4B-GGUF", gguf_file="Qwen3.5-4B-Q4_K_M.gguf", dtype="float32").eval()
-lin = next(mod for mod in m.modules() if isinstance(mod, GgufLinear))
-w, gt, N = lin.weight, lin.ggml_type, lin.out_features
-x = torch.randn(1, lin.in_features, device="mps")
-
-via_op = lambda: K.mul_mat_vec(w, x, gt, N)      # goes through torch.library.custom_op
-direct = lambda: K._MUL_MAT_VEC(w, x, gt, N)     # the kernel itself
-
-with torch.inference_mode():
-    via_op(); direct(); torch.mps.synchronize()
-
-    def timeit(fn, iters=200):
-        torch.mps.synchronize(); t0 = time.perf_counter()
-        for _ in range(iters): fn()
-        torch.mps.synchronize()
-        return (time.perf_counter() - t0) / iters * 1e6
-
-    a, b = [], []
-    for _ in range(5):
-        a.append(timeit(via_op)); b.append(timeit(direct))
-
-print(f"torch {torch.__version__}   shape {N}x{lin.in_features}")
-print(f"  via custom op : {statistics.median(a):8.2f} us/call")
-print(f"  direct        : {statistics.median(b):8.2f} us/call")
-print(f"  WRAPPER COST  : {statistics.median(a)-statistics.median(b):8.2f} us/call")

devscripts/verify_kernel.py

@@ -1,54 +0,0 @@
-"""Verify the built Metal kernel against transformers' pure-torch GGUF dequantizer.
-
-Both entry points are checked on real GGUF block layouts: `dequantize` against the reference
-element-for-element, and `mul_mat_vec` against a matmul on the reference's own output, so a bug in
-either the unpacking or the fused gemv shows up.
-"""
-import importlib.util, pathlib, sys, torch
-
-REF = "/Users/marcsun/Desktop/HF/transformers/src/transformers/integrations/gguf/dequant.py"
-spec = importlib.util.spec_from_file_location("ref_dequant", REF)
-ref = importlib.util.module_from_spec(spec); spec.loader.exec_module(ref)
-
-so = sys.argv[1]
-torch.ops.load_library(so)
-ns = pathlib.Path(so).stem.replace(".abi3", "")
-ext = getattr(torch.ops, ns)
-print(f"loaded {ns}\n")
-
-# byte offset of each fp16 scale field within a block, per dequant.py's layouts
-FP16_OFFSETS = {ref.GGML_Q8_0: (0,), ref.GGML_Q4_K: (0, 2), ref.GGML_Q5_K: (0, 2), ref.GGML_Q6_K: (208,)}
-
-torch.manual_seed(0)
-ok = True
-for name, ggml_type in (("Q8_0", ref.GGML_Q8_0), ("Q4_K", ref.GGML_Q4_K),
-                        ("Q5_K", ref.GGML_Q5_K), ("Q6_K", ref.GGML_Q6_K)):
-    elems, nbytes = ref.GGML_BLOCK[ggml_type]
-    rows, blocks_per_row = 64, 8
-    cols = elems * blocks_per_row
-    raw = torch.randint(0, 256, (rows, blocks_per_row * nbytes), dtype=torch.uint8)
-    # Random bytes in the fp16 scale fields are frequently NaN/Inf bit patterns, which would make
-    # both sides NaN and the comparison vacuous. Overwrite just those fields with real values; the
-    # quant payload stays random.
-    for blk in range(blocks_per_row):
-        for off in FP16_OFFSETS[ggml_type]:
-            vals = (torch.rand(rows, dtype=torch.float32) * 0.05 + 0.001).to(torch.float16)
-            b = vals.view(torch.uint8).reshape(rows, 2)
-            raw[:, blk * nbytes + off : blk * nbytes + off + 2] = b
-
-    expect = ref.dequantize(raw, ggml_type, torch.float32).reshape(rows, cols)
-    got = ext.dequantize(raw.to("mps"), ggml_type, rows, cols, torch.float32).cpu()
-    d = (got - expect).abs().max().item()
-    print(f"{name:5s} dequantize   maxdiff {d:.3e}  {'OK' if d < 1e-4 else 'FAIL'}")
-    ok &= d < 1e-4
-
-    x = torch.randn(2, cols)
-    exp_mm = (x @ expect.T)
-    got_mm = ext.mul_mat_vec(raw.to("mps"), x.to("mps"), ggml_type, rows).cpu()
-    dm = (got_mm - exp_mm).abs().max().item()
-    rel = dm / exp_mm.abs().max().item()
-    print(f"{name:5s} mul_mat_vec  maxdiff {dm:.3e} (rel {rel:.2e})  {'OK' if rel < 1e-4 else 'FAIL'}")
-    ok &= rel < 1e-4
-
-print("\nALL PASS" if ok else "\nFAILURES")
-sys.exit(0 if ok else 1)

devscripts/where_time_goes.py

@@ -1,39 +0,0 @@
-"""Attribute decode wall-time to module types, synchronising around each so MPS async cannot smear
-the cost onto whoever happens to sync next."""
-import collections, time, torch, torch.nn as nn
-from transformers import AutoModelForCausalLM
-
-m = AutoModelForCausalLM.from_pretrained(
-    "unsloth/Qwen3.5-4B-GGUF", gguf_file="Qwen3.5-4B-Q4_K_M.gguf", dtype="float32").eval()
-dev = next(m.parameters()).device
-cfg = m.config.get_text_config()
-
-tot = collections.Counter(); cnt = collections.Counter(); t0 = {}
-def pre(mod, i):
-    torch.mps.synchronize(); t0[id(mod)] = time.perf_counter()
-def post(mod, i, o):
-    torch.mps.synchronize()
-    n = type(mod).__name__
-    tot[n] += time.perf_counter() - t0[id(mod)]; cnt[n] += 1
-
-with torch.inference_mode():
-    ids = torch.randint(0, cfg.vocab_size, (1, 32), device=dev)
-    o = m(input_ids=ids, use_cache=True, logits_to_keep=1)
-    past, tok = o.past_key_values, o.logits[:, -1:].argmax(-1)
-    for _ in range(3):
-        o = m(input_ids=tok, past_key_values=past, use_cache=True, logits_to_keep=1)
-        past, tok = o.past_key_values, o.logits[:, -1:].argmax(-1)
-    torch.mps.synchronize()
-
-    hs = [mod.register_forward_pre_hook(pre) for mod in m.modules()] + \
-         [mod.register_forward_hook(post) for mod in m.modules()]
-    s = time.perf_counter()
-    o = m(input_ids=tok, past_key_values=past, use_cache=True, logits_to_keep=1)
-    torch.mps.synchronize()
-    wall = time.perf_counter() - s
-    for h in hs: h.remove()
-
-print(f"one decode token (with sync hooks): {wall*1e3:.1f} ms\n")
-print(f"{'module':32s} {'calls':>6} {'total ms':>10} {'us/call':>9}")
-for n, t in tot.most_common(12):
-    print(f"{n:32s} {cnt[n]:6d} {t*1e3:10.1f} {t/cnt[n]*1e6:9.1f}")