Encode on the MPS stream's own queue

torch owns the command buffer from its serial dispatch queue. Asking for the
encoder on the calling thread races with torch committing that buffer, and
encoding into a committed buffer aborts with

failed assertion _status < MTLCommandBufferStatusCommitted

which reproduces reliably during model loading, where many copies are committing
concurrently. Isolated calls never hit it, which is why the op-level tests passed.

devscripts/ab_compiled_vs_intree.py

@@ -0,0 +1,63 @@
+"""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

@@ -0,0 +1,66 @@
+"""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

@@ -0,0 +1,48 @@
+"""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

@@ -0,0 +1,40 @@
+"""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

@@ -0,0 +1,15 @@
+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

@@ -0,0 +1,37 @@
+"""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

@@ -0,0 +1,54 @@
+"""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

@@ -0,0 +1,39 @@
+"""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}")

gguf_metal/ggml_dispatch.mm

@@ -146,15 +146,20 @@ extern "C" int gguf_metal_mul_mat(void *blocks, size_t blocks_off, void *x, size
   const int32_t ne12 = 1;
 
   // torch keeps one encoder open on its command buffer and coalesces into it; opening a second is
-  // an error. Reusing it is also what puts these dispatches in the same submission as torch's own.
+  // an error, so its encoder is reused -- which also puts these dispatches in the same submission
+  // as torch's own. That buffer is owned by the stream's serial queue though, so asking for the
+  // encoder from the calling thread races with torch committing it, and encoding into a committed
+  // buffer aborts (`_status < MTLCommandBufferStatusCommitted`). Everything below therefore runs on
+  // that queue, which is also what torch's own extension guidance prescribes.
   at::mps::MPSStream *stream = at::mps::getCurrentMPSStream();
+  __block int rc = 0;
+  dispatch_sync(stream->queue(), ^{
   id<MTLComputeCommandEncoder> enc = stream->commandEncoder();
 
   id<MTLBuffer> b_blocks = (__bridge id<MTLBuffer>)blocks;
   id<MTLBuffer> b_x = (__bridge id<MTLBuffer>)x;
   id<MTLBuffer> b_out = (__bridge id<MTLBuffer>)out;
 
-  int rc = 0;
   if (M <= 8) {
     // gemv: one token (or a handful), the decode path
     const int nsg = info->mv_nsg, nr0 = info->mv_nr0;
@@ -242,6 +247,7 @@ extern "C" int gguf_metal_mul_mat(void *blocks, size_t blocks_off, void *x, size
           threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
     }
   }
+  });
 
   return rc;
 }
@@ -280,7 +286,9 @@ extern "C" int gguf_metal_get_rows(void *blocks, size_t blocks_off, void *indice
   const NSUInteger nth = MIN((NSUInteger)ne00t, max_threads);
   const NSUInteger nw0 = (ne00t + nth - 1) / nth;
 
+  // Same as the matmul: encode on the stream's own queue so this cannot race torch's commits.
   at::mps::MPSStream *stream = at::mps::getCurrentMPSStream();
+  dispatch_sync(stream->queue(), ^{
   id<MTLComputeCommandEncoder> enc = stream->commandEncoder();
   [enc setComputePipelineState:pso];
   [enc setBytes:&args length:sizeof(args) atIndex:0];
@@ -289,5 +297,6 @@ extern "C" int gguf_metal_get_rows(void *blocks, size_t blocks_off, void *indice
   [enc setBuffer:(__bridge id<MTLBuffer>)out offset:out_off atIndex:3];
   [enc dispatchThreadgroups:MTLSizeMake(nw0 * rows, 1, 1)
       threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
+  });
   return 0;
 }