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 commentaires · Sur 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)
Sur le site
Liens install, modèles, releases.