CPU torch.topk slower than expected for contiguous small k over a large dimension
- 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
Assessment
This issue has not been assessed yet.