Repository navigation
[CUDA][Performance] Topk is slow #3064
Description
MLX topk on CUDA is pretty slow in some cases (especially compared to PyTorch).
Here is a benchmark:
import time
import mlx.core as mx
b = 2048
v = 8192
k = 32
q = mx.random.normal(shape=(b, v)).astype(mx.bfloat16)
def fun(q):
for _ in range(50):
idx = mx.argpartition(-q, kth=k-1, axis=-1)[:, :k]
values = mx.take_along_axis(q, idx, axis=-1)
q = mx.put_along_axis(q, idx, values, axis=-1)
mx.eval(q)
for _ in range(20):
fun(q)
tic = time.time()
for _ in range(20):
fun(q)
toc = time.time()
ms = 1e3 * (toc - tic)
print(f"MLX {ms=:.3f}")
import torch
q = torch.randn(size=(b, v)).to("cuda").to(torch.bfloat16)
def topk_old(q):
return idx, values
def fun(q):
for _ in range(50):
values, idx = torch.topk(q, k=k, axis=-1)
q = torch.scatter(q, -1, idx, values)
torch.cuda.synchronize()
for _ in range(20):
fun(q)
tic = time.time()
for _ in range(20):
fun(q)
toc = time.time()
ms = 1e3 * (toc - tic)
print(f"PyTorch {ms=:.3f}")On a spark:
MLX ms=2975.014
PyTorch ms=919.764
Activity
NripeshN commentedon Jan 26, 2026
At sort.cu, we use gpu_sort() which performs a radix sort which could be inefficient(O(n log n)) compared to even our CPU backend where we use std::nth_element which is O(n).
We could use NVIDIA's RAFT library which already has an efficient top-k selection
awni commentedon Jan 26, 2026
The sort in
sort.cuis a merge sort not a radix sort. Also a radix sort is notO(n log n))butO(n)with a constant based on the bit width. For an efficient top-k (esp for half precision) a radix-based selection probably makes sense. Though we will have to play around with that.NripeshN commentedon Jan 27, 2026
Ah you are right, I was on an earlier commit where we were using CUB's radix sort. The problem then was we were doing full radix sort rather than Radix select. Also if I am not wrong right now we are at
O(n log n))😅 . You are right radix sort isO(w.n).Can I take on this issue?👉🏽👈🏽
awni commentedon Jan 27, 2026
Feel free to work on it!
awni commentedon Jan 27, 2026
A standalone MLX benchmark for various shape and type combinations: