metadata
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
import torch
from kernels import get_kernel
topk = get_kernel("marcsun13/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