add the README and tests

README.md

@@ -0,0 +1,33 @@
+---
+library_name: kernels
+license: mit
+tags:
+  - kernel
+---
+
+## topk-metal
+
+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("marcsun13/topk-metal", 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
+```

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)