Pull requests / #743

#743 qsa_select: coalesced streaming top-k for long Turing prompts (+23% 248K prefill)

closed · @konijiwa110 · 0 comentarios · En GitHub

BenchmarksNVIDIA / CUDAModels & quants

Descripción

## What

The register top-k holds at most `TK_T * TK_PER` blocks (~135K cells). Past that the prompt path falls back to
`block_topk_kernel`, whose per-thread runs read the scores uncoalesced: 12.8 ms per 256 queries at 250K on an RTX
2080 Ti, the largest prefill phase there at that length.

`block_topk_stream_kernel` reads the keys from memory on every radix pass with block `t + i * TK_T` (coalesced), so
it has no register limit, then emits the cells in ascending chunks of `TK_T` blocks with two block scans (ties to the
lowest index), giving `block_topk_kernel`'s ids.

It is taken only where the active bound is already used (Turing prefill, or `STRATA_TOPK_ACTIVE_ANY=1` in tests) above
22,528 blocks (~90K cells). Decode and other cards are unchanged. `STRATA_TOPK_STREAM=0` keeps the previous dispatch.

`qsa_topk_active_parity` gains cases either side of the 22,528-block bound, heavily tied scores at 200K and all-zero
scores at 262K (20 cases), comparing against the capacity-dispatch reference.

## Measured

RTX 2080 Ti, 256 queries, capacity 262144, top-k with the active bound (`STRATA_TOPK_STREAM=0` / unset):

| context | before | stream |
|---|---|---|
| 90K | 0.49 ms | 0.48 ms |
| 131K | 1.03 ms | 0.75 ms |
| 250K | 12.8 ms | 1.13 ms |

Selections identical to the reference kernel 256/256 at each length. `qsa_topk_active_parity`: PASS (20 cases, ids
bitwise identical), with the stream on and off.

End to end, Swift 1.5 IQ3_XXS, same build A/B, prefill tok/s: 248K 951 -> 1167 (+23%), 155K 1166 -> 1210 (+4%),
128K and below unchanged. `qsa select` at 248K: 70.9 s -> 22.2 s. (These runs had #742's scorer on as well.)

Only tested on sm_75. Independent of #741 and #742 (merges cleanly with both).

Developed with an AI coding assistant; all numbers measured on an RTX 2080 Ti 22 GB / i5-12490F / 64 GB.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

En el sitio

Enlaces a install, modelos, releases.