pytorch / pytorch/pytorch

CPU torch.topk slower than expected for contiguous small k over a large dimension

Open
#197,196 1 comment 1 reaction 1 assignee Claimed by @frost-intel View on GitHub
bot-triaged module: cpu module: performance module: sorting and selection triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### Problem

On CPU, `torch.topk` over a contiguous `[batch, vocab]` tensor with small `k` costs more than the shape suggests it should. This is the filter step in top-k sampling, so on CPU inference it runs once per generated token.

Measured single-threaded on a from-source build of `cd96aa6ae45` (CPU-only, `USE_CUDA=0`, gcc 15.2, Python 3.14, Ubuntu on WSL2, AMD Ryzen 9 7945HX), median over 5 rounds x 300 iterations:

| shape | k | median |
| --- | --- | --- |
| (1, 50257) | 50 | 175 us |
| (32, 50257) | 50 | 4.40 ms |

Profiling a DistilGPT2 CPU generation loop attributes the cost to `topk_impl_loop` in `aten/src/ATen/native/TopKImpl.h`.

### Versions

`2.15.0a0+gitcd96aa6`, built from source, CPU-only.

### Note

I have a patch that cuts the (1, 50257) case roughly in half and leaves every other shape on the existing code path. I'd rather confirm this is worth addressing before sending it.

This was found by [CodeWeave](https://github.com/HighTech-Innovators/CodeWeave), an automated pipeline I run that profiles CPU inference workloads for energy-reduction opportunities and flags hot operators for review; the originating run is at [HighTech-Innovators/CodeWeave-PyTorch](https://github.com/HighTech-Innovators/CodeWeave-PyTorch). I reproduced the profile and re-ran the measurements above myself before filing.

*Disclosure: prepared with AI assistance (Claude Code). The measurements were run on my machine and reviewed by me.*

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.