[rafatxf]

Builds on #2072 (it sits next to that PR's split top-k in qsa_select.cu); the first commit here is #2072's, the second is this one.

Past ~115K cells the prompt path's top-k is DRAM traffic on Turing: block_topk_wide_kernel reads a query's scores six times (four radix passes, a count, the emit), 400 KB a query at 400K, 256 queries a call.

block_topk_sample_kernel reads them once:

1. a bound LB: the KS-th largest key among every 32nd block (KS = 2 × the selection's blocks / 32 + 16);
2. one read of all blocks, keeping those with key ≥ LB in shared memory;
3. when those hold at least width cells, the width-th largest cell key (block_topk_kernel's threshold) is ≥ LB, so every block at or above the threshold is a candidate: the threshold, above and the tie budget computed over the candidates with block_topk_kernel's rule are the reference's, and the candidates sorted by block (bitonic) emit the reference's cells in its order;
4. otherwise (too few cells above LB, or more than 4,096 candidates) the query is flagged and block_topk_wide_kernel<true> runs it - the all-equal-scores case of the parity test takes that path.

Used for prompt calls above 28,672 active blocks (below, the wide kernel is as fast); STRATA_TOPK_SAMPLE=0 keeps the wide kernel.

Setup

RTX 2080 Ti 22 GB (sm_75), i7-6700, CUDA 13.0, driver 595, SM 1650 MHz, 240 W; v0.1.43 + #2072, -DSTRATA_BUILD_TESTS=ON. Model Q2_0, --max-context 524288, --kv int8 --kv-resident 32768, STRATA_SELECT_SIMT=2 (#2071).

qsa_topk_parity CTX 256 20 524288 1 (a prompt call), dispatched time

context	wide	sample
98K	0.397 ms	0.396 (wide: below the bound)
131K	0.513	0.452
262K	0.952	0.589
400K	1.476	0.706
520K	1.840	0.787

qsa_topk_parity --selftest: PASS (continuous, 64-level and equal scores; identical ids). qsa_topk_active_parity (STRATA_TOPK_ACTIVE_ANY=1): 20 cases, identical.

Engine

First logits after a 196K prompt: bitwise identical to the run without it (STRATA_DUMP_FIRST_LOGITS). The gain sits at depth: with this, #2071 and #2091 and #2092, appending 16K tokens to a 390K conversation (--serve) went 15.9 → 12.9 s here.

🤖 Generated with Claude Code

[Niko1221]

Thanks rafatxf, good work. We ported this with your authorship kept. We have no Turing card, so we could not confirm the speed. What we did check: the default path is byte-identical to before (identity gate on RTX 3060, P100 and A4000), and with the new path forced on, including a build that runs the real sm_75 device code on an sm_86 card, the gate is identical too. Because we cannot measure it on sm_75 it ships as opt-in (STRATA_SM75_PERF=1, plus your own variable STRATA_TOPK_SAMPLE), documented in docs/DETAILS.md. If a Turing owner confirms the gain we will consider turning it on by default.

[rafatxf]

Ported in 0.1.45, thanks! Turing numbers for STRATA_SM75_PERF (2x RTX 2080 Ti, Q2_0 and IQ3_S) are in #2069. Closing.