topk-metal: top-k over a small row for MoE routers

.gitattributes

@@ -0,0 +1,39 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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
+# Built extensions are binaries, tracked by pattern rather than by name so a rebuild cannot
+# introduce an untracked one. See the LFS-pointer trap in gguf-kernels/SKILL.md before committing:
+# every shipped .so must be a real library, not a 132-byte pointer.
+build/**/*.so filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,8 @@
+localbuild/out/
+build/
+result
+result-*
+__pycache__/
+.pytest_cache/
+.DS_Store
+localbuild/work/

build.toml

@@ -0,0 +1,32 @@
+# One backend per `[kernel.*]` section, all implementing the entry points declared in
+# torch-ext/torch_binding.h.
+#
+# 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-metal"
+version = 1
+license = "MIT"
+backends = ["metal"]
+
+[general.hub]
+repo-id = "marcsun13/topk-metal"
+
+[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",
+]

flake.lock

@@ -0,0 +1,117 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1767039857,
+        "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs",
+        "rust-overlay": "rust-overlay"
+      },
+      "locked": {
+        "lastModified": 1785942216,
+        "narHash": "sha256-OlefmZzVX8KRYFwqo5Mx2MuG9LHjhwkehaN6vXu2c1c=",
+        "owner": "huggingface",
+        "repo": "kernels",
+        "rev": "81f55ea30fd8f819dcf93a3c934dd584c895bd2f",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "kernels",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1783284758,
+        "narHash": "sha256-tiQ8/qi8I45OOaBBYlVbXoAVkeQzvvTQOv5I45rMw5o=",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d",
+        "type": "github"
+      },
+      "original": {
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "rust-overlay": {
+      "inputs": {
+        "nixpkgs": [
+          "kernel-builder",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1783320166,
+        "narHash": "sha256-l7C/OsjcnWDOk2K3ssj+SBduwL67LashjBqis9+t468=",
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "rev": "20ee15370c9256669d66968b89ee20a4b0a4e673",
+        "type": "github"
+      },
+      "original": {
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "type": "github"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

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_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);
+}

topk_metal/dispatch.mm

@@ -0,0 +1,116 @@
+/* Metal side: load the metallib and encode the kernel on torch's own stream.
+ *
+ * What this has that a `torch.mps.compile_shader` caller does not is encoding into torch's *current*
+ * command buffer -- so the dispatch joins the work already queued rather than becoming its own
+ * submission -- and doing that on the stream's own serial queue, which is what keeps it legal.
+ */
+
+#import <Metal/Metal.h>
+
+#include <ATen/mps/MPSDevice.h>
+#include <ATen/mps/MPSStream.h>
+
+#include <string>
+#include <unordered_map>
+
+#include "common.h"
+
+#ifdef EMBEDDED_METALLIB_HEADER
+#include EMBEDDED_METALLIB_HEADER
+#endif
+
+namespace {
+
+id<MTLDevice> device() { return at::mps::MPSDevice::getInstance()->device(); }
+
+id<MTLLibrary> library() {
+  static id<MTLLibrary> lib = nil;
+  if (lib != nil) {
+    return lib;
+  }
+  NSError *error = nil;
+#ifdef EMBEDDED_METALLIB_HEADER
+  lib = EMBEDDED_METALLIB_NAMESPACE::createLibrary(device(), &error);
+#else
+  // Local builds point at a metallib on disk; the packaged build embeds it instead.
+  const char *path = getenv("TOPK_METALLIB");
+  if (path == nullptr) {
+    NSLog(@"topk-metal: TOPK_METALLIB is unset and no metallib is embedded");
+    return nil;
+  }
+  NSURL *url = [NSURL fileURLWithPath:[NSString stringWithUTF8String:path]];
+  lib = [device() newLibraryWithURL:url error:&error];
+#endif
+  if (lib == nil) {
+    NSLog(@"topk-metal: could not load the metallib: %@", error);
+  } else {
+    [lib retain];
+  }
+  return lib;
+}
+
+// Pipelines are cached under the same key ggml uses: the kernel name plus the constants it was
+// specialised with, since a different specialisation is a different pipeline.
+id<MTLComputePipelineState> pipeline(const std::string &key, const char *fn_name,
+                                     MTLFunctionConstantValues *constants) {
+  static std::unordered_map<std::string, id<MTLComputePipelineState>> cache;
+  auto it = cache.find(key);
+  if (it != cache.end()) {
+    return it->second;
+  }
+  id<MTLLibrary> lib = library();
+  if (lib == nil) {
+    return nil;
+  }
+  NSError *error = nil;
+  NSString *name = [NSString stringWithUTF8String:fn_name];
+  id<MTLFunction> fn = constants == nil
+                           ? [lib newFunctionWithName:name]
+                           : [lib newFunctionWithName:name constantValues:constants error:&error];
+  if (fn == nil) {
+    NSLog(@"topk-metal: no function %s in the metallib: %@", fn_name, error);
+    return nil;
+  }
+  id<MTLComputePipelineState> pso = [device() newComputePipelineStateWithFunction:fn error:&error];
+  [fn release];
+  if (pso == nil) {
+    NSLog(@"topk-metal: could not build a pipeline for %s: %@", fn_name, error);
+    return nil;
+  }
+  cache[key] = pso;
+  return pso;
+}
+
+}  // namespace
+
+extern "C" 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) {
+  at::mps::MPSStream *stream = at::mps::getCurrentMPSStream();
+  __block int rc = 0;
+  dispatch_sync(stream->queue(), ^{
+    id<MTLComputeCommandEncoder> enc = stream->commandEncoder();
+    id<MTLComputePipelineState> pso = pipeline("kernel_top_k_f32", "kernel_top_k_f32", nil);
+    if (pso == nil) {
+      rc = 2;
+      return;
+    }
+    const int32_t n32 = (int32_t)n, k32 = (int32_t)k, sm32 = softmax;
+    // One threadgroup per row; a power of two so the reduction below can halve cleanly.
+    NSUInteger nth = 32;
+    while (nth * 2 <= (NSUInteger)n && nth < 256) {
+      nth *= 2;
+    }
+    [enc setComputePipelineState:pso];
+    [enc setBuffer:(__bridge id<MTLBuffer>)logits offset:logits_off atIndex:0];
+    [enc setBuffer:(__bridge id<MTLBuffer>)indices offset:indices_off atIndex:1];
+    [enc setBuffer:(__bridge id<MTLBuffer>)values offset:values_off atIndex:2];
+    [enc setBytes:&n32 length:sizeof(n32) atIndex:3];
+    [enc setBytes:&k32 length:sizeof(k32) atIndex:4];
+    [enc setBytes:&sm32 length:sizeof(sm32) atIndex:5];
+    [enc setThreadgroupMemoryLength:nth * (sizeof(float) + sizeof(int)) atIndex:0];
+    [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)rows, 1, 1)
+        threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
+  });
+  return rc;
+}

topk_metal/top_k.metal

@@ -0,0 +1,84 @@
+#include <metal_stdlib>
+
+using namespace metal;
+
+/* Top-k over a small row, for a MoE router.
+ *
+ * Not ggml's: ggml sorts (a bitonic pass plus a merge) because its `argsort` is general. A router
+ * wants the largest `k` of a few hundred logits, on one row per token, and torch's MPS `topk` is a
+ * full sort too -- 71us for k=8 of 256, against 21us for a plain argmax over the same data. Over 40
+ * layers that is ~2.8ms per token spent selecting 8 numbers.
+ *
+ * One threadgroup per row, one pass per output: each pass reduces the row to its maximum, records
+ * it, and masks it out. k*n comparisons, but k and n are small and the launch dominates either way.
+ */
+kernel void kernel_top_k_f32(
+        device const float * logits    [[buffer(0)]],
+        device       int   * indices   [[buffer(1)]],
+        device       float * values    [[buffer(2)]],
+        constant     int   & n         [[buffer(3)]],
+        constant     int   & k         [[buffer(4)]],
+        constant     int   & softmax   [[buffer(5)]],
+        uint  row  [[threadgroup_position_in_grid]],
+        uint  tid  [[thread_position_in_threadgroup]],
+        uint  nth  [[threads_per_threadgroup]],
+        threadgroup float * shared_val [[threadgroup(0)]]) {
+    device const float * src = logits + (ulong)row * n;
+    threadgroup int * shared_idx = (threadgroup int *)(shared_val + nth);
+
+    // `taken` is tracked by writing -INFINITY into a per-thread running copy of the best candidate,
+    // so no scratch buffer is needed for the mask: each pass re-scans and skips already-chosen ones.
+    for (int picked = 0; picked < k; ++picked) {
+        float best = -INFINITY;
+        int   best_i = -1;
+        for (int i = tid; i < n; i += nth) {
+            float v = src[i];
+            bool used = false;
+            for (int j = 0; j < picked; ++j) {
+                used = used || (indices[(ulong)row * k + j] == i);
+            }
+            if (!used && v > best) {
+                best = v;
+                best_i = i;
+            }
+        }
+        shared_val[tid] = best;
+        shared_idx[tid] = best_i;
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        for (uint stride = nth / 2; stride > 0; stride >>= 1) {
+            if (tid < stride) {
+                // ties go to the lower index, which is what a stable top-k does
+                bool take = shared_val[tid + stride] > shared_val[tid] ||
+                            (shared_val[tid + stride] == shared_val[tid] &&
+                             shared_idx[tid + stride] < shared_idx[tid]);
+                if (take) {
+                    shared_val[tid] = shared_val[tid + stride];
+                    shared_idx[tid] = shared_idx[tid + stride];
+                }
+            }
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+        }
+
+        if (tid == 0) {
+            indices[(ulong)row * k + picked] = shared_idx[0];
+            values [(ulong)row * k + picked] = shared_val[0];
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    // A router softmaxes exactly the k it selected, and that is another dispatch for k numbers --
+    // as expensive as a softmax over the whole row, since it is launch-bound. Done here instead.
+    if (softmax != 0 && tid == 0) {
+        device float * out = values + (ulong)row * k;
+        float top = out[0];  // passes are in descending order, so the first is the maximum
+        float sum = 0.0f;
+        for (int i = 0; i < k; ++i) {
+            out[i] = exp(out[i] - top);
+            sum += out[i];
+        }
+        for (int i = 0; i < k; ++i) {
+            out[i] /= sum;
+        }
+    }
+}

topk_metal/topk_metal.cpp

@@ -0,0 +1,41 @@
+/* Metal implementation of the entry point.
+ *
+ * No Metal header is included here: all Metal contact sits behind the extern "C" boundary in
+ * dispatch.mm, and every output comes from torch's allocator so the op composes with the rest of a
+ * model's memory.
+ *
+ * A torch MPS tensor's storage is a whole MTLBuffer that the tensor may be a view into, so each
+ * buffer crosses the boundary as a (buffer, byte offset) pair.
+ */
+
+#include <torch/torch.h>
+
+#include <vector>
+
+#include "common.h"
+#include "torch_binding.h"
+
+namespace {
+
+void *mtl_buffer(const at::Tensor &t) { return const_cast<void *>(t.storage().data()); }
+
+size_t byte_offset(const at::Tensor &t) {
+  return static_cast<size_t>(t.storage_offset()) * t.element_size();
+}
+
+}  // namespace
+
+std::vector<at::Tensor> top_k(const at::Tensor &logits, int64_t k, bool softmax) {
+  TORCH_CHECK(logits.is_mps() && logits.dim() == 2, "logits must be a 2D mps tensor");
+  const at::Tensor x =
+      logits.scalar_type() == at::kFloat ? logits.contiguous() : logits.to(at::kFloat).contiguous();
+  const int64_t rows = x.size(0), n = x.size(1);
+  TORCH_CHECK(k > 0 && k <= n, "k must be within the row");
+  at::Tensor indices = at::empty({rows, k}, x.options().dtype(at::kInt));
+  at::Tensor values = at::empty({rows, k}, x.options());
+  const int status = topk_metal_top_k(mtl_buffer(x), byte_offset(x), mtl_buffer(indices),
+                                      byte_offset(indices), mtl_buffer(values), byte_offset(values),
+                                      rows, n, k, softmax ? 1 : 0);
+  TORCH_CHECK(status == 0, "top_k: no kernel for this build (", status, ")");
+  return {values, indices};
+}

torch-ext/topk_metal/__init__.py

@@ -0,0 +1,39 @@
+"""Top-k over a small row, as a torch op.
+
+For a MoE router. torch's MPS `topk` is a full sort, and so is ggml's -- `GGML_OP_TOP_K` dispatches
+`kernel_argsort_f32_i32_desc` -- which is why this is written rather than ported: selecting the
+largest 8 of 256 logits does not need the row ordered. 26us against 71us, once per layer per token.
+"""
+
+import os
+from pathlib import Path
+
+import torch
+
+
+# A local (non-nix) build leaves the metallib on disk beside this module instead of embedding it, and
+# `dispatch.mm` looks it up through this variable when nothing is embedded. Set before the extension
+# is imported, and never over an existing value, so an explicit choice still wins.
+_METALLIB = Path(__file__).parent / "ggml-metal.metallib"
+if _METALLIB.is_file():
+    os.environ.setdefault("TOPK_METALLIB", str(_METALLIB))
+
+from ._ops import add_op_namespace_prefix, ops  # noqa: E402
+
+
+def top_k(logits: torch.Tensor, k: int, softmax: bool = False):
+    """`(values, indices)` of the largest `k` per row of a 2D f32 tensor.
+
+    `indices` is int32, which is what an expert-routed matmul wants, so routing them onward costs no
+    cast. With `softmax`, `values` is softmaxed over the selected k.
+    """
+    return ops.top_k(logits, k, softmax)
+
+
+@torch.library.register_fake(add_op_namespace_prefix("top_k"))
+def _top_k_fake(logits, k, softmax):
+    rows = logits.shape[0]
+    return [logits.new_empty((rows, k)), logits.new_empty((rows, k), dtype=torch.int32)]
+
+
+__all__ = ["top_k"]

torch-ext/torch_binding.cpp

@@ -0,0 +1,14 @@
+#include <torch/library.h>
+
+#include "registration.h"
+#include "torch_binding.h"
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("top_k(Tensor logits, int k, bool softmax) -> Tensor[]");
+
+#if defined(METAL_KERNEL)
+  ops.impl("top_k", torch::kMPS, &top_k);
+#endif
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h

@@ -0,0 +1,12 @@
+#pragma once
+
+#include <torch/torch.h>
+
+#include <vector>
+
+// Top-k over each row: returns `{values, indices}`, both `rows x k`, with `indices` as int32.
+//
+// For a MoE router, where torch's MPS top-k is a full sort -- 71us for k=8 of 256, against 26us
+// here -- and a router runs once per layer per token. With `softmax` set, `values` comes back
+// softmaxed over the k that were selected, which is what a router wants and saves a dispatch.
+std::vector<at::Tensor> top_k(const at::Tensor &logits, int64_t k, bool softmax);