add the README, which build-and-upload does not copy
Browse files
README.md
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
library_name: kernels
|
| 3 |
+
license: mit
|
| 4 |
+
tags:
|
| 5 |
+
- kernel
|
| 6 |
+
---
|
| 7 |
+
|
| 8 |
+
## topk
|
| 9 |
+
|
| 10 |
+
Top-k over a small row, for a MoE router. One threadgroup per row, one reduction pass per output:
|
| 11 |
+
`k*n` comparisons, but `k` and `n` are small and the launch dominates either way. Optionally softmaxes
|
| 12 |
+
the `k` it selected, which is what a router wants and saves a second dispatch.
|
| 13 |
+
|
| 14 |
+
Not a port. ggml has no top-k of its own — `GGML_OP_TOP_K` dispatches `kernel_argsort_f32_i32_desc`,
|
| 15 |
+
a full bitonic sort — and torch's MPS `topk` is a full sort too. Selecting the largest 8 of 256 logits
|
| 16 |
+
does not need the row ordered: 26 us here against 71 us for either sort, once per layer per token.
|
| 17 |
+
|
| 18 |
+
`indices` comes back as int32, which is what an expert-routed matmul wants, so routing them onward
|
| 19 |
+
costs no cast.
|
| 20 |
+
|
| 21 |
+
## Usage
|
| 22 |
+
|
| 23 |
+
```python
|
| 24 |
+
import torch
|
| 25 |
+
from kernels import get_kernel
|
| 26 |
+
|
| 27 |
+
topk = get_kernel("marcsun13/topk", version=1)
|
| 28 |
+
|
| 29 |
+
logits = torch.randn(1, 256, device="mps") # one row of router logits
|
| 30 |
+
|
| 31 |
+
values, indices = topk.top_k(logits, 8) # (1, 8) f32, (1, 8) int32
|
| 32 |
+
weights, experts = topk.top_k(logits, 8, True) # values softmaxed over the selected 8
|
| 33 |
+
```
|