trim the metal shader to the dispatched kernels

.gitattributes

@@ -56,3 +56,5 @@ build/torch212-cxx11-cu132-x86_64-linux/_gguf_kernels_cuda_1df3b99.abi3.so filte
 build/torch213-cxx11-cu126-x86_64-linux/_gguf_kernels_cuda_1df3b99.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch213-cxx11-cu130-x86_64-linux/_gguf_kernels_cuda_1df3b99.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch213-cxx11-cu132-x86_64-linux/_gguf_kernels_cuda_1df3b99.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch212-metal-aarch64-darwin/_gguf_kernels_metal_7b10656.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch213-metal-aarch64-darwin/_gguf_kernels_metal_7b10656.abi3.so filter=lfs diff=lfs merge=lfs -text

build.toml

@@ -87,9 +87,11 @@ src = [
     "gguf_metal/gguf_metal.cpp",
     "gguf_metal/ggml_dispatch.mm",
     "gguf_metal/common.h",
-    # ggml's shader, compiled into the embedded metallib. Its mul_mv_* kernels are what the gemv
-    # dispatches to; the K-quants have no dequantize-to-dense kernel upstream, so that one is ours.
-    "vendor/src/ggml-metal/ggml-metal.metal",
+    # ggml's shader, cut by `trim_shader.py` to the three kernels the dispatch names per quantization
+    # type -- 628 exported kernels down to 56, since none of ggml's attention, norm, conv or sampling
+    # kernels are reachable from here. `vendor/` keeps the untouched copy it is generated from.
+    "gguf_metal/ggml-metal-quant.metal",
+    "gguf_metal/quant_types.h",
     "vendor/src/ggml-metal/ggml-metal-impl.h",
     "vendor/src/ggml-common.h",
 ]

build/torch212-metal-aarch64-darwin/_gguf_kernels_metal_7b10656.abi3.so

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "gguf-kernels",
-  "id": "_gguf_kernels_metal_0becb0a",
+  "id": "_gguf_kernels_metal_7b10656",
   "version": 1,
   "license": "MIT",
   "python-depends": [],
@@ -11,8 +11,8 @@
     "algorithm": "sha256",
     "files": {
       "__init__.py": "GoFcHJbbJ9C8b40XNznMxsTcC+CV9kaFE8a1DP763kU=",
-      "_gguf_kernels_metal_0becb0a.abi3.so": "s2nMW0aOL9O5rIT1xjiXpSvCXgOBw3EE0R0JVvWtTdk=",
-      "_ops.py": "3htLDPFBO8zVroZkTqPmryb9pDPjZ5OuMItEBg94Jp4="
+      "_gguf_kernels_metal_7b10656.abi3.so": "1eoCqrNTAf9aN42G3mJHe+zWYEzVbvYACALBVnfJN10=",
+      "_ops.py": "nwCRMS2OmSxach0Edkf8TLdxZXmrrMhMJntgxLN2MAI="
     }
   },
   "provenance": {
@@ -22,7 +22,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "0becb0aafe12b27da05fa08f4b6cc7c72b288fbb",
+      "sha": "7b106564513dead0e4416ea0c7d3898eb23e0ca2",
       "dirty": false
     }
   }

build/torch213-metal-aarch64-darwin/_gguf_kernels_metal_7b10656.abi3.so

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "gguf-kernels",
-  "id": "_gguf_kernels_metal_0becb0a",
+  "id": "_gguf_kernels_metal_7b10656",
   "version": 1,
   "license": "MIT",
   "python-depends": [],
@@ -11,8 +11,8 @@
     "algorithm": "sha256",
     "files": {
       "__init__.py": "GoFcHJbbJ9C8b40XNznMxsTcC+CV9kaFE8a1DP763kU=",
-      "_gguf_kernels_metal_0becb0a.abi3.so": "WdDCcX9EMIvJZbguUgzORbCm9+SWvJ6V38I2fbx19EE=",
-      "_ops.py": "3htLDPFBO8zVroZkTqPmryb9pDPjZ5OuMItEBg94Jp4="
+      "_gguf_kernels_metal_7b10656.abi3.so": "5tuy6u/VjXaPip4w6JsI8fehOWSLEvuJJKNMrMvY0GY=",
+      "_ops.py": "nwCRMS2OmSxach0Edkf8TLdxZXmrrMhMJntgxLN2MAI="
     }
   },
   "provenance": {
@@ -22,7 +22,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "0becb0aafe12b27da05fa08f4b6cc7c72b288fbb",
+      "sha": "7b106564513dead0e4416ea0c7d3898eb23e0ca2",
       "dirty": false
     }
   }

gguf_metal/quant_types.h

@@ -0,0 +1,8 @@
+#pragma once
+// Generated by trim_shader.py from llama.cpp @ 432d7ffe2c3b -- do not edit.
+
+namespace gguf_quant_types {
+
+inline constexpr const char *built[] = {"iq1_m", "iq1_s", "iq2_s", "iq2_xs", "iq2_xxs", "iq3_s", "iq3_xxs", "iq4_nl", "iq4_xs", "mxfp4", "q2_K", "q3_K", "q4_0", "q4_1", "q4_K", "q5_0", "q5_1", "q5_K", "q6_K", "q8_0"};
+
+}  // namespace gguf_quant_types

trim_shader.py

@@ -0,0 +1,198 @@
+"""Cut ggml's Metal shader down to the kernels this package can actually dispatch, and emit the
+table `ggml_dispatch.mm` uses to decide what it supports.
+
+ggml's Metal backend is one translation unit: `vendor/src/ggml-metal/ggml-metal.metal` exports 603
+kernels, of which we name 25. There is no upstream target for a subset, but the bulk of the file is
+templates plus a block of `template [[host_name(...)]]` instantiations -- and an uninstantiated
+template emits no code. So dropping the instantiations we cannot name is enough: 6.84 MB -> 0.64 MB
+of metallib, with the same 52 tests passing.
+
+`vendor/` stays a byte-for-byte copy of upstream; the trimmed shader is generated beside the backend
+source and committed, because a hub `kernels` build compiles the `src` list in build.toml and has
+nowhere to run a pre-step. Both outputs are generated -- never hand-edit them, re-run this instead
+(`vendor.py` does, after every pin bump).
+
+The dispatcher names kernels as `kernel_flash_attn_ext[_vec]_f32_dk<K>_dv<V>`, so the f16, bf16 and
+five quantized-KV families upstream also instantiates are unreachable here. The dk/dv pairs that do
+exist are parsed out rather than hand-copied, which is what keeps `supports_flash_attn` from
+claiming a shape whose pipeline was never built.
+
+Usage: python trim_shader.py [--no-trim]
+"""
+
+import argparse
+import os
+import re
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+SHADER_IN = os.path.join(HERE, "vendor", "src", "ggml-metal", "ggml-metal.metal")
+SHADER_OUT = os.path.join(HERE, "gguf_metal", "ggml-metal-quant.metal")
+HEADER_OUT = os.path.join(HERE, "gguf_metal", "quant_types.h")
+
+# What `ggml_dispatch.mm` builds with snprintf. Everything else upstream exports is dropped.
+# The three kernels `ggml_dispatch.mm` names per quantization type, and the types its `type_table()`
+# carries -- parsed from the dispatch rather than repeated here, so the two cannot drift.
+def _dispatched_types():
+    source = open(os.path.join(HERE, "gguf_metal", "ggml_dispatch.mm")).read()
+    table = source[source.index("const std::unordered_map<int, TypeInfo> &type_table()") :]
+    return re.findall(r'\{GGML_\w+, \{"(\w+)"', table[: table.index("\n}")])
+
+
+TYPES = _dispatched_types()
+DISPATCHABLE = re.compile(
+    r"^kernel_(?:mul_mv|mul_mm)_(" + "|".join(TYPES) + r")_f32$|^kernel_get_rows_(" + "|".join(TYPES) + r")$"
+)
+GENERATED_BY = "trim_shader.py"
+
+
+# Kernels declared plainly, with no `host_name` to instantiate them -- rwkv, ssm, pooling, upscale,
+# the adamw steps, the quantized mul_mv family. Every `kernel void` is an exported entry point, so
+# they survive on their own; nothing can call them from here, since a kernel cannot call a kernel.
+#
+# A templated one is a different thing with the same spelling: it emits nothing until instantiated,
+# and its body is what the `host_name` lines above refer to. Dropping those leaves the `template<>`
+# prefix dangling on whatever follows, so they are left alone and cost nothing.
+PLAIN_KERNEL = re.compile(r"^kernel void (\w+)\(")
+
+
+def _referenced_names(lines):
+    """Kernel names other lines instantiate or alias: templates the file still needs, not leaf kernels."""
+    referring = "".join(l for l in lines if "decltype(" in l or "host_name" in l)
+    return set(re.findall(r"\b(kernel_\w+)\b", referring))
+
+
+def _declaration_start(out, name):
+    """Rewind over the `template<...>` header and comments that belong to the definition being cut.
+
+    A dropped definition whose header stays behind reparents that header onto whatever follows, so
+    the header has to go with it -- as do the comment lines written for it.
+    """
+    while out:
+        last = out[-1].strip()
+        if last.startswith("//") or not last:
+            out.pop()
+            continue
+        # A `template<...>` header belongs to the definition being cut. A `template [[host_name(..)]] ... ;`
+        # instantiation does not -- it is a whole declaration of its own, and eating it silently drops a
+        # kernel that was meant to be kept.
+        if last.startswith("template") and "host_name" not in last and not last.endswith(";"):
+            out.pop()
+            continue
+        break
+    return out
+
+
+def _skip_statement(lines, i):
+    """Advance past the statement starting at `i`: to its `;`, or past the body if it has one."""
+    depth, in_body = 0, False
+    while i < len(lines):
+        depth += lines[i].count("{") - lines[i].count("}")
+        in_body = in_body or "{" in lines[i]
+        ends = (in_body and depth <= 0) or (not in_body and lines[i].rstrip().endswith(";"))
+        i += 1
+        if ends:
+            return i
+    return i
+
+
+def trim(lines):
+    """Drop each exported kernel we cannot name, statement by statement.
+
+    Four shapes export a kernel: a one-line template instantiation; `template [[host_name(..)]]`
+    with the instantiation on following lines; `[[host_name(..)]]` on a full kernel definition; and
+    a plain `kernel void name(...)` definition with no attribute at all.
+    """
+    referenced = _referenced_names(lines)
+    kept_names, out, i = [], [], 0
+    while i < len(lines):
+        plain = PLAIN_KERNEL.match(lines[i])
+        if plain and not any(k in plain.group(1) for k in ("mul_mv", "mul_mm", "get_rows")) and plain.group(1) not in referenced:
+            _declaration_start(out, plain.group(1))
+            i = _skip_statement(lines, i)
+            # `typedef decltype(name<...>) name_t;` exists only to name the thing just dropped.
+            while i < len(lines) and f"decltype({plain.group(1)}" in lines[i]:
+                i += 1
+            continue
+        found = re.search(r'host_name\("([^"]+)"', lines[i])
+        if found and not DISPATCHABLE.match(found.group(1)):
+            if lines[i].rstrip().endswith(";"):
+                i += 1
+                continue
+            i = _skip_statement(lines, i)
+            continue
+        if found:
+            kept_names.append(found.group(1))
+        out.append(lines[i])
+        i += 1
+    return out, kept_names
+
+
+def header(names, revision):
+    """The quantization types this build carries kernels for, parsed from the shader."""
+    types = sorted({m.group(1) or m.group(2) for name in names if (m := DISPATCHABLE.match(name))})
+    entries = ", ".join(f'"{t}"' for t in types)
+    return f"""\
+#pragma once
+// Generated by {GENERATED_BY} from llama.cpp @ {revision} -- do not edit.
+
+namespace gguf_quant_types {{
+
+inline constexpr const char *built[] = {{{entries}}};
+
+}}  // namespace gguf_quant_types
+"""
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    ap.add_argument(
+        "--no-trim",
+        action="store_true",
+        help="copy the shader verbatim instead of cutting it. The trim is an optimisation, so it must "
+        "never be what blocks a pin bump: if upstream reshapes its declarations and the cut stops "
+        "compiling, ship the whole 6.84 MB shader with this and fix the cut afterwards.",
+    )
+    args = ap.parse_args()
+
+    with open(SHADER_IN) as f:
+        lines = f.readlines()
+    revision = "unknown"
+    upstream = os.path.join(HERE, "vendor", "UPSTREAM")
+    if os.path.exists(upstream):
+        with open(upstream) as f:
+            revision = f.read().split("\n")[1][:12]
+
+    out, kept = (lines, [m.group(1) for l in lines
+                         for m in [re.search(r'host_name\("([^"]+)"', l)]
+                         if m and DISPATCHABLE.match(m.group(1))]) if args.no_trim else trim(lines)
+
+    # The dispatcher builds names for both paths and calls three helpers by name; a keep-list that
+    # lost any of them would build clean and fail at runtime, so refuse to emit it.
+    # Every type the dispatch can ask for must have all three of its kernels, or a model using that
+    # quantization fails at its first forward instead of at build time.
+    # Against the emitted shader, not the instantiation list: some of these are plain definitions rather
+    # than `host_name` instantiations, so they never appear in `kept`.
+    body = "".join(out)
+    missing = [f"{prefix}{t}{suffix}"
+               for t in TYPES
+               for prefix, suffix in (("kernel_mul_mv_", "_f32"), ("kernel_mul_mm_", "_f32"),
+                                      ("kernel_get_rows_", ""))
+               if f"{prefix}{t}{suffix}" not in body]
+    if missing:
+        raise SystemExit(f"refusing to write: kept {len(kept)} kernels, missing {len(missing)}: "
+                         f"{missing[:4]}. Upstream's shader has changed shape -- rerun with --no-trim.")
+    with open(SHADER_OUT, "w") as f:
+        f.write(f"// Generated by {GENERATED_BY} from vendor/src/ggml-metal/ggml-metal.metal "
+                f"(llama.cpp @ {revision}) -- do not edit.\n")
+        f.writelines(out)
+    with open(HEADER_OUT, "w") as f:
+        f.write(header(kept, revision))
+
+    dropped = sum("host_name" in line for line in lines) - len(kept)
+    print(f"kept {len(kept)} kernels over {len(TYPES)} quantization types, dropped {dropped}")
+    print(f"  {SHADER_OUT}")
+    print(f"  {HEADER_OUT}")
+
+
+if __name__ == "__main__":
+    main()