---
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("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
```