Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+build/torch210-metal-aarch64-darwin/_gguf_dequant_530e374.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch211-metal-aarch64-darwin/_gguf_dequant_530e374.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch28-metal-aarch64-darwin/_gguf_dequant_530e374.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-metal-aarch64-darwin/_gguf_dequant_530e374.abi3.so filter=lfs diff=lfs merge=lfs -text

build/torch210-metal-aarch64-darwin/__init__.py

@@ -0,0 +1,600 @@
+"""GGUF dequantization kernels for Apple Silicon (Metal).
+
+Six standalone kernels covering the GGUF formats most common in the wild:
+
+    Q4_0   Q8_0   Q4_K   Q5_K   Q6_K   IQ4_NL   IQ4_XS
+
+Two API levels are exposed for each:
+
+  - ``dequantize_<format>_bytes(quant_bytes) -> np.ndarray``
+    One-shot path used by the tests for bit-exact validation against a
+    NumPy reference. Allocates input/output buffers internally.
+
+  - ``allocate_<format>_plan(nblocks)`` + ``dispatch_<format>(plan)``
+    Low-level path for benchmarks. Buffers are pre-allocated; ``dispatch``
+    only encodes + dispatches + syncs, so the timing isolates GPU work
+    from CPU<->buffer copies.
+
+The runtime path uses PyObjC + Metal.framework so the correctness loop
+works on any Apple Silicon Mac without Xcode 26 / Nix / kernel-builder.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+# Lazy globals — populated on first use.
+_DEVICE = None
+_COMMAND_QUEUE = None
+_LIBRARY = None
+_PIPELINES: dict = {}
+
+_KERNEL_SOURCE_PATH = Path(__file__).resolve().parent.parent.parent / "gguf_dequant_metal" / "dequantize.metal"
+
+
+def _ensure_metal():
+    global _DEVICE, _COMMAND_QUEUE, _LIBRARY
+    if _DEVICE is not None:
+        return
+    import Metal  # pyobjc-framework-metal
+
+    _DEVICE = Metal.MTLCreateSystemDefaultDevice()
+    if _DEVICE is None:
+        raise RuntimeError("No default Metal device available")
+    _COMMAND_QUEUE = _DEVICE.newCommandQueue()
+
+    src = _KERNEL_SOURCE_PATH.read_text()
+    options = Metal.MTLCompileOptions.new()
+    lib, err = _DEVICE.newLibraryWithSource_options_error_(src, options, None)
+    if lib is None:
+        raise RuntimeError(f"Metal compile failed: {err}")
+    _LIBRARY = lib
+
+
+def _pipeline(fn_name: str):
+    _ensure_metal()
+    pso = _PIPELINES.get(fn_name)
+    if pso is not None:
+        return pso
+
+    fn = _LIBRARY.newFunctionWithName_(fn_name)
+    if fn is None:
+        raise RuntimeError(f"Kernel function not found: {fn_name}")
+    pso, err = _DEVICE.newComputePipelineStateWithFunction_error_(fn, None)
+    if pso is None:
+        raise RuntimeError(f"Pipeline state creation failed: {err}")
+    _PIPELINES[fn_name] = pso
+    return pso
+
+
+# (kernel_name, bytes_per_block, elems_per_block, threads_per_block_group)
+_FORMATS = {
+    "q4_0":   ("dequantize_q4_0",   18, 32, 16),
+    "q8_0":   ("dequantize_q8_0",   34, 32, 32),
+    "q4_K":   ("dequantize_q4_K",  144, 256, 256),
+    "q5_K":   ("dequantize_q5_K",  176, 256, 256),
+    "q6_K":   ("dequantize_q6_K",  210, 256, 256),
+    "iq4_nl": ("dequantize_iq4_nl", 18, 32, 16),
+    "iq4_xs": ("dequantize_iq4_xs", 136, 256, 256),
+}
+
+
+class Plan:
+    """Pre-allocated I/O buffers for a single (format, nblocks) shape."""
+
+    __slots__ = ("fmt", "nblocks", "in_buf", "out_buf", "_nblocks_bytes")
+
+    def __init__(self, fmt: str, nblocks: int, in_buf, out_buf):
+        self.fmt = fmt
+        self.nblocks = nblocks
+        self.in_buf = in_buf
+        self.out_buf = out_buf
+        self._nblocks_bytes = np.array([nblocks], dtype=np.uint32).tobytes()
+
+
+def allocate_plan(fmt: str, nblocks: int) -> Plan:
+    import Metal
+
+    _ensure_metal()
+    _, bpb, epb, _ = _FORMATS[fmt]
+    in_buf = _DEVICE.newBufferWithLength_options_(nblocks * bpb, Metal.MTLResourceStorageModeShared)
+    out_buf = _DEVICE.newBufferWithLength_options_(nblocks * epb * 4, Metal.MTLResourceStorageModeShared)
+    return Plan(fmt=fmt, nblocks=nblocks, in_buf=in_buf, out_buf=out_buf)
+
+
+def dispatch(plan: Plan, sync: bool = True):
+    import Metal
+
+    kernel_name, _, _, tg = _FORMATS[plan.fmt]
+    pso = _pipeline(kernel_name)
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.in_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.out_buf, 0, 1)
+    enc.setBytes_length_atIndex_(plan._nblocks_bytes, len(plan._nblocks_bytes), 2)
+    enc.dispatchThreads_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(plan.nblocks * tg, 1, 1),
+        Metal.MTLSizeMake(tg, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+def dequantize_bytes(fmt: str, quant_bytes) -> np.ndarray:
+    """One-shot bytes-in / float32-out for the named format."""
+    _ensure_metal()
+    _, bpb, epb, _ = _FORMATS[fmt]
+
+    buf = bytes(quant_bytes) if not isinstance(quant_bytes, (bytes, bytearray)) else quant_bytes
+    nbytes = len(buf)
+    if nbytes % bpb != 0:
+        raise ValueError(f"{fmt} input not a multiple of {bpb} bytes: got {nbytes}")
+    nblocks = nbytes // bpb
+
+    plan = allocate_plan(fmt, nblocks)
+    in_view = np.frombuffer(plan.in_buf.contents().as_buffer(nbytes), dtype=np.uint8)
+    in_view[:] = np.frombuffer(buf, dtype=np.uint8)
+    dispatch(plan, sync=True)
+    out_nbytes = nblocks * epb * 4
+    return np.frombuffer(plan.out_buf.contents().as_buffer(out_nbytes), dtype=np.float32).copy()
+
+
+# --- Per-format convenience wrappers ----------------------------------------
+
+def dequantize_q4_0_bytes(qb):   return dequantize_bytes("q4_0",   qb)
+def dequantize_q8_0_bytes(qb):   return dequantize_bytes("q8_0",   qb)
+def dequantize_q4_K_bytes(qb):   return dequantize_bytes("q4_K",   qb)
+def dequantize_q5_K_bytes(qb):   return dequantize_bytes("q5_K",   qb)
+def dequantize_q6_K_bytes(qb):   return dequantize_bytes("q6_K",   qb)
+def dequantize_iq4_nl_bytes(qb): return dequantize_bytes("iq4_nl", qb)
+def dequantize_iq4_xs_bytes(qb): return dequantize_bytes("iq4_xs", qb)
+
+def allocate_q4_0_plan(nblocks):   return allocate_plan("q4_0",   nblocks)
+def allocate_q8_0_plan(nblocks):   return allocate_plan("q8_0",   nblocks)
+def allocate_q4_K_plan(nblocks):   return allocate_plan("q4_K",   nblocks)
+def allocate_q5_K_plan(nblocks):   return allocate_plan("q5_K",   nblocks)
+def allocate_q6_K_plan(nblocks):   return allocate_plan("q6_K",   nblocks)
+def allocate_iq4_nl_plan(nblocks): return allocate_plan("iq4_nl", nblocks)
+def allocate_iq4_xs_plan(nblocks): return allocate_plan("iq4_xs", nblocks)
+
+def dispatch_q4_0(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q8_0(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q4_K(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q5_K(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q6_K(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_iq4_nl(plan, **kw): return dispatch(plan, **kw)
+def dispatch_iq4_xs(plan, **kw): return dispatch(plan, **kw)
+
+
+# ---------------------------------------------------------------------------
+# Torch op entry points — the production path used by transformers (no PyObjC,
+# no runtime compile). Only available in builds produced by kernel-builder,
+# which generates `_ops.py` exporting `ops = torch.ops._gguf_dequant_<rev>`.
+# In a dev-mode checkout (running from torch-ext/ directly) `_ops` isn't
+# present, so the torch_* entry points raise on call.
+# ---------------------------------------------------------------------------
+
+try:
+    from ._ops import ops as _ops  # generated by kernel-builder
+except ImportError:
+    _ops = None
+
+
+def _require_ops():
+    if _ops is None:
+        raise RuntimeError(
+            "torch op binding not available — this build was loaded from the dev tree. "
+            "Use the *_bytes / *_plan APIs (PyObjC path) or install a kernel-builder build."
+        )
+    return _ops
+
+
+def torch_dequantize_q4_0(input, output):   _require_ops().dequantize_q4_0(input, output)
+def torch_dequantize_q8_0(input, output):   _require_ops().dequantize_q8_0(input, output)
+def torch_dequantize_q4_K(input, output):   _require_ops().dequantize_q4_K(input, output)
+def torch_dequantize_q5_K(input, output):   _require_ops().dequantize_q5_K(input, output)
+def torch_dequantize_q6_K(input, output):   _require_ops().dequantize_q6_K(input, output)
+def torch_dequantize_iq4_nl(input, output): _require_ops().dequantize_iq4_nl(input, output)
+def torch_dequantize_iq4_xs(input, output): _require_ops().dequantize_iq4_xs(input, output)
+
+
+# =============================================================================
+# Q4_K × fp32 matmul — `GgufLinear` ceiling experiment.
+# =============================================================================
+# Layout convention (matches llama.cpp's kernel_mul_mm):
+#   A (Q4_K weights):  shape (M, K)  M = out_features, K = in_features.
+#                      M and K must be multiples of 64 and 256 respectively for
+#                      the simdgroup tile geometry to land cleanly.
+#   B (fp32 input):    shape (N, K)  N = batch * seqlen, contiguous in K.
+#   C (fp32 output):   shape (M, N)  column-major-ish (matches llama.cpp).
+
+_BLOCK_SIZE_M = 64
+_BLOCK_SIZE_N = 32
+_THREAD_PER_BLOCK = 128
+_SHARED_BYTES = 8192
+
+
+class MatMulQ4KPlan:
+    __slots__ = ("M", "N", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, M, N, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.M = M
+        self.N = N
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def allocate_matmul_q4_K_plan(M: int, N: int, K: int) -> MatMulQ4KPlan:
+    """Pre-allocate weight (Q4_K), input (fp32), output (fp32) MTLBuffers."""
+    import Metal
+
+    _ensure_metal()
+    assert M % _BLOCK_SIZE_M == 0, f"M must be multiple of {_BLOCK_SIZE_M}, got {M}"
+    assert N % _BLOCK_SIZE_N == 0, f"N must be multiple of {_BLOCK_SIZE_N}, got {N}"
+    assert K % 256 == 0, f"K must be multiple of 256, got {K}"
+
+    a_bytes = M * (K // 256) * 144  # M rows × (K/QK_K) super-blocks × 144 bytes
+    b_bytes = N * K * 4             # fp32
+    c_bytes = M * N * 4             # fp32
+
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+
+    # Pack ne00, ne01, ne11, nb01, nb11 into a uniforms buffer.
+    nb01 = (K // 256) * 144  # bytes per A row
+    nb11 = K * 4             # bytes per B row
+    dims = np.array([K], dtype=np.int64).tobytes()  # ne00
+    dims += np.array([M], dtype=np.int64).tobytes()  # ne01
+    dims += np.array([N], dtype=np.int64).tobytes()  # ne11
+    dims += np.array([nb01], dtype=np.uint64).tobytes()
+    dims += np.array([nb11], dtype=np.uint64).tobytes()
+
+    return MatMulQ4KPlan(M=M, N=N, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def dispatch_matmul_q4_K(plan: MatMulQ4KPlan, sync: bool = True):
+    """Dispatch mul_mat_q4_K_f32 on the pre-allocated buffers."""
+    import Metal
+
+    pso = _pipeline("mul_mat_q4_K_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    # ne00 (K), ne01 (M), ne11 (N), nb01, nb11 via setBytes — 5 scalars
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8],   8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16],  8, 4)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[16:24], 8, 5)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[24:32], 8, 6)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[32:40], 8, 7)
+    enc.setThreadgroupMemoryLength_atIndex_(_SHARED_BYTES, 0)
+
+    grid_x = (plan.N + _BLOCK_SIZE_N - 1) // _BLOCK_SIZE_N
+    grid_y = (plan.M + _BLOCK_SIZE_M - 1) // _BLOCK_SIZE_M
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, grid_y, 1),
+        Metal.MTLSizeMake(_THREAD_PER_BLOCK, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+# =============================================================================
+# mul_mat_vec_q4_K — batch-1 decode (memory-bandwidth-bound regime).
+# =============================================================================
+
+_N_DST = 4
+_SIMD_WIDTH = 32
+
+
+class MatVecQ4KPlan:
+    __slots__ = ("M", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, M, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.M = M
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def allocate_matvec_q4_K_plan(M: int, K: int) -> MatVecQ4KPlan:
+    import Metal
+
+    _ensure_metal()
+    assert M % _N_DST == 0, f"M must be a multiple of {_N_DST}, got {M}"
+    assert K % 256 == 0, f"K must be a multiple of 256, got {K}"
+
+    a_bytes = M * (K // 256) * 144
+    b_bytes = K * 4
+    c_bytes = M * 4
+
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+
+    dims = np.array([K], dtype=np.int64).tobytes() + np.array([M], dtype=np.int64).tobytes()
+    return MatVecQ4KPlan(M=M, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def dispatch_matvec_q4_K(plan: MatVecQ4KPlan, sync: bool = True):
+    import Metal
+
+    pso = _pipeline("mul_mat_vec_q4_K_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8], 8, 3)   # ne00 = K
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16], 8, 4)  # ne01 = M
+
+    grid_x = (plan.M + _N_DST - 1) // _N_DST
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, 1, 1),
+        Metal.MTLSizeMake(_SIMD_WIDTH, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+# =============================================================================
+# Q4_0 / Q8_0 matmul + matvec — same layout as Q4_K, smaller blocks (32 elems).
+# =============================================================================
+
+_N_SIMDGROUP = 2  # Q4_0/Q8_0 matvec uses 2 simdgroups per threadgroup
+
+# Per quant type for the mat path: (block bytes, block elems, matvec_M_alignment).
+# block_elems = 32 for Q4_0/Q8_0/IQ4_NL, 256 for all K-quants and IQ4_XS.
+# matvec_M_alignment:
+#   Q4_0/Q8_0  →  2 simdgroups × N_DST(4) = 8 rows per threadgroup
+#   Q4_K       →  1 simdgroup  × N_DST(4) = 4 rows per threadgroup
+#   Q5_K       →  2 simdgroups × 2 rows   = 4 rows per threadgroup
+#   Q6_K       →  2 simdgroups × 1 row    = 2 rows per threadgroup
+_MAT_FORMATS = {
+    "q4_0":   (18,  32,  8),
+    "q8_0":   (34,  32,  8),
+    "q4_K":   (144, 256, 4),
+    "q5_K":   (176, 256, 4),
+    "q6_K":   (210, 256, 2),
+    "iq4_nl": (18,  32,  4),       # 2 sg × 2 rows/sg = 4 rows/threadgroup
+    "iq4_xs": (136, 256, 4),
+}
+
+
+class MatMulPlan:
+    """Pre-allocated I/O buffers for `mul_mat_<fmt>_f32`. Block element size
+    depends on the format (32 for Q4_0/Q8_0, 256 for Q4_K)."""
+
+    __slots__ = ("fmt", "M", "N", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, fmt, M, N, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.fmt = fmt
+        self.M = M
+        self.N = N
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def _quant_block_elems(fmt: str) -> int:
+    return _MAT_FORMATS[fmt][1]
+
+
+def _quant_block_bytes(fmt: str) -> int:
+    return _MAT_FORMATS[fmt][0]
+
+
+def _matvec_row_alignment(fmt: str) -> int | None:
+    return _MAT_FORMATS[fmt][2]
+
+
+def allocate_matmul_plan(fmt: str, M: int, N: int, K: int) -> MatMulPlan:
+    import Metal
+
+    _ensure_metal()
+    assert fmt in _MAT_FORMATS, f"unsupported fmt {fmt}"
+    block_elems = _quant_block_elems(fmt)
+    block_bytes = _quant_block_bytes(fmt)
+    assert M % _BLOCK_SIZE_M == 0, f"M must be multiple of {_BLOCK_SIZE_M}, got {M}"
+    assert N % _BLOCK_SIZE_N == 0, f"N must be multiple of {_BLOCK_SIZE_N}, got {N}"
+    assert K % block_elems == 0, f"K must be multiple of {block_elems} for {fmt}, got {K}"
+
+    a_bytes = M * (K // block_elems) * block_bytes
+    b_bytes = N * K * 4
+    c_bytes = M * N * 4
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+
+    nb01 = (K // block_elems) * block_bytes
+    nb11 = K * 4
+    dims = (
+        np.array([K], dtype=np.int64).tobytes()
+        + np.array([M], dtype=np.int64).tobytes()
+        + np.array([N], dtype=np.int64).tobytes()
+        + np.array([nb01], dtype=np.uint64).tobytes()
+        + np.array([nb11], dtype=np.uint64).tobytes()
+    )
+    return MatMulPlan(fmt=fmt, M=M, N=N, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def dispatch_matmul(plan: MatMulPlan, sync: bool = True):
+    import Metal
+
+    pso = _pipeline(f"mul_mat_{plan.fmt}_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8],   8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16],  8, 4)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[16:24], 8, 5)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[24:32], 8, 6)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[32:40], 8, 7)
+    enc.setThreadgroupMemoryLength_atIndex_(_SHARED_BYTES, 0)
+
+    grid_x = (plan.N + _BLOCK_SIZE_N - 1) // _BLOCK_SIZE_N
+    grid_y = (plan.M + _BLOCK_SIZE_M - 1) // _BLOCK_SIZE_M
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, grid_y, 1),
+        Metal.MTLSizeMake(_THREAD_PER_BLOCK, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+class MatVecPlan:
+    __slots__ = ("fmt", "M", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, fmt, M, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.fmt = fmt
+        self.M = M
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def allocate_matvec_plan(fmt: str, M: int, K: int) -> MatVecPlan:
+    """Allocate I/O for `mul_mat_vec_<fmt>_f32`. Format-specific row alignment:
+    Q4_0/Q8_0: 8, Q4_K: 4, Q5_K: 4, Q6_K: 2."""
+    import Metal
+
+    _ensure_metal()
+    block_elems = _quant_block_elems(fmt)
+    block_bytes = _quant_block_bytes(fmt)
+    align = _matvec_row_alignment(fmt)
+    if align is None:
+        raise ValueError(f"mul_mat_vec not implemented for {fmt}")
+    assert M % align == 0, f"M must be a multiple of {align} for {fmt}, got {M}"
+    assert K % block_elems == 0, f"K must be a multiple of {block_elems} for {fmt}, got {K}"
+
+    a_bytes = M * (K // block_elems) * block_bytes
+    b_bytes = K * 4
+    c_bytes = M * 4
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+    dims = np.array([K], dtype=np.int64).tobytes() + np.array([M], dtype=np.int64).tobytes()
+    return MatVecPlan(fmt=fmt, M=M, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def encode_matvec(plan: MatVecPlan, enc):
+    """Encode a single matvec dispatch onto an existing MTLComputeCommandEncoder.
+    Caller is responsible for committing the parent command buffer. Threadgroup
+    memory for the IQ4 lookup table is set via setThreadgroupMemoryLength here."""
+    import Metal
+
+    pso = _pipeline(f"mul_mat_vec_{plan.fmt}_f32")
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8], 8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16], 8, 4)
+
+    if plan.fmt == "q4_K":
+        rows_per_tg = _N_DST
+        tg_threads = _SIMD_WIDTH
+    elif plan.fmt in ("q4_0", "q8_0"):
+        rows_per_tg = _N_SIMDGROUP * _N_DST
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q5_K":
+        rows_per_tg = _N_SIMDGROUP * 2
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q6_K":
+        rows_per_tg = _N_SIMDGROUP * 1
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt in ("iq4_nl", "iq4_xs"):
+        rows_per_tg = _N_SIMDGROUP * 2
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+        enc.setThreadgroupMemoryLength_atIndex_(256, 0)
+    else:
+        raise ValueError(f"matvec not wired for {plan.fmt}")
+    grid_x = (plan.M + rows_per_tg - 1) // rows_per_tg
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, 1, 1),
+        Metal.MTLSizeMake(tg_threads, 1, 1),
+    )
+
+
+def new_command_buffer():
+    """Start a fresh MTLCommandBuffer. Caller batches `encode_matvec(plan, enc)`
+    onto its encoder, then calls `cb.commit()` and `cb.waitUntilCompleted()`
+    once for the whole batch — matching llama.cpp's per-step semantics."""
+    _ensure_metal()
+    return _COMMAND_QUEUE.commandBuffer()
+
+
+def dispatch_matvec(plan: MatVecPlan, sync: bool = True):
+    import Metal
+
+    pso = _pipeline(f"mul_mat_vec_{plan.fmt}_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8], 8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16], 8, 4)
+
+    needs_shared = plan.fmt in ("iq4_nl", "iq4_xs")
+    if plan.fmt == "q4_K":
+        rows_per_tg = _N_DST
+        tg_threads = _SIMD_WIDTH
+    elif plan.fmt in ("q4_0", "q8_0"):
+        rows_per_tg = _N_SIMDGROUP * _N_DST
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q5_K":
+        rows_per_tg = _N_SIMDGROUP * 2
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q6_K":
+        rows_per_tg = _N_SIMDGROUP * 1
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt in ("iq4_nl", "iq4_xs"):
+        rows_per_tg = _N_SIMDGROUP * 2            # 2 sg × 2 rows/sg
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    else:
+        raise ValueError(f"matvec not wired for {plan.fmt}")
+    grid_x = (plan.M + rows_per_tg - 1) // rows_per_tg
+    if needs_shared:
+        enc.setThreadgroupMemoryLength_atIndex_(256, 0)  # 64 floats per simdgroup × 2 sg
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, 1, 1),
+        Metal.MTLSizeMake(tg_threads, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd

build/torch210-metal-aarch64-darwin/_gguf_dequant_530e374.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e502779f760fa35dac6a8aea0c4e63e6dbac0641cb3daf0ae34e4115754899d8
+size 237040

build/torch210-metal-aarch64-darwin/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _gguf_dequant_530e374
+ops = torch.ops._gguf_dequant_530e374
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_gguf_dequant_530e374::{op_name}"

build/torch210-metal-aarch64-darwin/gguf_dequant/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+from pathlib import Path
+from types import ModuleType
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch210-metal-aarch64-darwin/metadata.json

@@ -0,0 +1,4 @@
+{
+  "version": 1,
+  "python-depends": []
+}

build/torch210-metal-aarch64-darwin/nn.py

@@ -0,0 +1,191 @@
+"""GgufLinear — drop-in nn.Linear replacement that runs quantized matmul directly.
+
+Forward pass:
+  - If the input has effective batch size 1 (single token), use mul_mat_vec
+    (memory-bandwidth-bound, faster than bf16 mat-vec on MPS).
+  - Otherwise use mul_mat (compute-bound, similar throughput to bf16 GEMM but
+    keeps weights at native quantization → 3–4× less RAM than bf16).
+
+Supported quant types today: Q4_0, Q8_0, Q4_K. Adding more is just a matter of
+porting the candle kernels — the dispatch logic is generic over `quant_type`.
+
+Prototype design: weights are stored as raw uint8 bytes in `qweight`. The
+forward path copies activations into a shared MTLBuffer, runs the Metal
+kernel, reads the fp32 output back. Production path (kernels-community +
+PR #767) skips the copies via direct MTLBuffer handoff.
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+import torch
+
+from . import (
+    allocate_matmul_plan,
+    allocate_matvec_plan,
+    dispatch_matmul,
+    dispatch_matvec,
+)
+
+
+# Per-quant-type geometry: (block_bytes, block_elems, matvec_M_alignment)
+_QUANT_INFO = {
+    "Q4_0": (18, 32, 8),   # matvec uses 2 simdgroups × N_DST=4 = 8 rows/threadgroup
+    "Q8_0": (34, 32, 8),
+    "Q4_K": (144, 256, 4), # matvec uses 1 simdgroup × N_DST=4 = 4 rows/threadgroup
+}
+
+# Map our Python-side quant name to the Metal kernel suffix.
+_KERNEL_FMT = {
+    "Q4_0": "q4_0",
+    "Q8_0": "q8_0",
+    "Q4_K": "q4_K",
+}
+
+
+class GgufLinear(torch.nn.Module):
+    """Linear layer backed by GGUF-quantized weights.
+
+    Args:
+        in_features: K dimension. Must be a multiple of the format's block size
+            (32 for Q4_0/Q8_0, 256 for Q4_K).
+        out_features: M dimension. Must be a multiple of the format's matvec
+            row alignment (8 for Q4_0/Q8_0, 4 for Q4_K) and of 64 for matmul.
+        quant_type: one of "Q4_0", "Q8_0", "Q4_K".
+        bias: whether to add a learnable bias after the matmul.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        quant_type: str = "Q4_K",
+        bias: bool = False,
+        device: Optional[torch.device | str] = None,
+    ):
+        super().__init__()
+        if quant_type not in _QUANT_INFO:
+            raise NotImplementedError(
+                f"GgufLinear supports {sorted(_QUANT_INFO)} today, got {quant_type}"
+            )
+        block_bytes, block_elems, matvec_align = _QUANT_INFO[quant_type]
+        if in_features % block_elems != 0:
+            raise ValueError(
+                f"in_features must be a multiple of {block_elems} for {quant_type}, got {in_features}"
+            )
+        if out_features % matvec_align != 0:
+            raise ValueError(
+                f"out_features must be a multiple of {matvec_align} for {quant_type}, got {out_features}"
+            )
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.quant_type = quant_type
+        self._fmt = _KERNEL_FMT[quant_type]
+        self._block_bytes = block_bytes
+        self._block_elems = block_elems
+
+        nblocks_per_row = in_features // block_elems
+        qweight_nbytes = out_features * nblocks_per_row * block_bytes
+        self.register_buffer(
+            "qweight",
+            torch.empty(qweight_nbytes, dtype=torch.uint8, device=device or "cpu"),
+            persistent=True,
+        )
+
+        if bias:
+            self.bias = torch.nn.Parameter(
+                torch.zeros(out_features, dtype=torch.float32, device=device or "cpu"),
+                requires_grad=False,
+            )
+        else:
+            self.register_parameter("bias", None)
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"quant_type={self.quant_type}, bias={self.bias is not None}"
+        )
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: shape (..., K). Returns (..., M) fp32."""
+        if x.shape[-1] != self.in_features:
+            raise ValueError(f"Last dim of x must be {self.in_features}, got {x.shape[-1]}")
+
+        batch_shape = x.shape[:-1]
+        x_flat = x.reshape(-1, self.in_features).contiguous().to(torch.float32)
+        N = x_flat.shape[0]
+
+        # CPU view of qweight for the bytes copy into shared MTLBuffers.
+        # Production path skips this once the kernels-community build is wired.
+        qweight_cpu = self.qweight.detach().to("cpu", copy=False)
+        x_cpu = x_flat.to("cpu", copy=False)
+
+        if N == 1:
+            out = self._forward_matvec(qweight_cpu, x_cpu.reshape(-1))
+        else:
+            out = self._forward_matmul(qweight_cpu, x_cpu)
+
+        out = out.reshape(*batch_shape, self.out_features)
+        if x.device != out.device:
+            out = out.to(x.device)
+        if self.bias is not None:
+            out = out + self.bias.to(out.device).to(out.dtype)
+        return out
+
+    def _forward_matvec(self, qweight_cpu: torch.Tensor, x_cpu: torch.Tensor) -> torch.Tensor:
+        M, K = self.out_features, self.in_features
+        plan = allocate_matvec_plan(self._fmt, M=M, K=K)
+        a_view = np.frombuffer(plan.a_buf.contents().as_buffer(qweight_cpu.numel()), dtype=np.uint8)
+        a_view[:] = qweight_cpu.numpy()
+        b_view = np.frombuffer(plan.b_buf.contents().as_buffer(K * 4), dtype=np.float32)
+        b_view[:] = x_cpu.numpy()
+        dispatch_matvec(plan, sync=True)
+        c_view = np.frombuffer(plan.c_buf.contents().as_buffer(M * 4), dtype=np.float32)
+        return torch.from_numpy(c_view.copy())
+
+    def _forward_matmul(self, qweight_cpu: torch.Tensor, x_cpu: torch.Tensor) -> torch.Tensor:
+        N, K = x_cpu.shape
+        M = self.out_features
+        plan = allocate_matmul_plan(self._fmt, M=M, N=N, K=K)
+        a_view = np.frombuffer(plan.a_buf.contents().as_buffer(qweight_cpu.numel()), dtype=np.uint8)
+        a_view[:] = qweight_cpu.numpy()
+        b_view = np.frombuffer(plan.b_buf.contents().as_buffer(N * K * 4), dtype=np.float32)
+        b_view[:] = x_cpu.numpy().reshape(-1)
+        dispatch_matmul(plan, sync=True)
+        c_view = np.frombuffer(plan.c_buf.contents().as_buffer(M * N * 4), dtype=np.float32)
+        # Kernel output layout is (N, M) in row-major; we want (N, M).
+        return torch.from_numpy(c_view.reshape(N, M).copy())
+
+    @classmethod
+    def from_quantized_bytes(
+        cls,
+        quant_type: str,
+        qweight_bytes: bytes,
+        in_features: int,
+        out_features: int,
+        bias: Optional[torch.Tensor] = None,
+        device: Optional[torch.device | str] = None,
+    ) -> "GgufLinear":
+        """Convenience for tests/loaders: hydrate from raw GGUF block bytes."""
+        layer = cls(in_features=in_features, out_features=out_features,
+                    quant_type=quant_type, bias=bias is not None, device=device)
+        expected = layer.qweight.numel()
+        if len(qweight_bytes) != expected:
+            raise ValueError(
+                f"qweight_bytes length {len(qweight_bytes)} != expected {expected} "
+                f"for {quant_type} M={out_features} K={in_features}"
+            )
+        layer.qweight.copy_(torch.frombuffer(bytearray(qweight_bytes), dtype=torch.uint8))
+        if bias is not None:
+            layer.bias.copy_(bias.to(torch.float32))
+        return layer
+
+    # Back-compat shim for the existing Q4_K-only test
+    @classmethod
+    def from_q4_K_bytes(cls, qweight_bytes, in_features, out_features, bias=None, device=None):
+        return cls.from_quantized_bytes(
+            "Q4_K", qweight_bytes, in_features, out_features, bias=bias, device=device
+        )

build/torch210-metal-aarch64-darwin/transformers_glue.py

@@ -0,0 +1,284 @@
+"""Glue for plugging `GgufLinear` into a transformers model.
+
+The flow is intentionally non-invasive: we read raw block bytes directly from
+the GGUF file via `gguf.GGUFReader`, walk the model, and for every `nn.Linear`
+whose corresponding GGUF tensor is one of our supported quant types we
+in-place replace the module with `GgufLinear` and copy the bytes into
+`qweight`. Tensors that don't map to a Linear (or are unsupported quants like
+F32 norms, embeddings) get dequantized via `gguf.dequantize` and copied into
+the existing parameter as fp32.
+
+This avoids modifying the transformers loader. The trade-off is that the
+model has to be constructed first with regular `nn.Linear` shells (which we
+then mutate), but those shells are on the meta device or freshly initialised
+and don't carry the dequantized weight at the time of swap — so peak memory
+during load is bounded by `(model_size_fp32) + (model_size_quantized_qweights)`,
+which fits in RAM for the small models we care about.
+
+Typical use:
+
+    from transformers import AutoConfig, AutoModelForCausalLM
+    cfg = AutoConfig.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct-GGUF",
+                                     gguf_file="qwen2.5-0.5b-instruct-q4_0.gguf")
+    model = AutoModelForCausalLM.from_config(cfg).eval()
+    load_gguf_into_model_with_gguf_linear(
+        model,
+        gguf_path="…/qwen2.5-0.5b-instruct-q4_0.gguf",
+    )
+    out = model.generate(input_ids=…, max_new_tokens=10)
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import gguf
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .nn import GgufLinear, _QUANT_INFO
+
+
+# Map gguf.GGMLQuantizationType.name → our GgufLinear quant_type string.
+_GGUF_NAME_TO_GGUF_LINEAR = {
+    "Q4_0": "Q4_0",
+    "Q8_0": "Q8_0",
+    "Q4_K": "Q4_K",
+}
+
+
+def _quant_supported_by_gguf_linear(qt: gguf.GGMLQuantizationType) -> bool:
+    return qt.name in _GGUF_NAME_TO_GGUF_LINEAR
+
+
+# --- Tensor-name mapping ------------------------------------------------------
+# GGUF tensors use llama.cpp-style names like `blk.0.attn_q.weight`. HF models
+# use `model.layers.0.self_attn.q_proj.weight`. We use a tiny version of the
+# transformers rename table here so the swap util can stay self-contained.
+# (Transformers' own GGUF loader has a much more thorough version of this.)
+
+_BASIC_RENAMES = [
+    (r"^blk\.", "model.layers."),
+    (r"\.attn_q\.weight$",      ".self_attn.q_proj.weight"),
+    (r"\.attn_k\.weight$",      ".self_attn.k_proj.weight"),
+    (r"\.attn_v\.weight$",      ".self_attn.v_proj.weight"),
+    (r"\.attn_output\.weight$", ".self_attn.o_proj.weight"),
+    (r"\.ffn_gate\.weight$",    ".mlp.gate_proj.weight"),
+    (r"\.ffn_up\.weight$",      ".mlp.up_proj.weight"),
+    (r"\.ffn_down\.weight$",    ".mlp.down_proj.weight"),
+    (r"\.attn_norm\.weight$",   ".input_layernorm.weight"),
+    (r"\.ffn_norm\.weight$",    ".post_attention_layernorm.weight"),
+    (r"\.attn_(q|k|v)\.bias$",  r".self_attn.\1_proj.bias"),
+    (r"^token_embd\.weight$",   "model.embed_tokens.weight"),
+    (r"^output_norm\.weight$",  "model.norm.weight"),
+    (r"^output\.weight$",       "lm_head.weight"),
+]
+
+
+def _gguf_name_to_hf_name(gguf_name: str) -> str:
+    """Apply our minimal rename table to map a GGUF tensor name to an HF parameter name."""
+    import re
+
+    out = gguf_name
+    for pat, repl in _BASIC_RENAMES:
+        out = re.sub(pat, repl, out)
+    return out
+
+
+def _is_attn_qk(gguf_name: str) -> bool:
+    return ".attn_q.weight" in gguf_name or ".attn_k.weight" in gguf_name
+
+
+def _reverse_permute_qk(w: np.ndarray, n_head: int, n_kv: int | None) -> np.ndarray:
+    """GGUF Q/K weights are permuted on quantize; undo before loading into HF.
+    Reference: https://github.com/ggerganov/llama.cpp/blob/.../convert_hf_to_gguf.py."""
+    if n_kv is not None and n_head != n_kv:
+        n_head = n_kv
+    dim = w.shape[0] // n_head // 2
+    w = w.reshape(n_head, dim, 2, *w.shape[1:])
+    return w.swapaxes(2, 1).reshape(*w.shape[:1], *w.shape[3:]) if False else w.reshape(-1, *w.shape[3:]).swapaxes(2, 1).reshape(n_head * dim * 2, *w.shape[3:])
+
+
+def _set_module_param(model: nn.Module, qualified_name: str, value: torch.Tensor) -> bool:
+    """Set a parameter or buffer on `model` by dotted name. Returns True on success."""
+    parts = qualified_name.split(".")
+    mod = model
+    for p in parts[:-1]:
+        if not hasattr(mod, p):
+            return False
+        mod = getattr(mod, p)
+    leaf = parts[-1]
+    if leaf in mod._parameters:
+        param = mod._parameters[leaf]
+        if param is None:
+            return False
+        with torch.no_grad():
+            mod._parameters[leaf] = nn.Parameter(value.to(dtype=param.dtype, device=param.device),
+                                                 requires_grad=param.requires_grad)
+        return True
+    if leaf in mod._buffers:
+        with torch.no_grad():
+            mod._buffers[leaf] = value
+        return True
+    return False
+
+
+def _get_module_by_name(model: nn.Module, qualified_name: str) -> Optional[nn.Module]:
+    """Return the leaf module matching `qualified_name` minus its terminal attribute,
+    or None if anything along the path is missing."""
+    parts = qualified_name.split(".")
+    mod = model
+    for p in parts[:-1]:
+        if not hasattr(mod, p):
+            return None
+        mod = getattr(mod, p)
+    return mod
+
+
+def _get_attr_by_path(model: nn.Module, qualified_name: str) -> Optional[nn.Module]:
+    """Walk dotted `qualified_name` on `model` and return the final attribute,
+    or None if anything along the path is missing."""
+    if not qualified_name:
+        return model
+    mod = model
+    for p in qualified_name.split("."):
+        if not hasattr(mod, p):
+            return None
+        mod = getattr(mod, p)
+    return mod
+
+
+def _swap_linear_to_gguf(
+    parent: nn.Module, attr: str, in_features: int, out_features: int, quant_type: str,
+    bias_tensor: Optional[torch.Tensor], device: Optional[torch.device | str] = None
+) -> GgufLinear:
+    """Replace `parent.attr` (an nn.Linear) with a new GgufLinear of the same shape."""
+    new = GgufLinear(in_features=in_features, out_features=out_features,
+                     quant_type=quant_type, bias=bias_tensor is not None, device=device)
+    if bias_tensor is not None and new.bias is not None:
+        new.bias.data.copy_(bias_tensor.to(torch.float32))
+    setattr(parent, attr, new)
+    return new
+
+
+def load_gguf_into_model_with_gguf_linear(
+    model: nn.Module,
+    gguf_path: str | Path,
+    skip_modules: Optional[set[str]] = None,
+    force_dequant: bool = False,
+) -> dict:
+    """Read a GGUF file and populate `model` in-place.
+
+    Two-phase load to keep the byte path simple and the bias path correct:
+
+      Phase 1 — full dequant: every tensor (linear weights, norms, embeddings,
+        biases) is dequantized to fp32 and written into the matching HF
+        parameter. After this pass the model would already produce correct
+        outputs (this is what `force_dequant=True` stops at).
+      Phase 2 — swap linears: walk every `nn.Linear` that maps to a tensor
+        whose GGUF quant type is supported by `GgufLinear`; swap it for a
+        `GgufLinear` and re-attach raw block bytes (preserving the bias that
+        Phase 1 just loaded).
+
+    Returns a dict with diagnostic counts.
+    """
+    reader = gguf.GGUFReader(str(gguf_path))
+
+    # Pull a few config bits for the Q/K reverse-permute.
+    def _read_int(field: str, default: int = 0) -> int:
+        f = reader.fields.get(field)
+        if f is None:
+            return default
+        return int(f.parts[f.data[0]][0]) if f.data else default
+
+    arch = ""
+    arch_field = reader.fields.get("general.architecture")
+    if arch_field is not None and arch_field.data:
+        arch = bytes(arch_field.parts[arch_field.data[0]]).decode("utf-8", errors="ignore")
+    n_head    = _read_int(f"{arch}.attention.head_count")
+    n_kv_head = _read_int(f"{arch}.attention.head_count_kv") or n_head
+
+    skip_modules = skip_modules or set()
+    stats = {"as_gguf_linear": 0, "dequantized": 0, "skipped": 0, "unmatched": []}
+
+    # Track which (hf_name → quant_type) pairs are eligible for GgufLinear swap
+    # after Phase 1. We'll re-quantize the loaded fp32 weight in Phase 2.
+    swap_candidates: list[tuple[str, str]] = []
+
+    # ---------------- Phase 1: dequant everything into the model. -------------
+    for t in reader.tensors:
+        hf_name = _gguf_name_to_hf_name(t.name)
+        if any(skip in hf_name for skip in skip_modules):
+            stats["skipped"] += 1
+            continue
+
+        parent = _get_module_by_name(model, hf_name)
+        if parent is None or not hasattr(parent, hf_name.rsplit(".", 1)[-1]):
+            stats["unmatched"].append(t.name)
+            continue
+        leaf = hf_name.rsplit(".", 1)[-1]
+        cur = getattr(parent, leaf)
+
+        fp32 = gguf.dequantize(
+            np.frombuffer(np.ascontiguousarray(t.data).tobytes(), dtype=np.uint8),
+            t.tensor_type,
+        ).astype(np.float32)
+        if _is_attn_qk(t.name) and n_head > 0:
+            shape = list(t.shape)
+            M, K = int(shape[1]), int(shape[0])
+            heads = n_head if ".attn_q." in t.name else n_kv_head
+            head_dim = M // heads
+            fp32 = fp32.reshape(M, K).reshape(heads, head_dim // 2, 2, K).swapaxes(2, 1).reshape(M, K)
+        target_shape = tuple(cur.shape) if hasattr(cur, "shape") else (fp32.size,)
+        if np.prod(target_shape) == fp32.size:
+            fp32 = fp32.reshape(target_shape)
+        tensor = torch.from_numpy(fp32)
+        if not _set_module_param(model, hf_name, tensor):
+            stats["unmatched"].append(t.name)
+            continue
+        stats["dequantized"] += 1
+
+        # Record candidacy for swap in Phase 2.
+        is_linear_weight = isinstance(parent, nn.Linear) and leaf == "weight"
+        if is_linear_weight and _quant_supported_by_gguf_linear(t.tensor_type) and not force_dequant:
+            swap_candidates.append((hf_name, _GGUF_NAME_TO_GGUF_LINEAR[t.tensor_type.name]))
+
+    # ---------------- Phase 2: swap Linear → GgufLinear. ----------------------
+    # Re-quantize the loaded fp32 weight via gguf-py (works for Q4_0/Q8_0;
+    # K-quants don't have a Python quantizer yet — those will skip swap).
+    for hf_name, quant_type in swap_candidates:
+        # hf_name ends in `.weight`; the Linear module is at hf_name[:-len(".weight")]
+        linear_path = hf_name[: -len(".weight")]
+        grandparent_path, _, linear_attr = linear_path.rpartition(".")
+        grandparent = model if not grandparent_path else _get_attr_by_path(model, grandparent_path)
+        linear = _get_attr_by_path(model, linear_path)
+        if grandparent is None or not isinstance(linear, nn.Linear):
+            continue
+
+        # Re-quantize the (now permuted-if-applicable) fp32 weight.
+        try:
+            qt = getattr(gguf.GGMLQuantizationType, quant_type)
+            qbytes = bytes(
+                gguf.quantize(linear.weight.detach().cpu().numpy().astype(np.float32), qt).tobytes()
+            )
+        except (NotImplementedError, Exception) as e:
+            # K-quants etc. — leave as Linear (fp32 weight is already loaded).
+            stats.setdefault("re_quantize_failed", []).append((hf_name, str(e)))
+            continue
+
+        bias_tensor = linear.bias.detach().clone() if linear.bias is not None else None
+        new = GgufLinear(
+            in_features=linear.in_features,
+            out_features=linear.out_features,
+            quant_type=quant_type,
+            bias=bias_tensor is not None,
+        )
+        if bias_tensor is not None:
+            new.bias.data.copy_(bias_tensor.to(torch.float32))
+        new.qweight.copy_(torch.frombuffer(bytearray(qbytes), dtype=torch.uint8))
+        setattr(grandparent, linear_attr, new)
+        stats["as_gguf_linear"] += 1
+        stats["dequantized"] -= 1  # this one's no longer "just" dequant — swapped
+
+    return stats

build/torch211-metal-aarch64-darwin/__init__.py

@@ -0,0 +1,600 @@
+"""GGUF dequantization kernels for Apple Silicon (Metal).
+
+Six standalone kernels covering the GGUF formats most common in the wild:
+
+    Q4_0   Q8_0   Q4_K   Q5_K   Q6_K   IQ4_NL   IQ4_XS
+
+Two API levels are exposed for each:
+
+  - ``dequantize_<format>_bytes(quant_bytes) -> np.ndarray``
+    One-shot path used by the tests for bit-exact validation against a
+    NumPy reference. Allocates input/output buffers internally.
+
+  - ``allocate_<format>_plan(nblocks)`` + ``dispatch_<format>(plan)``
+    Low-level path for benchmarks. Buffers are pre-allocated; ``dispatch``
+    only encodes + dispatches + syncs, so the timing isolates GPU work
+    from CPU<->buffer copies.
+
+The runtime path uses PyObjC + Metal.framework so the correctness loop
+works on any Apple Silicon Mac without Xcode 26 / Nix / kernel-builder.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+# Lazy globals — populated on first use.
+_DEVICE = None
+_COMMAND_QUEUE = None
+_LIBRARY = None
+_PIPELINES: dict = {}
+
+_KERNEL_SOURCE_PATH = Path(__file__).resolve().parent.parent.parent / "gguf_dequant_metal" / "dequantize.metal"
+
+
+def _ensure_metal():
+    global _DEVICE, _COMMAND_QUEUE, _LIBRARY
+    if _DEVICE is not None:
+        return
+    import Metal  # pyobjc-framework-metal
+
+    _DEVICE = Metal.MTLCreateSystemDefaultDevice()
+    if _DEVICE is None:
+        raise RuntimeError("No default Metal device available")
+    _COMMAND_QUEUE = _DEVICE.newCommandQueue()
+
+    src = _KERNEL_SOURCE_PATH.read_text()
+    options = Metal.MTLCompileOptions.new()
+    lib, err = _DEVICE.newLibraryWithSource_options_error_(src, options, None)
+    if lib is None:
+        raise RuntimeError(f"Metal compile failed: {err}")
+    _LIBRARY = lib
+
+
+def _pipeline(fn_name: str):
+    _ensure_metal()
+    pso = _PIPELINES.get(fn_name)
+    if pso is not None:
+        return pso
+
+    fn = _LIBRARY.newFunctionWithName_(fn_name)
+    if fn is None:
+        raise RuntimeError(f"Kernel function not found: {fn_name}")
+    pso, err = _DEVICE.newComputePipelineStateWithFunction_error_(fn, None)
+    if pso is None:
+        raise RuntimeError(f"Pipeline state creation failed: {err}")
+    _PIPELINES[fn_name] = pso
+    return pso
+
+
+# (kernel_name, bytes_per_block, elems_per_block, threads_per_block_group)
+_FORMATS = {
+    "q4_0":   ("dequantize_q4_0",   18, 32, 16),
+    "q8_0":   ("dequantize_q8_0",   34, 32, 32),
+    "q4_K":   ("dequantize_q4_K",  144, 256, 256),
+    "q5_K":   ("dequantize_q5_K",  176, 256, 256),
+    "q6_K":   ("dequantize_q6_K",  210, 256, 256),
+    "iq4_nl": ("dequantize_iq4_nl", 18, 32, 16),
+    "iq4_xs": ("dequantize_iq4_xs", 136, 256, 256),
+}
+
+
+class Plan:
+    """Pre-allocated I/O buffers for a single (format, nblocks) shape."""
+
+    __slots__ = ("fmt", "nblocks", "in_buf", "out_buf", "_nblocks_bytes")
+
+    def __init__(self, fmt: str, nblocks: int, in_buf, out_buf):
+        self.fmt = fmt
+        self.nblocks = nblocks
+        self.in_buf = in_buf
+        self.out_buf = out_buf
+        self._nblocks_bytes = np.array([nblocks], dtype=np.uint32).tobytes()
+
+
+def allocate_plan(fmt: str, nblocks: int) -> Plan:
+    import Metal
+
+    _ensure_metal()
+    _, bpb, epb, _ = _FORMATS[fmt]
+    in_buf = _DEVICE.newBufferWithLength_options_(nblocks * bpb, Metal.MTLResourceStorageModeShared)
+    out_buf = _DEVICE.newBufferWithLength_options_(nblocks * epb * 4, Metal.MTLResourceStorageModeShared)
+    return Plan(fmt=fmt, nblocks=nblocks, in_buf=in_buf, out_buf=out_buf)
+
+
+def dispatch(plan: Plan, sync: bool = True):
+    import Metal
+
+    kernel_name, _, _, tg = _FORMATS[plan.fmt]
+    pso = _pipeline(kernel_name)
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.in_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.out_buf, 0, 1)
+    enc.setBytes_length_atIndex_(plan._nblocks_bytes, len(plan._nblocks_bytes), 2)
+    enc.dispatchThreads_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(plan.nblocks * tg, 1, 1),
+        Metal.MTLSizeMake(tg, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+def dequantize_bytes(fmt: str, quant_bytes) -> np.ndarray:
+    """One-shot bytes-in / float32-out for the named format."""
+    _ensure_metal()
+    _, bpb, epb, _ = _FORMATS[fmt]
+
+    buf = bytes(quant_bytes) if not isinstance(quant_bytes, (bytes, bytearray)) else quant_bytes
+    nbytes = len(buf)
+    if nbytes % bpb != 0:
+        raise ValueError(f"{fmt} input not a multiple of {bpb} bytes: got {nbytes}")
+    nblocks = nbytes // bpb
+
+    plan = allocate_plan(fmt, nblocks)
+    in_view = np.frombuffer(plan.in_buf.contents().as_buffer(nbytes), dtype=np.uint8)
+    in_view[:] = np.frombuffer(buf, dtype=np.uint8)
+    dispatch(plan, sync=True)
+    out_nbytes = nblocks * epb * 4
+    return np.frombuffer(plan.out_buf.contents().as_buffer(out_nbytes), dtype=np.float32).copy()
+
+
+# --- Per-format convenience wrappers ----------------------------------------
+
+def dequantize_q4_0_bytes(qb):   return dequantize_bytes("q4_0",   qb)
+def dequantize_q8_0_bytes(qb):   return dequantize_bytes("q8_0",   qb)
+def dequantize_q4_K_bytes(qb):   return dequantize_bytes("q4_K",   qb)
+def dequantize_q5_K_bytes(qb):   return dequantize_bytes("q5_K",   qb)
+def dequantize_q6_K_bytes(qb):   return dequantize_bytes("q6_K",   qb)
+def dequantize_iq4_nl_bytes(qb): return dequantize_bytes("iq4_nl", qb)
+def dequantize_iq4_xs_bytes(qb): return dequantize_bytes("iq4_xs", qb)
+
+def allocate_q4_0_plan(nblocks):   return allocate_plan("q4_0",   nblocks)
+def allocate_q8_0_plan(nblocks):   return allocate_plan("q8_0",   nblocks)
+def allocate_q4_K_plan(nblocks):   return allocate_plan("q4_K",   nblocks)
+def allocate_q5_K_plan(nblocks):   return allocate_plan("q5_K",   nblocks)
+def allocate_q6_K_plan(nblocks):   return allocate_plan("q6_K",   nblocks)
+def allocate_iq4_nl_plan(nblocks): return allocate_plan("iq4_nl", nblocks)
+def allocate_iq4_xs_plan(nblocks): return allocate_plan("iq4_xs", nblocks)
+
+def dispatch_q4_0(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q8_0(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q4_K(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q5_K(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q6_K(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_iq4_nl(plan, **kw): return dispatch(plan, **kw)
+def dispatch_iq4_xs(plan, **kw): return dispatch(plan, **kw)
+
+
+# ---------------------------------------------------------------------------
+# Torch op entry points — the production path used by transformers (no PyObjC,
+# no runtime compile). Only available in builds produced by kernel-builder,
+# which generates `_ops.py` exporting `ops = torch.ops._gguf_dequant_<rev>`.
+# In a dev-mode checkout (running from torch-ext/ directly) `_ops` isn't
+# present, so the torch_* entry points raise on call.
+# ---------------------------------------------------------------------------
+
+try:
+    from ._ops import ops as _ops  # generated by kernel-builder
+except ImportError:
+    _ops = None
+
+
+def _require_ops():
+    if _ops is None:
+        raise RuntimeError(
+            "torch op binding not available — this build was loaded from the dev tree. "
+            "Use the *_bytes / *_plan APIs (PyObjC path) or install a kernel-builder build."
+        )
+    return _ops
+
+
+def torch_dequantize_q4_0(input, output):   _require_ops().dequantize_q4_0(input, output)
+def torch_dequantize_q8_0(input, output):   _require_ops().dequantize_q8_0(input, output)
+def torch_dequantize_q4_K(input, output):   _require_ops().dequantize_q4_K(input, output)
+def torch_dequantize_q5_K(input, output):   _require_ops().dequantize_q5_K(input, output)
+def torch_dequantize_q6_K(input, output):   _require_ops().dequantize_q6_K(input, output)
+def torch_dequantize_iq4_nl(input, output): _require_ops().dequantize_iq4_nl(input, output)
+def torch_dequantize_iq4_xs(input, output): _require_ops().dequantize_iq4_xs(input, output)
+
+
+# =============================================================================
+# Q4_K × fp32 matmul — `GgufLinear` ceiling experiment.
+# =============================================================================
+# Layout convention (matches llama.cpp's kernel_mul_mm):
+#   A (Q4_K weights):  shape (M, K)  M = out_features, K = in_features.
+#                      M and K must be multiples of 64 and 256 respectively for
+#                      the simdgroup tile geometry to land cleanly.
+#   B (fp32 input):    shape (N, K)  N = batch * seqlen, contiguous in K.
+#   C (fp32 output):   shape (M, N)  column-major-ish (matches llama.cpp).
+
+_BLOCK_SIZE_M = 64
+_BLOCK_SIZE_N = 32
+_THREAD_PER_BLOCK = 128
+_SHARED_BYTES = 8192
+
+
+class MatMulQ4KPlan:
+    __slots__ = ("M", "N", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, M, N, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.M = M
+        self.N = N
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def allocate_matmul_q4_K_plan(M: int, N: int, K: int) -> MatMulQ4KPlan:
+    """Pre-allocate weight (Q4_K), input (fp32), output (fp32) MTLBuffers."""
+    import Metal
+
+    _ensure_metal()
+    assert M % _BLOCK_SIZE_M == 0, f"M must be multiple of {_BLOCK_SIZE_M}, got {M}"
+    assert N % _BLOCK_SIZE_N == 0, f"N must be multiple of {_BLOCK_SIZE_N}, got {N}"
+    assert K % 256 == 0, f"K must be multiple of 256, got {K}"
+
+    a_bytes = M * (K // 256) * 144  # M rows × (K/QK_K) super-blocks × 144 bytes
+    b_bytes = N * K * 4             # fp32
+    c_bytes = M * N * 4             # fp32
+
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+
+    # Pack ne00, ne01, ne11, nb01, nb11 into a uniforms buffer.
+    nb01 = (K // 256) * 144  # bytes per A row
+    nb11 = K * 4             # bytes per B row
+    dims = np.array([K], dtype=np.int64).tobytes()  # ne00
+    dims += np.array([M], dtype=np.int64).tobytes()  # ne01
+    dims += np.array([N], dtype=np.int64).tobytes()  # ne11
+    dims += np.array([nb01], dtype=np.uint64).tobytes()
+    dims += np.array([nb11], dtype=np.uint64).tobytes()
+
+    return MatMulQ4KPlan(M=M, N=N, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def dispatch_matmul_q4_K(plan: MatMulQ4KPlan, sync: bool = True):
+    """Dispatch mul_mat_q4_K_f32 on the pre-allocated buffers."""
+    import Metal
+
+    pso = _pipeline("mul_mat_q4_K_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    # ne00 (K), ne01 (M), ne11 (N), nb01, nb11 via setBytes — 5 scalars
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8],   8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16],  8, 4)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[16:24], 8, 5)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[24:32], 8, 6)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[32:40], 8, 7)
+    enc.setThreadgroupMemoryLength_atIndex_(_SHARED_BYTES, 0)
+
+    grid_x = (plan.N + _BLOCK_SIZE_N - 1) // _BLOCK_SIZE_N
+    grid_y = (plan.M + _BLOCK_SIZE_M - 1) // _BLOCK_SIZE_M
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, grid_y, 1),
+        Metal.MTLSizeMake(_THREAD_PER_BLOCK, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+# =============================================================================
+# mul_mat_vec_q4_K — batch-1 decode (memory-bandwidth-bound regime).
+# =============================================================================
+
+_N_DST = 4
+_SIMD_WIDTH = 32
+
+
+class MatVecQ4KPlan:
+    __slots__ = ("M", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, M, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.M = M
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def allocate_matvec_q4_K_plan(M: int, K: int) -> MatVecQ4KPlan:
+    import Metal
+
+    _ensure_metal()
+    assert M % _N_DST == 0, f"M must be a multiple of {_N_DST}, got {M}"
+    assert K % 256 == 0, f"K must be a multiple of 256, got {K}"
+
+    a_bytes = M * (K // 256) * 144
+    b_bytes = K * 4
+    c_bytes = M * 4
+
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+
+    dims = np.array([K], dtype=np.int64).tobytes() + np.array([M], dtype=np.int64).tobytes()
+    return MatVecQ4KPlan(M=M, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def dispatch_matvec_q4_K(plan: MatVecQ4KPlan, sync: bool = True):
+    import Metal
+
+    pso = _pipeline("mul_mat_vec_q4_K_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8], 8, 3)   # ne00 = K
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16], 8, 4)  # ne01 = M
+
+    grid_x = (plan.M + _N_DST - 1) // _N_DST
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, 1, 1),
+        Metal.MTLSizeMake(_SIMD_WIDTH, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+# =============================================================================
+# Q4_0 / Q8_0 matmul + matvec — same layout as Q4_K, smaller blocks (32 elems).
+# =============================================================================
+
+_N_SIMDGROUP = 2  # Q4_0/Q8_0 matvec uses 2 simdgroups per threadgroup
+
+# Per quant type for the mat path: (block bytes, block elems, matvec_M_alignment).
+# block_elems = 32 for Q4_0/Q8_0/IQ4_NL, 256 for all K-quants and IQ4_XS.
+# matvec_M_alignment:
+#   Q4_0/Q8_0  →  2 simdgroups × N_DST(4) = 8 rows per threadgroup
+#   Q4_K       →  1 simdgroup  × N_DST(4) = 4 rows per threadgroup
+#   Q5_K       →  2 simdgroups × 2 rows   = 4 rows per threadgroup
+#   Q6_K       →  2 simdgroups × 1 row    = 2 rows per threadgroup
+_MAT_FORMATS = {
+    "q4_0":   (18,  32,  8),
+    "q8_0":   (34,  32,  8),
+    "q4_K":   (144, 256, 4),
+    "q5_K":   (176, 256, 4),
+    "q6_K":   (210, 256, 2),
+    "iq4_nl": (18,  32,  4),       # 2 sg × 2 rows/sg = 4 rows/threadgroup
+    "iq4_xs": (136, 256, 4),
+}
+
+
+class MatMulPlan:
+    """Pre-allocated I/O buffers for `mul_mat_<fmt>_f32`. Block element size
+    depends on the format (32 for Q4_0/Q8_0, 256 for Q4_K)."""
+
+    __slots__ = ("fmt", "M", "N", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, fmt, M, N, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.fmt = fmt
+        self.M = M
+        self.N = N
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def _quant_block_elems(fmt: str) -> int:
+    return _MAT_FORMATS[fmt][1]
+
+
+def _quant_block_bytes(fmt: str) -> int:
+    return _MAT_FORMATS[fmt][0]
+
+
+def _matvec_row_alignment(fmt: str) -> int | None:
+    return _MAT_FORMATS[fmt][2]
+
+
+def allocate_matmul_plan(fmt: str, M: int, N: int, K: int) -> MatMulPlan:
+    import Metal
+
+    _ensure_metal()
+    assert fmt in _MAT_FORMATS, f"unsupported fmt {fmt}"
+    block_elems = _quant_block_elems(fmt)
+    block_bytes = _quant_block_bytes(fmt)
+    assert M % _BLOCK_SIZE_M == 0, f"M must be multiple of {_BLOCK_SIZE_M}, got {M}"
+    assert N % _BLOCK_SIZE_N == 0, f"N must be multiple of {_BLOCK_SIZE_N}, got {N}"
+    assert K % block_elems == 0, f"K must be multiple of {block_elems} for {fmt}, got {K}"
+
+    a_bytes = M * (K // block_elems) * block_bytes
+    b_bytes = N * K * 4
+    c_bytes = M * N * 4
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+
+    nb01 = (K // block_elems) * block_bytes
+    nb11 = K * 4
+    dims = (
+        np.array([K], dtype=np.int64).tobytes()
+        + np.array([M], dtype=np.int64).tobytes()
+        + np.array([N], dtype=np.int64).tobytes()
+        + np.array([nb01], dtype=np.uint64).tobytes()
+        + np.array([nb11], dtype=np.uint64).tobytes()
+    )
+    return MatMulPlan(fmt=fmt, M=M, N=N, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def dispatch_matmul(plan: MatMulPlan, sync: bool = True):
+    import Metal
+
+    pso = _pipeline(f"mul_mat_{plan.fmt}_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8],   8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16],  8, 4)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[16:24], 8, 5)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[24:32], 8, 6)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[32:40], 8, 7)
+    enc.setThreadgroupMemoryLength_atIndex_(_SHARED_BYTES, 0)
+
+    grid_x = (plan.N + _BLOCK_SIZE_N - 1) // _BLOCK_SIZE_N
+    grid_y = (plan.M + _BLOCK_SIZE_M - 1) // _BLOCK_SIZE_M
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, grid_y, 1),
+        Metal.MTLSizeMake(_THREAD_PER_BLOCK, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+class MatVecPlan:
+    __slots__ = ("fmt", "M", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, fmt, M, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.fmt = fmt
+        self.M = M
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def allocate_matvec_plan(fmt: str, M: int, K: int) -> MatVecPlan:
+    """Allocate I/O for `mul_mat_vec_<fmt>_f32`. Format-specific row alignment:
+    Q4_0/Q8_0: 8, Q4_K: 4, Q5_K: 4, Q6_K: 2."""
+    import Metal
+
+    _ensure_metal()
+    block_elems = _quant_block_elems(fmt)
+    block_bytes = _quant_block_bytes(fmt)
+    align = _matvec_row_alignment(fmt)
+    if align is None:
+        raise ValueError(f"mul_mat_vec not implemented for {fmt}")
+    assert M % align == 0, f"M must be a multiple of {align} for {fmt}, got {M}"
+    assert K % block_elems == 0, f"K must be a multiple of {block_elems} for {fmt}, got {K}"
+
+    a_bytes = M * (K // block_elems) * block_bytes
+    b_bytes = K * 4
+    c_bytes = M * 4
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+    dims = np.array([K], dtype=np.int64).tobytes() + np.array([M], dtype=np.int64).tobytes()
+    return MatVecPlan(fmt=fmt, M=M, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def encode_matvec(plan: MatVecPlan, enc):
+    """Encode a single matvec dispatch onto an existing MTLComputeCommandEncoder.
+    Caller is responsible for committing the parent command buffer. Threadgroup
+    memory for the IQ4 lookup table is set via setThreadgroupMemoryLength here."""
+    import Metal
+
+    pso = _pipeline(f"mul_mat_vec_{plan.fmt}_f32")
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8], 8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16], 8, 4)
+
+    if plan.fmt == "q4_K":
+        rows_per_tg = _N_DST
+        tg_threads = _SIMD_WIDTH
+    elif plan.fmt in ("q4_0", "q8_0"):
+        rows_per_tg = _N_SIMDGROUP * _N_DST
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q5_K":
+        rows_per_tg = _N_SIMDGROUP * 2
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q6_K":
+        rows_per_tg = _N_SIMDGROUP * 1
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt in ("iq4_nl", "iq4_xs"):
+        rows_per_tg = _N_SIMDGROUP * 2
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+        enc.setThreadgroupMemoryLength_atIndex_(256, 0)
+    else:
+        raise ValueError(f"matvec not wired for {plan.fmt}")
+    grid_x = (plan.M + rows_per_tg - 1) // rows_per_tg
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, 1, 1),
+        Metal.MTLSizeMake(tg_threads, 1, 1),
+    )
+
+
+def new_command_buffer():
+    """Start a fresh MTLCommandBuffer. Caller batches `encode_matvec(plan, enc)`
+    onto its encoder, then calls `cb.commit()` and `cb.waitUntilCompleted()`
+    once for the whole batch — matching llama.cpp's per-step semantics."""
+    _ensure_metal()
+    return _COMMAND_QUEUE.commandBuffer()
+
+
+def dispatch_matvec(plan: MatVecPlan, sync: bool = True):
+    import Metal
+
+    pso = _pipeline(f"mul_mat_vec_{plan.fmt}_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8], 8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16], 8, 4)
+
+    needs_shared = plan.fmt in ("iq4_nl", "iq4_xs")
+    if plan.fmt == "q4_K":
+        rows_per_tg = _N_DST
+        tg_threads = _SIMD_WIDTH
+    elif plan.fmt in ("q4_0", "q8_0"):
+        rows_per_tg = _N_SIMDGROUP * _N_DST
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q5_K":
+        rows_per_tg = _N_SIMDGROUP * 2
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q6_K":
+        rows_per_tg = _N_SIMDGROUP * 1
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt in ("iq4_nl", "iq4_xs"):
+        rows_per_tg = _N_SIMDGROUP * 2            # 2 sg × 2 rows/sg
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    else:
+        raise ValueError(f"matvec not wired for {plan.fmt}")
+    grid_x = (plan.M + rows_per_tg - 1) // rows_per_tg
+    if needs_shared:
+        enc.setThreadgroupMemoryLength_atIndex_(256, 0)  # 64 floats per simdgroup × 2 sg
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, 1, 1),
+        Metal.MTLSizeMake(tg_threads, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd

build/torch211-metal-aarch64-darwin/_gguf_dequant_530e374.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e502779f760fa35dac6a8aea0c4e63e6dbac0641cb3daf0ae34e4115754899d8
+size 237040

build/torch211-metal-aarch64-darwin/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _gguf_dequant_530e374
+ops = torch.ops._gguf_dequant_530e374
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_gguf_dequant_530e374::{op_name}"

build/torch211-metal-aarch64-darwin/gguf_dequant/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+from pathlib import Path
+from types import ModuleType
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch211-metal-aarch64-darwin/metadata.json

@@ -0,0 +1,4 @@
+{
+  "version": 1,
+  "python-depends": []
+}

build/torch211-metal-aarch64-darwin/nn.py

@@ -0,0 +1,191 @@
+"""GgufLinear — drop-in nn.Linear replacement that runs quantized matmul directly.
+
+Forward pass:
+  - If the input has effective batch size 1 (single token), use mul_mat_vec
+    (memory-bandwidth-bound, faster than bf16 mat-vec on MPS).
+  - Otherwise use mul_mat (compute-bound, similar throughput to bf16 GEMM but
+    keeps weights at native quantization → 3–4× less RAM than bf16).
+
+Supported quant types today: Q4_0, Q8_0, Q4_K. Adding more is just a matter of
+porting the candle kernels — the dispatch logic is generic over `quant_type`.
+
+Prototype design: weights are stored as raw uint8 bytes in `qweight`. The
+forward path copies activations into a shared MTLBuffer, runs the Metal
+kernel, reads the fp32 output back. Production path (kernels-community +
+PR #767) skips the copies via direct MTLBuffer handoff.
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+import torch
+
+from . import (
+    allocate_matmul_plan,
+    allocate_matvec_plan,
+    dispatch_matmul,
+    dispatch_matvec,
+)
+
+
+# Per-quant-type geometry: (block_bytes, block_elems, matvec_M_alignment)
+_QUANT_INFO = {
+    "Q4_0": (18, 32, 8),   # matvec uses 2 simdgroups × N_DST=4 = 8 rows/threadgroup
+    "Q8_0": (34, 32, 8),
+    "Q4_K": (144, 256, 4), # matvec uses 1 simdgroup × N_DST=4 = 4 rows/threadgroup
+}
+
+# Map our Python-side quant name to the Metal kernel suffix.
+_KERNEL_FMT = {
+    "Q4_0": "q4_0",
+    "Q8_0": "q8_0",
+    "Q4_K": "q4_K",
+}
+
+
+class GgufLinear(torch.nn.Module):
+    """Linear layer backed by GGUF-quantized weights.
+
+    Args:
+        in_features: K dimension. Must be a multiple of the format's block size
+            (32 for Q4_0/Q8_0, 256 for Q4_K).
+        out_features: M dimension. Must be a multiple of the format's matvec
+            row alignment (8 for Q4_0/Q8_0, 4 for Q4_K) and of 64 for matmul.
+        quant_type: one of "Q4_0", "Q8_0", "Q4_K".
+        bias: whether to add a learnable bias after the matmul.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        quant_type: str = "Q4_K",
+        bias: bool = False,
+        device: Optional[torch.device | str] = None,
+    ):
+        super().__init__()
+        if quant_type not in _QUANT_INFO:
+            raise NotImplementedError(
+                f"GgufLinear supports {sorted(_QUANT_INFO)} today, got {quant_type}"
+            )
+        block_bytes, block_elems, matvec_align = _QUANT_INFO[quant_type]
+        if in_features % block_elems != 0:
+            raise ValueError(
+                f"in_features must be a multiple of {block_elems} for {quant_type}, got {in_features}"
+            )
+        if out_features % matvec_align != 0:
+            raise ValueError(
+                f"out_features must be a multiple of {matvec_align} for {quant_type}, got {out_features}"
+            )
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.quant_type = quant_type
+        self._fmt = _KERNEL_FMT[quant_type]
+        self._block_bytes = block_bytes
+        self._block_elems = block_elems
+
+        nblocks_per_row = in_features // block_elems
+        qweight_nbytes = out_features * nblocks_per_row * block_bytes
+        self.register_buffer(
+            "qweight",
+            torch.empty(qweight_nbytes, dtype=torch.uint8, device=device or "cpu"),
+            persistent=True,
+        )
+
+        if bias:
+            self.bias = torch.nn.Parameter(
+                torch.zeros(out_features, dtype=torch.float32, device=device or "cpu"),
+                requires_grad=False,
+            )
+        else:
+            self.register_parameter("bias", None)
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"quant_type={self.quant_type}, bias={self.bias is not None}"
+        )
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: shape (..., K). Returns (..., M) fp32."""
+        if x.shape[-1] != self.in_features:
+            raise ValueError(f"Last dim of x must be {self.in_features}, got {x.shape[-1]}")
+
+        batch_shape = x.shape[:-1]
+        x_flat = x.reshape(-1, self.in_features).contiguous().to(torch.float32)
+        N = x_flat.shape[0]
+
+        # CPU view of qweight for the bytes copy into shared MTLBuffers.
+        # Production path skips this once the kernels-community build is wired.
+        qweight_cpu = self.qweight.detach().to("cpu", copy=False)
+        x_cpu = x_flat.to("cpu", copy=False)
+
+        if N == 1:
+            out = self._forward_matvec(qweight_cpu, x_cpu.reshape(-1))
+        else:
+            out = self._forward_matmul(qweight_cpu, x_cpu)
+
+        out = out.reshape(*batch_shape, self.out_features)
+        if x.device != out.device:
+            out = out.to(x.device)
+        if self.bias is not None:
+            out = out + self.bias.to(out.device).to(out.dtype)
+        return out
+
+    def _forward_matvec(self, qweight_cpu: torch.Tensor, x_cpu: torch.Tensor) -> torch.Tensor:
+        M, K = self.out_features, self.in_features
+        plan = allocate_matvec_plan(self._fmt, M=M, K=K)
+        a_view = np.frombuffer(plan.a_buf.contents().as_buffer(qweight_cpu.numel()), dtype=np.uint8)
+        a_view[:] = qweight_cpu.numpy()
+        b_view = np.frombuffer(plan.b_buf.contents().as_buffer(K * 4), dtype=np.float32)
+        b_view[:] = x_cpu.numpy()
+        dispatch_matvec(plan, sync=True)
+        c_view = np.frombuffer(plan.c_buf.contents().as_buffer(M * 4), dtype=np.float32)
+        return torch.from_numpy(c_view.copy())
+
+    def _forward_matmul(self, qweight_cpu: torch.Tensor, x_cpu: torch.Tensor) -> torch.Tensor:
+        N, K = x_cpu.shape
+        M = self.out_features
+        plan = allocate_matmul_plan(self._fmt, M=M, N=N, K=K)
+        a_view = np.frombuffer(plan.a_buf.contents().as_buffer(qweight_cpu.numel()), dtype=np.uint8)
+        a_view[:] = qweight_cpu.numpy()
+        b_view = np.frombuffer(plan.b_buf.contents().as_buffer(N * K * 4), dtype=np.float32)
+        b_view[:] = x_cpu.numpy().reshape(-1)
+        dispatch_matmul(plan, sync=True)
+        c_view = np.frombuffer(plan.c_buf.contents().as_buffer(M * N * 4), dtype=np.float32)
+        # Kernel output layout is (N, M) in row-major; we want (N, M).
+        return torch.from_numpy(c_view.reshape(N, M).copy())
+
+    @classmethod
+    def from_quantized_bytes(
+        cls,
+        quant_type: str,
+        qweight_bytes: bytes,
+        in_features: int,
+        out_features: int,
+        bias: Optional[torch.Tensor] = None,
+        device: Optional[torch.device | str] = None,
+    ) -> "GgufLinear":
+        """Convenience for tests/loaders: hydrate from raw GGUF block bytes."""
+        layer = cls(in_features=in_features, out_features=out_features,
+                    quant_type=quant_type, bias=bias is not None, device=device)
+        expected = layer.qweight.numel()
+        if len(qweight_bytes) != expected:
+            raise ValueError(
+                f"qweight_bytes length {len(qweight_bytes)} != expected {expected} "
+                f"for {quant_type} M={out_features} K={in_features}"
+            )
+        layer.qweight.copy_(torch.frombuffer(bytearray(qweight_bytes), dtype=torch.uint8))
+        if bias is not None:
+            layer.bias.copy_(bias.to(torch.float32))
+        return layer
+
+    # Back-compat shim for the existing Q4_K-only test
+    @classmethod
+    def from_q4_K_bytes(cls, qweight_bytes, in_features, out_features, bias=None, device=None):
+        return cls.from_quantized_bytes(
+            "Q4_K", qweight_bytes, in_features, out_features, bias=bias, device=device
+        )

build/torch211-metal-aarch64-darwin/transformers_glue.py

@@ -0,0 +1,284 @@
+"""Glue for plugging `GgufLinear` into a transformers model.
+
+The flow is intentionally non-invasive: we read raw block bytes directly from
+the GGUF file via `gguf.GGUFReader`, walk the model, and for every `nn.Linear`
+whose corresponding GGUF tensor is one of our supported quant types we
+in-place replace the module with `GgufLinear` and copy the bytes into
+`qweight`. Tensors that don't map to a Linear (or are unsupported quants like
+F32 norms, embeddings) get dequantized via `gguf.dequantize` and copied into
+the existing parameter as fp32.
+
+This avoids modifying the transformers loader. The trade-off is that the
+model has to be constructed first with regular `nn.Linear` shells (which we
+then mutate), but those shells are on the meta device or freshly initialised
+and don't carry the dequantized weight at the time of swap — so peak memory
+during load is bounded by `(model_size_fp32) + (model_size_quantized_qweights)`,
+which fits in RAM for the small models we care about.
+
+Typical use:
+
+    from transformers import AutoConfig, AutoModelForCausalLM
+    cfg = AutoConfig.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct-GGUF",
+                                     gguf_file="qwen2.5-0.5b-instruct-q4_0.gguf")
+    model = AutoModelForCausalLM.from_config(cfg).eval()
+    load_gguf_into_model_with_gguf_linear(
+        model,
+        gguf_path="…/qwen2.5-0.5b-instruct-q4_0.gguf",
+    )
+    out = model.generate(input_ids=…, max_new_tokens=10)
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import gguf
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .nn import GgufLinear, _QUANT_INFO
+
+
+# Map gguf.GGMLQuantizationType.name → our GgufLinear quant_type string.
+_GGUF_NAME_TO_GGUF_LINEAR = {
+    "Q4_0": "Q4_0",
+    "Q8_0": "Q8_0",
+    "Q4_K": "Q4_K",
+}
+
+
+def _quant_supported_by_gguf_linear(qt: gguf.GGMLQuantizationType) -> bool:
+    return qt.name in _GGUF_NAME_TO_GGUF_LINEAR
+
+
+# --- Tensor-name mapping ------------------------------------------------------
+# GGUF tensors use llama.cpp-style names like `blk.0.attn_q.weight`. HF models
+# use `model.layers.0.self_attn.q_proj.weight`. We use a tiny version of the
+# transformers rename table here so the swap util can stay self-contained.
+# (Transformers' own GGUF loader has a much more thorough version of this.)
+
+_BASIC_RENAMES = [
+    (r"^blk\.", "model.layers."),
+    (r"\.attn_q\.weight$",      ".self_attn.q_proj.weight"),
+    (r"\.attn_k\.weight$",      ".self_attn.k_proj.weight"),
+    (r"\.attn_v\.weight$",      ".self_attn.v_proj.weight"),
+    (r"\.attn_output\.weight$", ".self_attn.o_proj.weight"),
+    (r"\.ffn_gate\.weight$",    ".mlp.gate_proj.weight"),
+    (r"\.ffn_up\.weight$",      ".mlp.up_proj.weight"),
+    (r"\.ffn_down\.weight$",    ".mlp.down_proj.weight"),
+    (r"\.attn_norm\.weight$",   ".input_layernorm.weight"),
+    (r"\.ffn_norm\.weight$",    ".post_attention_layernorm.weight"),
+    (r"\.attn_(q|k|v)\.bias$",  r".self_attn.\1_proj.bias"),
+    (r"^token_embd\.weight$",   "model.embed_tokens.weight"),
+    (r"^output_norm\.weight$",  "model.norm.weight"),
+    (r"^output\.weight$",       "lm_head.weight"),
+]
+
+
+def _gguf_name_to_hf_name(gguf_name: str) -> str:
+    """Apply our minimal rename table to map a GGUF tensor name to an HF parameter name."""
+    import re
+
+    out = gguf_name
+    for pat, repl in _BASIC_RENAMES:
+        out = re.sub(pat, repl, out)
+    return out
+
+
+def _is_attn_qk(gguf_name: str) -> bool:
+    return ".attn_q.weight" in gguf_name or ".attn_k.weight" in gguf_name
+
+
+def _reverse_permute_qk(w: np.ndarray, n_head: int, n_kv: int | None) -> np.ndarray:
+    """GGUF Q/K weights are permuted on quantize; undo before loading into HF.
+    Reference: https://github.com/ggerganov/llama.cpp/blob/.../convert_hf_to_gguf.py."""
+    if n_kv is not None and n_head != n_kv:
+        n_head = n_kv
+    dim = w.shape[0] // n_head // 2
+    w = w.reshape(n_head, dim, 2, *w.shape[1:])
+    return w.swapaxes(2, 1).reshape(*w.shape[:1], *w.shape[3:]) if False else w.reshape(-1, *w.shape[3:]).swapaxes(2, 1).reshape(n_head * dim * 2, *w.shape[3:])
+
+
+def _set_module_param(model: nn.Module, qualified_name: str, value: torch.Tensor) -> bool:
+    """Set a parameter or buffer on `model` by dotted name. Returns True on success."""
+    parts = qualified_name.split(".")
+    mod = model
+    for p in parts[:-1]:
+        if not hasattr(mod, p):
+            return False
+        mod = getattr(mod, p)
+    leaf = parts[-1]
+    if leaf in mod._parameters:
+        param = mod._parameters[leaf]
+        if param is None:
+            return False
+        with torch.no_grad():
+            mod._parameters[leaf] = nn.Parameter(value.to(dtype=param.dtype, device=param.device),
+                                                 requires_grad=param.requires_grad)
+        return True
+    if leaf in mod._buffers:
+        with torch.no_grad():
+            mod._buffers[leaf] = value
+        return True
+    return False
+
+
+def _get_module_by_name(model: nn.Module, qualified_name: str) -> Optional[nn.Module]:
+    """Return the leaf module matching `qualified_name` minus its terminal attribute,
+    or None if anything along the path is missing."""
+    parts = qualified_name.split(".")
+    mod = model
+    for p in parts[:-1]:
+        if not hasattr(mod, p):
+            return None
+        mod = getattr(mod, p)
+    return mod
+
+
+def _get_attr_by_path(model: nn.Module, qualified_name: str) -> Optional[nn.Module]:
+    """Walk dotted `qualified_name` on `model` and return the final attribute,
+    or None if anything along the path is missing."""
+    if not qualified_name:
+        return model
+    mod = model
+    for p in qualified_name.split("."):
+        if not hasattr(mod, p):
+            return None
+        mod = getattr(mod, p)
+    return mod
+
+
+def _swap_linear_to_gguf(
+    parent: nn.Module, attr: str, in_features: int, out_features: int, quant_type: str,
+    bias_tensor: Optional[torch.Tensor], device: Optional[torch.device | str] = None
+) -> GgufLinear:
+    """Replace `parent.attr` (an nn.Linear) with a new GgufLinear of the same shape."""
+    new = GgufLinear(in_features=in_features, out_features=out_features,
+                     quant_type=quant_type, bias=bias_tensor is not None, device=device)
+    if bias_tensor is not None and new.bias is not None:
+        new.bias.data.copy_(bias_tensor.to(torch.float32))
+    setattr(parent, attr, new)
+    return new
+
+
+def load_gguf_into_model_with_gguf_linear(
+    model: nn.Module,
+    gguf_path: str | Path,
+    skip_modules: Optional[set[str]] = None,
+    force_dequant: bool = False,
+) -> dict:
+    """Read a GGUF file and populate `model` in-place.
+
+    Two-phase load to keep the byte path simple and the bias path correct:
+
+      Phase 1 — full dequant: every tensor (linear weights, norms, embeddings,
+        biases) is dequantized to fp32 and written into the matching HF
+        parameter. After this pass the model would already produce correct
+        outputs (this is what `force_dequant=True` stops at).
+      Phase 2 — swap linears: walk every `nn.Linear` that maps to a tensor
+        whose GGUF quant type is supported by `GgufLinear`; swap it for a
+        `GgufLinear` and re-attach raw block bytes (preserving the bias that
+        Phase 1 just loaded).
+
+    Returns a dict with diagnostic counts.
+    """
+    reader = gguf.GGUFReader(str(gguf_path))
+
+    # Pull a few config bits for the Q/K reverse-permute.
+    def _read_int(field: str, default: int = 0) -> int:
+        f = reader.fields.get(field)
+        if f is None:
+            return default
+        return int(f.parts[f.data[0]][0]) if f.data else default
+
+    arch = ""
+    arch_field = reader.fields.get("general.architecture")
+    if arch_field is not None and arch_field.data:
+        arch = bytes(arch_field.parts[arch_field.data[0]]).decode("utf-8", errors="ignore")
+    n_head    = _read_int(f"{arch}.attention.head_count")
+    n_kv_head = _read_int(f"{arch}.attention.head_count_kv") or n_head
+
+    skip_modules = skip_modules or set()
+    stats = {"as_gguf_linear": 0, "dequantized": 0, "skipped": 0, "unmatched": []}
+
+    # Track which (hf_name → quant_type) pairs are eligible for GgufLinear swap
+    # after Phase 1. We'll re-quantize the loaded fp32 weight in Phase 2.
+    swap_candidates: list[tuple[str, str]] = []
+
+    # ---------------- Phase 1: dequant everything into the model. -------------
+    for t in reader.tensors:
+        hf_name = _gguf_name_to_hf_name(t.name)
+        if any(skip in hf_name for skip in skip_modules):
+            stats["skipped"] += 1
+            continue
+
+        parent = _get_module_by_name(model, hf_name)
+        if parent is None or not hasattr(parent, hf_name.rsplit(".", 1)[-1]):
+            stats["unmatched"].append(t.name)
+            continue
+        leaf = hf_name.rsplit(".", 1)[-1]
+        cur = getattr(parent, leaf)
+
+        fp32 = gguf.dequantize(
+            np.frombuffer(np.ascontiguousarray(t.data).tobytes(), dtype=np.uint8),
+            t.tensor_type,
+        ).astype(np.float32)
+        if _is_attn_qk(t.name) and n_head > 0:
+            shape = list(t.shape)
+            M, K = int(shape[1]), int(shape[0])
+            heads = n_head if ".attn_q." in t.name else n_kv_head
+            head_dim = M // heads
+            fp32 = fp32.reshape(M, K).reshape(heads, head_dim // 2, 2, K).swapaxes(2, 1).reshape(M, K)
+        target_shape = tuple(cur.shape) if hasattr(cur, "shape") else (fp32.size,)
+        if np.prod(target_shape) == fp32.size:
+            fp32 = fp32.reshape(target_shape)
+        tensor = torch.from_numpy(fp32)
+        if not _set_module_param(model, hf_name, tensor):
+            stats["unmatched"].append(t.name)
+            continue
+        stats["dequantized"] += 1
+
+        # Record candidacy for swap in Phase 2.
+        is_linear_weight = isinstance(parent, nn.Linear) and leaf == "weight"
+        if is_linear_weight and _quant_supported_by_gguf_linear(t.tensor_type) and not force_dequant:
+            swap_candidates.append((hf_name, _GGUF_NAME_TO_GGUF_LINEAR[t.tensor_type.name]))
+
+    # ---------------- Phase 2: swap Linear → GgufLinear. ----------------------
+    # Re-quantize the loaded fp32 weight via gguf-py (works for Q4_0/Q8_0;
+    # K-quants don't have a Python quantizer yet — those will skip swap).
+    for hf_name, quant_type in swap_candidates:
+        # hf_name ends in `.weight`; the Linear module is at hf_name[:-len(".weight")]
+        linear_path = hf_name[: -len(".weight")]
+        grandparent_path, _, linear_attr = linear_path.rpartition(".")
+        grandparent = model if not grandparent_path else _get_attr_by_path(model, grandparent_path)
+        linear = _get_attr_by_path(model, linear_path)
+        if grandparent is None or not isinstance(linear, nn.Linear):
+            continue
+
+        # Re-quantize the (now permuted-if-applicable) fp32 weight.
+        try:
+            qt = getattr(gguf.GGMLQuantizationType, quant_type)
+            qbytes = bytes(
+                gguf.quantize(linear.weight.detach().cpu().numpy().astype(np.float32), qt).tobytes()
+            )
+        except (NotImplementedError, Exception) as e:
+            # K-quants etc. — leave as Linear (fp32 weight is already loaded).
+            stats.setdefault("re_quantize_failed", []).append((hf_name, str(e)))
+            continue
+
+        bias_tensor = linear.bias.detach().clone() if linear.bias is not None else None
+        new = GgufLinear(
+            in_features=linear.in_features,
+            out_features=linear.out_features,
+            quant_type=quant_type,
+            bias=bias_tensor is not None,
+        )
+        if bias_tensor is not None:
+            new.bias.data.copy_(bias_tensor.to(torch.float32))
+        new.qweight.copy_(torch.frombuffer(bytearray(qbytes), dtype=torch.uint8))
+        setattr(grandparent, linear_attr, new)
+        stats["as_gguf_linear"] += 1
+        stats["dequantized"] -= 1  # this one's no longer "just" dequant — swapped
+
+    return stats

build/torch28-metal-aarch64-darwin/__init__.py

@@ -0,0 +1,600 @@
+"""GGUF dequantization kernels for Apple Silicon (Metal).
+
+Six standalone kernels covering the GGUF formats most common in the wild:
+
+    Q4_0   Q8_0   Q4_K   Q5_K   Q6_K   IQ4_NL   IQ4_XS
+
+Two API levels are exposed for each:
+
+  - ``dequantize_<format>_bytes(quant_bytes) -> np.ndarray``
+    One-shot path used by the tests for bit-exact validation against a
+    NumPy reference. Allocates input/output buffers internally.
+
+  - ``allocate_<format>_plan(nblocks)`` + ``dispatch_<format>(plan)``
+    Low-level path for benchmarks. Buffers are pre-allocated; ``dispatch``
+    only encodes + dispatches + syncs, so the timing isolates GPU work
+    from CPU<->buffer copies.
+
+The runtime path uses PyObjC + Metal.framework so the correctness loop
+works on any Apple Silicon Mac without Xcode 26 / Nix / kernel-builder.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+# Lazy globals — populated on first use.
+_DEVICE = None
+_COMMAND_QUEUE = None
+_LIBRARY = None
+_PIPELINES: dict = {}
+
+_KERNEL_SOURCE_PATH = Path(__file__).resolve().parent.parent.parent / "gguf_dequant_metal" / "dequantize.metal"
+
+
+def _ensure_metal():
+    global _DEVICE, _COMMAND_QUEUE, _LIBRARY
+    if _DEVICE is not None:
+        return
+    import Metal  # pyobjc-framework-metal
+
+    _DEVICE = Metal.MTLCreateSystemDefaultDevice()
+    if _DEVICE is None:
+        raise RuntimeError("No default Metal device available")
+    _COMMAND_QUEUE = _DEVICE.newCommandQueue()
+
+    src = _KERNEL_SOURCE_PATH.read_text()
+    options = Metal.MTLCompileOptions.new()
+    lib, err = _DEVICE.newLibraryWithSource_options_error_(src, options, None)
+    if lib is None:
+        raise RuntimeError(f"Metal compile failed: {err}")
+    _LIBRARY = lib
+
+
+def _pipeline(fn_name: str):
+    _ensure_metal()
+    pso = _PIPELINES.get(fn_name)
+    if pso is not None:
+        return pso
+
+    fn = _LIBRARY.newFunctionWithName_(fn_name)
+    if fn is None:
+        raise RuntimeError(f"Kernel function not found: {fn_name}")
+    pso, err = _DEVICE.newComputePipelineStateWithFunction_error_(fn, None)
+    if pso is None:
+        raise RuntimeError(f"Pipeline state creation failed: {err}")
+    _PIPELINES[fn_name] = pso
+    return pso
+
+
+# (kernel_name, bytes_per_block, elems_per_block, threads_per_block_group)
+_FORMATS = {
+    "q4_0":   ("dequantize_q4_0",   18, 32, 16),
+    "q8_0":   ("dequantize_q8_0",   34, 32, 32),
+    "q4_K":   ("dequantize_q4_K",  144, 256, 256),
+    "q5_K":   ("dequantize_q5_K",  176, 256, 256),
+    "q6_K":   ("dequantize_q6_K",  210, 256, 256),
+    "iq4_nl": ("dequantize_iq4_nl", 18, 32, 16),
+    "iq4_xs": ("dequantize_iq4_xs", 136, 256, 256),
+}
+
+
+class Plan:
+    """Pre-allocated I/O buffers for a single (format, nblocks) shape."""
+
+    __slots__ = ("fmt", "nblocks", "in_buf", "out_buf", "_nblocks_bytes")
+
+    def __init__(self, fmt: str, nblocks: int, in_buf, out_buf):
+        self.fmt = fmt
+        self.nblocks = nblocks
+        self.in_buf = in_buf
+        self.out_buf = out_buf
+        self._nblocks_bytes = np.array([nblocks], dtype=np.uint32).tobytes()
+
+
+def allocate_plan(fmt: str, nblocks: int) -> Plan:
+    import Metal
+
+    _ensure_metal()
+    _, bpb, epb, _ = _FORMATS[fmt]
+    in_buf = _DEVICE.newBufferWithLength_options_(nblocks * bpb, Metal.MTLResourceStorageModeShared)
+    out_buf = _DEVICE.newBufferWithLength_options_(nblocks * epb * 4, Metal.MTLResourceStorageModeShared)
+    return Plan(fmt=fmt, nblocks=nblocks, in_buf=in_buf, out_buf=out_buf)
+
+
+def dispatch(plan: Plan, sync: bool = True):
+    import Metal
+
+    kernel_name, _, _, tg = _FORMATS[plan.fmt]
+    pso = _pipeline(kernel_name)
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.in_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.out_buf, 0, 1)
+    enc.setBytes_length_atIndex_(plan._nblocks_bytes, len(plan._nblocks_bytes), 2)
+    enc.dispatchThreads_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(plan.nblocks * tg, 1, 1),
+        Metal.MTLSizeMake(tg, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+def dequantize_bytes(fmt: str, quant_bytes) -> np.ndarray:
+    """One-shot bytes-in / float32-out for the named format."""
+    _ensure_metal()
+    _, bpb, epb, _ = _FORMATS[fmt]
+
+    buf = bytes(quant_bytes) if not isinstance(quant_bytes, (bytes, bytearray)) else quant_bytes
+    nbytes = len(buf)
+    if nbytes % bpb != 0:
+        raise ValueError(f"{fmt} input not a multiple of {bpb} bytes: got {nbytes}")
+    nblocks = nbytes // bpb
+
+    plan = allocate_plan(fmt, nblocks)
+    in_view = np.frombuffer(plan.in_buf.contents().as_buffer(nbytes), dtype=np.uint8)
+    in_view[:] = np.frombuffer(buf, dtype=np.uint8)
+    dispatch(plan, sync=True)
+    out_nbytes = nblocks * epb * 4
+    return np.frombuffer(plan.out_buf.contents().as_buffer(out_nbytes), dtype=np.float32).copy()
+
+
+# --- Per-format convenience wrappers ----------------------------------------
+
+def dequantize_q4_0_bytes(qb):   return dequantize_bytes("q4_0",   qb)
+def dequantize_q8_0_bytes(qb):   return dequantize_bytes("q8_0",   qb)
+def dequantize_q4_K_bytes(qb):   return dequantize_bytes("q4_K",   qb)
+def dequantize_q5_K_bytes(qb):   return dequantize_bytes("q5_K",   qb)
+def dequantize_q6_K_bytes(qb):   return dequantize_bytes("q6_K",   qb)
+def dequantize_iq4_nl_bytes(qb): return dequantize_bytes("iq4_nl", qb)
+def dequantize_iq4_xs_bytes(qb): return dequantize_bytes("iq4_xs", qb)
+
+def allocate_q4_0_plan(nblocks):   return allocate_plan("q4_0",   nblocks)
+def allocate_q8_0_plan(nblocks):   return allocate_plan("q8_0",   nblocks)
+def allocate_q4_K_plan(nblocks):   return allocate_plan("q4_K",   nblocks)
+def allocate_q5_K_plan(nblocks):   return allocate_plan("q5_K",   nblocks)
+def allocate_q6_K_plan(nblocks):   return allocate_plan("q6_K",   nblocks)
+def allocate_iq4_nl_plan(nblocks): return allocate_plan("iq4_nl", nblocks)
+def allocate_iq4_xs_plan(nblocks): return allocate_plan("iq4_xs", nblocks)
+
+def dispatch_q4_0(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q8_0(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q4_K(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q5_K(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q6_K(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_iq4_nl(plan, **kw): return dispatch(plan, **kw)
+def dispatch_iq4_xs(plan, **kw): return dispatch(plan, **kw)
+
+
+# ---------------------------------------------------------------------------
+# Torch op entry points — the production path used by transformers (no PyObjC,
+# no runtime compile). Only available in builds produced by kernel-builder,
+# which generates `_ops.py` exporting `ops = torch.ops._gguf_dequant_<rev>`.
+# In a dev-mode checkout (running from torch-ext/ directly) `_ops` isn't
+# present, so the torch_* entry points raise on call.
+# ---------------------------------------------------------------------------
+
+try:
+    from ._ops import ops as _ops  # generated by kernel-builder
+except ImportError:
+    _ops = None
+
+
+def _require_ops():
+    if _ops is None:
+        raise RuntimeError(
+            "torch op binding not available — this build was loaded from the dev tree. "
+            "Use the *_bytes / *_plan APIs (PyObjC path) or install a kernel-builder build."
+        )
+    return _ops
+
+
+def torch_dequantize_q4_0(input, output):   _require_ops().dequantize_q4_0(input, output)
+def torch_dequantize_q8_0(input, output):   _require_ops().dequantize_q8_0(input, output)
+def torch_dequantize_q4_K(input, output):   _require_ops().dequantize_q4_K(input, output)
+def torch_dequantize_q5_K(input, output):   _require_ops().dequantize_q5_K(input, output)
+def torch_dequantize_q6_K(input, output):   _require_ops().dequantize_q6_K(input, output)
+def torch_dequantize_iq4_nl(input, output): _require_ops().dequantize_iq4_nl(input, output)
+def torch_dequantize_iq4_xs(input, output): _require_ops().dequantize_iq4_xs(input, output)
+
+
+# =============================================================================
+# Q4_K × fp32 matmul — `GgufLinear` ceiling experiment.
+# =============================================================================
+# Layout convention (matches llama.cpp's kernel_mul_mm):
+#   A (Q4_K weights):  shape (M, K)  M = out_features, K = in_features.
+#                      M and K must be multiples of 64 and 256 respectively for
+#                      the simdgroup tile geometry to land cleanly.
+#   B (fp32 input):    shape (N, K)  N = batch * seqlen, contiguous in K.
+#   C (fp32 output):   shape (M, N)  column-major-ish (matches llama.cpp).
+
+_BLOCK_SIZE_M = 64
+_BLOCK_SIZE_N = 32
+_THREAD_PER_BLOCK = 128
+_SHARED_BYTES = 8192
+
+
+class MatMulQ4KPlan:
+    __slots__ = ("M", "N", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, M, N, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.M = M
+        self.N = N
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def allocate_matmul_q4_K_plan(M: int, N: int, K: int) -> MatMulQ4KPlan:
+    """Pre-allocate weight (Q4_K), input (fp32), output (fp32) MTLBuffers."""
+    import Metal
+
+    _ensure_metal()
+    assert M % _BLOCK_SIZE_M == 0, f"M must be multiple of {_BLOCK_SIZE_M}, got {M}"
+    assert N % _BLOCK_SIZE_N == 0, f"N must be multiple of {_BLOCK_SIZE_N}, got {N}"
+    assert K % 256 == 0, f"K must be multiple of 256, got {K}"
+
+    a_bytes = M * (K // 256) * 144  # M rows × (K/QK_K) super-blocks × 144 bytes
+    b_bytes = N * K * 4             # fp32
+    c_bytes = M * N * 4             # fp32
+
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+
+    # Pack ne00, ne01, ne11, nb01, nb11 into a uniforms buffer.
+    nb01 = (K // 256) * 144  # bytes per A row
+    nb11 = K * 4             # bytes per B row
+    dims = np.array([K], dtype=np.int64).tobytes()  # ne00
+    dims += np.array([M], dtype=np.int64).tobytes()  # ne01
+    dims += np.array([N], dtype=np.int64).tobytes()  # ne11
+    dims += np.array([nb01], dtype=np.uint64).tobytes()
+    dims += np.array([nb11], dtype=np.uint64).tobytes()
+
+    return MatMulQ4KPlan(M=M, N=N, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def dispatch_matmul_q4_K(plan: MatMulQ4KPlan, sync: bool = True):
+    """Dispatch mul_mat_q4_K_f32 on the pre-allocated buffers."""
+    import Metal
+
+    pso = _pipeline("mul_mat_q4_K_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    # ne00 (K), ne01 (M), ne11 (N), nb01, nb11 via setBytes — 5 scalars
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8],   8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16],  8, 4)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[16:24], 8, 5)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[24:32], 8, 6)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[32:40], 8, 7)
+    enc.setThreadgroupMemoryLength_atIndex_(_SHARED_BYTES, 0)
+
+    grid_x = (plan.N + _BLOCK_SIZE_N - 1) // _BLOCK_SIZE_N
+    grid_y = (plan.M + _BLOCK_SIZE_M - 1) // _BLOCK_SIZE_M
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, grid_y, 1),
+        Metal.MTLSizeMake(_THREAD_PER_BLOCK, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+# =============================================================================
+# mul_mat_vec_q4_K — batch-1 decode (memory-bandwidth-bound regime).
+# =============================================================================
+
+_N_DST = 4
+_SIMD_WIDTH = 32
+
+
+class MatVecQ4KPlan:
+    __slots__ = ("M", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, M, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.M = M
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def allocate_matvec_q4_K_plan(M: int, K: int) -> MatVecQ4KPlan:
+    import Metal
+
+    _ensure_metal()
+    assert M % _N_DST == 0, f"M must be a multiple of {_N_DST}, got {M}"
+    assert K % 256 == 0, f"K must be a multiple of 256, got {K}"
+
+    a_bytes = M * (K // 256) * 144
+    b_bytes = K * 4
+    c_bytes = M * 4
+
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+
+    dims = np.array([K], dtype=np.int64).tobytes() + np.array([M], dtype=np.int64).tobytes()
+    return MatVecQ4KPlan(M=M, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def dispatch_matvec_q4_K(plan: MatVecQ4KPlan, sync: bool = True):
+    import Metal
+
+    pso = _pipeline("mul_mat_vec_q4_K_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8], 8, 3)   # ne00 = K
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16], 8, 4)  # ne01 = M
+
+    grid_x = (plan.M + _N_DST - 1) // _N_DST
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, 1, 1),
+        Metal.MTLSizeMake(_SIMD_WIDTH, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+# =============================================================================
+# Q4_0 / Q8_0 matmul + matvec — same layout as Q4_K, smaller blocks (32 elems).
+# =============================================================================
+
+_N_SIMDGROUP = 2  # Q4_0/Q8_0 matvec uses 2 simdgroups per threadgroup
+
+# Per quant type for the mat path: (block bytes, block elems, matvec_M_alignment).
+# block_elems = 32 for Q4_0/Q8_0/IQ4_NL, 256 for all K-quants and IQ4_XS.
+# matvec_M_alignment:
+#   Q4_0/Q8_0  →  2 simdgroups × N_DST(4) = 8 rows per threadgroup
+#   Q4_K       →  1 simdgroup  × N_DST(4) = 4 rows per threadgroup
+#   Q5_K       →  2 simdgroups × 2 rows   = 4 rows per threadgroup
+#   Q6_K       →  2 simdgroups × 1 row    = 2 rows per threadgroup
+_MAT_FORMATS = {
+    "q4_0":   (18,  32,  8),
+    "q8_0":   (34,  32,  8),
+    "q4_K":   (144, 256, 4),
+    "q5_K":   (176, 256, 4),
+    "q6_K":   (210, 256, 2),
+    "iq4_nl": (18,  32,  4),       # 2 sg × 2 rows/sg = 4 rows/threadgroup
+    "iq4_xs": (136, 256, 4),
+}
+
+
+class MatMulPlan:
+    """Pre-allocated I/O buffers for `mul_mat_<fmt>_f32`. Block element size
+    depends on the format (32 for Q4_0/Q8_0, 256 for Q4_K)."""
+
+    __slots__ = ("fmt", "M", "N", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, fmt, M, N, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.fmt = fmt
+        self.M = M
+        self.N = N
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def _quant_block_elems(fmt: str) -> int:
+    return _MAT_FORMATS[fmt][1]
+
+
+def _quant_block_bytes(fmt: str) -> int:
+    return _MAT_FORMATS[fmt][0]
+
+
+def _matvec_row_alignment(fmt: str) -> int | None:
+    return _MAT_FORMATS[fmt][2]
+
+
+def allocate_matmul_plan(fmt: str, M: int, N: int, K: int) -> MatMulPlan:
+    import Metal
+
+    _ensure_metal()
+    assert fmt in _MAT_FORMATS, f"unsupported fmt {fmt}"
+    block_elems = _quant_block_elems(fmt)
+    block_bytes = _quant_block_bytes(fmt)
+    assert M % _BLOCK_SIZE_M == 0, f"M must be multiple of {_BLOCK_SIZE_M}, got {M}"
+    assert N % _BLOCK_SIZE_N == 0, f"N must be multiple of {_BLOCK_SIZE_N}, got {N}"
+    assert K % block_elems == 0, f"K must be multiple of {block_elems} for {fmt}, got {K}"
+
+    a_bytes = M * (K // block_elems) * block_bytes
+    b_bytes = N * K * 4
+    c_bytes = M * N * 4
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+
+    nb01 = (K // block_elems) * block_bytes
+    nb11 = K * 4
+    dims = (
+        np.array([K], dtype=np.int64).tobytes()
+        + np.array([M], dtype=np.int64).tobytes()
+        + np.array([N], dtype=np.int64).tobytes()
+        + np.array([nb01], dtype=np.uint64).tobytes()
+        + np.array([nb11], dtype=np.uint64).tobytes()
+    )
+    return MatMulPlan(fmt=fmt, M=M, N=N, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def dispatch_matmul(plan: MatMulPlan, sync: bool = True):
+    import Metal
+
+    pso = _pipeline(f"mul_mat_{plan.fmt}_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8],   8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16],  8, 4)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[16:24], 8, 5)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[24:32], 8, 6)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[32:40], 8, 7)
+    enc.setThreadgroupMemoryLength_atIndex_(_SHARED_BYTES, 0)
+
+    grid_x = (plan.N + _BLOCK_SIZE_N - 1) // _BLOCK_SIZE_N
+    grid_y = (plan.M + _BLOCK_SIZE_M - 1) // _BLOCK_SIZE_M
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, grid_y, 1),
+        Metal.MTLSizeMake(_THREAD_PER_BLOCK, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+class MatVecPlan:
+    __slots__ = ("fmt", "M", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, fmt, M, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.fmt = fmt
+        self.M = M
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def allocate_matvec_plan(fmt: str, M: int, K: int) -> MatVecPlan:
+    """Allocate I/O for `mul_mat_vec_<fmt>_f32`. Format-specific row alignment:
+    Q4_0/Q8_0: 8, Q4_K: 4, Q5_K: 4, Q6_K: 2."""
+    import Metal
+
+    _ensure_metal()
+    block_elems = _quant_block_elems(fmt)
+    block_bytes = _quant_block_bytes(fmt)
+    align = _matvec_row_alignment(fmt)
+    if align is None:
+        raise ValueError(f"mul_mat_vec not implemented for {fmt}")
+    assert M % align == 0, f"M must be a multiple of {align} for {fmt}, got {M}"
+    assert K % block_elems == 0, f"K must be a multiple of {block_elems} for {fmt}, got {K}"
+
+    a_bytes = M * (K // block_elems) * block_bytes
+    b_bytes = K * 4
+    c_bytes = M * 4
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+    dims = np.array([K], dtype=np.int64).tobytes() + np.array([M], dtype=np.int64).tobytes()
+    return MatVecPlan(fmt=fmt, M=M, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def encode_matvec(plan: MatVecPlan, enc):
+    """Encode a single matvec dispatch onto an existing MTLComputeCommandEncoder.
+    Caller is responsible for committing the parent command buffer. Threadgroup
+    memory for the IQ4 lookup table is set via setThreadgroupMemoryLength here."""
+    import Metal
+
+    pso = _pipeline(f"mul_mat_vec_{plan.fmt}_f32")
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8], 8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16], 8, 4)
+
+    if plan.fmt == "q4_K":
+        rows_per_tg = _N_DST
+        tg_threads = _SIMD_WIDTH
+    elif plan.fmt in ("q4_0", "q8_0"):
+        rows_per_tg = _N_SIMDGROUP * _N_DST
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q5_K":
+        rows_per_tg = _N_SIMDGROUP * 2
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q6_K":
+        rows_per_tg = _N_SIMDGROUP * 1
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt in ("iq4_nl", "iq4_xs"):
+        rows_per_tg = _N_SIMDGROUP * 2
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+        enc.setThreadgroupMemoryLength_atIndex_(256, 0)
+    else:
+        raise ValueError(f"matvec not wired for {plan.fmt}")
+    grid_x = (plan.M + rows_per_tg - 1) // rows_per_tg
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, 1, 1),
+        Metal.MTLSizeMake(tg_threads, 1, 1),
+    )
+
+
+def new_command_buffer():
+    """Start a fresh MTLCommandBuffer. Caller batches `encode_matvec(plan, enc)`
+    onto its encoder, then calls `cb.commit()` and `cb.waitUntilCompleted()`
+    once for the whole batch — matching llama.cpp's per-step semantics."""
+    _ensure_metal()
+    return _COMMAND_QUEUE.commandBuffer()
+
+
+def dispatch_matvec(plan: MatVecPlan, sync: bool = True):
+    import Metal
+
+    pso = _pipeline(f"mul_mat_vec_{plan.fmt}_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8], 8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16], 8, 4)
+
+    needs_shared = plan.fmt in ("iq4_nl", "iq4_xs")
+    if plan.fmt == "q4_K":
+        rows_per_tg = _N_DST
+        tg_threads = _SIMD_WIDTH
+    elif plan.fmt in ("q4_0", "q8_0"):
+        rows_per_tg = _N_SIMDGROUP * _N_DST
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q5_K":
+        rows_per_tg = _N_SIMDGROUP * 2
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q6_K":
+        rows_per_tg = _N_SIMDGROUP * 1
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt in ("iq4_nl", "iq4_xs"):
+        rows_per_tg = _N_SIMDGROUP * 2            # 2 sg × 2 rows/sg
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    else:
+        raise ValueError(f"matvec not wired for {plan.fmt}")
+    grid_x = (plan.M + rows_per_tg - 1) // rows_per_tg
+    if needs_shared:
+        enc.setThreadgroupMemoryLength_atIndex_(256, 0)  # 64 floats per simdgroup × 2 sg
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, 1, 1),
+        Metal.MTLSizeMake(tg_threads, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd

build/torch28-metal-aarch64-darwin/_gguf_dequant_530e374.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28a3233811126e83e9a2e95f84828b33889973a5ef37c402949253da8ce3ffda
+size 236176

build/torch28-metal-aarch64-darwin/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _gguf_dequant_530e374
+ops = torch.ops._gguf_dequant_530e374
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_gguf_dequant_530e374::{op_name}"

build/torch28-metal-aarch64-darwin/gguf_dequant/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+from pathlib import Path
+from types import ModuleType
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch28-metal-aarch64-darwin/metadata.json

@@ -0,0 +1,4 @@
+{
+  "version": 1,
+  "python-depends": []
+}

build/torch28-metal-aarch64-darwin/nn.py

@@ -0,0 +1,191 @@
+"""GgufLinear — drop-in nn.Linear replacement that runs quantized matmul directly.
+
+Forward pass:
+  - If the input has effective batch size 1 (single token), use mul_mat_vec
+    (memory-bandwidth-bound, faster than bf16 mat-vec on MPS).
+  - Otherwise use mul_mat (compute-bound, similar throughput to bf16 GEMM but
+    keeps weights at native quantization → 3–4× less RAM than bf16).
+
+Supported quant types today: Q4_0, Q8_0, Q4_K. Adding more is just a matter of
+porting the candle kernels — the dispatch logic is generic over `quant_type`.
+
+Prototype design: weights are stored as raw uint8 bytes in `qweight`. The
+forward path copies activations into a shared MTLBuffer, runs the Metal
+kernel, reads the fp32 output back. Production path (kernels-community +
+PR #767) skips the copies via direct MTLBuffer handoff.
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+import torch
+
+from . import (
+    allocate_matmul_plan,
+    allocate_matvec_plan,
+    dispatch_matmul,
+    dispatch_matvec,
+)
+
+
+# Per-quant-type geometry: (block_bytes, block_elems, matvec_M_alignment)
+_QUANT_INFO = {
+    "Q4_0": (18, 32, 8),   # matvec uses 2 simdgroups × N_DST=4 = 8 rows/threadgroup
+    "Q8_0": (34, 32, 8),
+    "Q4_K": (144, 256, 4), # matvec uses 1 simdgroup × N_DST=4 = 4 rows/threadgroup
+}
+
+# Map our Python-side quant name to the Metal kernel suffix.
+_KERNEL_FMT = {
+    "Q4_0": "q4_0",
+    "Q8_0": "q8_0",
+    "Q4_K": "q4_K",
+}
+
+
+class GgufLinear(torch.nn.Module):
+    """Linear layer backed by GGUF-quantized weights.
+
+    Args:
+        in_features: K dimension. Must be a multiple of the format's block size
+            (32 for Q4_0/Q8_0, 256 for Q4_K).
+        out_features: M dimension. Must be a multiple of the format's matvec
+            row alignment (8 for Q4_0/Q8_0, 4 for Q4_K) and of 64 for matmul.
+        quant_type: one of "Q4_0", "Q8_0", "Q4_K".
+        bias: whether to add a learnable bias after the matmul.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        quant_type: str = "Q4_K",
+        bias: bool = False,
+        device: Optional[torch.device | str] = None,
+    ):
+        super().__init__()
+        if quant_type not in _QUANT_INFO:
+            raise NotImplementedError(
+                f"GgufLinear supports {sorted(_QUANT_INFO)} today, got {quant_type}"
+            )
+        block_bytes, block_elems, matvec_align = _QUANT_INFO[quant_type]
+        if in_features % block_elems != 0:
+            raise ValueError(
+                f"in_features must be a multiple of {block_elems} for {quant_type}, got {in_features}"
+            )
+        if out_features % matvec_align != 0:
+            raise ValueError(
+                f"out_features must be a multiple of {matvec_align} for {quant_type}, got {out_features}"
+            )
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.quant_type = quant_type
+        self._fmt = _KERNEL_FMT[quant_type]
+        self._block_bytes = block_bytes
+        self._block_elems = block_elems
+
+        nblocks_per_row = in_features // block_elems
+        qweight_nbytes = out_features * nblocks_per_row * block_bytes
+        self.register_buffer(
+            "qweight",
+            torch.empty(qweight_nbytes, dtype=torch.uint8, device=device or "cpu"),
+            persistent=True,
+        )
+
+        if bias:
+            self.bias = torch.nn.Parameter(
+                torch.zeros(out_features, dtype=torch.float32, device=device or "cpu"),
+                requires_grad=False,
+            )
+        else:
+            self.register_parameter("bias", None)
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"quant_type={self.quant_type}, bias={self.bias is not None}"
+        )
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: shape (..., K). Returns (..., M) fp32."""
+        if x.shape[-1] != self.in_features:
+            raise ValueError(f"Last dim of x must be {self.in_features}, got {x.shape[-1]}")
+
+        batch_shape = x.shape[:-1]
+        x_flat = x.reshape(-1, self.in_features).contiguous().to(torch.float32)
+        N = x_flat.shape[0]
+
+        # CPU view of qweight for the bytes copy into shared MTLBuffers.
+        # Production path skips this once the kernels-community build is wired.
+        qweight_cpu = self.qweight.detach().to("cpu", copy=False)
+        x_cpu = x_flat.to("cpu", copy=False)
+
+        if N == 1:
+            out = self._forward_matvec(qweight_cpu, x_cpu.reshape(-1))
+        else:
+            out = self._forward_matmul(qweight_cpu, x_cpu)
+
+        out = out.reshape(*batch_shape, self.out_features)
+        if x.device != out.device:
+            out = out.to(x.device)
+        if self.bias is not None:
+            out = out + self.bias.to(out.device).to(out.dtype)
+        return out
+
+    def _forward_matvec(self, qweight_cpu: torch.Tensor, x_cpu: torch.Tensor) -> torch.Tensor:
+        M, K = self.out_features, self.in_features
+        plan = allocate_matvec_plan(self._fmt, M=M, K=K)
+        a_view = np.frombuffer(plan.a_buf.contents().as_buffer(qweight_cpu.numel()), dtype=np.uint8)
+        a_view[:] = qweight_cpu.numpy()
+        b_view = np.frombuffer(plan.b_buf.contents().as_buffer(K * 4), dtype=np.float32)
+        b_view[:] = x_cpu.numpy()
+        dispatch_matvec(plan, sync=True)
+        c_view = np.frombuffer(plan.c_buf.contents().as_buffer(M * 4), dtype=np.float32)
+        return torch.from_numpy(c_view.copy())
+
+    def _forward_matmul(self, qweight_cpu: torch.Tensor, x_cpu: torch.Tensor) -> torch.Tensor:
+        N, K = x_cpu.shape
+        M = self.out_features
+        plan = allocate_matmul_plan(self._fmt, M=M, N=N, K=K)
+        a_view = np.frombuffer(plan.a_buf.contents().as_buffer(qweight_cpu.numel()), dtype=np.uint8)
+        a_view[:] = qweight_cpu.numpy()
+        b_view = np.frombuffer(plan.b_buf.contents().as_buffer(N * K * 4), dtype=np.float32)
+        b_view[:] = x_cpu.numpy().reshape(-1)
+        dispatch_matmul(plan, sync=True)
+        c_view = np.frombuffer(plan.c_buf.contents().as_buffer(M * N * 4), dtype=np.float32)
+        # Kernel output layout is (N, M) in row-major; we want (N, M).
+        return torch.from_numpy(c_view.reshape(N, M).copy())
+
+    @classmethod
+    def from_quantized_bytes(
+        cls,
+        quant_type: str,
+        qweight_bytes: bytes,
+        in_features: int,
+        out_features: int,
+        bias: Optional[torch.Tensor] = None,
+        device: Optional[torch.device | str] = None,
+    ) -> "GgufLinear":
+        """Convenience for tests/loaders: hydrate from raw GGUF block bytes."""
+        layer = cls(in_features=in_features, out_features=out_features,
+                    quant_type=quant_type, bias=bias is not None, device=device)
+        expected = layer.qweight.numel()
+        if len(qweight_bytes) != expected:
+            raise ValueError(
+                f"qweight_bytes length {len(qweight_bytes)} != expected {expected} "
+                f"for {quant_type} M={out_features} K={in_features}"
+            )
+        layer.qweight.copy_(torch.frombuffer(bytearray(qweight_bytes), dtype=torch.uint8))
+        if bias is not None:
+            layer.bias.copy_(bias.to(torch.float32))
+        return layer
+
+    # Back-compat shim for the existing Q4_K-only test
+    @classmethod
+    def from_q4_K_bytes(cls, qweight_bytes, in_features, out_features, bias=None, device=None):
+        return cls.from_quantized_bytes(
+            "Q4_K", qweight_bytes, in_features, out_features, bias=bias, device=device
+        )

build/torch28-metal-aarch64-darwin/transformers_glue.py

@@ -0,0 +1,284 @@
+"""Glue for plugging `GgufLinear` into a transformers model.
+
+The flow is intentionally non-invasive: we read raw block bytes directly from
+the GGUF file via `gguf.GGUFReader`, walk the model, and for every `nn.Linear`
+whose corresponding GGUF tensor is one of our supported quant types we
+in-place replace the module with `GgufLinear` and copy the bytes into
+`qweight`. Tensors that don't map to a Linear (or are unsupported quants like
+F32 norms, embeddings) get dequantized via `gguf.dequantize` and copied into
+the existing parameter as fp32.
+
+This avoids modifying the transformers loader. The trade-off is that the
+model has to be constructed first with regular `nn.Linear` shells (which we
+then mutate), but those shells are on the meta device or freshly initialised
+and don't carry the dequantized weight at the time of swap — so peak memory
+during load is bounded by `(model_size_fp32) + (model_size_quantized_qweights)`,
+which fits in RAM for the small models we care about.
+
+Typical use:
+
+    from transformers import AutoConfig, AutoModelForCausalLM
+    cfg = AutoConfig.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct-GGUF",
+                                     gguf_file="qwen2.5-0.5b-instruct-q4_0.gguf")
+    model = AutoModelForCausalLM.from_config(cfg).eval()
+    load_gguf_into_model_with_gguf_linear(
+        model,
+        gguf_path="…/qwen2.5-0.5b-instruct-q4_0.gguf",
+    )
+    out = model.generate(input_ids=…, max_new_tokens=10)
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import gguf
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .nn import GgufLinear, _QUANT_INFO
+
+
+# Map gguf.GGMLQuantizationType.name → our GgufLinear quant_type string.
+_GGUF_NAME_TO_GGUF_LINEAR = {
+    "Q4_0": "Q4_0",
+    "Q8_0": "Q8_0",
+    "Q4_K": "Q4_K",
+}
+
+
+def _quant_supported_by_gguf_linear(qt: gguf.GGMLQuantizationType) -> bool:
+    return qt.name in _GGUF_NAME_TO_GGUF_LINEAR
+
+
+# --- Tensor-name mapping ------------------------------------------------------
+# GGUF tensors use llama.cpp-style names like `blk.0.attn_q.weight`. HF models
+# use `model.layers.0.self_attn.q_proj.weight`. We use a tiny version of the
+# transformers rename table here so the swap util can stay self-contained.
+# (Transformers' own GGUF loader has a much more thorough version of this.)
+
+_BASIC_RENAMES = [
+    (r"^blk\.", "model.layers."),
+    (r"\.attn_q\.weight$",      ".self_attn.q_proj.weight"),
+    (r"\.attn_k\.weight$",      ".self_attn.k_proj.weight"),
+    (r"\.attn_v\.weight$",      ".self_attn.v_proj.weight"),
+    (r"\.attn_output\.weight$", ".self_attn.o_proj.weight"),
+    (r"\.ffn_gate\.weight$",    ".mlp.gate_proj.weight"),
+    (r"\.ffn_up\.weight$",      ".mlp.up_proj.weight"),
+    (r"\.ffn_down\.weight$",    ".mlp.down_proj.weight"),
+    (r"\.attn_norm\.weight$",   ".input_layernorm.weight"),
+    (r"\.ffn_norm\.weight$",    ".post_attention_layernorm.weight"),
+    (r"\.attn_(q|k|v)\.bias$",  r".self_attn.\1_proj.bias"),
+    (r"^token_embd\.weight$",   "model.embed_tokens.weight"),
+    (r"^output_norm\.weight$",  "model.norm.weight"),
+    (r"^output\.weight$",       "lm_head.weight"),
+]
+
+
+def _gguf_name_to_hf_name(gguf_name: str) -> str:
+    """Apply our minimal rename table to map a GGUF tensor name to an HF parameter name."""
+    import re
+
+    out = gguf_name
+    for pat, repl in _BASIC_RENAMES:
+        out = re.sub(pat, repl, out)
+    return out
+
+
+def _is_attn_qk(gguf_name: str) -> bool:
+    return ".attn_q.weight" in gguf_name or ".attn_k.weight" in gguf_name
+
+
+def _reverse_permute_qk(w: np.ndarray, n_head: int, n_kv: int | None) -> np.ndarray:
+    """GGUF Q/K weights are permuted on quantize; undo before loading into HF.
+    Reference: https://github.com/ggerganov/llama.cpp/blob/.../convert_hf_to_gguf.py."""
+    if n_kv is not None and n_head != n_kv:
+        n_head = n_kv
+    dim = w.shape[0] // n_head // 2
+    w = w.reshape(n_head, dim, 2, *w.shape[1:])
+    return w.swapaxes(2, 1).reshape(*w.shape[:1], *w.shape[3:]) if False else w.reshape(-1, *w.shape[3:]).swapaxes(2, 1).reshape(n_head * dim * 2, *w.shape[3:])
+
+
+def _set_module_param(model: nn.Module, qualified_name: str, value: torch.Tensor) -> bool:
+    """Set a parameter or buffer on `model` by dotted name. Returns True on success."""
+    parts = qualified_name.split(".")
+    mod = model
+    for p in parts[:-1]:
+        if not hasattr(mod, p):
+            return False
+        mod = getattr(mod, p)
+    leaf = parts[-1]
+    if leaf in mod._parameters:
+        param = mod._parameters[leaf]
+        if param is None:
+            return False
+        with torch.no_grad():
+            mod._parameters[leaf] = nn.Parameter(value.to(dtype=param.dtype, device=param.device),
+                                                 requires_grad=param.requires_grad)
+        return True
+    if leaf in mod._buffers:
+        with torch.no_grad():
+            mod._buffers[leaf] = value
+        return True
+    return False
+
+
+def _get_module_by_name(model: nn.Module, qualified_name: str) -> Optional[nn.Module]:
+    """Return the leaf module matching `qualified_name` minus its terminal attribute,
+    or None if anything along the path is missing."""
+    parts = qualified_name.split(".")
+    mod = model
+    for p in parts[:-1]:
+        if not hasattr(mod, p):
+            return None
+        mod = getattr(mod, p)
+    return mod
+
+
+def _get_attr_by_path(model: nn.Module, qualified_name: str) -> Optional[nn.Module]:
+    """Walk dotted `qualified_name` on `model` and return the final attribute,
+    or None if anything along the path is missing."""
+    if not qualified_name:
+        return model
+    mod = model
+    for p in qualified_name.split("."):
+        if not hasattr(mod, p):
+            return None
+        mod = getattr(mod, p)
+    return mod
+
+
+def _swap_linear_to_gguf(
+    parent: nn.Module, attr: str, in_features: int, out_features: int, quant_type: str,
+    bias_tensor: Optional[torch.Tensor], device: Optional[torch.device | str] = None
+) -> GgufLinear:
+    """Replace `parent.attr` (an nn.Linear) with a new GgufLinear of the same shape."""
+    new = GgufLinear(in_features=in_features, out_features=out_features,
+                     quant_type=quant_type, bias=bias_tensor is not None, device=device)
+    if bias_tensor is not None and new.bias is not None:
+        new.bias.data.copy_(bias_tensor.to(torch.float32))
+    setattr(parent, attr, new)
+    return new
+
+
+def load_gguf_into_model_with_gguf_linear(
+    model: nn.Module,
+    gguf_path: str | Path,
+    skip_modules: Optional[set[str]] = None,
+    force_dequant: bool = False,
+) -> dict:
+    """Read a GGUF file and populate `model` in-place.
+
+    Two-phase load to keep the byte path simple and the bias path correct:
+
+      Phase 1 — full dequant: every tensor (linear weights, norms, embeddings,
+        biases) is dequantized to fp32 and written into the matching HF
+        parameter. After this pass the model would already produce correct
+        outputs (this is what `force_dequant=True` stops at).
+      Phase 2 — swap linears: walk every `nn.Linear` that maps to a tensor
+        whose GGUF quant type is supported by `GgufLinear`; swap it for a
+        `GgufLinear` and re-attach raw block bytes (preserving the bias that
+        Phase 1 just loaded).
+
+    Returns a dict with diagnostic counts.
+    """
+    reader = gguf.GGUFReader(str(gguf_path))
+
+    # Pull a few config bits for the Q/K reverse-permute.
+    def _read_int(field: str, default: int = 0) -> int:
+        f = reader.fields.get(field)
+        if f is None:
+            return default
+        return int(f.parts[f.data[0]][0]) if f.data else default
+
+    arch = ""
+    arch_field = reader.fields.get("general.architecture")
+    if arch_field is not None and arch_field.data:
+        arch = bytes(arch_field.parts[arch_field.data[0]]).decode("utf-8", errors="ignore")
+    n_head    = _read_int(f"{arch}.attention.head_count")
+    n_kv_head = _read_int(f"{arch}.attention.head_count_kv") or n_head
+
+    skip_modules = skip_modules or set()
+    stats = {"as_gguf_linear": 0, "dequantized": 0, "skipped": 0, "unmatched": []}
+
+    # Track which (hf_name → quant_type) pairs are eligible for GgufLinear swap
+    # after Phase 1. We'll re-quantize the loaded fp32 weight in Phase 2.
+    swap_candidates: list[tuple[str, str]] = []
+
+    # ---------------- Phase 1: dequant everything into the model. -------------
+    for t in reader.tensors:
+        hf_name = _gguf_name_to_hf_name(t.name)
+        if any(skip in hf_name for skip in skip_modules):
+            stats["skipped"] += 1
+            continue
+
+        parent = _get_module_by_name(model, hf_name)
+        if parent is None or not hasattr(parent, hf_name.rsplit(".", 1)[-1]):
+            stats["unmatched"].append(t.name)
+            continue
+        leaf = hf_name.rsplit(".", 1)[-1]
+        cur = getattr(parent, leaf)
+
+        fp32 = gguf.dequantize(
+            np.frombuffer(np.ascontiguousarray(t.data).tobytes(), dtype=np.uint8),
+            t.tensor_type,
+        ).astype(np.float32)
+        if _is_attn_qk(t.name) and n_head > 0:
+            shape = list(t.shape)
+            M, K = int(shape[1]), int(shape[0])
+            heads = n_head if ".attn_q." in t.name else n_kv_head
+            head_dim = M // heads
+            fp32 = fp32.reshape(M, K).reshape(heads, head_dim // 2, 2, K).swapaxes(2, 1).reshape(M, K)
+        target_shape = tuple(cur.shape) if hasattr(cur, "shape") else (fp32.size,)
+        if np.prod(target_shape) == fp32.size:
+            fp32 = fp32.reshape(target_shape)
+        tensor = torch.from_numpy(fp32)
+        if not _set_module_param(model, hf_name, tensor):
+            stats["unmatched"].append(t.name)
+            continue
+        stats["dequantized"] += 1
+
+        # Record candidacy for swap in Phase 2.
+        is_linear_weight = isinstance(parent, nn.Linear) and leaf == "weight"
+        if is_linear_weight and _quant_supported_by_gguf_linear(t.tensor_type) and not force_dequant:
+            swap_candidates.append((hf_name, _GGUF_NAME_TO_GGUF_LINEAR[t.tensor_type.name]))
+
+    # ---------------- Phase 2: swap Linear → GgufLinear. ----------------------
+    # Re-quantize the loaded fp32 weight via gguf-py (works for Q4_0/Q8_0;
+    # K-quants don't have a Python quantizer yet — those will skip swap).
+    for hf_name, quant_type in swap_candidates:
+        # hf_name ends in `.weight`; the Linear module is at hf_name[:-len(".weight")]
+        linear_path = hf_name[: -len(".weight")]
+        grandparent_path, _, linear_attr = linear_path.rpartition(".")
+        grandparent = model if not grandparent_path else _get_attr_by_path(model, grandparent_path)
+        linear = _get_attr_by_path(model, linear_path)
+        if grandparent is None or not isinstance(linear, nn.Linear):
+            continue
+
+        # Re-quantize the (now permuted-if-applicable) fp32 weight.
+        try:
+            qt = getattr(gguf.GGMLQuantizationType, quant_type)
+            qbytes = bytes(
+                gguf.quantize(linear.weight.detach().cpu().numpy().astype(np.float32), qt).tobytes()
+            )
+        except (NotImplementedError, Exception) as e:
+            # K-quants etc. — leave as Linear (fp32 weight is already loaded).
+            stats.setdefault("re_quantize_failed", []).append((hf_name, str(e)))
+            continue
+
+        bias_tensor = linear.bias.detach().clone() if linear.bias is not None else None
+        new = GgufLinear(
+            in_features=linear.in_features,
+            out_features=linear.out_features,
+            quant_type=quant_type,
+            bias=bias_tensor is not None,
+        )
+        if bias_tensor is not None:
+            new.bias.data.copy_(bias_tensor.to(torch.float32))
+        new.qweight.copy_(torch.frombuffer(bytearray(qbytes), dtype=torch.uint8))
+        setattr(grandparent, linear_attr, new)
+        stats["as_gguf_linear"] += 1
+        stats["dequantized"] -= 1  # this one's no longer "just" dequant — swapped
+
+    return stats

build/torch29-metal-aarch64-darwin/__init__.py

@@ -0,0 +1,600 @@
+"""GGUF dequantization kernels for Apple Silicon (Metal).
+
+Six standalone kernels covering the GGUF formats most common in the wild:
+
+    Q4_0   Q8_0   Q4_K   Q5_K   Q6_K   IQ4_NL   IQ4_XS
+
+Two API levels are exposed for each:
+
+  - ``dequantize_<format>_bytes(quant_bytes) -> np.ndarray``
+    One-shot path used by the tests for bit-exact validation against a
+    NumPy reference. Allocates input/output buffers internally.
+
+  - ``allocate_<format>_plan(nblocks)`` + ``dispatch_<format>(plan)``
+    Low-level path for benchmarks. Buffers are pre-allocated; ``dispatch``
+    only encodes + dispatches + syncs, so the timing isolates GPU work
+    from CPU<->buffer copies.
+
+The runtime path uses PyObjC + Metal.framework so the correctness loop
+works on any Apple Silicon Mac without Xcode 26 / Nix / kernel-builder.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+# Lazy globals — populated on first use.
+_DEVICE = None
+_COMMAND_QUEUE = None
+_LIBRARY = None
+_PIPELINES: dict = {}
+
+_KERNEL_SOURCE_PATH = Path(__file__).resolve().parent.parent.parent / "gguf_dequant_metal" / "dequantize.metal"
+
+
+def _ensure_metal():
+    global _DEVICE, _COMMAND_QUEUE, _LIBRARY
+    if _DEVICE is not None:
+        return
+    import Metal  # pyobjc-framework-metal
+
+    _DEVICE = Metal.MTLCreateSystemDefaultDevice()
+    if _DEVICE is None:
+        raise RuntimeError("No default Metal device available")
+    _COMMAND_QUEUE = _DEVICE.newCommandQueue()
+
+    src = _KERNEL_SOURCE_PATH.read_text()
+    options = Metal.MTLCompileOptions.new()
+    lib, err = _DEVICE.newLibraryWithSource_options_error_(src, options, None)
+    if lib is None:
+        raise RuntimeError(f"Metal compile failed: {err}")
+    _LIBRARY = lib
+
+
+def _pipeline(fn_name: str):
+    _ensure_metal()
+    pso = _PIPELINES.get(fn_name)
+    if pso is not None:
+        return pso
+
+    fn = _LIBRARY.newFunctionWithName_(fn_name)
+    if fn is None:
+        raise RuntimeError(f"Kernel function not found: {fn_name}")
+    pso, err = _DEVICE.newComputePipelineStateWithFunction_error_(fn, None)
+    if pso is None:
+        raise RuntimeError(f"Pipeline state creation failed: {err}")
+    _PIPELINES[fn_name] = pso
+    return pso
+
+
+# (kernel_name, bytes_per_block, elems_per_block, threads_per_block_group)
+_FORMATS = {
+    "q4_0":   ("dequantize_q4_0",   18, 32, 16),
+    "q8_0":   ("dequantize_q8_0",   34, 32, 32),
+    "q4_K":   ("dequantize_q4_K",  144, 256, 256),
+    "q5_K":   ("dequantize_q5_K",  176, 256, 256),
+    "q6_K":   ("dequantize_q6_K",  210, 256, 256),
+    "iq4_nl": ("dequantize_iq4_nl", 18, 32, 16),
+    "iq4_xs": ("dequantize_iq4_xs", 136, 256, 256),
+}
+
+
+class Plan:
+    """Pre-allocated I/O buffers for a single (format, nblocks) shape."""
+
+    __slots__ = ("fmt", "nblocks", "in_buf", "out_buf", "_nblocks_bytes")
+
+    def __init__(self, fmt: str, nblocks: int, in_buf, out_buf):
+        self.fmt = fmt
+        self.nblocks = nblocks
+        self.in_buf = in_buf
+        self.out_buf = out_buf
+        self._nblocks_bytes = np.array([nblocks], dtype=np.uint32).tobytes()
+
+
+def allocate_plan(fmt: str, nblocks: int) -> Plan:
+    import Metal
+
+    _ensure_metal()
+    _, bpb, epb, _ = _FORMATS[fmt]
+    in_buf = _DEVICE.newBufferWithLength_options_(nblocks * bpb, Metal.MTLResourceStorageModeShared)
+    out_buf = _DEVICE.newBufferWithLength_options_(nblocks * epb * 4, Metal.MTLResourceStorageModeShared)
+    return Plan(fmt=fmt, nblocks=nblocks, in_buf=in_buf, out_buf=out_buf)
+
+
+def dispatch(plan: Plan, sync: bool = True):
+    import Metal
+
+    kernel_name, _, _, tg = _FORMATS[plan.fmt]
+    pso = _pipeline(kernel_name)
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.in_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.out_buf, 0, 1)
+    enc.setBytes_length_atIndex_(plan._nblocks_bytes, len(plan._nblocks_bytes), 2)
+    enc.dispatchThreads_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(plan.nblocks * tg, 1, 1),
+        Metal.MTLSizeMake(tg, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+def dequantize_bytes(fmt: str, quant_bytes) -> np.ndarray:
+    """One-shot bytes-in / float32-out for the named format."""
+    _ensure_metal()
+    _, bpb, epb, _ = _FORMATS[fmt]
+
+    buf = bytes(quant_bytes) if not isinstance(quant_bytes, (bytes, bytearray)) else quant_bytes
+    nbytes = len(buf)
+    if nbytes % bpb != 0:
+        raise ValueError(f"{fmt} input not a multiple of {bpb} bytes: got {nbytes}")
+    nblocks = nbytes // bpb
+
+    plan = allocate_plan(fmt, nblocks)
+    in_view = np.frombuffer(plan.in_buf.contents().as_buffer(nbytes), dtype=np.uint8)
+    in_view[:] = np.frombuffer(buf, dtype=np.uint8)
+    dispatch(plan, sync=True)
+    out_nbytes = nblocks * epb * 4
+    return np.frombuffer(plan.out_buf.contents().as_buffer(out_nbytes), dtype=np.float32).copy()
+
+
+# --- Per-format convenience wrappers ----------------------------------------
+
+def dequantize_q4_0_bytes(qb):   return dequantize_bytes("q4_0",   qb)
+def dequantize_q8_0_bytes(qb):   return dequantize_bytes("q8_0",   qb)
+def dequantize_q4_K_bytes(qb):   return dequantize_bytes("q4_K",   qb)
+def dequantize_q5_K_bytes(qb):   return dequantize_bytes("q5_K",   qb)
+def dequantize_q6_K_bytes(qb):   return dequantize_bytes("q6_K",   qb)
+def dequantize_iq4_nl_bytes(qb): return dequantize_bytes("iq4_nl", qb)
+def dequantize_iq4_xs_bytes(qb): return dequantize_bytes("iq4_xs", qb)
+
+def allocate_q4_0_plan(nblocks):   return allocate_plan("q4_0",   nblocks)
+def allocate_q8_0_plan(nblocks):   return allocate_plan("q8_0",   nblocks)
+def allocate_q4_K_plan(nblocks):   return allocate_plan("q4_K",   nblocks)
+def allocate_q5_K_plan(nblocks):   return allocate_plan("q5_K",   nblocks)
+def allocate_q6_K_plan(nblocks):   return allocate_plan("q6_K",   nblocks)
+def allocate_iq4_nl_plan(nblocks): return allocate_plan("iq4_nl", nblocks)
+def allocate_iq4_xs_plan(nblocks): return allocate_plan("iq4_xs", nblocks)
+
+def dispatch_q4_0(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q8_0(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q4_K(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q5_K(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_q6_K(plan, **kw):   return dispatch(plan, **kw)
+def dispatch_iq4_nl(plan, **kw): return dispatch(plan, **kw)
+def dispatch_iq4_xs(plan, **kw): return dispatch(plan, **kw)
+
+
+# ---------------------------------------------------------------------------
+# Torch op entry points — the production path used by transformers (no PyObjC,
+# no runtime compile). Only available in builds produced by kernel-builder,
+# which generates `_ops.py` exporting `ops = torch.ops._gguf_dequant_<rev>`.
+# In a dev-mode checkout (running from torch-ext/ directly) `_ops` isn't
+# present, so the torch_* entry points raise on call.
+# ---------------------------------------------------------------------------
+
+try:
+    from ._ops import ops as _ops  # generated by kernel-builder
+except ImportError:
+    _ops = None
+
+
+def _require_ops():
+    if _ops is None:
+        raise RuntimeError(
+            "torch op binding not available — this build was loaded from the dev tree. "
+            "Use the *_bytes / *_plan APIs (PyObjC path) or install a kernel-builder build."
+        )
+    return _ops
+
+
+def torch_dequantize_q4_0(input, output):   _require_ops().dequantize_q4_0(input, output)
+def torch_dequantize_q8_0(input, output):   _require_ops().dequantize_q8_0(input, output)
+def torch_dequantize_q4_K(input, output):   _require_ops().dequantize_q4_K(input, output)
+def torch_dequantize_q5_K(input, output):   _require_ops().dequantize_q5_K(input, output)
+def torch_dequantize_q6_K(input, output):   _require_ops().dequantize_q6_K(input, output)
+def torch_dequantize_iq4_nl(input, output): _require_ops().dequantize_iq4_nl(input, output)
+def torch_dequantize_iq4_xs(input, output): _require_ops().dequantize_iq4_xs(input, output)
+
+
+# =============================================================================
+# Q4_K × fp32 matmul — `GgufLinear` ceiling experiment.
+# =============================================================================
+# Layout convention (matches llama.cpp's kernel_mul_mm):
+#   A (Q4_K weights):  shape (M, K)  M = out_features, K = in_features.
+#                      M and K must be multiples of 64 and 256 respectively for
+#                      the simdgroup tile geometry to land cleanly.
+#   B (fp32 input):    shape (N, K)  N = batch * seqlen, contiguous in K.
+#   C (fp32 output):   shape (M, N)  column-major-ish (matches llama.cpp).
+
+_BLOCK_SIZE_M = 64
+_BLOCK_SIZE_N = 32
+_THREAD_PER_BLOCK = 128
+_SHARED_BYTES = 8192
+
+
+class MatMulQ4KPlan:
+    __slots__ = ("M", "N", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, M, N, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.M = M
+        self.N = N
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def allocate_matmul_q4_K_plan(M: int, N: int, K: int) -> MatMulQ4KPlan:
+    """Pre-allocate weight (Q4_K), input (fp32), output (fp32) MTLBuffers."""
+    import Metal
+
+    _ensure_metal()
+    assert M % _BLOCK_SIZE_M == 0, f"M must be multiple of {_BLOCK_SIZE_M}, got {M}"
+    assert N % _BLOCK_SIZE_N == 0, f"N must be multiple of {_BLOCK_SIZE_N}, got {N}"
+    assert K % 256 == 0, f"K must be multiple of 256, got {K}"
+
+    a_bytes = M * (K // 256) * 144  # M rows × (K/QK_K) super-blocks × 144 bytes
+    b_bytes = N * K * 4             # fp32
+    c_bytes = M * N * 4             # fp32
+
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+
+    # Pack ne00, ne01, ne11, nb01, nb11 into a uniforms buffer.
+    nb01 = (K // 256) * 144  # bytes per A row
+    nb11 = K * 4             # bytes per B row
+    dims = np.array([K], dtype=np.int64).tobytes()  # ne00
+    dims += np.array([M], dtype=np.int64).tobytes()  # ne01
+    dims += np.array([N], dtype=np.int64).tobytes()  # ne11
+    dims += np.array([nb01], dtype=np.uint64).tobytes()
+    dims += np.array([nb11], dtype=np.uint64).tobytes()
+
+    return MatMulQ4KPlan(M=M, N=N, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def dispatch_matmul_q4_K(plan: MatMulQ4KPlan, sync: bool = True):
+    """Dispatch mul_mat_q4_K_f32 on the pre-allocated buffers."""
+    import Metal
+
+    pso = _pipeline("mul_mat_q4_K_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    # ne00 (K), ne01 (M), ne11 (N), nb01, nb11 via setBytes — 5 scalars
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8],   8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16],  8, 4)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[16:24], 8, 5)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[24:32], 8, 6)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[32:40], 8, 7)
+    enc.setThreadgroupMemoryLength_atIndex_(_SHARED_BYTES, 0)
+
+    grid_x = (plan.N + _BLOCK_SIZE_N - 1) // _BLOCK_SIZE_N
+    grid_y = (plan.M + _BLOCK_SIZE_M - 1) // _BLOCK_SIZE_M
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, grid_y, 1),
+        Metal.MTLSizeMake(_THREAD_PER_BLOCK, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+# =============================================================================
+# mul_mat_vec_q4_K — batch-1 decode (memory-bandwidth-bound regime).
+# =============================================================================
+
+_N_DST = 4
+_SIMD_WIDTH = 32
+
+
+class MatVecQ4KPlan:
+    __slots__ = ("M", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, M, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.M = M
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def allocate_matvec_q4_K_plan(M: int, K: int) -> MatVecQ4KPlan:
+    import Metal
+
+    _ensure_metal()
+    assert M % _N_DST == 0, f"M must be a multiple of {_N_DST}, got {M}"
+    assert K % 256 == 0, f"K must be a multiple of 256, got {K}"
+
+    a_bytes = M * (K // 256) * 144
+    b_bytes = K * 4
+    c_bytes = M * 4
+
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+
+    dims = np.array([K], dtype=np.int64).tobytes() + np.array([M], dtype=np.int64).tobytes()
+    return MatVecQ4KPlan(M=M, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def dispatch_matvec_q4_K(plan: MatVecQ4KPlan, sync: bool = True):
+    import Metal
+
+    pso = _pipeline("mul_mat_vec_q4_K_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8], 8, 3)   # ne00 = K
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16], 8, 4)  # ne01 = M
+
+    grid_x = (plan.M + _N_DST - 1) // _N_DST
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, 1, 1),
+        Metal.MTLSizeMake(_SIMD_WIDTH, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+# =============================================================================
+# Q4_0 / Q8_0 matmul + matvec — same layout as Q4_K, smaller blocks (32 elems).
+# =============================================================================
+
+_N_SIMDGROUP = 2  # Q4_0/Q8_0 matvec uses 2 simdgroups per threadgroup
+
+# Per quant type for the mat path: (block bytes, block elems, matvec_M_alignment).
+# block_elems = 32 for Q4_0/Q8_0/IQ4_NL, 256 for all K-quants and IQ4_XS.
+# matvec_M_alignment:
+#   Q4_0/Q8_0  →  2 simdgroups × N_DST(4) = 8 rows per threadgroup
+#   Q4_K       →  1 simdgroup  × N_DST(4) = 4 rows per threadgroup
+#   Q5_K       →  2 simdgroups × 2 rows   = 4 rows per threadgroup
+#   Q6_K       →  2 simdgroups × 1 row    = 2 rows per threadgroup
+_MAT_FORMATS = {
+    "q4_0":   (18,  32,  8),
+    "q8_0":   (34,  32,  8),
+    "q4_K":   (144, 256, 4),
+    "q5_K":   (176, 256, 4),
+    "q6_K":   (210, 256, 2),
+    "iq4_nl": (18,  32,  4),       # 2 sg × 2 rows/sg = 4 rows/threadgroup
+    "iq4_xs": (136, 256, 4),
+}
+
+
+class MatMulPlan:
+    """Pre-allocated I/O buffers for `mul_mat_<fmt>_f32`. Block element size
+    depends on the format (32 for Q4_0/Q8_0, 256 for Q4_K)."""
+
+    __slots__ = ("fmt", "M", "N", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, fmt, M, N, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.fmt = fmt
+        self.M = M
+        self.N = N
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def _quant_block_elems(fmt: str) -> int:
+    return _MAT_FORMATS[fmt][1]
+
+
+def _quant_block_bytes(fmt: str) -> int:
+    return _MAT_FORMATS[fmt][0]
+
+
+def _matvec_row_alignment(fmt: str) -> int | None:
+    return _MAT_FORMATS[fmt][2]
+
+
+def allocate_matmul_plan(fmt: str, M: int, N: int, K: int) -> MatMulPlan:
+    import Metal
+
+    _ensure_metal()
+    assert fmt in _MAT_FORMATS, f"unsupported fmt {fmt}"
+    block_elems = _quant_block_elems(fmt)
+    block_bytes = _quant_block_bytes(fmt)
+    assert M % _BLOCK_SIZE_M == 0, f"M must be multiple of {_BLOCK_SIZE_M}, got {M}"
+    assert N % _BLOCK_SIZE_N == 0, f"N must be multiple of {_BLOCK_SIZE_N}, got {N}"
+    assert K % block_elems == 0, f"K must be multiple of {block_elems} for {fmt}, got {K}"
+
+    a_bytes = M * (K // block_elems) * block_bytes
+    b_bytes = N * K * 4
+    c_bytes = M * N * 4
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+
+    nb01 = (K // block_elems) * block_bytes
+    nb11 = K * 4
+    dims = (
+        np.array([K], dtype=np.int64).tobytes()
+        + np.array([M], dtype=np.int64).tobytes()
+        + np.array([N], dtype=np.int64).tobytes()
+        + np.array([nb01], dtype=np.uint64).tobytes()
+        + np.array([nb11], dtype=np.uint64).tobytes()
+    )
+    return MatMulPlan(fmt=fmt, M=M, N=N, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def dispatch_matmul(plan: MatMulPlan, sync: bool = True):
+    import Metal
+
+    pso = _pipeline(f"mul_mat_{plan.fmt}_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8],   8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16],  8, 4)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[16:24], 8, 5)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[24:32], 8, 6)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[32:40], 8, 7)
+    enc.setThreadgroupMemoryLength_atIndex_(_SHARED_BYTES, 0)
+
+    grid_x = (plan.N + _BLOCK_SIZE_N - 1) // _BLOCK_SIZE_N
+    grid_y = (plan.M + _BLOCK_SIZE_M - 1) // _BLOCK_SIZE_M
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, grid_y, 1),
+        Metal.MTLSizeMake(_THREAD_PER_BLOCK, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd
+
+
+class MatVecPlan:
+    __slots__ = ("fmt", "M", "K", "a_buf", "b_buf", "c_buf", "_dims_bytes")
+
+    def __init__(self, fmt, M, K, a_buf, b_buf, c_buf, dims_bytes):
+        self.fmt = fmt
+        self.M = M
+        self.K = K
+        self.a_buf = a_buf
+        self.b_buf = b_buf
+        self.c_buf = c_buf
+        self._dims_bytes = dims_bytes
+
+
+def allocate_matvec_plan(fmt: str, M: int, K: int) -> MatVecPlan:
+    """Allocate I/O for `mul_mat_vec_<fmt>_f32`. Format-specific row alignment:
+    Q4_0/Q8_0: 8, Q4_K: 4, Q5_K: 4, Q6_K: 2."""
+    import Metal
+
+    _ensure_metal()
+    block_elems = _quant_block_elems(fmt)
+    block_bytes = _quant_block_bytes(fmt)
+    align = _matvec_row_alignment(fmt)
+    if align is None:
+        raise ValueError(f"mul_mat_vec not implemented for {fmt}")
+    assert M % align == 0, f"M must be a multiple of {align} for {fmt}, got {M}"
+    assert K % block_elems == 0, f"K must be a multiple of {block_elems} for {fmt}, got {K}"
+
+    a_bytes = M * (K // block_elems) * block_bytes
+    b_bytes = K * 4
+    c_bytes = M * 4
+    a_buf = _DEVICE.newBufferWithLength_options_(a_bytes, Metal.MTLResourceStorageModeShared)
+    b_buf = _DEVICE.newBufferWithLength_options_(b_bytes, Metal.MTLResourceStorageModeShared)
+    c_buf = _DEVICE.newBufferWithLength_options_(c_bytes, Metal.MTLResourceStorageModeShared)
+    dims = np.array([K], dtype=np.int64).tobytes() + np.array([M], dtype=np.int64).tobytes()
+    return MatVecPlan(fmt=fmt, M=M, K=K, a_buf=a_buf, b_buf=b_buf, c_buf=c_buf, dims_bytes=dims)
+
+
+def encode_matvec(plan: MatVecPlan, enc):
+    """Encode a single matvec dispatch onto an existing MTLComputeCommandEncoder.
+    Caller is responsible for committing the parent command buffer. Threadgroup
+    memory for the IQ4 lookup table is set via setThreadgroupMemoryLength here."""
+    import Metal
+
+    pso = _pipeline(f"mul_mat_vec_{plan.fmt}_f32")
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8], 8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16], 8, 4)
+
+    if plan.fmt == "q4_K":
+        rows_per_tg = _N_DST
+        tg_threads = _SIMD_WIDTH
+    elif plan.fmt in ("q4_0", "q8_0"):
+        rows_per_tg = _N_SIMDGROUP * _N_DST
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q5_K":
+        rows_per_tg = _N_SIMDGROUP * 2
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q6_K":
+        rows_per_tg = _N_SIMDGROUP * 1
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt in ("iq4_nl", "iq4_xs"):
+        rows_per_tg = _N_SIMDGROUP * 2
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+        enc.setThreadgroupMemoryLength_atIndex_(256, 0)
+    else:
+        raise ValueError(f"matvec not wired for {plan.fmt}")
+    grid_x = (plan.M + rows_per_tg - 1) // rows_per_tg
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, 1, 1),
+        Metal.MTLSizeMake(tg_threads, 1, 1),
+    )
+
+
+def new_command_buffer():
+    """Start a fresh MTLCommandBuffer. Caller batches `encode_matvec(plan, enc)`
+    onto its encoder, then calls `cb.commit()` and `cb.waitUntilCompleted()`
+    once for the whole batch — matching llama.cpp's per-step semantics."""
+    _ensure_metal()
+    return _COMMAND_QUEUE.commandBuffer()
+
+
+def dispatch_matvec(plan: MatVecPlan, sync: bool = True):
+    import Metal
+
+    pso = _pipeline(f"mul_mat_vec_{plan.fmt}_f32")
+    cmd = _COMMAND_QUEUE.commandBuffer()
+    enc = cmd.computeCommandEncoder()
+    enc.setComputePipelineState_(pso)
+    enc.setBuffer_offset_atIndex_(plan.a_buf, 0, 0)
+    enc.setBuffer_offset_atIndex_(plan.b_buf, 0, 1)
+    enc.setBuffer_offset_atIndex_(plan.c_buf, 0, 2)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[0:8], 8, 3)
+    enc.setBytes_length_atIndex_(plan._dims_bytes[8:16], 8, 4)
+
+    needs_shared = plan.fmt in ("iq4_nl", "iq4_xs")
+    if plan.fmt == "q4_K":
+        rows_per_tg = _N_DST
+        tg_threads = _SIMD_WIDTH
+    elif plan.fmt in ("q4_0", "q8_0"):
+        rows_per_tg = _N_SIMDGROUP * _N_DST
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q5_K":
+        rows_per_tg = _N_SIMDGROUP * 2
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt == "q6_K":
+        rows_per_tg = _N_SIMDGROUP * 1
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    elif plan.fmt in ("iq4_nl", "iq4_xs"):
+        rows_per_tg = _N_SIMDGROUP * 2            # 2 sg × 2 rows/sg
+        tg_threads = _N_SIMDGROUP * _SIMD_WIDTH
+    else:
+        raise ValueError(f"matvec not wired for {plan.fmt}")
+    grid_x = (plan.M + rows_per_tg - 1) // rows_per_tg
+    if needs_shared:
+        enc.setThreadgroupMemoryLength_atIndex_(256, 0)  # 64 floats per simdgroup × 2 sg
+    enc.dispatchThreadgroups_threadsPerThreadgroup_(
+        Metal.MTLSizeMake(grid_x, 1, 1),
+        Metal.MTLSizeMake(tg_threads, 1, 1),
+    )
+    enc.endEncoding()
+    cmd.commit()
+    if sync:
+        cmd.waitUntilCompleted()
+    return cmd

build/torch29-metal-aarch64-darwin/_gguf_dequant_530e374.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2345a3efd557e8757ae50488e360d5bbf79e5163c658a497f918d2e9727082a0
+size 236288

build/torch29-metal-aarch64-darwin/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _gguf_dequant_530e374
+ops = torch.ops._gguf_dequant_530e374
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_gguf_dequant_530e374::{op_name}"

build/torch29-metal-aarch64-darwin/gguf_dequant/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+from pathlib import Path
+from types import ModuleType
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch29-metal-aarch64-darwin/metadata.json

@@ -0,0 +1,4 @@
+{
+  "version": 1,
+  "python-depends": []
+}

build/torch29-metal-aarch64-darwin/nn.py

@@ -0,0 +1,191 @@
+"""GgufLinear — drop-in nn.Linear replacement that runs quantized matmul directly.
+
+Forward pass:
+  - If the input has effective batch size 1 (single token), use mul_mat_vec
+    (memory-bandwidth-bound, faster than bf16 mat-vec on MPS).
+  - Otherwise use mul_mat (compute-bound, similar throughput to bf16 GEMM but
+    keeps weights at native quantization → 3–4× less RAM than bf16).
+
+Supported quant types today: Q4_0, Q8_0, Q4_K. Adding more is just a matter of
+porting the candle kernels — the dispatch logic is generic over `quant_type`.
+
+Prototype design: weights are stored as raw uint8 bytes in `qweight`. The
+forward path copies activations into a shared MTLBuffer, runs the Metal
+kernel, reads the fp32 output back. Production path (kernels-community +
+PR #767) skips the copies via direct MTLBuffer handoff.
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+import torch
+
+from . import (
+    allocate_matmul_plan,
+    allocate_matvec_plan,
+    dispatch_matmul,
+    dispatch_matvec,
+)
+
+
+# Per-quant-type geometry: (block_bytes, block_elems, matvec_M_alignment)
+_QUANT_INFO = {
+    "Q4_0": (18, 32, 8),   # matvec uses 2 simdgroups × N_DST=4 = 8 rows/threadgroup
+    "Q8_0": (34, 32, 8),
+    "Q4_K": (144, 256, 4), # matvec uses 1 simdgroup × N_DST=4 = 4 rows/threadgroup
+}
+
+# Map our Python-side quant name to the Metal kernel suffix.
+_KERNEL_FMT = {
+    "Q4_0": "q4_0",
+    "Q8_0": "q8_0",
+    "Q4_K": "q4_K",
+}
+
+
+class GgufLinear(torch.nn.Module):
+    """Linear layer backed by GGUF-quantized weights.
+
+    Args:
+        in_features: K dimension. Must be a multiple of the format's block size
+            (32 for Q4_0/Q8_0, 256 for Q4_K).
+        out_features: M dimension. Must be a multiple of the format's matvec
+            row alignment (8 for Q4_0/Q8_0, 4 for Q4_K) and of 64 for matmul.
+        quant_type: one of "Q4_0", "Q8_0", "Q4_K".
+        bias: whether to add a learnable bias after the matmul.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        quant_type: str = "Q4_K",
+        bias: bool = False,
+        device: Optional[torch.device | str] = None,
+    ):
+        super().__init__()
+        if quant_type not in _QUANT_INFO:
+            raise NotImplementedError(
+                f"GgufLinear supports {sorted(_QUANT_INFO)} today, got {quant_type}"
+            )
+        block_bytes, block_elems, matvec_align = _QUANT_INFO[quant_type]
+        if in_features % block_elems != 0:
+            raise ValueError(
+                f"in_features must be a multiple of {block_elems} for {quant_type}, got {in_features}"
+            )
+        if out_features % matvec_align != 0:
+            raise ValueError(
+                f"out_features must be a multiple of {matvec_align} for {quant_type}, got {out_features}"
+            )
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.quant_type = quant_type
+        self._fmt = _KERNEL_FMT[quant_type]
+        self._block_bytes = block_bytes
+        self._block_elems = block_elems
+
+        nblocks_per_row = in_features // block_elems
+        qweight_nbytes = out_features * nblocks_per_row * block_bytes
+        self.register_buffer(
+            "qweight",
+            torch.empty(qweight_nbytes, dtype=torch.uint8, device=device or "cpu"),
+            persistent=True,
+        )
+
+        if bias:
+            self.bias = torch.nn.Parameter(
+                torch.zeros(out_features, dtype=torch.float32, device=device or "cpu"),
+                requires_grad=False,
+            )
+        else:
+            self.register_parameter("bias", None)
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"quant_type={self.quant_type}, bias={self.bias is not None}"
+        )
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: shape (..., K). Returns (..., M) fp32."""
+        if x.shape[-1] != self.in_features:
+            raise ValueError(f"Last dim of x must be {self.in_features}, got {x.shape[-1]}")
+
+        batch_shape = x.shape[:-1]
+        x_flat = x.reshape(-1, self.in_features).contiguous().to(torch.float32)
+        N = x_flat.shape[0]
+
+        # CPU view of qweight for the bytes copy into shared MTLBuffers.
+        # Production path skips this once the kernels-community build is wired.
+        qweight_cpu = self.qweight.detach().to("cpu", copy=False)
+        x_cpu = x_flat.to("cpu", copy=False)
+
+        if N == 1:
+            out = self._forward_matvec(qweight_cpu, x_cpu.reshape(-1))
+        else:
+            out = self._forward_matmul(qweight_cpu, x_cpu)
+
+        out = out.reshape(*batch_shape, self.out_features)
+        if x.device != out.device:
+            out = out.to(x.device)
+        if self.bias is not None:
+            out = out + self.bias.to(out.device).to(out.dtype)
+        return out
+
+    def _forward_matvec(self, qweight_cpu: torch.Tensor, x_cpu: torch.Tensor) -> torch.Tensor:
+        M, K = self.out_features, self.in_features
+        plan = allocate_matvec_plan(self._fmt, M=M, K=K)
+        a_view = np.frombuffer(plan.a_buf.contents().as_buffer(qweight_cpu.numel()), dtype=np.uint8)
+        a_view[:] = qweight_cpu.numpy()
+        b_view = np.frombuffer(plan.b_buf.contents().as_buffer(K * 4), dtype=np.float32)
+        b_view[:] = x_cpu.numpy()
+        dispatch_matvec(plan, sync=True)
+        c_view = np.frombuffer(plan.c_buf.contents().as_buffer(M * 4), dtype=np.float32)
+        return torch.from_numpy(c_view.copy())
+
+    def _forward_matmul(self, qweight_cpu: torch.Tensor, x_cpu: torch.Tensor) -> torch.Tensor:
+        N, K = x_cpu.shape
+        M = self.out_features
+        plan = allocate_matmul_plan(self._fmt, M=M, N=N, K=K)
+        a_view = np.frombuffer(plan.a_buf.contents().as_buffer(qweight_cpu.numel()), dtype=np.uint8)
+        a_view[:] = qweight_cpu.numpy()
+        b_view = np.frombuffer(plan.b_buf.contents().as_buffer(N * K * 4), dtype=np.float32)
+        b_view[:] = x_cpu.numpy().reshape(-1)
+        dispatch_matmul(plan, sync=True)
+        c_view = np.frombuffer(plan.c_buf.contents().as_buffer(M * N * 4), dtype=np.float32)
+        # Kernel output layout is (N, M) in row-major; we want (N, M).
+        return torch.from_numpy(c_view.reshape(N, M).copy())
+
+    @classmethod
+    def from_quantized_bytes(
+        cls,
+        quant_type: str,
+        qweight_bytes: bytes,
+        in_features: int,
+        out_features: int,
+        bias: Optional[torch.Tensor] = None,
+        device: Optional[torch.device | str] = None,
+    ) -> "GgufLinear":
+        """Convenience for tests/loaders: hydrate from raw GGUF block bytes."""
+        layer = cls(in_features=in_features, out_features=out_features,
+                    quant_type=quant_type, bias=bias is not None, device=device)
+        expected = layer.qweight.numel()
+        if len(qweight_bytes) != expected:
+            raise ValueError(
+                f"qweight_bytes length {len(qweight_bytes)} != expected {expected} "
+                f"for {quant_type} M={out_features} K={in_features}"
+            )
+        layer.qweight.copy_(torch.frombuffer(bytearray(qweight_bytes), dtype=torch.uint8))
+        if bias is not None:
+            layer.bias.copy_(bias.to(torch.float32))
+        return layer
+
+    # Back-compat shim for the existing Q4_K-only test
+    @classmethod
+    def from_q4_K_bytes(cls, qweight_bytes, in_features, out_features, bias=None, device=None):
+        return cls.from_quantized_bytes(
+            "Q4_K", qweight_bytes, in_features, out_features, bias=bias, device=device
+        )

build/torch29-metal-aarch64-darwin/transformers_glue.py

@@ -0,0 +1,284 @@
+"""Glue for plugging `GgufLinear` into a transformers model.
+
+The flow is intentionally non-invasive: we read raw block bytes directly from
+the GGUF file via `gguf.GGUFReader`, walk the model, and for every `nn.Linear`
+whose corresponding GGUF tensor is one of our supported quant types we
+in-place replace the module with `GgufLinear` and copy the bytes into
+`qweight`. Tensors that don't map to a Linear (or are unsupported quants like
+F32 norms, embeddings) get dequantized via `gguf.dequantize` and copied into
+the existing parameter as fp32.
+
+This avoids modifying the transformers loader. The trade-off is that the
+model has to be constructed first with regular `nn.Linear` shells (which we
+then mutate), but those shells are on the meta device or freshly initialised
+and don't carry the dequantized weight at the time of swap — so peak memory
+during load is bounded by `(model_size_fp32) + (model_size_quantized_qweights)`,
+which fits in RAM for the small models we care about.
+
+Typical use:
+
+    from transformers import AutoConfig, AutoModelForCausalLM
+    cfg = AutoConfig.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct-GGUF",
+                                     gguf_file="qwen2.5-0.5b-instruct-q4_0.gguf")
+    model = AutoModelForCausalLM.from_config(cfg).eval()
+    load_gguf_into_model_with_gguf_linear(
+        model,
+        gguf_path="…/qwen2.5-0.5b-instruct-q4_0.gguf",
+    )
+    out = model.generate(input_ids=…, max_new_tokens=10)
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import gguf
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .nn import GgufLinear, _QUANT_INFO
+
+
+# Map gguf.GGMLQuantizationType.name → our GgufLinear quant_type string.
+_GGUF_NAME_TO_GGUF_LINEAR = {
+    "Q4_0": "Q4_0",
+    "Q8_0": "Q8_0",
+    "Q4_K": "Q4_K",
+}
+
+
+def _quant_supported_by_gguf_linear(qt: gguf.GGMLQuantizationType) -> bool:
+    return qt.name in _GGUF_NAME_TO_GGUF_LINEAR
+
+
+# --- Tensor-name mapping ------------------------------------------------------
+# GGUF tensors use llama.cpp-style names like `blk.0.attn_q.weight`. HF models
+# use `model.layers.0.self_attn.q_proj.weight`. We use a tiny version of the
+# transformers rename table here so the swap util can stay self-contained.
+# (Transformers' own GGUF loader has a much more thorough version of this.)
+
+_BASIC_RENAMES = [
+    (r"^blk\.", "model.layers."),
+    (r"\.attn_q\.weight$",      ".self_attn.q_proj.weight"),
+    (r"\.attn_k\.weight$",      ".self_attn.k_proj.weight"),
+    (r"\.attn_v\.weight$",      ".self_attn.v_proj.weight"),
+    (r"\.attn_output\.weight$", ".self_attn.o_proj.weight"),
+    (r"\.ffn_gate\.weight$",    ".mlp.gate_proj.weight"),
+    (r"\.ffn_up\.weight$",      ".mlp.up_proj.weight"),
+    (r"\.ffn_down\.weight$",    ".mlp.down_proj.weight"),
+    (r"\.attn_norm\.weight$",   ".input_layernorm.weight"),
+    (r"\.ffn_norm\.weight$",    ".post_attention_layernorm.weight"),
+    (r"\.attn_(q|k|v)\.bias$",  r".self_attn.\1_proj.bias"),
+    (r"^token_embd\.weight$",   "model.embed_tokens.weight"),
+    (r"^output_norm\.weight$",  "model.norm.weight"),
+    (r"^output\.weight$",       "lm_head.weight"),
+]
+
+
+def _gguf_name_to_hf_name(gguf_name: str) -> str:
+    """Apply our minimal rename table to map a GGUF tensor name to an HF parameter name."""
+    import re
+
+    out = gguf_name
+    for pat, repl in _BASIC_RENAMES:
+        out = re.sub(pat, repl, out)
+    return out
+
+
+def _is_attn_qk(gguf_name: str) -> bool:
+    return ".attn_q.weight" in gguf_name or ".attn_k.weight" in gguf_name
+
+
+def _reverse_permute_qk(w: np.ndarray, n_head: int, n_kv: int | None) -> np.ndarray:
+    """GGUF Q/K weights are permuted on quantize; undo before loading into HF.
+    Reference: https://github.com/ggerganov/llama.cpp/blob/.../convert_hf_to_gguf.py."""
+    if n_kv is not None and n_head != n_kv:
+        n_head = n_kv
+    dim = w.shape[0] // n_head // 2
+    w = w.reshape(n_head, dim, 2, *w.shape[1:])
+    return w.swapaxes(2, 1).reshape(*w.shape[:1], *w.shape[3:]) if False else w.reshape(-1, *w.shape[3:]).swapaxes(2, 1).reshape(n_head * dim * 2, *w.shape[3:])
+
+
+def _set_module_param(model: nn.Module, qualified_name: str, value: torch.Tensor) -> bool:
+    """Set a parameter or buffer on `model` by dotted name. Returns True on success."""
+    parts = qualified_name.split(".")
+    mod = model
+    for p in parts[:-1]:
+        if not hasattr(mod, p):
+            return False
+        mod = getattr(mod, p)
+    leaf = parts[-1]
+    if leaf in mod._parameters:
+        param = mod._parameters[leaf]
+        if param is None:
+            return False
+        with torch.no_grad():
+            mod._parameters[leaf] = nn.Parameter(value.to(dtype=param.dtype, device=param.device),
+                                                 requires_grad=param.requires_grad)
+        return True
+    if leaf in mod._buffers:
+        with torch.no_grad():
+            mod._buffers[leaf] = value
+        return True
+    return False
+
+
+def _get_module_by_name(model: nn.Module, qualified_name: str) -> Optional[nn.Module]:
+    """Return the leaf module matching `qualified_name` minus its terminal attribute,
+    or None if anything along the path is missing."""
+    parts = qualified_name.split(".")
+    mod = model
+    for p in parts[:-1]:
+        if not hasattr(mod, p):
+            return None
+        mod = getattr(mod, p)
+    return mod
+
+
+def _get_attr_by_path(model: nn.Module, qualified_name: str) -> Optional[nn.Module]:
+    """Walk dotted `qualified_name` on `model` and return the final attribute,
+    or None if anything along the path is missing."""
+    if not qualified_name:
+        return model
+    mod = model
+    for p in qualified_name.split("."):
+        if not hasattr(mod, p):
+            return None
+        mod = getattr(mod, p)
+    return mod
+
+
+def _swap_linear_to_gguf(
+    parent: nn.Module, attr: str, in_features: int, out_features: int, quant_type: str,
+    bias_tensor: Optional[torch.Tensor], device: Optional[torch.device | str] = None
+) -> GgufLinear:
+    """Replace `parent.attr` (an nn.Linear) with a new GgufLinear of the same shape."""
+    new = GgufLinear(in_features=in_features, out_features=out_features,
+                     quant_type=quant_type, bias=bias_tensor is not None, device=device)
+    if bias_tensor is not None and new.bias is not None:
+        new.bias.data.copy_(bias_tensor.to(torch.float32))
+    setattr(parent, attr, new)
+    return new
+
+
+def load_gguf_into_model_with_gguf_linear(
+    model: nn.Module,
+    gguf_path: str | Path,
+    skip_modules: Optional[set[str]] = None,
+    force_dequant: bool = False,
+) -> dict:
+    """Read a GGUF file and populate `model` in-place.
+
+    Two-phase load to keep the byte path simple and the bias path correct:
+
+      Phase 1 — full dequant: every tensor (linear weights, norms, embeddings,
+        biases) is dequantized to fp32 and written into the matching HF
+        parameter. After this pass the model would already produce correct
+        outputs (this is what `force_dequant=True` stops at).
+      Phase 2 — swap linears: walk every `nn.Linear` that maps to a tensor
+        whose GGUF quant type is supported by `GgufLinear`; swap it for a
+        `GgufLinear` and re-attach raw block bytes (preserving the bias that
+        Phase 1 just loaded).
+
+    Returns a dict with diagnostic counts.
+    """
+    reader = gguf.GGUFReader(str(gguf_path))
+
+    # Pull a few config bits for the Q/K reverse-permute.
+    def _read_int(field: str, default: int = 0) -> int:
+        f = reader.fields.get(field)
+        if f is None:
+            return default
+        return int(f.parts[f.data[0]][0]) if f.data else default
+
+    arch = ""
+    arch_field = reader.fields.get("general.architecture")
+    if arch_field is not None and arch_field.data:
+        arch = bytes(arch_field.parts[arch_field.data[0]]).decode("utf-8", errors="ignore")
+    n_head    = _read_int(f"{arch}.attention.head_count")
+    n_kv_head = _read_int(f"{arch}.attention.head_count_kv") or n_head
+
+    skip_modules = skip_modules or set()
+    stats = {"as_gguf_linear": 0, "dequantized": 0, "skipped": 0, "unmatched": []}
+
+    # Track which (hf_name → quant_type) pairs are eligible for GgufLinear swap
+    # after Phase 1. We'll re-quantize the loaded fp32 weight in Phase 2.
+    swap_candidates: list[tuple[str, str]] = []
+
+    # ---------------- Phase 1: dequant everything into the model. -------------
+    for t in reader.tensors:
+        hf_name = _gguf_name_to_hf_name(t.name)
+        if any(skip in hf_name for skip in skip_modules):
+            stats["skipped"] += 1
+            continue
+
+        parent = _get_module_by_name(model, hf_name)
+        if parent is None or not hasattr(parent, hf_name.rsplit(".", 1)[-1]):
+            stats["unmatched"].append(t.name)
+            continue
+        leaf = hf_name.rsplit(".", 1)[-1]
+        cur = getattr(parent, leaf)
+
+        fp32 = gguf.dequantize(
+            np.frombuffer(np.ascontiguousarray(t.data).tobytes(), dtype=np.uint8),
+            t.tensor_type,
+        ).astype(np.float32)
+        if _is_attn_qk(t.name) and n_head > 0:
+            shape = list(t.shape)
+            M, K = int(shape[1]), int(shape[0])
+            heads = n_head if ".attn_q." in t.name else n_kv_head
+            head_dim = M // heads
+            fp32 = fp32.reshape(M, K).reshape(heads, head_dim // 2, 2, K).swapaxes(2, 1).reshape(M, K)
+        target_shape = tuple(cur.shape) if hasattr(cur, "shape") else (fp32.size,)
+        if np.prod(target_shape) == fp32.size:
+            fp32 = fp32.reshape(target_shape)
+        tensor = torch.from_numpy(fp32)
+        if not _set_module_param(model, hf_name, tensor):
+            stats["unmatched"].append(t.name)
+            continue
+        stats["dequantized"] += 1
+
+        # Record candidacy for swap in Phase 2.
+        is_linear_weight = isinstance(parent, nn.Linear) and leaf == "weight"
+        if is_linear_weight and _quant_supported_by_gguf_linear(t.tensor_type) and not force_dequant:
+            swap_candidates.append((hf_name, _GGUF_NAME_TO_GGUF_LINEAR[t.tensor_type.name]))
+
+    # ---------------- Phase 2: swap Linear → GgufLinear. ----------------------
+    # Re-quantize the loaded fp32 weight via gguf-py (works for Q4_0/Q8_0;
+    # K-quants don't have a Python quantizer yet — those will skip swap).
+    for hf_name, quant_type in swap_candidates:
+        # hf_name ends in `.weight`; the Linear module is at hf_name[:-len(".weight")]
+        linear_path = hf_name[: -len(".weight")]
+        grandparent_path, _, linear_attr = linear_path.rpartition(".")
+        grandparent = model if not grandparent_path else _get_attr_by_path(model, grandparent_path)
+        linear = _get_attr_by_path(model, linear_path)
+        if grandparent is None or not isinstance(linear, nn.Linear):
+            continue
+
+        # Re-quantize the (now permuted-if-applicable) fp32 weight.
+        try:
+            qt = getattr(gguf.GGMLQuantizationType, quant_type)
+            qbytes = bytes(
+                gguf.quantize(linear.weight.detach().cpu().numpy().astype(np.float32), qt).tobytes()
+            )
+        except (NotImplementedError, Exception) as e:
+            # K-quants etc. — leave as Linear (fp32 weight is already loaded).
+            stats.setdefault("re_quantize_failed", []).append((hf_name, str(e)))
+            continue
+
+        bias_tensor = linear.bias.detach().clone() if linear.bias is not None else None
+        new = GgufLinear(
+            in_features=linear.in_features,
+            out_features=linear.out_features,
+            quant_type=quant_type,
+            bias=bias_tensor is not None,
+        )
+        if bias_tensor is not None:
+            new.bias.data.copy_(bias_tensor.to(torch.float32))
+        new.qweight.copy_(torch.frombuffer(bytearray(qbytes), dtype=torch.uint8))
+        setattr(grandparent, linear_attr, new)
+        stats["as_gguf_linear"] += 1
+        stats["dequantized"] -= 1  # this one's no longer "just" dequant — swapped
+
+    return stats