top k

A row of logits is half a megabyte, so the kernel is not bandwidth-bound and the cost is in how many passes you make over it. Why a full sort does far more work than the question asks, the two shapes that fit small and large k, and the four-pass radix select with a reference that matches a sort exactly.

The classification is a small function; the parser around it is where the problem is. Generators rather than lists, aggregation bounded by the number of distinct instructions rather than by the file, a top-k that never sorts the whole thing, and the malformed lines that stop the job at hour three.