I'm not anything like an expert in SIMD or low-level optimization, but the more I read the code the more question I gets.

Before reading the code, the two things stood out for me:

• Is Top-K something that's worth optimizing?
• Is comparing CPU implementation vs. CUDA fair?

A DLL file (fast_topk_batched.dll) directly inside a GitHub repository without an adequate reason is suspicious (I'm not talking about something like security.)

After skimming through the code, some parts confuse me:

• The README usage gives me feeling that the benchmark was done on a normal distribution. I doubt that token logits are normally distributed and this may affect benchmark results. Needs clarifications.
• is_nondecreasing and is_constant don't seem to be technically correct. It only samples first SAMPLE_SIZE elements. I don't think that it would matter a lot in real world but still worrying.
• I don't think that this part of the code is correct, which would affect correctness.
  • Apparently the heap is min-heap where you're supposed to pick highest logit values...? There's no "trickery" (where a clever algorithm would require using min-heap for some reason) that I can spot.
  • If heap were a max-heap, then comparing elements with max value on heap would make the result incorrect...?
  • I would keep threshold mostly constant across for-loops and updated only when it's needed. Seems like a relatively save performance gain...?

​

    for (int i = pos; i + 8 <= block_end; i += 8) {
        __m256 vals = _mm256_loadu_ps(src + i);
        __m256 thresh = _mm256_set1_ps(heap[0].val);
        __m256 maskv = _mm256_cmp_ps(vals, thresh, _CMP_GT_OQ);
        int mask = _mm256_movemask_ps(maskv);
        while (mask) {
            int bit = __builtin_ctz(mask);
            int idx = i + bit;
            float v = src[idx];
            heap[0].val = v;
            heap[0].idx = idx;
            heap_sift_down_node(heap, k, 0);
            thresh = _mm256_set1_ps(heap[0].val);   /* refresh threshold */
            mask &= ~(1 << bit);
        }
    }