topk: add top-k kernel for MoE routers on Metal

.github/kernel-maintainers.json

@@ -51,6 +51,7 @@
   "sgl-flash-attn3": [],
   "sonic-moe": ["U08Q42K8MGW"],
   "tinygrad-rms": [],
+  "topk": ["U0586KAH2E9"],
   "trimul-gpumode": [],
   "triton-kernels": [],
   "vllm-flash-attn3": [],

scripts/check_kernel_freshness.py

@@ -19,6 +19,7 @@
     "einops": "https://github.com/arogozhnikov/einops",
     "finegrained-fp8": "",
     "flash-attn-ops": "https://github.com/Dao-AILab/flash-attention",
+    "topk": "",
     "flash-attn2": "https://github.com/Dao-AILab/flash-attention",
     "flash-attn3": "https://github.com/Dao-AILab/flash-attention",
     "flash-attn4": "https://github.com/Dao-AILab/flash-attention",

topk/CARD.md

@@ -0,0 +1,55 @@
+---
+library_name: kernels
+{% if license %}license: {{ license }}
+{% endif %}---
+
+This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
+
+## How to use
+{% if functions %}
+
+```python
+# make sure `kernels` is installed: `pip install -U kernels`
+from kernels import get_kernel
+
+kernel_module = get_kernel("{{ repo_id }}", version={{ version }})
+{{ functions[0] }} = kernel_module.{{ functions[0] }}
+
+{{ functions[0] }}(...)
+```
+{% else %}
+
+Usage example not available.
+{% endif %}
+
+## Available functions
+{% if functions %}
+{% for func in functions %}
+- `{{ func }}`
+{% endfor %}
+{% else %}
+
+Function list not available.
+{% endif %}
+{% if layers %}
+
+## Available layers
+{% for layer in layers %}
+- `{{ layer }}`
+{% endfor %}
+{% endif %}
+
+## Benchmarks
+{% if has_benchmark %}
+
+Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`.
+{% else %}
+
+No benchmark available yet.
+{% endif %}
+{% if upstream %}
+
+## Source code
+
+Source code of this kernel originally comes from {{ upstream }} and it was repurposed for compatibility with `kernels`.
+{% endif %}

topk/README.md

@@ -0,0 +1,33 @@
+---
+library_name: kernels
+license: mit
+tags:
+  - kernel
+---
+
+## topk
+
+Top-k over a small row, for a MoE router. One threadgroup per row, one reduction pass per output:
+`k*n` comparisons, but `k` and `n` are small and the launch dominates either way. Optionally softmaxes
+the `k` it selected, which is what a router wants and saves a second dispatch.
+
+Not a port. ggml has no top-k of its own — `GGML_OP_TOP_K` dispatches `kernel_argsort_f32_i32_desc`,
+a full bitonic sort — and torch's MPS `topk` is a full sort too. Selecting the largest 8 of 256 logits
+does not need the row ordered: 26 us here against 71 us for either sort, once per layer per token.
+
+`indices` comes back as int32, which is what an expert-routed matmul wants, so routing them onward
+costs no cast.
+
+## Usage
+
+```python
+import torch
+from kernels import get_kernel
+
+topk = get_kernel("kernels-community/topk", version=1)
+
+logits = torch.randn(1, 256, device="mps")          # one row of router logits
+
+values, indices = topk.top_k(logits, 8)             # (1, 8) f32, (1, 8) int32
+weights, experts = topk.top_k(logits, 8, True)      # values softmaxed over the selected 8
+```

topk/build.toml

@@ -0,0 +1,35 @@
+# One backend per `[kernel.*]` section, all implementing the entry points declared in
+# torch-ext/torch_binding.h.
+#
+# The backend lives in `topk_metal/` while the package is just `topk`: a CUDA port would be a
+# `topk_cuda/` section beside it, with no change to the bindings' schema or the Python API.
+#
+# Unlike the `ggml-*` packages there is no `vendor/`: this kernel is not a port of anything upstream.
+# ggml has no top-k of its own -- `GGML_OP_TOP_K` dispatches `kernel_argsort_f32_i32_desc`, a full
+# bitonic sort -- so there was nothing to vendor and this is written against `metal_stdlib` alone.
+
+[general]
+name = "topk"
+version = 1
+license = "MIT"
+backends = ["metal"]
+
+[general.hub]
+repo-id = "kernels-community/topk"
+
+[torch]
+src = [
+    "torch-ext/torch_binding.cpp",
+    "torch-ext/torch_binding.h",
+]
+
+[kernel.topk_metal]
+backend = "metal"
+depends = ["torch"]
+include = ["topk_metal", "torch-ext"]
+src = [
+    "topk_metal/topk_metal.cpp",
+    "topk_metal/dispatch.mm",
+    "topk_metal/common.h",
+    "topk_metal/top_k.metal",
+]

topk/flake.nix

@@ -0,0 +1,12 @@
+{
+  description = "Fused sequence-mixing kernels from ggml: the gated delta rule, as one kernel";
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernels";
+  };
+  outputs =
+    { self, kernel-builder }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

topk/tests/test_top_k.py

@@ -0,0 +1,57 @@
+"""`top_k` against `torch.topk`, which is what it replaces.
+
+The kernel exists to cut a sort down to a selection, so the test that matters is that it picks the
+same elements, in the same order, as the sort it stands in for.
+"""
+
+import os
+
+import pytest
+import torch
+
+
+DEV = "mps"
+LIB = os.environ.get("TOPK_LOCAL_LIB")
+
+if LIB:
+    torch.ops.load_library(LIB)
+    ops = getattr(torch.ops, os.path.basename(LIB).removesuffix(".so"))
+else:
+    from pathlib import Path
+
+    try:
+        from kernels import get_local_kernel
+
+        ops = get_local_kernel(Path(__file__).resolve().parent.parent, "metal")
+    except Exception as error:  # pragma: no cover
+        pytest.skip(f"no kernel to test ({error})", allow_module_level=True)
+
+pytestmark = pytest.mark.skipif(not torch.backends.mps.is_available(), reason="needs mps")
+
+
+@pytest.mark.parametrize(("rows", "n", "k"), [(1, 256, 8), (1, 128, 4), (7, 512, 8), (1, 64, 1), (3, 320, 16)])
+def test_matches_torch_topk(rows, n, k):
+    torch.manual_seed(0)
+    logits = torch.randn(rows, n, device=DEV)
+    values, indices = ops.top_k(logits, k, False)
+    ref_values, ref_indices = torch.topk(logits, k, dim=-1)
+    assert torch.equal(indices.long(), ref_indices)
+    assert torch.equal(values, ref_values)
+    assert indices.dtype == torch.int32
+
+
+def test_softmax_is_over_the_selected_k():
+    torch.manual_seed(0)
+    logits = torch.randn(4, 256, device=DEV)
+    values, _ = ops.top_k(logits, 8, True)
+    reference = torch.softmax(torch.topk(logits, 8, dim=-1).values.float(), dim=-1)
+    assert torch.allclose(values, reference, atol=1e-6)
+    assert torch.allclose(values.sum(-1), torch.ones(4, device=DEV), atol=1e-6)
+
+
+def test_k_equal_to_the_row():
+    """Degenerate but legal: selecting everything is a sort, and must still agree."""
+    torch.manual_seed(0)
+    logits = torch.randn(2, 64, device=DEV)
+    _, indices = ops.top_k(logits, 64, False)
+    assert torch.equal(indices.long(), torch.topk(logits, 64, dim=-1).indices)

topk/topk_metal/common.h

@@ -0,0 +1,24 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+/* The Metal-facing boundary.
+ *
+ * Everything that touches Metal lives behind this function, so torch headers and Metal headers never
+ * meet in one translation unit. Buffers arrive as (MTLBuffer, byte offset) pairs because a torch
+ * tensor's storage is a whole MTLBuffer that the tensor may only be a view into.
+ *
+ * Returns 0 on success, or non-zero when the metallib has no such kernel, which the caller reports
+ * rather than faulting.
+ */
+
+extern "C" {
+
+// Top-k over each row of `logits` (`rows x n` f32) -> `indices` (i32) and `values` (f32), both
+// `rows x k`. With `softmax` set, `values` is softmaxed over the k that were selected, which is what
+// a router wants and saves the caller a second dispatch.
+int topk_metal_top_k(void *logits, size_t logits_off, void *indices, size_t indices_off,
+                     void *values, size_t values_off, int64_t rows, int64_t n, int64_t k,
+                     int softmax);
+}