Pull requests / #742

#742 qsa_select: FP32 tiled block scores below sm_80 (+22% 128K / +33% 248K prefill on a 2080 Ti)

closed · @konijiwa110 · 0 comments · View on GitHub

BenchmarksNVIDIA / CUDAModels & quants

Description

## What

Below sm_80 there is no TF32 tensor-core block scorer, so the prompt path's QSA selection scores blocks on the warp
kernel (4 products per lane, 5 shuffles): ~1 TFLOPS on an RTX 2080 Ti, 24-30% of a 128K prompt there.

`qsa_block_scores_tc` now runs an FP32 tiled GEMM on those cards (`block_scores_simt_kernel`): rows = query x indexer
head, K = 128 through shared memory, each thread 2 queries x 4 heads x 4 blocks, relu per head summed in registers.
~6.9 TFLOPS. The tail block keeps the tail kernel's arithmetic. Like the tensor-core scorer it is FP32 in another
summation order, so not bitwise equal to the warp kernel.

`STRATA_SELECT_SIMT=0` (or the existing `STRATA_SELECT_OLD=1`) keeps the warp kernel.

`qsa_select_bench` is now also registered as a CUDA test: at 32K, and at 131K with `STRATA_QSA_WARP=1` so the FP32
tiled scorer is exercised on any card. Gate: error vs an FP64 reference no worse than 4x the warp kernel's.

## Measured

RTX 2080 Ti, `qsa_select_bench` (256 queries, capacity 262144):

| context | warp kernel | FP32 tiled |
|---|---|---|
| 32K | 1.953 ms | 0.303 ms (6.5x) |
| 131K | 7.831 ms | 1.117 ms (7.0x) |

Selections identical to the warp kernel's 256/256; FP64 error 5e-5 on a score scale of ~268. 4K-250K: 4.5-6.8x.

End to end, Swift 1.5 IQ3_XXS, same build A/B with `STRATA_SELECT_SIMT=0` / unset, prefill tok/s: 128K 1016 -> 1242
(+22%), 248K 715 -> 952 (+33%), 32K +3%, 4K unchanged. `qsa select` at 128K: 29.9 s -> 6.4 s. Decode unchanged. (These runs had #741's FP16 path on as well.)

ctest on the 2080 Ti: `qsa_topk_active_parity`, `qsa_select_bench`, `qsa_select_bench_fp32_tiled` pass.

Only tested on sm_75. Independent of #741 and #743 (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)

Related on strata.com

Editorial links to help you install, pick models, or read release notes — not part of the upstream thread.