Metal: cover every quant type upstream's metallib has a gemv for

The type table was four entries; the metallib always held kernels for 22. The gap was host-side:
CUDA reaches upstream's `mul_mat_vec_q_switch_type`, one function that switches over every type,
while ggml's Metal type->pipeline mapping lives in ggml-metal-device.cpp, which pulls in the whole
backend and is not compiled here. So each type has to be registered by hand.

Transcribed from ggml_metal_library_get_pipeline_mul_mv: the legacy quants, the K quants, the IQ
quants and MXFP4 -- 20 types, parity with the CUDA port apart from NVFP4/Q1_0/Q2_0, which the Metal
metallib has no kernels for. The one field not available as a macro is the grid: ggml-metal-ops.cpp
dispatches the dense types and q8_0 on (ne01 + nr0 - 1)/nr0 and everything else divides by nsg too,
so among quantized types q8_0 stays the only `mv_reduce_across_sgs`.

`GEMV_TYPES` becomes a `gemv_types()` op each backend answers for itself. It was one hardcoded CUDA
list, which Metal inherited and over-claimed by 19 types -- a type routed into a gemv with no kernel
for it faults rather than falling back.

tests: run on MPS as well as CUDA and cover every supported type. Dequantize is compared against
the block scale rather than each element -- 17 of the 20 come back bit-exact, Q4_0/Q4_1/Q4_K land
within an ulp of f32, and asking a near-zero element to match to its own magnitude asks for an
exactness a different summation order breaks. MXFP4's scale is one E8M0 byte, not an fp16 field, so
the harness's mask left the reference itself overflowing to inf.

test_vendor_drift: the table transcribes constants out of vendor/, and `vendor.py --rev` replaces
those files. A renamed N_SG_*/N_R0_* fails the build, but a changed smem or grid rule compiles and
silently returns garbage. This parses upstream and fails loudly instead.

gguf_cuda/gguf_cuda.cu

@@ -35,6 +35,13 @@ int dtype_code(at::ScalarType t) {
 
 }  // namespace
 
+// The cases upstream's `mul_mat_vec_q_switch_type` dispatches (mmvq.cu): Q4_0/Q4_1/Q5_0/Q5_1/Q8_0,
+// the K quants, the IQ quants, MXFP4/NVFP4 and Q1_0/Q2_0. Kept as a literal because that switch is
+// static and exposes no way to enumerate it; a type added upstream shows up here when the pin moves.
+std::vector<int64_t> gemv_types() {
+  return {2, 3, 6, 7, 8, 10, 11, 12, 13, 14, 16, 17, 18, 19, 20, 21, 22, 23, 29, 39, 40, 41, 42};
+}
+
 at::Tensor dequantize(const at::Tensor &blocks, int64_t ggml_type, int64_t rows, int64_t cols,
                       at::ScalarType dtype) {
   TORCH_CHECK(blocks.is_cuda() && blocks.scalar_type() == at::kByte, "blocks must be cuda uint8");

gguf_metal/common.h

@@ -28,4 +28,8 @@ int gguf_metal_get_rows(void *blocks, size_t blocks_off, void *indices, size_t i
 
 // True when the type has a gemv/gemm kernel in the metallib.
 int gguf_metal_supports(int ggml_type);
+
+// The ggml type ids above, written into `out` (up to `max`); returns how many there are. Called
+// once at import to publish `GEMV_TYPES`, so the caller never has to guess this backend's coverage.
+int gguf_metal_gemv_types(int *out, int max);
 }

gguf_metal/ggml_dispatch.mm

@@ -30,7 +30,10 @@ namespace {
 
 // ggml type ids, matching ggml.h
 enum { GGML_Q4_0 = 2, GGML_Q4_1 = 3, GGML_Q5_0 = 6, GGML_Q5_1 = 7, GGML_Q8_0 = 8,
-       GGML_Q2_K = 10, GGML_Q3_K = 11, GGML_Q4_K = 12, GGML_Q5_K = 13, GGML_Q6_K = 14 };
+       GGML_Q2_K = 10, GGML_Q3_K = 11, GGML_Q4_K = 12, GGML_Q5_K = 13, GGML_Q6_K = 14,
+       GGML_IQ2_XXS = 16, GGML_IQ2_XS = 17, GGML_IQ3_XXS = 18, GGML_IQ1_S = 19,
+       GGML_IQ4_NL = 20, GGML_IQ3_S = 21, GGML_IQ2_S = 22, GGML_IQ4_XS = 23,
+       GGML_IQ1_M = 29, GGML_MXFP4 = 39 };
 
 struct TypeInfo {
   const char *name;   // the infix in kernel_mul_mv_<name>_f32
@@ -38,21 +41,47 @@ struct TypeInfo {
   int block_bytes;
   int mv_nsg;         // N_SG_*  from ggml-metal-impl.h
   int mv_nr0;         // N_R0_*
-  // Two gemv conventions upstream. The K-quants give each simdgroup its own rows
-  // (first_row = (tgpig.x*NSG + sgitg)*nr0), so the grid is divided by nsg and nothing is shared.
-  // q8_0 instead splits the K reduction across simdgroups (r0 = tgpig.x*NR0) and combines them
-  // through threadgroup memory in helper_mv_reduce_and_write -- so its grid must NOT be divided by
-  // nsg, and it needs smem. Getting either wrong silently returns garbage.
+  // Two gemv grids upstream, chosen by type in ggml-metal-ops.cpp's mul_mv dispatch: the dense
+  // types and q8_0 take (ne01 + nr0 - 1)/nr0, everything else divides by nsg as well. q8_0 splits
+  // the K reduction across simdgroups (r0 = tgpig.x*NR0) and combines them through threadgroup
+  // memory in helper_mv_reduce_and_write, so dividing its grid by nsg would drop rows; the others
+  // give each simdgroup its own rows (first_row = (tgpig.x*NSG + sgitg)*nr0) and share nothing.
+  // Among quantized types q8_0 is the only one on the first grid. Getting it wrong returns garbage
+  // rather than failing.
   bool mv_reduce_across_sgs;
+  // ggml_metal_library_get_pipeline_mul_mv's `smem`, per type. Mostly 0; the IQ types cache their
+  // lookup grid in threadgroup memory and q8_0 needs a reduction buffer.
   size_t mv_smem;
 };
 
+// Transcribed from ggml_metal_library_get_pipeline_mul_mv (ggml-metal-device.cpp): every type it
+// implements, minus the dense ones (a GGUF weight this reaches is always quantized) and nvfp4,
+// which upstream's Metal backend has no kernel for. Block sizes are ggml's GGML_QUANT_SIZES.
 const std::unordered_map<int, TypeInfo> &type_table() {
   static const std::unordered_map<int, TypeInfo> table = {
+      // legacy quants
+      {GGML_Q4_0, {"q4_0", 32, 18, N_SG_Q4_0, N_R0_Q4_0, false, 0}},
+      {GGML_Q4_1, {"q4_1", 32, 20, N_SG_Q4_1, N_R0_Q4_1, false, 0}},
+      {GGML_Q5_0, {"q5_0", 32, 22, N_SG_Q5_0, N_R0_Q5_0, false, 0}},
+      {GGML_Q5_1, {"q5_1", 32, 24, N_SG_Q5_1, N_R0_Q5_1, false, 0}},
       {GGML_Q8_0, {"q8_0", 32, 34, N_SG_Q8_0, N_R0_Q8_0, true, 32 * sizeof(float) * N_R0_Q8_0}},
+      // K quants
+      {GGML_Q2_K, {"q2_K", 256, 84, N_SG_Q2_K, N_R0_Q2_K, false, 0}},
+      {GGML_Q3_K, {"q3_K", 256, 110, N_SG_Q3_K, N_R0_Q3_K, false, 0}},
       {GGML_Q4_K, {"q4_K", 256, 144, N_SG_Q4_K, N_R0_Q4_K, false, 0}},
       {GGML_Q5_K, {"q5_K", 256, 176, N_SG_Q5_K, N_R0_Q5_K, false, 0}},
       {GGML_Q6_K, {"q6_K", 256, 210, N_SG_Q6_K, N_R0_Q6_K, false, 0}},
+      // IQ quants
+      {GGML_IQ2_XXS, {"iq2_xxs", 256, 66, N_SG_IQ2_XXS, N_R0_IQ2_XXS, false, 256 * 8 + 128}},
+      {GGML_IQ2_XS, {"iq2_xs", 256, 74, N_SG_IQ2_XS, N_R0_IQ2_XS, false, 512 * 8 + 128}},
+      {GGML_IQ3_XXS, {"iq3_xxs", 256, 98, N_SG_IQ3_XXS, N_R0_IQ3_XXS, false, 256 * 4 + 128}},
+      {GGML_IQ1_S, {"iq1_s", 256, 50, N_SG_IQ1_S, N_R0_IQ1_S, false, 0}},
+      {GGML_IQ4_NL, {"iq4_nl", 32, 18, N_SG_IQ4_NL, N_R0_IQ4_NL, false, 32 * sizeof(float)}},
+      {GGML_IQ3_S, {"iq3_s", 256, 110, N_SG_IQ3_S, N_R0_IQ3_S, false, 512 * 4}},
+      {GGML_IQ2_S, {"iq2_s", 256, 82, N_SG_IQ2_S, N_R0_IQ2_S, false, 0}},
+      {GGML_IQ4_XS, {"iq4_xs", 256, 136, N_SG_IQ4_XS, N_R0_IQ4_XS, false, 32 * sizeof(float)}},
+      {GGML_IQ1_M, {"iq1_m", 256, 56, N_SG_IQ1_M, N_R0_IQ1_M, false, 0}},
+      {GGML_MXFP4, {"mxfp4", 32, 17, N_SG_MXFP4, N_R0_MXFP4, false, 32 * sizeof(float)}},
   };
   return table;
 }
@@ -134,6 +163,17 @@ void set_bool(MTLFunctionConstantValues *cv, bool value, NSUInteger index) {
 
 extern "C" int gguf_metal_supports(int ggml_type) { return lookup(ggml_type) != nullptr; }
 
+extern "C" int gguf_metal_gemv_types(int *out, int max) {
+  int n = 0;
+  for (const auto &entry : type_table()) {
+    if (n < max) {
+      out[n] = entry.first;
+    }
+    ++n;
+  }
+  return n;
+}
+
 extern "C" int gguf_metal_mul_mat(void *blocks, size_t blocks_off, void *x, size_t x_off, void *out,
                                   size_t out_off, int ggml_type, int64_t K, int64_t N, int64_t M) {
   const TypeInfo *info = lookup(ggml_type);

gguf_metal/gguf_metal.cpp

@@ -25,6 +25,13 @@ size_t byte_offset(const at::Tensor &t) {
 
 }  // namespace
 
+std::vector<int64_t> gemv_types() {
+  int ids[64];
+  const int n = gguf_metal_gemv_types(ids, 64);
+  TORCH_CHECK(n <= 64, "gguf-kernels: the Metal type table outgrew the buffer here");
+  return std::vector<int64_t>(ids, ids + n);
+}
+
 at::Tensor dequantize(const at::Tensor &blocks, int64_t ggml_type, int64_t rows, int64_t cols,
                       at::ScalarType dtype) {
   TORCH_CHECK(blocks.is_mps() && blocks.scalar_type() == at::kByte, "blocks must be mps uint8");

tests/test_gguf_kernels.py

@@ -10,57 +10,86 @@ import numpy as np
 import pytest
 import torch
 
-from gguf_kernels import MAX_GEMV_ROWS, dequantize, mul_mat_vec
+from gguf_kernels import GEMV_TYPES, MAX_GEMV_ROWS, dequantize, mul_mat_vec
 
 
 gguf = pytest.importorskip("gguf", reason="the reference unpacker comes from the `gguf` package")
 
-# name -> (ggml type id, values per block, bytes per block)
+DEVICE = "cuda" if torch.cuda.is_available() else "mps"
+
+# name -> (ggml type id, values per block, bytes per block), taken from gguf's own quant sizes so a
+# block layout is never restated here. Q1_0/Q2_0/NVFP4 are missing: the `gguf` release this tests
+# against cannot unpack them, so there is no reference to compare a kernel to.
 QUANT_TYPES = {
-    "Q4_K": (12, 256, 144),
-    "Q5_K": (13, 256, 176),
-    "Q6_K": (14, 256, 210),
-    "Q8_0": (8, 32, 34),
+    name: (int(t), *gguf.GGML_QUANT_SIZES[t])
+    for name, t in (
+        (n, getattr(gguf.GGMLQuantizationType, n, None))
+        for n in (
+            "Q4_0", "Q4_1", "Q5_0", "Q5_1", "Q8_0",
+            "Q2_K", "Q3_K", "Q4_K", "Q5_K", "Q6_K",
+            "IQ2_XXS", "IQ2_XS", "IQ3_XXS", "IQ1_S", "IQ4_NL",
+            "IQ3_S", "IQ2_S", "IQ4_XS", "IQ1_M", "MXFP4",
+        )
+    )
+    if t is not None
 }
 
+# Which of those this build actually implements a gemv for. Asked rather than assumed: the two
+# backends do not cover the same set, and a type routed into a gemv it has no kernel for faults.
+SUPPORTED = {name: v for name, v in QUANT_TYPES.items() if v[0] in GEMV_TYPES}
 
-def random_blocks(rows: int, cols: int, ggml_name: str, device="cuda"):
+
+def random_blocks(rows: int, cols: int, ggml_name: str, device=DEVICE):
     """Random blocks and the values `gguf` reads out of them."""
     _, block_values, block_bytes = QUANT_TYPES[ggml_name]
     generator = np.random.default_rng(0)
     packed = generator.integers(0, 256, (rows, cols // block_values * block_bytes), dtype=np.uint8)
-    # every fp16 scale sits at an even offset in its block, so clearing bit 6 of each odd byte keeps
-    # every possible fp16 field finite whatever the rest of the pattern is
-    packed[:, 1::2] &= 0xBF
+    if ggml_name == "MXFP4":
+        # The only type whose scale is not an fp16 field: one E8M0 byte per block, worth
+        # 2^(byte-127), so a random byte is worth up to 2^128 and the block overflows to inf in the
+        # reference before a kernel is even involved. Held near 2^0 instead.
+        packed[:, ::block_bytes] = generator.integers(120, 135, packed[:, ::block_bytes].shape)
+    else:
+        # every fp16 scale sits at an even offset in its block, so clearing bit 6 of each odd byte
+        # keeps every possible fp16 field finite whatever the rest of the pattern is
+        packed[:, 1::2] &= 0xBF
 
     quant_type = getattr(gguf.GGMLQuantizationType, ggml_name)
     reference = gguf.quants.dequantize(packed.reshape(-1), quant_type).reshape(rows, cols)
+    assert np.isfinite(reference).all(), f"the {ggml_name} reference is not finite"
     return torch.from_numpy(packed).to(device), torch.from_numpy(reference).to(device)
 
 
 @pytest.mark.kernels_ci
-@pytest.mark.parametrize("ggml_name", QUANT_TYPES)
+@pytest.mark.parametrize("ggml_name", SUPPORTED)
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
 def test_dequantize_matches_reference(ggml_name, dtype):
-    ggml_type = QUANT_TYPES[ggml_name][0]
+    ggml_type = SUPPORTED[ggml_name][0]
     rows, cols = 64, 512
     blocks, reference = random_blocks(rows, cols, ggml_name)
 
     out = dequantize(blocks, ggml_type, rows, cols, dtype)
 
     assert out.shape == (rows, cols) and out.dtype == dtype
-    # the kernel writes `dtype` directly, so the tolerance is that dtype's own resolution
-    torch.testing.assert_close(out.float(), reference, rtol=torch.finfo(dtype).eps * 4, atol=0)
+    # The kernel writes `dtype` directly, so the tolerance is that dtype's own resolution -- but
+    # measured against the scale of the values, not each element. Every value in a block shares one
+    # scale, so a dequantizer's error is a fraction of that scale; asking a near-zero element to
+    # match to its own magnitude asks for an exactness even a different summation order breaks.
+    # Most types do come back bit-exact; Q4_0/Q4_1/Q4_K land within an ulp or so of f32.
+    scale = reference.abs().max()
+    torch.testing.assert_close(
+        out.float(), reference, rtol=0, atol=torch.finfo(dtype).eps * 4 * scale
+    )
 
 
 @pytest.mark.kernels_ci
-@pytest.mark.parametrize("ggml_name", QUANT_TYPES)
+@pytest.mark.parametrize("ggml_name", SUPPORTED)
 @pytest.mark.parametrize("n_rows", [1, MAX_GEMV_ROWS])
 def test_mul_mat_vec_matches_matmul(ggml_name, n_rows):
-    ggml_type = QUANT_TYPES[ggml_name][0]
+    ggml_type = SUPPORTED[ggml_name][0]
     out_features, in_features = 128, 512
     blocks, reference = random_blocks(out_features, in_features, ggml_name)
-    x = torch.randn(n_rows, in_features, dtype=torch.bfloat16, device="cuda")
+    x = torch.randn(n_rows, in_features, dtype=torch.bfloat16, device=DEVICE)
 
     out = mul_mat_vec(blocks, x, ggml_type, out_features)
 
@@ -73,10 +102,10 @@ def test_mul_mat_vec_matches_matmul(ggml_name, n_rows):
 @pytest.mark.kernels_ci
 def test_gemv_is_compileable():
     """A graph break here would cost more than the kernel saves, so the fake has to be right."""
-    ggml_type = QUANT_TYPES["Q4_K"][0]
+    ggml_type = SUPPORTED["Q4_K"][0]
     out_features, in_features = 128, 512
     blocks, reference = random_blocks(out_features, in_features, "Q4_K")
-    x = torch.randn(1, in_features, dtype=torch.bfloat16, device="cuda")
+    x = torch.randn(1, in_features, dtype=torch.bfloat16, device=DEVICE)
 
     compiled = torch.compile(
         lambda t: mul_mat_vec(blocks, t, ggml_type, out_features), fullgraph=True

tests/test_vendor_drift.py

@@ -0,0 +1,111 @@
+"""The Metal type table against the upstream code it was transcribed from.
+
+`gguf_metal/ggml_dispatch.mm` re-implements, for the two ops here, what ggml's Metal backend does in
+`ggml_metal_library_get_pipeline_mul_mv` and `ggml_metal_op_mul_mat`: which types have a gemv, how
+much threadgroup memory each wants, and which of the two grids it is dispatched on. Those files are
+vendored, so `vendor.py --rev` replaces them, and a pin bump can change any of it.
+
+Most of that drift is caught by the compiler -- the nsg/nr0 values are used through upstream's own
+`N_SG_*`/`N_R0_*` macros, so a rename or removal fails the build. What is not caught is a changed
+`smem` expression or a changed grid rule: both still compile and both silently return wrong numbers.
+This parses upstream and fails loudly instead.
+
+No device and no built extension needed, so it runs anywhere.
+"""
+
+import re
+from pathlib import Path
+
+import pytest
+
+
+ROOT = Path(__file__).resolve().parent.parent
+DEVICE_CPP = ROOT / "vendor/src/ggml-metal/ggml-metal-device.cpp"
+OPS_CPP = ROOT / "vendor/src/ggml-metal/ggml-metal-ops.cpp"
+DISPATCH_MM = ROOT / "gguf_metal/ggml_dispatch.mm"
+
+pytestmark = pytest.mark.skipif(not DEVICE_CPP.exists(), reason="vendor/ is not checked out")
+
+
+def upstream_mul_mv_table():
+    """{type name: smem expression} from ggml_metal_library_get_pipeline_mul_mv's switch."""
+    body = DEVICE_CPP.read_text()
+    start = body.index("ggml_metal_library_get_pipeline_mul_mv(")
+    end = body.index("ggml_metal_library_get_pipeline_mul_mm_id_map0", start)
+    switch = body[start:end]
+
+    table = {}
+    for case in re.finditer(
+        r"case GGML_TYPE_(\w+):\s*\{(.*?)\}\s*break;", switch, re.S
+    ):
+        name, arm = case.group(1), case.group(2)
+        smem = re.search(r"smem\s*=\s*([^;]+);", arm)
+        table[name] = " ".join(smem.group(1).split()) if smem else "0"
+    return table
+
+
+def ours():
+    """{type name: smem expression} from our table, keyed the way upstream names its types."""
+    body = DISPATCH_MM.read_text()
+    start = body.index("const std::unordered_map<int, TypeInfo> &type_table()")
+    end = body.index("const TypeInfo *lookup(", start)
+
+    table = {}
+    for entry in re.finditer(
+        r"\{GGML_(\w+),\s*\{\"[\w]+\",\s*\d+,\s*\d+,\s*N_SG_\w+,\s*N_R0_\w+,\s*(true|false),\s*([^}]+)\}\}",
+        body[start:end],
+    ):
+        name, reduce_flag, smem = entry.group(1), entry.group(2), entry.group(3)
+        table[name] = (" ".join(smem.split()).rstrip(","), reduce_flag == "true")
+    return table
+
+
+def normalize(expr: str) -> str:
+    """`32 * sizeof(float) * N_R0_Q8_0` and `32*sizeof(float)*N_R0_Q8_0` are the same expression."""
+    return expr.replace(" ", "").replace("(float)", "(float)")
+
+
+def test_every_type_we_claim_still_exists_upstream():
+    upstream = upstream_mul_mv_table()
+    missing = sorted(set(ours()) - set(upstream))
+    assert not missing, (
+        f"{missing} are in our table but no longer in upstream's mul_mv switch; the pin moved under "
+        "gguf_metal/ggml_dispatch.mm"
+    )
+
+
+def test_threadgroup_memory_matches_upstream():
+    upstream = upstream_mul_mv_table()
+    mismatched = {
+        name: (smem, upstream[name])
+        for name, (smem, _) in ours().items()
+        if normalize(smem) != normalize(upstream[name])
+    }
+    assert not mismatched, (
+        "upstream changed the threadgroup memory for these types (ours, theirs): "
+        f"{mismatched}. Too little is a fault, too much silently costs occupancy."
+    )
+
+
+def test_grid_rule_still_singles_out_q8_0():
+    """The one place a type's grid divisor is decided, and the one field we cannot get from a macro.
+
+    Upstream dispatches the dense types and q8_0 on `(ne01 + nr0 - 1)/nr0` and everything else on
+    `(ne01 + nr0*nsg - 1)/(nr0*nsg)`. Our `mv_reduce_across_sgs` is exactly that predicate, so if
+    upstream ever moves a quantized type across, this catches it.
+    """
+    body = OPS_CPP.read_text()
+    start = body.index("ggml_metal_library_get_pipeline_mul_mv(lib, op)")
+    branch = body[start : body.index("return 1;", start)]
+
+    condition = re.search(r"if \(op->src\[0\]->type ==(.*?)\) \{", branch, re.S)
+    assert condition, "upstream's mul_mv grid branch is no longer shaped as a type comparison"
+    quantized_on_nr0_grid = {
+        t for t in re.findall(r"GGML_TYPE_(\w+)", condition.group(1))
+    } - {"F32", "F16", "BF16"}
+
+    ours_on_nr0_grid = {name for name, (_, reduce_flag) in ours().items() if reduce_flag}
+    assert ours_on_nr0_grid == quantized_on_nr0_grid, (
+        f"upstream now puts {quantized_on_nr0_grid} on the undivided grid, we assume "
+        f"{ours_on_nr0_grid}; mv_reduce_across_sgs in gguf_metal/ggml_dispatch.mm has to follow"
+    )

torch-ext/gguf_kernels/__init__.py

@@ -20,9 +20,16 @@ __all__ = ["GEMV_TYPES", "MAX_GEMV_ROWS", "dequantize", "mul_mat_vec"]
 # caller with more (prefill) dequantizes and uses an ordinary matmul.
 MAX_GEMV_ROWS = 8
 
-# ggml type ids `mul_mat_vec` implements: Q4_0/Q4_1/Q5_0/Q5_1/Q8_0, the K quants, and the IQ quants.
-# `dequantize` covers more, so check this one before choosing the fused path.
-GEMV_TYPES = frozenset({2, 3, 6, 7, 8, 10, 11, 12, 13, 14, 16, 17, 18, 19, 20, 21, 22, 23, 29, 39, 40, 41, 42})
+# ggml type ids `mul_mat_vec` implements, as this build reports them -- it is backend-specific, and
+# a type routed into a gemv that has no kernel for it is a fault, not a fallback. CUDA covers
+# Q4_0/Q4_1/Q5_0/Q5_1/Q8_0, the K quants, the IQ quants, MXFP4/NVFP4 and Q1_0/Q2_0; Metal is the
+# same minus NVFP4/Q1_0/Q2_0, which its metallib has no kernels for. `dequantize` covers more, so
+# check this before choosing the fused path.
+try:
+    GEMV_TYPES = frozenset(ops.gemv_types())
+except AttributeError:
+    # A build published before `gemv_types` existed; those are CUDA-only, so its list is theirs.
+    GEMV_TYPES = frozenset({2, 3, 6, 7, 8, 10, 11, 12, 13, 14, 16, 17, 18, 19, 20, 21, 22, 23, 29, 39, 40, 41, 42})
 
 
 def dequantize(

torch-ext/torch_binding.cpp

@@ -6,6 +6,10 @@
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("dequantize(Tensor blocks, int ggml_type, int rows, int cols, ScalarType dtype) -> Tensor");
   ops.def("mul_mat_vec(Tensor blocks, Tensor x, int ggml_type, int out_features) -> Tensor");
+  // Takes no tensor, so it has no device to dispatch on and is registered as a catch-all. Each
+  // backend's shared object is its own library namespace, so there is one implementation per build.
+  ops.def("gemv_types() -> int[]");
+  ops.impl("gemv_types", &gemv_types);
 
   // The schema is the same for every backend; only the implementation differs.
 #if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)

torch-ext/torch_binding.h

@@ -2,6 +2,14 @@
 
 #include <torch/torch.h>
 
+#include <vector>
+
+// The ggml type ids this build's `mul_mat_vec` implements. Backend-specific: the CUDA port reaches
+// everything upstream's `mul_mat_vec_q_switch_type` covers, while Metal's list is the type table in
+// gguf_metal/ggml_dispatch.mm. A caller must ask rather than assume, or it will route a type into a
+// gemv that has no kernel for it.
+std::vector<int64_t> gemv_types();
+
 // The two entry points every backend implements. Both take a GGUF weight exactly as it is stored
 // in the file — `(rows, bytes_per_row)` uint8 blocks — so nothing has to be unpacked to use them.