Pull requests / #1825
#1825 qsa_select: prompt-path block scores with a thread per query head (below sm_80, same bits)
open · @xrip · 0 comentarios · En GitHub
Setup & installMulti-GPUAMD / HIPNVIDIA / CUDAModels & quants
Descripción
## What The prompt path's QSA block scores on cards without the TF32 scorer (below sm_80; measured on Turing, sm_75) run on a new kernel that is 4.6-5.2x faster per call, with bitwise the same scores. sm_80 and newer, HIP, and decode run the same code as before. `STRATA_SCORES_THREAD=0` turns it off. Two commits: 1. **A thread per (query, head)** (`block_scores_thread_kernel`). The old kernel scored every (query, key block) with a whole warp: 20 shuffles and five 512-byte loads per score. Now each lane keeps one 128-value query head in registers, reads each key once from shared memory as a broadcast, and adds its 32 four-value partials in the same tree that lane 0 of the old xor-shuffle sum used. Same pairs, same order, so the same bits. The heads are summed 0..3 from 0.0f as before, and the tail block keeps `block_scores_tail_kernel`. 2. **Four lanes per query.** Nsight Compute showed the first version bound by the shared-memory pipe (`mio_throttle`): one key load per 4 FMAs. The old shuffle tree splits into four quarters that only meet at the end, so now four lanes share a query, each holds one quarter of all four heads, one key load feeds 16 FMAs, and two xor shuffles finish the same tree. Still the same bits. New test: `qsa_scores_parity` checks every score of the new kernel against `block_scores_kernel` (normal values, wide magnitudes with zeros, small integers with ties; up to 262K context), then times both. ## Measured CMP 50HX (TU102, sm_75, unlocked to RTX 2080 speed), CUDA 13.3, GCC 14.2. - `qsa_scores_parity`: 141 M scores bitwise identical on both cards. - Time per call, old warp kernel -> this PR: | queries | context | before | after | |---:|---:|---:|---:| | 1,024 | 140K | 35.9 ms | 6.2 ms | | 1,024 | 200K | 52.4 ms | 8.5 ms | - Why it matters: in an nsys profile of one +4.4K-token prompt at 140K (Flash-Next IQ3_XXS, 2x CMP 50HX, `--layer-split auto`), `block_scores_kernel` was the largest prompt kernel: 1,045 ms per card, 17 % of the time. - Whole engine, same setup: a coding-agent session from 34K to 190K tokens (36 turns, +~4.4K tokens each, greedy, fresh server per run) gave word-for-word the same text as before. Time to first token at 190K went from 10.49 s to 7.13 s, but that number also includes the tokenizer change in #1823; the scorer's own share was not split out. Note: measured on a fork based on 0.1.38. The patch applies to current `main` without changes; re-measuring on 0.1.41 is still to do. Volta and Pascal (the experimental build) also take this kernel and were not tested. ## How to check On an sm_75 card: `qsa_scores_parity` (also in ctest). It prints `PASS` when every score matches, and the time of both kernels. 🤖 Generated with [Claude Code](https://claude.com/claude-code)
En el sitio
Enlaces a install, modelos, releases.